Graph Attention over Feature Maps: TensorGATLayer Internals and Honest Limits

Graph attention networks let each node weigh its neighbours instead of treating them equally — a learned, content-dependent average. The original GAT formulation assumes vector node features. TGraphX ships two attention layers that work when nodes are feature maps of shape [C, H, W], and they make different trade-offs. This note explains the math of TensorGATLayer, contrasts it with the lighter AttentionMessagePassing, and is candid about which attention granularities are implemented and which are deliberately deferred.

This builds on the message-passing skeleton and the tensor-valued edges discussion; read those first if the φ/AGG/ψ notation or edge biases are unfamiliar.

The split-head attention math

tgraphx/layers/gat.py documents its own formulation as a split-attention form, equivalent to the classic concat-and-dot GAT. For head k and edge i → j:

score_ij^k  = attention logit for head k on edge (i → j)
        α_ij^k      = softmax_i( score_ij^k )      # normalised over incoming edges to j
        o_j^k       = Σ_{ i ∈ N(j) }  α_ij^k · value_i^k
        o_j         = concat_k o_j^k   or   mean_k o_j^k
        

The important word is softmax_i: the attention weights are normalised over each destination's incoming edges, so Σ_i α_ij^k = 1. That is what makes the output a genuine weighted average of neighbours, and it is exactly the property the test suite checks (see the attention-normalisation invariant in permutation equivariance). Heads are combined either by concatenation (concat_heads=True, which requires out_channels divisible by num_heads) or by averaging.

Two attention modes — and two that are planned

TensorGATLayer exposes an attention_mode argument with two shipped values:

The source is explicit that 'spatial' and 'voxel' modes "are planned for a future release," and the constructor raises a ValueError if you request them. This matches the maturity note in the README: per-pixel and per-voxel GAT scores are not shipped because an [E, K, H, W] score tensor is memory-prohibitive at realistic edge counts. So today you can attend per-channel, but not yet per-pixel — and the framework tells you so rather than silently approximating.

import tgraphx as tgx
        layer = tgx.TensorGATLayer(in_channels=16, out_channels=32, num_heads=4,
                                   attention_mode="channel")
        

The lighter alternative: AttentionMessagePassing

tgraphx/layers/attention_message.py offers a different, cheaper attention. It projects query, key, and value with 1×1 convolutions, computes a score by summing the query–key product over the channel axis and scaling by 1/√d, then applies sigmoid(leaky_relu(·)). The result is a spatial gate of shape [E, 1, H, W]: one coefficient per spatial location per edge. Crucially this is a sigmoid gate, not a softmax over neighbours — it modulates each message independently rather than forming a normalised competition among a node's neighbours. It is a reasonable choice when you want spatially-resolved gating and can accept that the coefficients do not sum to one; it is not a drop-in replacement for the normalised GAT semantics above. Knowing which one you are using prevents a subtle misreading of your model.

Bounding memory with chunked attention

Attention materialises per-edge intermediates, which is where memory pressure comes from on dense graphs. TensorGATLayer.forward() accepts a chunk_size=K argument that processes edges in groups: peak memory for intermediates scales as O(K × heads × C_head × spatial) instead of O(E × …), and the documented guarantee is that the chunked output matches the unchunked output within floating-point tolerance. examples/gat_chunking_demo.py demonstrates the pattern. Two honest notes: chunking trades time for memory rather than making attention cheaper overall, and you should profile your own graph before assuming a given chunk_size is necessary — the framework does not ship a benchmark that would let anyone claim a particular setting is optimal in general.

Edge features as an attention bias

As covered in the edges article, a vector edge feature enters TensorGATLayer as an additive bias on the attention logit per (edge, head); spatial edge tensors are mean-pooled to vectors first. This lets the relation between two nodes nudge how much attention flows along that edge — for example, a typed or geometric edge feature making some neighbours more relevant — without changing the softmax normalisation that keeps the weights summing to one.

How to describe this accurately

A conservative summary: TGraphX provides multi-head graph attention for tensor-valued nodes via TensorGATLayer, with scalar and channel attention modes and softmax normalisation over incoming edges; per-pixel/voxel attention is planned but not shipped. A lighter AttentionMessagePassing layer offers a 1×1-conv sigmoid spatial gate. Every clause there is backed by tgraphx/layers/gat.py and tgraphx/layers/attention_message.py.

Related guides

• Inside Tensor Message Passing: The φ → AGG → ψ Pipeline
• Tensor-Valued Edges in TGraphX
• Tensor-Valued Nodes in Graph Neural Networks

Conclusion

TGraphX brings attention to feature-map nodes in two flavours: a properly normalised multi-head TensorGATLayer with scalar and channel modes, and a lightweight sigmoid spatial gate in AttentionMessagePassing. The framework is upfront that per-pixel attention is a future item, and it guards the memory-saving chunked path with a numerical-equivalence test. Reading the two layers side by side — softmax-normalised versus sigmoid-gated — is the difference between using attention correctly and misreporting what your model does.