From Tokens to Concepts: Leveraging SAE for SPLADE

1. Problem Statement

Learned Sparse Retrieval (LSR) models, such as SPLADE, have become state-of-the-art for first-stage information retrieval (IR) by offering an excellent trade-off between effectiveness and efficiency. They achieve this by learning to expand queries and documents into sparse vectors based on the pre-trained language model's (PLM) vocabulary.

However, SPLADE and similar models suffer from inherent limitations due to their reliance on the fixed token vocabulary of the underlying PLM (e.g., BERT, DistilBERT):

• Vocabulary Constraints: They struggle with polysemy (one token having multiple meanings) and synonymy (different tokens having the same meaning) because they are bound to specific token IDs.
• Multilingual/Multi-modal Challenges: Adapting to new languages or modalities is difficult because it requires retraining or fine-tuning the entire vocabulary projection layer.
• Suboptimal Representations: The final Transformer layer used for the vocabulary projection is not always the most suitable for IR tasks, and the model often relies on term expansion to handle vocabulary mismatches, which is not fully satisfactory.

The authors propose replacing the token-based vocabulary with a latent space of semantic concepts learned via Sparse Auto-Encoders (SAEs).

2. Methodology: SAE-SPLADE

The authors introduce SAE-SPLADE, a novel architecture that decouples the retrieval vocabulary from the PLM's token set. The approach involves a two-stage training pipeline:

A. Architecture Design

1. SAE Pre-training:
   • A Sparse Auto-Encoder is trained to reconstruct contextualized token embeddings from a frozen PLM (DistilBERT).
   • The SAE encoder maps dense hidden states $h$ into a sparse latent vector $z(h) \in \mathbb{R}^M$, where $M$ (latent dimension) is much larger than the hidden dimension (overcomplete basis).
   • The authors utilize TopK SAE, which enforces sparsity by selecting only the top-$k$ active latents, avoiding the need for explicit $\ell_1$ regularization which can cause "shrinkage bias."
2. SAE-SPLADE Fine-tuning:
   • The SAE decoder is removed.
   • The trained SAE encoder replaces the standard MLM (Masked Language Modeling) projection head of the SPLADE model.
   • The model is fine-tuned on IR tasks (MS MARCO) using a combination of KL-divergence and MarginMSE distillation objectives, similar to SPLADEv3.
   • Sparsity Control: In addition to the TopK constraint during SAE encoding, the authors apply FLOPs regularization ($L_{flops}$) during IR fine-tuning to control the sparsity of the aggregated document/query representation.

B. Training Strategy

• Stage 1: Train the SAE on MS MARCO passages to learn a semantic vocabulary (latent concepts) from the PLM's hidden states.
• Stage 2: Fine-tune the combined SAE-SPLADE model. The PLM weights are updated alongside the SAE encoder weights to align the latent space with relevance modeling.
• Key Hyperparameters: The authors experiment with the number of active latents ($k$) and the size of the latent vocabulary ($M$). They found that $k=8$ and $M=2^{16}$ offer the best trade-off.

3. Key Contributions

1. Novel Architecture: The first work to successfully integrate SAEs into a Learned Sparse Retrieval framework, replacing the token vocabulary with a learned semantic concept space.
2. New Evaluation Metric ($E^2$): Introduction of an Efficiency-Effectiveness ($E^2$) score. This metric balances retrieval performance (MRR) against computational cost (Query-Document FLOPs), allowing for a more nuanced comparison of sparse models than accuracy alone.
3. Comprehensive Ablation Studies:
   • Identified that using the last hidden layer (before the transform head) of the PLM yields the best results.
   • Demonstrated that TopK SAE is sufficient; more complex variants (Hierarchical, Matryoshka) did not yield significant gains for IR.
   • Showed that token-level TopK constraints during training and inference significantly boost efficiency without sacrificing effectiveness.
4. Multilingual Adaptability: Demonstrated that SAE-based vocabularies generalize better across languages than token-based vocabularies, as concepts are shared across languages rather than being tied to specific tokens.

4. Experimental Results

The model was evaluated on MS MARCO (In-Domain), TREC-DL, LoTTE, and BEIR (Out-of-Domain), as well as multilingual datasets (mMARCO, MIRACL).

• Effectiveness vs. Efficiency:
  • SAE-SPLADE achieves comparable effectiveness to the state-of-the-art SPLADEv3 model on In-Domain and Out-of-Domain tasks.
  • It offers significantly improved efficiency. For example, with $k=8$, SAE-SPLADE achieves a QD-FLOPs of 0.67 compared to 1.47 for standard SPLADE, while maintaining similar retrieval scores.
  • The $k=16$ variant matches SPLADEv3's effectiveness almost exactly but with a slightly larger index size.
• Multilingual Performance:
  • On multilingual datasets (mMARCO, MIRACL), SAE-SPLADE outperforms token-based SPLADE and other sparse baselines (like M3-Sparse) in both effectiveness and efficiency.
  • Sparsity Analysis: SAE-SPLADE exhibits higher overlap of activated latents across different languages (11.5% overlap vs. 8.6% for SPLADE), indicating better cross-lingual semantic alignment.
• Qualitative Analysis:
  • The latent space successfully disentangles semantics. A single latent can cover synonyms (e.g., "fry", "grill", "roast" mapped to a "Cooking methods" concept), and polysemous tokens (e.g., "spring" or "bank") are disambiguated by mapping to different latents depending on context.

5. Significance and Conclusion

• Decoupling Vocabulary: The paper proves that the rigid token vocabulary of PLMs is not strictly necessary for high-performance sparse retrieval. Learning a semantic concept space via SAEs is a viable and often superior alternative.
• Efficiency Gains: By replacing the MLM head with an SAE encoder and utilizing TopK constraints, the model reduces the computational cost of retrieval (fewer posting list entries to access) while maintaining high accuracy.
• Future Directions: The authors suggest that while SAE-SPLADE is promising, future work could explore:
  • Integrating cross-lingual alignment techniques to close the gap with dense multilingual models.
  • Developing SAE variants specifically tailored to IR objectives rather than general LLM interpretability.
  • Combining SAE concepts with lexical tokens to unify semantic expansion and exact matching.

In summary, SAE-SPLADE represents a significant step forward in Learned Sparse Retrieval, offering a more flexible, efficient, and semantically robust alternative to traditional token-based models.