SlowGuess/NeurIPS_2023_OCR huggingface dataset:
<sup>1</sup> School of Computing, National University of Singapore
$^{3}$  School of Computer Science, The University of Auckland
5 Graduate School of Frontier Sciences, The University of Tokyo
xuxilie@comp.nus.edu.sg jingfeng.zhang@auckland.ac.nz
fengliu.ml@gmail.com sugi@k.u-tokyo.ac.jp
Adversarial contrastive learning (ACL) is a technique that enhances standard contrastive learning (SCL) by incorporating adversarial data to learn a robust representation that can withstand adversarial attacks and common corruptions without requiring costly annotations. To improve transferability, the existing work introduced the standard invariant regularization (SIR) to impose style-independence property to SCL, which can exempt the impact of nuisance style factors in the standard representation. However, it is unclear how the style-independence property benefits ACL-learned robust representations. In this paper, we leverage the technique of causal reasoning to interpret the ACL and propose adversarial invariant regularization (AIR) to enforce independence from style factors. We regulate the ACL using both SIR and AIR to output the robust representation. Theoretically, we show that AIR implicitly encourages the representational distance between different views of natural data and their adversarial variants to be independent of style factors. Empirically, our experimental results show that invariant regularization significantly improves the performance of state-of-the-art ACL methods in terms of both standard generalization and robustness on downstream tasks. To the best of our knowledge, we are the first to apply causal reasoning to interpret ACL and develop AIR for enhancing ACL-learned robust representations. Our source code is at https://github.com/GodXuxilie/Enhancing_ACL_via_AIR.
The attention towards pre-trained models that can be easily finetuned for various downstream applications has significantly increased recently [17, 18, 19, 41]. Notably, foundation models [4] via self-supervision on large-scale unlabeled data, such as GPT [5] and CLAP [20], can be adapted to a wide range of downstream tasks. Due to the high cost of annotating large-scale data, unsupervised learning techniques [21, 34, 49, 16] are commonly used to obtain generalizable representations, in which standard contrastive learning (SCL) has been shown as the most effective way [9, 38, 10].
Adversarial contrastive learning (ACL) [30, 29], that incorporates adversarial data [37] with SCL [9], can yield robust representations that are both generalizable and robust against adversarial attacks [23]
Figure 1: Causal graph of data generation procedure in SCL [38] (left panel) and ACL (right panel).  $s$  is style variable,  $c$  is content variable,  $x$  is unlabeled data,  $\tilde{x}$  is the generated adversarial data, and  $\theta$  is the parameter of representation. The solid arrows are causal, but the dash arrows are not causal. The proxy label  $y^{R} \in \mathcal{Y}^{R}$  is a refinement of the target labels  $y_{t} \in \mathcal{Y} = \{y_{i}\}_{i=1}^{T}$  that represents an unknown downstream task. In the causal view, each augmented view of a data point has its own proxy label, where both SCL and ACL aim to minimize the differences between the representation outputs of different views of the same data.
and common corruptions [26]. ACL is attracting increasing popularity since the adversarial robustness of the pre-trained models is essential to safety-critical applications [6, 33]. Many studies have tried to improve ACL from various perspectives including increasing contrastive views and leveraging pseudo labels [22], leveraging hard negative sampling [42, 50], and dynamically scheduling the strength of data augmentations [36].
The style-independence property of learned representations, which eliminates the effects of nuisance style factors in SCL, has been shown to improve the transferability of representations [38]. To achieve style-independence property in learned representations, Mitrovic et al. [38] proposed a standard invariant regularization (SIR) that uses causal reasoning [39, 40] to enforce the representations of natural data to be invariant to style factors. SIR has been shown effective in enhancing representation transferability. However, it remains unclear in the literature how to achieve style independence in robust representations learned via ACL, especially when adversarial data is used during pre-training.
Therefore, we leverage the technique of causal reasoning to enforce robust representations learned via ACL to be style-independent. As shown in the right panel of Figure 1, we construct the causal graph of ACL (details refer to Section 3.1). Different from the causal graph of SCL [38], ACL has an extra path  $x \to \tilde{x}$  since ACL will generate the adversarial data  $\tilde{x}$  given the unlabeled data  $x$  (i.e.,  $x \to \tilde{x}$ ). Then, ACL learns representations by maximizing both the probability of the proxy label  $y^{R}$  given the natural data and that given the adversarial data (i.e.,  $x \to y^{R}$  and  $\tilde{x} \to y^{R}$ ). Theorem 1 shows that maximizing the aforementioned probability in the causal view is equivalent to the learning objective of ACL [29], which justifies the rationality of our constructed causal graph of ACL.
To enforce robust representations to be style-independent, we propose an adversarial invariant regularization (AIR). Specifically, AIR (see Eq. (7)) aims to penalize the Kullback-Leibler divergence between the robust representations of the unlabeled data augmented via two different data augmentation functions during the learning procedure of  $x \rightarrow \tilde{x} \rightarrow y^R$ . Note that SIR [38] is a special case of AIR in the standard context where  $x = \tilde{x}$ . Then, we propose to learn robust representations by minimizing the adversarial contrastive loss [29, 22, 50, 36] together with a weighted sum of SIR and AIR and show the learning algorithm in Algorithm 1.
Furthermore, we give a theoretical understanding of the proposed AIR and show that AIR is beneficial to the robustness against corruptions on downstream tasks. Based on the decomposition of AIR shown in Lemma 2, we propose Proposition 3 which explains that AIR implicitly enforces the representational distance between original as well as augmented views of natural data and their adversarial variants to be independent of the style factors. In addition, we theoretically show that the style-independence property of robust representations learned via ACL will still hold on the downstream classification tasks in Proposition 4, which could be helpful to improve robustness against input perturbations on downstream tasks [3, 45, 28, 38].
Empirically, we conducted comprehensive experiments on various datasets including CIFAR-10 [31], CIFAR-100 [31], STL-10 [12], CIFAR-10-C [26], and CIFAR-100-C [26] to show the effectiveness of our proposed method in improving ACL methods [29, 22, 50, 36]. We demonstrate that our proposed method can achieve the new state-of-the-art (SOTA) performance on various downstream tasks by significantly enhancing standard generalization and robustness against adversarial attacks [13, 14, 2] and common corruptions [26]. Notably, compared with the prior SOTA method DynACL [36], AIR improves the standard and robust test accuracy by  $2.71\%$  (from  $69.59 \pm 0.08\%$  to  $72.30 \pm 0.10\%$ ) and  $1.17\%$  (from  $46.49 \pm 0.05\%$  to  $47.66 \pm 0.06\%$ ) on the STL-10 task and increases the test accuracy under common corruptions by  $1.15\%$  (from  $65.60 \pm 0.18\%$  to  $66.75 \pm 0.10\%$ ) on CIFAR-10-C.
2 Related Works and Preliminaries
Here, we introduce the related works in ACL and causal reasoning and provide the preliminaries.
2.1 Related Works
Adversarial contrastive learning (ACL). Contrastive learning (CL) is frequently used to leverage large unlabeled datasets for learning useful representations. Chen et al. [9] presented SimCLR that leverages contrastive loss for learning useful representations and achieved significantly improved standard test accuracy on downstream tasks. Recently, adversarial contrastive learning (ACL) [30, 29, 27, 22, 50, 52, 36] has become the most effective unsupervised approaches to learn robust representations. Jiang et al. [29], which is the seminal work of ACL that incorporates adversarial data [37] with contrastive loss [9], showed that ACL can exhibit better robustness against adversarial attacks [23, 13] and common corruptions [26] on downstream tasks compared with SCL [9].
One research line focuses on improving the performance of ACL. AdvCL [22] leverages a third contrastive view for high-frequency components and pseudo labels generated by clustering to improve ACL. Yu et al. [50] proposed the asymmetric InfoNCE objective (A-InfoNCE) that utilizes the hard negative sampling method [42] to further enhance AdvCL [22]. Recently, DynACL [36] dynamically schedules the strength of data augmentations and correspondingly schedules the weight for standard and adversarial contrastive losses, achieving the SOTA performance among existing ACL methods. Another research line focuses on improving the efficiency of ACL since ACL requires tremendous running time to generate the adversarial training data during pre-training. Xu et al. [48] proposed a robustness-aware coreset selection (RCS) method to speed up ACL by decreasing the number of training data. Our paper, belonging to the former research line, leveraged the technique of causal reasoning to further enhance ACL and achieved new SOTA results. We treat the incorporation of our proposed method with RCS [48] as future work.
Causal reasoning. Causal reasoning [39, 53, 40, 8, 56, 38, 7, 44, 43] has been widely applied to machine learning to identify causal relations and ignore nuisance factors by intervention. In particular, Zhang et al. [53, 56] investigated the causality in supervised adversarial training [37, 55] where label information is required and proposed the regularizer to eliminate the difference between the natural and adversarial distributions. Mitrovic et al. [38] built the causal graph of SCL in the standard context and introduced the standard invariant regularization (SIR) which aims to enforce the representation of unlabeled natural data to be invariant of the data augmentations. SIR can significantly improve robustness transferability to downstream tasks in terms of common corruptions [26], which is empirically validated by Mitrovic et al. [38].
Standard contrastive learning (SCL) [9]. Let  $h_\theta : \mathcal{X} \to \mathcal{Z}$  be a representation extractor parameterized by  $\theta$ ,  $g : \mathcal{Z} \to \mathcal{V}$  be a projection head that maps representations to the space where the contrastive loss is applied, and  $\tau_i, \tau_j : \mathcal{X} \to \mathcal{X}$  be two transformation operations randomly sampled from a pre-defined transformation set  $\mathcal{T}$ . Given a minibatch  $B \sim \mathcal{X}^\beta$  consisting of  $\beta$  original samples, we denote the augmented minibatch  $B^u = \{\tau_u(x_k) | \forall x_k \in B\}$  via augmentation function  $\tau_u(\cdot)$ . We take  $f_\theta(\cdot) = g \circ h_\theta(\cdot)$  and  $x_k^u = \tau_u(x_k)$  for any  $x_k \sim \mathcal{X}$  and  $u \in \{i,j\}$ . Given a positive pair  $(x_k^i, x_k^j)$ , the standard contrastive loss proposed by SimCLR [9] is formulated as follows:
$$
\ell_ {\mathrm {C L}} \left(x _ {k} ^ {i}, x _ {k} ^ {j}; \theta\right) = - \sum_ {u \in \{i, j \}} \log \frac {e ^ {\operatorname* {s i m} \left(f _ {\theta} \left(x _ {k} ^ {i}\right) , f _ {\theta} \left(x _ {k} ^ {j}\right)\right) / t}}{\sum_ {x \in B ^ {i} \cup B ^ {j} \backslash \left\{x _ {k} ^ {u} \right\}}} e ^ {\operatorname* {s i m} \left(f _ {\theta} \left(x _ {k} ^ {u}\right) , f _ {\theta} (x)\right) / t}}, \tag {1}
$$
ACL [29] and DynACL [36]. Let  $(\mathcal{X}, d_{\infty})$  be the input space  $\mathcal{X}$  with the infinity distance metric  $d_{\infty}(x, x') = \|x - x'\|_{\infty}$ , and  $\mathcal{B}_{\epsilon}[x] = \{x' \in \mathcal{X} \mid d_{\infty}(x, x') \leq \epsilon\}$  be the closed ball of radius  $\epsilon > 0$  centered at  $x \in \mathcal{X}$ . Given a data point  $x_k \in \mathcal{X}$ , the adversarial contrastive loss is as follows:
$$
\ell_ {\mathrm {A C L}} \left(x _ {k}; \theta\right) = (1 + \omega) \cdot \ell_ {\mathrm {C L}} \left(\tilde {x} _ {k} ^ {i}, \tilde {x} _ {k} ^ {j}; \theta\right) + (1 - \omega) \cdot \ell_ {\mathrm {C L}} \left(x _ {k} ^ {i}, x _ {k} ^ {j}; \theta\right), \tag {2}
$$
$$
\text {w h e r e} \quad \tilde {x} _ {k} ^ {i}, \tilde {x} _ {k} ^ {j} = \arg \max  \ell_ {\mathrm {C L}} \left(\tilde {x} _ {k} ^ {i}, \tilde {x} _ {k} ^ {j}; \theta\right), \tag {3}
$$
$$
\begin{array}{c} \tilde {x} _ {k} ^ {i} \in \mathcal {B} _ {\epsilon} [ x _ {k} ^ {i} ] \\ \tilde {x} _ {k} ^ {j} \in \mathcal {B} _ {\epsilon} [ x _ {k} ^ {j} ] \end{array}
$$
in which  $\omega \in [0,1]$  is a hyperparameter and  $\tilde{x}_k^i$  and  $\tilde{x}_k^j$  are adversarial data generated via projected gradient descent (PGD) [37, 54, 55] within the  $\epsilon$ -balls centered at  $x_k^i$  and  $x_k^j$ . Note that ACL [29] fixes  $\omega = 0$  while DynACL [36] dynamically schedules  $\omega$  according to its dynamic augmentation scheduler that gradually anneals from a strong augmentation to a weak one. We leave the details of the data augmentation scheduler [36] in Appendix C due to the page limitation.
In this section, we first create the causal graph for ACL which is the fundamental premise for causal reasoning. Then, we introduce adversarial invariant regularization (AIR) according to the causal understanding in the context of ACL and the learning algorithm of our proposed method. Lastly, we demonstrate the theoretical understanding of AIR and theoretically show that the style-independence property is generalizable to downstream tasks.
3.1 Adversarial Contrastive Learning (ACL) in the View of Causality
A causal graph is arguably the fundamental premise for causal reasoning [39, 40]. Previously, Mitrovic et al. [38] constructed the causal graph of SCL as shown in the left panel of Figure 1. However, how to build causal graphs of ACL is still unclear in the literature. Therefore, we take the primitive step to construct the causal graph of ACL.
To begin with, we formalize the data generation procedure in the causal graph. Let  $x \in \mathcal{X}$  denote an unlabeled data point and  $\mathcal{Y} = \{y_t\}_{t=1}^T$  denote a set of labels in an unknown downstream task. Then, similar to Mitrovic et al. [38] and Zhang et al. [56] we make the following assumptions: 1) the data point  $x$  is generated from the content factor  $c$  and the style factor  $s$ , i.e.,  $c \to x \gets s$ , 2) the content is independent of the style, i.e.,  $c \perp s$ , and 3) only the content, which is the original image data, is related to the downstream task, i.e.,  $c \to y_t$ .
Then, we build the learning procedure of CL. According to the above assumptions, the content serves an essential role in downstream classification tasks. Therefore, CL can be regarded as estimating the probability of the label given the content, i.e.,  $p(y_{t}|c)$ . To learn the representations from the unlabeled data  $x$ , CL needs to maximize the conditional probability  $p(y_{t}|x)$ .
However, in practice, the label  $y_{t}$  is unknown. Thanks to the causal concept of refinement [8, 38], we can construct a proxy label  $y^{R} \in \mathcal{Y}^{R}$  which is a refinement of  $\mathcal{V}$ . This proxy label  $y^{R}$  should be consistent across different views of the unlabeled data  $x$ , which enables learning the representations. Intuitively, the proxy label  $y^{R}$  can be viewed as a more fine-grained variant of the original label  $y_{t}$ . For example, a refinement for recognizing cars would be the task of classifying various types of cars. Therefore, SCL targets maximizing the probability of the proxy label given the unlabeled natural data, i.e.,  $p(y^{R}|x)$  [38].
Besides  $x \to y^R$ , ACL has an extra path  $x \to \tilde{x} \to y^R$  since ACL needs to first generate the adversarial data (Eq. (3)) and then learns representations from the unlabeled natural and adversarial data (Eq. (2)). According to Eq. (3), we observe that ACL only uses the observed unlabeled data  $x$  for generating adversarial data, i.e.,  $x \to \tilde{x} \gets \theta$ . Then, according to the paths  $x \to y^R$  and  $\tilde{x} \to y^R$ , from the causal view, ACL learns representations by maximizing both the probability of the proxy label  $y^R$  given the natural data (i.e.,  $p(y^R | x)$ ) and that given adversarial data (i.e.,  $p(y^R | \tilde{x})$ ).
Lastly, we provide Theorem 1 to justify the rationality of our constructed causal graph of ACL.
Theorem 1. From the causal view, in ACL, maximizing the conditional probability both  $p(y^R | x)$  and  $p(y^R | \tilde{x})$  is equivalent to minimizing the learning objective of ACL [29] that is the sum of standard and adversarial contrastive losses.
The proof of Theorem 1 is in Appendix B.1.
3.2 Adversarial Invariant Regularization
According to the independence of causal mechanisms [40], performing interventions on the style variable  $s$  should not change the probability of the proxy label given the unlabeled data, i.e.,  $p(y^{R}|x)$ . According to the path  $x \to \tilde{x} \to y^{R}$  shown in the causal graph of ACL, we can obtain that
$$
p \left(y ^ {R} | x\right) = p \left(y ^ {R} | \tilde {x}\right) p (\tilde {x} | x), \tag {4}
$$
under the assumption that the process of  $x \to \tilde{x} \to y^R$  satisfies the Markov condition. Therefore, we should make the conditional probability  $p(y^R | \tilde{x}) p(\tilde{x} | x)$  that is learned via ACL become style-independent. It guides us to propose the following criterion that should be satisfied by ACL:
$$
p ^ {d o \left(\tau_ {i}\right)} \left(y ^ {R} | \tilde {x}\right) p ^ {d o \left(\tau_ {i}\right)} (\tilde {x} | x) = p ^ {d o \left(\tau_ {j}\right)} \left(y ^ {R} | \tilde {x}\right) p ^ {d o \left(\tau_ {j}\right)} (\tilde {x} | x) \quad \forall \tau_ {i}, \tau_ {j} \in \mathcal {T}, \tag {5}
$$
where  $do(\tau)$  as the short form of  $do(s = \tau)$  denotes that we perform the intervention on  $s$  via data augmentation function  $\tau (\cdot)$ . We leverage the representational distance between various views of natural data and their adversarial variants normalized by the softmax function to calculate the conditional probability  $p^{do(\tau_u)}(y^R |\tilde{x})$  and  $p^{do(\tau_u)}(\tilde{x} |x)$  where  $u\in \{i,j\}$  as follows:
$$
p ^ {d o \left(\tau_ {u}\right)} \left(y ^ {R} | \tilde {x}\right) = \frac {e ^ {\operatorname* {s i m} \left(f _ {\theta} \left(x\right) , f _ {\theta} \left(\tilde {x} ^ {u}\right)\right) / t}}{\sum_ {x _ {k} \in B} e ^ {\operatorname* {s i m} \left(f _ {\theta} \left(x _ {k}\right) , f _ {\theta} \left(\tilde {x} _ {k} ^ {u}\right)\right) / t}}, \quad p ^ {d o \left(\tau_ {u}\right)} (\tilde {x} | x) = \frac {e ^ {\operatorname* {s i m} \left(f _ {\theta} \left(\tilde {x} ^ {u}\right) , f _ {\theta} \left(x ^ {u}\right)\right) / t}}{\sum_ {x _ {k} \in B} e ^ {\operatorname* {s i m} \left(f _ {\theta} \left(\tilde {x} _ {k} ^ {u}\right) , f _ {\theta} \left(x _ {k} ^ {u}\right)\right) / t}}, \tag {6}
$$
in which  $x \in B \sim \mathcal{X}^{\beta}$  is the original view of natural data,  $x^{u}$  is the augmented view of natural data via augmentation  $\tau_{u}(\cdot)$ ,  $\tilde{x}^{u}$  is the adversarial variant generated via Eq. (3), and  $t$  is the temperature parameter. Note that we use  $x$  as the approximated surrogate of its true content  $c$  that incurs  $y^{R}$  (i.e.,  $c \to y^{R}$ ).
To achieve the above criterion in Eq. (5), we propose an adversarial invariant regularization (AIR) to regulate robust representations as follows:
$$
\mathcal {L} _ {\mathrm {A I R}} (B; \theta , \epsilon) = \mathrm {K L} \left(p ^ {d o (\tau_ {i})} \left(y ^ {R} | \tilde {x}\right) p ^ {d o (\tau_ {i})} (\tilde {x} | x) \| p ^ {d o (\tau_ {j})} \left(y ^ {R} | \tilde {x}\right) p ^ {d o (\tau_ {j})} (\tilde {x} | x); B\right), \tag {7}
$$
in which  $\epsilon \geq 0$  is the adversarial budget and  $\mathrm{KL}(p(x)\| q(x);B) = \sum_{x\in B}p(x)\log \frac{p(x)}{q(x)}$  denotes the Kullback-Leibler (KL) divergence. AIR can enforce the representation to satisfy the criterion in Eq. (5), thus eliminating the effect of the style factors on the representation. Therefore, when setting  $\epsilon$  as a positive constant,  $\mathcal{L}_{\mathrm{AIR}}(B;\theta ,\epsilon)$  can effectively regulate robust representations of adversarial data to be style-independent.
Besides, to explicitly regulate standard representations of natural data to be independent of style factors, we can simply set  $\epsilon = 0$  of AIR. We formulate AIR with  $\epsilon = 0$  as follows:
$$
\mathcal {L} _ {\mathrm {A I R}} (B; \theta , 0) = \mathrm {K L} \left(p ^ {d o \left(\tau_ {i}\right)} \left(y ^ {R} | x\right) \| p ^ {d o \left(\tau_ {j}\right)} \left(y ^ {R} | x\right); B\right), \tag {8}
$$
$$
\text {w h e r e} \quad p ^ {d o (\tau_ {u})} (y ^ {R} | x) = \frac {e ^ {\operatorname* {s i m} \left(f _ {\theta} (x) , f _ {\theta} \left(x ^ {u}\right)\right) / t}}{\sum_ {x _ {k} \in B} e ^ {\operatorname* {s i m} \left(f _ {\theta} \left(x _ {k}\right) , f _ {\theta} \left(x _ {k} ^ {u}\right)\right) / t}} \quad \forall u \in \{i, j \}.
$$
Note that  $\mathcal{L}_{\mathrm{AIR}}(B;\theta ,0)$  is the same formulation as the standard invariant regularization (SIR) [38], which can make the standard representation style-independent.
By incorporating adversarial contrastive loss with AIR, the learning objective function of our proposed method is shown as follows:
$$
\arg \min  _ {\theta} \sum_ {x \in U} \ell_ {\mathrm {A C L}} (x; \theta) + \lambda_ {1} \cdot \mathcal {L} _ {\mathrm {A I R}} (U; \theta , 0) + \lambda_ {2} \cdot \mathcal {L} _ {\mathrm {A I R}} (U; \theta , \epsilon) \tag {9}
$$
where  $U \sim \mathcal{X}^N$  refers to an unlabeled dataset consisting of  $N$  samples,  $\epsilon > 0$  is the adversarial budget, and  $\lambda_1 \geq 0$  and  $\lambda_2 \geq 0$  are two hyperparameters. The learning algorithm of ACL with AIR is demonstrated in Algorithm 1. Note that our proposed AIR is compatible with various learning objectives such as ACL [29], AdvCL [22], A-InfoNCE [50], and DynACL [36].
Algorithm 1 ACL with Adversarial Invariant Regularization (AIR)
1: Input: Unlabeled training set  $U$ , total training epochs  $E$ , learning rate  $\eta$ , batch size  $\beta$ , adversarial budget  $\epsilon > 0$ , hyperparameters  $\lambda_1$  and  $\lambda_2$   
2: Output: Pre-trained representation extractor  $h_\theta$   
3: Initialize parameters of model  $f_\theta = g \circ h_\theta$   
4: for  $e = 0$  to  $E - 1$  do  
5: for batch  $m = 1, \ldots, \lceil |U| / \beta \rceil$  do  
6: Sample a minibatch  $B_m$  from  $U$   
7: Update  $\theta \gets \theta - \eta \cdot \nabla_\theta \sum_{x_k \in B_m} \ell_{\mathrm{ACL}}(x_k; \theta) + \lambda_1 \cdot \mathcal{L}_{\mathrm{AIR}}(B_m; \theta, 0) + \lambda_2 \cdot \mathcal{L}_{\mathrm{AIR}}(B_m; \theta, \epsilon)$   
8: end for  
9: end for
We provide a theoretical understanding of AIR in Proposition 3 based on the decomposition of AIR shown in Lemma 2. We set  $\epsilon >0$  of AIR by default in this section.
Lemma 2 (Decomposition of AIR). AIR in Eq. (7) can be decomposed into two terms as follows:
$$
\begin{array}{l} \mathcal {L} _ {\mathrm {A I R}} (B; \theta , \epsilon) = \mathbb {E} _ {x \sim p ^ {d o (\tau_ {i})} (\tilde {x} | x)} [ \mathrm {K L} (p ^ {d o (\tau_ {i})} (y ^ {R} | \tilde {x}) \| p ^ {d o (\tau_ {j})} (y ^ {R} | \tilde {x})) ] \\ + \mathbb {E} _ {x \sim p ^ {d o (\tau_ {i})} (y ^ {R} | \tilde {x})} [ \mathrm {K L} (p ^ {d o (\tau_ {i})} (\tilde {x} | x) \| p ^ {d o (\tau_ {j})} (\tilde {x} | x)) ], \\ \end{array}
$$
where  $\mathbb{E}_{x\sim Q^3 (x)}[\mathrm{KL}(Q^1 (x)\| Q^2 (x))]$  is the expectation of KL divergence over  $Q^3$  and  $Q^{1},Q^{2},Q^{3}$  are probability distributions.
The proof of Lemma 2 is in Appendix B.2. Based on Lemma 2, we propose Proposition 3 that provides the theoretical understanding of AIR.