text:
\section{Adaptation for Multi-Task Learning}
Different tasks are conditionally independent from each other. Let $I(Z;Y_k)$ denote the mutual information for the downstream task $Y_k$. In the same manner as the information bottleneck method, we can formulate MTL as the optimization problem as follows
\begin{align}
\label{Eq:MIB_1}
\text{max}  & ~ \sum_{k=1}^K I(Z;Y_k),   \\
\label{Eq:MIB_2}
\text{s.t.} & ~ I(Z;X) \leq I_C.
\end{align}

Similar to Eq.~(\ref{Eq:IB_Lag}), instead of solving the above constrained optimization problem directly, we deal with the following Lagrangian objective function
\begin{align}
\label{Eq:MTIB_Lag}
\mathscr{L}_{\text{MTIB}} = \sum_{k=1}^K I(Z;Y_k) - \beta I(Z;X).
\end{align}

Direct computing the mutual information $I(Z;Y_k)$ and $I(Z;X)$ in Eq.~(\ref{Eq:MTIB_Lag}) is a challenging task. Variational inference can be employed to approximate them. Firstly, with Assumption 1, we can derive a variational lower bound for $I(Z;X)$ as follow

\begin{align}
\label{Eq:Prob_Weighted}
\log \left\{ p(y, z)\right\}  & ~ = \sum_{k=1}^K \log \left\{ p(y_k,z;\sigma_k) \right\}  \nonumber  \\
& ~ = \sum_{k=1}^K \log \left\{ p(y_k|z;\sigma_k)p(z) \right\},
\end{align}
where $\sigma_k$ is a positive scalar used to explicitly indicate the homoscedastic uncertainty of downstream task $k$.

Eq.~(\ref{Eq:MTIB_Lag}), the final loss function of the proposed MTVIB can be obtained
\begin{align}
\label{Eq:MTVIB_Loss_Final}
& \mathscr{L}_{\text{MTVIB}}  \nonumber \\
= & \frac{1}{N}\sum_{n=1}^{N}\mathop{{}\mathbb{E}_{\epsilon_z \sim p(\epsilon_z)}} \bigg [ \sum_{k=1}^{K} -\log \big\{ p_{\theta_k}\big(y_n|f_{\phi}(x_n,\epsilon_z)\big) \big\} \bigg]
\nonumber \\  & \qquad  + \beta D_{\text{KL}}(p_{\phi}\big(z|x_n)||q(z)\big) \nonumber \\
= & \sum_{k=1}^{K} \Big[\frac{1}{\sigma^2_k} \mathscr{E}_{k} + \log \{ \sigma_k \} \Big] + \beta D_{\text{KL}}(p_{\phi}\big(z|x)||q(z)\big).
\end{align}
where $N$ is the size of mini batch for Monte Carlo sampling. Similar to the VIB, in model testing, we sample the latent variable $z$ using encoder $p_{\phi}(z|x)$, then $z$ will be used as the input for decoders $p_{\theta_k}(y_k|z)$ whose outputs are our model's predictions.

\subsubsection{Results on MNIST}
\begin{table*}[ht!]
\caption{FGSM attack results on MNIST.}
\label{Tab:fgsm_mnist}
\centering
\begin{tabular}{|c|c|c|c|c|c|c|c|} \hline
 Method & Clean & $\epsilon=0.05$ & $\epsilon=0.1$ & $\epsilon=0.15$ & $\epsilon=0.2$ & $\epsilon=0.25$ & $\epsilon=0.3$ \\ \hline
 Baseline &\textbf{0.9933} &0.9339 &0.9088 &0.9024 &0.8917 &0.8767 &0.8637\\
 Distillation &0.9928 &0.9660 &0.9510 &0.9469 &0.9365 &0.9189 &0.8812 \\
 Adversarial &0.9922 &0.9856 &0.9799 &0.9735 &0.9651 &0.9531 &0.9441\\
 VIB-R &0.9918 &0.9721 & 0.9643 &0.9565 &0.9384 &0.9193 &0.8849 \\
 VIB-M &0.9922 &0.9699 &0.9627 &0.9556 &0.9406 &0.9141 &0.8948\\
 VIB-R(Ref) &0.9929 &0.9905 &0.9845 &0.9799 &0.9681 &0.9506 &0.9171 \\
 VIB-M(Ref) &0.9931 &0.9899 &0.9846 &0.9802 &0.9716 &0.9529 &0.9261 \\
 VIB-M(PGD) &0.9907 &0.9838 &0.9755 &0.9674 &0.9591 &0.9477 &0.9355\\

 VIB-M(Ref+PGD) &0.9926 &0.9900 &0.9891 &0.9876 &0.9860 &0.9845 &0.9817  \\
 VIB-M(Ref+FGSM) &0.9928 &\textbf{0.9916} &\textbf{0.9893} &\textbf{0.9890} &\textbf{0.9868} &\textbf{0.9857} &\textbf{0.9851} \\
 \hline
\end{tabular}
\end{table*}

The FGSM attack results on MNIST are exhibited in Table \ref{Tab:fgsm_mnist}.

\begin{table*}[ht!]
\caption{PGD attack results on MNIST.}
\label{Tab:pgd_mnist}
\centering
\begin{tabular}{|c|c|c|c|c|c|c|c|} \hline
 Method & Clean &  $\epsilon=0.05$ & $\epsilon=0.1$ & $\epsilon=0.15$ & $\epsilon=0.2$ & $\epsilon=0.25$ & $\epsilon=0.3$ \\ \hline
 Baseline &\textbf{0.9933} &0.9267 &0.8813 &0.7288 &0.5138 &0.3666 &0.3546 \\
 Distillation &0.9928 &0.9374 &0.7944 &0.4080 &0.0991 &0.0289 &0.0226\\
 Adversarial &0.9922 &0.9851 &0.9753 &0.9641 &0.9477 &0.9298 &0.9225  \\
 VIB-R &0.9918 &0.9256 &0.4884 & 0.0634 &0.0081 &0.0033 &0.0035   \\
 VIB-M &0.9922 &0.9311 &0.6415 &0.1363 &0.0111 &0.0027 &0.0021  \\
 VIB-R(Ref) &0.9929 &0.9501 &0.8773 &0.6405 &0.3697 &0.2285 &0.2059 \\
 VIB-M(Ref) &0.9931 &0.9827 &0.9196 &0.6127 &0.3728 &0.3309 &0.3290  \\
 VIB-M(PGD) &0.9907 &0.9826 &0.9724 &0.9566 &0.9418 &0.9239 &0.9174 \\

 VIB-M(Ref+PGD) &0.9926 &\textbf{0.9895} &\textbf{0.9858} &\textbf{0.9793} &\textbf{0.9755} &\textbf{0.9697}  &\textbf{0.9675} \\
 VIB-M(Ref+FGSM) &0.9928 &0.9857 &0.9751 &0.9606 &0.9424 &0.9225 &0.9116 \\
 \hline
\end{tabular}
\end{table*}

\begin{table*}[ht!]
\caption{PGD attack results on MNIST.}
\label{Tab:pgd_mnist}
\centering
\begin{tabular}{|c|c|c|c|c|c|c|c|} \hline
 Method & Clean &  $\epsilon=0.05$ & $\epsilon=0.1$ & $\epsilon=0.15$ & $\epsilon=0.2$ & $\epsilon=0.25$ & $\epsilon=0.3$ \\ \hline
 Baseline &\textbf{0.9933} &0.9264 &0.8707 &0.6437 &0.2541 &0.0530 &0.0174 \\
 Distillation &0.9928 &0.9376 &0.7679 &0.3021 &0.0177 &0.0004 &0.0001\\
 Adversarial &0.9922 &0.9839 &0.9742 &0.9627 &0.9448 &0.9152 &0.8913  \\
 VIB-R &0.9918 &0.9207 &0.4286 & 0.0241 &0.0004 &0.0000 &0.0000   \\
 VIB-M &0.9922 &0.9311 &0.6415 &0.1363 &0.0111 &0.0027 &0.0021  \\
 VIB-R(Ref) &0.9929 &0.9520 &0.8611 &0.5571 &0.1842 &0.0307 &0.0064 \\
 VIB-M(Ref) &0.9931 &0.9627 &0.7534 &0.3458 &0.1912 &0.1684 &0.1632 \\
 VIB-M(PGD) &0.9907 &0.9889 &0.9861 &0.9841 &0.9803 &0.9750 &0.9632 \\

 VIB-M(Ref+PGD) &0.9926 &0.9878 &0.9848 &0.9801 &0.9744 &0.9681 &0.9538 \\
 VIB-M(Ref+FGSM) &0.9928 &0.9857 &0.9740 &0.9549 &0.9313 &0.8915 &0.7926 \\
 \hline
\end{tabular}
\end{table*}

\begin{table*}[ht!]
\caption{PGD attack results on MNIST.}
\label{Tab:pgd_mnist}
\centering
\begin{tabular}{|c|c|c|c|c|c|c|c|} \hline
 Method & Clean &  $\epsilon=0.05$ & $\epsilon=0.1$ & $\epsilon=0.15$ & $\epsilon=0.2$ & $\epsilon=0.25$ & $\epsilon=0.3$ \\ \hline
 Baseline &\textbf{0.9933} &0.9267 &0.8820 &0.7282 &0.4489 &0.1938 &0.0738 \\%
 Distillation &0.9928 &0.9383 &0.7923 &0.4018 &0.0930 &0.0253 &0.0177\\ %
 Adversarial &0.9922 &0.9837 &0.9752 &0.9641 &0.9542 &0.9410 &0.9040  \\%
 VIB-R &0.9918 &0.9251 &0.4942 &0.0616 &0.0038 &0.0001 &0.0001 \\%
 VIB-M &0.9922 &0.9617 &0.7649 &0.3728 &0.2315 &0.1932 &0.1642  \\%
 VIB-R(Ref) &0.9929 &0.9508 &0.8780 &0.6298 &0.3188 &0.1061 &0.0292 \\%
 VIB-M(Ref) &0.9931 &0.9844 &0.9195 &0.6127 &0.3061 &0.2019 &0.1448  \\ %
 VIB-M(PGD) &0.9907 &0.9882 &\textbf{0.9875} &\textbf{0.9843} &\textbf{0.9818} &\textbf{0.9752} &\textbf{0.9499} \\ %
 VIB-M(Ref+PGD) &0.9926 &\textbf{0.9896} &0.9869 &0.9831 &0.9785 &0.9595 &0.9080 \\% all

 VIB-M(Ref+FGSM) &0.9928 &0.9843 &0.9754 &0.9608 &0.9406 &0.9120 &0.8298 \\
 \hline
\end{tabular}
\end{table*}

\iffalse
\begin{table}
\caption{Running time per epoch on CIFAR-10 (CPU: i7-9750H, GPU: RTX2060).}
\label{Tab:time_cifar10}
\centering
\begin{tabular}{|c|c|c|c|} \hline
 Method & Time per epoch & Total epoch & Total time\\ \hline
 Baseline &34.526s &200 &1.92h\\
 Distillation &66.351 &170 &5.05h \\
 Adversarial &502.349/273.049s &100 &9.77h \\
 Adv+Ref &349.456s & &h \\
 VIB-R &28.831 &250 &2.00h\\
 VIB-M & &\\
 VIB-R(Ref) &63.291 &300  \\
 VIB-M(Ref) &71.128 &200 \\
 VIB-M(PGD) &134.058 &300\\
 VIB-M(Ref+PGD) &376.235 &140\\
 VIB-M(Ref+FGSM) &143.916 &80 \\
 \hline
\end{tabular}
\end{table}
\fi

\iffalse
\begin{table*}[ht!]
\small
\caption{Models' classification test accuracy results ($\%$) with the FGSM attack on the MNIST dataset.}
\label{Tab:fgsm_mnist}
\centering
\begin{tabular}{|l|c|c|c|c|c|c|c|} \hline
 Method & Clean & $\epsilon=0.05$ & $\epsilon=0.1$ & $\epsilon=0.15$ & $\epsilon=0.2$ & $\epsilon=0.25$ & $\epsilon=0.3$ \\ \hline

 Baseline &$99.20_{\pm0.05}$ &$93.64_{\pm0.18}$ &$91.03_{\pm0.28}$ &$89.76_{\pm0.22}$ &$88.73_{\pm0.22}$ &$87.94_{\pm0.11}$ &$86.57_{\pm0.18}$\\%

 Baseline(KD) &$99.28_{\pm0.04}$ &$96.61_{\pm0.06}$ &$95.17_{\pm0.09}$ &$94.33_{\pm0.19}$ &$93.39_{\pm0.31}$ &$91.20_{\pm0.21}$ &$88.35_{\pm0.22}$ \\%

 Baseline(PGD) &$99.17_{\pm0.02}$ &\pmb{$98.54_{\pm0.03}$} &$97.92_{\pm0.08}$ &$97.18_{\pm0.06}$ &$96.40_{\pm0.09}$ &$95.40_{\pm0.12}$ &$94.36_{\pm0.06}$\\%

 Baseline(REF) &\pmb{$99.35_{\pm0.02}$} &$97.00_{\pm0.07}$ &$96.10_{\pm0.21}$ &$95.56_{\pm0.07}$ &$94.88_{\pm0.20}$ &$93.47_{\pm0.23}$ &$91.27_{\pm0.24}$\\%

 Baseline(PGD+REF) &\pmb{$99.35_{\pm0.03}$} &$97.00_{\pm0.11}$ &$96.02_{\pm0.12}$ &$95.58_{\pm0.14}$ &$94.75_{\pm0.26}$ &$93.52_{\pm0.38}$ &$91.35_{\pm0.18}$\\%

 VIB-R &$99.10_{\pm0.05}$ &$97.24_{\pm0.05}$ &$96.52_{\pm0.20}$ &$95.52_{\pm0.20}$ &$94.20_{\pm0.12}$ &$91.44_{\pm0.26}$ &$88.44_{\pm0.34}$ \\%

 VIB-M &$99.17_{\pm0.03}$ &$96.96_{\pm0.11}$ &$96.23_{\pm0.08}$ &$95.29_{\pm0.18}$ &$94.04_{\pm0.17}$ &$91.68_{\pm0.15}$ &$88.96_{\pm0.40}$\\%

 VIB-R(REF) &$99.24_{\pm0.03}$ &$97.24_{\pm0.07}$ &$96.58_{\pm0.10}$ &$95.45_{\pm0.12}$ &$93.35_{\pm0.27}$ &$90.45_{\pm0.33}$ &$86.33_{\pm0.50}$\\%

 VIB-M(REF) &$99.28_{\pm0.02}$ &$96.96_{\pm0.08}$ &$96.19_{\pm0.16}$ &$95.44_{\pm0.16}$ &$94.79_{\pm0.05}$ &$92.86_{\pm0.23}$ &$89.60_{\pm0.04}$\\ %

 VIB-M(PGD) &$99.04_{\pm0.05}$ &$98.37_{\pm0.04}$ &$97.61_{\pm0.08}$ &$96.79_{\pm0.08}$ &$95.86_{\pm0.13}$ &$94.73_{\pm0.06}$ &$93.41_{\pm0.10}$\\ %

 VIB-M(REF+FGSM) &$99.22_{\pm0.06}$ &$98.06_{\pm0.06}$ &$97.64_{\pm0.12}$ &$97.35_{\pm0.08}$ &$97.12_{\pm0.13}$ &$96.89_{\pm0.09}$ &$96.53_{\pm0.11}$\\

 VIB-M(REF+PGD) &$99.13_{\pm0.04}$ &$98.48_{\pm0.08}$ &\pmb{$98.03_{\pm0.04}$} &\pmb{$97.79_{\pm0.07}$} &\pmb{$97.59_{\pm0.08}$} &\pmb{$97.35_{\pm0.05}$} &\pmb{$97.12_{\pm0.12}$} \\% all

 \hline
\end{tabular}
\end{table*}

\begin{table*}[ht!]
\small
\caption{Models' classification test accuracy results ($\%$) with the PGD attack accuracy on the MNIST dataset.}
\label{Tab:pgd_mnist}
\centering
\begin{tabular}{|l|c|c|c|c|c|c|c|} \hline
 Method & Clean &  $\epsilon=0.05$ & $\epsilon=0.1$ & $\epsilon=0.15$ & $\epsilon=0.2$ & $\epsilon=0.25$ & $\epsilon=0.3$ \\ \hline

 Baseline &$99.20_{\pm0.05}$ &$92.60_{\pm0.06}$ &$87.36_{\pm0.23}$ &$64.54_{\pm0.39}$ &$22.88_{\pm0.14}$ &$3.09_{\pm0.13}$ &$0.36_{\pm0.05}$ \\%

 Baseline(KD) &$99.28_{\pm0.04}$ &$93.71_{\pm0.10}$ &$76.74_{\pm0.13}$ &$30.88_{\pm0.12}$ &$1.98_{\pm0.08}$ &$0.01_{\pm0.01}$ &$0.00_{\pm0.00}$ \\%

 Baseline(PGD) &$99.17_{\pm0.02}$ &\pmb{$98.44_{\pm0.07}$} &$97.51_{\pm0.05}$ &$96.47_{\pm0.01}$ &$94.68_{\pm0.09}$ &$92.57_{\pm0.07}$ &$88.34_{\pm0.09}$  \\%

 Baseline(REF) &\pmb{$99.35_{\pm0.02}$} &$95.45_{\pm0.05}$ &$88.99_{\pm0.18}$ &$60.79_{\pm0.17}$ &$15.14_{\pm0.11}$ &$1.39_{\pm0.06}$ &$0.14_{\pm0.01}$\\%

 Baseline(PGD+REF) &\pmb{$99.35_{\pm0.03}$} &$95.39_{\pm0.17}$ &$89.98_{\pm0.13}$ &$66.31_{\pm0.24}$ &$22.75_{\pm0.21}$ &$2.58_{\pm0.14}$ &$0.36_{\pm0.04}$ \\%

 VIB-R &$99.10_{\pm0.05}$ &$92.06_{\pm0.19}$ &$43.28_{\pm0.24}$ &$2.55_{\pm0.08}$ &$0.03_{\pm0.02}$ &$0.00_{\pm0.00}$ &$0.00_{\pm0.00}$ \\ %

 VIB-M &$99.17_{\pm0.03}$ &$93.39_{\pm0.08}$ &$61.25_{\pm0.19}$ &$8.81_{\pm0.11}$ &$0.03_{\pm0.01}$ &$0.00_{\pm0.00}$ &$0.00_{\pm0.00}$ \\%

VIB-R(REF) &$99.24_{\pm0.03}$ &$95.19_{\pm0.07}$ &$86.33_{\pm0.24}$ &$55.69_{\pm0.50}$ &$17.51_{\pm0.23}$ &$2.15_{\pm0.10}$ &$0.27_{\pm0.03}$ \\ %

VIB-M(REF) &$99.28_{\pm0.02}$ &$95.17_{\pm0.06}$ &$86.28_{\pm0.23}$ &$53.82_{\pm0.32}$ &$13.53_{\pm0.05}$ &$1.64_{\pm0.09}$ &$0.19_{\pm0.02}$ \\%

VIB-M(PGD) &$99.04_{\pm0.05}$ &$98.26_{\pm0.02}$ &$97.19_{\pm0.06}$ &$95.84_{\pm0.06}$ &$94.03_{\pm0.13}$ &$90.86_{\pm0.14}$ &$85.51_{\pm0.18}$ \\ %

VIB-M(REF+FGSM) &$99.22_{\pm0.06}$ &$97.41_{\pm0.12}$ &$96.79_{\pm0.07}$ &$95.99_{\pm0.10}$ &$95.03_{\pm0.09}$ &$92.99_{\pm0.27}$ &$86.95_{\pm0.16}$ \\%

VIB-M(REF+PGD) &$99.13_{\pm0.04}$ &$98.24_{\pm0.06}$ &\pmb{$97.70_{\pm0.08}$} &\pmb{$97.36_{\pm0.08}$} &\pmb{$96.87_{\pm0.09}$} &\pmb{$96.01_{\pm0.06}$} &\pmb{$93.64_{\pm0.30}$} \\% all

 \hline
\end{tabular}
\end{table*}

\begin{table*}[ht!]
\small
\caption{Models' classification test accuracy results ($\%$) with the FGSM attack accuracy on the CIFAR-10 dataset.}
\label{Tab:fgsm_cifar10}
\centering
\begin{tabular}{ |l|c|c|c|c|c|c|c|} \hline
 Method & Clean & $\epsilon=0.01$ & $\epsilon=0.02$ & $\epsilon=0.03$ & $\epsilon=0.04$ & $\epsilon=0.05$ & $\epsilon=0.07$ \\ \hline

 Baseline &$92.88_{\pm0.13}$ &$86.36_{\pm0.14}$ &$83.67_{\pm0.18}$ &$81.68_{\pm0.23}$ &$79.99_{\pm0.18}$ &$78.52_{\pm0.14}$ &$75.28_{\pm0.18}$ \\%

 Baseline(KD) &\pmb{$93.41_{\pm0.06}$} &$84.34_{\pm0.24}$ &$80.42_{\pm0.20}$ &$77.58_{\pm0.29}$ &$75.44_{\pm0.30}$ &$73.56_{\pm0.19}$ &$70.03_{\pm0.36}$ \\%

 Baseline(PGD) &$90.24_{\pm0.19}$ &$84.71_{\pm0.22}$ &$79.36_{\pm0.31}$ &$75.49_{\pm0.21}$ &$72.72_{\pm0.11}$ &$70.48_{\pm0.17}$ &$67.78_{\pm0.21}$ \\%

 Baseline(REF) &$90.51_{\pm0.22}$ &$84.71_{\pm0.34}$ &$81.81_{\pm0.21}$ &$79.52_{\pm0.17}$ &$77.31_{\pm0.26}$ &$75.15_{\pm0.29}$ &$71.96_{\pm0.16}$\\  %

 Baseline(PGD+REF) &$87.11_{\pm0.08}$ &$83.05_{\pm0.21}$ &$79.29_{\pm0.04}$ &$76.40_{\pm0.17}$ &$73.91_{\pm0.22}$ &$71.77_{\pm0.14}$ &$68.13_{\pm0.35}$ \\%

 VIB-R &$90.57_{\pm0.07}$ &$84.66_{\pm0.32}$ &$81.24_{\pm0.14}$ &$78.66_{\pm0.27}$ &$76.47_{\pm0.19}$ &$74.84_{\pm0.25}$ &$71.72_{\pm0.27}$ \\ %

 VIB-M &$91.47_{\pm0.14}$ &$84.79_{\pm0.24}$ &$81.94_{\pm0.15}$ &$79.98_{\pm0.22}$ &$78.16_{\pm0.13}$ &$76.52_{\pm0.11}$ &$73.87_{\pm0.26}$ \\%

 VIB-R(REF) &$90.98_{\pm0.07}$ &$86.82_{\pm0.16}$ &$84.80_{\pm0.23}$ &$83.09_{\pm0.19}$ &$81.62_{\pm0.20}$ &$80.37_{\pm0.20}$ &$77.97_{\pm0.25}$ \\ %

 VIB-M(REF) &$91.49_{\pm0.05}$ &\pmb{$87.42_{\pm0.17}$} &\pmb{$85.32_{\pm0.13}$} &\pmb{$84.01_{\pm0.32}$} &\pmb{$82.42_{\pm0.26}$} &\pmb{$80.97_{\pm0.24}$} &\pmb{$78.84_{\pm0.20}$} \\ %

 VIB-M(PGD) &$74.55_{\pm0.03}$ &$69.43_{\pm0.26}$ &$63.74_{\pm0.26}$ &$58.88_{\pm0.28}$ &$53.87_{\pm0.31}$ &$49.67_{\pm0.28}$ &$42.96_{\pm0.20}$ \\ %

 VIB-M(REF+FGSM) &$87.29_{\pm0.09}$ &$83.62_{\pm0.19}$ &$81.56_{\pm0.18}$ &$80.11_{\pm0.09}$ &$78.61_{\pm0.16}$ &$77.48_{\pm0.18}$ &$75.20_{\pm0.23}$ \\%

 VIB-M(REF+PGD) &$89.57_{\pm0.16}$ &$86.24_{\pm0.13}$ &$83.69_{\pm0.18}$ &$82.06_{\pm0.17}$ &$80.23_{\pm0.36}$ &$78.78_{\pm0.30}$ &$76.44_{\pm0.20}$ \\%

\hline
\end{tabular}
\end{table*}

\begin{table*}[ht!]
\small
\caption{Models' classification test accuracy results ($\%$) with the PGD attack accuracy on the CIFAR-10 dataset.}
\label{Tab:pgd_cifar10}
\centering
\begin{tabular}{ |l|c|c|c|c|c|c|c|} \hline
 Method & Clean & $\epsilon=0.01$ & $\epsilon=0.02$ & $\epsilon=0.03$ & $\epsilon=0.04$ & $\epsilon=0.05$ & $\epsilon=0.07$ \\ \hline

 Baseline &$92.88_{\pm0.13}$ &$77.28_{\pm0.18}$ &$63.96_{\pm0.26}$ &$46.42_{\pm0.28}$ &$27.18_{\pm0.25}$ &$15.49_{\pm0.38}$ &$5.58_{\pm0.13}$ \\%

 Baseline(KD) &\pmb{$93.41_{\pm0.06}$} &$76.23_{\pm0.13}$ &$51.26_{\pm0.09}$ &$32.97_{\pm0.21}$ &$24.64_{\pm0.19}$ &$20.23_{\pm0.23}$ &$14.93_{\pm0.23}$ \\ %

 Baseline(PGD) &$90.24_{\pm0.19}$ &$84.00_{\pm0.11}$ &$76.02_{\pm0.23}$ &$67.08_{\pm0.15}$ &$58.13_{\pm0.22}$ &$50.37_{\pm0.23}$ &$41.04_{\pm0.16}$ \\%

 Baseline(REF) &$90.51_{\pm0.22}$ &$71.88_{\pm0.32}$ &$57.66_{\pm0.40}$ &$41.69_{\pm0.26}$ &$28.29_{\pm0.17}$ &$19.80_{\pm0.22}$&$11.74_{\pm0.14}$\\%

 Baseline(PGD+REF) &$87.11_{\pm0.08}$ &$82.56_{\pm0.11}$ &$77.21_{\pm0.16}$ &$71.65_{\pm0.21}$ &$65.17_{\pm0.36}$ &$58.65_{\pm0.24}$ &$46.25_{\pm0.18}$ \\%

 VIB-R &$90.57_{\pm0.07}$ &$79.06_{\pm0.34}$ &$63.62_{\pm0.25}$ &$48.99_{\pm0.33}$ &$36.59_{\pm0.13}$ &$26.93_{\pm0.21}$ &$14.39_{\pm0.25}$ \\ %

 VIB-M &$91.47_{\pm0.14}$ &$79.61_{\pm0.15}$ &$68.18_{\pm0.29}$ &$62.47_{\pm0.41}$ &$57.98_{\pm0.10}$ &$52.23_{\pm0.16}$ &$40.10_{\pm0.33}$\\ %

 VIB-R(REF) &$90.98_{\pm0.07}$ &$79.87_{\pm0.06}$ &$69.35_{\pm0.30}$ &$62.58_{\pm0.23}$ &$56.57_{\pm0.15}$ &$51.90_{\pm0.44}$ &$44.82_{\pm0.34}$\\%

 VIB-M(REF) &$91.49_{\pm0.05}$ &$80.55_{\pm0.19}$ &$70.52_{\pm0.18}$ &$60.05_{\pm0.24}$ &$50.94_{\pm0.20}$ &$44.01_{\pm0.09}$ &$35.02_{\pm0.18}$ \\%

 VIB-M(PGD) &$74.55_{\pm0.03}$ &$68.74_{\pm0.12}$ &$62.92_{\pm0.27}$ &$56.32_{\pm0.37}$ &$50.29_{\pm0.19}$ &$44.46_{\pm0.32}$ &$33.91_{\pm0.22}$ \\%

 VIB-M(REF+FGSM) &$87.29_{\pm0.09}$ &$82.47_{\pm0.11}$ &$77.22_{\pm0.18}$ &$71.54_{\pm0.34}$ &$66.11_{\pm0.33}$ &$60.35_{\pm0.18}$ &$50.94_{\pm0.29}$ \\%

 VIB-M(REF+PGD) &$89.57_{\pm0.16}$ &\pmb{$85.65_{\pm0.08}$} &\pmb{$80.92_{\pm0.13}$} &\pmb{$75.78_{\pm0.20}$} &\pmb{$69.86_{\pm0.20}$} &\pmb{$63.89_{\pm0.20}$} &\pmb{$54.00_{\pm0.19}$} \\ %

 \hline
\end{tabular}
\end{table*}

\fi
\section{Introduction}

This short example shows a contrived example on how to format the authors' information for {\it IJCAI--PRICAI--20 Proceedings}.

\section{Author names}

Each author name must be followed by:
\begin{itemize}
    \item A newline {\tt \textbackslash{}\textbackslash{}} command for the last author.
    \item An {\tt \textbackslash{}And} command for the second to last author.
    \item An {\tt \textbackslash{}and} command for the other authors.
\end{itemize}

\section{Affiliations}

After all authors, start the affiliations section by using the {\tt \textbackslash{}affiliations} command.
Each affiliation must be terminated by a newline {\tt \textbackslash{}\textbackslash{}} command. Make sure that you include the newline on the last affiliation too.

\section{Mapping authors to affiliations}

If some scenarios, the affiliation of each author is clear without any further indication (\emph{e.g.}, all authors share the same affiliation, all authors have a single and different affiliation). In these situations you don't need to do anything special.

In more complex scenarios you will have to clearly indicate the affiliation(s) for each author. This is done by using numeric math superscripts {\tt \$\{\^{}$i,j, \ldots$\}\$}. You must use numbers, not symbols, because those are reserved for footnotes in this section (should you need them). Check the authors definition in this example for reference.

\section{Emails}

This section is optional, and can be omitted entirely if you prefer. If you want to include e-mails, you should either include all authors' e-mails or just the contact author(s)' ones.

Start the e-mails section with the {\tt \textbackslash{}emails} command. After that, write all emails you want to include separated by a comma and a space, following the same order used for the authors (\emph{i.e.}, the first e-mail should correspond to the first author, the second e-mail to the second author and so on).

You may ``contract" consecutive e-mails on the same domain as shown in this example (write the users' part within curly brackets, followed by the domain name). Only e-mails of the exact same domain may be contracted. For instance, you cannot contract ``person@example.com" and ``other@test.example.com" because the domains are different.

\end{document}

\section{Introduction}

The {\it IJCAI--20 Proceedings} will be printed from electronic
manuscripts submitted by the authors. These must be PDF ({\em Portable
Document Format}) files formatted for 8-1/2$''$ $\times$ 11$''$ paper.

\subsection{Length of Papers}

All paper {\em submissions} must have a maximum of six pages, plus at most one for references. The seventh page cannot contain {\bf anything} other than references.

The length rules may change for final camera-ready versions of accepted papers and will differ between tracks. Some tracks may include only references in the last page, whereas others allow for any content in all pages. Similarly, some tracks allow you to buy a few extra pages should you want to, whereas others don't.

If your paper is accepted, please carefully read the notifications you receive, and check the proceedings submission information website\footnote{\url{https://proceedings.ijcai.org/info}} to know how many pages you can finally use. That website holds the most up-to-date information regarding paper length limits at all times. Please notice that if your track allows for a special references-only page, the {\bf references-only page(s) cannot contain anything else than references} (i.e.: do not write your acknowledgments on that page or you will be charged for it).

\subsection{Word Processing Software}

As detailed below, IJCAI has prepared and made available a set of
\LaTeX{} macros and a Microsoft Word template for use in formatting
your paper. If you are using some other word processing software, please follow the format instructions given below and ensure that your final paper looks as much like this sample as possible.

\section{Style and Format}

\LaTeX{} and Word style files that implement these instructions
can be retrieved electronically. (See Appendix~\ref{stylefiles} for
instructions on how to obtain these files.)

\subsection{Layout}

Print manuscripts two columns to a page, in the manner in which these
instructions are printed. The exact dimensions for pages are:
\begin{itemize}
\item left and right margins: .75$''$
\item column width: 3.375$''$
\item gap between columns: .25$''$
\item top margin---first page: 1.375$''$
\item top margin---other pages: .75$''$
\item bottom margin: 1.25$''$
\item column height---first page: 6.625$''$
\item column height---other pages: 9$''$
\end{itemize}

All measurements assume an 8-1/2$''$ $\times$ 11$''$ page size. For
A4-size paper, use the given top and left margins, column width,
height, and gap, and modify the bottom and right margins as necessary.

\subsection{Format of Electronic Manuscript}

For the production of the electronic manuscript, you must use Adobe's
{\em Portable Document Format} (PDF). A PDF file can be generated, for
instance, on Unix systems using {\tt ps2pdf} or on Windows systems
using Adobe's Distiller. There is also a website with free software
and conversion services: \url{http://www.ps2pdf.com}. For reasons of
uniformity, use of Adobe's {\em Times Roman} font is strongly suggested.
In \LaTeX2e{} this is accomplished by writing
\begin{quote}
\mbox{\tt $\backslash$usepackage\{times\}}
\end{quote}
in the preamble.\footnote{You may want also to use the package {\tt
latexsym}, which defines all symbols known from the old \LaTeX{}
version.}

Additionally, it is of utmost importance to specify the {\bf
letter} format (corresponding to 8-1/2$''$ $\times$ 11$''$) when
formatting the paper. When working with {\tt dvips}, for instance, one
should specify {\tt -t letter}.

\subsection{Title and Author Information}

Center the title on the entire width of the page in a 14-point bold
font. The title must be capitalized using Title Case. Below it, center author name(s) in 12-point bold font. On the following line(s) place the affiliations, each affiliation on its own line using 12-point regular font. Matching between authors and affiliations can be done using numeric superindices. Optionally, a comma-separated list of email addresses follows the affiliation(s) line(s), using 12-point regular font.

\subsubsection{Blind Review}

In order to make blind reviewing possible, authors must omit their
names and affiliations when submitting the paper for review. In place
of names and affiliations, provide a list of content areas. When
referring to one's own work, use the third person rather than the
first person. For example, say, ``Previously,
Gottlob~\shortcite{gottlob:nonmon} has shown that\ldots'', rather
than, ``In our previous work~\cite{gottlob:nonmon}, we have shown
that\ldots'' Try to avoid including any information in the body of the
paper or references that would identify the authors or their
institutions. Such information can be added to the final camera-ready
version for publication.

\subsection{Abstract}

Place the abstract at the beginning of the first column 3$''$ from the
top of the page, unless that does not leave enough room for the title
and author information. Use a slightly smaller width than in the body
of the paper. Head the abstract with ``Abstract'' centered above the
body of the abstract in a 12-point bold font. The body of the abstract
should be in the same font as the body of the paper.

The abstract should be a concise, one-paragraph summary describing the
general thesis and conclusion of your paper. A reader should be able
to learn the purpose of the paper and the reason for its importance
from the abstract. The abstract should be no more than 200 words long.

\subsection{Text}

The main body of the text immediately follows the abstract. Use
10-point type in a clear, readable font with 1-point leading (10 on
11).

Indent when starting a new paragraph, except after major headings.

\subsection{Headings and Sections}

When necessary, headings should be used to separate major sections of
your paper. (These instructions use many headings to demonstrate their
appearance; your paper should have fewer headings.). All headings should be capitalized using Title Case.

\subsubsection{Section Headings}

Print section headings in 12-point bold type in the style shown in
these instructions. Leave a blank space of approximately 10 points
above and 4 points below section headings.  Number sections with
arabic numerals.

\subsubsection{Subsection Headings}

Print subsection headings in 11-point bold type. Leave a blank space
of approximately 8 points above and 3 points below subsection
headings. Number subsections with the section number and the
subsection number (in arabic numerals) separated by a
period.

\subsubsection{Subsubsection Headings}

Print subsubsection headings in 10-point bold type. Leave a blank
space of approximately 6 points above subsubsection headings. Do not
number subsubsections.

\paragraph{Titled paragraphs.} You should use titled paragraphs if and
only if the title covers exactly one paragraph. Such paragraphs should be
separated from the preceding content by at least 3pt, and no more than
6pt. The title should be in 10pt bold font and ended with a period.
After that, a 1em horizontal space should follow the title before
the paragraph's text.

In \LaTeX{} titled paragraphs should be typeset using
\begin{quote}
{\tt \textbackslash{}paragraph\{Title.\} text} .
\end{quote}

\subsubsection{Acknowledgements}

You may include an unnumbered acknowledgments section, including
acknowledgments of help from colleagues, financial support, and
permission to publish. If present, acknowledgements must be in a dedicated,
unnumbered section appearing after all regular sections but before any
appendices or references.

Use
\begin{quote}
    {\tt \textbackslash{}section*\{Acknowledgements\}})
\end{quote}
to typeset the acknowledgements section in \LaTeX{}.

\subsubsection{Appendices}

Any appendices directly follow the text and look like sections, except
that they are numbered with capital letters instead of arabic
numerals. See this document for an example.

\subsubsection{References}

The references section is headed ``References'', printed in the same
style as a section heading but without a number. A sample list of
references is given at the end of these instructions. Use a consistent
format for references. The reference list should not include unpublished
work.

\subsection{Citations}

Citations within the text should include the author's last name and
the year of publication, for example~\cite{gottlob:nonmon}.  Append
lowercase letters to the year in cases of ambiguity.  Treat multiple
authors as in the following examples:~\cite{abelson-et-al:scheme}
or~\cite{bgf:Lixto} (for more than two authors) and
\cite{brachman-schmolze:kl-one} (for two authors).  If the author
portion of a citation is obvious, omit it, e.g.,
Nebel~\shortcite{nebel:jair-2000}.  Collapse multiple citations as
follows:~\cite{gls:hypertrees,levesque:functional-foundations}.
\nocite{abelson-et-al:scheme}
\nocite{bgf:Lixto}
\nocite{brachman-schmolze:kl-one}
\nocite{gottlob:nonmon}
\nocite{gls:hypertrees}
\nocite{levesque:functional-foundations}
\nocite{levesque:belief}
\nocite{nebel:jair-2000}

\subsection{Footnotes}

Place footnotes at the bottom of the page in a 9-point font.  Refer to
them with superscript numbers.\footnote{This is how your footnotes
should appear.} Separate them from the text by a short
line.\footnote{Note the line separating these footnotes from the
text.} Avoid footnotes as much as possible; they interrupt the flow of
the text.

\section{Illustrations}

Place all illustrations (figures, drawings, tables, and photographs)
throughout the paper at the places where they are first discussed,
rather than at the end of the paper.

They should be floated to the top (preferred) or bottom of the page,
unless they are an integral part
of your narrative flow. When placed at the bottom or top of
a page, illustrations may run across both columns, but not when they
appear inline.

Illustrations must be rendered electronically or scanned and placed
directly in your document. They should be cropped outside latex, otherwise portions of the image could reappear during the post-processing of your paper. All illustrations should be understandable when printed in black and
white, albeit you can use colors to enhance them. Line weights should
be 1/2-point or thicker. Avoid screens and superimposing type on
patterns, as these effects may not reproduce well.

Number illustrations sequentially. Use references of the following
form: Figure 1, Table 2, etc. Place illustration numbers and captions
under illustrations. Leave a margin of 1/4-inch around the area
covered by the illustration and caption.  Use 9-point type for
captions, labels, and other text in illustrations. Captions should always appear below the illustration.

\section{Tables}

Tables are considered illustrations containing data. Therefore, they should also appear floated to the top (preferably) or bottom of the page, and with the captions below them.

\begin{table}
\centering
\begin{tabular}{lll}
\hline
Scenario  & $\delta$ & Runtime \\
\hline
Paris       & 0.1s  & 13.65ms     \\
Paris       & 0.2s  & 0.01ms      \\
New York    & 0.1s  & 92.50ms     \\
Singapore   & 0.1s  & 33.33ms     \\
Singapore   & 0.2s  & 23.01ms     \\
\hline
\end{tabular}
\caption{Latex default table}
\label{tab:plain}
\end{table}

\begin{table}
\centering
\begin{tabular}{lrr}
\toprule
Scenario  & $\delta$ (s) & Runtime (ms) \\
\midrule
Paris       & 0.1  & 13.65      \\
            & 0.2  & 0.01       \\
New York    & 0.1  & 92.50      \\
Singapore   & 0.1  & 33.33      \\
            & 0.2  & 23.01      \\
\bottomrule
\end{tabular}
\caption{Booktabs table}
\label{tab:booktabs}
\end{table}

If you are using \LaTeX, you should use the {\tt booktabs} package, because it produces better tables than the standard ones. Compare Tables \ref{tab:plain} and~\ref{tab:booktabs}. The latter is clearly more readable for three reasons:

\begin{enumerate}
    \item The styling is better thanks to using the {\tt booktabs} rulers instead of the default ones.
    \item Numeric columns are right-aligned, making it easier to compare the numbers. Make sure to also right-align the corresponding headers, and to use the same precision for all numbers.
    \item We avoid unnecessary repetition, both between lines (no need to repeat the scenario name in this case) as well as in the content (units can be shown in the column header).
\end{enumerate}

\section{Formulas}

IJCAI's two-column format makes it difficult to typeset long formulas. A usual temptation is to reduce the size of the formula by using the {\tt small} or {\tt tiny} sizes. This doesn't work correctly with the current \LaTeX{} versions, breaking the line spacing of the preceding paragraphs and title, as well as the equation number sizes. The following equation demonstrates the effects (notice that this entire paragraph looks badly formatted):
\begin{tiny}
\begin{equation}
    x = \prod_{i=1}^n \sum_{j=1}^n j_i + \prod_{i=1}^n \sum_{j=1}^n i_j + \prod_{i=1}^n \sum_{j=1}^n j_i + \prod_{i=1}^n \sum_{j=1}^n i_j + \prod_{i=1}^n \sum_{j=1}^n j_i
\end{equation}
\end{tiny}%

Reducing formula sizes this way is strictly forbidden. We {\bf strongly} recommend authors to split formulas in multiple lines when they don't fit in a single line. This is the easiest approach to typeset those formulas and provides the most readable output%
\begin{align}
    x =& \prod_{i=1}^n \sum_{j=1}^n j_i + \prod_{i=1}^n \sum_{j=1}^n i_j + \prod_{i=1}^n \sum_{j=1}^n j_i + \prod_{i=1}^n \sum_{j=1}^n i_j + \nonumber\\
    + & \prod_{i=1}^n \sum_{j=1}^n j_i
\end{align}%

If a line is just slightly longer than the column width, you may use the {\tt resizebox} environment on that equation. The result looks better and doesn't interfere with the paragraph's line spacing: %
\begin{equation}
\resizebox{.91\linewidth}{!}{$
    \displaystyle
    x = \prod_{i=1}^n \sum_{j=1}^n j_i + \prod_{i=1}^n \sum_{j=1}^n i_j + \prod_{i=1}^n \sum_{j=1}^n j_i + \prod_{i=1}^n \sum_{j=1}^n i_j + \prod_{i=1}^n \sum_{j=1}^n j_i
$}
\end{equation}%

This last solution may have to be adapted if you use different equation environments, but it can generally be made to work. Please notice that in any case:

\begin{itemize}
    \item Equation numbers must be in the same font and size than the main text (10pt).
    \item Your formula's main symbols should not be smaller than {\small small} text (9pt).
\end{itemize}

For instance, the formula
\begin{equation}
    \resizebox{.91\linewidth}{!}{$
    \displaystyle
    x = \prod_{i=1}^n \sum_{j=1}^n j_i + \prod_{i=1}^n \sum_{j=1}^n i_j + \prod_{i=1}^n \sum_{j=1}^n j_i + \prod_{i=1}^n \sum_{j=1}^n i_j + \prod_{i=1}^n \sum_{j=1}^n j_i + \prod_{i=1}^n \sum_{j=1}^n i_j
$}
\end{equation}
would not be acceptable because the text is too small.

\section{Examples, Definitions, Theorems and Similar}

Examples, definitions, theorems, corollaries and similar must be written in their own paragraph. The paragraph must be separated by at least 2pt and no more than 5pt from the preceding and succeeding paragraphs. They must begin with the kind of item written in 10pt bold font followed by their number (e.g.: Theorem 1), optionally followed by a title/summary between parentheses in non-bold font and ended with a period. After that the main body of the item follows, written in 10 pt italics font (see below for examples).

In \LaTeX{} We strongly recommend you to define environments for your examples, definitions, propositions, lemmas, corollaries and similar. This can be done in your \LaTeX{} preamble using \texttt{\textbackslash{newtheorem}} -- see the source of this document for examples. Numbering for these items must be global, not per-section (e.g.: Theorem 1 instead of Theorem 6.1).

\begin{example}[How to write an example]
Examples should be written using the example environment defined in this template.
\end{example}

\begin{theorem}
This is an example of an untitled theorem.
\end{theorem}

You may also include a title or description using these environments as shown in the following theorem.

\begin{theorem}[A titled theorem]
This is an example of a titled theorem.
\end{theorem}

\section{Proofs}

Proofs must be written in their own paragraph separated by at least 2pt and no more than 5pt from the preceding and succeeding paragraphs. Proof paragraphs should start with the keyword ``Proof." in 10pt italics font. After that the proof follows in regular 10pt font. At the end of the proof, an unfilled square symbol (qed) marks the end of the proof.

In \LaTeX{} proofs should be typeset using the \texttt{\textbackslash{proof}} environment.

\begin{proof}
This paragraph is an example of how a proof looks like using the \texttt{\textbackslash{proof}} environment.
\end{proof}

\section{Algorithms and Listings}

Algorithms and listings are a special kind of figures. Like all illustrations, they should appear floated to the top (preferably) or bottom of the page. However, their caption should appear in the header, left-justified and enclosed between horizontal lines, as shown in Algorithm~\ref{alg:algorithm}. The algorithm body should be terminated with another horizontal line. It is up to the authors to decide whether to show line numbers or not, how to format comments, etc.

In \LaTeX{} algorithms may be typeset using the {\tt algorithm} and {\tt algorithmic} packages, but you can also use one of the many other packages for the task.

\begin{algorithm}[tb]
\caption{Example algorithm}
\label{alg:algorithm}
\textbf{Input}: Your algorithm's input\\
\textbf{Parameter}: Optional list of parameters\\
\textbf{Output}: Your algorithm's output
\begin{algorithmic}[1]
\STATE Let $t=0$.
\WHILE{condition}
\STATE Do some action.
\IF {conditional}
\STATE Perform task A.
\ELSE
\STATE Perform task B.
\ENDIF
\ENDWHILE
\STATE \textbf{return} solution
\end{algorithmic}
\end{algorithm}

\section*{Acknowledgments}

The preparation of these instructions and the \LaTeX{} and Bib\TeX{}
files that implement them was supported by Schlumberger Palo Alto
Research, AT\&T Bell Laboratories, and Morgan Kaufmann Publishers.
Preparation of the Microsoft Word file was supported by IJCAI.  An
early version of this document was created by Shirley Jowell and Peter
F. Patel-Schneider.  It was subsequently modified by Jennifer
Ballentine and Thomas Dean, Bernhard Nebel, Daniel Pagenstecher,
Kurt Steinkraus, Toby Walsh and Carles Sierra. The current version
has been prepared by Marc Pujol-Gonzalez and Francisco Cruz-Mencia.

\section{Introduction}

From the information-theoretical view of deep learning, the internal representation of some hidden layers is considered as a stochastic encoding $Z$ of the input $X$, and the goal of learning is to obtain an encoding that is maximally informative about the target $Y$. Therefore, information bottleneck (IB) approaches \cite{tishby2000information} force $Z$ to be a minimal latent representation of $X$ for predicting $Y$.
The
IB
methods have recently attracted significant attention in their ability of learning good hidden representations. See \cite{goldfeld2020information} for a recent survey. The
IB methods
can also
be used in sparse neural networks using Dropout \cite{achille2018information} and to compress neural networks \cite{dai2018compressing}. \cite{peng2019variational} developed a Variational Discriminator Bottleneck which can be applied to many applications, such as imitation learning, inverse reinforcement learning, and generative adversarial networks.

Variational information bottleneck (VIB), started from \cite{alemi2017deep}, leverages variational inference to implement the IB method \cite{tishby2000information,tishby2015deep}. VIB models (or stochastic neural networks fit using VIB method), are believed to be more robust than deterministic models (e.g., conventional deep neural networks) because VIB learns the disentangled latent representations $Z$ which ignore the details of the input $X$ as many as possible \cite{alemi2017deep}. This has been improved by \cite{federici2020learning} with a multi-view approach where each view represents a task.
Specifically, the robustness of learned representations can be improved by eliminating task-related information. However, additional information about multiple tasks is required.
Moreover, \cite{belghazi2018mutual} leverages Mutual Information Neural Estimation to
help with the implementation of VIB, without concerning  robustness performance. \cite{goyal2020variational} proposes Variational Bandwidth Bottleneck which compresses only part of the input, and shows that it has better generalization for reinforcement learning tasks.

In this paper, we propose a novel method that learns robust VIB models that is radically different from the above methods (i.e., without taking the multi-view approach, is able to work with convolutional networks, and is robust against more involved attacks). The main idea is to use soft labels as the learning targets. The soft labels are pre-learned from a reference network. Specifically, they are the numerical outputs of the last hidden layer of the reference network but not the logits of the penultimate layer as in knowledge distillation \cite{hinton2015distilling,papernot2016distillation}. By using soft labels as training targets, the VIB training has been transformed from the original classification problem into a maximum likelihood estimation based regression problem so the outputs share the same manifold. Moreover, \cite{vincent2008extracting,vincent2010stacked} showed that an autoencoder structure trained with perturbed images as the input and the clean images as the target can be used to reconstruct \lq\lq{}repaired\rq\rq{} images from perturbed ones. Our method is inspired by this method to some extent, but with completely different purposes.

There are two major advantages of soft labels. On the one hand, they would lessen the confidence of the model and prevent it from overfitting~\cite{pmlr-v97-cohen19c}. On the other hand, they improve the efficiency of model training. It is widely acknowledged that model training with the single-step adversarial examples can easily lead to the risk of \emph{label leaking}~\cite{kurakin2016adversarial}, that is, the model's performance on adversarial examples is better than that with the clean examples. Therefore, many other studies employ iterative adversarial examples but the model training process becomes computationally expensive. Using the soft labels pre-learned from the reference network can help circumvent the risk of label leaking with the single-step adversarial examples, whose generation process such as FGSM \cite{goodfellow2014explaining} is much simpler and quicker.

Moreover, in the VIB, we replace the reparameterization method \cite{kingma2014auto} with the mutual information neural estimation (in short MINE) \cite{belghazi2018mutual}. The latter does not rely on Gaussian assumption for the posterior so it also improves the robustness of VIB. To validate our proposed method, extensive experiments are performed on MNIST and CIFAR-10 datasets. We examine both FGSM and PGD attacks at different scales, and compare the proposed method with several widely used methods such as adversarial training, knowledge distillation. The experimental results verify the robustness improvement of our method.

The overall contributions of this paper are summarized as follows:

\begin{itemize}
    \item We propose to use the soft labels learned by a reference network to cast all inputs onto the same manifold;
    improve the robustness of VIB models;

    \item We consider adversarial training for VIB to fully materialise the capability of pre-learned soft labels in guiding the learning of smooth VIB models. Experimentally, we find that this method is able to defend both FGSM attack and PGD attack \cite{}  effectively;
    We combine Adversarial Training with pre-learned soft labels to further the performance of VIB, and this method is able to defend both FGSM and PGD attack effectively;

    \item We utilize one-step adversarial samples to accelerate the training process without suffering the ``label leaking effect'' \cite{kurakin2016adversarial}.
    which says the testing accuracy on adversarial images are higher than the accuracy on clear images.
    \item Extensive experiments show the superior performance of our method over the state-of-the-art.
\end{itemize}

\section{Related Work}

\subsection{Adversarial Attacks}

\textbf{FGSM Attack.}
The fast gradient sign method (FGSM)~\cite{goodfellow2014explaining} is a simple gradient-based method used to produce adversarial samples, which can be described as follow:
\begin{align}
\label{Eq:FGSM}
x^\text{Adv} = x + \epsilon \cdot \text{sign}\big(\nabla_x J(w;x;y)\big)
\end{align}
where $J$ is the loss function, $w$ is the model parameter and $\epsilon $ is a positive scalar which indicates the magnitude of the perturbation.

\textbf{PGD Attack.}
FGSM is a simple one-step adversarial method, which is not powerful enough in some cases. Hence, Kurakin et al. \cite{kurakin2016adversarial} proposed a more powerful method to generate adversarial samples, called projected gradient descent (PGD). PGD is similar to FGSM, while it attacks clean samples recursively, thus it is able to produce more complex adversarial samples. Note that in this work, all the PGD attacks are $L^{\infty}$ norm. And PGD can be described as follows:
\begin{align}
\label{Eq:PGD}
x^\text{Adv}_0 & = x; \nonumber  \\
x^\text{Adv}_{t+1} & = \text{Clip}_{x, \epsilon} \Big \{x^\text{Adv}_t + \alpha \cdot \text{sign}\big(\nabla_x J(w;x^\text{Adv}_t ;y)\big) \Big \}
\end{align}

where $\alpha$ is the attack step size and $t$ is the attack step number.

\subsection{Variational Information Bottleneck}

Thanks to the ability of compressing and extracting information from input data, some previous works have shown the potential of variational information bottleneck to defend adversarial attacks.

For instance, \cite{alemi2017deep} tested the standard variational information bottleneck model on FSGM attacks \cite{goodfellow2014explaining}  and proved its effectiveness.

More recently, \cite{9088132} considers the robustness of VIB by adapting the variational bound of \cite{alemi2017deep}, but this method is only implemented via fully-connected layers.

However, the main weakness of the methods proposed in \cite{alemi2017deep} and \cite{9088132} is that their robustness against stronger attacks, e.g., projected gradient descent (PGD) \cite{kurakin2016adversarial}, is not verified.

\iffalse
DNNs are
susceptible to adversarial attacks \cite{goodfellow2014explaining}, \cite{madry2018towards}, \cite{carlini2017towards}. In order to enable DNNs
to be
resilient to adversarial attacks, researchers have PROP many approaches. For instance, \cite{kurakin2016adversarial} PROP to use adversarial samples, i.e. PGD samples, as the training input and this method is called Adversarial Training. \cite{papernot2016distillation} used Knowledge Distillation to acquire the soft labels as the training targets to improve the robustness of deep neural networks. Another approach PROP by \cite{raghunathan2018certified} is based on Adversarial Training and regularization.

\fi

\section{Model}

Let $X$ be the input, $Z$ the latent representation of $X$, and $Y$ the target. Without loss of generality, we assume an information Markov chain $X \rightarrow Z \rightarrow Y$. The information bottleneck~\cite{tishby2000information,tishby2015deep} solves the following optimisation problem:
\begin{align}
\label{Eq:IB_1}
\text{max}  & ~ I(Z;Y),   \\
\label{Eq:IB_2}
\text{s.t.} & ~ I(Z;X) \leq I_C,
\end{align}
where $I$ is the mutual information and $I_C$ is a constant representing the information constraint. This is equivalent to maximizing the Lagrangian objective function
\begin{align}
\label{Eq:IB_Lag}
\mathscr{L}_{\text{IB}} = I(Z;Y) - \beta \big( I(Z;X) - I_C \big).
\end{align}
Based on Eq.~(\ref{Eq:IB_Lag}), the VIB model~\cite{alemi2017deep} computes $I(Z;X)$ via an encoder neural network and computes $I(Z;Y)$ via a classifier neural network, respectively. Our study aims to enhance the VIB model's robustness in face of adversarial attacks such as the projected gradient
descent (PGD) attack \cite{kurakin2016adversarial} and the fast gradient sign method (FGSM) attack \cite{goodfellow2014explaining}.

\begin{figure}[t]
\centering
\includegraphics[width=1\linewidth]{robust_vib.pdf}
\caption{Overview of our PROP robust learning method for the VIB, where $X$ is the input, $Y$ is the output, $f_E$ is the encoder network, $XZ$ (red) represents the joint samples, $XZ$ (yellow) represents the independent samples, $f_D$ is the discriminator, $f_R$ is the reference network, $f_C$ is the classifier, and $o_{\text{REF}}$ and $o_{\text{VIB}}$ are the soft labels from the reference network and the classifier, respectively.}
\label{fig:framework}
\end{figure}

The overview of our proposed method is presented in Fig.\ref{fig:framework}. Our model has a reference network $f_R$, an encoder $f_E$, a discriminator $f_D$ and a classifier $f_C$. It consists of three stages. In Stage I, a reference network $f_{R}(x)$ is pre-trained with clean samples $x$ to learn the soft labels $o_{\text{REF}}$, which are the output of the last hidden layer. The Stage II aims to compute the mutual information loss $I(Z;X)$ using the Mutual Information Neural Estimation (MINE) method~\cite{belghazi2018mutual}. The input $x'$ is firstly perturbed with either random noise (normal training) or adversarial attacks (adversarial training) to train an encoder network $f_E$. Then, the joint samples $XZ_\text{Joint}$ obtained by the encoder $f_E$ are shuffled to obtain the independent samples $XZ_\text{Ind}$. Afterwards, the discriminator $f_D$ is used to discriminate $XZ_\text{Joint}$ from $XZ_\text{Ind}$. The purpose of using $XZ_\text{Ind}$ and $f_D$ is due to our adoption of the Donsker-Varadhan representation of the Kullback–Leibler (KL) divergence. In Stage III, the learned soft labels are leveraged to compute the mean squared error loss $||o_{\text{VIB}}-o_{\text{REF}}||^2$ as the approximation of $I(Z;Y)$, where $o_{\text{VIB}}$ is the output of the last hidden layer of the VIB classifier $f_C$.

We jointly train the encoder $f_E$, the discriminator $f_D$ and the classifier $f_C$ with the loss function $\mathscr{L}_{\text{PROP}}$. First, for a set of inputs $x_i$, we define the following loss function for the VIB model with reference:
\begin{align}
\label{Eq:RVIB_loss}
\mathscr{L} & =
\sum_{i}^{N} \bigg[ \Big\| f_R(x_i)- f_C \big (f_E(x'_i) \big) \Big\|^2  \nonumber  \\
& \qquad  + \beta \Big( \mathop{{}\mathbb{E}}{_{P_{XZ}}}\big [f_D \big] - \log \mathop{\mathbb{E}}{_{_{P_{X}\otimes P_{Z}}}\big[e^{f_D} \big]} \Big) \bigg ]
\end{align}
where $N$ is the size of mini batch, and $\beta$ is a non-negative Lagrangian multiplier determining the optimisation trade-off. Based on Eq. (\ref{Eq:RVIB_loss}), let $\mathscr{L}_{\text{Clean}}(x_i^{\text{Clean}})$ and $\mathscr{L}_{\text{Adv}}(x_i^{\text{Adv}})$ be the loss functions trained on clean samples and adversarial samples, respectively. Then, we can have the loss function $\mathscr{L}_{\text{PROP}}$ for the adversarial training:
\begin{align}
\label{Eq:Final_loss}
\mathscr{L}_{\text{PROP}} & = \frac{\sum_{i \in \text{Clean}}{\mathscr{L}_{\text{Clean}}(x_i^{\text{Clean}})}  +\lambda \sum_{i \in \text{Adv} }{\mathscr{L}_{\text{Adv}}(x_i^{\text{Adv}})}}{m-k+\lambda k},
\end{align}
where $m$ is the total batch size, $k$ is the adversarial sample size, $\lambda$ is the relative weight. We call our new VIB method REF-VIB\ to highlight the distinct feature of pre-learned soft labels. Algorithm~\ref{Alg:Algorithm} presents the overall training procedure. In the following, we explain the components of our REF-VIB\ method with details.

\begin{algorithm}[tb]
\caption{REF-VIB training}
\label{Alg:Algorithm}
\textbf{Requirements}: $x$ (images) , $y$ (labels)

\begin{algorithmic}[1]
\WHILE{ $\text{epoch} \leq \text{max epoch}$}
    \STATE{train a reference neural network $f_R(x)$ with clean samples;}
\ENDWHILE

\WHILE{ $\text{epoch} \leq \text{max epoch}$}

    \STATE{If $i$ $\in$ Clean:}
        \STATE{\quad $x'_i$ $\gets$ $x_i$ + $\epsilon$}
    \STATE{If $i$ $\in$ Adversarial:}
        \STATE{\quad $x'_i$ $\gets$ attack $x_i$}

        \STATE{$xz_{\text{Joint}}$ $\gets$ $f_E(x'_i)$}

        \STATE{$xz_{\text{Ind}}$ $\gets$ shuffle $xz_{\text{Joint}}$ along the batch axis}

        \STATE{$I(Z;X)$ $\gets$ $\mathop{{}\mathbb{E}}{_{P_{XZ}}}[f_D(xz_{\text{Joint}})]$ \\
        \qquad \qquad \qquad $ - \log\{ \mathop{\mathbb{E}}{_{_{P_{X}\otimes P_{Z}}}[e^{f_D(xz_{\text{Ind}})}]} \}$}

        \STATE{$I(Z;Y)$ $\gets$ $\|f_C(xz_{\text{Joint}})-f_R(x_i)\|^2$}

        \STATE{Minimize final loss function $\mathscr{L}_{\text{PROP}}$}

\ENDWHILE

\end{algorithmic}
\end{algorithm}

The VIB model~\cite{alemi2017deep} implements Eq. (\ref{Eq:IB_Lag}) with reparameterization method \cite{kingma2014auto} via deep neural networks.
It approximates $I(Z;X)$ and $I(Z;Y)$ with an encoder network and a decoder network, respectively,
by deriving the variational bounds for them.
The upper bound of $I(Z;X)$ yields
\begin{align}
\label{Eq:I_ZX_bound}
I(Z;X) \leq \mathbb{E}_x \left[D_\text{KL} \big(p(z|x)||q(z)\big)\right],
\end{align}
where $D_\text{KL}$ is the Kullback-Leibler (KL) divergence, $q(z)$ is an uninformative prior distribution.
Similarly, the lower bound of $I(Z;Y)$ can be computed by
\begin{align}
\label{Eq:I_ZY_bound}
I(Z;Y) \geq & \mathbb{E}_{p(y,z)} \big [ \log \left\{ p(y|z) \right\} \big].
\end{align}
Combining Eq. (\ref{Eq:I_ZX_bound}) and Eq. (\ref{Eq:I_ZY_bound}), the loss function becomes
\begin{align}
\label{Eq:OP_VIB_Repa}
 \mathscr{L}_{\text{VIB}}
 = &  \
 \mathbb{E}_{\mathscr{D}} \bigg[ \mathbb{E}_{z \sim p(z|x)} \Big[-\log\{ p(y|z)\} \Big] \nonumber \\
 & \hspace*{50pt} + \beta D_\text{KL}\big(p(z|x)|| q(z)\big ) \bigg],
\end{align}
where $p(z|x)$ is the encoder and $p(y|z)$ is the classifier.

In the following, we introduce how do we update Eqs. (\ref{Eq:I_ZX_bound})-(\ref{Eq:I_ZY_bound}) by taking the MINE method and the reference network, respectively. This will lead to the update of $\mathscr{L}_{\text{VIB}}$ into $\mathscr{L}$ of Eq. (\ref{Eq:RVIB_loss}), and VIB into the proposed REF-VIB\ method.

\paragraph{Mutual Information Neural Estimation}

As for estimating the mutual information $I(X;Z)$, apart from using the re-parameterization method, we can alternatively use the Mutual Information Neural Estimation (MINE) method~\cite{belghazi2018mutual}. By using the MINE method, the Gaussian assumption on the posterior $p(z|x)$ can be relaxed, so the approximation of the true posterior $p(z|x)$ can be more accurate.

First, we introduce the Donsker-Varadhan representation of the KL divergence \cite{donsker1983asymptotic,belghazi2018mutual}, defined as below
\begin{align}
\label{Eq:DV}
D_\text{KL}(P||Q) = \underset{T:\Omega \rightarrow \mathop{{}\mathbb{R}}}{\mathrm{sup}} \mathop{{}\mathbb{E}}{_{P}}\big[T\big] - \log \mathop{\mathbb{E}}{_{Q}\big[e^T\big]},
\end{align}
where $P$ is the posterior distribution, $Q$ is the prior distribution and $T$ is a discriminator network used to distinguish joint samples from independent samples. This can be implemented via a discriminator network using both joint and independent samples, as presented in the Stage II of Fig.~\ref{fig:framework}.

The reason why we consider the KL divergence is because $I(X;Z)$ can be defined as follow:
\begin{align}
\label{Eq:I_ZX}
I(Z;X) = & \ D_\text{KL} \big (\mathop{{}\mathbb{{P_{XZ}}}}||\mathop{{}\mathbb{{P_{X}}}} \otimes \mathop{{}\mathbb{{P_{Z}}}} \big )
\end{align}
where $\mathbb{{P_{XZ}}}$ is the joint distribution, $\mathbb{{P_{X}}}$ and $\mathbb{{P_{Z}}}$ are the marginal distributions
Therefore, as opposed to the variational upper bound in Eq. (\ref{Eq:I_ZX_bound}), we use bound $I(Z;X) \geq I_\theta(Z;X)$ to estimate $I(Z;X)$ according to Eq. (\ref{Eq:DV}), such that
\begin{equation}
\label{Eq:I_ZX_DV}
I_\theta(Z;X) = \underset{\theta \in \Theta}{\mathrm{sup}} \mathop{{}\mathbb{E}}{_{P_{XZ}}}\big[T_\theta\big] - \log \mathop{\mathbb{E}}{_{_{P_{X}\otimes P_{Z}}}\big[e^{T_\theta}\big]},
\end{equation}
where $T_\theta$ is a family of functions parametrized by a discriminator network (i.e., $f_D$ shown in Fig.~\ref{fig:framework}) with parameters $\theta \in \Theta$.  Hereafter, we omit $\theta$ for simplicity. To solve Eq. (\ref{Eq:I_ZX_DV}), we use the mini batches of the input to feed an encoder network so as to acquire the samples of the joint distribution $P_{XZ}$ (positive samples). While, the samples of $P_{X}\otimes P_{Z}$ can be obtained through shuffling the samples of $P_{XZ}$ along the batch axis (negative samples).

To simplify the discussion, we use \emph{VIB-M} to denote the VIB model trained using the MINE method and \emph{VIB-R} to denote the VIB model trained through the re-parameterization method. Combining Eq. (\ref{Eq:I_ZY_bound}) and Eq. (\ref{Eq:I_ZX_DV}), the loss function of the VIB-M model can be obtained as follows:
\begin{align}
\label{Eq:OP_VIB_Mine}
& \mathscr{L}_{\text{VIB-M}} =  \mathbb{E}_{\mathscr{f_D}} \bigg[ \mathbb{E}_{z \sim p(z|x)} \big[-\log\{ p(y|z)\} \big]  \nonumber \\
& \hspace*{45pt} + \beta \Big (\mathop{{}\mathbb{E}}{_{P_{XZ}}}\big[f_D\big] - \log \mathop{\mathbb{E}}{_{_{P_{X}\otimes P_{Z}}}\big[e^{f_D}\big]} \Big ) \bigg].
\end{align}

It should be noted that $\mathscr{L}_{\text{VIB-M}}$ is different from $\mathscr{L}$ in Eq. (\ref{Eq:RVIB_loss}) because of the term $\mathbb{E}_{z \sim p(z|x)} \big[-\log\{ p(y|z)\} \big]  \nonumber$, which will be replaced with $\|f_R(x_i)- f_C \big (f_E(x'_i) \big)\|^2 $ after using the reference network, discussed next.

\paragraph{Reference Network}
\label{sec:referencenetwork}

Another key novelty of our method is that we use a reference network to pre-learn the soft labels on clean samples, and then use the pre-learned  soft labels -- instead of original hard labels -- as the learning targets in the VIB model. These soft labels can be regarded as a set of numerical vectors corresponding to each input image in space. Note that the soft labels used in our method are s different from the labels used in Knowledge Distillation \cite{hinton2015distilling,mishra2017apprentice}, which are logits.

The reason of doing this is that hard labels only retain the discrete categorical information of each input but discard the more detailed information, which may impede the learning performance of VIB. Moreover, with the guidance of the soft labels learned on clean samples, the predictions of VIB model on perturbed samples are projected onto the correct manifold during the training process.

Now, according to Eq. (\ref{Eq:I_ZY_bound}), we use the pre-learned soft labels to compute the likelihood $p(y|z)$ to approximate the lower bound of $I(Z;Y)$. To this end, we let $f_R(x)$ be the output of the last hidden layer of the reference network and $f_C(xz_{\text{Joint}})=f_C \big (f_E(x'_i) \big)$ be the output of the last hidden layer of the classifier in the VIB. The distance between $f_R(\cdot)$ and $f_C(\cdot)$ can be measured by the Euclidean distance:
\begin{equation}
\label{Eq:Ref_loss}
\text{DIST}(o_{\text{REF}}, o_{\text{VIB}}) = \Big\|f_R(x_i)- f_C \big (f_E(x'_i) \big) \Big\|^2
\end{equation}

Then, Eq. (\ref{Eq:Ref_loss}) is the loss function of the downstream task in the VIB, and by replacing $\mathbb{E}_{z \sim p(z|x)} \big[-\log\{ p(y|z)\} \big]  \nonumber$ of Eq. (\ref{Eq:OP_VIB_Mine}) with it, we obtain $\mathscr{L}$ of Eq. (\ref{Eq:RVIB_loss}).

\paragraph{Adversarial Training with Reference}

Deep neural networks (DNNs) are known to be vulnerable to adversarial attacks. Even small perturbations on the clean images may significantly decrease the test performance of DNNs. The VIB model has been argued to be more robust than conventional neural networks \cite{alemi2017deep}.
Nevertheless, as shown in Section~\ref{sec:experiments}, the VIB models are still not robust enough to withstand strong adversarial attacks, such as PGD attacks. See the comparison between the VIB-M (REF) and the VIB-M (REF+FGSM) models. Therefore, we consider another effective way of defending attacks, i.e., adversarial training, which uses adversarial samples as the training input.

In this paper, we consider two popular gradient-based methods to generate adversarial samples, i.e., the FGSM attack~\cite{goodfellow2014explaining} and the PGD attack~\cite{kurakin2016adversarial}. Simply speaking, the FGSM attack is a single-step attack method and the PGD attack finds adversarial samples in an iterative way.

Let $x$ denote the clean samples and $x_\text{Adv}$ the adversarial samples. We can obtain the adversarial samples as follows:
\begin{align}
\label{Eq:Adv}
x^\text{Adv} = x + \delta
\end{align}
where $\delta$ is the perturbation. For any $x \in {X}$ with $l_p$-ball $\mathcal{B}(x, \epsilon)$ around $x: \{ x_\text{Adv} \in X: ||x_\text{Adv} - x_\text{Clean} || \leq \epsilon \}$ \cite{madaan2020adversarial}, in a VIB model, the related latent presentation $z$ is sampled from the input $x$ via an encoder and the $\hat{y}$ is corresponding prediction computed from $z$ via a classifier. This relation can be described as follow:
\begin{align}
\label{Eq:Prediction}
\hat{y} =  f_C(z) =  f_C(f_E(x))
\end{align}
where $f_E(\cdot)$ is the encoder and $f_C(\cdot)$ is the classifier.

According to Eq. (\ref{Eq:Ref_loss}), to train the VIB classifier $f_C$ with adversarial samples, we minimize the following loss function:
\begin{align}
\label{Eq:Adv_ref}
    & \ \mathbb{E}_{\{\mathscr{D}, \delta\}} \Big[ \Big\|f_C(xz^\text{Adv}_{\text{Joint}})-f_R(x)\Big\|^2 \Big] \nonumber \\
=  & \ \mathbb{E}_{\{\mathscr{D}, \delta\}} \Big[ \Big\|f_C(f_E(x^\text{Adv}))-f_R(x)\Big\|^2 \Big],
\end{align}
where $xz^\text{Adv}_{\text{Joint}}$ is the joint latent variables learned by adversarial samples. We substitute the term $\mathbb{E}_{z \sim p(z|x)} \big[-\log\{ p(y|z)\} \big]$ in Eq. (\ref{Eq:OP_VIB_Mine}) with Eq. (\ref{Eq:Adv_ref}), then according to Eq. (\ref{Eq:RVIB_loss}), we have the adversarial training loss function, $\mathscr{L}_{\text{Adv}}$ in Eq. (\ref{Eq:Final_loss}) which is the final loss function of the PROP model. Eq. (\ref{Eq:Adv_ref}) suggests that we can learn the encoder and classifier jointly for the VIB model, combining adversarial training and the reference network. Consequently, the distortion of latent variables caused by adversarial attack can be suppressed as well, which leads to better generalization so as to defend attacks.

The disadvantage of using PGD samples for training is its time consumption. One-step adversarial images, e.g., FGSM samples, are much faster to generate, however using such adversarial images to train conventional neural networks with hard labels will cause the \textbf{label leaking} effect \cite{kurakin2016adversarial}. Label leaking means the testing accuracy on adversarial images are much higher than the accuracy on clear images. This issue occurs because one-step adversarial images generated by simple transformations, and it is not hard for a deep neural network to learn such transformations, as a result, the deep neural network performs better on the adversarial images than the clean images. On the contrary, a more complex adversarial transformation, e.g., PGD, does not has this issue because its transformation is more difficult to learn. Fortunately, our method accelerates the learning process by leveraging FGSM samples without suffering from label leaking effect, as we use soft labels instead of hard labels.

\paragraph{Discussion}

We would like to discuss the rationale behind our selection of the three components in the proposed model.

\emph{Why Using the MINE Method?} We want to compute the posterior $p(z|x)$ as accurate as possible. Since the prior $q(z)$ is assumed to follow Gaussian, the reparameterization method assumes $p(z|x)$ follows Gaussian as well. However, the posterior can be more complex than the prior, and using a Gaussian distribution may not be accurate enough to approximate the prior. In the MINE method, we do not generate the posterior samples from a Gaussian distribution, instead, we generate the joint posterior samples with input perturbed by noise and use a discriminator to discriminate them from the shuffled samples. Once the model is trained, we can obtain the posterior samples with more complex shapes.

\emph{Why Need Reference Network?} In a VIB model, the input is compressed to a latent distribution by the encoder. We believe that the vulnerability of a VIB is caused by the distortion of the latent distribution. Therefore, to improve the robustness VIB, we need to minimise the distortion (a similar idea was PROP in \cite{madaan2020adversarial}, but their method is based on Bayesian neural networks and pruning). The latent distribution of VIB is Gaussian-like, and the pre-learned soft labels are numerical vectors which contain more information than the original hard labels. Hence, to suppress the distortion of the latent distribution, using the soft labels to minimize the loss function of VIB is more effusive than using the hard labels.

\emph{Why Performing Adversarial Training?} Adversarial training is known for improving the robustness of neural networks. The soft labels pre-learned by the reference network on the clean samples and the soft labels learned by the VIB model on the corrupted samples should share the same manifold. Consequently, we can leverage the pre-learned soft labels to guide the adversarial learning process of the VIB model to achieve a better performance. Our learning strategy is similar to denoising autoencoders \cite{vincent2008extracting,vincent2010stacked} to some extent. In denoising autoencoders, the corrupted images are used to train an autoencoder in order to recover
clean images, which is an unsupervised learning scheme. Our method is supervised, and we use adversarial samples as the corrupted samples and soft labels as the target instead of original clean images.

\begin{table}[t]
\centering
\small
\caption{Summary of the examined models.}
\label{Tab:model_summary}
\centering
\begin{tabular}{|m{1.1in}|m{1.9in}|}
\hline
Model & Description\\
\hline
Baseline & For the MNIST dataset, the baseline model architecture is CONV(64,3,1) + CONV(64,3,1) + FC(512) + FC(512) + FC(10); For the CIFAR-10 dataset, the baseline model is VGG16 \cite{simonyan2014very}. \\%
\hline
Baseline (KD) & The baseline model trained with KD \cite{papernot2016distillation}.\\%
\hline
Baseline (PGD) & The baseline model trained with PGD samples \cite{kurakin2016adversarial}. \\%
\hline
Baseline (REF) & The baseline model trained with a reference network.\\%
\hline
Baseline (PGD+REF) & The baseline model trained with trained with a reference network and PGD samples.\\%
\hline
VIB-R & The VIB model trained using the re-parameterization method~\cite{alemi2017deep}.\\%
\hline
VIB-M & The VIB model trained using the MINE method~\cite{belghazi2018mutual}.\\%
\hline
VIB-R (REF) & The VIB-R model trained with a reference network.\\%
\hline
VIB-M (REF) & The VIB-M model trained with a reference network.\\ %
\hline
VIB-M (PGD) & The VIB-M model trained with PGD samples.\\ %
\hline
VIB-M (REF+FGSM) & The VIB-M model trained with a reference network and FGSM samples.\\
\hline
VIB-M (REF+PGD) &  The VIB-M model trained with a reference network and PGD samples.\\% all
\hline
\end{tabular}
\end{table}

\begin{table}[t]
\small
\centering
\caption{Model training settings, where $T$ is the temperature, $\epsilon$ is the magnitude of perturbation, $\alpha$ is the attack step size, $t$ is the attack step number, and $\eta$ is the dimension of latent variables.}
\label{Tab:train_method}
\begin{tabular}{|m{1.0in}|m{2in}|}
\hline
Model & Description\\
\hline
KD & $T=1$.\\%
\hline
PGD (MNIST) & $\alpha = 0.01$, $m=100$, $k=50$, $\lambda = 0.3$, $\epsilon = 0.3$, $t = 20$. \\%
\hline
FGSM (MNIST) & $\epsilon = 0.3$. \\
\hline
PGD (CIFAR-10) & $\alpha = 0.007$, $m=100$, $k=50$, $\lambda = 0.3$, $\epsilon = 0.03$, $t = 10$. \\%
\hline
FGSM (CIFAR-10) & $\epsilon = 0.03$. \\
\hline
VIB-R & $\beta=0.001$, $\eta = 256$.\\%
\hline
VIB-M & $\beta=0.001$, $\eta = 256$.\\%
\hline
\end{tabular}
\vspace*{10pt}
\caption{Model testing settings, where $\epsilon$ is the magnitude of perturbation, $\alpha$ is the attack step size, $t$ is the attack step number.}
\label{Tab:test_method}
\centering
\begin{tabular}{|m{1.0in}|m{2in}|}
\hline
Model & Description\\
\hline
FGSM (MNIST) &  $\epsilon = 0.05$, $0.10$, $0.15$, $0.20$, $0.25$, $0.30$. \\
\hline
PGD (MNIST) & $\epsilon = 0.05$, $0.10$, $0.15$, $0.20$, $0.25$, $0.30$, $\alpha = 0.01$, $t$ = $40$, random restart. \\%
\hline
FGSM (CIFAR-10) &  $\epsilon =0.01$, $0.02$, $0.03$, $0.04$, $0.05$, $0.06$.\\
\hline
PGD (CIFAR-10) & $\epsilon=0.01$, $0.02$, $0.03$, $0.04$, $0.05$, $0.06$, $\alpha = 0.007$, $t$ = $40$, random restart.  \\%
\hline
\end{tabular}
\end{table}

\section{Experiments}
\label{sec:experiments}

\paragraph{Datasets}

Two publicly available datasets are used in the experiments: (i) the MNIST dataset~\cite{lecun1998mnist} which includes $60, 000$ gray level handwritten images of size $28 \times 28$; and (ii) the CIFAR-10 dataset~\cite{krizhevsky2009learning} which includes $60, 000$ color images of size $32 \times 32$. For each dataset, we randomly sample $50, 000$ images for model training and $10, 000$ images for model testing. The sampled images from the CIFAR-10 dataset are normalized into $[0,1]^{32\times 32\times 3}$.

\paragraph{Experimental Settings}

To have a comprehensive comparison, different machine learning algorithms using various training methods are examined, including the recent advances like knowledge distillation (KD)~\cite{papernot2016distillation}, PGD attack \cite{kurakin2016adversarial} and FGSM attack \cite{goodfellow2014explaining}. Also, as the VIB model can be estimated by either the reparameterization method or the MINE method, both training methods are included in our experiments. To simplify the discussion, an algorithm using a specific training method is considered as a model here and it is denoted by a unique name. For the reader's convenience, we present all the examined models and their descriptions in Table~\ref{Tab:model_summary}.

We adopt different settings of model architecture, training and testing. First, we use different neural networks as the baseline model and we have different architecture settings for the VIB model accordingly. For the MNIST dataset, the architecture of the baseline model is CONV(64,3,1) + CONV(64,3,1) + FC(512) + FC(512) + FC(10) and the architecture of the VIB model is CONV(32,3,1) + CONV(32,3,1) + FC(512) + FC(256) for encoder and FC(512) + FC(10) for classifier. For the CIFAR-10 dataset, the baseline models is VGG16 \cite{simonyan2014very} and the VIB model's architecture is CONV(64,3,1) + CONV(128,3,1) + CONV(256,3,1) + CONV(256,3,1) + CONV(512,3,1) + CONV(512,3,1) + FC(1024) + FC(512) + FC(256) for encoder and FC(512) + FC(1024) + FC(512) + FC(10) for classifier. Second, the baseline model and the VIB model are trained differently in these two datasets. Table~\ref{Tab:train_method} presents the details of model training settings. For the MNIST dataset, the optimizer is Adam \cite{kingma2014adam} and the learning rate is $0.001$. For the CIFAR-10 dataset, the optimizer is the stochastic gradient descent (SGD) with adaptive learning rate. Third, the trained models are evaluated under different attack settings in our datasets. As shown in Table~\ref{Tab:test_method}, the attack perturbation levels for the MNIST dataset range from 0.05 to 0.30 and are with the step size of 0.05 while they range from 0.01 to 0.06 and are with the step size of 0.01 for the CIFAR-10 dataset~\cite{kurakin2016adversarial}.

\paragraph{Results}

We run 5 randomised trials for each examined model on the MNIST and the CIFAR-10 datasets, and report the models' classification results (i.e., the average accuracy and its standard deviation) of the test data in Figs.~\ref{fig:mnist_fgsm}-\ref{fig:cifar10_pgd}.
Every curve represents a trained model, whose accuracy is gradually decreased with the increase of the parameter $\epsilon$ of either FGSM or PGD attack. A larger area under the curve (AUC) represents a more robust model.

\iffalse
\begin{figure}[t]
\centering
\includegraphics[width=1\linewidth]{mnist_fgsm}
\vspace*{-12pt}
\caption{Models' classification test accuracy results ($\%$) with the FGSM attacks on the MNIST dataset.}
\label{fig:mnist_fgsm}
\vspace*{10pt}
\centering
\includegraphics[width=1\linewidth]{mnist_pgd}
\vspace*{-12pt}
\caption{Models' classification test accuracy results ($\%$) with the PGD attacks on the MNIST dataset.}
\label{fig:mnist_pgd}
\end{figure}

\begin{figure}[t]
\centering
\includegraphics[width=1\linewidth]{cifar10_fgsm}
\vspace*{-12pt}
\caption{Models' classification test accuracy results ($\%$) with the FGSM attacks on the CIFAR-10 dataset.}
\label{fig:cifar10_fgsm}
\vspace*{10pt}
\includegraphics[width=1\linewidth]{cifar10_pgd}
\vspace*{-12pt}
\caption{Models' classification test accuracy results ($\%$) with the PGD attacks on the CIFAR-10 dataset. }
\label{fig:cifar10_pgd}
\end{figure}

\fi

\begin{table*}[ht!]
\small
\caption{Models' classification test accuracy results ($\%$) with the FGSM attack on the MNIST dataset.}
\label{Tab:fgsm_mnist}
\centering
\begin{tabular}{|l|c|c|c|c|c|c|c|} \hline
 Method & Clean & $\epsilon=0.05$ & $\epsilon=0.1$ & $\epsilon=0.15$ & $\epsilon=0.2$ & $\epsilon=0.25$ & $\epsilon=0.3$ \\ \hline

 Baseline &$99.20_{\pm0.05}$ &$93.64_{\pm0.18}$ &$91.03_{\pm0.28}$ &$89.76_{\pm0.22}$ &$88.73_{\pm0.22}$ &$87.94_{\pm0.11}$ &$86.57_{\pm0.18}$\\%

 Baseline(KD) &$99.28_{\pm0.04}$ &$96.61_{\pm0.06}$ &$95.17_{\pm0.09}$ &$94.33_{\pm0.19}$ &$93.39_{\pm0.31}$ &$91.20_{\pm0.21}$ &$88.35_{\pm0.22}$ \\%

 Baseline(PGD) &$99.17_{\pm0.02}$ &\pmb{$98.54_{\pm0.03}$} &$97.92_{\pm0.08}$ &$97.18_{\pm0.06}$ &$96.40_{\pm0.09}$ &$95.40_{\pm0.12}$ &$94.36_{\pm0.06}$\\%

 Baseline(REF) &\pmb{$99.35_{\pm0.02}$} &$97.00_{\pm0.07}$ &$96.10_{\pm0.21}$ &$95.56_{\pm0.07}$ &$94.88_{\pm0.20}$ &$93.47_{\pm0.23}$ &$91.27_{\pm0.24}$\\%

 Baseline(PGD+REF) &\pmb{$99.35_{\pm0.03}$} &$97.00_{\pm0.11}$ &$96.02_{\pm0.12}$ &$95.58_{\pm0.14}$ &$94.75_{\pm0.26}$ &$93.52_{\pm0.38}$ &$91.35_{\pm0.18}$\\%

 VIB-R &$99.10_{\pm0.05}$ &$97.24_{\pm0.05}$ &$96.52_{\pm0.20}$ &$95.52_{\pm0.20}$ &$94.20_{\pm0.12}$ &$91.44_{\pm0.26}$ &$88.44_{\pm0.34}$ \\%

 VIB-M &$99.17_{\pm0.03}$ &$96.96_{\pm0.11}$ &$96.23_{\pm0.08}$ &$95.29_{\pm0.18}$ &$94.04_{\pm0.17}$ &$91.68_{\pm0.15}$ &$88.96_{\pm0.40}$\\%

 VIB-R(REF) &$99.24_{\pm0.03}$ &$97.24_{\pm0.07}$ &$96.58_{\pm0.10}$ &$95.45_{\pm0.12}$ &$93.35_{\pm0.27}$ &$90.45_{\pm0.33}$ &$86.33_{\pm0.50}$\\%

 VIB-M(REF) &$99.28_{\pm0.02}$ &$96.96_{\pm0.08}$ &$96.19_{\pm0.16}$ &$95.44_{\pm0.16}$ &$94.79_{\pm0.05}$ &$92.86_{\pm0.23}$ &$89.60_{\pm0.04}$\\ %

 VIB-M(PGD) &$99.04_{\pm0.05}$ &$98.37_{\pm0.04}$ &$97.61_{\pm0.08}$ &$96.79_{\pm0.08}$ &$95.86_{\pm0.13}$ &$94.73_{\pm0.06}$ &$93.41_{\pm0.10}$\\ %

 VIB-M(REF+FGSM) &$99.22_{\pm0.06}$ &$98.06_{\pm0.06}$ &$97.64_{\pm0.12}$ &$97.35_{\pm0.08}$ &$97.12_{\pm0.13}$ &$96.89_{\pm0.09}$ &$96.53_{\pm0.11}$\\

 VIB-M(REF+PGD) &$99.13_{\pm0.04}$ &$98.48_{\pm0.08}$ &\pmb{$98.03_{\pm0.04}$} &\pmb{$97.79_{\pm0.07}$} &\pmb{$97.59_{\pm0.08}$} &\pmb{$97.35_{\pm0.05}$} &\pmb{$97.12_{\pm0.12}$} \\% all

 \hline
\end{tabular}
\end{table*}

\begin{table*}[ht!]
\small
\caption{Models' classification test accuracy results ($\%$) with the PGD attack accuracy on the MNIST dataset.}
\label{Tab:pgd_mnist}
\centering
\begin{tabular}{|l|c|c|c|c|c|c|c|} \hline
 Method & Clean &  $\epsilon=0.05$ & $\epsilon=0.1$ & $\epsilon=0.15$ & $\epsilon=0.2$ & $\epsilon=0.25$ & $\epsilon=0.3$ \\ \hline

 Baseline &$99.20_{\pm0.05}$ &$92.60_{\pm0.06}$ &$87.36_{\pm0.23}$ &$64.54_{\pm0.39}$ &$22.88_{\pm0.14}$ &$3.09_{\pm0.13}$ &$0.36_{\pm0.05}$ \\%

 Baseline(KD) &$99.28_{\pm0.04}$ &$93.71_{\pm0.10}$ &$76.74_{\pm0.13}$ &$30.88_{\pm0.12}$ &$1.98_{\pm0.08}$ &$0.01_{\pm0.01}$ &$0.00_{\pm0.00}$ \\%

 Baseline(PGD) &$99.17_{\pm0.02}$ &\pmb{$98.44_{\pm0.07}$} &$97.51_{\pm0.05}$ &$96.47_{\pm0.01}$ &$94.68_{\pm0.09}$ &$92.57_{\pm0.07}$ &$88.34_{\pm0.09}$  \\%

 Baseline(REF) &\pmb{$99.35_{\pm0.02}$} &$95.45_{\pm0.05}$ &$88.99_{\pm0.18}$ &$60.79_{\pm0.17}$ &$15.14_{\pm0.11}$ &$1.39_{\pm0.06}$ &$0.14_{\pm0.01}$\\%

 Baseline(PGD+REF) &\pmb{$99.35_{\pm0.03}$} &$95.39_{\pm0.17}$ &$89.98_{\pm0.13}$ &$66.31_{\pm0.24}$ &$22.75_{\pm0.21}$ &$2.58_{\pm0.14}$ &$0.36_{\pm0.04}$ \\%

 VIB-R &$99.10_{\pm0.05}$ &$92.06_{\pm0.19}$ &$43.28_{\pm0.24}$ &$2.55_{\pm0.08}$ &$0.03_{\pm0.02}$ &$0.00_{\pm0.00}$ &$0.00_{\pm0.00}$ \\ %

 VIB-M &$99.17_{\pm0.03}$ &$93.39_{\pm0.08}$ &$61.25_{\pm0.19}$ &$8.81_{\pm0.11}$ &$0.03_{\pm0.01}$ &$0.00_{\pm0.00}$ &$0.00_{\pm0.00}$ \\%

VIB-R(REF) &$99.24_{\pm0.03}$ &$95.19_{\pm0.07}$ &$86.33_{\pm0.24}$ &$55.69_{\pm0.50}$ &$17.51_{\pm0.23}$ &$2.15_{\pm0.10}$ &$0.27_{\pm0.03}$ \\ %

VIB-M(REF) &$99.28_{\pm0.02}$ &$95.17_{\pm0.06}$ &$86.28_{\pm0.23}$ &$53.82_{\pm0.32}$ &$13.53_{\pm0.05}$ &$1.64_{\pm0.09}$ &$0.19_{\pm0.02}$ \\%

VIB-M(PGD) &$99.04_{\pm0.05}$ &$98.26_{\pm0.02}$ &$97.19_{\pm0.06}$ &$95.84_{\pm0.06}$ &$94.03_{\pm0.13}$ &$90.86_{\pm0.14}$ &$85.51_{\pm0.18}$ \\ %

VIB-M(REF+FGSM) &$99.22_{\pm0.06}$ &$97.41_{\pm0.12}$ &$96.79_{\pm0.07}$ &$95.99_{\pm0.10}$ &$95.03_{\pm0.09}$ &$92.99_{\pm0.27}$ &$86.95_{\pm0.16}$ \\%

VIB-M(REF+PGD) &$99.13_{\pm0.04}$ &$98.24_{\pm0.06}$ &\pmb{$97.70_{\pm0.08}$} &\pmb{$97.36_{\pm0.08}$} &\pmb{$96.87_{\pm0.09}$} &\pmb{$96.01_{\pm0.06}$} &\pmb{$93.64_{\pm0.30}$} \\% all

\hline
\end{tabular}
\end{table*}

\begin{table*}[ht!]
\small
\caption{Models' classification test accuracy results ($\%$) with the FGSM attack accuracy on the CIFAR-10 dataset.}
\label{Tab:fgsm_cifar10}
\centering
\begin{tabular}{ |l|c|c|c|c|c|c|c|} \hline
 Method & Clean & $\epsilon=0.01$ & $\epsilon=0.02$ & $\epsilon=0.03$ & $\epsilon=0.04$ & $\epsilon=0.05$ & $\epsilon=0.06$ \\ \hline

 Baseline &$92.88_{\pm0.13}$ &$86.36_{\pm0.14}$ &$83.67_{\pm0.18}$ &$81.68_{\pm0.23}$ &$79.99_{\pm0.18}$ &$78.52_{\pm0.14}$ &$76.88_{\pm0.41}$ \\%

 Baseline(KD) &\pmb{$93.41_{\pm0.06}$} &$84.34_{\pm0.24}$ &$80.42_{\pm0.20}$ &$77.58_{\pm0.29}$ &$75.44_{\pm0.30}$ &$73.56_{\pm0.19}$ &$71.74_{\pm0.64}$ \\%

 Baseline(PGD) &$90.24_{\pm0.19}$ &$84.71_{\pm0.22}$ &$79.36_{\pm0.31}$ &$75.49_{\pm0.21}$ &$72.72_{\pm0.11}$ &$70.48_{\pm0.17}$ &$68.99_{\pm0.14}$ \\%

 Baseline(REF) &$90.51_{\pm0.22}$ &$84.71_{\pm0.34}$ &$81.81_{\pm0.21}$ &$79.52_{\pm0.17}$ &$77.31_{\pm0.26}$ &$75.15_{\pm0.29}$ &$73.63_{\pm0.22}$\\  %

 Baseline(PGD+REF) &$87.11_{\pm0.08}$ &$83.05_{\pm0.21}$ &$79.29_{\pm0.04}$ &$76.40_{\pm0.17}$ &$73.91_{\pm0.22}$ &$71.77_{\pm0.14}$ &$69.76_{\pm0.22}$ \\%

 VIB-R &$90.57_{\pm0.07}$ &$84.66_{\pm0.32}$ &$81.24_{\pm0.14}$ &$78.66_{\pm0.27}$ &$76.47_{\pm0.19}$ &$74.84_{\pm0.25}$ &$73.03_{\pm0.17}$ \\ %

 VIB-M &$91.47_{\pm0.14}$ &$84.79_{\pm0.24}$ &$81.94_{\pm0.15}$ &$79.98_{\pm0.22}$ &$78.16_{\pm0.13}$ &$76.52_{\pm0.11}$ &$75.26_{\pm0.16}$ \\%

 VIB-R(REF) &$90.98_{\pm0.07}$ &$86.82_{\pm0.16}$ &$84.80_{\pm0.23}$ &$83.09_{\pm0.19}$ &$81.62_{\pm0.20}$ &$80.37_{\pm0.20}$ &$79.44_{\pm0.21}$ \\ %

 VIB-M(REF) &$91.49_{\pm0.05}$ &\pmb{$87.42_{\pm0.17}$} &\pmb{$85.32_{\pm0.13}$} &\pmb{$84.01_{\pm0.32}$} &\pmb{$82.42_{\pm0.26}$} &\pmb{$80.97_{\pm0.24}$} &\pmb{$80.05_{\pm0.42}$} \\ %

 VIB-M(PGD) &$74.55_{\pm0.03}$ &$69.43_{\pm0.26}$ &$63.74_{\pm0.26}$ &$58.88_{\pm0.28}$ &$53.87_{\pm0.31}$ &$49.67_{\pm0.28}$ &$46.10_{\pm0.30}$ \\ %

 VIB-M(REF+FGSM) &$87.29_{\pm0.09}$ &$83.62_{\pm0.19}$ &$81.56_{\pm0.18}$ &$80.11_{\pm0.09}$ &$78.61_{\pm0.16}$ &$77.48_{\pm0.18}$ &$76.17_{\pm0.26}$ \\%

 VIB-M(REF+PGD) &$89.57_{\pm0.16}$ &$86.24_{\pm0.13}$ &$83.69_{\pm0.18}$ &$82.06_{\pm0.17}$ &$80.23_{\pm0.36}$ &$78.78_{\pm0.30}$ &$77.52_{\pm0.18}$ \\%

\hline
\end{tabular}
\end{table*}

\begin{table*}[ht!]
\small
\caption{Models' classification test accuracy results ($\%$) with the PGD attack accuracy on the CIFAR-10 dataset.}
\label{Tab:pgd_cifar10}
\centering
\begin{tabular}{ |l|c|c|c|c|c|c|c|} \hline
 Method & Clean & $\epsilon=0.01$ & $\epsilon=0.02$ & $\epsilon=0.03$ & $\epsilon=0.04$ & $\epsilon=0.05$ & $\epsilon=0.06$ \\ \hline

 Baseline &$92.88_{\pm0.13}$ &$77.28_{\pm0.18}$ &$63.96_{\pm0.26}$ &$46.42_{\pm0.28}$ &$27.18_{\pm0.25}$ &$15.49_{\pm0.38}$ &$8.92_{\pm0.16}$ \\%

 Baseline(KD) &\pmb{$93.41_{\pm0.06}$} &$76.23_{\pm0.13}$ &$51.26_{\pm0.09}$ &$32.97_{\pm0.21}$ &$24.64_{\pm0.19}$ &$20.23_{\pm0.23}$ &$17.15_{\pm0.33}$ \\ %

 Baseline(PGD) &$90.24_{\pm0.19}$ &$84.00_{\pm0.11}$ &$76.02_{\pm0.23}$ &$67.08_{\pm0.15}$ &$58.13_{\pm0.22}$ &$50.37_{\pm0.23}$ &$44.79_{\pm0.20}$ \\%

 Baseline(REF) &$90.51_{\pm0.22}$ &$71.88_{\pm0.32}$ &$57.66_{\pm0.40}$ &$41.69_{\pm0.26}$ &$28.29_{\pm0.17}$ &$19.80_{\pm0.22}$&$14.60_{\pm0.19}$\\%

 Baseline(PGD+REF) &$87.11_{\pm0.08}$ &$82.56_{\pm0.11}$ &$77.21_{\pm0.16}$ &$71.65_{\pm0.21}$ &$65.17_{\pm0.36}$ &$58.65_{\pm0.24}$ &$52.40_{\pm0.21}$ \\%

 VIB-R &$90.57_{\pm0.07}$ &$79.06_{\pm0.34}$ &$63.62_{\pm0.25}$ &$48.99_{\pm0.33}$ &$36.59_{\pm0.13}$ &$26.93_{\pm0.21}$ &$19.86_{\pm0.15}$ \\ %

 VIB-M &$91.47_{\pm0.14}$ &$79.61_{\pm0.15}$ &$68.18_{\pm0.29}$ &$62.47_{\pm0.41}$ &$57.98_{\pm0.10}$ &$52.23_{\pm0.16}$ &$46.08_{\pm0.69}$\\ %

 VIB-R(REF) &$90.98_{\pm0.07}$ &$79.87_{\pm0.06}$ &$69.35_{\pm0.30}$ &$62.58_{\pm0.23}$ &$56.57_{\pm0.15}$ &$51.90_{\pm0.44}$ &$47.70_{\pm0.18}$\\%

 VIB-M(REF) &$91.49_{\pm0.05}$ &$80.55_{\pm0.19}$ &$70.52_{\pm0.18}$ &$60.05_{\pm0.24}$ &$50.94_{\pm0.20}$ &$44.01_{\pm0.09}$ &$39.21_{\pm0.47}$ \\%

 VIB-M(PGD) &$74.55_{\pm0.03}$ &$68.74_{\pm0.12}$ &$62.92_{\pm0.27}$ &$56.32_{\pm0.37}$ &$50.29_{\pm0.19}$ &$44.46_{\pm0.32}$ &$38.94_{\pm0.18}$ \\%

 VIB-M(REF+FGSM) &$87.29_{\pm0.09}$ &$82.47_{\pm0.11}$ &$77.22_{\pm0.18}$ &$71.54_{\pm0.34}$ &$66.11_{\pm0.33}$ &$60.35_{\pm0.18}$ &$55.53_{\pm0.14}$ \\%

 VIB-M(REF+PGD) &$89.57_{\pm0.16}$ &\pmb{$85.65_{\pm0.08}$} &\pmb{$80.92_{\pm0.13}$} &\pmb{$75.78_{\pm0.20}$} &\pmb{$69.86_{\pm0.20}$} &\pmb{$63.89_{\pm0.20}$} &\pmb{$59.87_{\pm0.24}$} \\ %

 \hline
\end{tabular}
\end{table*}

There are a few insights we would like to share.
First of all, the VIB model shows a good predictive capability. However, it is still vulnerable to strong attacks like the PGD attack. From the results on the MNIST dataset, we can see that the regularly-trained VIB models, i.e., the VIB-R and the VIB-M models, outperform the corresponding baseline models under the FGSM attacks. However, they become worse under the PGD attacks. Our experimental results also show that the MINE method generally outperforms the reparamterization method for training the VIB model because the former can provide more complex posteriors.

It is generally believed that training a deep learning model with adversarial examples can improve the model's robustness, particularly, under the PGD attack. However, according to the MNIST results, the VIB (PGD) model does not achieve the best on any case. That indicates a naive adoption of adversarial training cannot improve the performance of the VIB model.  In fact, the results obtained by the VIB-M (REF+FGSM) and the VIB-M (REF+PGD) suggest only when the VIB model is trained with the support of reference network, the best performance can be achieved. \emph{By contrast, using our proposed methods, VIB-M(REF+FGSM/PGD), the robustness of the VIB models can be significantly improved under both FGSM and PGD attacks.} Between the VIB-M (REF+FGSM) and the VIB-M (REF+PGD) models, the former enjoys faster training while maintains comparable performance. From the results on CIFAR-10, we see that VIB-R/M do not necessarily outperform the baselines under either the FGSM or the PGD attacks. The VIB-R/M (REF) models have better performance under the FGSM attacks, but do not perform very well under the PGD attack. Our proposed methods, the VIB-M (REF+FGSM/PGD) models, have similar excellent performance against both FGSM and PDG attacks.

Overall, it can be seen that, first, learning with reference can improve the robustness of the VIB model; second, compared with other methods, adversarial training is more effective to defend the PGD attacks and using reference can improve the performance of adversarial training; third, as opposed to using hard labels, our method can use a single-step adversarial training method to accelerate the training process, i.e. the FGSM training, without causing the label leaking effect.

\section{Conclusion}

In this paper, we propose a novel method for training the VIB model which can produce accurate predictions as well as be robust under adversarial attacks. The idea is to train the VIB model by using the reference network and the MINE method. The former enables the mapping of both clean and perturbed inputs onto a smooth output manifold, so that the learned VIB model is smoother (and more robust). Also, this enables a fast training with FGSM-based method instead of PGD-based, without compromising the robustness performance. To the best of our knowledge, this is the very first study that integrates them into a VIB-based framework. The experimental results show that our proposed model outperforms the state-of-the-art benchmarked models.

\bibliographystyle{named}

\section{Introduction}

Deep neural networks (DNNs) have shown very powerful learning  ability for carrying on many kinds of tasks, such as image classification. However, DNNs are proved to be susceptible to adversarial attack \cite{goodfellow2014explaining} \cite{madry2018towards} \cite{carlini2017towards}. In order to enable DNNs to defend adversarial attack, researchers have proposed many approaches.

For instance, \cite{kurakin2016adversarial} proposes to use adversarial samples as the training sample and this method is called Adversarial Training. In \cite{papernot2016distillation}, the researchers use Knowledge Distillation to acquire soft labels to as the training targets to improve the robustness of deep neural networks. Another approach proposed in \cite{raghunathan2018certified} is based on Adversarial Training and regularization.

Variational Information Bottleneck (VIB) \cite{alemi2017deep} leverages Variational Inference to implement the Information Bottleneck method \cite{tishby2000information}, \cite{tishby2015deep}, which enable it to learn disentangled latent representations in which the latent variables corresponding to similar images are close to each other. Due to this, VIB has been proven to be robust under adversarial attack \cite{alemi2017deep}. However, the performance VIB can still be further improved. For example, \cite{federici2020learning} uses a multi-view Information Bottleneck to learn robust representations. Additionally, \cite{madaan2020adversarial} uses Adversarial Training and Vulnerability Suppression loss to prune the neural networks so as to increase their robustness.

In this work, we propose a novel method to improve the robustness of VIB. One of the main idea is to take advantage of  pre-learnt soft labels as the learning targets. Note that the soft labels we use here are the numerical outputs of the last hidden layer of the decoder not the probabilities as in Knowledge Distillation. Using soft labels as training targets cast the original classification into a maximum likelihood estimation (MLE)-based regression problem, therefore the output computed by clean samples in the training process and the output computed by adversarial samples in the training process share the same manifold because the adversarial samples can be seen as clean samples with larger variances. That's to say our learning goal is not the decision boundary for classification tasks which is prone to cause overfitting but a set of vectors in space similar to regression tasks. To some extent, this method is close to Denoising Autoencodes \cite{vincent2008extracting} \cite{vincent2010stacked}. As a result, in our model, training with clean samples and reference can improve the robustness of standard VIB.

Furthermore, if we use adversarial samples as the training samples, the performance of our method can be further improved. This is because, in this case, the VIB can learn more powerful modeling ability to transform the perturbed samples to the target manifold compare to clean samples. But thanks to the soft labels pre-learnt from clean images, the VIB model learnt via adversarial samples is still suitable for clean sample. Finally, combining these two approaches, the robustness of VIB is significantly improved.

The overall contributions in this paper are summarized as follows:

\begin{itemize}
    \item We propose to use the soft labels learnt by a reference net to improve the robustness of VIB models;

    \item We combine Adversarial Training with pre-learnt soft labels to further the performance of VIB, and this method is able to defend both FGSM and PGD attack effectively;

    \item We utilize one-step adversarial samples to accelerate the training process without suffering the label leaking effect.

\end{itemize}

\section{Proposed Model}

\begin{figure}[t]
\centering
\includegraphics[width=1\linewidth]{fig_framework.pdf}
\caption{Caption}
\label{fig:framework}
\end{figure}

In this work, we propose a novel VIB model to defend adversarial attack. Our method consists of a reference network $f_{\text{Ref}}(x)$, an encoder $f_E$, a discriminator $T$ and a classifier $f_C$. The reference net is used to learn the soft labels from the clean images. These soft labels can be used to measure the discrepancy between the clear manifold and the perturbed manifold. We estimate the mutual information $I(X;Z)$ by using MINE. Finally, we leverage Adversarial Training to further improve the performance.

The final loss function is
\begin{align}
\label{Eq:RVIB_loss}
\mathscr{L} & =
\sum_{i \in Clean } \bigg[ ||f_{\text{Ref}}(x_i)- f_C \big (f_E(x_i^\text{Pert}) \big)||^2  \nonumber  \\
& + \beta \big (\mathop{{}\mathbb{E}}{_{P_{XZ}}}\big [T \big] - \log \mathop{\mathbb{E}}{_{_{P_{X}\otimes P_{Z}}}\big[e^{T} \big]} \big ) \bigg ]
\end{align}

\begin{align}
\label{Eq:Final_loss}
\mathscr{L}_{\text{Proposed}} & = \frac{1}{m-k+\lambda k} \big ( \sum_{i \in Clean }{\mathscr{L}_{\text{Clean}}(x_i,y_i)}  \nonumber \\
& \qquad +\lambda \sum_{i \in Adv }{\mathscr{L}_{\text{Adv}}(x_i^{\text{Adv}}, y_i)} \big)
\end{align}
where $m$ is the batch size, $k$ is the adversarial sample size, $\lambda$ is the relative weight.

At the first step, we pre-train a reference net with the clean images; at the second step, we use PGD or FGSM to generate adversarial samples; at the third step, we use these adversarial samples to feed the encoder to obtain the latent variables; at the fourth step, we use a discriminator to distinguish the joint variables from the shuffled variables; at the fifth step, we use the joint variables to feed the classifier; at the sixth step, we use the clean samples to obtain the corespondent soft labels via the reference net; at the last step we minimize the final loss function.

The overall algorithm is summarized as follow:

\begin{algorithm*}
\caption{Algorithm}
\label{Alg:Algorithm}
\begin{algorithmic}[1]
\Require{$x$ (input), $y$ (target)}
\Statex

\While{ $\text{epoch} \leq \text{max epoch}$}
    \State{train a reference neural network $f_{\text{Ref}}(x)$;}
\EndWhile\\

\While{ $\text{epoch} \leq \text{max epoch}$} \Comment{train robust VIB}

\State{If $i$ $\in$ Clean:}
    \State{\quad $x_i^\text{Pert}$ $\gets$ $x_i$ + $\epsilon$}  \Comment{inject noise to provide gradients}
\State{If $i$ $\in$ Adversarial:}
    \State{\quad $x_i^\text{Pert}$ $\gets$ $x_i$}  \Comment{get adversarial samples according to Eq.(\ref{Eq:FGSM}) or Eq.(\ref{Eq:PGD})}

    \State{$xz_{\text{Joint}}$ $\gets$ $f_E(x_i^\text{Pert})$}  \Comment{$f_E$: encoder network; $xz_{\text{Joint}}$: joint samples $P_{XZ}$}

    \State{$xz_{\text{Ind}}$ $\gets$ shuffle $xz_{\text{Joint}}$}   \Comment{$xz_{\text{Ind}}$: independent samples $P_{X} \otimes P_{Z}$}

    \State{$\mathop{{}\mathbb{E}}{_{P_{XZ}}}[T] - \log \mathop{\mathbb{E}}{_{_{P_{X}\otimes P_{Z}}}[e^{T}]}$} \Comment{compute upper bound for $I(Z;X)$}

    \State{$o_{\text{Ref}}$ $\gets$ $f_{\text{Ref}}(x_i)$}   \Comment{compute target via reference net }

    \State{$o_{\text{VIB}}$ $\gets$ $f_C(xz_{\text{Joint}})$}
    \State{Minimize final loss function $\mathscr{L}_{\text{Proposed}}$}

\EndWhile

\end{algorithmic}
\end{algorithm*}

\subsection{Variational Information Bottleneck}
In this section, we briefly introduce Variational Information Bottleneck (VIB). Let $X$ be the input, $Y$ be the target and $Z$ be the latent representation, and assume the information Markov chain, $X \rightarrow Z \rightarrow Y$, then the optimization objective of the information bottleneck method \cite{tishby2000information}, \cite{tishby2015deep} can be described as follows:

\begin{align}
\label{Eq:IB_1}
\text{max}  & ~ I(Z;Y),   \\
\label{Eq:IB_2}
\text{s.t.} & ~ I(Z;X) \leq I_C,
\end{align}

where $I$ is the mutual information between two variables, and $I_C$ is a constant representing information constraint.

Equivalently, the above problem is to maximize the following Lagrangian objective function
\begin{align}
\label{Eq:IB_Lag}
\mathscr{L}_{\text{IB}} = I(Z;Y) - \beta \big( I(Z;X) - I_C \big),
\end{align}
where $\beta$ is a non-negative Lagrangian multiplier.

Since direct computing mutual information is difficult, we can derive the variational bounds for $I(Z;X)$ and $I(Z;Y)$, respectively. Hence, the upper bound of $I(Z;X)$ yields

\begin{align}
\label{Eq:I_ZX_bound}
I(Z;X) \leq \mathbb{E}_x \left[D_{\text{KL}} \big(p(z|x)||q(z)\big)\right]
\end{align}

where $q(z)$ is an uninformative prior distribution. This equation can be solved by the re-parameterization method~\cite{kingma2014auto}:
\begin{align}
\label{Eq:F_Statistic}
f_\phi(x,\epsilon_z) = \mu_z(x) + \sigma_z(x) \odot \epsilon_z
\end{align}
where $f_\phi(x,\epsilon_z)$ is the deterministic function used in encoder $p_{\phi}(z|x)$, $\odot$ is the Hadamard product, and $\epsilon_z$ is a random noise sampled from a standard diagonal Gaussian distribution $\mathcal{N}(0,\mathbb{I})$, $\mu_z$ and $\sigma_z$ are the deterministic functions to calculate the mean and variance of the latent Gaussian distribution, respectively.

The lower bound of $I(Z;Y)$ yields
\begin{align}
\label{Eq:I_ZY_bound}
I(Z;Y) \geq & \mathbb{E}_{p(y,z)} \big [ \log \left\{ p(y|z) \right\} \big]
\end{align}

The loss function of VIB is

\begin{align}
\label{Eq:OP_VIB_Repa}
\mathscr{L}_{\text{VIB}} & =  \mathbb{E}_{\mathscr{D}} \bigg[ \mathbb{E}_{z \sim p(z|x)} \big[-\log\{ p(y|z)\} \big] \nonumber  \\
& \qquad + \beta D_{\text{KL}}\big(p(z|x)|| q(z)\big ) \bigg].
\end{align}

where $D_{\text{KL}}$ is the Kullback-Leibler (KL) divergence, $p(z|x)$ is the encoder and $p(y|z)$ is the decoder.

Compared to conventional neural networks, the latent disentangled representation of VIB can improve the robustness of the neural network. But it is still not powerful enough to withstand strong adversarial attacks.

\subsection{Mutual Information Neural Estimation}

For estimating mutual information $I(X;Z)$, apart from using the re-parameterization method, alternatively, we can use Mutual Information Neural Estimation (MINE) \cite{belghazi2018mutual}. The mutual information $I(X;Z)$ can be defined as follow:
\begin{align}
\label{Eq:I_ZX}
I(Z;X) =  D_{KL} \big (\mathop{{}\mathbb{{P_{XZ}}}}||\mathop{{}\mathbb{{P_{X}}}} \otimes \mathop{{}\mathbb{{P_{Z}}}} \big ) \end{align}

where $\otimes$ denotes Kronecker product.

We can leverage the Donsker-Varadhan representation of the KL divergence \cite{belghazi2018mutual}

\begin{align}
\label{Eq:DV}
D_{KL}(P||Q) = \underset{T:\Omega \rightarrow \mathop{{}\mathbb{R}}}{\mathrm{sup}} \mathop{{}\mathbb{E}}{_{P}}\big[T\big] - \log \mathop{\mathbb{E}}{_{Q}\big[e^T\big]}
\end{align}

where $P$ is the posterior distribution, $Q$ is the prior distribution and $T$ is a discriminator net used to distinguish joint samples from independent samples.

Therefore, as opposed to the variational upper bound, Eq. (\ref{Eq:I_ZX_bound}), we can estimate $I(Z;X)$ according to Eq. (\ref{Eq:DV})

\begin{equation}
\label{Eq:I_ZX_DV}
I_\theta(Z;X) = \underset{\theta \in \Theta}{\mathrm{sup}} \mathop{{}\mathbb{E}}{_{P_{XZ}}}\big[T_\theta\big] - \log \mathop{\mathbb{E}}{_{_{P_{X}\otimes P_{Z}}}\big[e^{T_\theta}\big]}
\end{equation}

The loss function of VIB-MINE is
\begin{align}
\label{Eq:OP_VIB_Mine}
\mathscr{L}_{\text{VIB-M}} & =  \mathbb{E}_{\mathscr{D}} \bigg[ \mathbb{E}_{z \sim p_{\phi}(z|x)} \big[-\log\{ p_{\omega}(y|z)\} \big]  \nonumber \\
& \qquad + \beta \Big (\mathop{{}\mathbb{E}}{_{P_{XZ}}}\big[T_\theta\big] - \log \mathop{\mathbb{E}}{_{_{P_{X}\otimes P_{Z}}}\big[e^{T_\theta}\big]} \Big ) \bigg].
\end{align}

The advantage of using MINE is that we do not need to assume the posterior $p(z|x)$ is Gaussian, thus its shape can be more complex to approximate more accurately the true posterior $p(z|x)$. In the proposed method, we will use MINE instead of the re-parameterization method \cite{kingma2014auto}.

\subsection{Reference Network}

In order to learning robust outputs, we choose to use soft labels as the learning targets, which can be regarded a set of target vectors corresponding to each input image in the space, instead of original hard labels. The reason of doing so is that hard labels only keep the discrete categorical information but lose the more detailed information of each input. More over, this estimation method can be see as a MLE method, thus when the variance of the input increases, the model's output still can yield the original estimation.

To this end, we use a reference net to learn soft labels as the learning targets for the VIB model. Note that the soft labels in our method are s different from the labels used in Knowledge Distillation \cite{hinton2015distilling} \cite{mishra2017apprentice}, which are probabilities.

\begin{equation}
o_{\text{Ref}} = f_{\text{Ref}}(x)
\end{equation}

where $o_{\text{ref}}$ is the output of the reference network's last hidden layer (before the softmax operation).

The final output of the reference network
\begin{equation}
y_{\text{Ref}} = \text{softmax}(o_{\text{Ref}}) \nonumber
\end{equation}

Similarly, the prediction of our model is the output of the classifier network's last hidden layer whose output is numerical values (not probabilities)

\begin{equation}
o_{\text{VIB}} = f_{\text{VIB}}(x)
\end{equation}

The new reconstruction error is the Euclidean distance (or Manhattan distance) between the source network output and the target network output:
\begin{equation}
\text{Dist}(o_{\text{Ref}}, o_{\text{VIB}}) = ||o_{\text{Ref}}-o_{\text{VIB}}||^2
\end{equation}

This is maximum likelihood estimation (MLE), like denoising autoencoders \cite{vincent2008extracting} \cite{vincent2010stacked}, it can reduce the perturbation from the corrupted input, e.g., gradient-based adversarial attack.

\subsection{Adversarial Training}
Deep neural networks are well known for being vulnerable to adversarial attack. Even small perturbation on the clean images may significantly decrease the test performance of DNNs.

\subsubsection{Adversarial Attack}
\textbf{FGSM Attack.}
The fast gradient sign method (FGSM)~\cite{goodfellow2014explaining} is a simple gradient-based method used to produce the adversarial samples:
\begin{align}
\label{Eq:FGSM}
x^\text{Adv} = x + \epsilon \cdot \text{sign}\big(\nabla_x J(w;x;y)\big)
\end{align}
where $J$ is the loss function, $w$ is the model parameter and $\epsilon $ is a positive scalar which denotes the magnitude of the perturbation.

\textbf{PGD Attack}
Since FGSM is a simple one-step adversarial method, which is not powerful enough in some cases, researchers propose a more powerful method to generate adversarial samples, called Projected Gradient Descent (PGD) \cite{kurakin2016adversarial}. PGD is similar to FGSM, whereas it attacks the clean samples recursively, thus it produce more complex adversarial samples.

\begin{align}
\label{Eq:PGD}
x^\text{Adv}_0 & = x; \nonumber  \\
x^\text{Adv}_{t+1} & = \text{Clip}_{x, \epsilon} \Big \{x^\text{Adv}_t + \alpha \cdot \text{sign}\big(\nabla_x J(w;x^\text{Adv}_t ;y)\big) \Big \}
\end{align}

where $\alpha$ is the step size and $t$ is the step number.

\subsubsection{Training with Adversarial Samples}
Adversarial samples can fool the trained neural networks. Nevertheless, adversarial samples can also be used to help training more robust neural networks, this is known as Adversarial Training \cite{kurakin2016adversarial}.

Adversarial samples:
$x_\text{Clean}$ is clean samples and $x_\text{Adv}$ denotes adversarial samples.

\begin{align}
\label{Eq:Adv}
x_\text{Adv}=x_\text{Clean} + \delta
\end{align}

where $\delta$ is the perturbation.

For any $x \in {X}$ with $l_p$-ball $\mathcal{B}(x, \epsilon)$ around $x: \{ x_\text{Adv} \in X: ||x_\text{Adv} - x_\text{Clean} || \leq \epsilon \}$ \cite{madaan2020adversarial}, in a VIB, the related latent presentation $z$  from the input $x$ via a encoder and the $\hat{y}$ is related prediction computed from $z$ via a classifier. This relations can be described as follow:
\begin{align}
\label{Eq:Prediction}
\hat{y} =  f_C(z) =  f_C(f_E(x))
\end{align}

where $f_E(x)$ is the encoder and $f_C(z)$ is the classifier.

To make the neural network robust enough to adversarial attack, we need to ensure the output computed by perturbed samples and the output computed by clean samples share the same manifold. Therefore, we should minimize the distance between these manifolds. According to Eq. (\ref{Eq:Prediction}), this task can be cast into

\begin{align}
\label{Eq:Distance}
& \min \mathbb{E}_{\{\mathcal{D}, \epsilon\}} \big[ \text{Dist}(\tilde{y},y) \big] \nonumber \\
& \qquad \equiv \min \mathbb{E}_{\{\mathcal{D}, \epsilon\}} \big[\text{Dist}\big(f_C(\tilde{z}),f_C(z) \big)\big]\nonumber \\
& \qquad \equiv \min \mathbb{E}_{\{\mathcal{D}, \epsilon\}} \big[\text{Dist}(\tilde{z},z)\big]
\end{align}

where $\text{Dist}$ denotes some kind of distance. The posterior $p(z|x)$ is Gaussian-like thus $p(y|z)$ is Gaussian-like as well. Consequently, $\text{Dist}(\tilde{y},y)$ can be chosen to be Euclidean distance.

Based on Eq. (\ref{Eq:Distance}), using soft labels can keep the discrepancy information between adversarial latent representations and clean latent representations. Therefore, equivalently, minimizing $\text{Dist}(\tilde{y},y)$ will prevent the manifold of the latent representation being distorted. In \cite{madaan2020adversarial}, $\text{Dist}(\tilde{z},z)$ is called \textit{vulnerability}. However, since we use the soft labels learnt from clean samples, we do not need to minimize the discrepancy between the manifold learnt from clean samples and the manifold learnt from adversarial samples additionally. The loss function of Adversarial Training:

Adversarial samples can help learning robust neural network, how training with one-step adversarial images will cause the \textbf{label leaking} effect \cite{kurakin2016adversarial}. It says if the neural network is trained with one-step adversarial images, e.g., FGSM samples, and the true hard labels, the testing accuracy on adversarial images are much higher than the accuracy on clear images.  This issue appears because one-step adversarial images generated by simple transformations, and it is not hard for neural network model to learn such transformations, as a result, the neural network performs better on the adversarial images than the clean images. On the contrary, more complex adversarial transformation, e.g., PGD are more difficult to learn. However, the drawback of PGD is it it much more time-consuming. Fortunately, as \cite{kurakin2016adversarial} recommends, in our method, we use soft labels instead of hard labels to circumvent the label leaking effect. Consequently, we can use FGSM adversarial learning to expedite the training process and avoid the label leaking effect.

\section{Experiments}

\textit{Distillation} stands for Knowledge Distillation \cite{papernot2016distillation};
\textit{Adversarial} stands for Adversarial Training \cite{kurakin2016adversarial}; \textit{Reference} stands for training with reference; \textit{Adv + Ref} stands for adversarial training with reference; \textit{VIB-R} stands for the VIB models using the re-parameterization trick \cite{alemi2017deep}; \textit{VIB-M} stands for the VIB models using MINE \cite{belghazi2018mutual}; \textit{Ref} denotes the models learning with a reference net; \textit{PGD} denotes the models using PGD Adversarial Training; \textit{FGSM} denotes the models using FGSM Adversarial Training.

\subsection{Data}

We use following datasets to validate our method:
\textbf{MNIST.} This dataset includes $60, 000$ grey scale handwritten images \cite{lecun1998mnist}. The image number for training is $50, 000$ and the image number for training is $10, 000$. The sizes of these image are $28 \times 28$. The architecture for the baseline, Distillation Training, Adversarial Training, Reference, and Adv + Ref is CONV(64,3,1) + CONV(64,3,1) + FC(512) + FC(512) + FC(10).

\textbf{CIFAR-10.} This dataset includes $60, 000$ color images \cite{lecun1998mnist}. The image number for training is $50, 000$ and the image number for training is $10, 000$. The sizes of these image are $32 \times 32$. The images are normalized for training. The architecture for the baseline, Distillation Training, Adversarial Training, Reference, and Adv + Ref is VGG16 \cite{simonyan2014very}.

\subsection{Experimental Settings}

For MNIST, for Distillation Training: the temperature $T=1$; For Adversarial Training: PGD with step size $\alpha = 0.01$, batch size $m=100$, training adversarial size $k=50$, $\epsilon = 0.3$ and steps = $20$; FGSM: $\epsilon = 0.3$; For VIB models: $\beta=0.001$, latent dimension = $256$. The optimizers are Adam \cite{kingma2014adam} and the learning rate is $0.001$.

For CIFAR-10, for Distillation Training: the temperature $T=1$; for Adversarial Training: PGD with step size $\alpha = 0.007$, batch size $m=100$, training adversarial size $k=50$, $\epsilon = 0.03$ and the step number = $10$; FGSM: $\epsilon = 0.3$; for VIB models: $\beta=0.001$, latent dimension = $256$. The optimizers are Stochastic Gradient Descent (SGD) with adaptive learning rate.

For MNIST, we use values $0.05$, $0.1$, $0.15$, $0.2$, $0.25$, $0.3$ for $\epsilon$ for both FGSM attack and PGD attack. The step size of PGD, $\alpha$, is set to be $0.01$. The step number is $40$.

For CIFAR-10, we use values $0.01$, $0.02$, $0.03$, $0.04$, $0.05$, $0.07$ for $\epsilon $ for both FGSM attack and PGD attack with random start. The step size of PGD, $\alpha$, is set to be $0.007$. The step number is $40$.

We run $5$ trials with random seeds for each experiment, the results with means and standard deviations are exhibited in the tables.

\subsection{Results and Discussion}

\subsubsection{Results on MNIST}

\begin{table*}[ht!]
\small
\caption{FGSM attack results on MNIST ($\%$).}
\label{Tab:fgsm_mnist}
\centering
\begin{tabular}{|c|c|c|c|c|c|c|c|} \hline
 Method & Clean & $\epsilon=0.05$ & $\epsilon=0.1$ & $\epsilon=0.15$ & $\epsilon=0.2$ & $\epsilon=0.25$ & $\epsilon=0.3$ \\ \hline

 Baseline &$99.20_{\pm0.05}$ &$93.64_{\pm0.18}$ &$91.03_{\pm0.28}$ &$89.76_{\pm0.22}$ &$88.73_{\pm0.22}$ &$87.94_{\pm0.11}$ &$86.57_{\pm0.18}$\\%

 Distillation &$99.28_{\pm0.04}$ &$96.61_{\pm0.06}$ &$95.17_{\pm0.09}$ &$94.33_{\pm0.19}$ &$93.39_{\pm0.31}$ &$91.20_{\pm0.21}$ &$88.35_{\pm0.22}$ \\%

 Adversarial &$99.17_{\pm0.02}$ &\pmb{$98.54_{\pm0.03}$} &$97.92_{\pm0.08}$ &$97.18_{\pm0.06}$ &$96.40_{\pm0.09}$ &$95.40_{\pm0.12}$ &$94.36_{\pm0.06}$\\%

 Reference &\pmb{$99.35_{\pm0.02}$} &$97.00_{\pm0.07}$ &$96.10_{\pm0.21}$ &$95.56_{\pm0.07}$ &$94.88_{\pm0.20}$ &$93.47_{\pm0.23}$ &$91.27_{\pm0.24}$\\%

 Adv+Ref &\pmb{$99.35_{\pm0.03}$} &$97.00_{\pm0.11}$ &$96.02_{\pm0.12}$ &$95.58_{\pm0.14}$ &$94.75_{\pm0.26}$ &$93.52_{\pm0.38}$ &$91.35_{\pm0.18}$\\%

 VIB-R &$99.10_{\pm0.05}$ &$97.24_{\pm0.05}$ &$96.52_{\pm0.20}$ &$95.52_{\pm0.20}$ &$94.20_{\pm0.12}$ &$91.44_{\pm0.26}$ &$88.44_{\pm0.34}$ \\%

 VIB-M &$99.17_{\pm0.03}$ &$96.96_{\pm0.11}$ &$96.23_{\pm0.08}$ &$95.29_{\pm0.18}$ &$94.04_{\pm0.17}$ &$91.68_{\pm0.15}$ &$88.96_{\pm0.40}$\\%

 VIB-R(Ref) &$99.24_{\pm0.03}$ &$97.24_{\pm0.07}$ &$96.58_{\pm0.10}$ &$95.45_{\pm0.12}$ &$93.35_{\pm0.27}$ &$90.45_{\pm0.33}$ &$86.33_{\pm0.50}$\\%

 VIB-M(Ref) &$99.28_{\pm0.02}$ &$96.96_{\pm0.08}$ &$96.19_{\pm0.16}$ &$95.44_{\pm0.16}$ &$94.79_{\pm0.05}$ &$92.86_{\pm0.23}$ &$89.60_{\pm0.04}$\\ %

 VIB-M(PGD) &$99.04_{\pm0.05}$ &$98.37_{\pm0.04}$ &$97.61_{\pm0.08}$ &$96.79_{\pm0.08}$ &$95.86_{\pm0.13}$ &$94.73_{\pm0.06}$ &$93.41_{\pm0.10}$\\ %

 VIB-M(Ref+FGSM) &$99.22_{\pm0.06}$ &$98.06_{\pm0.06}$ &$97.64_{\pm0.12}$ &$97.35_{\pm0.08}$ &$97.12_{\pm0.13}$ &$96.89_{\pm0.09}$ &$96.53_{\pm0.11}$\\

 VIB-M(Ref+PGD) &$99.13_{\pm0.04}$ &$98.48_{\pm0.08}$ &\pmb{$98.03_{\pm0.04}$} &\pmb{$97.79_{\pm0.07}$} &\pmb{$97.59_{\pm0.08}$} &\pmb{$97.35_{\pm0.05}$} &\pmb{$97.12_{\pm0.12}$} \\% all

 \hline
\end{tabular}
\end{table*}

The FGSM attack results on MNIST are exhibited in Table \ref{Tab:fgsm_mnist}.

\begin{table*}[ht!]
\small
\caption{PGD attack results on MNIST ($\%$).}
\label{Tab:pgd_mnist}
\centering
\begin{tabular}{|c|c|c|c|c|c|c|c|} \hline
 Method & Clean &  $\epsilon=0.05$ & $\epsilon=0.1$ & $\epsilon=0.15$ & $\epsilon=0.2$ & $\epsilon=0.25$ & $\epsilon=0.3$ \\ \hline

 Baseline &$99.20_{\pm0.05}$ &$92.60_{\pm0.06}$ &$87.36_{\pm0.23}$ &$64.54_{\pm0.39}$ &$22.88_{\pm0.14}$ &$3.09_{\pm0.13}$ &$0.36_{\pm0.05}$ \\%

 Distillation &$99.28_{\pm0.04}$ &$93.71_{\pm0.10}$ &$76.74_{\pm0.13}$ &$30.88_{\pm0.12}$ &$1.98_{\pm0.08}$ &$0.01_{\pm0.01}$ &$0.00_{\pm0.00}$ \\%

 Adversarial &$99.17_{\pm0.02}$ &\pmb{$98.44_{\pm0.07}$} &$97.51_{\pm0.05}$ &$96.47_{\pm0.01}$ &$94.68_{\pm0.09}$ &$92.57_{\pm0.07}$ &$88.34_{\pm0.09}$  \\%

 Reference &\pmb{$99.35_{\pm0.02}$} &$95.45_{\pm0.05}$ &$88.99_{\pm0.18}$ &$60.79_{\pm0.17}$ &$15.14_{\pm0.11}$ &$1.39_{\pm0.06}$ &$0.14_{\pm0.01}$\\%

 Adv+Ref &\pmb{$99.35_{\pm0.03}$} &$95.39_{\pm0.17}$ &$89.98_{\pm0.13}$ &$66.31_{\pm0.24}$ &$22.75_{\pm0.21}$ &$2.58_{\pm0.14}$ &$0.36_{\pm0.04}$ \\%

 VIB-R &$99.10_{\pm0.05}$ &$92.06_{\pm0.19}$ &$43.28_{\pm0.24}$ &$2.55_{\pm0.08}$ &$0.03_{\pm0.02}$ &$0.00_{\pm0.00}$ &$0.00_{\pm0.00}$ \\ %

 VIB-M &$99.17_{\pm0.03}$ &$93.39_{\pm0.08}$ &$61.25_{\pm0.19}$ &$8.81_{\pm0.11}$ &$0.03_{\pm0.01}$ &$0.00_{\pm0.00}$ &$0.00_{\pm0.00}$ \\%

 VIB-R(Ref) &$99.24_{\pm0.03}$ &$95.19_{\pm0.07}$ &$86.33_{\pm0.24}$ &$55.69_{\pm0.50}$ &$17.51_{\pm0.23}$ &$2.15_{\pm0.10}$ &$0.27_{\pm0.03}$ \\ %

 VIB-M(Ref) &$99.28_{\pm0.02}$ &$95.17_{\pm0.06}$ &$86.28_{\pm0.23}$ &$53.82_{\pm0.32}$ &$13.53_{\pm0.05}$ &$1.64_{\pm0.09}$ &$0.19_{\pm0.02}$ \\%

 VIB-M(PGD) &$99.04_{\pm0.05}$ &$98.26_{\pm0.02}$ &$97.19_{\pm0.06}$ &$95.84_{\pm0.06}$ &$94.03_{\pm0.13}$ &$90.86_{\pm0.14}$ &$85.51_{\pm0.18}$ \\ %

 VIB-M(Ref+FGSM) &$99.22_{\pm0.06}$ &$97.41_{\pm0.12}$ &$96.79_{\pm0.07}$ &$95.99_{\pm0.10}$ &$95.03_{\pm0.09}$ &$92.99_{\pm0.27}$ &$86.95_{\pm0.16}$ \\%

 VIB-M(Ref+PGD) &$99.13_{\pm0.04}$ &$98.24_{\pm0.06}$ &\pmb{$97.70_{\pm0.08}$} &\pmb{$97.36_{\pm0.08}$} &\pmb{$96.87_{\pm0.09}$} &\pmb{$96.01_{\pm0.06}$} &\pmb{$93.64_{\pm0.30}$} \\% all

 \hline
\end{tabular}
\end{table*}

The PGD attack results on MNIST are exhibited in Table \ref{Tab:pgd_mnist}.

\subsubsection{Results on CIFAR-10}
\begin{table*}[ht!]
\small
\caption{FGSM attack results on CIFAR-10 ($\%$).}
\label{Tab:fgsm_cifar10}
\centering
\begin{tabular}{ |c|c|c|c|c|c|c|c|} \hline
 Method & Clean & $\epsilon=0.01$ & $\epsilon=0.02$ & $\epsilon=0.03$ & $\epsilon=0.04$ & $\epsilon=0.05$ & $\epsilon=0.07$ \\ \hline

 Baseline &$92.88_{\pm0.13}$ &$86.36_{\pm0.14}$ &$83.67_{\pm0.18}$ &$81.68_{\pm0.23}$ &$79.99_{\pm0.18}$ &$78.52_{\pm0.14}$ &$75.28_{\pm0.18}$ \\%

 Distillation &\pmb{$93.41_{\pm0.06}$} &$84.34_{\pm0.24}$ &$80.42_{\pm0.20}$ &$77.58_{\pm0.29}$ &$75.44_{\pm0.30}$ &$73.56_{\pm0.19}$ &$70.03_{\pm0.36}$ \\%

 Adversarial &$90.24_{\pm0.19}$ &$84.71_{\pm0.22}$ &$79.36_{\pm0.31}$ &$75.49_{\pm0.21}$ &$72.72_{\pm0.11}$ &$70.48_{\pm0.17}$ &$67.78_{\pm0.21}$ \\%

 Reference &$90.51_{\pm0.22}$ &$84.71_{\pm0.34}$ &$81.81_{\pm0.21}$ &$79.52_{\pm0.17}$ &$77.31_{\pm0.26}$ &$75.15_{\pm0.29}$ &$71.96_{\pm0.16}$\\  %

 Adv+Ref &$87.11_{\pm0.08}$ &$83.05_{\pm0.21}$ &$79.29_{\pm0.04}$ &$76.40_{\pm0.17}$ &$73.91_{\pm0.22}$ &$71.77_{\pm0.14}$ &$68.13_{\pm0.35}$ \\%

 VIB-R &$90.57_{\pm0.07}$ &$84.66_{\pm0.32}$ &$81.24_{\pm0.14}$ &$78.66_{\pm0.27}$ &$76.47_{\pm0.19}$ &$74.84_{\pm0.25}$ &$71.72_{\pm0.27}$ \\ %

 VIB-M &$91.47_{\pm0.14}$ &$84.79_{\pm0.24}$ &$81.94_{\pm0.15}$ &$79.98_{\pm0.22}$ &$78.16_{\pm0.13}$ &$76.52_{\pm0.11}$ &$73.87_{\pm0.26}$ \\%

 VIB-R(Ref) &$90.98_{\pm0.07}$ &$86.82_{\pm0.16}$ &$84.80_{\pm0.23}$ &$83.09_{\pm0.19}$ &$81.62_{\pm0.20}$ &$80.37_{\pm0.20}$ &$77.97_{\pm0.25}$ \\ %

 VIB-M(Ref) &$91.49_{\pm0.05}$ &\pmb{$87.42_{\pm0.17}$} &\pmb{$85.32_{\pm0.13}$} &\pmb{$84.01_{\pm0.32}$} &\pmb{$82.42_{\pm0.26}$} &\pmb{$80.97_{\pm0.24}$} &\pmb{$78.84_{\pm0.20}$} \\ %

 VIB-M(PGD) &$74.55_{\pm0.03}$ &$69.43_{\pm0.26}$ &$63.74_{\pm0.26}$ &$58.88_{\pm0.28}$ &$53.87_{\pm0.31}$ &$49.67_{\pm0.28}$ &$42.96_{\pm0.20}$ \\ %

 VIB-M(Ref+FGSM) &$87.29_{\pm0.09}$ &$83.62_{\pm0.19}$ &$81.56_{\pm0.18}$ &$80.11_{\pm0.09}$ &$78.61_{\pm0.16}$ &$77.48_{\pm0.18}$ &$75.20_{\pm0.23}$ \\%

 VIB-M(Ref+PGD) &$89.57_{\pm0.16}$ &$86.24_{\pm0.13}$ &$83.69_{\pm0.18}$ &$82.06_{\pm0.17}$ &$80.23_{\pm0.36}$ &$78.78_{\pm0.30}$ &$76.44_{\pm0.20}$ \\%
\hline
\end{tabular}
\end{table*}

The FGSM attack results on CIFAR-10 are exhibited in Table \ref{Tab:fgsm_cifar10}.

\begin{table*}[ht!]
\small
\caption{PGD attack results on CIFAR-10 ($\%$).}
\label{Tab:pgd_cifar10}
\centering
\begin{tabular}{ |c|c|c|c|c|c|c|c|} \hline
 Method & Clean & $\epsilon=0.01$ & $\epsilon=0.02$ & $\epsilon=0.03$ & $\epsilon=0.04$ & $\epsilon=0.05$ & $\epsilon=0.07$ \\ \hline

 Baseline &$92.88_{\pm0.13}$ &$77.28_{\pm0.18}$ &$63.96_{\pm0.26}$ &$46.42_{\pm0.28}$ &$27.18_{\pm0.25}$ &$15.49_{\pm0.38}$ &$5.58_{\pm0.13}$ \\%

 Distillation &\pmb{$93.41_{\pm0.06}$} &$76.23_{\pm0.13}$ &$51.26_{\pm0.09}$ &$32.97_{\pm0.21}$ &$24.64_{\pm0.19}$ &$20.23_{\pm0.23}$ &$14.93_{\pm0.23}$ \\ %

 Adversarial &$90.24_{\pm0.19}$ &$84.00_{\pm0.11}$ &$76.02_{\pm0.23}$ &$67.08_{\pm0.15}$ &$58.13_{\pm0.22}$ &$50.37_{\pm0.23}$ &$41.04_{\pm0.16}$ \\%

 Reference &$90.51_{\pm0.22}$ &$71.88_{\pm0.32}$ &$57.66_{\pm0.40}$ &$41.69_{\pm0.26}$ &$28.29_{\pm0.17}$ &$19.80_{\pm0.22}$&$11.74_{\pm0.14}$\\%

 Adv+Ref &$87.11_{\pm0.08}$ &$82.56_{\pm0.11}$ &$77.21_{\pm0.16}$ &$71.65_{\pm0.21}$ &$65.17_{\pm0.36}$ &$58.65_{\pm0.24}$ &$46.25_{\pm0.18}$ \\%

 VIB-R &$90.57_{\pm0.07}$ &$79.06_{\pm0.34}$ &$63.62_{\pm0.25}$ &$48.99_{\pm0.33}$ &$36.59_{\pm0.13}$ &$26.93_{\pm0.21}$ &$14.39_{\pm0.25}$ \\ %

 VIB-M &$91.47_{\pm0.14}$ &$79.61_{\pm0.15}$ &$68.18_{\pm0.29}$ &$62.47_{\pm0.41}$ &$57.98_{\pm0.10}$ &$52.23_{\pm0.16}$ &$40.10_{\pm0.33}$\\ %

 VIB-R(Ref) &$90.98_{\pm0.07}$ &$79.87_{\pm0.06}$ &$69.35_{\pm0.30}$ &$62.58_{\pm0.23}$ &$56.57_{\pm0.15}$ &$51.90_{\pm0.44}$ &$44.82_{\pm0.34}$\\%

 VIB-M(Ref) &$91.49_{\pm0.05}$ &$80.55_{\pm0.19}$ &$70.52_{\pm0.18}$ &$60.05_{\pm0.24}$ &$50.94_{\pm0.20}$ &$44.01_{\pm0.09}$ &$35.02_{\pm0.18}$ \\%

 VIB-M(PGD) &$74.55_{\pm0.03}$ &$68.74_{\pm0.12}$ &$62.92_{\pm0.27}$ &$56.32_{\pm0.37}$ &$50.29_{\pm0.19}$ &$44.46_{\pm0.32}$ &$33.91_{\pm0.22}$ \\%

 VIB-M(Ref+FGSM) &$87.29_{\pm0.09}$ &$82.47_{\pm0.11}$ &$77.22_{\pm0.18}$ &$71.54_{\pm0.34}$ &$66.11_{\pm0.33}$ &$60.35_{\pm0.18}$ &$50.94_{\pm0.29}$ \\%

 VIB-M(Ref+PGD) &$89.57_{\pm0.16}$ &\pmb{$85.65_{\pm0.08}$} &\pmb{$80.92_{\pm0.13}$} &\pmb{$75.78_{\pm0.20}$} &\pmb{$69.86_{\pm0.20}$} &\pmb{$63.89_{\pm0.20}$} &\pmb{$54.00_{\pm0.19}$} \\ %

 \hline
\end{tabular}
\end{table*}

The PGD attack results on CIFAR-10 are exhibited in Table \ref{Tab:pgd_cifar10}.

\iffalse
\begin{table}
\caption{Running time per epoch on CIFAR-10 (CPU: i7-9750H, GPU: RTX2060).}
\label{Tab:time_cifar10}
\centering
\begin{tabular}{|c|c|c|c|} \hline
 Method & Time per epoch & Total epoch & Total time\\ \hline
 Baseline &34.526s &200 &1.92h\\
 Distillation &66.351 &170 &5.05h \\
 Adversarial &502.349/273.049s &100 &9.77h \\
 Adv+Ref &349.456s & &h \\
 VIB-R &28.831 &250 &2.00h\\
 VIB-M & &\\
 VIB-R(Ref) &63.291 &300  \\
 VIB-M(Ref) &71.128 &200 \\
 VIB-M(PGD) &134.058 &300\\
 VIB-M(Ref+PGD) &376.235 &140\\
 VIB-M(Ref+FGSM) &143.916 &80 \\
 \hline
\end{tabular}
\end{table}
\fi

From the results, firstly, we can see that, when both training with clean images, the VIB models using pre-learnt soft labels are more robust than the VIB models using hard labels under both FGSM and PGD attack.

Secondly, applying Adversarial Training to our method can further improve the VIB's performance, in particular under PGD attack and does not damage the performance on the clean samples.

Thirdly, compared to standard Adversarial Training, our proposed method can be accelerated via being trained with one-step adversarial samples, which will not lead to label leaking effect, and the performance is still comparable to the very time-consuming PGD Adversarial Training.

\section{Ablation Study}

\begin{table*}[ht!]
\small
\caption{Performance improvements ($\%$).}
\label{Tab:Ablation}
\centering
\begin{tabular}{ |c|c|c|c|} \hline
Method & Using VIB & Using Adv & Using Reference\\ \hline
MNIST (FGSM)   &1.46 &3.47 &1.53 \\
MNIST (PGD)    &12.97 &131.22 &15.05\\
CIFAR10 (FGSM) &-4.94 &-10.85 &4.92 \\
CIFAR10 (PGD)  &20.41 &33.77  &22.36 \\
\hline
\end{tabular}
\end{table*}

\section{Conclusions}

In this work, we propose a novel robust VIB model for defending adversarial attacks. Our method uses the pre-learnt soft labels as the target and adversarial images as the input, which enables our model to more resilient to adversarial than other existing methods. For validation, we test our method on MNIST and CIFAR-10 under a series of FGSM attack and PGD attack, and the obtained results show that our method outperform other previous methods.

\bibliographystyle{plain}

\subsubsection{Results on CIFAR-10}

\begin{table*}[ht!]
\small
\caption{Models' classification test accuracy results ($\%$) with the FGSM attack accuracy on the CIFAR-10 dataset.}
\label{Tab:fgsm_cifar10}
\centering
\begin{tabular}{ |l|c|c|c|c|c|c|c|} \hline
 Method & Clean & $\epsilon=0.01$ & $\epsilon=0.02$ & $\epsilon=0.03$ & $\epsilon=0.04$ & $\epsilon=0.05$ & $\epsilon=0.06$ \\ \hline

 Baseline &$92.88_{\pm0.13}$ &$86.36_{\pm0.14}$ &$83.67_{\pm0.18}$ &$81.68_{\pm0.23}$ &$79.99_{\pm0.18}$ &$78.52_{\pm0.14}$ &$76.88_{\pm0.41}$ \\%

 Baseline(KD) &\pmb{$93.41_{\pm0.06}$} &$84.34_{\pm0.24}$ &$80.42_{\pm0.20}$ &$77.58_{\pm0.29}$ &$75.44_{\pm0.30}$ &$73.56_{\pm0.19}$ &$71.74_{\pm0.64}$ \\%

 Baseline(PGD) &$90.24_{\pm0.19}$ &$84.71_{\pm0.22}$ &$79.36_{\pm0.31}$ &$75.49_{\pm0.21}$ &$72.72_{\pm0.11}$ &$70.48_{\pm0.17}$ &$68.99_{\pm0.14}$ \\%

 Baseline(REF) &$90.51_{\pm0.22}$ &$84.71_{\pm0.34}$ &$81.81_{\pm0.21}$ &$79.52_{\pm0.17}$ &$77.31_{\pm0.26}$ &$75.15_{\pm0.29}$ &$73.63_{\pm0.22}$\\  %

 Baseline(PGD+REF) &$87.11_{\pm0.08}$ &$83.05_{\pm0.21}$ &$79.29_{\pm0.04}$ &$76.40_{\pm0.17}$ &$73.91_{\pm0.22}$ &$71.77_{\pm0.14}$ &$69.76_{\pm0.22}$ \\%

 VIB-R &$90.57_{\pm0.07}$ &$84.66_{\pm0.32}$ &$81.24_{\pm0.14}$ &$78.66_{\pm0.27}$ &$76.47_{\pm0.19}$ &$74.84_{\pm0.25}$ &$73.03_{\pm0.17}$ \\ %

 VIB-M &$91.47_{\pm0.14}$ &$84.79_{\pm0.24}$ &$81.94_{\pm0.15}$ &$79.98_{\pm0.22}$ &$78.16_{\pm0.13}$ &$76.52_{\pm0.11}$ &$75.26_{\pm0.16}$ \\%

 VIB-R(REF) &$90.98_{\pm0.07}$ &$86.82_{\pm0.16}$ &$84.80_{\pm0.23}$ &$83.09_{\pm0.19}$ &$81.62_{\pm0.20}$ &$80.37_{\pm0.20}$ &$79.44_{\pm0.21}$ \\ %

 VIB-M(REF) &$91.49_{\pm0.05}$ &\pmb{$87.42_{\pm0.17}$} &\pmb{$85.32_{\pm0.13}$} &\pmb{$84.01_{\pm0.32}$} &\pmb{$82.42_{\pm0.26}$} &\pmb{$80.97_{\pm0.24}$} &\pmb{$80.05_{\pm0.42}$} \\ %

 VIB-M(PGD) &$74.55_{\pm0.03}$ &$69.43_{\pm0.26}$ &$63.74_{\pm0.26}$ &$58.88_{\pm0.28}$ &$53.87_{\pm0.31}$ &$49.67_{\pm0.28}$ &$46.10_{\pm0.30}$ \\ %

 VIB-M(REF+FGSM) &$87.29_{\pm0.09}$ &$83.62_{\pm0.19}$ &$81.56_{\pm0.18}$ &$80.11_{\pm0.09}$ &$78.61_{\pm0.16}$ &$77.48_{\pm0.18}$ &$76.17_{\pm0.26}$ \\%

 VIB-M(REF+PGD) &$89.57_{\pm0.16}$ &$86.24_{\pm0.13}$ &$83.69_{\pm0.18}$ &$82.06_{\pm0.17}$ &$80.23_{\pm0.36}$ &$78.78_{\pm0.30}$ &$77.52_{\pm0.18}$ \\%

\hline
\end{tabular}
\end{table*}

\begin{table*}[ht!]
\small
\caption{Models' classification test accuracy results ($\%$) with the PGD attack accuracy on the CIFAR-10 dataset.}
\label{Tab:pgd_cifar10}
\centering
\begin{tabular}{ |l|c|c|c|c|c|c|c|} \hline
 Method & Clean & $\epsilon=0.01$ & $\epsilon=0.02$ & $\epsilon=0.03$ & $\epsilon=0.04$ & $\epsilon=0.05$ & $\epsilon=0.06$ \\ \hline

 Baseline &$92.88_{\pm0.13}$ &$77.28_{\pm0.18}$ &$63.96_{\pm0.26}$ &$46.42_{\pm0.28}$ &$27.18_{\pm0.25}$ &$15.49_{\pm0.38}$ &$8.92_{\pm0.16}$ \\%

 Baseline(KD) &\pmb{$93.41_{\pm0.06}$} &$76.23_{\pm0.13}$ &$51.26_{\pm0.09}$ &$32.97_{\pm0.21}$ &$24.64_{\pm0.19}$ &$20.23_{\pm0.23}$ &$17.15_{\pm0.33}$ \\ %

 Baseline(PGD) &$90.24_{\pm0.19}$ &$84.00_{\pm0.11}$ &$76.02_{\pm0.23}$ &$67.08_{\pm0.15}$ &$58.13_{\pm0.22}$ &$50.37_{\pm0.23}$ &$44.79_{\pm0.20}$ \\%

 Baseline(REF) &$90.51_{\pm0.22}$ &$71.88_{\pm0.32}$ &$57.66_{\pm0.40}$ &$41.69_{\pm0.26}$ &$28.29_{\pm0.17}$ &$19.80_{\pm0.22}$&$14.60_{\pm0.19}$\\%

 Baseline(PGD+REF) &$87.11_{\pm0.08}$ &$82.56_{\pm0.11}$ &$77.21_{\pm0.16}$ &$71.65_{\pm0.21}$ &$65.17_{\pm0.36}$ &$58.65_{\pm0.24}$ &$52.40_{\pm0.21}$ \\%

 VIB-R &$90.57_{\pm0.07}$ &$79.06_{\pm0.34}$ &$63.62_{\pm0.25}$ &$48.99_{\pm0.33}$ &$36.59_{\pm0.13}$ &$26.93_{\pm0.21}$ &$19.86_{\pm0.15}$ \\ %

 VIB-M &$91.47_{\pm0.14}$ &$79.61_{\pm0.15}$ &$68.18_{\pm0.29}$ &$62.47_{\pm0.41}$ &$57.98_{\pm0.10}$ &$52.23_{\pm0.16}$ &$46.08_{\pm0.69}$\\ %

 VIB-R(REF) &$90.98_{\pm0.07}$ &$79.87_{\pm0.06}$ &$69.35_{\pm0.30}$ &$62.58_{\pm0.23}$ &$56.57_{\pm0.15}$ &$51.90_{\pm0.44}$ &$47.70_{\pm0.18}$\\%

 VIB-M(REF) &$91.49_{\pm0.05}$ &$80.55_{\pm0.19}$ &$70.52_{\pm0.18}$ &$60.05_{\pm0.24}$ &$50.94_{\pm0.20}$ &$44.01_{\pm0.09}$ &$39.21_{\pm0.47}$ \\%

 VIB-M(PGD) &$74.55_{\pm0.03}$ &$68.74_{\pm0.12}$ &$62.92_{\pm0.27}$ &$56.32_{\pm0.37}$ &$50.29_{\pm0.19}$ &$44.46_{\pm0.32}$ &$38.94_{\pm0.18}$ \\%

 VIB-M(REF+FGSM) &$87.29_{\pm0.09}$ &$82.47_{\pm0.11}$ &$77.22_{\pm0.18}$ &$71.54_{\pm0.34}$ &$66.11_{\pm0.33}$ &$60.35_{\pm0.18}$ &$55.53_{\pm0.14}$ \\%

 VIB-M(REF+PGD) &$89.57_{\pm0.16}$ &\pmb{$85.65_{\pm0.08}$} &\pmb{$80.92_{\pm0.13}$} &\pmb{$75.78_{\pm0.20}$} &\pmb{$69.86_{\pm0.20}$} &\pmb{$63.89_{\pm0.20}$} &\pmb{$59.87_{\pm0.24}$} \\ %

 \hline
\end{tabular}
\end{table*}

\iffalse
\begin{table}
\caption{Running time per epoch on CIFAR-10 (CPU: i7-9750H, GPU: RTX2060).}
\label{Tab:time_cifar10}
\centering
\begin{tabular}{|c|c|c|c|} \hline
 Method & Time per epoch & Total epoch & Total time\\ \hline
 Baseline &34.526s &200 &1.92h\\
 Distillation &66.351 &170 &5.05h \\
 Adversarial &502.349/273.049s &100 &9.77h \\
 Adv+Ref &349.456s & &h \\
 VIB-R &28.831 &250 &2.00h\\
 VIB-M & &\\
 VIB-R(Ref) &63.291 &300  \\
 VIB-M(Ref) &71.128 &200 \\
 VIB-M(PGD) &134.058 &300\\
 VIB-M(Ref+PGD) &376.235 &140\\
 VIB-M(Ref+FGSM) &143.916 &80 \\
 \hline
\end{tabular}
\end{table}
\fi

\end{document}

\section{Introduction}

From the information-theoretical view of deep learning, the internal representation of some hidden layers is considered as a stochastic encoding $Z$ of the input $X$, and the goal of learning is to obtain an encoding that is maximally informative about the target $Y$. Therefore, information bottleneck (IB) approaches \cite{tishby2000information} force $Z$ to be a minimal latent representation of $X$ for predicting $Y$.
The
IB
methods have recently attracted significant attention in their ability of learning good hidden representations. See \cite{goldfeld2020information} for a recent survey. The
IB methods
can also
be used in sparse neural networks using Dropout \cite{achille2018information} and to compress neural networks \cite{dai2018compressing}. \cite{peng2019variational} developed a Variational Discriminator Bottleneck which can be applied to many applications, such as imitation learning, inverse reinforcement learning, and generative adversarial networks.

Variational information bottleneck (VIB), started from \cite{alemi2017deep}, leverages variational inference to implement the IB method \cite{tishby2000information,tishby2015deep}. VIB models (or stochastic neural networks fit using VIB method), are believed to be more robust than deterministic models (e.g., conventional deep neural networks) because VIB learns the disentangled latent representations $Z$ which ignore the details of the input $X$ as many as possible \cite{alemi2017deep}. This has been improved by \cite{federici2020learning} with a multi-view approach where each view represents a task.
Specifically, the robustness of learned representations can be improved by eliminating task-related information. However, additional information about multiple tasks is required.
Moreover, \cite{belghazi2018mutual} leverages Mutual Information Neural Estimation to
help with the implementation of VIB, without concerning  robustness performance. \cite{goyal2020variational} proposes Variational Bandwidth Bottleneck which compresses only part of the input, and shows that it has better generalization for reinforcement learning tasks. \cite{9088132} considers the robustness of VIB by adapting the variational bound of \cite{alemi2017deep}, but this method is only implemented via fully-connected layers and tested on FGSM attacks.

In this paper, we propose a novel method that learns robust VIB models that is radically different from the above methods (i.e., without taking the multi-view approach, is able to work with convolutional networks, and is robust against more involved attacks). The main idea is to use soft labels as the learning targets. The soft labels are pre-learned from a reference network. Specifically, they are the numerical outputs of the last hidden layer of the reference network but not the logits of the penultimate layer as in knowledge distillation \cite{hinton2015distilling,papernot2016distillation}. By using soft labels as training targets, the VIB training has been transformed from the original classification problem into a maximum likelihood estimation based regression problem so the outputs share the same manifold. Moreover, \cite{vincent2008extracting,vincent2010stacked} showed that an autoencoder structure trained with perturbed images as the input and the clean images as the target can be used to reconstruct \lq\lq{}repaired\rq\rq{} images from perturbed ones. Our method is inspired by this method to some extent, but with completely different purposes.

There are two major advantages of soft labels. On the one hand, they would lessen the confidence of the model and prevent it from overfitting~\cite{pmlr-v97-cohen19c}. On the other hand, they improve the efficiency of model training. It is widely acknowledged that model training with the single-step adversarial examples can easily lead to the risk of \emph{label leaking}~\cite{kurakin2016adversarial}, that is, the model's performance on adversarial examples is better than that with the clean examples. Therefore, many other studies employ iterative adversarial examples but the model training process becomes computationally expensive. Using the soft labels pre-learned from the reference network can help circumvent the risk of label leaking with the single-step adversarial examples, whose generation process such as FGSM \cite{goodfellow2014explaining} is much simpler and quicker.

Moreover, in the VIB, we replace the reparameterization method \cite{kingma2014auto} with the mutual information neural estimation (in short MINE) \cite{belghazi2018mutual}. The latter does not rely on Gaussian assumption for the posterior so it also improves the robustness of VIB. To validate our PROP method, extensive experiments are performed on MNIST and CIFAR-10 datasets. We examine both FGSM and PGD attacks at different scales, and compare the PROP method with several widely used methods such as adversarial training, knowledge distillation. The experimental results verify the robustness improvement of our method.

\iffalse
DNNs are
susceptible to adversarial attacks \cite{goodfellow2014explaining}, \cite{madry2018towards}, \cite{carlini2017towards}. In order to enable DNNs
to be
resilient to adversarial attacks, researchers have PROP many approaches. For instance, \cite{kurakin2016adversarial} PROP to use adversarial samples, i.e. PGD samples, as the training input and this method is called Adversarial Training. \cite{papernot2016distillation} used Knowledge Distillation to acquire the soft labels as the training targets to improve the robustness of deep neural networks. Another approach PROP by \cite{raghunathan2018certified} is based on Adversarial Training and regularization.

\fi

\section{Model}

Let $X$ be the input, $Z$ the latent representation of $X$, and $Y$ the target. Without loss of generality, we assume an information Markov chain $X \rightarrow Z \rightarrow Y$. The information bottleneck~\cite{tishby2000information,tishby2015deep} solves the following optimisation problem:
\begin{align}
\label{Eq:IB_1}
\text{max}  & ~ I(Z;Y),   \\
\label{Eq:IB_2}
\text{s.t.} & ~ I(Z;X) \leq I_C,
\end{align}
where $I$ is the mutual information and $I_C$ is a constant representing the information constraint. This is equivalent to maximizing the Lagrangian objective function
\begin{align}
\label{Eq:IB_Lag}
\mathscr{L}_{\text{IB}} = I(Z;Y) - \beta \big( I(Z;X) - I_C \big).
\end{align}
Based on Eq.~(\ref{Eq:IB_Lag}), the VIB model~\cite{alemi2017deep} computes $I(Z;X)$ via an encoder neural network and computes $I(Z;Y)$ via a classifier neural network, respectively. Our study aims to enhance the VIB model's robustness in face of adversarial attacks such as the projected gradient
descent (PGD) attack \cite{kurakin2016adversarial} and the fast gradient sign method (FGSM) attack \cite{goodfellow2014explaining}.

\begin{figure}[t]
\centering
\includegraphics[width=1\linewidth]{robust_vib.pdf}
\caption{Overview of our PROP robust learning method for the VIB, where $X$ is the input, $Y$ is the output, $f_E$ is the encoder network, $XZ$ (red) represents the joint samples, $XZ$ (yellow) represents the independent samples, $f_D$ is the discriminator, $f_R$ is the reference network, $f_C$ is the classifier, and $o_{\text{REF}}$ and $o_{\text{VIB}}$ are the soft labels from the reference network and the classifier, respectively.}
\label{fig:framework}
\end{figure}

The overview of our proposed method is presented in Fig.\ref{fig:framework}. Our model has a reference network $f_R$, an encoder $f_E$, a discriminator $f_D$ and a classifier $f_C$. It consists of three stages. In Stage I, a reference network $f_{R}(x)$ is pre-trained with clean samples $x$ to learn the soft labels $o_{\text{REF}}$, which are the output of the last hidden layer. The Stage II aims to compute the mutual information loss $I(Z;X)$ using the Mutual Information Neural Estimation (MINE) method~\cite{belghazi2018mutual}. The input $x'$ is firstly perturbed with either random noise (normal training) or adversarial attacks (adversarial training) to train an encoder network $f_E$. Then, the joint samples $XZ_\text{Joint}$ obtained by the encoder $f_E$ are shuffled to obtain the independent samples $XZ_\text{Ind}$. Afterwards, the discriminator $f_D$ is used to discriminate $XZ_\text{Joint}$ from $XZ_\text{Ind}$. The purpose of using $XZ_\text{Ind}$ and $f_D$ is due to our adoption of the Donsker-Varadhan representation of the Kullback–Leibler (KL) divergence. In Stage III, the learned soft labels are leveraged to compute the mean squared error loss $||o_{\text{VIB}}-o_{\text{REF}}||^2$ as the approximation of $I(Z;Y)$, where $o_{\text{VIB}}$ is the output of the last hidden layer of the VIB classifier $f_C$.

We jointly train the encoder $f_E$, the discriminator $f_D$ and the classifier $f_C$ with the loss function $\mathscr{L}_{\text{PROP}}$. First, for a set of inputs $x_i$, we define the following loss function for the VIB model with reference:
\begin{align}
\label{Eq:RVIB_loss}
\mathscr{L} & =
\sum_{i}^{N} \bigg[ \Big\| f_R(x_i)- f_C \big (f_E(x'_i) \big) \Big\|^2  \nonumber  \\
& \qquad  + \beta \Big( \mathop{{}\mathbb{E}}{_{P_{XZ}}}\big [f_D \big] - \log \mathop{\mathbb{E}}{_{_{P_{X}\otimes P_{Z}}}\big[e^{f_D} \big]} \Big) \bigg ]
\end{align}
where $N$ is the size of mini batch, and $\beta$ is a non-negative Lagrangian multiplier determining the optimisation trade-off. Based on Eq. (\ref{Eq:RVIB_loss}), let $\mathscr{L}_{\text{Clean}}(x_i^{\text{Clean}})$ and $\mathscr{L}_{\text{Adv}}(x_i^{\text{Adv}})$ be the loss functions trained on clean samples and adversarial samples, respectively. Then, we can have the loss function $\mathscr{L}_{\text{PROP}}$ for the adversarial training:
\begin{align}
\label{Eq:Final_loss}
\mathscr{L}_{\text{PROP}} & = \frac{\sum_{i \in \text{Clean}}{\mathscr{L}_{\text{Clean}}(x_i^{\text{Clean}})}  +\lambda \sum_{i \in \text{Adv} }{\mathscr{L}_{\text{Adv}}(x_i^{\text{Adv}})}}{m-k+\lambda k},
\end{align}
where $m$ is the total batch size, $k$ is the adversarial sample size, $\lambda$ is the relative weight. We call our new VIB method REF-VIB\ to highlight the distinct feature of pre-learned soft labels. Algorithm~\ref{Alg:Algorithm} presents the overall training procedure. In the following, we explain the components of our REF-VIB\ method with details.

\begin{algorithm}[tb]
\caption{REF-VIB training}
\label{Alg:Algorithm}
\textbf{Requirements}: $x$ (images) , $y$ (labels)

\begin{algorithmic}[1]
\WHILE{ $\text{epoch} \leq \text{max epoch}$}
    \STATE{train a reference neural network $f_R(x)$;}
\ENDWHILE

\WHILE{ $\text{epoch} \leq \text{max epoch}$}

    \STATE{If $i$ $\in$ Clean:}
        \STATE{\quad $x'_i$ $\gets$ $x_i$ + $\epsilon$}
    \STATE{If $i$ $\in$ Adversarial:}
        \STATE{\quad $x'_i$ $\gets$ attack $x_i$}

        \STATE{$xz_{\text{Joint}}$ $\gets$ $f_E(x'_i)$}

        \STATE{$xz_{\text{Ind}}$ $\gets$ shuffle $xz_{\text{Joint}}$ along the batch axis}

        \STATE{$I(Z;X)$ $\gets$ $\mathop{{}\mathbb{E}}{_{P_{XZ}}}[f_D(xz_{\text{Joint}})]$ \\
        \qquad \qquad \qquad $ - \log\{ \mathop{\mathbb{E}}{_{_{P_{X}\otimes P_{Z}}}[e^{f_D(xz_{\text{Ind}})}]} \}$}

        \STATE{$I(Z;Y)$ $\gets$ $\|f_C(xz_{\text{Joint}})-f_R(x_i)\|^2$}

        \STATE{Minimize final loss function $\mathscr{L}_{\text{PROP}}$}

\ENDWHILE

\end{algorithmic}
\end{algorithm}

The VIB model~\cite{alemi2017deep} implements Eq. (\ref{Eq:IB_Lag}) with reparameterization method \cite{kingma2014auto} via deep neural networks.
It approximates $I(Z;X)$ and $I(Z;Y)$ with an encoder network and a decoder network, respectively,
by deriving the variational bounds for them.
The upper bound of $I(Z;X)$ yields
\begin{align}
\label{Eq:I_ZX_bound}
I(Z;X) \leq \mathbb{E}_x \left[D_\text{KL} \big(p(z|x)||q(z)\big)\right],
\end{align}
where $D_\text{KL}$ is the Kullback-Leibler (KL) divergence, $q(z)$ is an uninformative prior distribution.
Similarly, the lower bound of $I(Z;Y)$ can be computed by
\begin{align}
\label{Eq:I_ZY_bound}
I(Z;Y) \geq & \mathbb{E}_{p(y,z)} \big [ \log \left\{ p(y|z) \right\} \big].
\end{align}
Combining Eq. (\ref{Eq:I_ZX_bound}) and Eq. (\ref{Eq:I_ZY_bound}), the loss function becomes
\begin{align}
\label{Eq:OP_VIB_Repa}
 \mathscr{L}_{\text{VIB}}
 = &  \
 \mathbb{E}_{\mathscr{D}} \bigg[ \mathbb{E}_{z \sim p(z|x)} \Big[-\log\{ p(y|z)\} \Big] \nonumber \\
 & \hspace*{50pt} + \beta D_\text{KL}\big(p(z|x)|| q(z)\big ) \bigg],
\end{align}
where $p(z|x)$ is the encoder and $p(y|z)$ is the classifier.

In the following, we introduce how do we update Eqs. (\ref{Eq:I_ZX_bound})-(\ref{Eq:I_ZY_bound}) by taking the MINE method and the reference network, respectively. This will lead to the update of $\mathscr{L}_{\text{VIB}}$ into $\mathscr{L}$ of Eq. (\ref{Eq:RVIB_loss}), and VIB into the proposed REF-VIB\ method.

\paragraph{Mutual Information Neural Estimation}

As for estimating the mutual information $I(X;Z)$, apart from using the re-parameterization method, we can alternatively use the Mutual Information Neural Estimation (MINE) method~\cite{belghazi2018mutual}. By using the MINE method, the Gaussian assumption on the posterior $p(z|x)$ can be relaxed, so the approximation of the true posterior $p(z|x)$ can be more accurate.

First, we introduce the Donsker-Varadhan representation of the KL divergence \cite{donsker1983asymptotic,belghazi2018mutual}, defined as below
\begin{align}
\label{Eq:DV}
D_\text{KL}(P||Q) = \underset{T:\Omega \rightarrow \mathop{{}\mathbb{R}}}{\mathrm{sup}} \mathop{{}\mathbb{E}}{_{P}}\big[T\big] - \log \mathop{\mathbb{E}}{_{Q}\big[e^T\big]},
\end{align}
where $P$ is the posterior distribution, $Q$ is the prior distribution and $T$ is a discriminator network used to distinguish joint samples from independent samples. This can be implemented via a discriminator network using both joint and independent samples, as presented in the Stage II of Fig.~\ref{fig:framework}.

The reason why we consider the KL divergence is because $I(X;Z)$ can be defined as follow:
\begin{align}
\label{Eq:I_ZX}
I(Z;X) = & \ D_\text{KL} \big (\mathop{{}\mathbb{{P_{XZ}}}}||\mathop{{}\mathbb{{P_{X}}}} \otimes \mathop{{}\mathbb{{P_{Z}}}} \big )
\end{align}
where $\mathbb{{P_{XZ}}}$ is the joint distribution, $\mathbb{{P_{X}}}$ and $\mathbb{{P_{Z}}}$ are the marginal distributions
Therefore, as opposed to the variational upper bound in Eq. (\ref{Eq:I_ZX_bound}), we use bound $I(Z;X) \geq I_\theta(Z;X)$ to estimate $I(Z;X)$ according to Eq. (\ref{Eq:DV}), such that
\begin{equation}
\label{Eq:I_ZX_DV}
I_\theta(Z;X) = \underset{\theta \in \Theta}{\mathrm{sup}} \mathop{{}\mathbb{E}}{_{P_{XZ}}}\big[T_\theta\big] - \log \mathop{\mathbb{E}}{_{_{P_{X}\otimes P_{Z}}}\big[e^{T_\theta}\big]},
\end{equation}
where $T_\theta$ is a family of functions parametrized by a discriminator network (i.e., $f_D$ shown in Fig.~\ref{fig:framework}) with parameters $\theta \in \Theta$.  Hereafter, we omit $\theta$ for simplicity. To solve Eq. (\ref{Eq:I_ZX_DV}), we use the mini batches of the input to feed an encoder network so as to acquire the samples of the joint distribution $P_{XZ}$ (positive samples). While, the samples of $P_{X}\otimes P_{Z}$ can be obtained through shuffling the samples of $P_{XZ}$ along the batch axis (negative samples).

To simplify the discussion, we use \emph{VIB-M} to denote the VIB model trained using the MINE method and \emph{VIB-R} to denote the VIB model trained through the re-parameterization method. Combining Eq. (\ref{Eq:I_ZY_bound}) and Eq. (\ref{Eq:I_ZX_DV}), the loss function of the VIB-M model can be obtained as follows:
\begin{align}
\label{Eq:OP_VIB_Mine}
& \mathscr{L}_{\text{VIB-M}} =  \mathbb{E}_{\mathscr{f_D}} \bigg[ \mathbb{E}_{z \sim p(z|x)} \big[-\log\{ p(y|z)\} \big]  \nonumber \\
& \hspace*{45pt} + \beta \Big (\mathop{{}\mathbb{E}}{_{P_{XZ}}}\big[f_D\big] - \log \mathop{\mathbb{E}}{_{_{P_{X}\otimes P_{Z}}}\big[e^{f_D}\big]} \Big ) \bigg].
\end{align}

It should be noted that $\mathscr{L}_{\text{VIB-M}}$ is different from $\mathscr{L}$ in Eq. (\ref{Eq:RVIB_loss}) because of the term $\mathbb{E}_{z \sim p(z|x)} \big[-\log\{ p(y|z)\} \big]  \nonumber$, which will be replaced with $\|f_R(x_i)- f_C \big (f_E(x'_i) \big)\|^2 $ after using the reference network, discussed next.

\paragraph{Reference Network}
\label{sec:referencenetwork}

Another key novelty of our method is that we use a reference network to pre-learn the soft labels on clean samples, and then use the pre-learned  soft labels -- instead of original hard labels -- as the learning targets in the VIB model. These soft labels can be regarded as a set of numerical vectors corresponding to each input image in space. Note that the soft labels used in our method are s different from the labels used in Knowledge Distillation \cite{hinton2015distilling,mishra2017apprentice}, which are logits.

The reason of doing this is that hard labels only retain the discrete categorical information of each input but discard the more detailed information, which may impede the learning performance of VIB. Moreover, with the guidance of the soft labels learned on clean samples, the predictions of VIB model on perturbed samples are projected onto the correct manifold during the training process.

Now, according to Eq. (\ref{Eq:I_ZY_bound}), we use the pre-learned soft labels to compute the likelihood $p(y|z)$ to approximate the lower bound of $I(Z;Y)$. To this end, we let $f_R(x)$ be the output of the last hidden layer of the reference network and $f_C(xz_{\text{Joint}})=f_C \big (f_E(x'_i) \big)$ be the output of the last hidden layer of the classifier in the VIB. The distance between $f_R(\cdot)$ and $f_C(\cdot)$ can be measured by the Euclidean distance:
\begin{equation}
\label{Eq:Ref_loss}
\text{DIST}(o_{\text{REF}}, o_{\text{VIB}}) = \Big\|f_R(x_i)- f_C \big (f_E(x'_i) \big) \Big\|^2
\end{equation}

Then, Eq. (\ref{Eq:Ref_loss}) is the loss function of the downstream task in the VIB, and by replacing $\mathbb{E}_{z \sim p(z|x)} \big[-\log\{ p(y|z)\} \big]  \nonumber$ of Eq. (\ref{Eq:OP_VIB_Mine}) with it, we obtain $\mathscr{L}$ of Eq. (\ref{Eq:RVIB_loss}).

\paragraph{Adversarial Training with Reference}

Deep neural networks (DNNs) are known to be vulnerable to adversarial attacks. Even small perturbations on the clean images may significantly decrease the test performance of DNNs. The VIB model has been argued to be more robust than conventional neural networks \cite{alemi2017deep}.
Nevertheless, as shown in Section~\ref{sec:experiments}, the VIB models are still not robust enough to withstand strong adversarial attacks, such as PGD attacks. See the comparison between the VIB-M (REF) and the VIB-M (REF+FGSM) models. Therefore, we consider another effective way of defending attacks, i.e., adversarial training, which uses adversarial samples as the training input.

In this paper, we consider two popular gradient-based methods to generate adversarial samples, i.e., the FGSM attack~\cite{goodfellow2014explaining} and the PGD attack~\cite{kurakin2016adversarial}. Simply speaking, the FGSM attack is a single-step attack method and the PGD attack finds adversarial samples in an iterative way.

Let $x$ denote the clean samples and $x_\text{Adv}$ the adversarial samples. We can obtain the adversarial samples as follows:
\begin{align}
\label{Eq:Adv}
x^\text{Adv} = x + \delta
\end{align}
where $\delta$ is the perturbation. For any $x \in {X}$ with $l_p$-ball $\mathcal{B}(x, \epsilon)$ around $x: \{ x_\text{Adv} \in X: ||x_\text{Adv} - x_\text{Clean} || \leq \epsilon \}$ \cite{madaan2020adversarial}, in a VIB model, the related latent presentation $z$ is sampled from the input $x$ via an encoder and the $\hat{y}$ is corresponding prediction computed from $z$ via a classifier. This relation can be described as follow:
\begin{align}
\label{Eq:Prediction}
\hat{y} =  f_C(z) =  f_C(f_E(x))
\end{align}
where $f_E(\cdot)$ is the encoder and $f_C(\cdot)$ is the classifier.

According to Eq. (\ref{Eq:Ref_loss}), to train the VIB classifier $f_C$ with adversarial samples, we minimize the following loss function:
\begin{align}
\label{Eq:Adv_ref}
    & \ \mathbb{E}_{\{\mathscr{D}, \delta\}} \Big[ \Big\|f_C(xz^\text{Adv}_{\text{Joint}})-f_R(x)\Big\|^2 \Big] \nonumber \\
=  & \ \mathbb{E}_{\{\mathscr{D}, \delta\}} \Big[ \Big\|f_C(f_E(x^\text{Adv}))-f_R(x)\Big\|^2 \Big],
\end{align}
where $xz^\text{Adv}_{\text{Joint}}$ is the joint latent variables learned by adversarial samples. We substitute the term $\mathbb{E}_{z \sim p(z|x)} \big[-\log\{ p(y|z)\} \big]$ in Eq. (\ref{Eq:OP_VIB_Mine}) with Eq. (\ref{Eq:Adv_ref}), then according to Eq. (\ref{Eq:RVIB_loss}), we have the adversarial training loss function, $\mathscr{L}_{\text{Adv}}$ in Eq. (\ref{Eq:Final_loss}) which is the final loss function of the PROP model. Eq. (\ref{Eq:Adv_ref}) suggests that we can learn the encoder and classifier jointly for the VIB model, combining adversarial training and the reference network. Consequently, the distortion of latent variables caused by adversarial attack can be suppressed as well, which leads to better generalization so as to defend attacks.

The disadvantage of using PGD samples for training is its time consumption. One-step adversarial images, e.g., FGSM samples, are much faster to generate, however using such adversarial images to train conventional neural networks with hard labels will cause the \textbf{label leaking} effect \cite{kurakin2016adversarial}. Label leaking means the testing accuracy on adversarial images are much higher than the accuracy on clear images. This issue occurs because one-step adversarial images generated by simple transformations, and it is not hard for a deep neural network to learn such transformations, as a result, the deep neural network performs better on the adversarial images than the clean images. On the contrary, a more complex adversarial transformation, e.g., PGD, does not has this issue because its transformation is more difficult to learn. Fortunately, our method accelerates the learning process by leveraging FGSM samples without suffering from label leaking effect, as we use soft labels instead of hard labels.

\paragraph{Discussion}

We would like to discuss the rationale behind our selection of the three components in the proposed model.

\emph{Why Using the MINE Method?} We want to compute the posterior $p(z|x)$ as accurate as possible. Since the prior $q(z)$ is assumed to follow Gaussian, the reparameterization method assumes $p(z|x)$ follows Gaussian as well. However, the posterior can be more complex than the prior, and using a Gaussian distribution may not be accurate enough to approximate the prior. In the MINE method, we do not generate the posterior samples from a Gaussian distribution, instead, we generate the joint posterior samples with input perturbed by noise and use a discriminator to discriminate them from the shuffled samples. Once the model is trained, we can obtain the posterior samples with more complex shapes.

\emph{Why Need Reference Network?} In a VIB model, the input is compressed to a latent distribution by the encoder. We believe that the vulnerability of a VIB is caused by the distortion of the latent distribution. Therefore, to improve the robustness VIB, we need to minimise the distortion (a similar idea was PROP in \cite{madaan2020adversarial}, but their method is based on Bayesian neural networks and pruning). The latent distribution of VIB is Gaussian-like, and the pre-learned soft labels are numerical vectors which contain more information than the original hard labels. Hence, to suppress the distortion of the latent distribution, using the soft labels to minimize the loss function of VIB is more effusive than using the hard labels.

\emph{Why Performing Adversarial Training?} Adversarial training is known for improving the robustness of neural networks. The soft labels pre-learned by the reference network on the clean samples and the soft labels learned by the VIB model on the corrupted samples should share the same manifold. Consequently, we can leverage the pre-learned soft labels to guide the adversarial learning process of the VIB model to achieve a better performance. Our learning strategy is similar to denoising autoencoders \cite{vincent2008extracting,vincent2010stacked} to some extent. In denoising autoencoders, the corrupted images are used to train an autoencoder in order to recover
clean images, which is an unsupervised learning scheme. Our method is supervised, and we use adversarial samples as the corrupted samples and soft labels as the target instead of original clean images.

\begin{table}[t]
\centering
\small
\caption{Summary of the examined models.}
\label{Tab:model_summary}
\centering
\begin{tabular}{|m{1.1in}|m{1.9in}|}
\hline
Model & Description\\
\hline
Baseline & For the MNIST dataset, the baseline model architecture is CONV(64,3,1) + CONV(64,3,1) + FC(512) + FC(512) + FC(10); For the CIFAR-10 dataset, the baseline model is VGG16 \cite{simonyan2014very}. \\%
\hline
Baseline (KD) & The baseline model trained with KD \cite{papernot2016distillation}.\\%
\hline
Baseline (PGD) & The baseline model trained with PGD samples \cite{kurakin2016adversarial}. \\%
\hline
Baseline (REF) & The baseline model trained with a reference network.\\%
\hline
Baseline (PGD+REF) & The baseline model trained with trained with a reference network and PGD samples.\\%
\hline
VIB-R & The VIB model trained using the re-parameterization method~\cite{alemi2017deep}.\\%
\hline
VIB-M & The VIB model trained using the MINE method~\cite{belghazi2018mutual}.\\%
\hline
VIB-R (REF) & The VIB-R model trained with a reference network.\\%
\hline
VIB-M (REF) & The VIB-M model trained with a reference network.\\ %
\hline
VIB-M (PGD) & The VIB-M model trained with PGD samples.\\ %
\hline
VIB-M (REF+FGSM) & The VIB-M model trained with a reference network and FGSM samples.\\
\hline
VIB-M (REF+PGD) &  The VIB-M model trained with a reference network and PGD samples.\\% all
\hline
\end{tabular}
\end{table}

\begin{table}[t]
\small
\centering
\caption{Model training settings, where $T$ is the temperature, $\epsilon$ is the magnitude of perturbation, $\alpha$ is the attack step size, $t$ is the attack step number, and $\eta$ is the dimension of latent variables.}
\label{Tab:train_method}
\begin{tabular}{|m{1.0in}|m{2in}|}
\hline
Model & Description\\
\hline
KD & $T=1$.\\%
\hline
PGD (MNIST) & $\alpha = 0.01$, $m=100$, $k=50$, $\lambda = 0.3$, $\epsilon = 0.3$, $t = 20$. \\%
\hline
FGSM (MNIST) & $\epsilon = 0.3$. \\
\hline
PGD (CIFAR-10) & $\alpha = 0.007$, $m=100$, $k=50$, $\lambda = 0.3$, $\epsilon = 0.03$, $t = 10$. \\%
\hline
FGSM (CIFAR-10) & $\epsilon = 0.03$. \\
\hline
VIB-R & $\beta=0.001$, $\eta = 256$.\\%
\hline
VIB-M & $\beta=0.001$, $\eta = 256$.\\%
\hline
\end{tabular}
\vspace*{10pt}
\caption{Model testing settings, where $\epsilon$ is the magnitude of perturbation, $\alpha$ is the attack step size, $t$ is the attack step number.}
\label{Tab:test_method}
\centering
\begin{tabular}{|m{1.0in}|m{2in}|}
\hline
Model & Description\\
\hline
FGSM (MNIST) &  $\epsilon = 0.05$, $0.10$, $0.15$, $0.20$, $0.25$, $0.30$. \\
\hline
PGD (MNIST) & $\epsilon = 0.05$, $0.10$, $0.15$, $0.20$, $0.25$, $0.30$, $\alpha = 0.01$, $t$ = $40$, random restart. \\%
\hline
FGSM (CIFAR-10) &  $\epsilon =0.01$, $0.02$, $0.03$, $0.04$, $0.05$, $0.06$.\\
\hline
PGD (CIFAR-10) & $\epsilon=0.01$, $0.02$, $0.03$, $0.04$, $0.05$, $0.06$, $\alpha = 0.007$, $t$ = $40$, random restart.  \\%
\hline
\end{tabular}
\end{table}

\section{Experiments}
\label{sec:experiments}

\paragraph{Datasets}

Two publicly available datasets are used in the experiments: (i) the MNIST dataset~\cite{lecun1998mnist} which includes $60, 000$ gray level handwritten images of size $28 \times 28$; and (ii) the CIFAR-10 dataset~\cite{krizhevsky2009learning} which includes $60, 000$ color images of size $32 \times 32$. For each dataset, we randomly sample $50, 000$ images for model training and $10, 000$ images for model testing. The sampled images from the CIFAR-10 dataset are normalized into $[0,1]^{32\times 32\times 3}$.

\paragraph{Experimental Settings}

To have a comprehensive comparison, different machine learning algorithms using various training methods are examined, including the recent advances like knowledge distillation (KD)~\cite{papernot2016distillation}, PGD attack \cite{kurakin2016adversarial} and FGSM attack \cite{goodfellow2014explaining}. Also, as the VIB model can be estimated by either the reparameterization method or the MINE method, both training methods are included in our experiments. To simplify the discussion, an algorithm using a specific training method is considered as a model here and it is denoted by a unique name. For the reader's convenience, we present all the examined models and their descriptions in Table~\ref{Tab:model_summary}.

We adopt different settings of model architecture, training and testing. First, we use different neural networks as the baseline model and we have different architecture settings for the VIB model accordingly. For the MNIST dataset, the architecture of the baseline model is CONV(64,3,1) + CONV(64,3,1) + FC(512) + FC(512) + FC(10) and the architecture of the VIB model is CONV(32,3,1) + CONV(32,3,1) + FC(512) + FC(256) for encoder and FC(512) + FC(10) for classifier. For the CIFAR-10 dataset, the baseline models is VGG16 \cite{simonyan2014very} and the VIB model's architecture is CONV(64,3,1) + CONV(128,3,1) + CONV(256,3,1) + CONV(256,3,1) + CONV(512,3,1) +CONV(512,3,1) + FC(1024) + FC(512) + FC(256) for encoder and FC(512) + FC(1024) + FC(512) + FC(10) for classifier. Second, the baseline model and the VIB model are trained differently in these two datasets. Table~\ref{Tab:train_method} presents the details of model training settings. For the MNIST dataset, the optimizer is Adam \cite{kingma2014adam} and the learning rate is $0.001$. For the CIFAR-10 dataset, the optimizer is the stochastic gradient descent (SGD) with adaptive learning rate. Third, the trained models are evaluated under different attack settings in our datasets. As shown in Table~\ref{Tab:test_method}, the attack perturbation levels for the MNIST dataset range from 0.05 to 0.30 and are with the step size of 0.05 while they range from 0.01 to 0.06 and are with the step size of 0.01 for the CIFAR-10 dataset~\cite{kurakin2016adversarial}.

\paragraph{Results}

We run 5 randomised trials for each examined model on the MNIST and the CIFAR-10 datasets, and report the models' classification results (i.e., the average accuracy and its standard deviation) of the test data in Figs.~\ref{fig:mnist_fgsm}-\ref{fig:cifar10_pgd}.
Every curve represents a trained model, whose accuracy is gradually decreased with the increase of the parameter $\epsilon$ of either FGSM or PGD attack. A larger area under the curve (AUC) represents a more robust model.

\begin{figure}[t]
\centering
\includegraphics[width=1\linewidth]{mnist_fgsm}
\vspace*{-12pt}
\caption{Models' classification test accuracy results ($\%$) with the FGSM attacks on the MNIST dataset.}
\label{fig:mnist_fgsm}
\vspace*{10pt}
\centering
\includegraphics[width=1\linewidth]{mnist_pgd}
\vspace*{-12pt}
\caption{Models' classification test accuracy results ($\%$) with the PGD attacks on the MNIST dataset.}
\label{fig:mnist_pgd}
\end{figure}

\begin{figure}[t]
\centering
\includegraphics[width=1\linewidth]{cifar10_fgsm}
\vspace*{-12pt}
\caption{Models' classification test accuracy results ($\%$) with the FGSM attacks on the CIFAR-10 dataset.}
\label{fig:cifar10_fgsm}
\vspace*{10pt}
\includegraphics[width=1\linewidth]{cifar10_pgd}
\vspace*{-12pt}
\caption{Models' classification test accuracy results ($\%$) with the PGD attacks on the CIFAR-10 dataset. }
\label{fig:cifar10_pgd}
\end{figure}

There are a few insights we would like to share.
First of all, the VIB model shows a good predictive capability. However, it is still vulnerable to strong attacks like the PGD attack. From the results on the MNIST dataset, we can see that the regularly-trained VIB models, i.e., the VIB-R and the VIB-M models, outperform the corresponding baseline models under the FGSM attacks. However, they become worse under the PGD attacks. Our experimental results also show that the MINE method generally outperforms the reparamterization method for training the VIB model because the former can provide more complex posteriors.

It is generally believed that training a deep learning model with adversarial examples can improve the model's robustness, particularly, under the PGD attack. However, according to the MNIST results, the VIB (PGD) model does not achieve the best on any case. That indicates a naive adoption of adversarial training cannot improve the performance of the VIB model.  In fact, the results obtained by the VIB-M (REF+FGSM) and the VIB-M (REF+PGD) suggest only when the VIB model is trained with the support of reference network, the best performance can be achieved. \emph{By contrast, using our proposed methods, VIB-M(REF+FGSM/PGD), the robustness of the VIB models can be significantly improved under both FGSM and PGD attacks.} Between the VIB-M (REF+FGSM) and the VIB-M (REF+PGD) models, the former enjoys faster training while maintains comparable performance. From the results on CIFAR-10, we see that VIB-R/M do not necessarily outperform the baselines under either the FGSM or the PGD attacks. The VIB-R/M (REF) models have better performance under the FGSM attacks, but do not perform very well under the PGD attack. Our proposed methods, the VIB-M (REF+FGSM/PGD) models, have similar excellent performance against both FGSM and PDG attacks.

Overall, it can be seen that, first, learning with reference can improve the robustness of the VIB model; second, compared with other methods, adversarial training is more effective to defend the PGD attacks and using reference can improve the performance of adversarial training; third, as opposed to using hard labels, our method can use a single-step adversarial training method to accelerate the training process, i.e. the FGSM training, without causing the label leaking effect.

\section{Conclusion}

In this paper, we propose a novel method for training the VIB model which can produce accurate predictions as well as be robust under adversarial attacks. The idea is to train the VIB model by using the reference network and the MINE method. The former enables the mapping of both clean and perturbed inputs onto a smooth output manifold, so that the learned VIB model is smoother (and more robust). Also, this enables a fast training with FGSM-based method instead of PGD-based, without compromising the robustness performance. To the best of our knowledge, this is the very first study that integrates them into a VIB-based framework. The experimental results show that our proposed model outperforms the state-of-the-art benchmarked models.

\bibliographystyle{plain}

redpajama_set_name: RedPajamaArXiv