SlowGuess/NeurIPS_2023_OCR huggingface dataset:
- "FFN": We add an additional Feed-Forward layer (with the same architecture as MEGA) after each SeqBoat layer.
- “ $\{-\{\mathrm{GAU,FFN}\}\}$ : We add a FFN module in parallel with the GAU module to have  $M = 2$  with two modules in the function space  $\mathcal{F} = \{\mathrm{GAU,FFN}\}$  for our SMA mechanism to control the activation.
- "Sigmoid": We use Sigmoid instead of Tempered Softmax for latent decision probability calculation.
- " -MoE": Instead of using our Tempered Softmax for module activation decisions, we adapt the X-MoE [JJNH91] routing mechanism to our use case by only considering two modules, a GAU module and a module that always outputs zero. We do not apply load balancing loss in this case because there is only one module that needs to do actual computation.
- " -local": We don't use SMA and always activate the GAU module with local attention.
- "chunk": We don't use SMA and always activate the GAU module to process the equal-sized chunks of the input sequence. The chunk size is set to be the same as the memory size of SeqBoat for fair comparison.
- "mem32": We restrict the memory size of SeqBoat to 32.
- " -local32": We restrict the sliding window size of SeqBoat-local to 32.
We first explore the application of SMA with multiple heterogeneous modules and identify the current engineering challenges. We can see that SeqBoat-{GAU,FFN} obtains better trade-offs than
SeqBoat-FFN (which simply appends an extra FFN layer after each SeqBoat layer) on both the Image and the ListOps tasks. However, SeqBoat-{GAU,FFN} still falls behind SeqBoat (which does not include FFN anywhere at all). This is because the current implementation of SMA is I/O bounded due to copying large tensors with the scatter operator, and for a fair comparison with MEGA, we do not include any fused kernel to optimize the memory bandwidth utilization. We acknowledge that it needs more engineering efforts to scale up SMA with multiple modules, and advocate future works for the empirical impacts of SMA on the MoE community. We also include additional ablation results with X-MoE routing as an alternative design choice for the latent configurator. We can see that SeqBoat-MoE performs much worse than our SeqBoat model which uses the Tempered Softmax for decision confidence calculation.
Comparing SeqBoat-Sigmoid with SeqBoat, we can see that while theoretically using Softmax is equivalent to using Sigmoid under the binary decision scenario, using Softmax for decision confidence calculation gives much better results under our hyper-parameter settings. Comparing SeqBoat with SeqBoat-local and SeqBoat-chunk, we can see that our model obtains better accuracy than the models without SMA. If we constrain the memory size to a smaller value (see SeqBoat-mem32 v.s. SeqBoat-local32), this performance gap becomes more substantial (more than 10 percent absolute difference on Pathfinder). These evidences demonstrate the effectiveness of our SMA mechanism for capturing long-term dependencies. Comparing SeqBoat-FFN with SeqBoat, we can see that adding extra FFN layers does not provide any empirical performance gains for SeqBoat, and this motivates us to not include FFNs anywhere in our architecture.
Dynamic Routing [RSVG20b] aims to alleviate the second-order complexity of self-attention by dynamically mapping each of the token to its nearest-neighbor clusters, and only do self-attention inside the clusters. Dynamic routing can be understood as the following form when it is applied for sequence modeling,
$$
\max  _ {\theta , \mathbf {x} _ {<   t} ^ {c}} P _ {f} (\mathbf {x}) = \prod_ {t = 1} ^ {n} P _ {f} (x _ {t} | \mathbf {x} _ {<   t} ^ {c}, \theta)
$$
where  $f$  is the Transformer architecture with self-attention layers, and it can be seen as a special case of our formulation described in Equation (1). Even though dynamic routing learns both the parameters and the corresponding context for each of the token prediction, it has a fixed architecture  $f$  that is non-adaptive for both training and the inference.
Mixture of Experts (MoE) [JJNH91; FZS22;  $\mathrm{LLX}^{+}20$ ; ZLL  $^{+}22]$  is designed to learn to select different parameters with the same architecture for each of the input elements. Specifically, it has the following form when applied to sequence modeling,
$$
\max  _ {\theta_ {t}} P _ {f} (\mathbf {x}) = \prod_ {t = 1} ^ {n} P _ {f} \left(x _ {t} | \mathbf {x} _ {<   t}, \theta_ {t}\right)
$$
where the model parameter  $\theta_{t}$  is selected from a pool of candidates at each time step  $t$  based on the input content. We can see that it is a special case of our formulation as in Section 3.1 because it has a fixed context  $\mathbf{x}_{< t}$  for each of the token prediction and the model architecture is non-adaptive.
Comparison between SMA and MoE: The motivation behind our Sparse Modular Activation (SMA) mechanism is to enable neural networks to contextually skip modules of any architectures for efficient sequence modeling, while MoE aims to efficiently scale up the models with more parameters (usually by adding homogeneous modules). This difference of motivation results in a fundamental mechanism difference:
- SMA leverages a latent configurator to decide if each module needs to be activated for each sequence element, while MoE is designed to choose a predefined number of modules from a large group of modules.
This difference on the core design of the mechanisms further leads to the different properties of the mechanisms:
- SMA supports a dynamic model architecture that can learn to drop a sub-module entirely (for better efficiency) based on the task it is trained on, while MoE only selects the parameters of the modules for the same architecture. As shown in Figure 3, when the SeqBoat-full model is trained on the Text task of LRA, the first three layers learn to have zero activation time of the GAU module. This means that these layers degenerate to pure SSM layers from the initial SSM+GAU architecture after adapting to the target task.
- SMA is guaranteed to have a full coverage of the combinatorial search space ranging from zero activation to full activation of modules, while MoE can only cover a subset of the space by choosing a fixed number of modules to activate.
To the best of our knowledge, SMA is the first mechanism that successfully enables a neural network to obtain practical efficiency and complexity gains from sparsely activating a self-attention-like module, while none of the previous works on MoE ever achieved this.
Comparison with Flash Attention  $\left[\mathrm{DFE}^{+}22\right]$ : Flash Attention (FA) performs exact self-attention efficiently through utilizing hardware-aware techniques. FA and SMA focus on improving model efficiency on orthogonal levels. SMA is a module-level activation strategy that is perpendicular to how the attention module is actually doing the computation. In fact, we can also apply FA to the GAU module in our SeqBoat architecture for more efficient self-attention computation, but we decide not to use any custom fused kernels in our implementation for a fair comparison with MEGA and its variant.
Comparison with Efficient Attention Strategies [KKL20; CLD+20; WLK+20; MKW+21]: Our SMA mechanism enables a dynamic architecture that can learn to drop a sub-module entirely for better efficiency when it is adapted to the target task, and we empirically observe such behavior for different tasks in LRA. This capability is not possible for the previous efficient attention strategies because their architectures are static and thus an efficient attention layer is always applied before the Feed-Forward layer. Also, SMA is an orthogonal mechanism to efficient attention approaches because it does not modify the attention itself but the input to the attention modules, and thus can be stacked with other efficient attention strategies.
Comparison with GShard  $\left[\mathbf{LLX}^{+}20\right]$ : GShard divides tokens into groups of the same size, while our SMA allows different modules to have different group size. This means SMA can support an adaptive neural architecture whose sub-modules can be completely dropped if no tokens are selected for that module. Also, while GShard is a routing mechanism, SMA is an activation mechanism that can still function when there is only one module under consideration.
Comparison with Perceiver  $\left[\mathrm{JGB}^{+}21\right]$ : Perceiver utilizes a pre-defined number of latent representations to repeatedly attend to the input array, which can be understood as conducting soft selection from the input array to a fixed subset. Our SMA operates on a different level to directly conduct a hard selection from the input to form a dynamic subset. Plus, our SMA can be applied to causal language modeling, while it is unclear how Perceiver can be adapted to this important use case in the era of large language model.
Connection with Recurrent Independent Mechanisms  $\left[\mathrm{GLH}^{+}19\right]$ : Generally, Recurrent Independent Mechanisms (RIMs) and our work both touch on how to learn sparsely interacting modules. However, our works focus more on the efficiency gains from the modular sparsity, while RIMs emphasize on the generalization benefits. Plus, while the GAU module in our work can be regarded as modeling the sparse interaction between the state representations of SSM at different time steps, our work does not explicitly try to use attention bottlenecks for modeling the interaction between different SSM modules. Instead, we apply a rather simple strategy of using linear aggregation to combine the representations from both the SSM and the GAU modules. It would be interesting to see if the attention bottleneck can help SMA to better scale up with multiple heterogeneous modules, and we leave it as a future work worth exploring.
G Limitations, Future Works & Broader Impact
Our study has some limitations that should be acknowledged. First, while our Sparse Modular Activation (SMA) mechanism is generally applicable for activating any sub-modules in neural networks, we have narrowed the scope of our paper to focus solely on developing a more efficient neural architecture with SSM and GAU for long sequence modeling, due to limited resources. Additionally, the current implementation of SMA is I/O bounded due to copying large tensors with the scatter operator. For a fair comparison with the previous models, we did not include any fused kernel to optimize memory bandwidth utilization. This limitation highlights the need for more engineering efforts to scale up SMA with multiple modules.
There are several promising directions for future research building on our work. One interesting avenue is to explore how SMA can boost the downstream performance of a pre-trained large model, given that pre-training a useful large language model from scratch can be prohibitively expensive. Additionally, since SMA can produce discrete and dynamic module activation patterns for each of the data sample, our work opens up new possibilities for interpretability in neural networks. Future research could investigate the correlations between the activation of specific modules and input tokens or model predictions, to gain a better understanding of the properties of tasks or the behavior of models. Furthermore, our mechanism may provide extra controllability over model predictions, which can be achieved by manually modifying the module activation during inference time. Finally, a promising future direction is to investigate how to scale up SMA with a large number of modules, potentially leading to empirical impacts of SMA on the research field of Mixture-of-Experts. Some initial ideas include setting an upper bound for the activation time of each module, or limiting the number of modules that can be activated at the same time.
The SMA and the SeqBoat architecture can bring about considerable societal impact by enhancing the efficiency and interpretability of sequence models. These advancements can result in better-performing automated systems, like open-domain chatbot and automatic speech recognition, potentially improving communication, fostering inclusivity, and making such services more accessible to various societal sectors. The enhanced interpretability of models can generate trust in AI systems, especially in sensitive sectors like healthcare and finance. Transparent AI models can assist professionals in these fields to make better-informed decisions, potentially leading to improved societal outcomes.
However, these improvements come with a notable limitation: our research was unable to conduct experiments on models with a large number of parameters due to computational resource constraints. As such, while the proposed mechanisms show promising results in the tested scenarios, their scalability and applicability to larger, more complex models remain uncertain. This limitation may restrict the broader applicability of our findings, especially for large-scale problems or when fine-grained control over the model's behavior is needed. Furthermore, ethical considerations such as data privacy need to be seriously considered as our work progresses. As the efficiency of these models increases, there's an increased potential for processing large amounts of personal data, which necessitates robust privacy safeguards. Future work can also focus on addressing these limitations and ethical concerns, to ensure the advantages of these models can be responsibly realized.