When Flat Vector Node Features Are Insufficient for Your GNN Task

The Day 1 article introduced tensor-valued nodes in TGraphX. This one is a decision guide: how do you tell whether your task actually benefits from richer feature structure, or whether you should stick with flat vectors and PyG?

The question matters because tensor-valued nodes are not free. They cost memory, layer complexity, and a smaller framework ecosystem. The benefit only materializes for certain kinds of tasks.

When flat vectors are sufficient

Flat node feature vectors [N, D] work very well for:

• Citation networks. Nodes are papers; features are bag-of-words or learned embeddings. The natural representation is vector.
• Social networks. Nodes are users; features are demographic or behavioral attributes. Vector.
• Molecular property prediction. Atom features are typed (carbon, hydrogen, etc.) with a few scalar properties. Vector.
• Recommendation graphs. Nodes are users/items; features are learned embeddings or content vectors. Vector.
• Knowledge graphs. Entities have learned embeddings, optionally augmented with text features. Mostly vector.

For these, PyG with flat features is the right choice. Do not introduce complexity that does not help.

When flat vectors lose information

Flat vectors lose information when the natural feature has internal structure:

Spatial structure. An image patch as a node feature is [C, H, W]. Flattening to [C·H·W] discards the spatial layout — pixel (3, 4) is no longer geometrically adjacent to pixel (3, 5) in the feature vector. Any convolutional inductive bias has to be relearned from scratch.

Temporal structure. A time series window per node is [T, D]. Flattening to [T·D] loses temporal order from the model's perspective. RNN-style or convolutional temporal layers cannot be applied inside the GNN.

Volumetric structure. A 3D medical scan region per node is [C, D, H, W]. Flattening is even more lossy than 2D.

Multi-channel structure. A node with multiple aligned feature maps (e.g., RGB and depth) needs to preserve the channel dimension if you want channel-mixing operations later.

If your node feature has any of these structures, flattening can hurt model quality.

A concrete test

Suppose you have image-patch graphs and you are debating tensor-valued nodes. Try this:

1. Build a baseline with PyG: flatten patches to vectors, train a GCN or GAT.
2. Build a TGraphX version: keep patches as tensors, use tensor_gcn.
3. Train both for the same compute budget on the same data.
4. Compare accuracy.

If the TGraphX version is meaningfully better, the structure matters. If they are within noise, flat vectors are fine and you should keep using PyG for the ecosystem benefits.

Code example: the same task two ways

import torch
        import tgraphx as tgx
        
        # Suppose: 1000 image patches, each [3, 8, 8], with 10 classes
        N = 1000
        x_tensor = torch.randn(N, 3, 8, 8)
        edge_index = tgx.knn_graph(x_tensor.view(N, -1), k=8)
        y = torch.randint(0, 10, (N,))
        
        # Path A: flatten and use a vector layer
        x_flat = x_tensor.view(N, -1)         # [N, 192]
        result_flat = tgx.classify_nodes(
            x=x_flat, edge_index=edge_index, labels=y,
            model="gcn", seed=42,
        )
        
        # Path B: keep tensor structure
        result_tensor = tgx.classify_nodes(
            x=x_tensor, edge_index=edge_index, labels=y,
            model="tensor_gcn", seed=42,
        )
        
        print(f"Flat:   {result_flat.metrics['val_accuracy']:.3f}")
        print(f"Tensor: {result_tensor.metrics['val_accuracy']:.3f}")
        

A measurable gap (more than the typical seed-to-seed variance) in favor of tensor justifies the additional complexity. If both are similar, prefer the flat version.

Tasks where the benefit usually exists

• Image-patch graphs for vision tasks where patches retain spatial coherence.
• Medical imaging with volumetric node features.
• Sequence-augmented citation graphs where each paper carries a windowed text embedding.
• Robotics graphs where nodes carry sensor tensor readings.

Tasks where the benefit usually does not exist

• Cora, Citeseer, PubMed (citation) — flat vectors are the right model.
• Molecular property prediction (atom features are naturally vector).
• Recommendation graphs (learned embeddings, not raw features).
• Standard KG completion benchmarks (entity ID embeddings).

Decision rule

If you cannot articulate what structure flattening discards, the flat version is the right answer. Use PyG.

If you can articulate it — "patches lose spatial layout," "time windows lose temporal order" — try both and let the measurement decide. The framework cost is not free, but for the right problem it pays for itself.

On the dual approach

Some projects use both. PyG layers for the parts of the workflow where flat features suffice (e.g., a graph-level pooling step) and TGraphX Graph for the data layer that preserves tensor structure. The two frameworks can coexist; what matters is that you are intentional about where the structure is preserved and where it is collapsed.

FAQ

Q: Is there a quick way to estimate whether structure helps?
A: Train a model that ignores graph structure entirely (e.g., a per-node CNN for image patches) and a model that uses graph structure with flat features. If the per-node CNN is much better, the structure inside the node matters. Combine both with TGraphX.

Q: Does TGraphX work for molecular graphs?
A: Yes, but most molecular GNN benchmarks use flat atom features and PyG has more pre-built layers for this case. Use TGraphX for molecular graphs only if you have rich atom-level tensor features.

Q: What about node features that are short sequences (T=5, D=8)?
A: That is a borderline case. Try both. For very short sequences, flattening loses little and is simpler.

Q: How do I convert between tensor and flat representations safely?
A: tgx.Graph accepts both. To flatten: x_flat = g.node_features.view(g.num_nodes, -1). To check rank: tgx.assert_tensor_native(g, min_rank=3).