torchcell.models.equivariant_cell_graph_transformer.CrossGeneMixing

class CrossGeneMixing(hidden_dim: int, rank: int = 64, dropout: float = 0.1)

Bases: Module

GEARS-style cross-gene mixing: one shared global state, then per-gene readout.

GEARS computes a single cross_gene_state from ALL gene embeddings and concatenates it to every gene before the gene-specific output layer; ablating that layer measurably hurt its differential-expression correlation. It is the cheapest form of cross-gene communication – one pooled vector, no pairwise term – and sits strictly between the per-token readout (no communication at all) and full attention.

GEARS flattens [N, d] into a single Linear, which here would be 6,607 x 90 -> 38M parameters. This uses the low-rank equivalent: project each gene to rank, mean across genes, broadcast back. Same information path (every gene sees a summary of every other gene), tractable parameter count.

NOTE ON WHAT THIS CAN AND CANNOT DO. Every post-perturbation token is g(h_i + c_b), so the pooled state is a function of c_b – pooling adds no INFORMATION. It does enlarge the FUNCTION CLASS: the map c_b -> profile is no longer “add c_b to each h_i and apply one shared MLP”. Since the reproducible signal has effective rank ~33 and c_b carries 90 dimensions, information is not the binding constraint here; expressivity is. That is the hypothesis this arm tests.

forward(H_genes_pert: Tensor) → Tensor

Mix a pooled cross-gene state back into every gene token.

Parameters:

H_genes_pert – [batch, N, d] post-perturbation gene embeddings.

Returns:

[batch, N, d] with the shared cross-gene state folded in.