SlowGuess/NeurIPS_2023_OCR huggingface dataset:
where  $f$  is the attention function,  $s = r$  is the normalizing factor,  $\mathbf{B}_{\mathrm{rel}} \in \mathbb{R}^{n \times n}$  is the relative positional bias which is set as the simple relative position bias [SUV18; MZK+23] by default. We also have the choice for using RoPE [SLP+21] relative position embedding, and in this case, we don't apply relative position bias but use RoPE in its original form to transform the Q and K sequences. Squared ReLU is applied [SML+21] for the attention function on the image and the speech tasks and Softmax for the text related tasks. By default, the relative positions are calculated based on the token positions in the original sentence H instead of the compressed sequence Hc. For working memory mechanism, we conduct local attention between Q, K, V with the sliding window size (or the working memory size) w, and we set the normalizing factor s = w. The final output of the GAU is computed with a gating mechanism,
$$
\mathbf {G} = \operatorname {S i L U} \left(\mathbf {H} ^ {c} \mathbf {W} _ {g} + \mathbf {b} _ {g}\right) \quad \in \mathbb {R} ^ {r \times d _ {v}} \tag {7}
$$
$$
\mathbf {Y} ^ {c} = (\mathbf {G} \odot \mathbf {O}) \mathbf {W} _ {h} + \mathbf {b} _ {h} \quad \in \mathbb {R} ^ {r \times d _ {m}} \tag {8}
$$
where  $\mathbf{W}_g\in \mathbb{R}^{d_m\times d_v}$ ,  $\mathbf{W}_h\in \mathbb{R}^{d_v\times d_m}$ ,  $\mathbf{b}_g\in \mathbb{R}^{d_v}$ ,  $\mathbf{b}_h\in \mathbb{R}^{d_m}$  are learnable parameters.
A.3 Details of Multi-dimensional Damped EMA
As a special kind of SSM, the Multi-dimensional Damped EMA (MD-EMA) operator is first proposed in the MEGA  $\left[\mathrm{MZK}^{+}23\right]$  paper, and we briefly introduce it here for the completeness of our work. In Equation (2), the kernel  $\mathbf{K}\in \mathbb{R}^{n\times d_m}$  of MD-EMA is parameterized as following,
$$
\mathbf {K} = \left(\sum_ {i = 1} ^ {h} \boldsymbol {\eta} _ {i} (\boldsymbol {\alpha} \odot \boldsymbol {\beta}) _ {i}, \sum_ {i = 1} ^ {h} \boldsymbol {\eta} _ {i} \left(\phi \odot \boldsymbol {\alpha} \odot \boldsymbol {\beta}\right) _ {i}, \dots , \sum_ {i = 1} ^ {h} \boldsymbol {\eta} _ {i} \left(\phi^ {n} \odot \boldsymbol {\alpha} \odot \boldsymbol {\beta}\right) _ {i}\right) \tag {9}
$$
where  $\pmb{\eta} \in \mathbb{R}^{h \times d_m}$ ,  $\pmb{\alpha} \in [0,1]^{h \times d_m}$ ,  $\pmb{\delta} \in [0,1]^{h \times d_m}$ ,  $\pmb{\beta} \in \mathbb{R}^{h \times d_m}$ ,  $\phi = 1 - \pmb{\alpha} \odot \pmb{\delta}$ , are learnable parameters,  $\odot$  means the element-wise multiplication and  $h$  is a hyperparameter denoted as the EMA dimension. With the pre-computed kernel  $\mathbf{K}$ , the output of MD-EMA can be efficiently computed with FFTs in a time complexity of  $O(n \log n)$ . During sequential decoding stage at the time step  $t$ , the output of MD-EMA is calculated as
$$
u _ {t, i j} = \beta_ {i j} s _ {t, j}, \mathbf {s} _ {t} = [ s _ {t, j} ] \in \mathbb {R} ^ {d _ {m}}
$$
$$
\mathbf {u} _ {t} = \left[ u _ {t, i j} \right] \in \mathbb {R} ^ {h \times d _ {m}}
$$
$$
\mathbf {z} _ {t} = \boldsymbol {\alpha} \odot \mathbf {u} _ {t} + (1 - \boldsymbol {\alpha} \odot \boldsymbol {\delta}) \odot \mathbf {z} _ {t - 1}
$$
$$
\mathbf {M D - E M A} (\mathbf {s} _ {t}) = \sum_ {i = 1} ^ {h} \boldsymbol {\eta} _ {i} \mathbf {z} _ {t, i} + \mathbf {D} \odot \mathbf {s} _ {t}
$$
where  $\mathbf{z}_0 = \mathbf{0}$ ,  $\mathbf{s}_t \in \mathbb{R}^{d_m}$  is the input representation at time  $t$ , and  $\mathbf{D} \in \mathbb{R}^{d_m}$  are learnable parameters. We can see that the inference of MD-EMA has  $O(1)$  complexity per time step that is independent of the sequence length  $n$ .
A.4 Details of Latent Configurator with Gumbel Softmax
To produce this baseline, we apply Gumbel Softmax to the probability computation in the latent configurator described in Section 3.3. Specifically, the modified probability calculation is as follow,
$$
\mathbf {p} _ {i} ^ {\prime} = \frac {e ^ {(\mathbf {H} \mathbf {w} _ {i} + b _ {i} + G _ {i}) / \tau^ {\prime}}}{e ^ {(\mathbf {H} \mathbf {w} _ {0} + b _ {0} + G _ {i}) / \tau^ {\prime}} + e ^ {(\mathbf {H} \mathbf {w} _ {1} + b _ {1} + G _ {i}) / \tau^ {\prime}}} \quad \text {f o r} i \in \{0, 1 \}
$$
where  $G_{i} \sim \mathrm{Gumbel}(0,1)$  is the Gumbel noise sampled from a standard Gumbel distribution and  $\tau'$  follows the following temperature annealing schedule as in the previous work [BZMA20],
$$
\tau^ {\prime} = \max  (2 \times 0. 9 9 9 9 9 5 ^ {T}, 0. 5),
$$
where  $T$  is the number of training steps.
Table 4: Hyper-parameter Settings of our SeqBoat model for the LRA benchmark and the Speech Command (SC) dataset. DP is the dropout rate, BSZ is batch size, LR is learning rate, WD is weight decay, and Pre-N is Pre-normalization. BN, LN and SN refer to Batch Normalization, Layer Normalization and Scale Normalization.
Task	Depth	dm	α	dz	dv	Attn-FN	Norm	Pre-N	BSZ	LR	DP	WD	Epochs
ListOps	6	80	0.3	64	160	softmax	LN	FALSE	64	0.004	0.1	0.001	60
Text	4	128	0.3	64	256	softmax	SN	FALSE	50	0.004	0.1	0.01	50
Retrieval	6	128	0.3	64	256	softmax	SN	FALSE	64	0.003	0.1	0.04	40
Image	8	160	0.4	96	320	relu2	BN	TRUE	50	0.01	0	0.02	200
Pathfinder	6	128	1.0	64	256	relu2	BN	TRUE	128	0.01	0	0.01	200
Path-X	6	128	1.0	64	256	relu2	BN	TRUE	128	0.01	0	0.01	100
SC-Raw	8	60	1.0	30	120	relu2	BN	TRUE	20	0.01	0	0.01	200
A.5 Hyper-parameter settings
For Long Range Arena (LRA) and Speech Command tasks, we use the AdamW [LH18] optimizer and the detailed settings together with our initial temperature scale  $\alpha$  for the latent configurator are shown in Table 4. Following S5 [SWL23], we use NLI-style fusion of encoder features for the Retrieval task of LRA. For SeqBoat model, we use a working memory size  $w = 512$  for Path-X and  $w = 256$  for other tasks in LRA. For Speech Command, a working memory size of 256 is used. And for language modeling tasks, we use the RAdam [LJH+19] optimizer and a working memory of size 1,024. We use RoPE relative position embeddings and the Softmax attention function for language modeling, and the detailed hyperparameters are shown in Table 5. For the speech command task and the Text task of LRA, the relative positions are calculated based on the positions in the compressed sequence.
A.6 Details of Efficiency Measurements
Due to the dynamic sparsity of our model, we first measure the training time as the average per step wall time across the full training stage. The training speed is then calculated as the inverse of the training time. The memory consumption is measured as the average per step GPU memory allocation across the full training stage. To ensure fair comparisons, the relative training speedup between different models are calculated based on the training time measured on the same hardware using the same batch size settings. All the experiments are conducted on a mixed cluster with 8 NVIDIA V100 32GB GPUs and 2 NVIDIA A5000 24GB GPUs.
An alternative way of efficiency measurement is to calculate the total Floating Point Operations (FLOPs) used for training. However, the FLOPs calculation does not take account of the time-consuming memory copying of large matrices (which is exactly the speed bottleneck of SMA due to the scatter operation), and thus we choose to follow the previous works to measure the speedup based on the actual wall time for a stricter and fairer comparison.
Table 5: Hyper-parameters of our SeqBoat model for language modeling.
enwik8
Sequence Length × Batch Size	8192 × 8
Optimizer	RAdam
Learning Rate	0.005
Initial Learning Rate	0.002
RAdam-β	(0.9, 0.98)
Learning Rate Decay	linear
Weight Decay	0.1
Dropout	0.15
Attention Dropout	0.0
Gradient Clipping	0.25
Warmup steps	24K
Total updates	400K
Decoder Layers	14
Model Size (dm)	616
Query and Key Sizes (dz)	128
Value Size (dv)	1232
EMA dimension (h)	16
Initial Temperature Scale (α)	1.0
Working Memory Size (w)	1024
Total Parameters	39M
Proof. Given any  $\mathcal{L}' \subseteq \mathcal{L} = \operatorname{span}\{f_1^l, \ldots, f_M^l\}$ , we want to show that there exists a pair  $(\mathbf{a}_t', \mathbf{c}_t')$  such that  $\mathcal{L}_{\mathrm{SMA}}(\mathbf{a}_t', \mathbf{c}_t') = \mathcal{L}'$ . Since  $\mathcal{L}' \subseteq \mathcal{L}$ , every function in  $\mathcal{L}'$  can be written as a linear combination
of the functions  $\{f_1^l,\ldots ,f_M^l\}$ . Let's consider a function  $g\in \mathcal{L}^{\prime}$ . Then, we can write
$$
g = \sum_ {i = 1} ^ {M} \beta_ {i} f _ {i} ^ {l}
$$
where  $\beta_{i}\in \mathbb{R}$ . Now, let's define  $\mathbf{a}_t^\prime$  and  $\mathbf{c}_t^\prime$  as follows:
-  $\mathbf{a}_t^{\prime} = (a_t^{\prime i})$  where  $a_{t}^{\prime i} = 1$  if  $\beta_{i}\neq 0$  and  $a_{t}^{\prime i} = 0$  otherwise.
-  $\mathbf{c}_t' = (c_t'^i)$  where  $c_t'^i = 1$  if  $\beta_i \neq 0$  and  $c_t'^i = 0$  otherwise.
$$
\mathcal {L} _ {\mathrm {S M A}} \left(\mathbf {a} _ {t} ^ {\prime}, \mathbf {c} _ {t} ^ {\prime}\right) = \left\{\sum_ {i \in I} \zeta_ {i} c _ {t} ^ {\prime i} f _ {i} ^ {l} \mid \forall \zeta_ {i} \in \mathbb {R}, I = \{i \mid a _ {t} ^ {\prime i} = 1, 1 \leq i \leq M \} \right\}
$$
Since  $a_{t}^{i} = 1$  and  $c_{t}^{i} = 1$  if and only if  $\beta_{i}\neq 0$ , then,  $I = \{i\mid \beta_i\neq 0,1\leq i\leq M\}$ , we have
$$
\mathcal {L} _ {\mathrm {S M A}} \left(\mathbf {a} _ {t} ^ {\prime}, \mathbf {c} _ {t} ^ {\prime}\right) = \left\{\sum_ {i, \beta_ {i} \neq 0} \zeta_ {i} f _ {i} ^ {l} \mid \forall \zeta_ {i} \in \mathbb {R} \right\} = \left\{\sum_ {i = 1} ^ {M} \beta_ {i} f _ {i} ^ {l} \mid \forall \beta_ {i} \in \mathbb {R} \right\} = \mathcal {L} ^ {\prime}
$$
C Theorem of Implicit Regularization for Latent Decision Making
We first give a brief introduction of the implicit regularization to provide some backgrounds. The implicit regularization of gradient descent [BD20] is a side effect caused by the discreteness of the gradient descent updates. Specifically, it has been proved that a discrete gradient update  $\theta' = \theta - h\nabla_{\theta}L(\theta)$  of the parameter  $\theta$  is implicitly minimizing a transformed loss,  $L' = L + \frac{h}{4} ||\nabla_{\theta}L(\theta)||^2$  where  $h$  is the learning rate and  $L$  is the training loss term. Assuming the usage of Stochastic Gradient Decent (SGD) for gradient updates, [DMRB22] have shown that the implicit gradient regularizer  $\| \nabla_{\theta}L(\theta)\|^2$  will conditionally put an implicit pressure on  $\| \nabla_{\theta}f_{\theta}(x)\|_F^2$  for a multi-class classification cross-entropy loss term, where  $f_{\theta}(x)$  is a neural network that takes the input  $x$  and outputs the logits for classification. A Transfer Theorem [DMRB22] is further derived to characterize the conditional transfer of the implicit regularization pressure from  $\| \nabla_{\theta}f_{\theta}(x)\|_F^2$  to the gradient norm  $\| \nabla_x f_{\theta}(x)\|_F^2$ , where  $\| \cdot \|_F$  means the Frobenius norm. The transfer theorem assumes an ordinary neural architecture  $f_{\theta}(x): \mathbb{R}^d \mapsto \mathbb{R}^k$  with  $l$  layers in the following form,
$$
h _ {i} (x) = a _ {i} \left(w _ {h _ {i}} h _ {i - 1} (x) + b _ {i}\right), \quad \text {f o r} i = 1, 2, \dots , l,
$$
where  $a_{i}$  is the activation function,  $h_i(x)$  indicates the sub-network from the input  $x \in \mathbb{R}^d$  up to the output of layer  $i$  and  $h_0(x) = x$ . This implies  $h_l(x) = f_\theta(x)$ .
In this work, we consider a neural architecture that has a per-layer latent configurator  $p_i: \mathbb{R}^{|h_{i-1}|} \mapsto [0,1], p_i(z) = y_i(w_iz)$ , that maps the hidden representation from  $h_{i-1}$  to a confidence probability of activating a submodule  $q_i: \mathbb{R}^{|h_{i-1}|} \mapsto \mathbb{R}^{|h_i|}$ , i.e.,
$$
h _ {i} (x) = a _ {i} \left(p _ {i} \left(h _ {i - 1} (x)\right) q _ {i} \left(h _ {i - 1} (x)\right)\right), \quad \text {f o r} i = 1, 2, \dots , l,
$$
where  $y_{i}$  is the normalizing function to produce the probability. We use a similar technique of perturbation argument [DMRB22;  $\mathrm{SLK}^{+}18$ ] for proving that the implicit pressure will also transfer to the latent configurator  $p_{i}$ , which leads to the implicit regularization of latent activation decisions.
Theorem 2 (Transfer Theorem for Latent Configurator). Let  $f_{\theta}:\mathbb{R}^{d}\to \mathbb{R}^{k}$  represent a deep neural network consisting of  $l$  layers, i.e.,  $f_{\theta}(x) = f_{l}\circ f_{l - 1}\circ \dots \circ f_{1}(x)$  where  $f_{i}(z) = a_{i}(p_{i}(z)q_{i}(z))$ ,  $p_i:\mathbb{R}^{|z|}\mapsto [0,1]$ ,  $p_i(z) = y_i(w_iz)$  is a latent configurator;  $q_{i}:\mathbb{R}^{|z|}\mapsto \mathbb{R}^{|h_{i}|}$  is a sub-module function,  $a_{i}$  is the activation function,  $y_{i}$  is the probability normalization function and  $w_{i}$  is a weight matrix. For  $i = 1,2,\ldots ,l$ , we have
$$
\left| \left| \nabla_ {w _ {i}} r _ {i} (x) \right| \right| _ {2} ^ {2} \frac {\left| \left| d _ {i} (x) \right| \right| _ {2} ^ {2} \left| \left| \nabla_ {x} f _ {\theta} (x) \right| \right| _ {2} ^ {2}}{\left| \left| d _ {i} (x) \nabla_ {x} r _ {i} (x) + r _ {i} (x) \nabla_ {x} d _ {i} (x) \right| \right| _ {2} ^ {2}} \leq \left| \left| \nabla_ {w _ {i}} f _ {\theta} (x) \right| \right| _ {2} ^ {2} \tag {10}
$$
where  $r_i(x) \coloneqq p_i(h_{i-1}(x))$  and  $d_i(x) \coloneqq q_i(h_{i-1}(x))$ ,  $h_i(x)$  indicates the sub-network from the input  $x \in \mathbb{R}^d$  up to the output of layer  $i$  with  $h_0(x) = x$  and  $||\cdot||_2$  denotes the L2 operator norm for matrix and L2 norm for vector.
Proof. Following the above notations, we first rewrite the model function as
$$
f _ {w _ {i}} (x) = g _ {i} \left(r _ {i} (x) d _ {i} (x)\right),
$$
where  $g_{i}(z)$  is the remaining network of the full model  $f_{\theta}(x)$  following the  $i$ -th layer. Considering a small perturbation  $x + \delta x$  of the input  $x$ , there exists a corresponding perturbation  $w_{i} + u(x)$  for the weight matrix in the latent configurator  $p_{i}$  so that
$$
f _ {w _ {i}} (x + \delta x) = f _ {w _ {i} + u (x)} (x). \tag {11}
$$
For sufficiently small  $\delta x$ , we can identify  $r_i(x + \delta x)d_i(x + \delta x)$  with the linear approximation around  $x$ , i.e.,
$$
r _ {i} (x + \delta x) d _ {i} (x + \delta x) = r _ {i} (x) d _ {i} (x) + d _ {i} (x) \nabla_ {x} r _ {i} (x) \delta x + r _ {i} (x) \nabla_ {x} d _ {i} (x) \delta x
$$
where  $\nabla_{x}d_{i}:\mathbb{R}^{d}\mapsto \mathbb{R}^{|h_{i}|}$  and  $\nabla_{x}r_{i}:\mathbb{R}^{d}\mapsto [0,1]$  are the total derivatives. We have
$$
\begin{array}{l} f _ {w _ {i}} (x + \delta x) = g _ {i} \left(r _ {i} (x + \delta x) d _ {i} (x + \delta x)\right) \\ = g _ {i} \left(r _ {i} (x) d _ {i} (x) + d _ {i} (x) \nabla_ {x} r _ {i} (x) \delta x + r _ {i} (x) \nabla_ {x} d _ {i} (x) \delta x\right), \\ \end{array}
$$
and similarly, for  $r_i(x) = y_i(w_i h_{i-1}(x))$ , and the perturbation  $w_i \to w_i + u(\delta x)$ , we can have the linear approximation  $r_i'(x) = y_i((w_i + u(\delta x)) h_{i-1}(x)) = r_i(x) + \nabla_{w_i} r_i(x) u(\delta x)$ . This implies
$$
f _ {w _ {i} + u (\delta x)} (x) = g _ {i} \left(r _ {i} (x) d _ {i} (x) + \nabla_ {w _ {i}} r _ {i} (x) u (\delta x) d _ {i} (x)\right).
$$
Therefore, substituting both sides of Equation (11) with the equations above, and we get
$$
\nabla_ {w _ {i}} r _ {i} (x) u (\delta x) d _ {i} (x) = d _ {i} (x) \nabla_ {x} r _ {i} (x) \delta x + r _ {i} (x) \nabla_ {x} d _ {i} (x) \delta x.
$$
Note the fact that  $\nabla_{w_i}r_i(x)^T\nabla_{w_i}r_i(x) = ||\nabla_{w_i}r_i(x)||_2^2$ . This implies
$$
u (\delta x) = \frac {\nabla_ {w _ {i}} r _ {i} (x) ^ {T} \left(d _ {i} (x) \nabla_ {x} r _ {i} (x) \delta x + r _ {i} (x) \nabla_ {x} d _ {i} (x) \delta x\right) d _ {i} (x) ^ {T}}{\left| \left| \nabla_ {w _ {i}} r _ {i} (x) \right| \right| _ {2} ^ {2} \left| \left| d _ {i} (x) \right| \right| _ {2} ^ {2}}. \tag {12}
$$
Defining  $u(\delta x)$  as above and taking the derivative of both sides of Equation (11) with respect to  $\delta x$  at  $\delta x = 0$ , we have
$$
\nabla_ {x} f _ {w _ {i}} (x) = \nabla_ {w _ {i}} f _ {\theta} (x) \frac {\nabla_ {w _ {i}} r _ {i} (x) ^ {T} (d _ {i} (x) \nabla_ {x} r _ {i} (x) + r _ {i} (x) \nabla_ {x} d _ {i} (x)) d _ {i} (x) ^ {T}}{| | \nabla_ {w _ {i}} r _ {i} (x) | | _ {2} ^ {2} | | d _ {i} (x) | | _ {2} ^ {2}}.
$$
Applying the squared L2 operator norm  $||\cdot ||_2^2$ , at both sides of the equation, we get
$$
\begin{array}{l} \left| \left| \nabla_ {x} f _ {w _ {i}} (x) \right| \right| _ {2} ^ {2} = \left\| \nabla_ {w _ {i}} f _ {\theta} (x) \frac {\nabla_ {w _ {i}} r _ {i} (x) ^ {T} \left(d _ {i} (x) \nabla_ {x} r _ {i} (x) + r _ {i} (x) \nabla_ {x} d _ {i} (x)\right) d _ {i} (x) ^ {T}}{\left| \left| \nabla_ {w _ {i}} r _ {i} (x) \right| \right| _ {2} ^ {2} \left| \left| d _ {i} (x) \right| \right| _ {2} ^ {2}} \right\| _ {2} ^ {2} \\ \leq | | \nabla_ {w _ {i}} f _ {\theta} (x) | | _ {2} ^ {2} \frac {| | d _ {i} (x) \nabla_ {x} r _ {i} (x) + r _ {i} (x) \nabla_ {x} d _ {i} (x) | | _ {2} ^ {2}}{| | \nabla_ {w _ {i}} r _ {i} (x) | | _ {2} ^ {2} | | d _ {i} (x) | | _ {2} ^ {2}}. \\ \end{array}
$$
Since  $\nabla_x f_{w_i}(x) = \nabla_x f_\theta (x)$ , this completes the proof.
Remarks Considering the architecture of the latent configurator in the SeqBoat layer, i.e., the activation confidence probability is calculated as
$$
p _ {i} (z) = \frac {e ^ {z \cdot w _ {i} ^ {1} / \tau_ {i}}}{e ^ {z \cdot w _ {i} ^ {0} / \tau_ {i}} + e ^ {z \cdot w _ {i} ^ {1} / \tau_ {i}}},
$$
then we can derive the squared L2 norms of its derivatives,
$$
\begin{array}{l} | | \nabla_ {w _ {i} ^ {1}} p _ {i} (z) | | _ {2} ^ {2} = | | \nabla_ {w _ {i} ^ {0}} p _ {i} (z) | | _ {2} ^ {2} = (p _ {i} (z) (1 - p _ {i} (z)) | | z | | _ {2} / \tau_ {i}) ^ {2}, \\ | | \nabla_ {z} p _ {i} (z) | | _ {2} ^ {2} = (p _ {i} (z) (1 - p _ {i} (z)) | | w _ {i} ^ {0} - w _ {i} ^ {1} | | _ {2} / \tau_ {i}) ^ {2}. \\ \end{array}
$$
Substituting the above gradient values to Equation (10), we can see that
$$
\begin{array}{l} \left| \left| \nabla_ {w _ {i}} r _ {i} (x) \right| \right| _ {2} ^ {2} \frac {\left| \left| d _ {i} (x) \right| \right| _ {2} ^ {2} \left| \left| \nabla_ {x} f _ {\theta} (x) \right| \right| _ {2} ^ {2}}{\left| \left| d _ {i} (x) \nabla_ {x} r _ {i} (x) + r _ {i} (x) \nabla_ {x} d _ {i} (x) \right| \right| _ {2} ^ {2}} \leq \left| \left| \nabla_ {w _ {i}} f _ {\theta} (x) \right| \right| _ {2} ^ {2} \\ \Longrightarrow \frac {\left| \left| h _ {i - 1} (x) \right| \right| _ {2} ^ {2} \left| \left| d _ {i} (x) \right| \right| _ {2} ^ {2} \left| \left| \nabla_ {x} f _ {\theta} (x) \right| \right| _ {2} ^ {2}}{\left| \left| d _ {i} (x) \right| \right| _ {2} ^ {2} \left| \left| \nabla_ {x} h _ {i - 1} (x) \right| \right| _ {2} ^ {2} \left| \left| w _ {i} ^ {0} - w _ {i} ^ {1} \right| \right| _ {2} ^ {2} + \frac {\left| \left| \nabla_ {x} d _ {i} (x) \right| \right| _ {2} ^ {2} \tau_ {i} ^ {2}}{\left(1 - p _ {i} \left(h _ {i - 1} (x)\right)\right) ^ {2}}} \leq \left| \left| \nabla_ {w _ {i}} f _ {\theta} (x) \right| \right| _ {2} ^ {2} \\ \end{array}
$$
This implies that the implicit regularization pressure will tend to maximize the temperature  $\tau_{i}$  to increase the exploration of the latent decisions, while at the same time maximize the probability of modular activation or de-activation through maximizing  $||w_i^0 - w_i^1||_2^2$  and  $p_i(h_{i-1}(x))$  for exploiting the decisions. Therefore, there will be an implicit balancing between the exploration and the exploitation of the latent decisions, and we conjecture that this kind of balancing is the cause of the dynamic sparse activation patterns we observed from the experiments.
Table 6: Per-task speed up and memory allocation reduction on the LRA benchmark.
Models (Input Length)	ListOps (2,048)	Text (4,096)	Retrieval (4,000)	Image (1,024)	Path (1,024)	Path-X (16,384)
Training Step Speed Up (↑)
MEGA	1.00×	1.00×	1.00×	1.00×	1.00×	1.00×
MEGA-chunk	2.07×	2.42×	2.66×	1.39×	1.46×	1.36×
SeqBoat-full	1.79×	2.14×	2.44×	1.45×	1.27×	1.24×
SeqBoat	2.26×	3.61×	2.50×	1.33×	1.31×	1.30×
Training Step Memory Allocation (↓)
MEGA	1.00×	1.00×	1.00×	1.00×	1.00×	1.00×
MEGA-chunk	0.34×	0.29×	0.31×	0.75×	0.69×	0.82×
SeqBoat-full	0.45×	0.23×	0.27×	0.74×	0.68×	0.59×
SeqBoat	0.20×	0.16×	0.23×	0.74×	0.70×	0.89×
Evaluation Step Speed Up (↑)
MEGA	1.00×	1.00×	1.00×	1.00×	1.00×	1.00×
MEGA-chunk	2.15×	2.33×	2.57×	1.48×	1.50×	2.52×
SeqBoat-full	1.70×	2.47×	2.59×	1.52×	1.30×	1.95×
SeqBoat	2.67×	3.55×	2.69×	1.43×	1.38×	1.93×
D.1 Per-task Efficiency Measures for LRA Benchmark
As shown in Table 6, we compare the efficiency of our SeqBoat-full and SeqBoat models with the Mega and its chunk variants on all six tasks of the LRA benchmark. The evaluation speed is measured as the average inference throughput of the best performing model on the validation set. More details of the measurement are included in Appendix A.6. We can see that our SeqBoat model is consistently and substantially more efficient than Mega, and it also has comparable efficiency with Mega-chunk but significantly better average accuracy on the LRA benchmark.
E Additional Analysis
How do different model design choices affect the performance? We do the ablation study of our SeqBoat-full model on the Image and ListOps tasks of the LRA benchmark, as in Table 7, to have a general understanding of the influence of the design of the latent configurator and the SSM. "Always Activated" means that we don't use SMA and always activate GAU for conducting second-order self-attention. We can see that although the validation accuracy is slightly improved on both Image and Listops, the training efficiency drops significantly. We also explore different parameterization of the convolution kernel for the linear SSM and apply S4D [GGGR22] to replace the original MD-EMA  $\left[\mathrm{MZK}^{+}23\right]$  kernel. The S4D variant of our model has much better training efficiency but the quality
Table 7: Ablation study of the SeqBoat-full model on Image and ListOps of the LRA benchmark. We report the validation accuracy on the Image and ListOps tasks, together with the training speed up for each of the tasks compared with MEGA.
Model	Image	Speed (↑)	ListOps	Speed (↑)
SeqBoat-full	90.16	1.45×	60.95	1.79×
- Always Activated	90.70	1.25×	62.10	1.11×
- Gumbel Softmax	89.94	1.29×	59.75	1.76×
- SSM with S4D	82.84	1.88×	59.25	2.45×
are significantly worse than SeqBoat-full with the default hyperparameter setting. It is possible that a well-tuned S4D kernel can have better performance than the MD-EMA for the use case of SMA but we leave a more comprehensive study as in future work. We also explore different differentiable approximation of latent decision making by replacing Tempered Softmax with Gumbel Softmax [JGP17]. However, we don't observe the performance improvements over Tempered Softmax for our setting. More details of the implementation of Gumbel Softmax based latent configurator is shown in Appendix A.4.
Table 8: Ablation study of the SeqBoat model on Image, ListOps and Pathfinder of the LRA benchmark. We report the validation accuracy on each of the tasks, together with the training speed up compared with MEGA.
Model	Image	Speed (↑)	ListOps	Speed (↑)	Path.	Speed (↑)
SeqBoat	90.36	1.33×	59.05	2.26×	96.34	1.31×
SeqBoat-FFN	89.60	1.13×	56.45	2.05×	96.62	0.97×
SeqBoat-{GAU, FFN}	89.74	1.20×	59.10	1.97×	96.10	1.06×
SeqBoat-Sigmoid	87.34	1.77×	58.20	2.52×	91.35	1.67×
SeqBoat-MoE	85.60	1.66×	58.10	2.19×	91.41	1.80×
SeqBoat-local	90.12	1.46×	59.00	1.78×	96.08	1.54×
SeqBoat-chunk	86.46	2.17×	58.35	2.79×	93.91	2.97×
SeqBoat-mem32	78.16	1.73×	58.70	2.14×	92.97	1.86×
SeqBoat-local32	75.58	1.67×	53.90	2.19×	82.05	1.80×
We further present additional ablation studies on our SeqBoat model to justify the effectiveness of both the proposed working memory mechanism and the design choices of our architecture. We explain the meaning of the postfixes in the table as following: