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.
翻译:线性状态空间模型(SSMs)因其对循环结构的高效编码,在多种序列建模任务中展现出强劲性能。然而,在语言建模和机器翻译等更复杂的任务中,基于自注意力的模型仍优于SSMs。同时采用SSM与自注意力的混合模型通常表现优异,但现有方法将注意力模块静态且均匀地应用于输入序列的所有元素,导致质量与效率的权衡次优。本文提出稀疏模块激活(SMA)机制,该通用机制能以可微方式使神经网络对序列元素稀疏动态地激活子模块。通过允许各元素跳过未激活子模块,SMA在序列建模的训练与推理阶段均可降低计算开销与内存消耗。作为SMA的特定实例,我们设计了新型神经网络架构SeqBoat,其通过SMA基于SSM学习到的状态表征稀疏激活门控注意力单元(GAU)。通过约束GAU仅对激活输入进行局部注意力计算,SeqBoat能以理论上无限注意力跨度实现线性推理复杂度,并提供显著优于基于分块模型的质量-效率权衡。在语言建模、语音分类及长程竞技场等广泛任务的实验中,SeqBoat在线性复杂度混合模型中取得了最新的最优结果,并通过学习到的稀疏激活模式揭示了各任务所需的注意力规模。