Semigroup-JEPA: Better Features, Not Better Dynamics

Paper: 2609.10464 Authors: Andy Zeyi Liu, Haoran Sun, Lucas Baker, Randall Balestriero, John Sous Categories: cs.AI, cs.CV, cs.LG

The Gap

JEPA world models learn a compact latent representation of the world that supports prediction and planning. That is a behavioural claim, and it leaves a question unasked: their capability to learn physics and generate physically realistic dynamics remains untested. A latent space can support good predictions without the model representing the physical law — the same way a curve fit can predict well without containing the mechanism.

The failure mode that exposes this is out-of-distribution generalisation within a law. The paper’s testbed is dynamical tasks under different gravitational fields that, despite obeying the same physical law, exhibit qualitatively different dynamics — floating motion in weak fields to rapid bouncing in strong ones. One law, very different behaviour, so a model that has learned the law should transfer and a model that has learned the surface statistics should not.

   JEPA WORLD MODELS: A BEHAVIOURAL CLAIM, AN UNASKED QUESTION

   JEPA world models learn a COMPACT LATENT REPRESENTATION that
   supports PREDICTION and PLANNING
        |
        v
   [THE UNASKED QUESTION]
     their capability to LEARN PHYSICS and generate PHYSICALLY
     REALISTIC dynamics remains UNTESTED
        <- a latent space can support good predictions WITHOUT the
           model representing the LAW
        <- the same way a CURVE FIT can predict well without
           containing the MECHANISM

   [THE FAILURE MODE THAT EXPOSES THIS]
     OUT-OF-DISTRIBUTION GENERALISATION **WITHIN** A LAW
        testbed: dynamical tasks under DIFFERENT GRAVITATIONAL FIELDS
          that, DESPITE OBEYING THE SAME PHYSICAL LAW, exhibit
          QUALITATIVELY DIFFERENT DYNAMICS
          -> from FLOATING MOTION in weak fields to RAPID BOUNCING in
             strong ones
       -> one LAW, very different BEHAVIOUR
       -> a model that learned the LAW should TRANSFER
          a model that learned the SURFACE STATISTICS should NOT

The Increment

One sentence: Before this paper, JEPA world models had not been tested on physics; after it, supplying the governing parameter to the temporal model cuts open-loop error up to 2× and lifts control success up to 2.5×, and the gain is traced to the encoder rather than the predictor.

Core Mechanism

The method has two changes, and the second is the one that produces the result.

The parameter governing the physics is supplied to the temporal model via action-conditioning. So gravity is not something the model must infer from the observation stream — it is given, as a control input. That reframing matters for what is being tested: the question becomes whether the model can use a known physical parameter, rather than whether it can infer one.

The encoder and predictor are jointly trained through an autoregressive latent rollout. This is the structural change, and its significance is not obvious until the explanation arrives: joint training over multi-step rollouts is what allows the rollout loss to reach the encoder.

The evaluation spans two settings with different demands: two-dimensional datasets where SG-JEPA reduces open-loop prediction error by up to 2×, and three-dimensional robotic datasets where it increases control success rate up to 2.5×, with independently trained diffusion policies — so the control result is not an artefact of a policy co-adapted to the world model. The comparison is against DINO-WM, an existing model in the same family.

Then the part that makes the paper more than a benchmark result. The authors develop a linear feature model that separates local law-conditioned error from its recursive amplification under rollout. That decomposition is the analytic instrument: total error at long horizon is not just the per-step error, it is the per-step error compounded by recurrence, and separating them lets you ask which part the method improved.

Guided by that model, the explanation is that back-propagating the multi-step rollout loss into the representation trains the encoder to keep the features that the predictor can carry forward, and those are the features the dynamics depend on. So the mechanism is a selection pressure: the rollout loss, applied to the representation, favours encoders whose features remain usable across many steps — and features that survive many steps of a rollout under a fixed physical law are, plausibly, the ones the law acts on.

And the conclusion drawn from that: most of the gain comes from the encoder learning better features rather than from the predictor learning better dynamics. This is a claim that could easily have been left as “our architecture is better”, and the paper identifies which of the two components moved. It also implies the interesting bottleneck for this class of models is representational rather than dynamic.

   TWO CHANGES, AND THE SECOND PRODUCES THE RESULT

   [1] THE PARAMETER GOVERNING THE PHYSICS IS SUPPLIED to the temporal
       model VIA ACTION-CONDITIONING
         -> gravity is NOT something the model must INFER from the
            observation stream: it is GIVEN, as a control input
         <- reframing matters for what is TESTED: the question becomes
            whether the model can USE a known physical parameter,
            rather than whether it can INFER one

   [2] THE ENCODER AND PREDICTOR ARE JOINTLY TRAINED through an
       AUTOREGRESSIVE LATENT ROLLOUT
         <- the STRUCTURAL change, and its significance is not obvious
            until the explanation arrives: JOINT TRAINING OVER
            MULTI-STEP ROLLOUTS is what allows the ROLLOUT LOSS TO
            REACH THE ENCODER

   EVALUATION ACROSS TWO SETTINGS WITH DIFFERENT DEMANDS
     2D datasets: OPEN-LOOP PREDICTION ERROR reduced by up to 2x
     3D robotic datasets: CONTROL SUCCESS RATE up to 2.5x, with
       INDEPENDENTLY TRAINED DIFFUSION POLICIES
       <- so the control result is NOT an artefact of a policy
          CO-ADAPTED to the world model
     comparison: DINO-WM, an existing model IN THE SAME FAMILY

   THE ANALYTIC INSTRUMENT
     a LINEAR FEATURE MODEL that SEPARATES
       LOCAL law-conditioned error
     FROM
       its RECURSIVE AMPLIFICATION UNDER ROLLOUT
       <- total error at long horizon is not just PER-STEP error; it
          is per-step error COMPOUNDED BY RECURRENCE
       <- separating them lets you ask WHICH PART the method improved

   THE EXPLANATION, GUIDED BY THAT MODEL
     BACK-PROPAGATING THE MULTI-STEP ROLLOUT LOSS INTO THE
     REPRESENTATION trains the ENCODER to keep THE FEATURES THAT THE
     PREDICTOR CAN CARRY FORWARD, and THOSE ARE THE FEATURES THE
     DYNAMICS DEPEND ON
       <- the mechanism is a SELECTION PRESSURE: the rollout loss,
          applied to the representation, favours encoders whose
          features REMAIN USABLE ACROSS MANY STEPS
       <- features that SURVIVE MANY STEPS OF A ROLLOUT UNDER A FIXED
          PHYSICAL LAW are, plausibly, THE ONES THE LAW ACTS ON

   CONCLUSION DRAWN
     MOST OF THE GAIN COMES FROM THE ENCODER LEARNING BETTER FEATURES,
     NOT FROM THE PREDICTOR LEARNING BETTER DYNAMICS
       <- could easily have been left as "our architecture is better";
          the paper identifies WHICH OF THE TWO COMPONENTS MOVED
       -> implies the interesting bottleneck for this class of models
          is REPRESENTATIONAL rather than DYNAMIC

Think of it as learning to predict where a ball goes by first learning a good coordinate system. If your representation of the scene mixes up the properties that the law acts on — position, velocity, the strength of the field — then no amount of better extrapolation on top will save you: the predictor is working with numbers that do not separate the things the law cares about. What the rollout loss does, pushed back into the representation, is penalise encodings that cannot be carried forward many steps — and an encoding that survives many steps is one where the law’s variables stay cleanly separated. So the improvement shows up as better features, and the predictor looks barely different because it was already adequate for the coordinates it was given.

Key Concepts

• Testing physics rather than prediction quality: transferring across fields that share a law but differ qualitatively. It distinguishes representing the law from fitting surface statistics.
• Supplying the governing parameter as control input: making gravity an action-conditioning signal rather than something to infer. It reframes the test as “can the model use a known parameter”.
• Joint training through autoregressive rollout: the structural change that lets the rollout loss reach the encoder. It is the mechanism, and it is not visible from the architecture alone.
• Separating local error from recursive amplification: the linear feature model. It is what makes it possible to attribute the gain between the per-step and compounded components.
• The representational bottleneck: most of the gain coming from the encoder. It implies that improving this class of models is primarily a feature-learning problem.

Framework Shift

Before (JEPA trained for prediction and planning):
  learn a compact latent supporting prediction
  -> physics untested: a curve fit can predict without containing
     the mechanism
  -> the governing parameter must be inferred from the stream
  -> gain, if any, is not attributed between components

After (parameter supplied, rollout loss reaches the representation):
  gravity given via action-conditioning
  encoder and predictor jointly trained through latent rollout
  -> open-loop error down up to 2x; control success up to 2.5x
  -> linear model separates local error from recursive amplification
  -> most of the gain is the ENCODER: features the predictor can
     carry forward, which are the ones the dynamics depend on

From asking whether a world model predicts well, to asking whether it represents the law well enough to transfer across the law’s regimes, the core shift is that the improvement turned out to be representational, and the diagnostic that showed it separated per-step error from its compounding.

Expert Assessment

Problem choice: Excellent, and the testbed is what makes it decisive. Using one law with qualitatively different regimes is exactly the right construction, because it separates a model that has learned the law from one that has learned the statistics of the training conditions — and the paper notes that this had not been tested, which is a real gap given how JEPA models are motivated.

Method maturity: Two aspects stand out. Supplying the governing parameter is a deliberate simplification of the question, and the paper is clear that it tests use rather than inference. More importantly, the linear feature model that separates local error from recursive amplification is a genuine analytic contribution: attributing an architectural gain between encoder and predictor is exactly the question such papers usually leave as an assertion, and answering it here yields a claim about where the bottleneck lies. Training the comparison policy independently rather than jointly is the right control for the control-task result.

Experimental integrity: The attribution result is the honest part, because it is not the flattering reading — an encoder-centric explanation is less actionable for people hoping to swap in a better predictor. Reporting the diffusion policies as independently trained rules out the co-adaptation confound. The limitations are that the linear feature model is a simplification whose separation is approximate, and that the setting is simulated dynamics under varied gravity, so whether the representational story holds for contact-rich or partially observed manipulation is not established.

Writing quality: The mechanism claim is stated with its reasoning, which is what makes it checkable rather than a slogan. Because the practical consequence is “improve the encoder”, a short passage on what that implies for architecture choices — for instance whether any change to the predictor would matter at all — would make the finding directly usable.

Verdict: strong accept — it tests JEPA models on physics with a construction that separates law-learning from statistics-fitting, and it traces the improvement to the encoder using an analytic decomposition rather than asserting an architecture-level win.

Takeaways

• Test transfer within a law, not just across tasks. Regimes that share a mechanism and differ qualitatively separate law-learning from statistics-fitting.
• Ask which component moved. Attaching a gain to the encoder rather than the predictor changes what you would try next.
• Separate per-step error from its recursive amplification. Long-horizon error is compounded, and the decomposition is what allows attribution.
• Train the comparison policy independently. A policy co-adapted to the world model will flatter it for reasons unrelated to the model’s quality.