PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Frederik Berenz

arXiv:2608.27367v1cs.CVcs.AI

TL;DR

Fixed JEPA encoders can waste capacity on simple tasks while lacking capacity for complex ones. SCG grows ViT encoders in width or depth through function-preserving trials selected by predictive-loss verification. Across tested tasks, it matched or exceeded the fixed-small baseline, with up to 56× greater parameter efficiency and improvements of 49% on a 2D task and 20.3% on a 60D task.

  • Problem

    Fixed-size JEPA encoders can be over-provisioned for simple tasks and under-provisioned for complex ones, while exhibiting redundant attention heads and layers.

  • Method

    SCG starts from a minimal ViT encoder and incrementally adds attention heads or transformer blocks through function-preserving, task-agnostic test-and-verify trials, with SIGReg maintaining independent semantic dimensions.

  • Results

    Across tested environments, SCG matched or exceeded the fixed-small baseline, achieving 49% improvement on a 2D task, 20.3% on a 60D task, and up to 56× greater parameter efficiency.

  • Takeaways & Limitations

    JEPA world-model encoders need not be pre-allocated at maximum capacity and can grow successively as task demands change.

  • Takeaways & Limitations

    Experiments tested models up to 5.7M parameters on a single GPU; scaling to ViT-Base at 86M parameters was untested.

Abstract

from arXiv · show

Joint-Embedding Predictive Architectures (JEPAs) for world modeling typically employ fixed-size Vision Transformer encoders that are over-provisioned for simple tasks and under-provisioned for complex ones, with significant redundancy across attention heads. We propose Successive Capacity Growth (SCG), a method that starts from a minimal encoder (1 head, 2 layers, 283K parameters) and grows incrementally in width (adding attention heads for low-level semantic capacity) or depth (adding transformer blocks for higher-order semantic abstraction), driven by a task-agnostic test-and-verify mechanism that exploits function-preserving expansion to safely trial architectural changes and roll back if they do not improve prediction loss. The Sketched Isotropic Gaussian Regularizer (SIGReg) ensures that all learned semantic dimensions remain statistically independent and aligned with the predictive objective, preventing collapse even as the architecture grows. On a 60-dimensional multi-object dynamics task, SCG naturally triggers depth expansion, improving prediction loss by 20.3% over the fixed small baseline with 56 times greater parameter efficiency than scaling to the fixed large model; on a 2D navigation task, a single width expansion yields even an 23% improvement over the fixed large model. Across all three tested environments of increasing complexity, the adaptive encoder matches or exceeds the fixed small baseline, with zero false-positive expansions and bit-exact function preservation (ratio = 1.0, absolute difference = 0.0). The take-away is that JEPA world model encoders need not be pre-allocated at maximum capacity - they can grow successively as the task demands, achieving significant compute and data efficiency while maintaining representation quality.

1 Introduction

JEPA world-model encoders are typically fixed despite task-dependent capacity needs and substantial redundancy. SCG addresses this by growing width or depth only when predictive evidence warrants expansion.

  • Motivation: LeWM’s fixed ViT-Tiny is over-provisioned for many tasks, with redundant attention heads and diminishing contributions from deeper layers.The three heads have pairwise attention-map cosine similarity ≈0.0001, while residual ratios decline from 0.9 to 0.15 across layers 3–12.
  • Motivation: A 2D navigation task and a 60-dimensional dynamics task use the same encoder despite differing capacity requirements.SCG instead expands width for additional low-level semantic dimensions or depth for higher-order abstraction when demanded by the task.
  • Existing approaches: The fixed large encoder reaches prediction loss 0.0004 versus 0.0005 for fixed small on the 2D Two-Room task, requiring 20× more parameters.The reported redundancy includes three effectively replicated attention heads.
  • Related work: Function-preserving expansion and task-agnostic triggering address a gap left by prior progressive, recursive, and function-preserving growth methods.The stated gap is deciding when expansion is warranted without per-dataset tuning.
  • Contribution: SCG extends LeWM with an adaptive encoder while retaining its stable JEPA training objective and SIGReg-based anti-collapse setup.LeWM trains jointly from raw pixels without stop-gradient or EMA, providing the testbed for mid-training architectural growth.

2 Method

SCG monitors predictive loss and safely trials width or depth expansions that preserve the encoder’s function initially. Accepted expansions add semantic dimensions or hierarchical computation, while failed trials are rolled back.

  • Growth policy: SCG starts from a 1-head, 2-layer, 283K-parameter encoder and detects plateaus using prediction loss over recent epochs.A plateau is defined as no more than 2% relative improvement over two epochs compared with the preceding two.
  • Width expansion: Width expansion adds one attention head and increases dmodel by 64 for additional independent semantic factors.SIGReg is intended to keep the new dimensions statistically independent from existing ones.
  • Depth expansion: Depth expansion adds an identity-initialized transformer block for learning hierarchical abstractions over existing semantic factors.The new block initially maps its input to itself through zeroed output projections and identity LayerNorm settings.
  • Test and verify: Each expansion preserves the model output immediately, then is retained only if Lpred improves by more than 2% after two training epochs.The decision uses prediction loss rather than total loss because SIGReg fluctuates independently of predictive capacity.
  • Function preservation: Width expansion initializes new-head QKV weights by copying an existing head while zeroing its output projection and padding downstream inputs.This lets the head begin from a meaningful initialization without initially changing the model output.
  • Rollback: Failed expansion trials restore the saved model and optimizer state, preserving the pre-expansion training trajectory.The algorithm deep-copies the model and optimizer before trying width and then depth expansion.

3 Experiments

Across three environments, SCG adapts encoder growth to task complexity: it expands width or depth when prediction loss plateaus, while preserving function exactly and avoiding unnecessary growth.

  • Experimental setup: SCG starts from Fixed Small and grows through test-and-verify expansion, whereas Fixed Small and Fixed Large remain static capacity baselines.The adaptive model begins with 1 head, 2 layers, d_model = 64, and 283K parameters.
  • Main results: 76.5% prediction loss improvement over Fixed Small occurs on Two-Room after a single width expansion, while 17.8-38.8% improvement occurs on 30-Object after depth expansion.On Push-T, SCG makes no expansion and converges.
  • Main results: 20.3% prediction loss improvement over Fixed Small on 30-Object uses 50K additional parameters, versus 5.4M for Fixed Large, yielding approximately 56× greater parameter efficiency.SCG approaches Fixed Large performance while using substantially fewer added parameters.
  • Expansion behavior: SCG’s expansion decisions vary by task and seed: depth is retained on some 30-Object runs, width on one Two-Room run, and none on Push-T.Some runs remain learning without plateau, showing that expansion is not forced at a fixed schedule.
  • Function preservation: Width and depth expansions achieve bit-exact function preservation, while subsequent training makes the added capacity immediately usable.All six validation runs report identical values; new-head cosine similarity changes from 1.0 to 0.91–0.99, and new-layer residual ratio from 0.0 to 0.22.
  • Ablations: Using prediction loss alone for plateau detection eliminates false-positive expansions across 9 runs, whereas effective rank remains approximately 3 for both 2D and 60D tasks.The ablations identify prediction loss as the more reliable capacity signal.
  • Over-provisioning analysis: On Two-Room, SCG outperforms Fixed Large with prediction loss 0.000176 versus 0.000362 while using only 11% of its parameters.The result is reported as evidence of an over-provisioning penalty for simple dynamics.

4 Limitations

The evaluation is limited to synthetic environments, prediction loss, modest model scales, and finite training horizons for expansion detection.

  • All experiments use synthetic environments, so validation on real data with higher visual complexity remains necessary.
  • The study evaluates prediction loss but not planning success rate, the ultimate downstream metric for world models.
  • Models up to 5.7M parameters were tested, leaving scaling to ViT-Base at 86M parameters untested.
  • One of three 30-Object seeds did not plateau within 50 epochs, preventing expansion under the tested schedule.

5 Discussion & Outlook

SCG frames encoder capacity as successive, task-driven growth: width adds independent semantic dimensions, while depth adds hierarchical abstraction. The discussion links this strategy to reduced redundancy, efficiency, and future extensions.

  • Successive learning: SCG grows encoder capacity successively, using width for independent semantic dimensions and depth for higher-order semantic abstraction.
  • Successive learning: SIGReg’s characteristic-function matching keeps grown dimensions statistically independent and aligned with the predictive objective.
  • Redundancy in fixed ViT: SCG reduces fixed-architecture redundancy while closing the performance gap to the fixed large baseline with up to 56× greater parameter efficiency.
  • Compute and data efficiency: Starting from 283K rather than 5.7M parameters improves per-epoch training efficiency and reduces overfitting risk on small datasets.
  • The over-provisioning penalty: SCG surpassed Fixed Large on the 2D task, suggesting that excessive capacity can introduce redundant parameters that capture noise rather than signal.
  • Future directions: Future work includes real-data validation, planning evaluation, pruning, dynamic inference adjustment, and extension to ViT-Base.

6 Conclusion

SCG grows ViT encoders through function-preserving, task-triggered width or depth expansions, starting minimally and retaining only useful capacity. Across the tested tasks, it improves over Fixed Small while using capacity more efficiently.

  • SCG starts from 283K parameters and expands width or depth only when task demands trigger its test-and-verify mechanism.
  • 49% improvement on a 2D task and 20.3% on a 60D task were achieved over the Fixed Small baseline, with up to 56× greater parameter efficiency.
  • SIGReg’s complex characteristic-function matching keeps grown dimensions statistically independent and non-redundant.

A Hyperparameters

The supplied hyperparameter section identifies Table 6 as the complete hyperparameter table.

  • Table 6 is identified as the complete hyperparameter table.

B Reproducibility Checklist

The experiments use fixed seeds, specified hardware and software, parallel execution, and 50 effective epochs, with code publicly available for reproducibility.

  • Experiments use seeds 3072, 42, and 123, with data generation seed 42 fixed across all runs.
  • Runs use an NVIDIA RTX A6000 with 48GB VRAM, PyTorch 2.4.1, and CUDA 12.4.
  • Three configurations per task and seed run in parallel via ThreadPoolExecutor, with 50 effective epochs excluding rollback test epochs.
  • All code is available at the project’s GitHub repository.
Loading 2608.27367v1…