torchcell.models.cell_graph_transformer_metabolism.perturbed_gene_pool

perturbed_gene_pool(H_genes_pert: Tensor, perturbation_indices: Tensor, perturbation_indices_batch: Tensor) → Tensor[source]

Mean over ONLY the perturbed gene tokens, per sample. Returns [batch, d].

WHY THIS EXISTS, with the numbers MEASURED rather than assumed (run main() in this file to reproduce them on a real batch):

head input term

across-batch std

h_CLS

0.000e+00

mean over ALL 6607 gene tokens

1.309e-02

mean over PERTURBED genes only

3.217e-02

Two conclusions, one of which corrects an earlier guess:

  1. h_CLS is the encoding of the UNPERTURBED reference graph (the parent encodes G once at batch size 1) and is broadcast identically across the batch. Its across-batch variance is EXACTLY zero – it cannot distinguish two strains, yet it occupies a full hidden_dim of the head’s input width. That part is confirmed.

  2. The genome-wide mean is NOT attenuated ~1/6607, as a naive “only one of 6,607 tokens differs” argument predicts. The equivariant perturbation operator cross-attends, so it has already propagated the perturbation into many gene representations before any pooling happens. The perturbed-gene pool carries only ~2.5x the across-batch signal of the genome-wide pool, not ~6607x.

So this pool is a real but MODEST improvement in signal-to-background, and it is NOT a sufficient explanation for the observed mean collapse (job 1331: val pearson exactly 0.00000, val loss ~0.97 = the variance of the standardized target). The genotype signal reaching the head is measurable either way; the model’s failure to map it onto betaxanthin is a separate problem.

Keeping it is still justified – for a single-KO strain the perturbed gene’s representation IS the genotype, which is the S0 argument from the decoder note applied to a scalar target – but do not cite it as the fix for the nulls.

Index convention matches the parent operator (equivariant_cell_graph_transformer.py:1649): perturbation_indices is a flat [total_pert_genes] tensor of gene indices and perturbation_indices_batch assigns each entry to its sample.