torchcell.models.cell_graph_transformer.PerturbationHead

class PerturbationHead(hidden_dim: int, num_heads: int = 4, dropout: float = 0.1, use_cross_attention: bool = True)

Bases: Module

Perturbation head with switchable attention mechanisms.

Implements g_ψ(h_CLS, H_genes, M(S)) from weekly report using either: - Cross-attention (default): perturbation summary attends to all genes - HyperSAGNN: within-set self-attention for perturbation genes

forward(H: Tensor, perturbation_indices: Tensor, batch_assignment: Tensor) → Tensor

Forward pass of perturbation head.

Parameters:

• H – [N+1, d] transformer output (CLS at position 0, genes at 1:N+1)

• 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