torchcell.models.equivariant_cell_graph_transformer.CrossAttnHead

class CrossAttnHead(hidden_dim: int, output_dim: int, num_heads: int = 4, dropout: float = 0.1, param_dim: int = 1, use_ffn: bool = True)[source]

Bases: Module

S3 cross-attention readout head (DETR / Perceiver style) for vector phenotypes.

The S1 GlobalHead collapses the token set to ONE pooled vector u = [h_CLS || mean_i h_i_pert] that every output feature must read from. A mean over ~6000 gene tokens dilutes a 1-2 gene perturbation by ~3 orders of magnitude, so the strain-specific signal is largely gone BEFORE the readout – the pool-dilution bottleneck. That is the leading hypothesis for morphology sitting at r~0.04 against a 0.61 noise ceiling while expression (which reads its own token, S0) is at its ceiling.

S3 removes the shared bottleneck: each of the F output features gets a LEARNED QUERY q_k that cross-attends over the FULL token set {h_CLS_pert, h_1_pert, ..., h_N_pert}:

\[A = \mathrm{softmax}\!\Big(\frac{(QW_Q)(H_{\mathrm{pert}}W_K)^\top}{\sqrt{d_k}}\Big), \qquad C = A\,(H_{\mathrm{pert}}W_V),\]

so output k reads its OWN weighted mixture c_k – per-output support is recovered WITHOUT needing a 1-to-1 feature<->token map (morphology features are not genes). Feature k can place attention mass directly on the perturbed genes instead of averaging them away. CLS is kept as a key so a feature can still read the whole-cell summary.

Queries are parameters (like the CLS token), NOT label values, so nothing about the target leaks in. The readout is DistHead-agnostic: param_dim widens the final projection to emit point / gaussian / quantile parameters per feature.

forward(h_CLS: Tensor, H_genes_pert: Tensor) → Tensor[source]

Forward pass of the cross-attention head.

Parameters:

  • h_CLS – [d] whole-cell CLS representation.

  • H_genes_pert – [batch, N, d] equivariant perturbed gene embeddings.

Returns:

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

Return type:

predictions