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

World models typically rely on fixed-size Vision Transformer (ViT) encoders, which are often either over-provisioned for simple tasks—wasting compute and potentially overfitting—or under-provisioned for complex tasks. This paper asks whether an encoder can adaptively grow its capacity during training, adding width (attention heads) or depth (transformer blocks) only when the task demands it, without requiring manual tuning or sacrificing performance.

The author introduces Successive Capacity Growth (SCG), a framework that begins with a minimal encoder (283K parameters) and monitors the prediction loss. When the model hits a plateau, SCG triggers a test-and-verify mechanism. It attempts to expand the architecture using function-preserving transformations—ensuring the model's output remains identical immediately after the change. If the expansion leads to a significant improvement in prediction loss, it is kept; otherwise, the model rolls back to its previous state. The Sketched Isotropic Gaussian Regularizer (SIGReg) is used to ensure that newly added semantic dimensions remain statistically independent and aligned with the predictive objective, preventing representational collapse.

SCG demonstrates that encoders do not need to be pre-allocated at maximum capacity. Across three environments of varying complexity, SCG naturally triggered expansions only when necessary: it added width for low-level semantic capacity on a 2D navigation task and depth for higher-order abstraction on a 60-dimensional multi-object dynamics task. On the 2D task, the adaptive encoder outperformed the fixed large baseline by 23%, suggesting that over-provisioning can actually hinder optimization. Overall, SCG achieved performance comparable to or better than fixed large models while using up to 56 times fewer parameters, proving that adaptive growth is a highly efficient alternative to static architectural design.