将Transformer与Mamba混合，高效处理超长文本。
MaBERT:A Padding Safe Interleaved Transformer Mamba Hybrid Encoder for Efficient Extended Context Masked Language Modeling
- 交替使用Transformer与Mamba层，兼顾全局依赖与计算效率。
- 在4096词长下，训练时间减少2.36倍，推理延迟降低2.43倍。
- 适合需要长上下文建模的NLP任务，如文本理解与推理。
自注意力编码器如BERT随序列长度呈平方级增长，使长上下文建模成本高昂。线性时间状态空间模型如Mamba虽高效，但难以建模全局交互，且易受填充带来的状态污染影响。我们提出MaBERT，一种混合编码器，将用于全局依赖建模的Transformer层与用于线性时间状态更新的Mamba层交错设计。该结构交替进行全局上下文整合与快速状态累积，实现对长输入的高效训练与推理。为稳定变长批量处理，引入安全填充掩码，阻止填充位置的状态传播，并采用掩码感知注意力池化，仅聚合有效标记信息。在GLUE基准上，MaBERT在八项任务中的五项取得最优平均得分，尤其在CoLA与句子对推理任务表现突出。当上下文从512扩展至4096词时，相比编码器基线平均值，训练时间减少2.36倍，推理延迟降低2.43倍，验证了其作为实用长上下文高效编码器的潜力。
原文摘要 · Abstract (English)
Self attention encoders such as Bidirectional Encoder Representations from Transformers(BERT) scale quadratically with sequence length, making long context modeling expensive. Linear time state space models, such as Mamba, are efficient; however, they show limitations in modeling global interactions and can suffer from padding induced state contamination. We propose MaBERT, a hybrid encoder that interleaves Transformer layers for global dependency modeling with Mamba layers for linear time state updates. This design alternates global contextual integration with fast state accumulation, enabling efficient training and inference on long inputs. To stabilize variable length batching, we introduce paddingsafe masking, which blocks state propagation through padded positions, and mask aware attention pooling, which aggregates information only from valid tokens. On GLUE, MaBERT achieves the best mean score on five of the eight tasks, with strong performance on the CoLA and sentence pair inference tasks. When extending the context from 512 to 4,096 tokens, MaBERT reduces training time and inference latency by 2.36x and 2.43x, respectively, relative to the average of encoder baselines, demonstrating a practical long context efficient encoder.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。