torchcell.models.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, adaptive_loss_weighting: bool = False, graph_reg_scale: float = 0.001)[source]
Bases:
ModuleCell Graph Transformer model.
Architecture: 1. Gene embeddings + CLS token 2. Transformer encoder with graph-regularized attention 3. Perturbation head with cross-attention
- 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) tuple[Tensor, dict[str, Any]][source]
Forward pass of Cell Graph Transformer.
- Parameters:
cell_graph – Full wildtype graph structure (not used directly, genes indexed by order)
batch – Perturbation data with indices and phenotypes
- Returns:
[batch_size, 1] gene interaction predictions representations: Dict with embeddings, attention weights, and losses
- Return type:
predictions