SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Huafeng Kuang $^{1}$ , Hong Liu $^{2*}$ , YongJian Wu $^{3}$ , Shin'ichi Satoh $^{2}$ , Rongrong Ji $^{1\dagger}$
<sup>1</sup>Key Laboratory of Multimedia Trusted Perception and Efficient Computing, Ministry of Education of China, Xiamen University, 361005, P.R. China
$^{2}$ National Institute of Informatics, Tokyo, 101-8430, Japan
$^{3}$ Youtu Laboratory, Tencent, Shanghai, 200233, China
skykuang@stu.xmu.edu.cn, lynnliu.xmu@gmail.com
littlekenwu@tencent.com, satoh@nii.ac.jp, rrji@xmu.edu.cn
Previous studies have shown that optimizing the information bottleneck can significantly improve the robustness of deep neural networks. Our study closely examines the information bottleneck principle and proposes an Information Bottleneck Distillation approach. This specially designed, robust distillation technique utilizes prior knowledge obtained from a robust pre-trained model to boost information bottlenecks. Specifically, we propose two distillation strategies that align with the two optimization processes of the information bottleneck. Firstly, we use a robust soft-label distillation method to increase the mutual information between latent features and output prediction. Secondly, we introduce an adaptive feature distillation method that automatically transfers relevant knowledge from the teacher model to the student model, thereby reducing the mutual information between the input and latent features. We conduct extensive experiments to evaluate our approach's robustness against state-of-the-art adversarial attackers such as PGD-attack and AutoAttack. Our experimental results demonstrate the effectiveness of our approach in significantly improving adversarial robustness. Our code is available at https://github.com/SkyKuang/IBD.
Numerous works have shown that deep neural networks (DNNs) are easily attacked by adversarial examples [56, 41, 4, 14], which involve adding imperceptible noise to inputs and causing incorrect outputs. This vulnerability of DNNs raises significant security concerns when deploying DNNs in safety-critical applications. To address this potential threat, various adversarial defense strategies have been proposed [44, 41, 72, 43, 72]. Among these defenses, adversarial training (AT) [41, 71] is a general solution for defending against adversarial attacks by incorporating adversarial examples generated by an adversarial attack into the training process. While many AT techniques can defend against sophisticated attacks, such as AutoAttack [14], a big robust generalization gap still exists between the training data and testing data [12].
Recent years have witnessed the growing popularity of the studies of the use of Information Bottleneck (IB) [57] in training robust DNNs [54, 2, 64, 68]. IB involves finding a trade-off in intermediate features  $Z$  between relevant information for the prediction  $Y$  and nuisance information about input  $X$ . The overall objective of IB is formulated as follows:
$$
\max  I (Z; Y) - \beta I (X; Z), \tag {1}
$$
* Project lead
Figure 1: Schematic of successive extraction of relevant information in IBD. The inputs are passed through a robust teacher and a student to get intermediate features (Z). These features are then passed through an adaptive feature distillation module, where an attention-based model determines similarities between the teacher and student features. Robust information from each teacher feature is transferred to the student with the attention values. The student model also minimizes cross-entropy (CE) loss using the teacher's soft output label. Consult the section 3.2 for more details.
where  $I$  denotes mutual information and  $\beta$  controls the trade-off between the two terms. Studies in [54, 64] show that IB can produce concise representations leading to better generalization, and can make intermediate features of DNNs less sensitive to input perturbations. On the other hand, variational information bottleneck (VIB) has also been found to significantly improve the robustness of models [2, 19]. In addition, [40] replaces mutual information with the Hilbert Schmidt Independence Criterion (HSIC) and calls this the HSIC Bottleneck, and [64] further investigates the HSIC bottleneck as a regularization to enhance the robustness of DNNs.
The crucial step to optimizing the IB is to calculate mutual information in an efficient and effective manner. An approach in [19] proposes a conditional entropy bottleneck (CEB) (see Eq.(6)), which enhances IB by introducing label priors in variational inference. This approach employs a backward encoder to approximate the true distribution  $p(z|y)$ . However, DNN processing can form a Markov chain  $(X \to Z \to Y)$  and obey the data processing inequality:  $H(X) \geq H(Z) \geq H(Y)$ , and the label  $Y$  often cannot satisfy the intrinsic dimension requirement [37]. As a result, directly using label information to approximate  $p(z|y)$  is imprecise. Furthermore, the backward encoder usually provides uniform priors for inputs with the same label and cannot provide customized priors based on different inputs, particularly for high-dimensional image data.
Fortunately, prior research has shown that robust pre-trained models can provide robust prior information to enhance model robustness and certainty [27, 53, 21]. Motivated by this, we revisit VIB through the lens of robustness distillation and further introduce a new IB objective: Information Bottleneck Distillation (IBD) (see Eq.(7)). To make IBD practical, we put forward two distillation strategies meant to target the two optimization processes of the information bottleneck respectively. Firstly, we utilize soft-label distillation to maximize the mutual information between intermediate features and output prediction. Secondly, we use adaptive feature distillation to restrict the mutual information between the input and intermediate features, which facilitates the transfer of appropriate knowledge from the teacher model to the student model, ultimately resulting in a more accurate estimation of the distribution of student features. A schematic diagram illustrating the IBD is shown in Figure 1.
Overall, we make the following contributions:
- Theoretically, we utilize conditional variational inference to construct a lower bound to estimate the mutual information and reformat the IB principle by using the adversarial robustness as the prior for learning features, which is termed Information Bottleneck Distillation (IBD).
- Algorithmically, to realize IBD, we propose two distillation strategies: robust soft-label and adaptive feature distillation, to match the two optimization processes of the information bottleneck, respectively.
- Experimentally, we conducted extensive experiments on various benchmark datasets such as CIFAR and ImageNet. The results show the effectiveness of our IBD in improving the robustness of DNNs against most attacks (e.g., PGD-attack and AutoAttack), and our IBD behaves more robustly than state-of-the-art methods.
Adversarial Attacks. Adversarial examples are observed when an attacker adds human-imperceptible perturbation to the inputs, which misleads DNN's output [56]. And a series of attacking algorithms are proposed [22, 4, 41, 18, 14, 38]. One of the most popular attacks is the Project Gradient Descent (PGD) attack [41], which uses multi-step projected gradient descent to generate adversarial examples:
$$
x _ {k} = \Pi \left(x _ {k - 1} + \alpha \cdot \operatorname {s i g n} \left(\nabla_ {x} \mathcal {L} _ {\mathcal {C E}} \left(h _ {\theta} \left(x _ {k - 1}\right), y\right)\right)\right), \tag {2}
$$
where  $\alpha$  is the step size,  $\Pi$  is the project function, and  $x_{k}$  is the adversarial example at the  $k$ -th step,  $h_\theta$  is a DNN model with weight  $\theta$ , and  $\mathcal{L}_{CE}$  is cross-entropy function. Another widely-used adversarial attack method is the C&W attack [4], which applies a rectifier function regularizer to generate adversarial examples near the original input. AutoAttack (AA) [14], an ensemble of different parameter-free attacks, is currently regarded as the most reliable evaluation of the adversarial robustness of the model. Black-box attacks, which do not require full knowledge of DNNs, are more practical, and previous works [67, 18] have shown that adversarial examples generated by one model can attack other models with a high probability. Query-based black-box attacks [3, 9] update the perturbation iteratively according to the attack objective. In this work, we test the robust performance under both white- and black-box attacks to verify the robustness of our proposed method.
Adversarial Robustness. To counter the threat of adversarial attacks, researchers have proposed various defense methods, including defensive distillation [44], manifold-projection [51], pre-processing [24, 69], provable defenses [45, 50], and adversarial training [11, 22, 41]. Among these, adversarial training (AT) [41] is a highly effective and widely-used solution. It involves adding adversarial examples, generated by the adversarial attack scheme, to the training data, which can be formulated by a min-max optimization problem:
$$
\min  _ {\theta} \mathbb {E} \left[ \max  _ {x _ {\mathrm {a d v}}} \mathcal {L} _ {C E} \left(h _ {\theta} \left(x _ {\mathrm {a d v}}\right), y\right) \right]. \tag {3}
$$
where  $x_{\mathrm{adv}}$  is adversarial example generated by a special attack algorithm. Many variants of AT have been proposed. [71] proposed the TRADES, which characterizes the trade-off between accuracy and robustness. [66] proposed adversarial weight perturbation (AWP), which uses weight perturbation to increase the robustness of the model. In addition, some new efforts have also been devoted from different aspects including designing new training strategies [43, 33], adversarial regularization [42, 48, 39], robustness architecture search [8, 25, 29, 30], and data augmentation [5, 46, 63].
Robustness Distillation. Knowledge distillation (KD) [28] is an effective method, wherein a well-trained teacher model is used to supervise the learning of the student model. The overall objective of KD is formulated as:
$$
\min  _ {\theta_ {S}} \mathbb {E} \left[ \mathcal {L} _ {\text {t a s k}} \left(h _ {s} (x), y\right) + \alpha \cdot L _ {\text {D i s t}} \left(h _ {s} (x), h _ {t} (x)\right) \right], \tag {4}
$$
where  $\mathcal{L}_{task}$  is the task-specific loss function for the target task and  $\mathcal{L}_{Dist}$  is the loss function that penalizes the difference between the teacher and the student. Recent studies show that KD can also be used to enhance the adversarial robustness of a student network [44]. Then, the adversarial robust distillation (ARD) method, as introduced in [21], combines adversarial training with KD to further increase the robustness of the student model under the supervision of an adversarial robust teacher network. Introspective Adversarial Distillation (IAD) [73] propose proposes a method to automatically realize the intuition of the previous reliable guidance during the adversarial distillation. And, RSLAD [74] improves the performance of ARD by using robust soft labels generated by a pre-trained teacher network to replace hard labels used in [21].
3 Information Bottleneck Distillation
In this section, we revisit the concept of information bottleneck from the perspective of robust distillation and establish the IBD objective. We further propose two distillation strategies to perform the two optimization processes of the information bottleneck, respectively. Finally, we combine IBD with AT to learn a robust model.
3.1 A Distillation View of Information Bottleneck
The information bottleneck theory has been used to explain deep models [58], compress models [16], and learn robust models [2]. To make the IB practical, a major challenge is to accurately estimate the mutual information. In [2], Variational Information Bottleneck  $(\mathrm{VIB})^{\top}$  utilizes variational inference to construct the lower bound to estimate the mutual information. VIB is formulated as:
$$
I (Z; Y) - \beta I (X; Z) \geq \underset {p (x, y) p (z | x)} {\mathbb {E}} \left[ \log q (y | z) - \beta \log \frac {p (z | x)}{q (z)} \right], \tag {5}
$$
where  $p(z|x)$  is feature distribution,  $q(y|z)$  is a variational approximation to the true distribution  $p(y|z)$ , and  $q(z)$  is the variational approximation to the true distribution  $p(z)$ . To optimize Eq.(5), VIB uses neural networks to parameterize the Gaussian densities  $p(z|x)$  and  $q(y|z)$ , and treats  $q(z)$  as a fixed  $K$ -dimensional spherical Gaussian distribution  $q(z) = \mathcal{N}(0,I)$  (where  $K$  is the size of the bottleneck). Then, the reparameterization trick and Monte Carlo sampling are used to get an unbiased estimate of the gradient. This process thus allows DNNs to handle high-dimensional data, such as images, avoiding previous limitations on discrete or Gaussian distribution cases.
Next, how to determine the approximate  $q(z)$  is important in variational inference. VIB empirically treats  $q(z)$  as a fixed Gaussian distribution. For a better approximation of  $q(z)$ , [19] proposes the Conditional Entropy Bottleneck (CEB), which introduces label prior information to approximate  $q(z)$ . The CEB can be formulated as:
$$
I (Z; Y) - \beta I (X; Z | Y) \geq \underset {p (x, y) p (z | x)} {\mathbb {E}} \left[ \log q (y | z) - \beta \log \frac {p (z | x)}{q (z | y)} \right]. \tag {6}
$$
In detail, CEB utilizes a linear mapping layer to parameterize  $q(z|y)$  that takes a one-hot label  $y$  as input and outputs a vector  $\mu_y$  as the mean of the Gaussian distribution  $q(z|y) = \mathcal{N}(\mu_y, I)$ . CEB also uses an identity matrix  $I$  for the variance of both  $q(z|y)$  and  $p(z|x)$ . Note that CEB and VIB differ in terms of the presence of a class conditional or unconditional variational marginal, and whether  $q(z|y)$  and  $p(z|x)$  have a fixed variance. Thus, by introducing such a class prior, CBE can learn an approximate distribution better. However, using the one-hot label alone to approximate high-dimensional distribution  $q(z)$  is inaccurate enough due to data processing inequality and intrinsic dimension [37], particularly for high-dimensional image data.
To address this challenge, we propose leveraging adversarial robustness distillation [21] to build the information bottleneck objective. This approach involves providing robust prior information obtained from a robust pre-trained model, which aims to improve model robustness and reduce model uncertainty [27, 53]. Specifically, we replace conventional one-hot label prior information with robust prior information derived from such a robust pre-trained model. We call this method Information Bottleneck Distillation (IBD). To implement IBD, we approximate  $q(z)$  by utilizing the intermediate features extracted by an adversially pre-trained model, and our final objective is formulated as:
$$
I (Z; Y) - \beta I (X; Z | T) \geq \underset {p (x, y) p (z | x)} {\mathbb {E}} \left[ \log q (y | z) - \beta \log \frac {p (z | x)}{q (z | t)} \right], \tag {7}
$$
where  $T$  is the random variable of intermediate features extracted by the adversarially pre-trained teacher model. We employ a Gaussian distribution  $\mathcal{N}(\mu_y,\delta)$  with mean  $\mu (\cdot)$  and variance  $\delta$  as the variational distribution  $q(z|t)$ . Here, the mean  $\mu (\cdot)$  is a function of the intermediate feature from the robust pre-trained model, and the variance  $\delta$  is set to an identity matrix as the same CEB. Eq. (7) provides a variational lower bound on IBD, and a more detailed derivation is available in Appendix B.
For more details of the bound, we refer the reader to [2], or see Appendix A.
3.2 Optimization for IBD
In a standard classification task, we denote  $h_\theta = (g \circ f)$  as a deep neural network with parameter  $\theta$ . Here,  $f: \mathbb{R}^{d_X} \to \mathbb{R}^{d_Z}$  maps the inputs  $X$  to the intermediate features  $Z$ , and  $g: \mathbb{R}^{d_Z} \to \mathbb{R}^{d_Y}$  further maps the intermediate features  $Z$  to the final outputs  $Y$ , so that  $Z = f(X)$  and  $g(Z) = h_\theta(X)$ . Then, we denote the target student network as  $h_s = (g_s \circ f_s)$ , and the adversarially pre-trained teacher network as  $h_t = (g_t \circ f_t)$ . Furthermore, a set of intermediate features from different layers of the student model is denoted as  $Z_s = f_s(X) = \{z_s^1, z_s^2, \ldots, z_s^m\}$ , and a set of intermediate features from different layers of the teacher model is denoted as  $Z_t = f_t(X) = \{z_t^1, z_t^2, \ldots, z_t^n\}$ . Here,  $n$  and  $m$  represent the number of layers of the student and teacher models, respectively. Each feature has its own feature map size and channel dimension, denoted as  $z \in \mathbb{R}^{C \times H \times W}$ , where  $C, H$  and  $W$  represent the channel numbers, feature height, and width, respectively.
1) Maximizing  $I(Z,Y)$  via Soft Label Distillation. To optimize Eq.(7) using the gradient descent algorithm, we simplify the first term on the right-hand side of Eq.(7) to the expected log-likelihood in the form of cross-entropy, which can be formulated as follows:
$$
\begin{array}{l} \underset {p (x, y) p (z \mid x)} {\mathbb {E}} \left[ \log q (y \mid z) \right] = \int p (y, z) \log q (y \mid z) d y d z \\ = \int p (z) p (y | z) \log q (y | z) d y d z \\ = \int p (z, x) p (y | z) \log q (y | z) d y d z d x \tag {8} \\ = \int p (x) p (z | x) p (y | z) \log q (y | z) d y d z d x \\ = \underset {p (x) p (z | x)} {\mathbb {E}} \left[ \int p (y | z) \log q (y | z) d y \right], \\ \end{array}
$$
where  $p(y|z)$  indicates the true likelihood considered as a target label  $y$  corresponding to  $z$ . The  $q(y|z)$  is modeled by the classifier of the student model  $h_s$ . During the training phase, we can use the output probability of the robust teacher model  $h_t$  as the soft label  $y_t = h_t(x_{\mathrm{nat}})$  to approximate the distribution of  $p(y|z)$ , despite knowing the label  $y$ . According to [74], using robust soft-labels is crucial in enhancing robustness. Thus, we rewrite the first term as:
$$
\underset {p (x, y) p (z \mid x)} {\mathbb {E}} [ \log q (y | z) ] = \underset {p (x) p (z \mid x)} {\mathbb {E}} [ \int p (y | z) \log q (y | z) d y ] = \underset {p (x)} {\mathbb {E}} [ - \mathcal {L} _ {C E} (h _ {s} (x), y _ {t}) ]. \tag {9}
$$
2) Maximizing  $-I(Z,X|T)$  via Adaptive Feature Distillation. To optimize the second term on the right-hand side of Eq.(7), the formulation can be written as:
$$
\begin{array}{l} \underset {p (x) p (z \mid x)} {\mathbb {E}} \left[ \log \frac {p (z \mid x)}{q (z \mid t)} \right] = \int p (z, x) \log \frac {p (z \mid x)}{q (z \mid t)} d z d x \\ = \int p (x) p (z | x) \log \frac {p (z | x)}{q (z | t)} d z d x \tag {10} \\ = \underset {p (x)} {\mathbb {E}} \left[ \int p (z | x) \log \frac {p (z | x)}{q (z | t)} d z \right]. \\ \end{array}
$$
Thus, we optimize the KL divergence between the feature likelihood  $p(z|x)$  and the appropriate feature probability  $q(z|t)$ . Here, we parameterize Gaussian densities  $p(z|x)$  and  $q(z|t)$  using neural networks, where the mean of  $p(z|x)$  and  $q(z|t)$  are the intermediate features of  $f_{s}$  and  $f_{t}$ , respectively. Both variances are set to an identity matrix. As a result, the second term can be calculated by:
$$
\underset {p (x, y) p (z \mid x)} {\mathbb {E}} \left[ \log \frac {p (z \mid x)}{q (z \mid t)} \right] = \underset {p (x) p (z \mid x)} {\mathbb {E}} \left[ \mathrm {K L} \left(p (z \mid x) \| q (z \mid t)\right) \right] = \underset {p (x)} {\mathbb {E}} \left[ \left(f _ {t} (x) - f _ {s} (x)\right) ^ {2} + c \right], \tag {11}
$$
where  $c$  is a constant. When optimizing DNNs using Eq.(11) as the objective, a challenge arises in selecting suitable intermediate features from models to calculate the loss. This is due to the fact that different intermediate features tend to have different information, especially when the student and teacher models have different architectures. Motivated by [6, 32], we leverage an attention-based feature distillation strategy to achieve cross-layer information transfer.
Given two sets of the intermediate features,  $Z_{t} = f_{t}(x) = \{z_{t}^{1}, z_{t}^{2}, \ldots, z_{t}^{n}\}$  and  $Z_{s} = f_{s}(x) = \{z_{s}^{1}, z_{s}^{2}, \ldots, z_{s}^{m}\}$ , our aim is to identify similarities for all possible combinations and transfer relevant information from the teacher to the student. We compare the features of both teacher and student models by using two different pooling methods - global average pooling and channel-wise pooling. The similarity determined by two globally pooled features is used as the weight for transferring information over the distance defined by the channel-wisely averaged features. In order to identify the similarity between  $Z_{t}$  and  $Z_{s}$ , we adopt a query-key concept of the attention mechanism [61]. In detail, each teacher features generate a query set  $Q_{t} = \{q_{t}^{1}, q_{t}^{2}, \ldots, q_{t}^{n}\}$ , and each student feature generate a key set  $K_{s} = \{k_{s}^{1}, k_{s}^{2}, \ldots, k_{s}^{m}\}$ , where the  $q_{t}^{n}$  and  $k_{s}^{m}$  are calculated as:
$$
\mathbf {q} _ {t} ^ {n} = R e \left(W _ {t} ^ {n} \cdot \operatorname {G A P} \left(z _ {t} ^ {n}\right)\right),
$$
$$
\mathbf {k} _ {s} ^ {m} = R e \left(W _ {s} ^ {m} \cdot \operatorname {G A P} \left(z _ {s} ^ {m}\right)\right), \tag {12}
$$
where GAP denotes the global average pooling,  $Re$  is Relu activation function,  $W_{t}^{n}$  and  $W_{s}^{m}$  are linear transition parameters. It should be noted that these features possess varying transition weights, as they convey different levels of information.
By utilizing the queries and keys, attention values that represent the relation between teacher and student features are calculated with a softmax function:
$$
A t t n = \operatorname {s o f t m a x} \left(\frac {Q _ {T} W ^ {Q - k} K _ {S} ^ {\top}}{\sqrt {d}}\right), \tag {13}
$$
where  $W^{Q - k} \in \mathbb{R}^{d \times d}$  is a bilinear weight. The bilinear weighting is utilized to generalize the attention values derived from different source ranks, as queries and keys are identified within features of distinct dimensions [34].  $Attn_{i,j}$  indicates the attention weight that captures the relation between the  $i$ -th teacher feature and the  $j$ -th student feature. Therefore,  $Attn$  can make the teacher feature  $z_t^n$  transmit its corresponding information selectively and adaptively to different student features. Finally, the second term on the right-hand side of Eq. (7) can be written as:
$$
\underset {p (x, y) p (z \mid x)} {\mathbb {E}} \left[ \log \frac {p (z \mid x)}{q (z \mid t)} \right] = \underset {p (x)} {\mathbb {E}} \left[ \sum_ {i} ^ {n} \sum_ {j} ^ {m} A t t n _ {i, j} \left(\mathbf {T} _ {t} ^ {i} \left(z _ {t} ^ {i}\right) - \mathbf {T} _ {s} ^ {j} \left(z _ {s} ^ {j}\right)\right) ^ {2} \right], \tag {14}
$$
where  $\mathbf{T}$  is a transform function with a feature map size alignment (up-sampled or down-sampled), and then apply a channel-wise average pooling to get an average feature map for loss computation.
Combing Eq.9 and Eq.14, the final objective of IBD can be defined as follows:
$$
L _ {I B D} = \min  _ {p (x)} \mathbb {E} \left[ \mathcal {L} _ {C E} \left(h _ {s} (x), y _ {t}\right) + \beta \sum_ {i} ^ {n} \sum_ {j} ^ {m} A t t n _ {i, j} \left(\mathbf {T} _ {t} ^ {i} \left(z _ {t} ^ {i}\right) - \mathbf {T} _ {s} ^ {j} \left(z _ {s} ^ {j}\right)\right) ^ {2} \right]. \tag {15}
$$
3.3 Applying IBD to Robust Learning
IBD can be naturally applied in combination with adversarial training. The final objective function is formulated as follows:
$$
\begin{array}{l} L _ {o b j} = \min  _ {p (x)} \mathbb {E} \left[ \left(1 - \alpha\right) \mathcal {L} _ {C E} \left(h _ {s} \left(x _ {\mathrm {n a t}}\right), y _ {t}\right) + \alpha \mathcal {L} _ {C E} \left(h _ {s} \left(x _ {\mathrm {a d v}}\right), y _ {t}\right) \right. \\ \left. + \beta \sum_ {i} ^ {n} \sum_ {j} ^ {m} A t t n _ {i, j} \left(\mathbf {T} _ {t} ^ {i} \left(z _ {t} ^ {i}\right) - \mathbf {T} _ {s} ^ {j} \left(z _ {s} ^ {j}\right)\right) ^ {2} \right], \\ \end{array}
$$