torchcell.models.equivariant_cell_graph_transformer.CellGraphTransformer

class CellGraphTransformer(gene_num: int, hidden_channels: int, num_transformer_layers: int, num_attention_heads: int, cell_graph: HeteroData, graph_regularization_config: dict[str, Any] | None = None, perturbation_head_config: dict[str, Any] | None = None, dropout: float = 0.1, graph_reg_lambda: float = 0.0, node_embeddings: dict[str, Any] | None = None, learnable_embedding_config: dict[str, Any] | None = None, heads_config: dict[str, Any] | None = None, perturbation_propagation_config: dict[str, Any] | None = None, post_perturbation_mixing_config: dict[str, Any] | None = None, attention_mask_config: dict[str, Any] | None = None, observed_label_config: dict[str, Any] | None = None, cross_gene_config: dict[str, Any] | None = None)[source]

Bases: Module

Equivariant Cell Graph Transformer model.

Architecture: 1. Gene embeddings + CLS token 2. Transformer encoder with graph-regularized attention 3. EQUIVARIANT perturbation transformation (Type I Virtual Instrument) 4. Perturbation readout head (Type II Virtual Instrument)

compute_graph_regularization_loss(attention_weights: Tensor, layer_idx: int) → Tensor[source]

Compute graph regularization loss using KL divergence.

Parameters:

  • attention_weights – [batch, heads, N, N] gene-gene attention weights

  • layer_idx – Current transformer layer index

Returns:

Total regularization loss for this layer

forward(cell_graph: HeteroData, batch: HeteroData, return_attention: bool = False, observed_values: Tensor | None = None, observed_mask: Tensor | None = None) → tuple[Tensor, dict[str, Tensor]][source]

Forward pass of Equivariant Cell Graph Transformer.

Parameters:

  • cell_graph – Full wildtype graph structure (not used directly, genes indexed by order)

  • batch – Perturbation data with indices and phenotypes

  • return_attention – If True, store and return attention weights (memory intensive: it forces the manual softmax path, which materializes a [1, heads, N+1, N+1] matrix per layer instead of using the fused kernel)

  • observed_values – [batch, N] partially observed measured labels for masked prediction, or None. Only consumed when the model was built with an observed_label_config.

  • observed_mask – [batch, N] boolean companion to observed_values, True where a value is observed. Passing an all-False mask is NOT the unconditioned model: the forced-on encoder still adds proj([0, 0]) to every token. Pass observed_values=None to skip the encoder.

Returns:

[batch_size, 1] gene interaction predictions representations: Dict with embeddings, attention weights (if requested), losses, and H_genes_pert

Return type:

predictions

property num_parameters: dict[str, int]

Count parameters in each component.

Every entry is that block’s own parameters() size (trainable or not), and total is sum(p.numel() for p in self.parameters()), so a size table read from total is the module’s real size. gene_embedding and embedding_preprocessor are always present (0 when disabled); every other optional block appears only when it is built. On this class the block entries sum to total; a subclass that registers more modules adds its own entries.