torchcell.models.equivariant_cell_graph_transformer.PerceiverMixing
- class PerceiverMixing(hidden_dim: int, num_latents: int = 32, num_heads: int = 9, dropout: float = 0.1, gate_mode: str = 'rezero')[source]
Bases:
ModulePOST-perturbation gene-gene mixing through a learned latent bottleneck.
WHY. The encoder is the ONLY place gene representations interact, and it runs on the wildtype graph at batch 1 – so that mixing is identical for every strain. After
EquivariantPerturbationTransforminjects the perturbation there is NO gene-gene interaction anywhere: the transform’s keys/values are the perturbed genes only,PerGeneHeadis a per-token MLP, and the global/perturbation heads merely pool. Gene i’s post-perturbation state therefore depends on \((h_i, \{h_p\})\) and never on \(h_j\) for any other gene, which makes downstream cascades (“delete p -> changes j -> changes i”) structurally unrepresentable.Naive fix – dense self-attention over the N tokens AFTER the perturbation – is per-instance, so it costs B times the encoder: ~503 GFLOP and a
[32, 9, 6607, 6607]attention tensor (~25 GB in bf16). Not affordable.This block routes every gene through
num_latentslearned latents instead: genes -> latents -> genes. Every gene still influences every other, but through an M-dim bottleneck, so cost is O(N*M) not O(N^2): ~4.9 GFLOP and a ~0.12 GB attention tensor at M=32. M=32 is also the MEASURED participation-ratio effective rank (32.8) of the reproducible residual gene-gene correlation.Graph-FREE – unlike graph-masked attention it needs no adjacency, so it transfers to organisms with no curated interaction networks. Permutation-equivariant: the latents are shared across genes, never indexed by gene. ReZero-gated at ZERO so the block is an exact identity at init and the arm is a clean ablation.