Successive Capacity Growth: Task-Complexity-Driven Width and Depth Expansion for Vision Transformer Encoders in JEPA World Models

1 Introduction

• The fixed 12-layer, 3-head encoder shows redundant attention heads and diminishing layer contributions, while the same 5M-parameter encoder is used for simple and complex tasks.Attention-map cosine similarity is ≈0.0001; residual ratios decline from 0.9 to 0.15 and inter-layer cosine similarities reach 0.99.Source: §1 · p. 2

1.1 Existing Approaches and Their Limitations

• Fixed encoders over-provision simple tasks: on 2D Two-Room, the 283K-parameter small encoder has prediction loss 0.0005 versus 0.0004 for the 5.7M-parameter large encoder.Prior expansion methods lack SCG’s online, task-driven trigger: they use predefined schedules, adjust existing architectures, require separate growth training, or provide no trigger policy.Source: §1.1 · p. 2

1.2 Motivation

• SCG targets the gap between function-preserving expansion and a task-agnostic trigger requiring no per-dataset tuning, with SIGReg keeping grown semantic dimensions independent and prediction-aligned.Trial expansions can be rolled back if they fail, at no cost beyond compute; the adaptive encoder is reported to achieve up to 56× greater parameter efficiency and avoid over- and under-fitting.Source: §1.2 · p. 2

1.3 Contributions

• 49% loss improvement on the 2D task and 20.3% improvement on the 60D task accompany width and depth expansion, respectively; SCG converges without expansion on the 5D task.SCG also achieves bit-exact function preservation for both expansion types, with ratio = 1.0 and absolute difference = 0.0.Source: §1.3 · p. 3

1.4 Related Work

• Unlike recursive vision transformers [9] or LiGO [8], SCG grows during JEPA training from predictive-loss signals; it adapts Net2Net [3] and transformer transformations and retains SIGReg [1].Its new attention head copies existing QKV weights but starts with a zero output projection, preserving function while allowing independent learning; SIGReg uses characteristic-function matching to ensure statistical independence.Source: §1.4 · p. 3

2.1 Problem Formulation and Preliminaries

• The custom pre-norm ViT uses patch size 14, 257 tokens, 64-dimensional heads, a 192-dimensional CLS projector, and a fixed MLP predictor; gradients flow through predicted and target embeddings.LSIGReg is described as preventing collapse.Source: §2.1 · pp. 3–4

2.2 Overview of the Approach

• Starting at 1 head, 2 layers, and 283K parameters, SCG expands after a 2% loss plateau, testing added width then depth and retaining only expansions yielding >2% prediction-loss gains.A trial trains for 2 full epochs; if it does not improve prediction loss enough, SCG rolls back and tries the other expansion type, stopping if neither succeeds.Source: §2.2 · p. 4

2.3 Core Components

• The 2% threshold defines both plateau detection and expansion verification rather than a tuned hyperparameter, while test epochs are excluded from equal effective training budgets.A two-epoch cooldown follows each successful expansion.Source: §2.3 · p. 5

2.4 Theoretical Foundation

• Width addresses representational bottlenecks, while depth addresses computational bottlenecks; function-preserving trials and optimizer-state rollback keep failed expansions from disrupting the optimization trajectory.The approach tests width first and then depth without explicitly diagnosing which bottleneck is active.Source: §2.4 · pp. 5–6

3.2 Main Results

• On 30-Object, SCG improves prediction loss 20.3% over Fixed Small and is 56× more parameter-efficient; on Two-Room it surpasses Fixed Large, while Push-T needs no expansion.On Two-Room, SCG also surpasses Fixed Small by 49%. For 30-Object, SCG reduces loss by 0.0017 with 50K additional parameters, while Fixed Large uses 5.4M additional parameters for a 0.0033 reduction.Source: §3.2 · pp. 6–7, Table 1 · p. 7, Table 5 · p. 9
• Across seeds, depth expansion was kept on 30-Object, width expansion on Two-Room, and all Push-T runs converged without expansion.On 30-Object, seeds 3072 and 123 kept depth expansions while seed 42 continued learning; on Two-Room, seed 42 kept a width expansion while seeds 3072 and 123 converged.Source: §3.2 · p. 7, Table 3 · p. 7

3.3 Function Preservation

• In three validation runs per dataset, width and depth expansions preserved all internal representations and final loss bit-exactly, with zero difference.Table 4 reports identical values across six validation runs and an fp_ratio defined as post-loss / pre-loss.Source: Table 4 · p. 8, §3.3 · p. 8
• After expansion, loss continued decreasing, the new head’s cosine similarity fell from 1.0 to 0.91–0.99, and the new layer’s residual ratio rose from 0.0 to 0.22.These changes indicate the added capacity became usable during continued training.Source: §3.3 · p. 8

3.4 Ablation Studies

• Using prediction loss alone for plateau triggers eliminated false-positive expansions (0/9 runs); total loss caused false positives, and effective rank was ≈3 on both 2D and 60D tasks.The authors report loss-based test-and-verify as more reliable than effective rank as a capacity signal.Source: §3.4 · p. 8
• Fixed Large’s three heads had pairwise cosine similarity ≈0.0001, indicating near-identical attention patterns and redundant processing capacity.This ablation identifies redundancy in the standard ViT-Tiny encoder configuration.Source: §3.4 · p. 9

3.5 Analysis

• The expansion choices aligned with the proposed bottlenecks: Two-Room needed more representational dimensions, 30-Object needed greater depth, and Push-T needed neither.The 30-Object depth expansions occurred with dmodel = 64; the authors interpret them as addressing hierarchical modeling of multi-object interactions.Source: §3.5 · p. 9
• On Two-Room, SCG’s prediction loss was 0.000176 versus Fixed Large’s 0.000362 despite using only 11% of its parameters.The authors describe this as an over-provisioning penalty and suggest excess capacity can capture noise or make optimization harder.Source: §3.5 · p. 9

• Experiments use synthetic environments, so the expansion mechanism still requires validation on real data with higher visual complexity.Source: §4 · p. 10
• LeWM uses planning via Cross-Entropy Method, but evaluation measures prediction loss rather than planning success rate, the ultimate world-model metric.Source: §4 · p. 10
• Models up to 5.7M parameters were tested, while scaling to ViT-Base (86M) remains untested.Source: §4 · p. 10
• One of three seeds on the 30-Object task did not plateau within 50 epochs, preventing expansion.Longer training or a lower plateau threshold could address this, but would require more test epochs.Source: §4 · p. 10

Successive learning

• Width expansion adds independent representation dimensions, depth expansion adds hierarchical transformations, and SIGReg keeps grown dimensions independent and aligned with prediction.The discussion characterizes this as successive, task-demand-driven capacity growth.Source: §5 · p. 10

The over-provisioning penalty

• Up to 56× greater parameter efficiency lets SCG close the performance gap to the fixed large baseline, and SCG surpassed that baseline on the 2D task.The discussion argues that over-provisioned parameters can capture noise rather than signal and degrade prediction.Source: §5 · p. 10

Redundancy in fixed ViT

• LeWM’s ViT-Tiny shows redundancy, with three heads at cosine similarity 0.0001 and layers 3–12 contributing progressively less, as residual ratios fall from 0.9 →0.15.The discussion presents these findings as evidence that fixed large encoders spend compute on redundant processing.Source: §5 · p. 10

Compute and data efficiency

• Starting from 283K rather than 5.7M parameters makes SCG faster per epoch and less prone to overfitting on small datasets, with at most 4 extra epochs per plateau for testing.The plateau tests use 2 epochs for width and 2 for depth.Source: §5 · p. 10

Broader impact

• SCG adds width for new independent semantic factors and depth for abstractions over existing ones, suggesting demand-driven, target-ordered representation growth for transformer models.The discussion contrasts this with fixed models that distribute information across redundant heads and layers.Source: §5 · p. 10

Future directions

• Future work includes testing on real data, evaluating the full LeWM planning pipeline, pruning unused heads or layers, dynamic inference-time width and depth adjustment [9], and scaling to ViT-Base.Source: §5 · p. 10

6 Conclusion

• 49% improvement on a 2D task with width expansion and 20.3% on a 60D task with depth expansion over the fixed small baseline, closing the gap to the fixed large baseline with up to 56× greater parameter efficiency.SCG starts from 283K parameters and grows only when the task demands.Source: §6 · p. 11
• SCG expands ViT encoders in width for low-level semantic capacity or in depth for higher-order abstraction, using a task-agnostic test-and-verify mechanism.Expansions preserve the function, and SCG grows only when the task demands.Source: §6 · p. 11
• SIGReg’s complex characteristic function matching helps keep grown dimensions statistically independent and non-redundant.The method starts from a 283K-parameter encoder.Source: §6 · p. 11

A Hyperparameters

• Table 6 is identified as the complete hyperparameter table.The supplied passage does not include the table’s settings or values.Source: Table 6 · p. 12