TGraphX › Insights › Tensor-Valued Edge Features in Graph Neural Networks
← Back to Insights
Tutorial · · 10 min read

Tensor-Valued Edge Features in Graph Neural Networks

Target keyword: tensor-valued edge features graph neural networks

Tensor-Valued Edge Features in Graph Neural Networks

Most GNN tutorials treat edge features as optional scalars — a weight, a distance, a type label. This works for many problems. But there are research domains where edges themselves carry structured multi-dimensional data: pair-wise image comparisons, spatial relationship maps between objects, volumetric connectivity descriptions in brain networks, or interaction tensors in molecular simulations.

Treating these structured edge features as flat vectors discards the same spatial structure that researchers are careful to preserve in node features. TGraphX's message-passing layers support tensor-valued edge features natively, in two distinct modes that preserve their spatial layout through aggregation.

Why Edge Features Are Different from Node Features

In standard message passing, the update for node v is:

h_v' = UPDATE( h_v, AGGREGATE( {MSG(h_u, e_uv) : u ∈ N(v)} ) )

The edge feature e_uv appears inside the message function MSG. For node features, the shape [N, ...] is easy to track because the node index is the first dimension throughout. For edge features, the challenge is that AGGREGATE discards the edge dimension — after aggregation, you have one value per destination node, not per edge.

This means edge feature tensors must be projected into node feature space before aggregation. The shape of the projection matters for structured edge features.

Edge Feature Modes in TGraphX

TGraphX's tensor-aware layers (TensorGINLayer, TensorGraphSAGELayer, TensorGATLayer) support two edge feature modes via the edge_features_kind argument:

Mode 1: "spatial" — structured edge features matching the node spatial layout.

Edge features have shape [E, edge_dim, *spatial] — for 2D: [E, edge_dim, H, W]. A 1×1 convolution projects edge_dim → in_channels, preserving H and W. The result is added to the source node's feature map before aggregation:

python
import torch
        from tgraphx.layers.gin import TensorGINLayer
        
        N, C, H, W = 50, 16, 8, 8
        E = 200
        
        layer = TensorGINLayer(
            in_channels=C,
            out_channels=32,
            use_edge_features=True,
            edge_dim=8,                     # channels in edge feature tensor
            edge_features_kind="spatial",   # [E, 8, H, W] edge features
            spatial_rank=2,
        )
        
        x = torch.randn(N, C, H, W)
        edge_index = torch.randint(0, N, (2, E), dtype=torch.long)
        edge_feats = torch.randn(E, 8, H, W)  # spatial edge features
        
        out = layer(x, edge_index, edge_features=edge_feats)
        print(out.shape)  # [50, 32, 8, 8]

Mode 2: "vector" — flat edge features projected into spatial space.

Edge features have shape [E, edge_dim]. A linear projection maps edge_dim → in_channels, then the result is unsqueezed to [E, in_channels, 1, 1] and broadcast over the spatial grid:

python
layer_vec = TensorGINLayer(
            in_channels=C,
            out_channels=32,
            use_edge_features=True,
            edge_dim=4,
            edge_features_kind="vector",    # [E, 4] edge features
            spatial_rank=2,
        )
        
        edge_feats_vec = torch.randn(E, 4)  # vector edge features
        out2 = layer_vec(x, edge_index, edge_features=edge_feats_vec)
        print(out2.shape)  # [50, 32, 8, 8]

The distinction matters: in "vector" mode, the edge feature acts as a spatially uniform bias over the node feature map. In "spatial" mode, the edge feature is a full spatial map that interacts with the source node's spatial content.

SAGE and GAT with Edge Features

The same two modes apply to TensorGraphSAGELayer and TensorGATLayer:

python
from tgraphx.layers.sage import TensorGraphSAGELayer
        from tgraphx.layers.gat import TensorGATLayer
        
        # SAGE with spatial edge features
        sage_layer = TensorGraphSAGELayer(
            in_channels=16,
            out_channels=32,
            use_edge_features=True,
            edge_dim=8,
            edge_features_kind="spatial",
            spatial_rank=2,
        )
        
        # GAT with vector edge features (as attention bias)
        gat_layer = TensorGATLayer(
            in_channels=16,
            out_channels=32,
            num_heads=4,
            use_edge_features=True,
            edge_dim=4,
            edge_features_kind="vector",
            spatial_rank=2,
        )

In GAT, vector edge features are projected to a per-head attention bias, which modulates the attention score for each edge. This is consistent with the GINE/GAT-with-edge-bias literature.

Use Cases: When Tensor-Valued Edges Make Sense

Scene graph edges. In scene graphs, a node is an object (represented as an image patch [C, H, W]), and an edge captures the spatial relationship between two objects. This spatial relationship could itself be represented as a difference map, an oriented bounding-box overlap, or a learned interaction matrix — naturally a tensor.

Brain connectivity. Structural connectivity between brain regions can be described by diffusion tensor imaging (DTI) tensors, which are 3×3 symmetric positive-definite matrices per fiber tract. Treating these as [E, 3, 3] edge features (a special case of edge_dim=3, spatial_rank=1, H=3) captures the full anisotropy.

Molecular bond features. Bond type, bond order, and computed electronic density along a bond can form a rich edge representation beyond a simple scalar.

Knowledge graph relation embeddings. In multimodal KGs, relation triples can carry visual evidence — the edge from an "image" entity to a "object" entity might carry the visual context of the recognition event as an image tensor. See knowledge graph embedding with tensor features for the KG-specific treatment.

Constructing a Graph with Edge Features

python
from tgraphx import Graph
        
        N, C, H, W = 50, 16, 8, 8
        E = 200
        
        x = torch.randn(N, C, H, W)
        edge_index = torch.randint(0, N, (2, E), dtype=torch.long)
        edge_feats = torch.randn(E, 8, H, W)  # [E, edge_dim, H, W]
        
        g = Graph(
            node_features=x,
            edge_index=edge_index,
            edge_features=edge_feats,   # stored as g.edge_features
        )
        
        print(g.node_features.shape)   # [50, 16, 8, 8]
        print(g.edge_features.shape)   # [200, 8, 8, 8]
        
        # PyG-style alias
        print(g.edge_attr.shape)       # [200, 8, 8, 8] — same tensor

The Graph object validates that edge_features.shape[0] == edge_index.shape[1] at construction time. Passing a mismatched edge feature tensor raises ValueError with shape details.

Edge Features in Mini-Batch Sampling

When using NeighborLoader or GraphSAINTLoader with edge features, TGraphX samples both the subgraph edges and their corresponding edge features:

python
from tgraphx import NeighborLoader
        
        loader = NeighborLoader(g, fanouts=[10, 5], batch_size=16, seed=42)
        
        for batch in loader:
            # batch.edge_features has the correct subset of edge feature tensors
            # aligned with batch.edge_index
            out = layer(batch.node_features, batch.edge_index,
                        edge_features=batch.edge_features)

Edge features with E_batch entries in the mini-batch correspond exactly to the E_batch edges in batch.edge_index. No manual index alignment is needed.

The Graph Object: Edge Features vs Edge Weight vs Edge Labels

TGraphX distinguishes three separate edge-level tensors:

Attribute Shape Purpose
edge_features [E, ...] Per-edge input features for message passing
edge_weight [E] Scalar multiplier applied after aggregation
edge_labels [E, ...] Supervision targets for edge-level prediction

These are separate attributes and cannot be substituted for each other. A common mistake is passing classification targets as edge_features — TGraphX stores them separately and does not silently mix them.

Limitations

Memory scaling. Spatial edge features scale as E × edge_dim × H × W. For a dense graph with 10K edges and [8, 8, 8] edge features, this is 5.2M float32 values (~20 MB). For sparse graphs this is manageable; for dense graphs or large spatial dims, memory becomes the dominant constraint. Use chunk_size in TensorGINLayer.forward() to process edges in chunks.

The "vector" mode broadcasts uniformly over spatial dims. This is intentional but means the edge feature does not interact locally with spatial positions in the node feature map. For use cases where spatial position within the node feature should interact with the edge content, "spatial" mode is more expressive but requires edge features to match node spatial dims.

No built-in edge feature pooling. TGraphX does not provide a dedicated edge-feature pooling layer for graph-level edge summarization. Edge features participate in node-level message passing but are not separately pooled at the graph level.

Related Articles