Skip to content

Directed diffusion experiment

This experiment asks whether the existing sparse graph operations can express source-normalized propagation in both directions of one directed graph.

Decision

At tinygrad revision dd16d5a, Tinymesh composed bidirectional diffusion without a new kernel or Graph operation:

positive affinity[E]
       |
       +--> outgoing sum by source --> forward weight[E] --> G.sum
       |
       +--> outgoing sum in G.T ----> reverse weight[E] --> G.T.sum

tinymesh.nn.DirectedDiffusion owns the reverse graph and two normalized edge fields. This experiment retains the host reference and gradient witness. A fixed recurrent caller realizes those fields once before reuse.

Operator

For every original edge e: u -> v:

p_forward[e] = a[e] / sum(a[k] for k with source[k] = u)
p_reverse[e] = a[e] / sum(a[k] for k with target[k] = v)

forward(X)[v] += p_forward[e] X[u]
reverse(X)[u] += p_reverse[e] X[v]

The second denominator is the outgoing degree of v in the reversed graph. The reverse Graph(nodes, target, source) keeps the original COO edge order, so both normalized fields preserve one identity.

Each edge has a positive affinity and therefore belongs to a row with a positive denominator. Isolated nodes own no edge, require no division convention, and return zero.

DirectedDiffusion validates shape, floating dtype, and device without realizing the affinity tensor. The caller or data boundary owns the positive value invariant.

Callers may subtract the root and concatenate directions as [forward(X) - X, reverse(X) - X]. That is ordinary tensor composition and remains model-owned: the operator does not need another method to express it.

Composition

The current public operations are sufficient:

ones[N,1]
    |
    +--> G.T.sum(ones, a) --> outgoing[N,1]
    |                              |
    |                       edge_values(source)
    |                              |
    |                       a / degree --> p_forward[E]
    |
    +--> G.sum(ones, a) ---> incoming[N,1]
                                   |
                            edge_values(target)
                                   |
                            a / degree --> p_reverse[E]

Preparing the fixed weights performs two sparse sums and two endpoint projections once. Realizing them marks the boundary between differentiable affinity construction and a fixed recurrent cache. Each later application performs exactly two sparse sums. Leading batch axes reuse the same topology and weights.

Exact witness

An independent Python edge loop computes both normalized fields, both directions, and gradients with respect to node values and raw affinity. The tinygrad matches that reference within 1e-5 on CPU and Metal.

The fixture contains asymmetric degrees, duplicate edges, and two isolated nodes. Reordering COO edges reorders the two weight fields and affinity gradients while preserving node outputs and node gradients.

Sparse work

After preparation, the only model-time intermediates are the two [N,H] outputs and fixed [E] weight fields:

forward       one csr_sum
reverse       one csr_sum
topology      O(N + E)
node fields   O(NH)
edge fields   O(E)

UOp inspection rejects [N,N] and [N,E] carriers. This is a structural complexity claim, not a performance claim.

Reference boundary

The operator follows the forward and reverse random-walk idea in DCRNN. The pinned PyG Temporal implementation constructs a dense adjacency before propagation. Tinymesh keeps both directions sparse and removes the two redundant learned local blocks; this experiment does not claim framework parity.

Limits

The result covers fixed directed topology, positive scalar affinity, shared graph batch axes, first-order gradients, and one-step diffusion. It does not cover learned affinity, zero or negative affinity, higher-order walks, changing topology, vector edge messages, or recurrent model quality.

Reproduce

DEV=CPU uv run --locked python -m unittest tests.test_directed_diffusion
DEV=METAL uv run --locked python -m unittest tests.test_directed_diffusion
uv run --locked python -m experiments.run directed_diffusion DEV=CPU
uv run --locked python -m experiments.run directed_diffusion DEV=METAL