torchcell.models.equivariant_cell_graph_transformer.PerturbationHead

class PerturbationHead(hidden_dim: int, dropout: float = 0.1, pooling: str = 'sum')

Bases: Module

Perturbation readout head for gene interaction prediction.

Operates on equivariant perturbed gene embeddings H_genes_pert [batch, N, d] and produces scalar gene interaction predictions per sample.

Implements Type II Virtual Instrument: H_genes_pert → y_GI

forward(h_CLS: Tensor, H_genes_pert: Tensor, perturbation_indices: Tensor, batch_assignment: Tensor) → Tensor

Forward pass of perturbation head.

Parameters:

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

• H_genes_pert – [batch, N, d] - perturbed gene embeddings (EQUIVARIANT)

• perturbation_indices – [total_pert_genes] - indices of perturbed genes

• batch_assignment – [total_pert_genes] - batch index for each perturbed gene

Returns:

[batch_size, 1] gene interaction predictions

Return type:

predictions