What the model learned about genes

What each coloring reveals

The paper visualizes the K = 10 embeddings with PCA on the left and cosine UMAP on the right. This interactive version uses the UMAP view and lets you switch among the same biological annotations.

• HDBSCAN clusters. Genes are grouped using their original 10-dimensional embeddings, before they are projected into this two-dimensional view. HDBSCAN uses a minimum cluster size of 150, and each cluster is labeled by the screen phenotype most strongly associated with it. Several clusters are associated with the DNA-damage response without being further disambiguated.
• CORUM complex families. Colors show the largest coarse families of physical protein complexes, with smaller families pooled into other. Their weak spatial grouping shows that physical-complex membership does not organize this embedding particularly well.
• Reactome categories. Colors show broad, top-level Reactome categories rather than hundreds of specific pathways. Apart from the common-essential region, they show little distinct spatial clustering.
• Hit rate. A gene's hit rate is its hits divided by the number of public training screens it appears in. Genes in few screens have a noisy rate, which is what the minimum screens present filter is for.
• Common-essential genes. The DepMap coloring highlights genes required across many cell lines, revealing the common-essential region referred to above. These annotations come from DepMap Public 26Q1 and remain subject to the DepMap Terms and Conditions; they are not covered by AssayLoop's MIT License.
• Grey points remain visible. Genes without a CORUM family or assigned Reactome category are shown as unassigned, rather than dropped.

Which embedding is worth starting from?

Recovering textbook biology does not predict task usefulness. For each initial embedding, cosine similarity is used to score gene–gene relationships from four established databases: STRING, CORUM, SIGNOR, and Reactome. A gene pair is labeled positive when its relationship appears in the database and negative otherwise. AUROC is calculated for each database and then averaged across all four. Surprisingly, embeddings that recover these databases most closely, such as GenePT and K562, are harder to train AssayFormer on. Initializations learned directly from historical screen data, such as matrix factorization, are the most useful for the downstream task.

How training changes the embeddings

Initialized BPMF embeddings barely change during training. In the BPMF row, each point in the left panel represents one of 300,000 randomly sampled gene pairs. Its horizontal position is the pair's cosine similarity before training, and its vertical position is the pair's cosine similarity after RL. The right panel measures rank-biased overlap (RBO) between each gene's 100 nearest neighbors before and after SFT, from SFT to RL, and from initialization to RL. Most BPMF gene neighborhoods change very little during SFT, and RL produces almost no additional change.

Other initializations can reorganize substantially. The full figure shows that embedding drift depends strongly on the source: BPMF nearly preserves its original gene neighborhoods, while several other initialization types change markedly during supervised training.