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

LLM Long Context NLP VDIT ML Mamba/SSM RNN
基于Transformer的嵌入模型存在计算复杂度为平方级、内存消耗为线性级的问题,因而难以有效处理长序列。为此,我们提出以循环架构作为高效替代方案,并引入一种“垂直分块推理策略”(vertically chunked inference strategy):该策略可实现快速嵌入生成,且当输入长度超过垂直分块大小后,内存占用即稳定为常数。通过对Mamba2模型进行微调,我们验证了其作为通用文本嵌入器的可行性——在多项基准测试中均取得具有竞争力的性能表现,同时相较基于Transformer的同类模型,显著降低了内存开销。我们通过实验进一步证实,该推理策略可成功应用于Mamba2、RWKV及xLSTM等多种循环架构模型,各类模型均展现出一致且可预期的运行时—内存权衡特性;由此确立了循环模型作为高效嵌入生成任务中Transformer有力替代方案的地位。
Transformer-based embedding models suffer from quadratic computational and linear memory complexity, limiting their utility for long sequences. We propose recurrent architectures as an efficient alternative, introducing a vertically chunked inference strategy that enables fast embedding generation with memory usage that becomes constant in the input length once it exceeds the vertical chunk size. By fine-tuning Mamba2 models, we demonstrate their viability as general-purpose text embedders, achieving competitive performance across a range of benchmarks while maintaining a substantially smaller memory footprint compared to transformer-based counterparts. We empirically validate the applicability of our inference strategy to Mamba2, RWKV, and xLSTM models, confirming consistent runtime-memory trade-offs across architectures and establishing recurrent models as a compelling alternative to transformers for efficient embedding generation.
许愿