Skip to content

Controlled node-time JEPA

This stage asks whether latent future prediction improves a reusable representation of a known spatiotemporal process—and whether the correct joint mesh matters. It is a controlled mechanism test, not a traffic benchmark.

Data and causal boundary

The existing controlled transport generator evolves independent zero-mean fields over one directed 24-node graph:

x[t+1] = 0.25 x[t] + 0.55 forward(x[t]) + 0.20 reverse(x[t])

context                   hidden target
x[0] x[1] x[2] x[3]  ->  x[4] x[5] x[6] x[7]

Training, validation, and test contain 128, 32, and 32 independent trajectories from the already pinned data seeds. Self-supervision sees training trajectories only. The downstream probe sees training targets only; validation decides and test confirms.

Product mesh

Each four-step block lowers through time.cartesian(space):

node token       (t, v) -> t * 24 + v

relation 0       (t-1, v) -> (t, v)             temporal
relation 1       (t, u)   -> (t, v)             forward affinity
relation 2       (t, v)   -> (t, u)             reverse affinity

Forward and reverse edge values use the same source- and destination-normalized coefficients as the stated transport law. They multiply source messages; they are never added as standalone features. A zero field therefore produces zero state, preventing static topology from becoming the prediction shortcut.

Arm Time edges Space edges Product edges
true 72 288 true 360
permuted 72 288 degree-matched false 360
temporal 72 0 72
spatial 0 288 true 288

temporal tests node-local history. spatial gives the frozen probe every time token but forbids messages between them. The matched four arms distinguish correct space, false space, no space, and no learned time edge.

Objective and frozen evaluation

context product -> online 3xMeshConv -> LayerNorm -> predictor -- L1 --+
                                                                    |
future product  -> EMA target 3xMeshConv -> LayerNorm -> stopgrad ---+

MeshConv is a research-only relation-weighted sum composed from Graph.edge_values and Graph.sum_edges. Three layers span the four-step path. The final token-wise LayerNorm removes scale as an objective shortcut. The online encoder and predictor receive gradients; the full target encoder starts identically and follows the online encoder with EMA 0.998.

All four context token rows remain in the frozen node representation. One shared linear head sees [node, 4 * hidden] and predicts the four future field values. Its inputs use training-only mean and standard deviation. Every arm is compared with its identical target encoder before pretraining and uses the same probe initialization and budget. Persistence and a standardized raw-history probe are non-neural references.

V-JEPA v1 grounds latent L1 prediction, an EMA target, masked spatiotemporal tokens, and frozen evaluation. This experiment deliberately uses its causal mask ablation rather than the paper's stronger non-causal multi-block mask: the whole future block is unavailable to the context encoder. It replaces video attention with sparse relation messages and is not a V-JEPA reproduction.

Registered decision

Seed 17 was used before registration to verify execution, eliminate a final ReLU collapse, reject additive topology-only messages, retain all token rows, and select V-JEPA's slower target regime. It is excluded from evidence below. Seeds 0, 1, and 2 are untouched.

The representation earns continuation only if true improves validation RMSE over its own random encoder on mean and in at least two of three seeds, with zero target gradient and final embedding variation at least half its initial value. The explicit joint mesh is supported only if trained true also beats trained permuted, temporal, and spatial in at least two of three paired validation runs. Test repeats the validation decision; it does not select one. No seed, arm, budget, feature, or threshold changes after clean formal results are visible.

Decision

At revision d7dadea, three Metal runs produced the registered seed-level results below. RMSE values are mean and population standard deviation across seeds. Gain is each trained target encoder's random-initialization RMSE minus its trained RMSE, so positive is better.

Arm Random validation Trained validation Gain Seed wins Trained test Test gain
true 0.047428 ± 0.003955 0.046152 ± 0.004810 +0.001276 ± 0.001287 3/3 0.047617 ± 0.005229 +0.001617 ± 0.001282
permuted 0.052101 ± 0.005740 0.053413 ± 0.006710 -0.001312 ± 0.001507 0.053861 ± 0.005884 -0.000783 ± 0.000806
temporal 0.079379 ± 0.002475 0.078719 ± 0.003006 +0.000661 ± 0.000939 0.080902 ± 0.003426 +0.000530 ± 0.000928
spatial 0.044436 ± 0.004942 0.043257 ± 0.002213 +0.001179 ± 0.003353 0.043544 ± 0.001709 +0.000794 ± 0.002884

Persistence validation RMSE was 0.097644; the standardized raw-history probe reached 0.055889 ± 0.000010. The true representation beat both in every seed. Its validation gains were +0.000166, +0.003081, and +0.000580, and all three test gains remained positive. Mean latent loss fell from 0.894219 to 0.656218. Target gradients were exactly zero, the EMA target moved by 2.840953 ± 0.230082, and embedding variation changed from 0.961556 ± 0.003974 to 0.962472 ± 0.004822 rather than collapsing.

The representation gate passes. Causal latent prediction consistently improves the true product encoder over its identical random initialization, and test repeats the validation direction.

The joint-mesh gate fails. Trained true beat permuted and temporal in all three validation seeds by mean RMSE margins of 0.007260 and 0.032566, but beat spatial in only one. spatial was better on mean validation RMSE by 0.002895 and on mean test RMSE by 0.004073, again winning two seeds. The probe already sees all four spatially encoded history rows, so these results do not justify learned temporal edges on this task.

Continue causal JEPA as a research mechanism, with spatial as the architecture incumbent and the full product as a control. Do not promote Mesh, MeshConv, encoders, predictor, objective, data policy, probes, or controls. Graph.cartesian, batched endpoint gathering, and batched edge sums already own the reusable contracts.

Reproduce

uv run --locked python -m experiments.run transport_jepa DEV=METAL EMA=0.998 HIDDEN=8 HISTORY=4 HORIZON=4 LR=0.001 PROBE_LR=0.05 PROBE_STEPS=150 SEED=0 STEPS=100
uv run --locked python -m experiments.run transport_jepa DEV=METAL EMA=0.998 HIDDEN=8 HISTORY=4 HORIZON=4 LR=0.001 PROBE_LR=0.05 PROBE_STEPS=150 SEED=1 STEPS=100
uv run --locked python -m experiments.run transport_jepa DEV=METAL EMA=0.998 HIDDEN=8 HISTORY=4 HORIZON=4 LR=0.001 PROBE_LR=0.05 PROBE_STEPS=150 SEED=2 STEPS=100