SlowGuess/NAACL_2022_OCR huggingface dataset:
$$
\begin{array}{l} \sigma \left(\frac {g _ {l} ^ {d} (i) + \Phi_ {l} ^ {d} (i) - \left(g _ {l} ^ {d} \left(i _ {k}\right) + \Phi_ {l} ^ {d} \left(i _ {k}\right)\right)}{\tau}\right) <   \hat {m} _ {l} ^ {d} (\tau) (i) \\ <   \sigma \left(\frac {g _ {l} ^ {d} (i) + \Phi_ {l} ^ {d} (i) - \left(g _ {l} ^ {d} \left(i _ {k + 1}\right) + \Phi_ {l} ^ {d} \left(i _ {k + 1}\right)\right)}{\tau}\right) \\ \end{array}
$$
By taking the limit on both sides, we get the following results:
$$
\lim  _ {\tau \to 0} \hat {m} _ {l} ^ {d} (\tau) (i _ {u}) = \left\{ \begin{array}{l l} 1, & \text {i f} u > k \\ 0, & \text {i f} u <   k \end{array} \right.
$$
And, because  $\sum_{u = 1}^{n_p}\hat{m}_l^d (\tau)(i_u) = k,$  by taking the limit on both sides, we will have  $\lim_{\tau \to 0}\hat{m}_l^d (\tau)(i_k) = 1$  Finally, we have
$$
\lim  _ {\tau \to 0} \hat {m} _ {l} ^ {d} (\tau) (i _ {u}) = \left\{ \begin{array}{l l} 1, & \text {i f} u \geqslant k \\ 0, & \text {i f} u <   k \end{array} \right.
$$
which is equivalent to  $\lim_{\tau \to 0}\hat{m}_l^d (\tau) = \tilde{m}_l^d$
In this section, we give a simple proof of inequality (3). In fact, we only need to prove  $\mathbb{H}\big[P(i_1,\dots ,i_k|\Phi_l^d)\big]\geqslant \mathbb{H}\big[P(i_1|\Phi_l^d)\big]$ . The proof is as follows:
$$
\begin{array}{l} \mathbb {H} \left[ P \left(i _ {1}, \dots , i _ {k} \mid \Phi_ {l} ^ {d}\right) \right] = - \underset {i _ {1}, \dots , i _ {k} \mid \Phi_ {l} ^ {d}} {\mathbb {E}} \left[ \log P \left(i _ {1}, \dots , i _ {k} \mid \Phi_ {l} ^ {d}\right) \right] ^ {5:} \\ = - \underset {i _ {1}, \dots , i _ {k} \mid \Phi_ {l} ^ {d}} {\mathbb {E}} \left[ \sum_ {j = 2} ^ {k} \log P \left(i _ {j} \mid i _ {1}, \dots , j _ {j - 1}, \Phi_ {l} ^ {d}\right) + \log P \left(i _ {1} \mid \Phi_ {l} ^ {d}\right) \right] \\ \geqslant - \mathbb {E} _ {i _ {1}, \dots , i _ {k} | \Phi_ {l} ^ {d}} \left[ \log P \left(i _ {1} \mid \Phi_ {l} ^ {d}\right) \right] \\ = - \underset {i _ {1} | \Phi_ {l} ^ {d}} {\mathbb {E}} \left[ \log P \left(i _ {1} \mid \Phi_ {l} ^ {d}\right) \right] = \mathbb {H} \left[ P \left(i _ {1} \mid \Phi_ {l} ^ {d}\right) \right] \\ \end{array}
$$
Figure 4: Visualization of domains according to their dropout masks (a large vector concatenating the dropping masks of all the layers of the model) constructed by PCA.
Algorithm 1 Training LaMGD
-  $n_d$  corpora  $C^d, d \in [1, \ldots, m]$  mains equipped by an emtion  $D_d(x)$
number of groups:  $n_p$  ; number of retained groups:  $k$
-  $i = 0$ ; iter_num
1: repeat
2: Pick a batch from domain  $d$
3: Sample  $\forall l, \forall p : g_l^d(p)^{\mathrm{i.i.d}}$  Gumbel(0,1)
4: Solve the equation  $\forall l$
using binary search
Compute mask of each layer
$$
\forall l, \hat {m} _ {l} ^ {d} (\tau) = \sigma \left(\frac {g _ {l} ^ {d} + \Phi_ {l} ^ {d} + \bar {\nu}}{\tau}\right)
$$
Apply masks to their corresponding layer
7: Compute gradient of training loss over the underlying Transformer
$$
\Delta_ {\theta} = \frac {\partial L}{\partial \theta}
$$
8: Compute gradient over the Soft-Top-K masks
$$
\frac {\partial D}{\partial \hat {m} _ {l} ^ {d} (\tau)}
$$
9: Compute implicit gradient of the Soft-Top-K masks over  $\Phi_l^d$
$$
\frac {\partial \bar {v}}{\partial \Phi_ {l} ^ {d}} = - \left(\frac {\partial g}{\partial \bar {v}}\right) ^ {- 1} \times \frac {\partial g}{\partial \Phi_ {l} ^ {d}}
$$
$$
\frac {\partial \hat {m} _ {l} ^ {d} (\tau)}{\partial \Phi_ {l} ^ {d}} = \frac {\partial \hat {m} _ {l} ^ {d} (\tau)}{\partial v} \times \frac {\partial v}{\partial \Phi_ {l} ^ {d}} + \frac {1}{\tau} \frac {\exp \left(\frac {g _ {l} ^ {d} (i) + \Phi_ {l} ^ {d} (i) + v}{\tau}\right)}{(1 + \exp \left(\frac {g _ {l} ^ {d} (i) + \Phi_ {l} ^ {d} (i) + v}{\tau}\right)) ^ {2}}
$$
10: Compute the gradient the training over  $\Phi_l^d$
$$
\Delta_ {\Phi_ {l} ^ {d}} = \frac {\partial D}{\partial \hat {m} _ {l} ^ {d} (\tau)} \times \frac {\partial \hat {m} _ {l} ^ {d} (\tau)}{\partial \Phi_ {l} ^ {d}} + \frac {\partial \mathbb {H} [ \operatorname {s o f t m a x} (\Phi_ {l} ^ {d}) ]}{\partial \Phi_ {l} ^ {d}}
$$
11: Update  $\theta$  and  $\Phi_l^d$  with their gradients
12:  $i = i + 1$
13: until  $i > \text{iter\_num}$