PaperlayerView exact revision on arXiv

Source-linked AI summary

Predictive Closure: A Set-Level Characterization of JEPA's Feature Selection

Mosh Levy, Aly Lidayan, Yoav Goldberg

arXiv:2610.27364v1cs.LGcs.AI

TL;DR

JEPA learns both the representation and its prediction target, so which features it encodes is unclear, despite the representation’s role as a world model for downstream tasks. The paper reformulates the objective as joint selection over feasible feature sets and tests the resulting predictive-closure account. Experiments find that unpredictable features can be better represented when they help predict others, and that interventions can steer representation indirectly through predictive partners.

  • Problem

    Because JEPA learns both sides of its prediction problem and does not specify latent feature content, it remains unclear which features it encodes, although absent features cannot inform downstream planning.

    Source: §1 · p. 1

  • Method

    The paper reformulates JEPA’s objective as selecting a feasible feature set with maximal mutual information between its current and next values, and tests predictions in controlled experiments.

    Source: §1 · p. 2, §3 · p. 4, §3.1 · p. 4

  • Results

    Unpredictable features become more decodable when they help predict others; changing predictive relationships can affect other features’ representation, and auxiliary objectives or predictor side information can strengthen predictive partners.

    Source: §1 · pp. 2–3

  • Takeaways & Limitations

    Targeting one feature can strengthen the representation of its predictive partners, so steering feature selection can be indirect.

    Source: §1 · p. 3

    Claims about changes to observations concern the counterfactual case where JEPA is trained on the sequence after the change.

    Source: §5 · p. 6

Abstract

from arXiv · show

In joint-embedding predictive architectures (JEPA), an encoder and a predictor are trained together to produce a representation that minimizes the error in predicting the representation of the next observation. Because the representations of both the current and the next observation are learned, it is unclear which features they will include. Earlier characterizations describe the features JEPA selects in terms of their individual predictability, raising the concern that unpredictable features, like a traffic light, will not be encoded into the representation. However, the prediction error is also affected by the relationships between features, e.g., the light's color helps predict whether the traffic will stop. We propose the predictive closure characterization, which accounts for these relationships. Based on a reformulation of the JEPA training objective, we show that selection depends on properties of feature sets. This view agrees with earlier works when features evolve independently, but makes new predictions when they do not, for instance that a feature like the traffic light can still be retained. We test this characterization in controlled experiments, which show that unpredictable features can become better represented when they help predict other features, and that changing one feature's predictive relationships can affect how well others are represented. We also demonstrate implications for steering feature selection: targeting one feature with an auxiliary objective, or supplying it to the predictor as side information, can strengthen the representation of another feature whose future values it helps predict, so steering can be indirect.

1 INTRODUCTION

JEPA learns both the representation and the prediction target, leaving unclear which features it will encode; this matters because its representation serves as a world model for downstream planning. The paper proposes predictive closure, which evaluates feature sets jointly, and tests its predictions and implications for steering.

1 INTRODUCTION

  • Earlier characterizations assess features by their own predictability, but the traffic light’s current color can help predict whether a cyclist rides or stops.Thus, considering a feature’s own future alone can miss its predictive value for other features.Source: §1 · p. 2, Figure 1 · p. 2
  • Predictive closure evaluates feature sets jointly: an unpredictable feature may be retained when its current value helps predict other represented features.The reformulated objective favors feasible sets with the greatest mutual information between current and next feature values.Source: §1 · p. 2
  • Unpredictable features become more decodable when they help predict other features; changing predictive relationships can also strengthen one feature set while weakening a competing set.A highly predictive set can still be poorly represented when it exceeds latent capacity.Source: §1 · p. 2
  • An auxiliary reconstruction objective can strengthen an untargeted predictive partner while reducing representation of competing feature sets.The experiments also find that supplying a feature directly to the predictor strengthens some features and weakens others.Source: §1 · pp. 2–3
  • Auxiliary objectives and predictor side information can strengthen a targeted feature’s predictive partners, so steering feature selection can be indirect.Viewing actions as side information suggests action-conditioned JEPA favors features the agent has an impact on.Source: §1 · p. 3

2 BACKGROUND

JEPA predicts the latent representation of the next observation, while its objective does not specify which features the latent should contain. The paper describes features as recoverable variables and contrasts earlier per-feature accounts with a joint account that considers predictive relationships.

2 BACKGROUND

  • JEPA learns a representation by predicting the latent representation of the next observation, and its latent can omit features unlike a reconstruction-trained latent.The paper focuses on which features JEPA encodes.Source: §2 · p. 3

2.1 JEPA DOESN’T EXPLICITLY SPECIFY FEATURES

  • JEPA’s encoder supplies both the predictor’s input and target, while regularization constrains the latent distribution without specifying its feature content.Features are variables of an observation; a feature is represented when its value can be recovered from the latent.Source: §2.1 · p. 3

2.2 EARLIER CHARACTERIZATIONS OF FEATURE SELECTION

  • Earlier accounts describe selection through feature speed, predictability, or a subspace, with the subspace account assuming linear encoder and predictor and independently evolving features.These per-feature descriptions disregard predictive relationships between features; the paper instead assesses features jointly.Source: §2.2 · p. 3

3 PREDICTIVE CLOSURE

Predictive closure characterizes JEPA as selecting a feasible feature set with the greatest predictive information between its present and future values. It explains selection through joint relationships among features, not only their individual predictability.

3.1 JEPA SELECTS SETS, NOT INDIVIDUAL FEATURES

  • The reformulation treats JEPA’s objective as selecting a feasible set whose present values provide the most information about its own future.Feasibility is constrained by the entropy imposed by regularization on the latent representation.Source: §3 · p. 4, §3.1 · p. 4
  • An ideally trained encoder represents a feasible feature set with the greatest predictive information between its current and next values.The paper calls the optimizing set predictive closure.Source: §3.1 · p. 4
  • Unlike individual feature ranking, predictive closure compares whole sets because a feature’s contribution depends on which other features are included.The two views coincide when features evolve independently, but can differ when features are predictively related.Source: §3.1 · p. 4

3.2 ALEATORIC UNCERTAINTY AND PREDICTIVE LEAKAGE

  • The decomposition distinguishes uncertainty in a set’s next values from predictive leakage: what the full observation predicts beyond the set’s current values.A feature that helps predict the set’s next values lowers leakage; a noisy feature adds uncertainty unless it helps predict other represented features.Source: §3.2 · p. 5

4 IMPLICATIONS OF PREDICTIVE CLOSURE

Predictive closure accounts for earlier findings and predicts how relationships among features affect representation. It also implies that auxiliary losses and predictor side information can steer feature selection indirectly.

4.1 EXPLAINING EARLIER FINDINGS

  • A slowly changing score can be represented over a fast-moving ball, while resampling a visual cue or pattern can leave it poorly represented or ignored.The cited findings are explained by how features contribute to the selection objective; what counts as slow depends on the time step.Source: §4.1 · p. 5

4.2 PREDICTING FEATURE SELECTION

  • An unpredictable feature can still be represented when its inclusion reduces predictive leakage for other represented features.The traffic light is predicted to be represented whenever the cyclist’s features are represented in the described scenario.Source: §4.2 · p. 5
  • Strengthening predictive relationships within one feature set can eventually displace a disjoint competing set from limited latent capacity.The competing set can be displaced even if it remains unchanged.Source: §4.2 · p. 6
  • Expanding a mutually predictive feature set beyond latent capacity can reduce its representation, even when its features remain equally predictable.The expansion can lead to an entirely different set being represented.Source: §4.2 · p. 6
  • An auxiliary loss on some features can strengthen their predictive partners at the expense of competing features.The paper gives reconstructing wind measurements as an example that can favor representing a flag’s state.Source: §4.2 · p. 6
  • Supplying side information to the predictor can strengthen features whose next values it helps predict, without requiring that information in the latent.The paper gives wind measurements supplied to the predictor as an example that can favor representing a flag’s state.Source: §4.2 · p. 6

5 TESTING THE SELECTION PREDICTIONS

Controlled image-sequence experiments test whether JEPA feature selection changes with predictive relationships while holding architecture and training budget fixed. Probes measure which features are represented in a capacity-limited latent.

5.1 CROSS-SUBSIDIZATION: REPRESENTING UNPREDICTABLE FEATURES

  • As q increases, the unpredictable light becomes better represented because its current value helps predict which images persist.The next light is independently sampled, so its future remains unpredictable.Source: Figure 2 · p. 7, §5.1 · p. 7

5.2 SET SWITCHING

  • Increasing pD shifts representation toward C/D and away from A/B, although A/B’s predictive relationships remain unchanged.The two pairs do not help predict each other, and decoding within each pair is similar.Source: §5.2 · p. 7, Figure 3 · p. 8

5.3 FEASIBILITY DEPENDENCE

  • Competing-cycle decoding decreases at every increase in N, and at N = 16 A/B is better represented in all three seeds.Average copying probabilities stay unchanged as N grows, although tile scale and layout also change; one sensitivity configuration places the A/B increase within error bars.Source: §5.3 · pp. 7–8, Figure 4 · p. 8

6 TESTING THE STEERING PREDICTIONS

Steering experiments target C while keeping the predictive relationships and main latent setup fixed, then measure image-detection accuracy from the frozen main encoder. They test whether steering affects predictive partners as well as directly targeted features.

6.1 AN AUXILIARY RECONSTRUCTION OBJECTIVE

  • Reconstructing C alone can strengthen D’s representation and weaken A/B’s; at the largest reconstruction weight, D is decoded almost as accurately as C across pD values.The decoder is discarded before probing, and this experiment uses one optimization seed.Source: §6.1 · p. 9, Figure 5 · p. 9

6.2 PREDICTOR SIDE INFORMATION

  • Supplying C to the predictor generally increases D’s and A’s representation while decreasing B’s, a redistribution different from reconstruction.D’s representation increases for pD ≥0.75 but decreases at pD = 0.60; low main-latent decoding does not establish that information is unavailable to the predictor.Source: §6.2 · p. 9, Figure 6 · p. 10

6.3 IMPLICATIONS FOR REPRESENTATION LEARNING

  • If action is treated as side information, action-conditioned JEPA is predicted to favor features the agent affects, while genuine action-conditioning consequences remain untested.In the side-information experiment, supplying C strengthens D for pD = 0.75, 0.85, and 0.95.Source: §6.3 · pp. 9–10
  • Supplying an extracted feature can strengthen a predictive partner without directly supervising it, but steering should be evaluated against all downstream-needed features.Auxiliary objectives can target a feature that predicts the desired feature’s future or one whose future it predicts; side information supports only the former direction.Source: §6.3 · p. 10

7 CONCLUSION

The conclusion frames JEPA representation learning as selection based on the joint predictive value of feasible feature sets. Experiments support indirect steering through predictive partners, while evaluation at scale on natural data remains future work.

7 CONCLUSION

  • Predictive closure selects a feasible feature set maximizing mutual information between its current and next feature values.The set-level account agrees with earlier findings when features evolve independently and predicts new outcomes when they interact.Source: §7 · p. 10
  • Experiments match predictions that unpredictable features can become more decodable when they predict others, and that changing a competing set can weaken representation despite unchanged individual relationships.Auxiliary objectives and predictor side information can also strengthen predictive partners, making steering indirect.Source: §7 · p. 10
  • The experiments test effect directions, but measuring their size in large JEPAs trained on natural data remains future work.The conclusion also identifies possible relevance to biological representation-learning models linked to prediction.Source: §7 · p. 10

A FORMALIZATION OF PREDICTIVE CLOSURE

The formalization treats JEPA’s represented features as feasible sets, with selection governed by the set’s predictive information rather than fixed individual feature scores. It derives how feature relationships and predictor side information alter that objective.

A.1 DEFINITIONS AND ASSUMPTIONS

  • A feasible set fits the fixed latent entropy constraint; because the constraint is below observation entropy, no encoder can represent every feature.Adding a feature that contributes new information requires removing others to remain feasible.Source: §A.1 · p. 14

A.2 PREDICTION LOSS AND INFORMATION

  • Under the shared latent distribution and Gaussian prediction-error assumption, minimizing expected squared prediction error selects the encoder with greatest predictive information.Without the Gaussian assumption, minimizing prediction loss maximizes a lower bound on mutual information.Source: §A.2 · pp. 14–15

A.3 FROM REPRESENTATIONS TO FEATURE SETS

  • The latent and its complete represented feature set determine one another, so an ideally trained encoder represents a feasible set with maximal predictive information.This connects the representation-level objective to optimization over feature sets.Source: §A.3 · p. 15

A.4 ALEATORIC UNCERTAINTY AND PREDICTIVE LEAKAGE

  • Among feasible sets, maximizing predictive information is equivalent to minimizing aleatoric uncertainty plus predictive leakage.The derivation rewrites the set’s future uncertainty using uncertainty given the full observation and information about the future that leaks beyond the selected set.Source: §A.4 · p. 16

A.5 RELATION TO PER-FEATURE SELECTION

  • A feature’s contribution and entropy cost generally depend on the other selected features, making predictive information a non-modular set objective.The contribution includes what the feature’s current value adds about the set’s future and what the present set reveals about the feature’s future beyond its own future.Source: §A.5 · p. 16
  • When features evolve independently, selection reduces to ranking individual predictive information if their entropies match; with unequal entropies, it becomes a knapsack problem.In dependent examples, a traffic light contributes only when the selected set includes images whose persistence it controls, and a partial cycle retains only links whose endpoints are included.Source: §A.5 · p. 16

A.6 SIDE INFORMATION

  • Side information favors sets whose next values it predicts beyond their current values, and the objective minimizes conditional future uncertainty given both the set and side information.When side information comes from the observation, it reduces predictive leakage; when it is not determined by the observation, as with actions, it can also reduce aleatoric uncertainty.Source: §A.6 · p. 17

B EXPERIMENTAL SETUP AND EVALUATION DETAILS

The experiments use constructed CIFAR-10 collage sequences, train JEPA end to end, then freeze the encoder and assess feature recovery with probes. The traffic-light setup varies how strongly the light predicts which image changes.

B.1 SHARED TRAINING AND PROBING PROTOCOL

  • Across experiments, JEPA is trained from scratch on collage sequences with a three-dimensional latent, then frozen-latent probes are evaluated on newly sampled held-out images.Training uses 5,000 class-balanced training images and a 0.9Lpred + 0.1LSIGReg objective.Source: §B.1 · p. 17

B.2 TWO-IMAGE TRAFFIC-LIGHT SETUP AND READOUTS

  • The traffic-light experiment varies q from 0 to 1 while each image’s marginal persistence remains 0.5; the light identifies the replaced position with probability (1 + q)/2.Three runs are trained at each q value.Source: §B.2 · p. 18
  • A width-16 probe decodes the current light from the frozen current-frame latent, with 50% chance accuracy.Probe training, validation, and test sets contain 20,000, 5,000, and 5,000 examples, respectively.Source: §B.2 · p. 18

B.3 SETUPS AND READOUTS FOR THE TWO-PAIR AND CYCLE EXPERIMENTS

The two-pair and cycle experiments vary predictive links or group size while probing image information in frozen latents. Steering interventions either reward retaining C or supply C directly to the predictor.

B.3 SETUPS AND READOUTS FOR THE TWO-PAIR AND CYCLE EXPERIMENTS

  • The two-pair experiment varies only the probability that C is copied to D’s next value, while holding the other specified copy relationships fixed.The varied pD values are 0.60, 0.75, 0.85, 0.95, and 1.Source: §B.3 · p. 18
  • The group-size experiment changes the number of C/D identities while keeping average copy probabilities fixed and assessing whether the unchanged A/B group’s representation changes.Displayed C/D area stays within 1.2% across sizes, but tile scale and layout still change, so the control does not isolate representation cost from image-extraction difficulty.Source: §B.3 · p. 18
  • Position-specific probes detect target-image presence, but because negatives exclude targets everywhere, passing both probes does not establish that the latent distinguishes positions.The protocol averages accuracy over targets, then positions within groups, and reports means and standard errors across optimization seeds.Source: §B.3 · pp. 18–19

B.4 STEERING INTERVENTIONS

  • Steering either adds a reconstruction objective targeting C or supplies a separately encoded C image to the predictor while keeping the main latent three-dimensional.The side-information intervention adds encoder capacity; the auxiliary reconstruction decoder is discarded before probing.Source: §B.4 · p. 19

C NEXT-OBSERVATION PREDICTION

The evaluation tests whether a system’s prediction is closer to the actual next-observation latent than to an incompatible alternative. In every tested setting, systems identify the actual next observation above chance, while shuffled predictions remain near chance.

C NEXT-OBSERVATION PREDICTION

  • Predictions are scored by whether they are closer in squared Euclidean distance to the actual next latent than to an incompatible alternative, across 5,000 held-out transitions.Alternatives are drawn from the observation distribution and cannot follow the current observation; shuffled predictions serve as a control.Source: C NEXT-OBSERVATION PREDICTION · pp. 19–20
  • In every tested setting, systems identify the actual next observation well above chance, while shuffled predictions remain near chance.Table 1 reports percentage correct as mean ± standard error across three training runs, with 50% chance performance.Source: C NEXT-OBSERVATION PREDICTION · p. 20, Table 1 · p. 20

D REGULARIZATION SENSITIVITY

Across tested regularization weights, the qualitative findings for cross-subsidization, predictability-driven set switching, and steering recur. Group-size effects remain qualified: the enlarged group becomes less represented at smaller λ, but the switch favoring the competing group is inconsistent.

D REGULARIZATION SENSITIVITY

  • Across tested λ values, the qualitative findings for cross-subsidization, predictability-driven set switching, and steering through reconstruction or side information recur.At smaller λ, enlarging the group still reduces its representation, but the switch favoring the competing group appears in only one of three runs.Source: D REGULARIZATION SENSITIVITY · p. 20
  • As q increases, the current light becomes more strongly represented, although stronger regularization weakens this effect.The dashed lines mark 50% chance light-decoding accuracy.Source: Figure 7 · p. 20
  • Increasing pD from 0.60 to 1 switches the better-represented pair from A/B to C/D across all three runs and regularization weights.Stronger regularization reduces the separation between pairs, and at r = 0.20, mean C/D decoding accuracy falls from pD = 0.95 to pD = 1.Source: Figure 8 · p. 21
  • Enlarging C/D reduces its decoding accuracy at every tested size, run, and regularization weight.At N = 16, A/B is better represented in all three runs at r = 0.10 and 0.20, but in only one of three at r = 0.05.Source: D REGULARIZATION SENSITIVITY · p. 20, Figure 9 · p. 21

E WORST- AND BEST-POSITION READOUTS FOR GROUP SIZE

Worst- and best-position readouts preserve the endpoint reversal in group-size preferences: the competing cycle is favored at the smallest size, while A/B is favored at the largest. Across sizes, C/D accuracy falls consistently, but A/B and intermediate-size preferences are less uniform.

E WORST- AND BEST-POSITION READOUTS FOR GROUP SIZE

Related papers

Loading 2610.27364v1…