SAE Encoder Embeddings: End-to-End Sparse Autoencoder Bottleneck for Retrieval

Status: Research & Architecture Design Phase Goal: Build the first encoder-only embedding model where the representation layer IS a Sparse Autoencoder, trained end-to-end with contrastive loss.

🎯 What This Is

• ModernBERT backbone (SOTA encoder-only with LLM innovations)
• TopK Sparse Autoencoder as the embedding bottleneck layer
• End-to-end contrastive training (not post-hoc SAE on frozen embeddings)

1. Interpretable — each active dimension corresponds to a learned semantic concept
2. Steerable — suppress/amplify specific features to control retrieval
3. Sparse-indexable — native sparse vector search (inverted index, not ANN)
4. Competitive — trained with modern contrastive objectives + hard negatives

🔬 Why This Is Novel

📂 Repository Structure

├── README.md                    # This file
├── ARCHITECTURE.md              # Detailed architecture design
├── PAPERS.md                    # Papers bibliography + key findings
├── TRAINING_RECIPE.md           # Full training recipe with hyperparameters
├── src/                         # (future) Implementation code
│   ├── model.py                 # SAE bottleneck + ModernBERT
│   ├── loss.py                  # Combined loss functions
│   └── train.py                 # Training script
└── experiments/                 # (future) Training logs and results

🏗️ Architecture Overview

Input text
    │
    ▼
┌─────────────────────────────────┐
│  ModernBERT-base (768-dim)      │  ← Backbone (trainable)
│  - RoPE positional embeddings   │
│  - FlashAttention 2             │
│  - GeGLU activations            │
│  - Alternating local/global attn│
│  - 8192 token context           │
└──────────────┬──────────────────┘
               │ mean-pool → v ∈ ℝ^768
               ▼
┌─────────────────────────────────┐
│  TopK Sparse Autoencoder        │  ← SAE Bottleneck (trainable)
│                                 │
│  Encoder: z = TopK(W_enc(v-b) + b_enc)
│           z ∈ ℝ^16384, ||z||_0 = k (32-128 active)
│                                 │
│  Decoder: v̂ = W_dec·z + b       │  ← For reconstruction loss only
│           (not used at inference)│
└──────────────┬──────────────────┘
               │
               ▼
        z (sparse embedding)
        Used for retrieval via sparse dot product

📊 Key Design Decisions

Why TopK (not L1)?

• Exact control of sparsity (k active features guaranteed)
• No shrinkage bias — L1 pushes all activations toward 0
• Better Pareto frontier at scale (OpenAI, arxiv:2406.04093)
• Dead latent prevention via AuxK loss

Why End-to-End (not frozen backbone)?

• Backbone learns to produce optimally decomposable representations
• CSR/SPLARE show frozen backbone limits retrieval performance
• Joint training enables the SAE to develop features that are useful for retrieval, not just reconstructive

Why ModernBERT?

• SOTA encoder-only architecture (surpasses BERT/RoBERTa/DeBERTa)
• LLM innovations: RoPE, FlashAttn, GeGLU, 8k context
• 768-dim base / 1024-dim large — good SAE input dimensions
• Hardware-aware design (efficient on T4/A10/A100)

🔗 Key References

⚡ Quick Links

• Backbone model: answerdotai/ModernBERT-base
• Training data: sentence-transformers/msmarco-bm25 + sentence-transformers/all-nli
• Evaluation: MTEB benchmark
• SAE reference impl: OpenAI sparse_autoencoder
• SPLARE (closest prior): arxiv:2603.13277
• CSR code: github.com/Mhz1y/CSR

📈 Expected Outcomes

1. Retrieval quality: Competitive with dense ModernBERT embeddings on MTEB retrieval tasks
2. Interpretability: Each active SAE feature maps to a human-interpretable concept
3. Steerability: Users can boost/suppress features to control search results
4. Efficiency: Sparse dot product with inverted index — potentially faster than dense ANN
5. Novel contribution: First end-to-end jointly-trained SAE embedding encoder

Generated by ML Intern

• Try ML Intern: https://smolagents-ml-intern.hf.space
• Source code: https://github.com/huggingface/ml-intern

Usage