README.md

---
language: en
license: apache-2.0
tags:
  - splade
  - sparse-retrieval
  - information-retrieval
  - beir
  - code-search
  - static-query
  - distilled
pipeline_tag: feature-extraction
library_name: transformers
---

# 3-Layer Static-Query SPLADE-v3-doc (Pruned & Retrained)

<!-- Evaluated on In-Domain BEIR and Official BEIR SciFact -->

This repository contains the 3-layer pruned and retrained static-query SPLADE-v3-doc model (45.75M parameters). It is designed for inference-free client-side semantic search across code documentation and technical corpora.

---

## Model Overview

- **Architecture**: 3-Layer Pruned Transformer Encoders (retaining layers `[0, 6, 11]` from `naver/splade-v3-doc`).
- **Parameter Count**: **45.75 Million** (**58.4% parameter reduction** compared to 110M SPLADE-v3-doc).
- **Query Mechanism**: **Zero-Inference Static Query Table** (`static_query_weights.pt`). At query time, terms are looked up directly from a learned static token weight table requiring 0 GPU executions and 0 neural forward passes.
- **Document Encoder**: Sparse `log1p(ReLU(logits)).max()` document-side pooling generating weighted vocabulary posting lists for static inverted indexes.
- **Selection Criterion**: Saved from `best_proxy` checkpoint maximizing top-1 ranking performance while maintaining sparsity budget targets.

---

## Training Setup & Hyperparameters

The model was retrained on Modal A10G GPUs via teacher-student distillation from a 12-layer fine-tuned teacher (`Akshat131/splade-multi-static-doc`) using the following command:

```bash
python -m modal run --detach pruningscript.py::train_pruned_static_student \
    --run-name combined_beir_v1 \
    --scored-file-name train_multi_static_teacher_scores.jsonl \
    --student-checkpoint "pruned_inits/pruned_3layer_init" \
    --checkpoint-name pruned_multi_static_3layer_final \
    --epochs 4 \
    --learning-rate 3.5e-5 \
    --static-query-learning-rate 1e-3 \
    --distill-temperature 10.0 \
    --margin-mse-weight 0.8 \
    --target-doc-active-dims 140.0 \
    --target-query-active-dims 80.0 \
    --lambda-doc-flops 1.5e-5 \
    --lambda-query-l1 1e-8 \
    --sparsity-zero-fraction 0.10
```

### Key Training Parameters

1. **Teacher Distillation**: Softmax KL Distillation (`--distill-temperature 10.0`) and MarginMSE loss (`--margin-mse-weight 0.8`) aligning score margins between positive and hard-negative document pairs.
2. **Static Query Learning Rate**: Boosted static query term learning rate (`--static-query-learning-rate 1e-3`) allowing fast adaptation of non-linear IDF weights for technical documentation syntax.
3. **Ramped Sparsity Schedule**: 10% zero-sparsity warmup phase (`--sparsity-zero-fraction 0.10`) followed by quadratic FLOPS regularization (`1.5e-5`), preventing dimensional collapse while enforcing target active dimensions (`T_doc = 140.0`).

---

## Evaluation Benchmark Results

The model was evaluated across multiple benchmarks using Modal A10G GPUs against previous 3-layer checkpoints and baselines:

### 1. In-Domain Technical Code Documentation Benchmark (`combined_beir` / `my_eval_dataset`)
*NumPy, Pandas, PyBind11 documentation corpora (14,352 queries across 10,386 documents).*

| Metric | NEW Trained 3-Layer Model | Previous Model (`Akshat131/static-splade-trained-pruned`) | Previously Used Model (`Arvind0101/static-query-splade-code-docs`) |
| :--- | :---: | :---: | :---: |
| **NDCG@10** | **0.5019** (combined) / **0.4119** (my_eval) | 0.3887 | 0.3838 |
| **MRR@10** | **0.6653** (combined) / **0.5423** (my_eval) | 0.5204 | 0.5161 |
| **Recall@10** | **0.4786** | 0.4563 | 0.4473 |
| **Recall@100** | **0.7091** | 0.6945 | 0.6866 |

### 2. Official BEIR Benchmark (`BEIR SciFact`)
*Standard Out-of-Domain Zero-Shot Information Retrieval Benchmark.*

| Metric | NEW Trained 3-Layer Model | Previous Model (`Akshat131/static-splade-trained-pruned`) | Previously Used Model (`Arvind0101/static-query-splade-code-docs`) |
| :--- | :---: | :---: | :---: |
| **NDCG@10** | **0.5815** | 0.5518 | 0.5690 |
| **MRR@10** | **0.5576** | 0.5134 | 0.5273 |
| **NDCG@5** | **0.5656** | 0.5325 | 0.5394 |
| **NDCG@100** | **0.6218** | 0.5892 | 0.5997 |

---

## How to Use

```python
import torch
from transformers import AutoTokenizer, AutoModelForMaskedLM
from huggingface_hub import hf_hub_download

repo_id = "Akshat131/splade-multi-static-pruned-v2"

# 1. Load Tokenizer & Document Encoder
tokenizer = AutoTokenizer.from_pretrained(repo_id)
doc_encoder = AutoModelForMaskedLM.from_pretrained(repo_id)

# 2. Load Static Query Weights for Inference-Free Query Scoring
weights_path = hf_hub_download(repo_id=repo_id, filename="static_query_weights.pt")
static_weights = torch.load(weights_path, map_location="cpu")["static_query_weights"]

# Query scoring requires 0 forward passes:
query_str = "how to drop NaN missing values in array"
query_tokens = tokenizer(query_str)["input_ids"]
query_vector = {token_id: static_weights[token_id].item() for token_id in set(query_tokens)}
```