MaBERT:A Padding Safe Interleaved Transformer Mamba Hybrid Encoder for Efficient Extended Context Masked Language Modeling

📄 arXiv: 2603.03001v1 📥 PDF

作者: Jinwoong Kim, Sangjin Park

分类: cs.CL

发布日期: 2026-03-03

备注: 8 pages

💡 一句话要点

MaBERT：一种Padding安全的交错Transformer-Mamba混合编码器，用于高效的扩展上下文掩码语言建模

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 长文本建模 混合编码器 Transformer Mamba 状态空间模型 padding安全 高效计算

📋 核心要点

  1. 现有Transformer模型在处理长序列时计算复杂度高，限制了其在长文本建模中的应用。
  2. MaBERT通过交错Transformer和Mamba层，结合全局依赖建模和线性时间状态更新，实现高效长文本处理。
  3. 实验结果表明，MaBERT在GLUE基准测试中表现优异，并在长上下文任务中显著降低了训练和推理时间。

📝 摘要（中文）

Transformer等自注意力编码器（如BERT）的计算复杂度随序列长度呈平方增长，使得长上下文建模成本高昂。线性时间状态空间模型（如Mamba）虽然高效，但在建模全局交互方面存在局限性，并且可能受到padding引入的状态污染。我们提出了MaBERT，一种混合编码器，它将Transformer层与Mamba层交错，Transformer层用于全局依赖建模，Mamba层用于线性时间状态更新。这种设计将全局上下文集成与快速状态累积交替进行，从而实现长输入的高效训练和推理。为了稳定可变长度批处理，我们引入了padding安全掩码，它阻止状态通过padding位置传播，以及感知掩码的注意力池化，它仅聚合来自有效token的信息。在GLUE上，MaBERT在八个任务中的五个上取得了最佳平均分数，在CoLA和句子对推理任务上表现出色。当将上下文从512个token扩展到4096个token时，MaBERT相对于编码器基线的平均值，分别将训练时间和推理延迟降低了2.36倍和2.43倍，展示了一种实用的长上下文高效编码器。

🔬 方法详解

问题定义：Transformer模型，如BERT，在处理长序列时，自注意力机制的计算复杂度随序列长度呈平方级增长，导致计算资源消耗巨大，难以有效处理长文本。此外，使用padding来处理不同长度的序列时，padding位置的信息可能会影响模型的状态更新，造成状态污染，降低模型性能。

核心思路：MaBERT的核心思路是将Transformer层和Mamba层交错使用。Transformer层擅长捕捉全局依赖关系，而Mamba层具有线性时间复杂度，能够高效地处理长序列。通过交替使用这两种层，MaBERT旨在兼顾全局建模能力和计算效率，从而实现对长文本的高效处理。

技术框架：MaBERT的整体架构是一个由Transformer层和Mamba层交错组成的混合编码器。输入序列首先经过嵌入层，然后依次通过Transformer层和Mamba层。Transformer层使用自注意力机制来捕捉全局上下文信息，Mamba层则通过状态空间模型进行线性时间的状态更新。为了处理变长序列，MaBERT还引入了padding安全掩码和感知掩码的注意力池化。padding安全掩码阻止状态通过padding位置传播，而感知掩码的注意力池化仅聚合来自有效token的信息。

关键创新：MaBERT的关键创新在于混合编码器的设计，它将Transformer的全局建模能力与Mamba的线性时间复杂度相结合。此外，padding安全掩码和感知掩码的注意力池化也是重要的创新，它们解决了padding引入的状态污染问题，提高了模型在变长序列上的性能。

关键设计：MaBERT的关键设计包括：1) Transformer层和Mamba层的交错方式，具体层数比例未知；2) padding安全掩码的实现方式，通过mask阻止padding位置的状态传播；3) 感知掩码的注意力池化，只对有效token进行信息聚合；4) 损失函数和优化器的选择，论文中未明确说明，未知。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

MaBERT在GLUE基准测试中，在八个任务中的五个上取得了最佳平均分数，尤其在CoLA和句子对推理任务上表现出色。在长上下文任务中，将上下文长度从512扩展到4096时，MaBERT相对于编码器基线的平均值，分别将训练时间和推理延迟降低了2.36倍和2.43倍，证明了其在长文本处理方面的效率。

🎯 应用场景

MaBERT具有广泛的应用前景，包括长文本分类、摘要生成、机器翻译、问答系统等。其高效的长文本处理能力使其能够应用于处理法律文档、医学记录、金融报告等需要理解长上下文信息的领域。MaBERT的出现有望推动自然语言处理技术在长文本领域的应用和发展。

📄 摘要（原文）

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.