Learning Retrieval Models with Sparse Autoencoders

Abstract

Sparse autoencoders (SAEs) provide a powerful mechanism for decomposing the dense representations produced by Large Language Models (LLMs) into interpretable latent features. We posit that SAEs constitute a natural foundation for Learned Sparse Retrieval (LSR), whose objective is to encode queries and documents into high-dimensional sparse representations optimized for efficient retrieval. In contrast to existing LSR approaches that project input sequences into the vocabulary space, SAE-based representations offer the potential to produce more semantically structured, expressive, and language-agnostic features. Building on this insight, we introduce SPLARE, a method to train SAE-based LSR models. Our experiments, relying on recently released open-source SAEs, demonstrate that this technique consistently outperforms vocabulary-based LSR in multilingual and out-of-domain settings. SPLARE-7B, a multilingual retrieval model capable of producing generalizable sparse latent embeddings for a wide range of languages and domains, achieves top results on MMTEB's multilingual and English retrieval tasks. We also developed a 2B-parameter variant with a significantly lighter footprint.

Figures (6)

• Figure 1: Overview of SPLARE. A pre-trained SAE can be inserted at any layer $l$ of the LLM to get sparse latent representations of input tokens. These token-level representations are then aggregated into a single sparse vector using a pooling mechanism analogous to SPLADE. During training, we only fine-tune the LLM parameters (via LoRA adapters) while keeping the SAE frozen. We show the top-5 activated features for the English query—for our final SPLARE model.
• Figure 2: (Left) Performance across layers on Llama Scope (Llama-3.1-8B) and Gemma Scope (Gemma-2-2B). (Right) Performance with increasing SAE width on Gemma-2. Evaluation done with $\text{Top-K}=(40,400)$.
• Figure 3: (Left) Impact of pruning documents with Top-K (with $k=40$ for queries). (Right) MS MARCO index distribution for SPLARE and SPLADE (8.8$M$ documents).
• Figure 5: Retrieval Latency (ms) when pooling documents (Left) or query (Right) representations with Top-K. In low-latency settings, SPLARE enables higher accuracy for a given level of latency.
• Figure :