torchcell.models.equivariant_cell_graph_transformer.ResponseBasisHead

class ResponseBasisHead(hidden_dim: int, rank: int = 32, param_dim: int = 1, dropout: float = 0.1)

Bases: Module

Low-rank response-basis decoder: the perturbation activates shared programs.

\[\hat y_i = \hat y_i^{\text{local}} + \sum_{j=1}^{r} b_{ij}(h_i)\, a_j(z_S)\]

The perturbation picks an amplitude for each of r global response PROGRAMS, and each gene carries a learned sensitivity to each program. This is a conditional factor model, not attention – and it is the multiplicative interaction between gene identity and perturbation that the additive g(h_i + c_b) form cannot represent, since h_i and c_b meet exactly once by addition and the next op is a LayerNorm that removes the scale in which their inner product lives.

rank=32 is MEASURED, not chosen: residual_covariance_diagnostic.py puts the participation-ratio effective rank of the reproducible residual gene-gene correlation at 32.8 (split-half r = 0.869 vs a 1e-4 null). If the response really is that low-rank, this parameterizes it directly for ~2*d*r parameters.

Supersedes LowRankBilinear, whose first version applied a LayerNorm to the product and so normalized away the magnitude it existed to expose.

forward(H_genes_pert: Tensor, pert_context: Tensor) → Tensor

Compute the low-rank response contribution.

Parameters:

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

• pert_context – [batch, 2d] the [z_S ; h_CLS] perturbation summary.

Returns:

[batch, N] when param_dim == 1 else [batch, N, param_dim].