torchcell.models.equivariant_cell_graph_transformer.LowRankBilinear

class LowRankBilinear(hidden_dim: int, rank: int = 32)

Bases: Module

Explicit rank-R multiplicative interaction between gene identity and perturbation.

WHY. h_i (what gene i is) and c (what was deleted) meet exactly once, by ADDITION, and the very next op is a LayerNorm. The natural pair-term for expression is an inner product – “how related is output gene i to the deleted gene” – and the sum does contain it:

\[\|h_i + c\|^2 = \|h_i\|^2 + 2\langle h_i, c\rangle + \|c\|^2\]

but LayerNorm divides by the per-channel standard deviation, DISCARDING exactly the magnitude in which \(\langle h_i, c\rangle\) is encoded. So the interaction is not merely hard to read off the sum – the operation immediately after the sum removes it.

This computes the interaction directly as a rank-R bilinear form,

\[b_r = \langle u_r, h_i\rangle \cdot \langle v_r, c\rangle, \qquad r = 1 \ldots R,\]

yielding R scalars per (strain, gene) that are appended to the per-gene head input. R=32 matches the measured effective rank (32.8) of the reproducible residual gene-gene correlation. Cost is two d x R projections – negligible – and it is the parameter-lean alternative to the full concat_context arm: if concat wins because the head is re-deriving an inner product, this should recover most of it for 2*d*R parameters.

Permutation-equivariant (no gene-indexed parameters), graph-free.

forward(H_genes: Tensor, context: Tensor) → Tensor

Compute the rank-R interaction per (strain, gene).

Parameters:

• H_genes – [N, d] strain-invariant wildtype gene embeddings.

• context – [batch, N, d] attended perturbation context.

Returns:

[batch, N, rank] bilinear interaction features.