Skip to content

Instantly share code, notes, and snippets.

@HSILA
Created July 15, 2026 04:08
Show Gist options
  • Select an option

  • Save HSILA/06535f3f7ae577e9cc284e0d7360f564 to your computer and use it in GitHub Desktop.

Select an option

Save HSILA/06535f3f7ae577e9cc284e0d7360f564 to your computer and use it in GitHub Desktop.
Chembed

Why ChEmbed improves on ChemRxiv retrieval but not on encyclopedic benchmarks: an embedding-geometry analysis

1. The question

Our evaluation shows a pattern that a reviewer will immediately question:

Benchmark Text type Metric Base model ChEmbed Change
ChemRxiv Retrieval (ours) chemistry literature paragraphs MRR@10 0.796 0.879 up
ChemTEB encyclopedic retrieval tasks chemistry-topic Wikipedia-style passages MRR@10 0.728 0.672 down
MTEB retrieval (general) general web/encyclopedic text nDCG@10 0.544 0.467 down

A model fine-tuned on chemistry got worse on a chemistry benchmark. Without an explanation, this invites two damaging readings: the fine-tuning damaged the model, or we overfit to our own benchmark. The analysis below rules out both by measuring what actually happened in the embedding space. The core claim it supports: topic and text distribution are different things. A ChemRxiv paragraph and a Wikipedia passage about the same compound differ systematically in vocabulary, sentence structure, information density, and length. ChEmbed specialized in the former, so no improvement should be expected on the latter, and the observed pattern is the predictable signature of specialization rather than degradation.

The argument has two halves, each carried by established statistical methods applied to our corpora: first, that literature and encyclopedic chemistry text occupy measurably distinct regions of embedding space, independent of our training; second, that fine-tuning reorganized the space almost exclusively in the region where the training data lies.

2. Setup

Two models: nomic-embed-text-v1 (the base) and ChEmbed-vanilla (fine-tuned from it on 1.7M retrieval-style pairs, over 98% of which pair an LLM-generated query with a real chemistry passage). Two corpora: 10,000 passages from the ChemRxiv Retrieval corpus, and 2,000 passages sampled evenly from Natural Questions and HotpotQA, the same sources from which the ChemTEB retrieval tasks are built. Only corpus passages are used, never queries, so that synthetic-query generation artifacts cannot influence any result. Where a measurement requires balanced groups, both corpora are subsampled to 2,000 passages per side; all sampling is fixed-seed and reproducible.

One design decision matters for the credibility of the first half: the distribution-shift measurements are performed in the base model's embedding space, not ChEmbed's. The base model is neutral ground; if the two text types are already separated there, the separation cannot be an artifact of our fine-tuning. (Run in ChEmbed's space, the same tests give nearly identical results, so the shift is model-independent; the base-space numbers are the ones we report.)

3. Is the distribution shift real? Two-sample tests

The statistical question underneath "are these two corpora the same kind of text?" is a classical one: given two samples of points, decide whether they were drawn from the same distribution. We apply two standard tests for it.

3.1 Classifier two-sample test

The method. Train a simple classifier to predict which sample a point came from, and evaluate it on held-out points. If the two samples come from the same distribution, no classifier can beat chance (50% with balanced groups); accuracy significantly above chance rejects the same-distribution hypothesis. This is the classifier two-sample test of Lopez-Paz and Oquab (ICLR 2017). Using a linear classifier makes the test conservative: it can only detect separation that is simple enough to cut with a hyperplane, so a high score cannot be attributed to an over-powerful model finding incidental patterns.

Our application. Logistic regression on the first 64 principal components of the base-model embeddings, 2,000 passages per source, stratified five-fold cross-validation.

Result and implication. 99.8% held-out accuracy against a 50% chance level. Literature and encyclopedic chemistry passages are not merely somewhat different; they are almost perfectly linearly separable in the base model's own representation. The distribution shift is a property of the data that predates our training.

3.2 Nearest-neighbor two-sample test

The method. Pool the two samples and ask, for each point, what fraction of its k nearest neighbors come from its own sample. Under the same-distribution hypothesis this fraction is near chance; a high fraction means the samples occupy distinct regions. This nearest-neighbor coincidence statistic underlies the multivariate two-sample tests of Schilling (JASA, 1986) and Henze (Annals of Statistics, 1988). It complements the classifier test because it is sensitive to local structure: two distributions could overlap globally yet still be locally segregated, or vice versa.

Our application. Cosine nearest neighbors, k = 10, balanced pooled sample of 2,000 + 2,000.

Result and implication. 97.8% of each passage's ten nearest neighbors come from its own source (chance ≈ 50%). This is the retrieval-relevant version of the shift: retrieval is nearest-neighbor search, so if a corpus forms its own neighborhood island, whatever the model learned about ranking within one island transfers poorly to the other.

3.3 A global summary: centroid distance

The mean embedding of each corpus gives a single direction summarizing it. The cosine distance between the two corpus centroids is 0.146 in the base model's space and 0.571 in ChEmbed's space. Two readings: the corpora drift apart even for the base model, and fine-tuning actively increased the separation between the text type it was trained on and everything else, which previews the second half of the argument.

4. What did fine-tuning change? Representation comparison

The second half of the argument compares the two models on the same texts, using two standard tools from the representation-analysis literature. The precedent for this kind of study is Merchant et al. (2020), who used probing and representational similarity to show that fine-tuning reshapes BERT representations selectively rather than globally; our finding is the retrieval-space version of the same phenomenon.

4.1 Neighbor overlap

The method. Embed the same passages with both models. For each passage, take its ten nearest neighbors under model A and under model B and measure the overlap of the two sets. This directly measures the quantity retrieval depends on: whether the two models agree on what counts as similar.

Result and implication. For encyclopedic passages the two models share 64% of neighbors on average. For ChemRxiv passages they share 0.1%: essentially none. Fine-tuning completely rewired similarity for literature text while leaving general text largely as the base model represents it. The training data (98%+ chemistry-literature pairs) predicts exactly this footprint.

4.2 Centered kernel alignment (CKA)

The method. Linear CKA (Kornblith et al., ICML 2019) is the standard index for comparing two representations of the same inputs, on a 0–1 scale. Its key property here is invariance to rotation and uniform scaling: a high CKA rules out the objection that the neighbor disagreement in 4.1 is a trivial re-orientation of the same geometry, and a low CKA confirms the geometry itself changed.

Result and implication. CKA is 0.020 between the two models on ChemRxiv passages and 0.867 on encyclopedic passages. Same contrast, independent method: the literature region was rebuilt, the general region was largely preserved (substantially, though not perfectly, which matters below).

4.3 How the literature region changed: from crowding to resolution

Average pairwise cosine similarity between passages tells us how crowded each region is. In the base model, two random ChemRxiv paragraphs have mean similarity 0.58, while two unrelated passages (one literature, one encyclopedic) sit at 0.41; in ChEmbed's space the same numbers are 0.02 and 0.02. In other words, the base model sees all chemistry literature as near-duplicates of one another, huddled in a narrow cone, while ChEmbed spreads them apart until two chemistry papers are as distinguishable from each other as chemistry is from Wikipedia. At first glance the ChEmbed numbers can look like a loss of structure; they are the opposite. Ranking passages against a query requires resolving differences within the corpus being searched, and a space where every candidate looks alike has no resolution to offer. The observation that transformer embeddings are natively crowded into narrow cones is itself well documented (Ethayarajh, EMNLP 2019); our fine-tuning effectively undid that crowding precisely where it mattered.

4.4 Why no PCA scatter plot

A 2D projection of a 768-dimensional space is not evidence in either direction: clusters that overlap in the top two principal components can be perfectly separable in the full space (our 99.8% linear probe demonstrates exactly this), and apparent separation can be a projection accident. The tests in Section 3 are the defensible form of the question a scatter plot only gestures at, which is why we report those and omit the picture.

5. Putting it together

The three benchmark results now read as one mechanism rather than three surprises. ChemRxiv Retrieval improves by 8.3 points MRR@10 because fine-tuning rebuilt the geometry of exactly that region: neighbor structure 99.9% new, crowding removed. The encyclopedic ChemTEB tasks fall by 5.6 points and MTEB retrieval by 7.7 points nDCG@10 because those corpora live in the region the training barely touched, where agreement with the base model is substantial (64% neighbors, 0.87 CKA) but not complete; partial preservation yields a moderate, not catastrophic, cost. And the reason gains do not transfer despite the shared chemistry topic is quantified in Section 3: the two kinds of chemistry text are nearly disjoint distributions, separable at 99.8% by a linear test in the base model's own space. Poor transfer of retrieval models across text distributions is independently documented at benchmark scale (BEIR; Thakur et al., 2021). The practical conclusion we draw in the paper follows: embedding models should be selected to match the text distribution of the target corpus rather than its domain label, and ChEmbed is positioned as a model for chemistry literature retrieval, not as a universally superior chemistry embedder.

6. Scope of the claims

The geometry explains why the benchmark pattern is expected; it does not by itself prove the retrieval scores, which remain the primary evidence. The analysis uses the vanilla ChEmbed variant; the other variants share the same data and recipe and are assumed to behave similarly, but were not measured. Passage length and formatting differ between the corpora and were not controlled for, so part of the measured separation may be carried by such surface features; this is acknowledged in the paper's limitations. None of these caveats affects the direction of the argument, only how far it should be stretched.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment