💬 NLP

Linear-Time and Constant-Memory Text Embeddings Based on Recurrent Language Models

This paper proposes a vertically chunked inference strategy for recurrent language models like Mamba2, RWKV, and xLSTM that enables linear-time, constant-memory text embedding generation, offering a memory-efficient and competitive alternative to quadratic-complexity transformer-based approaches.

Original authors: Tobias Grantner, Emanuel Sallinger, Martin Flechl

Published 2026-04-21

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Tobias Grantner, Emanuel Sallinger, Martin Flechl

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a massive library of books, articles, and documents. You want to build a super-smart librarian who can instantly find the perfect book for any question you ask, even if the question is about a 100-page novel.

In the world of Artificial Intelligence, this "smart librarian" is called a Text Embedding Model. It turns words into a list of numbers (a "fingerprint") so the computer can understand the meaning and find similar things.

However, the current best librarians (based on Transformer models, like the ones powering many famous AI chatbots) have a big problem: They get exhausted.

The Problem: The "Quadratic" Library

Think of a Transformer model like a librarian who reads a document by comparing every single word to every other word in that document to understand the context.

If the document has 10 words, they make 100 comparisons.
If the document has 1,000 words, they make 1,000,000 comparisons.
If the document has 10,000 words, they need to make 100,000,000 comparisons!

This is called quadratic complexity. As the document gets longer, the librarian's brain (computer memory) fills up instantly, and they get incredibly slow. They simply can't handle long books without running out of space.

The Solution: The "Recurrent" Librarian

The authors of this paper propose a different kind of librarian: one based on Recurrent Models (specifically a new type called Mamba2, along with RWKV and xLSTM).

Instead of comparing every word to every other word, this librarian reads the document one word at a time, like a human reading a book. They keep a small, constant-sized "notebook" (memory) where they summarize what they've read so far.

The Magic: No matter if the book is 10 pages or 10,000 pages, the librarian only needs to keep that one small notebook open. Their memory usage stays constant. They don't get overwhelmed by long texts.

The Innovation: The "Vertical Chunk" Strategy

You might ask: "If they read one word at a time, isn't that slow? Can't they read faster?"

Yes, reading one word at a time is slow on modern super-fast computers that love to do many things at once (parallel processing). The authors' big breakthrough is a strategy they call Vertically Chunked Inference.

Imagine the document is a long train.

Old Way (Horizontal): The librarian tries to process the entire train at once. If the train is too long, the station (memory) collapses.
New Way (Vertical Chunks): The authors cut the train into small, manageable carriages (chunks).
- They process one carriage through the entire depth of the librarian's brain (all layers) before moving to the next carriage.
- Crucially, they pass a "hand-off note" (the state) from one carriage to the next, so the meaning isn't lost.

The Analogy:
Think of it like a factory assembly line.

Transformers try to build the whole car at once on a giant table. If the car is huge, the table breaks.
The New Method builds the car in sections. They build the engine, then the chassis, then the wheels, passing the partially built car down the line. They only need a small workbench (constant memory) because they aren't holding the whole car at once. But because they use powerful machines, they still build it incredibly fast.

What Did They Find?

The researchers took these "Recurrent Librarians" (Mamba2, RWKV, xLSTM) and taught them to be embedding models using the same training methods as the best Transformers.

Performance: They are just as good as the Transformers at understanding text. They can find the right answers in the library just as accurately.
Efficiency: When the documents get long (like 32,000 words), the Transformers crash or slow to a crawl because they run out of memory. The Recurrent models? They keep chugging along at the same speed, using the same amount of memory.
The Sweet Spot: They found that you don't need to make the "carriages" (chunks) huge to get the speed benefits. Even small chunks allow the computer to work fast while keeping memory usage low.

Why Does This Matter?

This is a game-changer for anyone working with long documents (legal contracts, medical records, entire books) or for people with limited computer resources (like running AI on a laptop instead of a massive server farm).

It proves that we don't need to choose between "smart" and "efficient." We can have a librarian that is both brilliant at understanding long stories and efficient enough to fit in your pocket.

In short: The paper introduces a new way to process text that is as smart as the current AI giants but as memory-efficient as a simple notebook, making it perfect for the era of long-context AI.

1. Problem Statement

Transformer-based text embedding models (e.g., BERT, E5-Mistral) are the current state-of-the-art but suffer from significant scalability limitations:

Quadratic Computational Complexity: The self-attention mechanism scales as $O(T^2)$ with respect to sequence length $T$ , making inference slow for long documents.
Linear Memory Complexity: Storing intermediate activations for attention mechanisms requires memory that scales linearly with sequence length ( $O(T)$ ), often causing out-of-memory (OOM) errors on long sequences or large batch sizes.
Resource Constraints: These limitations restrict the utility of transformers in resource-constrained environments and for applications requiring long-context retrieval (e.g., full-document analysis).

While recurrent architectures like Mamba2, RWKV, and xLSTM offer linear time complexity and constant memory usage for text generation, their potential as general-purpose embedding models has remained largely unexplored.

2. Methodology

The authors propose a framework to adapt recurrent architectures for high-quality text embeddings while maintaining their efficiency advantages.

A. Model Training

Architecture: The authors fine-tune pretrained recurrent language models (specifically Mamba2, RWKV, and xLSTM) using the E5-Mistral-7b training recipe.
Procedure:
- They remove the language modeling head and append an End-of-Sequence (EOS) token to the input.
- The final hidden state at the EOS position serves as the text embedding.
- Loss Function: They utilize the InfoNCE loss (contrastive learning) to distinguish positive query-document pairs from negative examples.
- Instruction Tuning: To ensure versatility across downstream tasks, they prepend task-specific instructions to queries (but not documents) during training, allowing the model to adapt to various retrieval and classification tasks without retraining.

B. Inference Strategy: Vertically Chunked Inference

The core technical contribution is a novel inference strategy designed to exploit the "duality" of modern recurrent models (which can be formulated as both linear recurrences and parallelizable matrix operations).

The Duality: Models like Mamba2 can switch between a sequential recurrence (memory efficient, $O(1)$ state) and a parallel matrix formulation (fast, $O(T^2)$ if fully unrolled).
Horizontal Inference (Standard): Processes the entire sequence through one layer at a time. While it allows intra-layer parallelization, it requires storing full-sequence activations for the current layer, leading to $O(T)$ memory usage.
Vertically Chunked Inference (Proposed):
- Concept: Instead of processing the whole sequence layer-by-layer, the input sequence is divided into fixed-size vertical chunks of size $V$ .
- Process: The model processes one vertical chunk through all $L$ layers before moving to the next chunk.
- State Management: Between vertical chunks, the model maintains a recurrent state for each layer (analogous to time-step recurrence).
- Memory Efficiency: This decouples memory usage from the total sequence length $T$ . For $T > V$ , memory usage becomes constant ( $O(L + Q \cdot V)$ ), depending only on the model depth ( $L$ ) and chunk size ( $V$ ), not the input length.
- Parallelism: Within each vertical chunk, the model utilizes the parallelizable matrix formulation to accelerate computation.

3. Key Contributions

Demonstration of Viability: The paper establishes that recurrent models (specifically Mamba2) can serve as competitive, general-purpose text embedders, achieving performance comparable to similarly sized transformer baselines.
Vertically Chunked Inference Algorithm: A novel scheduling strategy that enables constant memory usage for embedding generation in recurrent architectures. This allows processing of arbitrarily long sequences without increasing memory footprint beyond a fixed threshold.
Cross-Architecture Validation: The inference strategy is empirically validated across three distinct recurrent architectures: Mamba2, RWKV, and xLSTM, confirming that the memory-runtime trade-off is consistent across different SSM (Structured State Space Model) and RNN variants.
Open Source Release: The authors released fine-tuned Mamba2-based model checkpoints and the inference implementation to support reproducibility.

4. Experimental Results

The models were evaluated on MTEB (English and Multilingual), MTEB v2, and the LongEmbed benchmark (long-context retrieval).

Embedding Quality:
- Mamba2 vs. Transformers: Fine-tuned Mamba2 models (1.3B and 7B) achieved competitive results against transformer baselines (Mistral 7B, Qwen2).
- Performance Gaps: On English tasks, Mamba2 models lagged slightly (1–2 percentage points) behind Mistral 7B, likely due to differences in pretraining data scale and composition. However, on Multilingual tasks, Mamba2 models often outperformed or matched the transformers.
- Long Context: On the LongEmbed benchmark (up to 32k tokens), recurrent models achieved performance parity with transformers, proving their efficacy for long documents.
Inference Efficiency (Runtime & Memory):
- Memory: Recurrent models with vertically chunked inference maintained constant memory usage for sequences longer than the vertical chunk size (e.g., >4096 tokens). In contrast, transformer models (even with FlashAttention) showed linear memory growth, quickly exceeding GPU memory limits.
- Runtime: For single-sequence inference on long documents, recurrent models were significantly faster than transformers. For batched inference, parallelization saturated at small vertical chunk sizes (e.g., $V=256$ to $512$), yielding runtimes comparable to horizontal inference but with drastically lower memory costs.
- Hardware: The efficiency gains were consistent across NVIDIA L40S, A100, and H100 GPUs.

5. Significance and Impact

Scalability: This work provides a practical solution for generating embeddings for long documents (e.g., legal contracts, scientific papers, books) that are currently infeasible with standard transformer embeddings due to memory constraints.
Resource Efficiency: By enabling constant-memory inference, these models make high-quality NLP accessible in resource-constrained environments (e.g., edge devices, consumer-grade GPUs).
Future Direction: The paper suggests that as recurrent base models are pretrained on larger, more diverse datasets (closing the quality gap with transformers), they will become the preferred architecture for scalable, long-context embedding generation.

In summary, the paper successfully bridges the gap between the efficiency of recurrent models and the high-quality requirements of modern text embedding tasks, offering a scalable alternative to the transformer paradigm for long-sequence processing.