- The paper presents a recurrent architecture using SSMs that achieves linear-time computation and constant memory via a novel vertically chunked inference strategy.
- It adapts contrastive training with dual operation modes to optimize long-context embedding tasks while ensuring hardware-aligned runtime and resource efficiency.
- Empirical results indicate competitive performance with transformers, particularly in long-document retrieval and deployments under constrained memory conditions.
Linear-Time and Constant-Memory Text Embeddings Based on Recurrent LLMs
Introduction
The paper "Linear-Time and Constant-Memory Text Embeddings Based on Recurrent LLMs" (2604.18199) systematically analyzes the role of modern recurrent architectures, specifically structured state space models (SSMs) such as Mamba2, RWKV, and xLSTM, for efficient and scalable text embedding generation. Contrary to transformer-based models, which impose quadratic computational and linear memory complexity with respect to input length, recurrent approaches retain linear computational scaling and—given the introduction of vertically chunked inference—constant activation memory with respect to sequence length for sufficiently large inputs. This facilitates their application to long-document embedding and deployment in resource-constrained environments.
While existing embedding models predominantly leverage transformers due to their expressive self-attention mechanisms, these models face fundamental hardware- and cost-based scalability barriers as input lengths increase. Transformer inference optimizations, such as FlashAttention and sliding-window attention, only partially alleviate memory and computation constraints by, respectively, reducing activation footprints or trading global context for locality. Recurrent models, historically under-explored for general-purpose embedding, fundamentally differ in their strict sequence-based, constant-state update property, making them better suited to long-context encoding.
This work harnesses the dual formulation of modern recurrent SSMs (Mamba2), which enables two equivalent modes:
- Sequential recurrence (O(T) time, O(1) memory with respect to sequence length),
- Matrix-based parallelization (O(T2) time and memory, exploiting hardware acceleration).
By partitioning the input sequence into fixed-size horizontal and vertical chunks, the authors present a hybrid "vertically chunked" inference stratagem, yielding O(Tâ‹…Q) computational overhead with constant activation memory beyond the vertical chunk threshold V (for all T>V).
Training and Embedding Protocol
The embedding approach adapts the Model-e5 contrastively-supervised training protocol—removing the generative LM head from a pretrained SSM (e.g., Mamba2) and treating the final-layer hidden state at the EOS position as the sequence embedding. Contrastive InfoNCE loss optimizes the encoder to maximize positive query-document similarity against a large set of hard negatives. Notably, instruction templates are prepended only to queries for retrieval-style tasks, avoiding repeated traversal of documents for each new instruction.
Instruction tuning is critical: it ensures models generalize across diverse downstream semantic tasks without retraining, thus efficiently supporting retrieval, classification, and semantic similarity estimation.
Chunked Inference Strategies
The key technical contribution is the vertically chunked cross-layer inference (VCI) method:
- Horizontal chunking processes full sequences layer-by-layer, maximizing parallelism but demanding memory that scales linearly with input length.
- Vertically chunked inference reroutes this computation to process all layers on fixed-size vertical sequence segments, storing only those chunk-wise activations and one recurrent state per layer. For T>V, memory costs are constant with respect to T.
Empirical analysis reveals that parallelization saturates at small V (e.g., $4096$ tokens), so vertical chunking allows parallel runtime convergence with a fixed memory requirement. This mechanism is generalized to any SSM architecture with compatible recurrence/parallelism duality (shown for Mamba2, RWKV, xLSTM).
Empirical Results
Fine-tuning Mamba2 and transformer-based architectures (Mistral, Qwen) with identical data/protocols, the models are evaluated on three key benchmarks:
- English Massive Text Embedding Benchmark (MTEB): 41 tasks (retrieval, classification, clustering, etc.)
- Multilingual MTEB: extended task/language coverage
- LongEmbed: retrieval in sequences up to 32K tokens
Main results:
- On multilingual tasks, recurrent models (Codestral, Mamba2) are at parity or outperform transformer baselines.
- On English tasks, recurrent models lag behind by only 1--2 percentage points; the gap is largest on retrieval sub-tasks.
- On long-context retrieval, all architectures (transformer and recurrent) perform similarly, with differences O(1)0 percentage points for like-sized models.
- Runtime/memory: For input length O(1)1, Codestral and other recurrent models surpass transformers (including those with efficient attention) both in throughput and memory efficiency. Memory usage saturates and remains flat for sequences longer than O(1)2, contrasting with the continual linear growth for transformers. Acceleration is maintained for smaller vertical chunk sizes on modern hardware given batch processing.
- All findings are robust across hardware variants (L40S, A100, H100) and generalize across SSM types.
The authors hypothesize that pretraining data scale, not architectural deficiency, accounts for the remaining quality gap on English tasks—the Transformer baseline models are pre-trained on O(1)3 larger corpora.
Implications and Future Directions
This work decisively demonstrates that SSM-based recurrent models can serve as competitive, high-quality, and resource-scalable text embedders suitable for production retrieval or representation tasks, especially where document lengths substantially exceed classic attention-window limits or where deployment must operate under constrained memory. Recurrent architectures impose a new trade-off envelope between expressiveness and cost: for long-context applications, vertically chunked inference unlocks constant memory scaling without runtime compromise. The approach is architecture-agnostic within the family of dualizable SSM-style RNNs.
This unlocks practical deployment scenarios for long-context semantic search, document understanding, and retrieval augmentation previously infeasible for transformer variants, and lays groundwork for new research on cross-architecture contrastive pretraining with SSMs. As larger, more robustly pretrained SSM base models become available, the remaining quality gap on English tasks should narrow, making recurrent encoders universally competitive.
Limitations
- Embedding quality assessment focused predominantly on Mamba2; extension to other architectures and more exhaustive domain/task coverage is warranted.
- Pretraining scale was not strictly controlled—further fair benchmarking requires recurrent models pretrained on corpora and durations equivalent to strong transformer baselines.
- Input length distributions and specific task granularity in benchmarking were not varied or dissected in detail.
The duality formulation and VCI scheduling concept, however, are architecturally generic for RNN-SSMs.
Conclusion
"Linear-Time and Constant-Memory Text Embeddings Based on Recurrent LLMs" (2604.18199) establishes the empirical and theoretical viability of high-performance, resource-scalable, general-purpose text embedding with recurrent SSM architectures. The vertically chunked inference strategy delivers constant-memory scaling and hardware-aligned runtime, overcoming key deployment and scaling bottlenecks of transformer-based models. This advances the state of efficient representation learning, with significant implications for long-context and resource-sensitive NLP.