Skip to content

T-GCN experiment

This record asks whether one immutable Graph can carry a trainable recurrent state across ordered node snapshots without a temporal kernel or container.

Decision

At tinygrad revision c9e1154, Tinymesh unrolls a T-GCN cell across fixed-topology snapshots and differentiates through space and time on CPU and Metal.

At that revision the result added no public API. The same checked equation now lives in tinymesh.nn.TGCN; one graph is reused, snapshots and hidden state are ordinary tinygrad tensors, and this experiment retains the temporal learning witness.

Why T-GCN first

PyTorch Geometric Temporal offers two adjacent gated designs:

  • T-GCN runs three GCN operations over the current node field, then combines each result with node-local hidden state (source).
  • GConvGRU graph-convolves both input and hidden state for all three gates, requiring six Chebyshev convolutions per step (source).

T-GCN is the smaller standard test of the temporal boundary because the existing experimental GCN caller is sufficient. The subsequent GConvGRU experiment adds hidden-state propagation and compares its parameter and sparse-call cost on the same temporal fixture.

The reference expresses its three input projections as separate GCN calls. They share node input, topology, and normalization and differ only in output weights, so linearity gives:

[GCN_z(X), GCN_r(X), GCN_h(X)] = GCN_[W_z | W_r | W_h](X)

At the recorded revision Tinymesh made one width-3H sparse call and sliced its output into the three gates. The public GCNConv now commutes the shared linear map after normalized aggregation: it makes one width-F sparse call, then one linear map to 3H. Both factorizations retain three independent weight blocks without repeating topology traversal.

Cell

For current node field X, previous hidden state H, and graph convolution GCN_G:

Z       = sigmoid(Linear_z([GCN_G,z(X), H]))
R       = sigmoid(Linear_r([GCN_G,r(X), H]))
H_tilde = tanh(Linear_h([GCN_G,h(X), R * H]))
H_next  = Z * H + (1 - Z) * H_tilde

Z decides how much prior state survives. R decides how much prior state enters the candidate. The candidate combines the current graph signal with reset hidden state.

X_t [N, F]
    |
 GCN [N, 3H]
    |
 split Z, R, candidate projections
    |              |              |
    +-- H_t-1      +-- H_t-1      +-- R * H_t-1
    |              |              |
 sigmoid Z      sigmoid R       tanh H_tilde
    |                             |
    +---------- gated update -----+
                  |
                  v
              H_t [N, H]

The graph caller uses (D^-1/2 A D^-1/2 X)W. Self-loops are explicit edges; the cell does not silently change topology.

Temporal data boundary

PyTorch Geometric Temporal's static signal container yields one ordinary graph snapshot at a time while reusing one edge index (source).

Tinymesh needs no container for this proof:

Graph G                    one fixed topology and cache owner
snapshot X_t [N, F]        one ordered node field
hidden H_t [N, H]          one recurrent node field

The first call initializes hidden state to zero. Later calls validate its node count, feature width, dtype, and device. Reusing the same Graph reuses its lowered CSR buffers.

Evidence

An independent host implementation evaluates all three gates over two snapshots. Tinymesh matches it on CPU and Metal. Reversing the snapshots changes the final state, which rejects order-insensitive aggregation. A UOp check sees one input-width csr_sum call in one cell step, enforcing the fused graph projection.

The learning witness is narrower:

graph:       0 -> 0, 1 -> 1, 0 -> 1
snapshot 0: X_0 = [1, 0]
snapshot 1: X_1 = [0, 0]
target:      final node-1 state = 1

Node 1 receives the signal only through edge 0 -> 1 at the first snapshot. The second snapshot contains no signal, so the prediction also requires temporal retention. One SGD step updates the candidate graph parameter:

initial loss                0.718740
candidate graph gradient   -0.188622
final loss                  0.686385
candidate graph weight      1.188622

Both backends return the same float32 values. This proves a parameter gradient crosses one spatial edge and one temporal transition. It does not establish forecast quality.

The formulation follows T-GCN.

Limits

The result covers one fixed graph, ordered snapshots without interval metadata, zero initial hidden state, explicit first-order unrolling, and one device. It does not cover timestamps, irregular intervals, missingness masks, changing edge values, changing topology, state detachment, truncated backpropagation, multi-step losses, or long-horizon stability.

The cell now also accepts shared-graph batch axes. The Chickenpox forecast records that execution and its limited predictive evidence; this toy witness still proves only the gradient path.

Reproduce

DEV=CPU uv run --locked python -m unittest tests.test_tgcn
DEV=METAL uv run --locked python -m unittest tests.test_tgcn
uv run --locked python -m experiments.run tgcn DEV=CPU
uv run --locked python -m experiments.run tgcn DEV=METAL