torchcell.models.equivariant_cell_graph_transformer.ObservedLabelEncoder

class ObservedLabelEncoder(hidden_dim: int, dropout: float = 0.1, gate_mode: str = 'on')[source]

Bases: Module

Encode PARTIALLY OBSERVED labels so the decoder can condition on them.

WHY. Every arm in this round changes what information reaches the per-gene readout; none changes the OUTPUT FACTORIZATION. The decoder emits 6,127 INDEPENDENT marginals in one shot, even though the reproducible residual gene-gene correlation has a measured participation-ratio effective rank of 32.8 (split-half r = 0.869 against a 1e-4 null). That joint structure is simply discarded at the output.

Masked prediction – the scGPT / BERT objective – restores it: mask a fraction of the measured genes, feed the REST in as input, and predict the held-out ones. Gene i’s prediction then depends on gene j’s observed value, which is exactly the dependence a per-token readout cannot express.

IT ONLY WORKS WITH CROSS-GENE MIXING. If nothing routes information between gene tokens after the perturbation, an observed value at gene j can never reach gene i and the conditioning is inert. This module therefore pairs with PerceiverMixing: the Perceiver is the channel, masking is the objective that forces the model to use it. (Measured: the Perceiver alone moved seed 0 by only +0.0045 – a pathway with nothing requiring it.)

A 100%-MASKED FORWARD IS NOT THE UNCONDITIONED MODEL. With no label observed the features are [0, 0], but this module still adds gate * proj([0, 0]) to every gene token (see __init__). In eval that offset is the same vector on every token; in train the dropout inside proj makes it differ per token; under gate_mode="rezero" it is zero at init (gate 0) and grows with the learned gate. Only a call with observed_values=None skips the encoder and is the unconditioned model, so a metric from an all-masked pass is the unconditioned model plus that offset (issue #566).

forward(H_genes_pert: Tensor, observed_values: Tensor, observed_mask: Tensor) → Tensor[source]

Add an encoding of the observed labels to the gene tokens.

Parameters:

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

  • observed_values – [batch, N] label values, ZERO where unobserved.

  • observed_mask – [batch, N] 1.0 where the label is observed, else 0.0.

Returns:

[batch, N, d] embeddings carrying the observed-label context.