het-gnn — heterogeneous-GNN traffic assignment (act three)

What. het-gnn is a lean variant of Liu & Meidani’s (2024) heterogeneous- graph-neural-network traffic assignment surrogate ([liu2024end], docs/REFERENCES.md): real road links and virtual OD links are two edge types over one node set, typed attention passes messages, and an edge MLP predicts each real link’s flow/capacity ratio. Unlike implicit-ue-nn, the equilibrium is NOT architectural — flow conservation enters only as a soft training-loss penalty — so the raw emission is a per-link regression that routes no demand exactly.

Why it is in the benchmark. It completes a gradient of feasibility mechanisms across the benchmark’s three learned models: learned-surrogate has no conservation and is censored; het-gnn has SOFT conservation and is censored raw, but recovers feasibility by an explicit decode; implicit-ue-nn has conservation BY CONSTRUCTION. See docs/design/adr-026-het-gnn.md for the full derivation, the honest sourcing (the canon paper is paywalled and was attributed unread; the formulation was cross-verified with zero discrepancies from the authors’ own preprint and NSF accepted manuscript), and every measured anchor.

Scope. The size-agnostic node-kernel’s exact permutation equivariance (A4), the raw-censored / decode-recovered two-checkpoint story (A2 — the central content of this notebook), the in-family training and conservation ablation, and the honest held-out headline on Sioux Falls.

How this notebook is graded

A notebook never claims a number it does not compute in that cell. Every scored quantity below — feasibility, the certified relative gap, the node-balance residual, the wmape flow errors, the permutation-equivariance diff — is recomputed live by the P1 Evaluator (or an explicit numeric check) from the flows the model emitted, in the cell where it is claimed. The certified- gap ordering against implicit-ue-nn is reported as a STABLE structural bound, never a tight inequality — the adr-026 review proved that particular ordering is not a CI invariant, and pinning it tightly would be exactly the kind of claim this benchmark exists to catch (README, Certified, not self-reported).

# Setup. `het-gnn` needs the optional `torch` extra (`pip install tabench[torch]`);
# the guard mirrors `src/tabench/models/__init__.py`'s own import-guard pattern.
#
# The inline backend is Agg-based: figures render headlessly into the notebook,
# so CI can execute tutorials without a display. NEVER matplotlib.use("Agg")
# in-kernel — it silently suppresses inline figure capture.
%matplotlib inline
try:
    import torch  # noqa: F401  (the optional torch extra; absence -> ModuleNotFoundError)
except ModuleNotFoundError as exc:
    if exc.name != "torch":
        raise
    raise ModuleNotFoundError(
        "het-gnn needs the optional 'torch' extra: pip install tabench[torch]"
    ) from exc

import numpy as np

from tabench import (
    BiconjugateFrankWolfeModel,
    Budget,
    Evaluator,
    LearnedSurrogateModel,
    RngBundle,
    Trace,
    braess_scenario,
    load_scenario,
    viz,
)
from tabench.core.scenario import Demand, Network, Scenario
from tabench.metrics.gaps import node_balance_residual
from tabench.models import het_gnn as M
from tabench.models._paths import PathEngine
from tabench.models.het_gnn import HetGNNModel
from tabench.models.implicit_ue import ImplicitUENNModel
from tabench.models.learned import _random_network_scenario

A gradient of feasibility mechanisms, and a decisive coincidence

model

conservation

raw emission

feasible score

learned-surrogate

none

censored

never

het-gnn

soft loss (w_c=0.05)

censored

by an explicit decode

implicit-ue-nn

architectural

feasible

by construction

The paper’s own conservation metric L~_c is the L1/D form of the harness’s censoring statistic node_balance_residual — the paper trains toward exactly the quantity the audit thresholds. Its own best reported values still land 3–5 orders of magnitude above the 1e-6*D feasibility tolerance, so the paper-faithful raw emission is censored with certainty, trained or untrained. That censored row is the honest headline; the decode (a REPO EXTENSION, not in the paper) is what puts a certified gap on the leaderboard.

scenario = braess_scenario()
print(f"scenario      : {scenario.name}")
print(f"content hash  : {scenario.content_hash()[:16]}…")
print(f"total demand  : {scenario.demand.total}")

scenario      : braess
content hash  : cf00f411cdccec88…
total demand  : 6.0

A4 — the size-agnostic node-kernel is EXACTLY permutation equivariant

The paper’s node feature is that node’s entire OD-demand row plus two geographic coordinates — machine-verified to change by 21.5 under a consistent node relabeling (NOT permutation equivariant), and its dimension is |V|, forcing transfer learning or dummy-node padding to change graph size. The lean substitution here is the intrinsic per-node [production, attraction, out_degree, in_degree]: exactly equivariant, so ONE trained model runs on every graph size — link-kernel in implicit-ue-nn, node-kernel here, both size-agnostic by different routes. We also check the mirror of implicit-ue-nn’s A4: at RANDOM untrained weights, the DECODE is demand-feasible by construction (the raw regression is not).

def _permute_scenario(base: Scenario, perm: np.ndarray) -> Scenario:
    # Relabel node u -> perm[u] (0-based), edge ORDER fixed: edge a is the same
    # road segment with relabeled endpoints in both scenarios.
    net = base.network
    nz = net.n_zones
    od = base.demand.matrix
    od_p = np.zeros_like(od)
    for i in range(nz):
        for j in range(nz):
            od_p[perm[i], perm[j]] = od[i, j]
    net_p = Network(
        name="perm", n_nodes=net.n_nodes, n_zones=nz, first_thru_node=net.first_thru_node,
        init_node=perm[net.init_node - 1] + 1, term_node=perm[net.term_node - 1] + 1,
        capacity=net.capacity, length=net.length, free_flow_time=net.free_flow_time,
        b=net.b, power=net.power, toll=net.toll, link_type=net.link_type,
    )
    return Scenario(name="perm", network=net_p, demand=Demand(od_p), family="fuzz-perm")


base = _random_network_scenario(2, 10, 4, 6)
n, nz = base.network.n_nodes, base.network.n_zones
rng = np.random.default_rng(0)
perm = np.arange(n)
perm[:nz] = rng.permutation(nz)
perm[nz:] = nz + rng.permutation(n - nz)
sc_perm = _permute_scenario(base, perm)

torch.manual_seed(7)
untrained_gnn = M._HetGNN()  # equivariance is architectural, holds at any fixed weights
with torch.no_grad():
    a = untrained_gnn(M._het_graph(base.network, base.demand))
    a_p = untrained_gnn(M._het_graph(sc_perm.network, sc_perm.demand))
max_diff = float((a - a_p).abs().max())
print(f"max |raw ratio(base) - raw ratio(permuted)| : {max_diff:.2e}  (paper's featurization: 21.5)")
assert max_diff < 1e-8

# The decode mirror of implicit-ue-nn's A4: random untrained weights, feasible decode.
engine = PathEngine(scenario.network)
rs = M._build_routes(scenario.network, scenario.demand, engine, M._N_CG)
with torch.no_grad():
    v_raw_untrained = torch.clamp(untrained_gnn(M._het_graph(scenario.network, scenario.demand)), min=0.0)
    v_raw_untrained = v_raw_untrained * M._het_graph(scenario.network, scenario.demand)["cap"]
h_untrained, _, _ = M._decode(rs, v_raw_untrained, M._N_DECODE)
v_dec_untrained = (rs.delta.t() @ h_untrained).numpy()
print(f"decoded (random weights) node_balance_residual : {node_balance_residual(scenario, v_dec_untrained):.2e}")
assert node_balance_residual(scenario, v_dec_untrained) < 1e-9
assert Evaluator(scenario).evaluate(v_dec_untrained)["feasible"] == 1.0

max |raw ratio(base) - raw ratio(permuted)| : 2.78e-16  (paper's featurization: 21.5)
decoded (random weights) node_balance_residual : 8.88e-16

A2 — the two-checkpoint story: raw censored, decode recovered

solve records TWO harness-certified checkpoints from ONE call, both recomputed by the same P1 Evaluator — nothing self-attested. On Sioux Falls (a disjoint TNTP scenario, a different topology AND congestion regime than the 8–14-node synthetic training family): the RAW emission needs no shortest path (sp_calls=0, a genuinely new budget point below the ridge’s 1) and is censored; the DECODE runs n_cg column-generation sweeps and earns a real certified gap.

sc = load_scenario("siouxfalls")
assert sc.family != M.TRAINING_FAMILY  # disjoint from the training family, by hash

trace = Trace()
HetGNNModel().solve(sc, Budget(iterations=M._N_DECODE), RngBundle(0), trace)
assert len(trace) == 2
raw, dec = trace.checkpoints[0], trace.checkpoints[1]

m_raw = Evaluator(sc).evaluate(raw.link_flows)
print(f"raw    : sp_calls={raw.coords.sp_calls}  iterations={raw.coords.iterations}  "
      f"feasible={m_raw['feasible']:.0f}  node_balance_residual={node_balance_residual(sc, raw.link_flows):.1f}")
assert raw.coords.sp_calls == 0 and raw.coords.iterations == 0
assert m_raw["feasible"] == 0.0
assert node_balance_residual(sc, raw.link_flows) > 1e-6 * sc.demand.total

m_dec = Evaluator(sc).evaluate(dec.link_flows)
print(f"decode : sp_calls={dec.coords.sp_calls}  feasible={m_dec['feasible']:.0f}  "
      f"relative_gap={m_dec['relative_gap']:.4f}")
assert dec.coords.sp_calls == M._N_CG
assert m_dec["feasible"] == 1.0
assert np.isfinite(m_dec["relative_gap"]) and 0.0 < m_dec["relative_gap"] < 1.0
assert dec.self_report["training_sp_calls"] > 0
assert "decode_residual" in dec.self_report

raw    : sp_calls=0  iterations=0  feasible=0  node_balance_residual=295.9
decode : sp_calls=6  feasible=1  relative_gap=0.2586

In-family training, and the conservation ablation

Two in-sample checks, scoped honestly to the training family (never claimed to transfer): Adam genuinely reduces the composite flow loss below an untrained head’s, and training WITH the soft conservation loss (w_c=0.05) yields a lower in-family aggregate node-balance residual than training WITHOUT it (w_c=0) — the paper’s contribution made measurable. This ablation does NOT transfer to the held-out net (adr-026, measured backwards on Sioux Falls) — the identifiability caveat again, pinned honestly rather than forced.

cases = M._training_cases()


def _family_loss(model) -> float:
    total = 0.0
    with torch.no_grad():
        for c in cases:
            total += float(((model(c["g"]) * c["cap"] - c["v_obs"]) ** 2).sum()) / c["scale"]
    return total


def _in_family_cons_residual(model) -> float:
    total = 0.0
    with torch.no_grad():
        for c in cases:
            f = model(c["g"]) * c["cap"]
            total += float(M._conservation_residual(f, c["g"], c["expected"])) / c["scale"]
    return total


trained, _ = M._train()  # cached; the default w_c=0.05 head
torch.manual_seed(M._TRAIN_SEED)
untrained = M._HetGNN()
print(f"family loss   trained={_family_loss(trained):.1f}  untrained={_family_loss(untrained):.1f}")
assert _family_loss(trained) < _family_loss(untrained)

head_wc = M._fit(cases, w_cons=M._W_CONS)
head_0 = M._fit(cases, w_cons=0.0)
res_wc, res_0 = _in_family_cons_residual(head_wc), _in_family_cons_residual(head_0)
print(f"in-family node-balance residual   w_c={M._W_CONS}: {res_wc:.1f}   w_c=0: {res_0:.1f}")
assert res_wc < res_0

family loss   trained=282.1  untrained=652.5
in-family node-balance residual   w_c=0.05: 6.4   w_c=0: 11.7

The honest held-out headline

Every axis NAMED and MEASURED, never a pre-committed flattering direction: the raw emission’s flow error is WORSE than the ridge surrogate’s (censored either way); the decode’s is BETTER (projecting onto the feasible polytope recovers accuracy); implicit-ue-nn and het-gnn are both feasible with gaps in a sane band at matched route sets (the adr-026 review proved the tight ordering between them is NOT a CI invariant — reported as a stable structural bound only); and a converged bfw wins the wall/convergence axis by orders, as always.

oracle = sc.reference.link_flows
wmape = lambda v: float(np.abs(v - oracle).sum() / np.abs(oracle).sum())

v_ridge = LearnedSurrogateModel().solve(sc, Budget(iterations=1), RngBundle(0), Trace()).final.link_flows
m_ridge = Evaluator(sc).evaluate(v_ridge)
assert m_ridge["feasible"] == 0.0  # a per-link regressor routes nobody either

print(f"wmape   raw={wmape(raw.link_flows):.4f}  decoded={wmape(dec.link_flows):.4f}  "
      f"ridge={wmape(v_ridge):.4f}")
assert wmape(raw.link_flows) > wmape(v_ridge)   # raw transfers WORSE than the ridge
assert wmape(dec.link_flows) < wmape(v_ridge)   # decoded transfers BETTER

v_impl = ImplicitUENNModel().solve(sc, Budget(iterations=3000), RngBundle(0), Trace()).final.link_flows
m_impl = Evaluator(sc).evaluate(v_impl)
print(f"certified gap   implicit-ue-nn={m_impl['relative_gap']:.4f}  het-gnn(decoded)={m_dec['relative_gap']:.4f}"
      "   (STABLE structural bound only — not a tight ordering, adr-026)")
assert m_impl["feasible"] == 1.0
assert 0.02 < m_impl["relative_gap"] < 0.6
assert 0.02 < m_dec["relative_gap"] < 0.6

v_bfw = BiconjugateFrankWolfeModel().solve(
    sc, Budget(iterations=300, target_relative_gap=1e-6), RngBundle(0), Trace()
).final.link_flows
m_bfw = Evaluator(sc).evaluate(v_bfw)
print(f"bfw (converged) : relative_gap={m_bfw['relative_gap']:.3e}  (the wall/convergence axis)")
assert m_bfw["relative_gap"] < 1e-4 < m_dec["relative_gap"]

wmape   raw=0.5603  decoded=0.1676  ridge=0.2820
certified gap   implicit-ue-nn=0.1684  het-gnn(decoded)=0.2586   (STABLE structural bound only — not a tight ordering, adr-026)
bfw (converged) : relative_gap=3.228e-06  (the wall/convergence axis)

Visualize

The certified artifact (the decoded checkpoint) is per-link flows on a road Network, the same shape every static model here emits, so tabench.viz applies directly (adr-035’s viz rule).

display(viz.plot_network_flows(sc.network, dec.link_flows))
display(viz.plot_flow_scatter(("bfw (converged)", v_bfw), {"het-gnn (decoded)": dec.link_flows}))
../../_images/c0cb74b4fccc9f83715e5ee41e4e901f526288f92ca41e3b86abe85ee127cba7.png ../../_images/9a19d0f2165b43946ca7b255ccc3a4fda1ec643134c68cae2099acb56ba4dee9.png

Takeaways & pointers

  • A decisive coincidence, not a bug in the paper. Its own conservation metric IS the harness’s censoring statistic, orders above tolerance even at its best reported values — the raw emission is censored with certainty, which is why the decode (a REPO EXTENSION) exists at all.

  • Two checkpoints, nothing self-attested. The runner certifies both from the emitted flows; the paper-faithful censored row stays visible in the same trace as the certified decode.

  • Size-agnostic by a different route than implicit-ue-nn. Link-kernel there, node-kernel here — both exactly equivariant, both transfer across graph sizes with no retraining.

  • The certified-gap ordering against implicit-ue-nn is provenance, not a claim. It flips under retraining and weight perturbation (adr-026); this notebook asserts only the stable structure.

  • Where next. implicit-ue-nn (01-implicit-ue-nn.ipynb) for the architectural (not decoded) feasibility mechanism; the full derivation and every measured anchor in docs/design/adr-026-het-gnn.md.