You can deploy this model and start earning money today!

soyuj/
llama3.2-1b-sae-splade-l12-16k

Name: llama3.2-1b-sae-splade-l12-16k
Availability: InStock

information-retrieval

msmarco

meta-llama/Llama-3.2-1B

finetune

llama3.2

Views

No views yet

SAE-SPLADE concept retriever — Llama-3.2-1B, layer-12 SAE (16k latents)

A SPLADE-style learned sparse retriever whose vocabulary is SAE latents instead of wordpieces. Backbone: Llama-3.2-1B made bidirectional + MNTP-adapted (soyuj/llama3.2-1b-bidirectional-mntp-msmarco), fine-tuned with LoRA for retrieval; concept space: a TopK SAE trained on the backbone's layer-12 activations. This is the post-fix clean retrain (v2) of our layer-12 concept channel.

Results (MS MARCO dev-small, 6,980 queries, full 8.8M-passage corpus, Anserini impact index)

model	MRR@10	R@1000
this model (layer-12 SAE, 16,384 latents)	0.3933	0.9854
sibling: layer-11 SAE, 65,536 latents (repo)	0.3941	0.9869

Representation

For a query or document, token hidden states are taken from the output of decoder block 12, which is hidden_states[13] with output_hidden_states=True (index 0 is the embedding layer — off-by-one here silently ruins everything). Then:

rep = alpha * log1p( maxpool_over_tokens( TopK_64( ReLU( W_enc · normalize(h) ) ) ) )

normalize = the SAE's activation normalizer (center + RMS; sae/normalizer.pt — required)
W_enc = the SAE encoder (16,384 × 2048), fine-tuned during retrieval training — use the weights in checkpoint/concept_head.pt, not the raw sae/sae.pt
TopK per token (k=64), then amax-pool over real token positions only — padding and special tokens (including the BOS that the Llama tokenizer prepends) are masked out before pooling; the repo code handles this, manual re-implementations must too. Then log1p saturation, times the trained scalar alpha (= 0.9872 in this checkpoint)
score(q, d) = sparse dot product — serve from a standard inverted index over latent ids

Files

checkpoint/
  concept_head.pt      # trained SAE-encoder weight + alpha + metadata
                       #   {concept_topk: 64, concept_dim: 16384, input_dim: 2048,
                       #    hidden_state_index: 13}
  llm_adapter/         # the trained LoRA adapter (the MNTP adapter fine-tuned in place)
  trainer_state.json
sae/                   # the underlying SAE checkpoint (needed to CONSTRUCT the model)
  sae.pt               # frozen-init SAE weights (encoder gets overwritten by concept_head.pt)
  normalizer.pt        # activation normalizer — REQUIRED at inference
  trainer_state.json   # SAE pretraining state (39,962 steps / ~1.31B tokens)
  config.json          # SAE training config (layer_index: 12, expansion 8, topk 64)
llama_concept_splade_config.json   # full retrieval-training config

How to use (with the `sae_splade` research repo)

python

from sae.llama_concept_splade import LlamaConceptSplade

model = LlamaConceptSplade.from_pretrained_components(
    sae_checkpoint_dir="<this repo>/sae",
    adapter_path="<this repo>/checkpoint/llm_adapter",   # ONE adapter: the trained one
    base_model_name_or_path="meta-llama/Llama-3.2-1B",
    concept_topk=64,
    torch_dtype="bfloat16",
    train_lora=False,
).to("cuda")
model.load_concept_head("<this repo>/checkpoint")        # trained encoder + alpha
model.eval()

q = model.encode_queries(["what is hypertension"])       # [1, 16384] sparse-ish
d = model.encode_documents(["Hypertension is high blood pressure ..."])
score = (q * d).sum()

⚠️ The three mistakes that will silently break it

Do not stack adapters. Pass checkpoint/llm_adapter as THE adapter at construction. Wrapping a second PeftModel around an already-PEFT-wrapped backbone double-prefixes every key and PEFT drops all trained weights with only a UserWarning — the model runs, exports, and evaluates — all on a backbone missing its retrieval-trained LoRA (frozen at the plain-MNTP initialization). (This exact bug invalidated a week of our fused- index results.)
hidden_states[13], not [12]. The SAE hooks the output of block 12; hidden_states[0] is the embedding layer, so block-i output is hidden_states[i+1]. concept_head.pt records hidden_state_index: 13 and the repo's loader enforces it; manual re-implementations must respect it.
A missing sae/normalizer.pt does not error. The loader silently falls back to identity normalization and produces plausible-looking but wrong vectors. After construction, verify model.concept_encoder.normalizer is an ActivationNormalizer, not nn.Identity.

Standalone: extract latents (no research-repo code)

extract_latents.py in this repo extracts per-token SAE latents directly — for interpretability / analysis, independent of the SPLADE pooling (needs only torch transformers peft + the files here, incl. bidirectional_llama.py):

bash

# raw SAE latents (the SAE as pretrained):
python extract_latents.py --sae-dir sae --hs-index 13
# latents from the RETRIEVAL-TUNED encoder (this model's checkpoint):
python extract_latents.py --sae-dir sae --hs-index 13 --concept-head checkpoint/concept_head.pt

Core recipe: bidirectional Llama-3.2-1B + MNTP adapter (merged) -> hidden_states[13] -> normalize (h - mean) / rms (from sae/normalizer.pt) -> ReLU(h @ W_enc.T) -> optional per-token TopK-64. Position 0 (BOS) latents are out-of-distribution — skip them. Sanity check we ran: the tuned encoder shares ~7 of the top-8 latent ids per token with the raw SAE, with shifted activations — a fine-tune of the same dictionary, as expected.

Training

KL distillation (weight 1.0) + MarginMSE (0.05) against cross-encoder (ms-marco-MiniLM-L-6-v2) scores over SPLADE EnsembleDistil hard negatives (nway 8), plus query/doc FLOPS regularization (0.06 / 0.04, 6k-step warmup). 40k steps, batch 6 × grad-accum 3, lr 2e-5 (LoRA) / 1e-5 (SAE encoder), bf16. Both the LoRA and the SAE encoder train; the scalar alpha is trainable. Indexing/eval used quantization ×100 into an Anserini -impact -pretokenized index over latent tokens.

soyuj/llama3.2-1b-sae-splade-l12-16k