Course project · Team ANLP Doomsday · mid submission

When does counting help you choose?

Our project, Predicting When Collection Statistics Help Top-k Selection: From Retrieval to Attention and Expert Routing, asks whether an old idea from search engines can improve three quite different parts of modern NLP systems, and whether we can tell in advance when it will. This page walks through the mid report from the ground up: the idea, what we built, what the experiments showed, why our first predictor failed, and what we learnt from working out the reason.

mid · 2 Oct 2026 182 primary + 1,226 sensitivity rows 14 representations · 13 collections BM25 · Quest · OLMoE

Aviral Gupta · Mohit Kumar Singh · Anirudh Sankar · Arihant Tripathy — IIIT Hyderabad · code ↗

§0The short version

A search engine, the attention layer of a language model and the router in a mixture-of-experts model look like unrelated pieces of machinery, but each of them does the same two things. It scores an input against a set of features using a dot product, and then it keeps only the k highest-scoring items. Search engines learnt long ago that the raw dot product can be improved with a correction called BM25. The other two mechanisms have never used it. We wanted to know when that correction actually helps, and whether it is possible to predict the answer without any relevance labels.

BM25 is really three corrections bundled together. IDF gives rare features more weight than common ones, saturation stops a feature from counting ever more as it repeats, and length normalisation keeps long documents from winning simply because they are long. We turned each of these into an independent on/off switch, which gives eight scorers ranging from the plain dot product (everything off) to full BM25 (everything on). We then built fourteen sparse "concept" vocabularies from a single frozen encoder, multilingual E5, and scored every one of them on thirteen English and Indic retrieval collections under all eight scorers. The same scorer is also wired into KV-cache page selection on Llama-3.2-1B and expert routing on OLMoE. Those two adapters are built and tested, but they have not yet been run at scale.

The retrieval half is complete. The full correction improves nDCG@10 in 90% of the primary rows, but mostly because the plain dot product over these codes is close to useless. The predictor we committed to in advance, built on the shape of each vocabulary's frequency curve, did no better than always guessing "the correction helps". Working out why turned into the main finding. An exact Shapley decomposition shows that 62% of the gain comes from saturation, which reads no collection statistics at all, while IDF supplies only about a tenth. How much saturation helps is set by how unequal the stored feature values are, and that account passed six tests we fixed before running them. Whether IDF and length normalisation help depends on relevance, which no statistic of the index can see. Agreement with a reference retriever stands in for relevance well enough to predict it in exploratory analysis, although a decision rule built on it gains only 0.005 nDCG@10.

90%
primary rows where full BM25 beats the plain dot product
62 · 28 · 10
percent of the gain from saturation, length and IDF (Shapley)
ρ = 0.87, 0.90
pre-registered test: value dispersion predicts saturation at unseen k1
R² 0.52–0.66
held-out fit of the exploratory pseudo-label predictor of the collection gain

§1Background you need

§1.1 · Retrieval as a dot product

A retriever is given a query and a large pile of documents, called the collection or corpus, and has to return the documents most likely to answer the query, in order. Most modern retrievers do this by turning the query into a vector xq and every document into a vector xd. Each document is scored with the dot product ⟨xq, xd⟩, and the top k are returned.

These vectors come in two broad kinds. A dense retriever such as DPR or E5 produces vectors of a few hundred numbers, all of them active and any of them possibly negative. A sparse retriever produces much longer vectors, with one slot per feature, where almost every slot is zero and the few non-zero ones are positive. Traditional keyword search is the classic sparse retriever, with one slot for each word in the vocabulary. Learned sparse retrieval keeps that sparse format but lets a model decide what the features should be, as SPLADE and the recent SAE-based retrievers do. This project works within that family.

§1.2 · What BM25 adds to the dot product

BM25 is the classical scoring function for keyword search, and after several decades it remains a strong baseline. Its three parts arrived one at a time: IDF as a statistical account of how specific a term is, saturation as an approximation to a model of how often a term repeats within a document, and pivoted length normalisation so that long documents are neither favoured nor penalised. Each part was designed for words, which is exactly why it is an open question whether any of them helps features that are not words.

ComponentThe problem it fixesWhat it doesReads
IDF
inverse document frequency
A word like "the" appears in every document, so matching it tells you nothing. Multiplies each feature by a weight that is large when few documents contain it and near zero when all of them do.the collection
Saturation A document that says "cat" 50 times is not 50 times more about cats than one that says it once. Passes each feature's value through a curve that rises quickly and then flattens out, capped at k1+1.only the document
Length normalisation Long documents match more query terms simply because they contain more text. Divides by the document's length relative to the average.the collection

The last column turns out to matter a great deal. IDF and length normalisation need statistics gathered over the whole collection, while saturation looks only at the document in front of it. The project set out to study collection statistics, and the largest part of the gain came from the one component that does not use them.

a worked example

Suppose the collection has a million documents. A feature that appears in only 10 of them receives an IDF weight of about 11.5, one that appears in half of them receives ln 2 ≈ 0.69, and one that appears in every document receives about 0.0000005, which effectively removes it from the score.

For saturation, take the standard setting k1 = 1.2 and a document of average length. A raw value of 1 stays at 1.00, a value of 5 becomes 1.77, and a value of 50 becomes only 2.15. However often a feature repeats, its contribution can never exceed 2.2.

For BM25 to make sense, a representation has to provide three things. Its values must be non-negative, since saturating a negative number has no meaning. Each feature must have a meaningful count of the documents it appears in, which is what IDF is computed from. And each document must have a length. A dense embedding provides none of these, because every coordinate is active in every document and any of them can be negative, which is why dense retrievers have always stuck to the plain dot product. A sparse representation, by contrast, provides all three.

§1.3 · Concept vocabularies from a dense encoder

There are several ways to turn a dense embedding into sparse features. You can cluster the embeddings and use cluster membership as the features, compress them into discrete codes, or train a sparse autoencoder (SAE), a small network that rewrites each vector as a handful of active features chosen from several thousand. Several recent papers do exactly this and then treat the resulting features as a vocabulary for retrieval. We use seven such constructions, which §5 describes.

§2The observation: three things that are one thing

The whole project rests on one observation. Three mechanisms, developed by three different research communities, perform the same computation: they give each feature a non-negative score and then make a hard top-k selection.

SettingThe "query"The "documents"What is kept
A · Learned sparse retrievalquery vector xq document vectors xdtop-k documents
B · KV-page selection (Quest)current attention query qt summaries of pages of 16 cached keystop pages to attend to
C · MoE routingtoken hidden state ht expert embeddings wetop-k experts
Retrieval query · document KV selection query · KV page Routing token · expert score per feature, φj ≥ 0 BM25-style correction IDF · saturation · length each switched on or off top-k at a fixed budget documents / KV pages / experts selection counts dfj over corpus or calibration set label-free features ↓ predicted gain
The computation the three settings share, with the BM25 correction inserted between scoring and top-k. The same per-feature counts that IDF needs also supply the diagnostics used to predict whether the correction will help.

None of the three asks how often a feature is selected across the collection, how evidence builds up within an item, or how large the item is. Lexical retrieval asks all three questions. Inside transformers only fragments of the answers have reappeared. Loss-free balancing, a technique for training mixture-of-experts models, adds a bias to each expert that shrinks when the expert is used too often, which amounts to a frequency discount with no saturation or length term. StreamingLLM always keeps the first few tokens in the KV cache because they absorb a large share of attention, a phenomenon known as the attention sink. Each is a local fix for one mechanism, and neither is framed as a general correction.

§3The puzzle, and the gap

Does the BM25 correction still help when the vocabulary is learned rather than made of words? The published evidence is contradictory: sometimes it helps a great deal, and sometimes it does real harm.

Vocabulary (Clavié et al., App. C)BM25Inner productWinner
SAE features of a frozen retriever0.5820.381BM25, by 20 points
SPLADE-v30.4550.568inner product, by 11 points

Both rows use the same benchmark suite and the same two scorers. The only thing that changed is the vocabulary, and that alone was enough to reverse which scorer comes out ahead. Other work points the same way. Park et al. (CL-SR) lose 11 to 18 percent of their nDCG@10 on TREC-DL when they remove IDF, which is the only published ablation of the correction. Han et al., working with image patches, find the two scorers within about two points, and when they raise the number of active features from 16 to 128, R@1 falls from 0.508 to 0.038.

The explanation offered is that BM25 helps as long as the vocabulary stays Zipfian, meaning a few features are very common and there is a long tail of rare ones, as with words in ordinary text. If that account is right, statistics of the frequency curve should predict the sign of the gain before anyone runs the experiment, and our proposal made that prediction the central test. The account also leaves three things open. None of the papers varies the vocabulary systematically or asks which of the three components produces the gain. The explanation has never been turned into a quantity that can be computed in advance. And it has never been tested inside a model, where frequency and usefulness may well diverge: a KV page that is selected constantly could be an attention sink the model depends on, and a rarely used expert could be one of the "super experts" whose removal breaks the model.

what this project contributes

A controlled grid of fourteen representations scored under all eight configurations, in 182 primary and 1,226 sensitivity rows. A negative result for the pre-registered predictor, together with the reason it could not have worked. An exact attribution of the gain to its three components, with a mechanism for saturation that passed tests fixed in advance. An exploratory label-free predictor of when collection statistics help. And adapters for KV-page selection and expert routing, ready for the frozen predictors.

§4The problem, stated precisely

§4.1 · The switchable scorer

For query and document vectors with non-negative coordinates, the score is S(q, d) = Σj xqj · wj · T(xdj, ℓd). Each of its three pieces is a switch. The weight wj is either 1 or the smoothed IDF,

idfj=ln(1+N−dfj+0.5dfj+0.5)

where dfj of the N documents contain feature j. The "1 +" inside the logarithm keeps the weight strictly positive; the textbook form of IDF turns negative for features in more than half the documents and would flip the sign of a score. The length factor ℓd is either 1 or 1 − b + b·Ld/avgL, where Ld = Σj xdj is the document's length. The transform T(f, ℓ) is either f/ℓ or the saturated form, so that written out in full the document-side weight is

gjσ(d)=idfjσI⋅{(k1+1)φj(d)φj(d)+k1ℓ(d)σS=1φj(d)/ℓ(d)σS=0

with k1 = 1.2 and b = 0.75 in the primary grid, the standard BM25 defaults. A configuration is a bit string in the order IDF, saturation, length.

Config (I S L)What is onName
000nothingplain inner product
100IDFIDF-weighted dot product
010saturationsaturation only, no collection statistics
001length
110, 101, 011two of three011 is the most frequent winner
111all threeBM25 over learned features
why eight configurations rather than two

Comparing only the inner product with BM25 would tell us whether the bundle helps, but not which part of it is responsible. Running all eight corners of the 2 × 2 × 2 cube lets us split the gain exactly among the three components (§4.3), and that is how we found that saturation does most of the work.

§4.2 · What counts as one row

Collection frequency belongs to a corpus, so the same representation has a different frequency curve on MS MARCO than on Telugu MIRACL. One row of the analysis therefore pairs a representation with a collection, at a given support budget (the caps on query and document features), BM25 parameters (k1, b) and basis seed. Each row carries eight nDCG@10 values Qc, one per configuration. This framing came from the TA's feedback on the interim proposal, which noted that counting only representations would have left 16 data points to fit 6 regressors.

§4.3 · Three things to predict

The pre-registered target is the margin, the gain of full BM25 over the plain inner product divided by the standard deviation of quality across the eight configurations:

m(R,c)=Q111−Q000sdσQσ

Dividing makes it unitless, so a margin of 2 means the same thing whether the metric is nDCG@10, RULER accuracy or MMLU accuracy, and a predictor fitted on retrieval can be scored elsewhere. §8.2 shows that this target turned out to be badly shaped for prediction, which led us to two targets defined on the same eight values.

The first is the Shapley value φs of each switch, its marginal effect averaged over every order in which the switches could be turned on. Writing QA for the configuration with the switches in A on,

φs=∑A∌s|A|!(2−|A|)!6(QA∪{s}−QA)

With three switches the weights are simple: a third for turning s on first, a sixth for each of the two ways of turning it on second, and a third for turning it on last. Because our grid contains all eight configurations, the values are exact, with no sampling, and the three of them add up to exactly Q111 − Q000.

The second is the collection gain G = Q111 − Q010, the value of IDF and length normalisation once saturation is already on. Its sign answers the practical question a practitioner faces: should I use collection statistics for this representation on this collection? We score predictors by sign accuracy against the majority sign and by held-out R² and Spearman correlation. For G we also measure the nDCG@10 reached when collection statistics are used only where the prediction is positive.

§4.4 · Transfer to attention and routing

The scorer applies to any selection that has a non-negative score and a selection count. For KV page i at decoding step t, the feature value is the softplus of the Quest criticality score. Its frequency is the fraction of calibration queries for which the page enters the top-k among the pages that exist at that point, and its length is the norm of the page's concatenated key minima and maxima. For expert e at a token, the value is the router probability, the frequency is the share of calibration tokens routed to e, and the length is the norm of the router logits. In attention a row is a model, RULER task, context length and token budget, scored by RULER accuracy. In routing a row is a model and task, scored by negative log-likelihood on WikiText and zero-shot MMLU accuracy.

A predictor fitted in retrieval is applied there without refitting. The registered criterion asks for a sign accuracy of at least 70% and a mean absolute error no more than 1.5 times the corpus-held-out error in retrieval. Since always predicting a gain already reaches 90% sign accuracy in retrieval, we added a third condition before any transfer run: the predictor must also beat the majority sign.

§5What we built: Setting A, retrieval

§5.1 · One frozen encoder

Every representation starts from multilingual E5-base at a pinned revision, run frozen in fp16. A pooled embedding is the mean over the content tokens, and the individual token states are kept for the token-level variant. Holding the encoder fixed is what makes the comparison controlled: any difference between two representations comes from how the vocabulary was built, not from a different model underneath.

§5.2 · Seven bases in four families

A basis is a recipe for turning a dense vector into sparse, non-negative features. All seven were fitted on the same 200,000 MS MARCO training passages and nowhere else.

FamilyBasisHow it works
LinearIdentityThe raw 768 E5 dimensions.
Gaussian random projectionMultiplies by a random matrix into 768 dimensions.
PCAKeeps the 384 directions of largest variance.
Clusteredsoft k-means4,096 cluster centres, with each vector given soft weights over its nearest centres.
QuantisedProduct quantisation (PQ)Splits the vector into 8 pieces and replaces each piece with the nearest of 256 codes, giving one active code per piece.
Residual quantisation (RQ)Uses 8 stages of 256 codes, where each stage encodes what the earlier stages missed.
LearnedTop-K SAEA sparse autoencoder with 8,192 features, at most 32 of which are active for any input.
split rectification

The linear bases produce negative values, which BM25 cannot handle. We therefore split each value z into two features, ([z]+, [−z]+), holding its positive and negative parts separately. The 768 signed numbers become 1,536 non-negative ones.

Each basis is used at two granularities. At the pooled level, the basis is applied to the mean of the token states. At the token level, it is applied to each token state and the codes are summed. Seven bases at two granularities give fourteen representations. A secondary SAE shares one dictionary across both granularities, which lets us separate the effect of how codes are aggregated from the effect of what the dictionary was fitted on. SPLADE has its own backbone rather than E5, so it stays outside the grid as an external comparison.

§5.3 · Budgets and statistics

Before any statistic is computed, supports are capped at 32 query features and 128 document features in the primary budget, so no scorer can win simply by seeing more features. Document frequencies, lengths and IDF are then computed from the capped vectors of each target corpus, and retrieval is exact over the whole corpus.

a note on linear bases

For a linear map C with diagonal weights D, weighted scoring q⊤C⊤DC d gives the same ranking as the plain inner product whenever C⊤DC is a multiple of the identity. Rectification, truncation, saturation and length normalisation are all non-linear, so none of our representations meets that condition, and IDF can genuinely change their rankings.

§5.4 · Thirteen collections

GroupCollectionsWhat it tests
MS MARCOMS MARCO dev, TREC-DL 2019, TREC-DL 2020 (one shared set of 8.8M passages) in-domain performance, since the bases were fitted on MS MARCO passages
BEIRArguAna, FiQA, NFCorpus, SCIDOCS, SciFact, TREC-COVIDother domains: science, medicine, finance and argument retrieval
LIMITLIMITbuilt so that the sets of documents its queries need exceed what a low-dimensional embedding can represent
MIRACLHindi, Bengali and Telugu, each searched within its own language Indic languages, where the frequency curve can look very different

§5.5 · Experiment families

The primary family covers all thirteen collections. Each sensitivity family varies one setting at a time on the ten collections outside MS MARCO, which is what lets us check whether a finding at the default settings survives when the settings change.

FamilyRows
primary (14 representations × 13 collections)182
support budget 16/64, 32/128, 64/256, native462
k1 ∈ {0.6, 1.2, 2.4}, b ∈ {0, 0.5, 0.75}630
basis seeds 1 and 2100
shuffled and uniform df controls84
shared SAE dictionary20
union on the ten non-MS MARCO collections1,226

Every row has eight scoring configurations. Families share their default cells, so the union is smaller than the sum. The MS MARCO budget rows are still pending.

§5.6 · Reference systems and significance

Two reference systems sit outside the grid. Lexical BM25 over ordinary word tokens is the classic keyword baseline, and dense E5 with exact search shows what the encoder achieves on its own. For significance, a paired bootstrap over queries with 10,000 draws recomputes both the gain and its standard deviation, and Holm correction is applied within each pre-declared family of tests: full BM25 against the inner product, and the three component main effects. Significance is reported on the ten collections outside MS MARCO.

§6What the predictors can see

No feature describing a row ever reads a relevance label. Within that rule, there are three tiers of information a predictor can use, each adding to the one before it.

§6.1 · Tier 1: the pre-registered frequency diagnostics

These are the six numbers we committed to in the proposal, all describing the shape of the feature-frequency curve. The primary ridge model uses the three marked ★, chosen before any result was seen; a secondary model uses all six.

DiagnosticWhat it measures
Active fractionThe share of all features that are ever switched on.
Zipf slope ★The slope of the log-log frequency curve, describing how steep it is.
Gini coefficientHow unequally activation mass is spread across features.
Head frequency massThe share of selections taken by the 1% most frequent features.
Frequent activation share ★The share of activation mass held by that same top 1%.
IDF variance ★How much room IDF has to reweight anything at all.

§6.2 · Tier 2: index statistics

These describe the values actually stored in the index rather than how often features occur. The most important is cvx, the coefficient of variation of the stored document values: their standard deviation divided by their mean. It measures how unequal the stored values are, and §8.4 explains why that is exactly what governs saturation.

§6.3 · Tier 3: ranking features and pseudo-labels

The richest tier looks at the rankings themselves. For every configuration and query we keep the top-100 list, and measure how much each switch changes the top 10 and how the top-100 scores are shaped, in the spirit of query performance prediction. We then build a pseudo-quality Q̃c: the mean nDCG@10 of each configuration's ranking when the top 10 of a reference retriever on the same query is treated as if it were relevant. The references are text BM25, dense E5, and the two fused by reciprocal rank. Putting Q̃c in place of the true Qc in the Shapley formula and in G gives label-free estimates of every target with no fitted parameter at all. This borrows an old idea from IR evaluation, ranking systems without judgements by comparing them to pooled output.

§6.4 · Validation

All predictors are ridge regressions, linear regression with a penalty on large weights, with the scaling and penalty chosen inside the training folds so that the held-out data never influences them. We validate along three axes, each answering a different question about generalisation.

AxisHeld outQuestion
Familya whole basis family (linear, clustered, quantised, learned)Does it work on a new kind of basis?
Corpusa whole corpusDoes it work on a new collection?
Paramsfit on primary rows, score on rows at other k1 and bDoes it survive a change of BM25 settings?

§7Settings B and C: inside a transformer

Both adapters are built, and six tests confirm that they match the native code exactly when every switch is off, that calibration is causal, that the token budget is charged correctly (including for sink pages) and that the expert mixture is right. Neither has been run at scale yet; that is the work for October.

§7.1 · Setting B, KV-page selection (Llama-3.2-1B-Instruct)

When a model decodes over a long context, attending to every cached key at every step is expensive. Quest reduces the cost by grouping the keys into pages of 16 tokens and storing the elementwise minimum and maximum key of each page. For the current query it computes an upper bound on the attention a page could receive, and then attends only to the pages with the highest bounds until a token budget is filled. Prefill stays dense; only decoding is sparse.

Our corrected version treats pages as documents. The Quest score is passed through a softplus so that it is positive, and then through the BM25 scorer. A page's frequency is how often it entered the top-k over causal prefill positions, and its length is the norm of its key summary. The runs will cover all thirteen RULER tasks with 500 examples each, at 4K and 8K context, with budgets of 512 and 1,024 tokens, against full attention, plain Quest and H2O, which keeps recent tokens together with the "heavy hitters" that have accumulated the most attention.

which direction the correction should push

H2O uses accumulated attention to decide which items to keep. IDF does the opposite, discounting items that are selected often. Attention sinks are the most frequently selected keys of all, yet the model cannot do without them, so this is the setting where a frequency correction is most likely to do harm.

§7.2 · Setting C, expert routing (OLMoE-1B-7B)

OLMoE has 64 experts per layer, of which 8 are active for each token, and that is enough experts for their selection frequencies to form a meaningful curve. We apply the correction to the router probabilities to decide which experts run, and then combine the chosen experts using their original gate weights, so only the selection changes. The all-off configuration calls the native forward pass and reproduces the unmodified model exactly. Quality is measured by negative log-likelihood on WikiText-2 and zero-shot MMLU accuracy, and an inference-time analogue of loss-free balancing is included as a comparison.

§7.3 · What gets carried over

Two predictors will be carried into these settings. The pre-registered predictor goes as registered. The pseudo-label predictor needs a label-free reference native to each setting: for KV pages that reference is the exact attention mass, and for routing it is the dense softmax mixture over all experts. Both come from the uncorrected model. Both predictors and their thresholds will be frozen before any corrected outcome is seen, and every result will be reported against the majority sign as well as the registered thresholds.

§8Results for Setting A

The results are told in the order they came: the correction helped, the predictor failed, and the attempt to explain that failure produced the main finding.

§8.1 · The correction helps because the inner product fails

CollectionText BM25Dense E5mean · 000mean · 111best · 111
MS MARCO dev.222.425.013.130.318
TREC-DL 2019.456.704.036.255.604
TREC-DL 2020.484.697.035.251.556
ArguAna.411.444.021.130.371
FiQA.236.381.013.085.266
NFCorpus.307.326.064.115.285
SCIDOCS.152.172.009.044.120
SciFact.662.694.089.197.419
TREC-COVID.582.695.134.246.508
LIMIT.939.015.001.004.032
MIRACL hi.357.582.011.124.328
MIRACL bn.407.700.025.216.490
MIRACL te.331.751.019.259.579

nDCG@10 of the two reference systems, the mean over the fourteen sparse representations under the inner product (000) and full correction (111), and the best representation under 111. The best is the token-level SAE everywhere except LIMIT, where it is token k-means. Highlighted cells beat text BM25.

Under the plain inner product, the capped sparse codes are close to useless on every collection. Full correction lifts their mean by a factor of between 1.8 and 14, and the margin is positive in 90% of the primary rows where it is defined. After correction the best representation becomes competitive with the text baseline: the token-level SAE is the strongest on twelve of the thirteen collections, beats text BM25 on six of them, and comes within 0.10 nDCG@10 of dense E5 on TREC-DL 2019. Token granularity is stronger than pooled throughout, and the shared-dictionary rows trace most of the pooled deficit to fitting the dictionary on pooled vectors. LIMIT is the one collection where text BM25 wins outright, as its construction predicts.

the catch in the good news

Full correction is the best of the eight configurations in fewer than a quarter of the primary rows. The most frequent winner is saturation with length normalisation (011), which leaves IDF out. The inner product scores below 0.01 nDCG@10 in almost half of the primary rows, and almost any correction beats a scorer that broken. That was the first hint of trouble for the predictor.

§8.2 · The pre-registered predictor fails

The margin was meant to be predicted from the shape of the frequency curve, and it could not be. The primary ridge model explains almost none of the held-out variance, with R² of 0.04 across basis families and 0.02 across corpora, and its sign accuracy is exactly that of always predicting a gain. The variant with all six diagnostics does no better. By the criterion we registered, the frequency diagnostics carry no usable information about the margin.

part of the reason: the target saturates

Most margins fall between 1.5 and 3.5, because the margin mostly records that the uncorrected scorer is broken. To see why, suppose the inner product scores Δ below the other seven configurations and those seven are roughly equal. The mean then sits Δ/8 below the seven, the population variance is 7Δ²/64, and the standard deviation is Δ√7/8, so the margin is

m = Δ / (Δ√7/8) = 8/√7 ≈ 3.02

whatever the size of Δ. A gain of 0.01 and a gain of 0.5 receive the same margin, which leaves little variation for any diagnostic to explain. The rest of the failure lies in what the correction is made of.

§8.3 · Where the gain comes from

To see what a predictor would need to capture, we split each row's gain among the three switches with the Shapley formula of §4.3.

Share of the all-on minus all-off gain in the primary rows.

RowsIDFSat.Len.IDF + len.
primary, 13 collections.10.62.28.38 (.36–.41)
budget 16/64.16.60.24.40
budget 32/128.10.61.29.39 (.36–.42)
budget 64/256.08.57.35.43
native, 3 collections−.01.54.47.46

Share of the gain assigned to each switch. Intervals resample whole collections. The budget rows cover the ten collections outside MS MARCO; native support covers ArguAna, NFCorpus and SciFact.

Saturation carries most of the gain in every set of rows. IDF, the component the Zipfian account is about, carries the least at the primary budget. Reading down the budget rows, IDF's share shrinks as supports grow while length's share grows, and at native support IDF's share is slightly negative. Han et al. report the same trend for image features, and the Holm-corrected main effects agree, with saturation significant most often and IDF least often.

IDF is not empty, though. When the document frequencies are shuffled across features, IDF's mean Shapley value turns negative, and the real value is larger than the shuffled one in 83% of rows (one-sided Wilcoxon p < 10−6). Making all frequencies uniform, which removes IDF altogether, leaves all-on quality no lower. So the frequencies carry real information, but little of it reaches nDCG@10.

§8.4 · Why saturation can be predicted and the rest cannot

BasisGran.111cvxφsatIDF + len.G > 0
identitypooled.1470.37.043.681.00
token.1420.60.060.55.77
Gaussianpooled.0980.29.020.77.85
token.0880.55.029.64.92
PCApooled.2220.75.138.24.69
token.3111.35.155.45.46
k-meanspooled.0180.04.000.94.31
token.1851.71.125.11.85
PQpooled.0120.00.0001.00.62
token.1441.44.099.09.69
RQpooled.0560.00.0001.00.92
token.2922.10.176.18.69
SAEpooled.1270.18.014.66.77
token.3733.31.193.32.46

Per representation: mean all-on nDCG@10, the coefficient of variation of stored values, the mean saturation Shapley value, the share of the gain from IDF and length, and the fraction of collections where collection statistics help after saturation (G > 0). cvx is averaged over the ten collections whose indexes were measured.

Read the cvx and φsat columns together and a pattern appears: the more unequal the stored values, the more saturation contributes. This follows directly from the form of the transform. With length normalisation off, saturation maps a stored value f to g(f) = (k1+1)f / (f + k1). The ratio g(f)/f falls as f grows, so saturation shrinks large values relative to small ones. If every stored value were equal, g would rescale all scores alike and change no ranking; the more unequal the values, the more the ranking moves. A smaller k1 bends the curve harder and so moves it further.

Pooled PQ and RQ are the limiting case. They store one active code per piece, all with the same value, and every document has the same length, so their saturation and length effects are exactly zero in every one of their rows, which also serves as an end-to-end check of the pipeline. Outside these two codes saturation almost always helps, and how much it moves the ranking is a property of the index that can be measured without any labels.

IDF and length normalisation behave differently. They too move the ranking by an amount the index fixes, but whether the move helps depends on whether heavy documents were over-rewarded relative to their relevance, and on whether rare features separate relevant documents from the rest. Both questions are about relevance, which no statistic of the index records. That is why a predictor built on the frequency curve could not tell when collection statistics help. The decision matters most for the strong representations: collection statistics supply most of the gain for the weak pooled linear codes, but for the token-level SAE and PCA they help after saturation on fewer than half the collections.

§8.5 · Testing the saturation account in advance

An explanation found by looking at the data needs a test the data did not shape. Before running any test we committed six hypotheses to the repository, all using the k1 and b sweep, which the exploration had not touched. H2 to H6 are Holm-corrected together. H1 uses the default settings, which coincide with the primary rows, so it checks consistency only.

HypothesisResult
H1IDF + length share at the default settings lies in [.36, .42].393
H2φsat at k1 = 0.6 exceeds φsat at k1 = 2.4.093 > .080
H3within-corpus Spearman ρ(cvx, φsat) ≥ 0.3.89, .89
H4a cvx model fitted at k1 = 1.2 ranks φsat at ρ ≥ 0.5.87, .90
H5φlen = 0 exactly at b = 0max 0.0
H6φlen at b = 0.75 exceeds φlen at b = 0.5.045 > .029

H2 and H6 have Holm-corrected p = 3.6 × 10−5. H3 and H4 are tested at k1 = 0.6 and 2.4. H5 covers 210 rows.

All six pass, but they are not equally informative. H5 holds by construction, since the length factor is exactly one at b = 0, so it only checks the code. H2 and H6 test directions that follow from the form of the scorer. H3 and H4 carry the real evidence: a model of saturation fitted on cvx at one value of k1 keeps its ranking at values of k1 it never saw.

§8.6 · Supplying relevance without labels

If the missing ingredient is relevance, a stand-in for relevance should restore the prediction. Everything in this subsection came after the confirmatory tests were fixed, so it is exploratory. We scored every configuration against a reference retriever's top 10, treated as relevant, and compared the three tiers of features from §6.

TargetFeaturesFamilyCorpusParams
φsatpre-registered−.46.14.02
+ index.47.42.50
+ pseudo-labels.56.62.76
φlenpre-registered−.30.14.00
+ index−.16.25−.02
+ pseudo-labels.49.70.77
φidfpre-registered−.01.00.02
+ index.02.02.07
+ pseudo-labels.40.21.73
Gpre-registered−.07−.02.16
+ index−.05−.09.15
+ pseudo-labels (BM25).38.52.62
+ pseudo-labels (E5).51.66.80

Held-out R² of ridge regression fitted on the 139 primary rows of the ten collections outside MS MARCO, holding out basis families or corpora, or predicting the 560 rows at other k1 and b. Each tier adds features to the one above it.

The pattern is clean. The frequency diagnostics predict none of the targets. Index statistics predict saturation and nothing else, as §8.4 would lead you to expect. Only the pseudo-label tier reaches positive held-out R² on every target and every axis. The mechanism is simple: every target is a fixed linear combination of the eight Qc, so a pseudo-quality that orders the configurations the same way as the true quality yields correlated targets, and it does, with a median within-row Spearman correlation of 0.98.

The practical test is a decision. A rule with no fitted parameter uses collection statistics only where the pseudo-label estimate of G is positive. It raises mean nDCG@10 from 0.143 (always using them) to 0.148, against 0.151 for an oracle that knows the true answer, which recovers 61% of a headroom that is itself below 0.01. The gain holds at other budgets and parameters, though across seeds its interval reaches zero. The gains are small because always using collection statistics is already close to the oracle on most rows, and we chose the rule during exploration.

the references have biases of their own

The correction under study is BM25 applied to other features, so agreement with text BM25 may reward configurations simply for behaving like BM25, and lexical systems also contributed to the judgement pools of several collections. The dense reference errs the other way: the sparse codes come from the same encoder, so agreement with it partly measures how much of the embedding a configuration keeps. The prediction of G holds with either reference, but since neither is free of bias, the result stays exploratory until it is tested on data fixed in advance.

§9What happens next

The retrieval half leaves two predictors to carry into the transformer, as §7.3 describes. The first week of October also commits a second pre-registration: the pseudo-label decision rule will be tested on the MS MARCO budget rows, which nobody has seen yet, so that the exploratory result gets a test the data did not shape.

§9.1 · Schedule to 31 October

WeekWorkOwner
1 · 3–9 OctResume the MS MARCO budget sweep (84 rows) after cluster maintenance. Run SPLADE-v3 as the external comparison and the conceptual replication of the Appendix C reversal. Validate both adapters at scale, with native parity, calibration cost and memory.Anirudh
Commit the second pre-registration: the pseudo-label rule on the unseen MS MARCO budget rows, the frozen transfer predictors and the transfer thresholds.Mohit
2 · 10–16 OctSetting B on Llama-3.2-1B-Instruct over all 13 RULER tasks with 500 examples, 4K and 8K contexts, budgets of 512 and 1,024 tokens, with H2O.Anirudh
Second confirmatory round in retrieval.Mohit
3 · 17–23 OctSetting C on OLMoE-1B-7B: WikiText-2 negative log-likelihood and zero-shot MMLU, with the loss-free-balancing analogue.Anirudh
Write-up of the bilinear identity and its ranking conditions. Draft of the final report from the retrieval results.Arihant
4 · 24–28 OctTransfer scoring with uncertainty for both predictors, per setting, against the majority sign.Mohit
Figures and tables from the saved artifacts. Viva slides.Arihant, Aviral
5 · 29–31 OctFinal report.Arihant
Slides and presentation. Code archive.Aviral, Anirudh

If time runs short, experiments are dropped in a fixed order: GSM8K and expert preservation first, then the 8K context in Setting B, then MMLU in Setting C (keeping WikiText), and finally the Appendix C replication.

§9.2 · Who does what

MemberRole
Aviral GuptaProposed the idea, and makes the presentations.
Anirudh SankarWrites and runs the experiment code.
Mohit Kumar SinghAnalyses the results and owns the pre-registrations.
Arihant TripathyWrites the report and relates the results to the literature.

§10Limitations

One encoder underlies every representation, so the effect of the basis cannot be fully separated from the effect of E5 itself. E5 has also seen MIRACL in training, which means the Indic results say nothing about languages the encoder has never encountered. The sparse codes are far weaker than the encoder they come from, so the attribution of the gain may not hold for strong learned sparse retrievers such as SPLADE, which is one reason the SPLADE-v3 comparison is on the schedule.

The label-free predictor is not free to run: it needs queries and a run of a second retriever. Routing is tested only at inference, so nothing here speaks to training. And the transfer test inherits the weakness of the margin described in §8.2, which is why every transfer result will sit next to the majority-sign baseline.

the most informative outcome

If the correction turns out to help retrieval but hurt inside a model, that would be the most revealing result of all. It would show that frequency and importance genuinely come apart inside a transformer, for instance because the pages selected most often are attention sinks the model cannot do without.

§11Engineering for the hardware we had

This section comes from the project's execution notes rather than the report itself, but it explains how a grid this size fitted on the hardware. Our nodes have four GPUs with 11 GB each and support neither bf16 nor FlashAttention-2. The full grid takes days to run, so every stage writes a manifest recording its configuration hash, the checksums of its inputs and the code version, resumes from checkpoints, and verifies checksums before reusing anything it has already produced.

BottleneckProblemFixEffect
EncodingThe E5 forward pass over MS MARCO was repeated for every basis, seven times over. Run E5 once per batch and apply every basis to the same GPU activations, sharded across four workers. 182 encoding tasks became 57 shared units, and the full corpus encoded in under seven hours.
PostingsPostings were re-transformed for every query batch under each of the eight configurations. Cache the transformed postings on local disk, keyed by the parent artefact, the collection statistics and the switch setting. Loading takes 1.1 s against 7.4 s to rebuild, and all 512 ranking comparisons matched exactly.
ScoringCPU scoring became the last bottleneck. A GPU path that scores 512 queries per pass in fp32, adds products in feature order without fused multiply-add, and breaks ties by document ID. 512 queries went from 24.0 s to 1.8 s, matching the CPU path exactly on 960 comparisons.
why exact matches matter

With a top-k cut-off, even a tiny floating-point difference can swap two documents near the boundary and change nDCG. If the faster GPU path did not reproduce the CPU rankings exactly, we would have no way of telling a genuine effect of a switch from arithmetic noise. That is why the GPU path fixes both the order of additions and the tie-breaking rule.

§12How the design changed since the proposal

ItemProposal (28 Aug)Mid report, and why
BasesEight, SPLADE includedSeven fitted on E5. SPLADE-v3 is external, since it shares no encoder with the rest.
Representations1614, plus a shared SAE dictionary across both granularities as a secondary check.
SwitchesIDF, saturation, length, top-kThree switches. Top-k is a fixed matched budget, varied separately in the budget family.
Collections14, one code-mixed13, with no code-mixed set, since we use neither synthetic code-mixing nor new relevance annotation.
Rows224182 primary, plus 1,226 sensitivity rows over budgets, k1/b, seeds and controls.
TargetsThe marginThe margin as registered, plus exact Shapley values and the collection gain G, after the margin proved poorly shaped.
FeaturesSix frequency diagnosticsThe six as registered, plus index statistics and exploratory pseudo-labels.
ValidationFamily and corpusFamily, corpus and held-out k1/b settings, plus six pre-registered confirmatory tests.
Setting B modelsLlama-3.2-1B/3B, Qwen2.5-7BLlama-3.2-1B-Instruct on RULER, because 11 GB cards limit context length and model size.
Setting C4-bit OLMoEOLMoE-1B-7B on WikiText-2 and MMLU, with GSM8K first in line to be cut.

§13Glossary

TermMeaning
nDCG@10A score from 0 to 1 for the top ten results, which rewards relevant documents more the higher they are ranked and normalises by the best possible ordering.
MS MARCO, TREC-DLWeb-passage retrieval from Bing. TREC-DL 2019 and 2020 are smaller query sets over the same passages with deeper human judgements.
BEIRA suite of retrieval datasets from different domains, used to test out-of-domain generalisation.
MIRACLA multilingual Wikipedia retrieval benchmark, of which we use Hindi, Bengali and Telugu.
LIMITA benchmark built to expose the limits of single-vector embedding retrieval.
ZipfianA distribution in which frequency falls off as a power of rank, with a few very common items and a long tail of rare ones, as with words in text.
SAE (Top-K)A sparse autoencoder that rewrites a vector as at most K active features out of many.
PQ / RQProduct and residual quantisation, two ways of compressing a vector into a few discrete codes.
Shapley valueA fair split of a total effect among contributors, averaging each one's marginal effect over every order in which they could be added. Exact here because all eight configurations were run.
Coefficient of variationStandard deviation divided by mean; a unit-free measure of how spread out values are.
Pseudo-labelsTreating another system's top results as if they were relevant, to estimate quality without human judgements.
Query performance predictionEstimating how well a retriever did on a query from the query and its ranking alone, with no labels.
KV cacheThe stored keys and values of earlier tokens, reused at every decoding step.
QuestQuery-aware selection of KV pages, which uses per-page minimum and maximum keys to bound attention.
H2OKV-cache eviction that keeps recent tokens plus the "heavy hitters" with the most accumulated attention.
Attention sinkEarly tokens that absorb a large share of attention, and whose removal breaks generation.
RULERA synthetic long-context benchmark built from needle-in-a-haystack tasks and harder variants.
OLMoEA fully open mixture-of-experts language model with 64 experts per layer, 8 of them active.
Ridge regressionLinear regression with an L2 penalty that shrinks the weights towards zero.
BootstrapEstimating uncertainty by resampling the data with replacement many times.
Holm correctionA way of adjusting p-values across many tests so that false positives stay under control.
Spearman ρThe correlation between two rankings, from −1 to 1.
Wilcoxon testA test of whether paired differences tend to be positive, without assuming they are normal.
Pre-registeredFixed before the results are seen, so the analysis cannot be tuned to the outcome.

§14Viva questions

Questions that are likely to come up about the mid submission, each with an answer short enough to give out loud.

  1. Explain the project in two sentences.

    "Retrieval, KV-page selection and MoE routing all take a dot product against non-negative features and keep the top k, and only lexical retrieval corrects that score, through BM25. We test when each of BM25's three components helps across a controlled family of learned vocabularies, try to predict that without relevance labels, and then carry the predictors into attention and routing."

  2. What is the main finding so far?

    "The correction helps in 90% of rows, but 62% of the gain comes from saturation, which uses no collection statistics. How much saturation helps is set by how unequal the stored values are, which we confirmed with pre-registered tests. Whether IDF and length normalisation help depends on relevance, which is why our frequency-based predictor failed."

  3. Why use Shapley values instead of ordinary main effects?

    "Because the components interact, so the effect of turning IDF on depends on what else is on. Shapley values average each switch's marginal effect over every order of turning them on, and they add up exactly to the total gain. With all eight configurations run, they are exact, with no sampling."

  4. Your pre-registered predictor failed. Why should we believe anything else?

    "We report the failure as registered and explain it. Part of it is the target: the margin saturates near 3 whenever the inner product is the outlier, whatever the size of the gain. The rest is that the frequency curve cannot see relevance. The saturation account was then tested on data the exploration never touched, and the pseudo-label result is clearly labelled exploratory, with a second pre-registration planned on the unseen MS MARCO budget rows."

  5. Why does saturation depend on the coefficient of variation?

    "Saturation maps f to (k1+1)f/(f+k1), and the ratio of output to input falls as f grows, so it shrinks large values relative to small ones. If every stored value were equal it would rescale every score alike and change no ranking. The more unequal the values, the more it reorders. Pooled PQ and RQ store equal values, and their saturation effect is exactly zero."

  6. Why is IDF so weak here when CL-SR loses up to 18% without it?

    "At a cap of 128 document features, the top-k truncation has already removed most of the very common features that IDF would down-weight, which is the explanation Han et al. give. IDF's share shrinks as the budget grows and is slightly negative at native support. It isn't empty, though: shuffling the frequencies makes it worse in 83% of rows."

  7. How does the pseudo-label predictor work, and what is wrong with it?

    "We treat a reference retriever's top 10 as relevant and compute nDCG for each configuration against it. Plugging those pseudo-qualities into the Shapley and G formulas gives label-free estimates, which reach held-out R² of 0.52 to 0.66 for G. The weaknesses are that both references are biased, BM25 towards BM25-like behaviour and E5 towards the encoder, that it needs a second retriever run, and that the decision rule built on it gains only 0.005 nDCG@10."

  8. Why does the inner product over the sparse codes score so badly?

    "The codes are capped and, at token level, summed, so a few features build up very large values and dominate the dot product. Saturation was designed for exactly that case, which is why the token-level SAE, with the most unequal values, also has the largest saturation effect."

  9. How do you get frequencies for KV pages and experts, which aren't documents?

    "We use selection counts on a calibration set. For a page, it is the fraction of calibration queries for which it enters the top-k among the pages that exist then. For an expert, it is the share of calibration tokens routed to it. Scores go through a positive map first, a softplus for Quest and the router probabilities for OLMoE, so that saturation is well defined."

  10. What is the biggest limitation?

    "Everything comes from one encoder, and the sparse codes are much weaker than it, so the attribution might not hold for a strong retriever like SPLADE. E5 has also seen MIRACL, so the Indic results aren't about unseen languages."

Course project · ANLP Doomsday · ANLP Monsoon '26
home · github ↗