Rotate Before You Update: Deeper Test-Time Adaptation for Graph Models

TL;DR for operators

A deployed graph model may encounter a new network, molecule population, or sensor regime and still need to predict before any verified label exists. The operational question is whether it should remain fixed or make a tightly bounded adjustment using only information available in the incoming unlabeled sample.

T3R offers the second option. Across nine DiTEC-WDN benchmarks, one update reduced average RMSE from 0.5023 under ERM to 0.3095, a 38.38% reduction, although it did not lead the average NSE metric. The method learns during training how an auxiliary objective can provide proxy updates for deeper prediction layers when the true supervised correction is unavailable.

That result is not permission for unrestricted live learning. The same update can help substantially, do almost nothing, or damage performance depending on which auxiliary objective supplies the gradient and how long adaptation continues. Operators therefore need a policy covering whether live updates are allowed, which signal is trusted, how many updates may occur, and whether the accuracy gain justifies the added inference latency.

A model must decide before the new labels exist

A graph model deployed on a water network, molecule collection, or communications topology can encounter a changed environment before anyone can verify the correct output. Pipes are reconfigured, sensors drift, or the target molecules differ from those used in training. The prediction is needed now; labels for ordinary fine-tuning may arrive much later.

The incoming graph still contains structure. A model can mask part of it and reconstruct the missing information, creating a learning signal without a target label. Test-time training uses that auxiliary objective to adjust the model on the test sample before making the main prediction.

Conventional versions usually update the shared encoder and leave the main decoder fixed. T3R, proposed by Huy Truong, Alexander Lazovik, and Victoria Degeler, asks whether the auxiliary signal can be redirected deeper into the prediction branch without pretending that the missing supervised label is available.¹

Across nine DiTEC-WDN regression benchmarks, one T3R step produced the best reported average $R^2$, RMSE, Pearson correlation, and MAE among ERM, joint training, Tent, conventional TTT, and TTT with one RotoGrad layer. Average RMSE fell from 0.5023 under ERM to 0.3095. Nash-Sutcliffe efficiency was the exception, and pooled variability remained substantial.

The same adaptation process can still improve, barely move, or damage performance depending on the auxiliary task and update count.

The technical contribution is a learned route for proxy updates

T3R uses a Y-shaped model: a shared graph encoder feeds a supervised main decoder and a self-supervised decoder. During joint training, both branches have real losses, so their gradients can be observed together.

Two objectives may push shared parameters in compatible or conflicting directions. The paper measures compatibility with cosine similarity:

Here, $\mathbf{g}_m$ is the main-task gradient and $\mathbf{g}_s$ is the self-supervised gradient. A positive value means the updates point in broadly similar directions. The paper treats this as an empirical task-affinity signal, not proof that one objective causes better adaptation.

T3R then inserts orthogonal rotation matrices into both decoder branches and trains them layer by layer so that the two task gradients rotate toward a shared direction. At inference, the main-task gradient is unavailable because the label is unknown. The self-supervised gradient remains observable. T3R transforms it out of the auxiliary orientation and into the learned main-task orientation, yielding a surrogate gradient for the next main-decoder layer.

A useful mental model is a calibrated coordinate adapter: training establishes how two update directions relate, and inference reuses that mapping when only one can be measured.

This extends adaptation beyond the encoder, but not to every parameter. The main rotation matrices and final classifier stay fixed. “Deeper” describes where proxy updates can reach, not an unrestricted rewrite of the model.

Auxiliary self-supervision must be qualified

The auxiliary-task experiment is mechanism evidence rather than the main benchmark. Its purpose is to test whether gradient alignment can identify a useful signal for adaptation.

Node masking stood apart. Its training gradient alignment was 0.12001, and one conventional TTT step reduced test MAE from 5.8103 to 3.9924, a 31.29% gain. Node-embedding and graph-embedding hinge losses had smaller positive alignments and gains of 5.85% and 4.10%. Several objectives had near-zero or negative alignment and delivered negligible or negative changes; edge masking worsened MAE by 4.71%.

This does not establish a universal safety threshold. It shows that choosing self-supervision by familiarity or implementation convenience is weak deployment practice. In these experiments, the auxiliary gradient’s direction was associated with whether one update helped.

Cognaptus inference: teams should treat the auxiliary objective as a qualified production dependency. Qualification should include gradient-alignment measurement during training, one-step and multi-step validation under representative shifts, and rejection criteria for unstable tasks.

The benchmarks support depth, but not unlimited optimization

The water-network average is the clearest main result. With one step, T3R reached average $R^2$ of 0.9492, RMSE of 0.3095, PCC of 0.9721, and MAE of 0.2196. It led on those four metrics across nine datasets, though not on average NSE. Per-dataset results were heterogeneous, so the average is not a guarantee for every topology.

The five-fold held-out-network analysis is a robustness and sensitivity test. It asks whether gains persist when the target network is excluded from training and the step count changes. T3R was strongest at shallow adaptation for RMSE and MAE, especially at one step. Additional steps often degraded regression error even where average $R^2$ and correlation improved.

The molecular experiments are a cross-domain extension. From ogbg-molbace to ogbg-molbbbp, T3R improved ROC-AUC from its own 0.4651 step-zero score to 0.5608 at 40 steps, a 20.57% relative gain. In the reverse direction, it rose from 0.3941 to a best 0.6682 at 10 steps, a 69.56% gain. These percentages use each T3R model’s own no-adaptation score as the denominator, not a common ERM baseline.

One direction favored 40 steps; the other peaked at 10 and then declined. The update count depends on the shift and cannot be inferred from adaptation depth.

Deployment is a policy decision

T3R changes four operational decisions.

First, decide whether live adaptation is permitted. A batch molecular-screening pipeline may tolerate several backward passes; a real-time infrastructure alarm may not.

Second, approve the auxiliary task before deployment. The system should not discover on live traffic that its self-supervised objective pushes the model in an unhelpful direction.

Third, separate adaptation depth from duration. T3R’s value comes from allowing proxy updates to reach more of the decoder. A one-step policy may capture much of the value in latency-sensitive regression, while larger budgets may be reserved for high-value batch cases.

Fourth, retain a rollback boundary. Teams may restore pretrained weights between samples, adapt within a bounded session, or maintain audited state. The paper evaluates benchmark adaptation, not a long-running production stream with delayed labels and cumulative updates.

The runtime numbers make this concrete. On the reported OGB setup, T3R took 4.05 seconds before adaptation, 1.35 times ERM’s 2.99 seconds. At one step it reached 5.24 seconds, or 1.75 times ERM. At 40 steps it required 28.27 seconds, a 9.45-fold slowdown. The measurements come from one NVIDIA RTX 3060 Laptop GPU, but they establish that the overhead is structural rather than negligible.

What remains uncertain

The paper provides moderate benchmark evidence, not a causal or theoretical guarantee. The link between positive gradient alignment and successful adaptation remains empirical. The surrogate-gradient argument also assumes sufficiently converged rotations and an auxiliary gradient that remains a reasonable guide at test time.

Coverage is limited to two domains, selected GAT and GIN backbones, simulated or cross-dataset shifts, and a narrow auxiliary-task set. There is no evidence here about persistent online adaptation with delayed label audits, adversarially corrupted graphs, or shifts that break the auxiliary objective.

The next validation burden is concrete: test the auxiliary task under realistic shifts, choose a bounded stopping policy, measure latency on the production stack, and compare the gain against waiting for labeled retraining data.

Deeper adaptation needs tighter controls

T3R supplies a mechanism for constructing deeper updates from a signal whose compatibility was learned in advance. Its strongest evidence supports a controlled operating model: qualify the auxiliary objective, permit a bounded number of updates, keep selected parameters fixed, monitor the accuracy-latency curve, and preserve rollback.

Under those conditions, graph models may respond to topology, sensor, or dataset changes before new labels are available. The practical value will depend less on the maximum benchmark gain than on whether a team can define when the model may update, how far the update may travel, and when it must stop.

Cognaptus: Automate the Present, Incubate the Future.