SAE-SPLADE: Sparse Semantic Retrieval

• The paper introduces SAE-SPLADE, a framework that replaces the fixed token vocabulary with a learned latent semantic space using sparse auto-encoders to overcome synonymy and polysemy issues.
• It employs a two-stage pipeline with SAE pre-training and SPLADE fine-tuning, using TopK sparsification to produce high-dimensional, compact sparse representations for effective retrieval.
• Empirical results on benchmarks like MS MARCO and BEIR show that SAE-SPLADE achieves competitive retrieval effectiveness while significantly reducing computational FLOPs and enhancing multilingual generalization.

SAE-SPLADE is a sparse retrieval framework that substitutes the standard fixed token vocabulary of SPLADE with a learned latent space of semantic concepts using Sparse Auto-Encoders (SAE). It is designed to address limitations of token-based projections, such as synonymy, polysemy, and cross-lingual or multi-modal challenges, commonly encountered in learned sparse IR models. SAE-SPLADE achieves comparable retrieval effectiveness to classical SPLADE on both in-domain and out-of-domain tasks, while significantly improving efficiency and transferring naturally to multilingual and potentially multi-modal settings (Zong et al., 23 Apr 2026).

1. Architectural Overview

The core architectural innovation in SAE-SPLADE is the replacement of the backbone vocabulary with a trainable set of latent concepts. This is realized through a two-stage pipeline:

• SAE Pre-training:
  • A frozen pretrained LLM (PLM; e.g., DistilBERT) is used to generate final-layer hidden vectors $h_i \in \mathbb{R}^d$ for each token.
  • A sparse auto-encoder (SAE), composed of an encoder $E \in \mathbb{R}^{M \times d}$ and a decoder $D \in \mathbb{R}^{d \times M}$, projects $h_i$ into a high-dimensional latent concept space ($M \gg d$, typical $M=2^{16}\approx 65536$), yielding

  $$z_i = S(E h_i + b_{\text{enc}}) \in \mathbb{R}^M,\quad \widehat{h}_i = D z_i + b_{\text{dec}}$$ - The sparsification function $S(\cdot)$ enforces that only $k$ largest elements of $z_i$ are nonzero (“TopK SAE”).

• SAE-SPLADE Fine-tuning:

  • The SAE decoder is discarded post-training.
  • The encoder and sparsification step, now serving as the SPLADE "projection head,” are used to produce per-token sparse codes.
  • Per-document ($E \in \mathbb{R}^{M \times d}$0) and per-query ($E \in \mathbb{R}^{M \times d}$1) embeddings are created by max-pooling across token-level latents; retrieval is based on sparse dot products.

Token-level representations are thus mapped not to a token vocabulary but to a semantic latent vocabulary, allowing more flexible and compact cross-linguistic and semantic grouping.

2. Key Mathematical Formulations

The mathematical machinery of SAE-SPLADE includes three main components:

2.1 SAE Pre-training Objective

Given a collection $E \in \mathbb{R}^{M \times d}$2 of PLM-generated hidden states:

• Encoder output: $E \in \mathbb{R}^{M \times d}$3
• Decoder reconstruction: $E \in \mathbb{R}^{M \times d}$4
• SAE loss:

$E \in \mathbb{R}^{M \times d}$5

• For TopK SAE, the $E \in \mathbb{R}^{M \times d}$6 norm is replaced by the hard constraint $E \in \mathbb{R}^{M \times d}$7.

2.2 SPLADE-Style Retrieval Scoring

Each query or document receives a pooled sparse vector $E \in \mathbb{R}^{M \times d}$8:

$E \in \mathbb{R}^{M \times d}$9

Retrieval score:

$D \in \mathbb{R}^{d \times M}$0

This is efficiently computed with an inverted index due to the high sparsity of $D \in \mathbb{R}^{d \times M}$1.

2.3 IR Fine-tuning Loss

Fine-tuning is performed using a combination of distillation and sparsity control losses:

$D \in \mathbb{R}^{d \times M}$2

Here, $D \in \mathbb{R}^{d \times M}$3 and $D \in \mathbb{R}^{d \times M}$4 are teacher-student distillation losses; $D \in \mathbb{R}^{d \times M}$5 and $D \in \mathbb{R}^{d \times M}$6 regularize representation sparsity.

3. Training Procedure and Optimization

SAE is pretrained on 8.8 million MS MARCO v1 passages using a frozen PLM. Optimization uses AdamW with a learning rate of $D \in \mathbb{R}^{d \times M}$7 for 160K steps, batch size $D \in \mathbb{R}^{d \times M}$860K tokens, $D \in \mathbb{R}^{d \times M}$9 TopK sparsity, and $h_i$0.

During retrieval fine-tuning on MS MARCO (using cross-encoder distillation, batch size 32 queries × 8 hard negatives), document and query lengths are truncated to 256 and 32 tokens, respectively. Hyperparameters are: $h_i$1 learning rate, 240K steps, $h_i$2, $h_i$3, $h_i$4, and $h_i$5. Aggregation by max-pooling maintains sparsity at both per-token and per-document levels.

Sparsity is enforced by construction via TopK selection per token ($h_i$6) and via regularization terms controlling global FLOPs across queries and documents.

4. Empirical Results and Ablation Analysis

Effectiveness and efficiency evaluations are performed on both in-domain benchmarks (MS MARCO dev, TREC-DL’19/’20) and out-of-domain (LoTTE, 13 BEIR datasets):

Model	MSM MRR@10	TREC-DL nDCG@10	LoTTE S@5	BEIR nDCG@10	QD-FLOPs	Avg doc nnz
BM25	18.3	49.2	51.0	43.7	0.13	39
SPLADEv3 (DistilBERT)	38.7	74.8	70.3	50.0	1.40	165
SPLADE (DistilBERT)	37.7	72.0	68.9	48.8	1.47	118
SAE-SPLADE (k=8)	37.6	72.1	68.8	49.2	0.67	109

SAE-SPLADE matches SPLADE in retrieval effectiveness but reduces query-document FLOPs by approximately 55% and decreases index size moderately (average doc nnz: 109 vs. 118). The $h_i$7 trade-off metric confirms a superior balance of efficiency and effectiveness.

Ablation studies demonstrate:

• Pre-training the SAE is essential—random projections do not suffice without strong regularization.
• Using the last-layer PLM hidden state (rather than a transformed output) yields optimal performance.
• The chosen latent size ($h_i$8) and sparsity level ($h_i$9) are empirically optimal.

5. Multilingual and Cross-Modal Generalization

SAE-SPLADE generalizes naturally to multilingual settings by learning language-agnostic latent concepts. On mMARCO (Arabic, Spanish, French, Japanese, Russian, Chinese, English) and MIRACL benchmarks, SAE-SPLADE attains slight improvements over SPLADE at reduced FLOPs:

• mMARCO dev (MRR@10): SAE-SPLADE $M \gg d$0 (FLOPs $M \gg d$1); SPLADE $M \gg d$2 (FLOPs $M \gg d$3)
• MIRACL (nDCG@10): SAE-SPLADE $M \gg d$4 (FLOPs $M \gg d$5); SPLADE $M \gg d$6 (FLOPs $M \gg d$7)

A plausible implication is increased cross-lingual alignment, as the model’s concept space captures semantic regularities irrespective of orthography or morphology.

The architecture suggests extensibility to multi-modal retrieval by jointly training an SAE over representations from both textual and visual backbones, unifying them into the sparse latent concept space.

6. Qualitative Characteristics and Implications

Quantitative results are complemented by qualitative analysis of the learned latent space. Examination of co-occurrence over 10,000 MS MARCO passages demonstrates that:

• A single SAE latent often groups synonyms (e.g., latent 6470 triggering for "dish," "soup," "ingredients," "sauce").
• Polysemous tokens are disambiguated across different latents (e.g., "spring" triggering latent 5162 for "Florida springs, box springs" and latent 48743 for "season").

This suggests that semantic concepts, rather than surface tokens, become the effective indexing units, reducing problems related to synonymy and polysemy.

These properties imply that the model is less sensitive to literal token overlap and is well positioned for tasks demanding robust semantic and multilingual retrieval. The work opens directions for optimizing SAE variants for IR objectives, refinement of cross-lingual alignment, and fully multi-modal sparse retrieval architectures (Zong et al., 23 Apr 2026).