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

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:

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).

2. Official BEIR Benchmark (BEIR SciFact)

Standard Out-of-Domain Zero-Shot Information Retrieval Benchmark.

How to Use

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)}