Universal Test-Time Training

1 INTRODUCTION

• Layer-private TTT confines each layer’s memory to its own writes, preventing layers from reading state written by other layers.This applies to layerwise KV-binding TTT, LaCT, and TTT-E2E, which keeps separate weight tensors across depth.Source: §1 · p. 2
• uTTT shares one fast-weight state or pool across layers while keeping backbone parameters layer-specific, making layer-private TTT a partitioned special case.Its memory recurs across chunks and layers, allowing a shallow layer to read a deep layer’s previous-chunk write.Source: §1 · p. 2
• Unlike architectures that share slow parameters across depth, uTTT shares bounded fast-weight memory updated by test-time optimization within each sequence.Prior shared-depth approaches include Universal Transformers, ALBERT, and shared trained expert parameters.Source: §1 · p. 2
• uTTT-MoE routes each token head to K of E shared experts, while uTTT-Dense applies one shared state at every layer.The layer-private TTT-MoE control uses E/L experts per layer; shared writes also merge into one operation per chunk, speeding training.Source: §1 · p. 2
• 2.6 and 2.1 RULER points over TTT-MoE at 124M and 760M; sharing adds 0.92 and 0.76 dB object PSNR in routed and dense view synthesis.uTTT-MoE also trains up to 1.39× faster than TTT-MoE.Source: §1 · p. 2
• Interventions that remove cross-layer reads or cut cross-depth gradients degrade performance, indicating trained models use routes enabled by sharing.Analysis treats layer-private TTT as a partitioned form of the shared time-and-depth memory recurrence.Source: §1 · pp. 2–3

2 PRELIMINARY

• Sequences are processed as contiguous chunks in causal order, with layer-head projections producing query, key, and value vectors for each token.The fast-weight state is a collection of SwiGLU matrix triples; reads provide prediction features without changing the state.Source: §2 · p. 3
• A read evaluates the fast-weight function on a query, while a write differentiates an associative loss that aligns the function’s key output with its value.Chunk-level writes aggregate token contributions at boundaries, amortizing update costs.Source: §2 · p. 3
• Fast weights reset to their learned initialization between sequences, and the final language-modeling chunk needs no write.Effective step sizes come from slow-parameter predictors; language models use momentum, Muon, and row renormalization, while NVS omits momentum and Muon.Source: §2 · p. 3
• During training, task-loss gradients backpropagate through earlier writes to train slow projections, coefficient predictors, and initialization; inference updates only sequence-specific state.The associative loss determines online writes, whereas the task loss supplies the outer learning signal.Source: §2 · p. 3

3 UNIVERSAL FAST WEIGHTS

• uTTT shares fast-weight state across depth while keeping backbone weights, projections, and local attention layer-specific.The shared state evolves over chunks and layers rather than maintaining a separate memory for each layer.Source: §3 · p. 3

3.1 SHARED FAST WEIGHTS

• Layers are grouped to share state; at p = L groups are layer-private, while at p = 1 all layers share one state.Layers use the chunk-start state, combine their write contributions at the boundary, and apply one update per group; no layer reads a write from its current chunk.Source: §3.1 · pp. 3–4

3.2 UTTT-MOE

• uTTT-MoE routes each token head to K experts in its group’s visible pool, which grows from E/L to E experts at p = 1.With expert size u, total state remains Eu and active state per token head remains Ku; the layer-private routed control uses the same routing but restricts each layer to E/L experts.Source: §3.2 · pp. 4–5

3.3 UTTT-DENSE

• uTTT-Dense applies every head block in one shared state without expert selection; its size-matched design preserves total state LS.The fixed-compute dense design instead stores pS at unchanged per-token fast-weight compute; the table compares these designs across group count p.Source: Table 1 · p. 4, §3.3 · p. 5

3.4 UNIFYING TIME AND DEPTH

• Sharing extends memory recurrence across time and depth: a deep layer’s write in one chunk can support a shallow layer’s next-chunk read, and vice versa.Later-read task-loss gradients propagate through inner updates to earlier writes at other depths; in routed pools, a direct path requires writer and reader to select a common expert.Source: §3.4 · p. 5, Figure 2 · p. 6

4.1 SETUP

• Models are trained at 124M and 760M on 32K-token context with 4,096-token chunks and sliding-window attention; PTL measures loss beyond 30K and RULER averages retrieval accuracy over 4K–32K.The four fast-weight models share total state and core training settings; uTTT-MoE and TTT-MoE use loss-free balancing and an aggregated update schedule.Source: §4.1 · p. 6, Table 2 · p. 7

4.2 RESULTS

• Sharing adds 2.6 and 2.1 RULER points over TTT-MoE at 124M and 760M with the same state and no loss penalty; uTTT-MoE leads bounded-state models while matching or beating full attention’s per-token loss.At 124M it leads TTT-MoE at every tested length; at 760M it leads at 4K and 8K but trails by less than two points at 16K and 32K.Source: §4.2 · p. 7, Figure 3 · p. 7
• Without load balancing, the busiest expert receives 5.9× and 9.6× its even share at 124M and 760M; at 124M, removing balancing costs uTTT-MoE 4.6 RULER points.At 760M, the unbalanced run’s higher mean comes from one easy needle task, while it loses 8–9 points on five others and reads mostly from 22 of 96 experts. The read-output probe reaches all 77 earlier deep-layer write locations in uTTT-MoE and none in TTT-MoE; freezing state writes raises loss by 7.98 mnat in Full and 1.13 mnat in Cut.Source: §4.2 · pp. 7–8, Figure 4 · p. 8

4.3 MEASURED COST

• uTTT-MoE trains fastest in every measured setting and reaches TTT-Dense’s final training loss in 1,540 H100 GPU-hours, 15% fewer than TTT-Dense and 27% fewer than TTT-MoE.The authors attribute the speedup as consistent with fewer, larger kernel launches from aggregating cross-layer contributions into one shared-pool update per non-final chunk boundary.Source: §4.3 · pp. 8–9, Figure 5 · p. 8

5.1 SETUP

• The experiments train object models on Objaverse and evaluate on Google Scanned Objects, using eleven training-view writes and 23 evaluation views; the setup also includes DL3DV scenes.A posed source view reads and writes the state, while a target-camera chunk reads state and renders the held-out view.Source: §5 · p. 9, §5.1 · pp. 9–10

5.2 FROM PRIVATE POOLS TO ONE SHARED POOL

• At fixed per-layer compute, full sharing gains 0.92 dB for routed models and 0.76 dB for dense models at view 23, with disjoint intervals for all-view means.The gains persist beyond the trained horizon, and the dense-model gain requires neither routing nor more state.Source: §5.2 · p. 10, Table 3 · p. 10
• A shared pool of 8 experts beats private pools by 0.99 dB at view 23 despite storing one eighth of the state and offering the same eight candidates per token head.This comparison isolates shared writes from a wider routing choice. Removing layer-8 writes to layer-3 reads reduces PSNR by 0.69 dB, versus at most 0.06 dB for same-layer routes.Source: §5.2 · p. 10

5.3 MORE COMPUTE AT FIXED STATE

• At fixed total state, routed view-23 PSNR improves monotonically by 1.07 dB as per-layer compute rises; dense models gain 1.29 dB, including 0.76 dB from sharing and 0.53 dB from widening.At eight routed compute units, view-23 PSNR is 26.55 dB versus 26.54 dB for full attention; all three metrics improve, including LPIPS falling from 0.185 for TTT-Dense to 0.158 for uTTT-MoE.Source: §5.3 · pp. 10–11

6 CONCLUSION

• uTTT shares fast-weight memory across depth while keeping backbone parameters layer-specific.The shared memory evolves along time and depth together.Source: §6 · p. 11
• Layer-private TTT is a partitioned special case of the broader shared-memory design space.The conclusion frames depth as contributing to test-time optimization instead of multiplying separate memories.Source: §6 · p. 11
• At equal or smaller state, uTTT improves language modeling and view synthesis; its routed form trains faster, and trained models depend on cross-layer routes.These are the conclusion’s stated scope-level outcomes across the paper’s two evaluation settings.Source: §6 · p. 11

A.1 FAST WEIGHTS AND TEST-TIME TRAINING

• In this paper’s TTT layers, chunk writes form the inner loop, while outer-loop training learns initial fast weights and update predictors, as in MAML.The passage also distinguishes this long-memory use from test-time optimization for generalization on novel instances.Source: §A.1 · p. 17

A.2 SHARING PARAMETERS AND EXPERTS ACROSS DEPTH

• Related depth-sharing designs tie slow layer parameters or share experts, whereas uTTT shares fast-weight memory updated within each sequence and retains its updates in the training graph.The related-work discussion contrasts this with routed, looked-up, or shared components that stay fixed at inference, apart from MoUE’s routing state.Source: §A.2 · p. 17

A.3 NOVEL VIEW SYNTHESIS

• Novel-view synthesis includes per-scene fitted representations and models that infer scene structure from one or a few images, including latent-token approaches.The passage notes that decoder-only LVSM runs input-view and target-ray tokens through one transformer, and that LaCT applies layer-private fast weights to this task.Source: §A.3 · p. 17

B DATASETS AND LICENSES

• Table B1 lists the external datasets used and their license or access terms.Source: B DATASETS AND LICENSES · p. 17, Table B1 · p. 18

C.1 READ NETWORK AND WRITE UPDATE

• Routed language models predict token momentum and average it over chunks, with shared pools also averaging across their layers; dense models predict per-head coefficients and average over tokens.Source: §C.1 · p. 18
• The dense update accumulates write directions into a master state, row-normalizes the fast copy, and in NVS lets source-view chunks write while target pose-only chunks read.Source: §C.1 · p. 19

C.2 WRITE SCHEDULES AND EXECUTION ORDER

• The trained aggregated schedule sums layer contributions evaluated at the chunk-start state, reducing optimizer updates per write-enabled chunk from L to p versus L successive sequential updates.Source: §C.2 · p. 19

D.1 LANGUAGE MODELING

• Language models share one recipe and differ in their token mixer beside 4K-window attention, while the fast-weight block is specified separately.Source: §D.1 · p. 20, Table D1 · p. 20, Table D3 · p. 21
• Model comparisons align parameters or state where specified; uTTT-Dense uses one shared state with selectable width multipliers and layer-specific update and output modules.Source: §D.1 · p. 20

D.2 NOVEL VIEW SYNTHESIS

• Novel-view fast-weight models share a backbone and recipe, while the language-modeling and view-synthesis blocks differ in projection layout, sublayer order, head dimensions, and update optimizers.Source: §D.2 · pp. 20, 22
• Novel-view references change the memory sublayer and compensating FFN width; Figure E1 also plots retrieval by task and validation loss by sequence position at both model scales.Source: §D.2 · p. 22, Figure E1 · p. 23

E.1 LOAD BALANCING

• The balancing comparisons vary mechanism and scope across routed runs, with figures resolving results by task and tables identifying the reported configurations.Source: Figure E2 · p. 24, Table E1 · p. 24, Table E2 · p. 24

E.2 A DENSE SHARED STATE AT 124M

• Table E4 reports a 124M dense shared-state width ladder, where r=1 matches TTT-Dense in per-token fast-weight compute and r=12 matches it in state size.Source: Table E4 · p. 25, §E.2 · p. 25

E.1 LOAD BALANCING

• At 124M and 760M, the reported runs leave 0 and 3 experts below a quarter of an even read-routing share, versus 3 and 74 for unbalanced routers; pool scope still leaves many experts underused within individual layers.Source: §E.1 · p. 25
• At 124M, six uTTT-MoE seeds span 14.47–16.60 RULER and 3.0753–3.0939 per-token loss; at 760M, the selected configuration uses the 124M mechanism and scope with the best of five tested bias-update rates.Source: §E.1 · p. 25

• The appendix tests seven 124M models on loss to 128K and retrieval to 64K; fixed-target and reset-segment controls address document-depth effects in the per-position comparisons.Source: F LENGTH EXTRAPOLATION · p. 26, Table F2 · p. 27
• From the first to fourth reset 32K segment, every model improves by 0.010–0.015 nats, while fixed-target Transformer-SWA remains at 3.222 nats as its visible prefix grows from 4K to 120K.Source: F LENGTH EXTRAPOLATION · p. 26, Table F1 · p. 27
• About 0.035 nats below SWA, both expert pools retain their advantage through 128K; uTTT-MoE has the lowest curve within 32K, while LaCT crosses zero near 85K and full attention deteriorates beyond training length.The 64K-trained full-attention reference deteriorates later than the other full-attention model.Source: F LENGTH EXTRAPOLATION · p. 26, Figure F1 · p. 27

G EFFICIENCY MEASUREMENTS

• Fused Triton kernels improve uTTT-MoE throughput by 15% at 124M and 18% at 760M; with fusion, uTTT-MoE remains 19% and 15% faster than TTT-MoE, respectively.Throughput is training tokens per second per GPU, measured as the median of three repeats at batch 1 × 32K on one A100.Source: G EFFICIENCY MEASUREMENTS · p. 28, Table G1 · p. 28
• At 760M on one A100, uTTT-MoE prefills at batch 1 as fast as TTT-Dense and 1.3× faster than TTT-MoE at 16K and 32K prompts; decoding matches TTT-MoE but is 0.72–0.74× TTT-Dense across tested contexts.Time to first token is 29–33% lower than TTT-MoE at 32K and 128K; chunk-boundary pauses are 26–27% shorter at tested contexts.Source: G EFFICIENCY MEASUREMENTS · p. 28, Figure G1 · p. 29
• On H100, uTTT-MoE is 1.14–1.16× faster than TTT-Dense across tested batches; at batch 1, peak memory is 1.29–1.37× TTT-Dense’s in the eight-GPU runs.Source: G EFFICIENCY MEASUREMENTS · p. 28, Figure G1 · p. 29

H ADDITIONAL NOVEL VIEW SYNTHESIS COMPARISONS

• The appendix supplies additional novel-view-synthesis comparisons, including configurations, intervals, and a comparison of the two write schedules.Source: H ADDITIONAL NOVEL VIEW SYNTHESIS… · p. 29

H.1 SHARING AND COMPUTE ON SCENES

• 0.54 dB higher PSNR over all views comes from raising routed activation from a=1 to a=8 on scenes; widening dense state adds 0.31 dB.The scene-sharing levels differ by at most 0.10 dB, so the sharing conclusion rests on objects, whose test set is seven times larger; the compute conclusion holds on both datasets.Source: §H.1 · p. 29

H.2 TOTAL EXPERTS AND POOL SIZE

• At 32 total experts, quality rises at every merge on Objaverse/GSO, and activation sweeps are monotone on both datasets.End-point intervals do not overlap for both object sweeps and for the scene activation sweep, supporting the same orderings at a second total pool size.Source: §H.2 · p. 29
• Across four sharing levels, object quality stays within 0.25 dB and scene quality within 0.16 dB; the 8-expert routed pool leads TTT-MoE-e8 on Objaverse/GSO over all views.The object comparison has non-overlapping intervals and uses one eighth as many experts with the same eight candidates per token head; widening dense state raises compute along with state.Source: §H.2 · p. 29

H.3 EVERY CONFIGURATION

• The configuration tables list layer-private pools, universal pools, depth partitions, activation sweeps, dense states, and dense partitions for each dataset.The supplied table description identifies the configuration coverage but does not provide cell values.Source: §H.3 · p. 31

H.4 WRITE SCHEDULES

• At target view 23, the aggregated schedule is about 0.1 dB lower on Objaverse/GSO and 0.01 dB lower on DL3DV, but trains 2.1% faster and uses 26.3% less peak memory.The canonical sweeps use the aggregated schedule; cost measurements use one H100 80GB.Source: §H.4 · p. 31, Table H3 · p. 33

I ANALYSIS OF SHARED MEMORY

• A shared-memory analysis derives possible cross-layer dependencies, probes them in trained models, and reports preferential expert reuse for repeated surface observations.The analysis distinguishes the dependency mechanism from its validation and observed routing behavior.Source: I ANALYSIS OF SHARED MEMORY · p. 31

I.1 MECHANISM: RECURRENCE AND CROSS-DEPTH REACH

• A shared state advances over chunk-layer pairs, letting a write by any layer in one chunk reach any layer’s read in the next chunk; private memories exclude deeper-writer-to-shallower-reader paths.For a universal pool, the memory-chain derivative is generically nonzero for dense memory, or for routed memory when writer and reader select a common expert; no chunk reads its own writes.Source: §I.1 · pp. 32, 34

I.2 VALIDATION: PROBE, TRAINING CONTROL AND INTERVENTIONS

• Inference interventions freeze pool writes or restrict reads to own-layer writes in language modeling, and remove writer–reader routes individually in view synthesis.The supplied passages specify intervention setups and measurements but do not report their outcome values.Source: §I.2 · pp. 35–36
• A read-output gradient probe perturbs each chunk-layer write input and measures its squared-gradient share, counting top-down cells as reached when nonzero in every document.The probe uses 300 documents and tests the layer-1 read in 124M TTT-MoE and uTTT-MoE models; random-model and finite-difference checks validate tested zero cells.Source: §I.2 · p. 35
• The cross-layer gradient cut preserves probe-model logits while removing gradients from boxed top-down cells and retaining the reading layer’s own row.The trained Cut model changes only the backward pass; the forward configuration and evaluation stream remain the same as the Full model.Source: §I.2 · p. 35

I.3 BEHAVIOR: EXPERT REUSE BY REDUNDANT TOKENS

• Same-surface token pairs reuse the same expert in 2.63% of comparisons, versus 1.88% for controls and 1.56% for uniform routing; their same-surface/control ratio is 1.40 [1.39, 1.41].Across sharing partitions, the ratio is 1.40×, 1.39×, and 1.15× for e64-p1, e32-p2, and e16-p4; the layer-private e8-p8 shares no physical expert across layers.Source: §I.3 · p. 36