Linear State Space Models (SSMs) have demonstrated strong performance in a variety of sequence modeling tasks due to their efficient encoding of the recurrent structure. However, in more comprehensive tasks like language modeling and machine translation, self-attention-based models still outperform SSMs. Hybrid models employing both SSM and self-attention generally show promising performance, but current approaches apply attention modules statically and uniformly to all elements in the input sequences, leading to sub-optimal quality-efficiency trade-offs. In this work, we introduce Sparse Modular Activation (SMA), a general mechanism enabling neural networks to sparsely and dynamically activate sub-modules for sequence elements in a differentiable manner. Through allowing each element to skip non-activated sub-modules, SMA reduces computation and memory consumption at both training and inference stages of sequence modeling. As a specific instantiation of SMA, we design a novel neural architecture, SeqBoat, which employs SMA to sparsely activate a Gated Attention Unit (GAU) based on the state representations learned from an SSM. By constraining the GAU to only conduct local attention on the activated inputs, SeqBoat can achieve linear inference complexity with theoretically infinite attention span, and provide substantially better quality-efficiency trade-off than the chunking-based models. With experiments on a wide range of tasks, including language modeling, speech classification and long-range arena, SeqBoat brings new state-of-the-art results among hybrid models with linear complexity and reveals the amount of attention needed for each task through the learned sparse activation patterns.
翻译:线性状态空间模型(Linear State Space Models, SSMs)因其对循环结构的高效编码,在多种序列建模任务中表现出强劲性能。然而,在语言建模和机器翻译等更复杂的任务中,基于自注意力的模型仍优于SSMs。同时采用SSM和自注意力的混合模型通常展现出有前景的性能,但当前方法将注意力模块静态且均匀地应用于输入序列中的所有元素,导致质量-效率权衡次优。在这项工作中,我们提出了稀疏模块化激活(Sparse Modular Activation, SMA),这是一种通用机制,使神经网络能够以可微的方式对序列元素进行稀疏且动态的子模块激活。通过允许每个元素跳过未激活的子模块,SMA在序列建模的训练和推理阶段均能减少计算量和内存消耗。作为SMA的具体实例化,我们设计了一种新颖的神经架构SeqBoat,它利用SMA基于从SSM学习到的状态表示,稀疏地激活门控注意力单元(Gated Attention Unit, GAU)。通过将GAU限制为仅对激活的输入执行局部注意力,SeqBoat能够实现理论上无限注意力跨度的线性推理复杂度,并在质量-效率权衡上显著优于基于分块的模型。通过在包括语言建模、语音分类和长程竞技场(Long-Range Arena)在内的广泛任务上进行实验,SeqBoat在线性复杂度混合模型中取得了新的最先进结果,并通过学习到的稀疏激活模式揭示了每个任务所需的注意力量。