torchcell.models.equivariant_cell_graph_transformer.PerturbationGraphPropagation

class PerturbationGraphPropagation(hidden_dim: int, graph_names: list[str], hops: int = 2, dropout: float = 0.1, gate_mode: str = 'rezero')[source]

Bases: Module

Route the perturbation to graph NEIGHBOURS – the missing pair-(p, i) term.

WHY THIS EXISTS. The encoder runs at batch 1 on the WILDTYPE graph, so H_genes and h_CLS are identical for every strain; the only strain-dependent step is EquivariantPerturbationTransform, whose K/V set is just the |S_b| perturbed tokens. For a SINGLE deletion the softmax is over one key, so the attention weight is exactly 1 for every query gene and the attended vector is query-INDEPENDENT:

\[H^{\mathrm{pert}}_{b,i} = g\big(h_i + c_b\big), \qquad c_b = W_O W_V h_p .\]

Verified numerically (max \(|attended_i - attended_0| = 0\)). ~96% of the expression rows are single deletions (Kemmeren2014 1,484 singles + Sameith2015, whose double-KO sub-study is n=72), so for almost the whole training set the model’s prediction factorizes ADDITIVELY: gene identity enters as \(h_i\), strain identity as one 90-d vector, and no term depends on the pair (p, i). By data processing every head output is then a function of \(c_b\) alone, which is why no readout change (S0 -> S1 -> S3, any scoring rule) can enlarge the hypothesis class. The nine graphs shape \(h_i\) but never CARRY the deletion.

WHAT THIS DOES. For each graph g and hop t, spread unit mass outward from the perturbed genes over the row-normalized adjacency:

\[r^{g,1} = \hat A_g^\top p_b, \qquad r^{g,t} = \hat A_g^\top r^{g,t-1},\]

giving len(graphs) * hops scalars per gene that answer “how strongly is gene i reached from the deletion, along graph g, in t hops”. That quantity depends on the PAIR, so it is exactly the term the additive form lacks. The features are projected to d and added to H_genes_pert through a ReZero gate initialized at ZERO – so the module starts as an exact identity and this is a clean ablation: any gain is attributable to propagation, not to the extra parameters.

Cost is a sparse mat-vec per (graph, hop), i.e. O(|E|) – negligible next to the encoder’s dense attention over 6,607 tokens.

forward(H_genes_pert: Tensor, adjacencies_T: dict[str, Tensor], perturbation_indices: Tensor, batch_assignment: Tensor) → Tensor[source]

Add graph-routed perturbation features to the perturbed gene embeddings.

Parameters:

  • H_genes_pert – [batch, N, d] equivariant perturbed gene embeddings.

  • adjacencies_T – relation name -> TRANSPOSED row-normalized sparse adjacency [N, N], so A_T @ x spreads mass outward from the perturbed genes.

  • perturbation_indices – [total_pert_genes] indices of perturbed genes.

  • batch_assignment – [total_pert_genes] batch index per perturbed gene.

Returns:

[batch, N, d] embeddings with the pair-(p, i) signal added.