torchcell.models.equivariant_cell_graph_transformer.GraphRegularizedTransformerLayer

class GraphRegularizedTransformerLayer(hidden_dim: int, num_heads: int, dropout: float = 0.1)[source]

Bases: Module

Transformer layer with graph-regularized attention heads.

Uses manual attention computation to get both output and attention weights for graph regularization loss.

forward(x: Tensor, return_attention: bool = False, head_mask: Tensor | None = None) → tuple[Tensor, Tensor | None][source]

Forward pass with manual attention computation.

Parameters:

  • x – [batch, N+1, d] where N+1 includes CLS token at position 0

  • return_attention – Whether to return attention weights for regularization

  • head_mask –

    Optional [heads, N+1, N+1] bool mask, True where attention is ALLOWED. When given (and weights are not requested) attention runs through scaled_dot_product_attention.

    WHY THAT MATTERS FOR MEMORY. This layer computes attention MANUALLY – explicit matmul then softmax – because the graph-regularization KL needs the weights. That materializes a [batch, heads, N+1, N+1] tensor every layer: 9 x 6608^2 x 4B = 1.57 GB, and it is the bulk of the measured 19.3 GB/run. Turning the KL off does NOT free it; the manual path materializes regardless. Masking replaces the KL as the way to inject graph structure, which removes the need for the weights, which finally allows the fused/flash SDPA kernel that never forms the matrix.

Returns:

[batch, N+1, d] transformed features gene_attention_weights: [batch, heads, N, N] gene-gene attention weights (if return_attention=True)

Return type:

output