We investigated whether ChEmbed training produces a more chemically meaningful representation space than the original Nomic model. Our initial idea was that descriptions of chemically related entities, identified using ChEBI relationships, might become more clearly grouped after training. Such a result could provide an intuitive visual explanation of what the model learned.
We connected PubChem entity descriptions to structured ChEBI relationships and compared their representations before and after training. We examined several complementary behaviours:
- whether related chemical descriptions became closer together;
- whether chemistry questions retrieved their corresponding descriptions more reliably;
- whether queries about chemical roles retrieved compounds annotated with those roles;
- whether the behaviour survived removal of compound names or explicit role terminology; and
- whether a query written for one description could retrieve an independent description of the same compound.
The comparisons included Nomic, ChEmbed-vanilla, and ChEmbed-progressive.
The original "tighter chemical clusters" hypothesis was not supported. By tighter clusters, we mean that descriptions of related compounds would be positioned closer to one another and more clearly separated from unrelated descriptions. However, a more training-aligned and useful pattern emerged.
ChEmbed, particularly the progressive model, became substantially better at connecting chemical-role questions with descriptions of compounds having those roles. For example, for the query asking which chemicals can function as coagulation factor Xa inhibitors, four of the eight ChEBI-linked inhibitors appeared in Nomic's first ten results, compared with all eight after progressive training.
This was not limited to one example. Across 170 ChEBI chemical-role families:
- the overall ranking improved for 143 families; and
- the fraction of known family members recovered in the first ten results increased, on average, from approximately 28% to 41%.
The factor Xa example is intentionally a strong illustrative case, rather than a typical one, and is labelled accordingly in the figure. The complete-cohort result is included alongside it to show that the effect is broader but not universal.
Another consistent observation helped explain the geometry: related descriptions generally became more dispersed after training, while chemistry queries became more selective about which description they retrieved. In other words, ChEmbed did not create tighter chemical clusters; it made descriptions easier to distinguish according to a particular search intent.
The clearest use of this result is as a before-and-after retrieval example, rather than as a two-dimensional embedding map.
We started from the frozen catalog used in the earlier analysis: 4,750
CID-linked target descriptions and 5,000 fixed PubChem background descriptions,
giving 9,750 descriptions in total. From the ChEBI relationships, we extracted
all 170 eligible has_role families and the PubChem compounds linked to each
role. No family was included or excluded according to model performance.
For every ChEBI role name, we created the same three query formulations:
- T1: "Which compounds act as [role]?"
- T2: "Find chemical entities with the role of [role]."
- T3: "What chemicals can function as [role]?"
This produced 510 queries, consisting of three phrasings for each of the 170 roles. Every query was embedded by Nomic, ChEmbed-vanilla, and ChEmbed-progressive, then compared with the same 9,750 embedded descriptions. The descriptions were ranked by query similarity. ChEBI supplied the list of known members for each role, so we could ask whether those members appeared earlier in the ranking after training. Other descriptions were left unlabeled and were not automatically treated as chemically incorrect.
The cohort figure uses known-member average precision, which is high when the ChEBI-linked members of a role are repeatedly found near the beginning of the ranked list. The factor Xa figure uses the simpler count of known members found within the first ten results.
The main figure can show the factor Xa example in two columns: Nomic before training and ChEmbed-progressive after training. Each column should display the top retrieved compounds using their molecular structures, names, and ranks, with the eight ChEBI-linked inhibitors visually highlighted. This makes the change immediately readable: the relevant compounds move from being only partially represented in the first ten results to all being present after training.
-
Panel A, consistency across query wording: The bars show average retrieval quality for T1, T2, and T3. ChEmbed-progressive improves over Nomic for all three formulations, indicating that the result is not caused by one particularly favourable wording.
-
Panel B, changes across the 170 families: Each point represents one ChEBI role family and shows its change relative to Nomic, averaged across the three queries. Points above zero improved. Progressive training improves 84% of the families, while vanilla is much less consistent.
-
Panel C, distribution over all 510 queries: This panel retains every role and every phrasing. A curve shifted farther to the right means that more queries retrieve their known ChEBI members effectively. The progressive curve shows a broad improvement rather than a result produced only by the selected factor Xa example.
-
Panel D, shuffled-label control: We randomly reassigned compound-to-role labels 999 times while leaving the model rankings fixed. The shaded distributions show the changes expected from arbitrary labels, and the diamonds show the changes actually observed. The observed progressive change is far larger than the shuffled-label expectation.
Together, the two figures communicate complementary levels of evidence. The factor Xa figure provides an intuitive example using recognizable compounds and literal ranks. The cohort figure shows that the direction of improvement extends across most role families, although it is not universal.
This presentation should be described as evidence that training changed the model's retrieval behaviour for explicit chemical-role queries. It should not be presented as evidence that the full embedding space became a better map of chemical similarity.
The results do not support the stronger claim that ChEmbed learned a generally improved chemical-entity hierarchy or molecular-similarity space.
Three controls establish this boundary:
-
Compound-name masking: Exact query-to-description retrieval dropped substantially after compound names were removed. This indicates that entity-name overlap contributes strongly to the near-perfect retrieval of exposed PubChem pairs.
-
Chemical-role masking: Removing explicit role terms from descriptions eliminated much of the role-retrieval effect. A smaller progressive advantage remained, but it could reflect surrounding context, synonyms, compound names, or training exposure rather than deeper chemical reasoning.
-
Alternate-description transfer: When a question generated from one PubChem description was used to retrieve an independent description of the same compound, ChEmbed performed worse than Nomic. The learned representation is therefore not clearly source-independent at the entity level.
The signal was also much stronger for ChEmbed-progressive than for vanilla. We should consequently avoid presenting the result as a uniform property of every ChEmbed variant.
The experiment was productive, but it changed the scientific claim.
The supported interpretation is:
ChEmbed training reorganizes PubChem retrieval around explicit chemical search intents, making descriptions associated with many chemical roles easier to retrieve and distinguish.
The unsupported interpretation is:
ChEmbed creates a universally better chemical ontology or molecular representation space.
I recommend using the simplified factor Xa figure as an intuitive behavioural illustration. It shows actual molecular structures and literal before/after ranks, while the accompanying cohort summary demonstrates that the selected example reflects a broader, though non-universal, trend. The more technical clustering, masking, and complete-cohort analyses can remain in the supplementary material.
Because the PubChem exposure is intentional, this figure should be presented as an explanation of learned model behaviour, not as a held-out evaluation of chemical knowledge.