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

TL;DR

MaBERT is a hybrid transformer model that combines global dependency modeling with efficient linear-time state updates, enabling effective long-context processing with reduced training and inference costs.

Contribution

It introduces a novel interleaved hybrid encoder architecture with paddingsafe masking and mask aware attention pooling for efficient long-context language modeling.

Findings

Achieves top scores on five GLUE tasks, especially CoLA and sentence inference.

Reduces training time and inference latency by over 2.3 times for extended contexts.

Demonstrates effective long-context modeling with 4096 tokens.

Abstract

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.