DSReg rotates mixed JEPA features into individual latent factors

Yujia Zheng, David Klindt, Randall Balestriero, and Bernhard Schölkopf propose DSReg, a lightweight regularizer that disentangles representations from a frozen LeJEPA encoder. Under the paper’s assumptions, the method recovers individual latent factors without labels, a decoder, or an end-to-end reconstruction objective.

Representation learners often preserve the information needed for downstream tasks while distributing each real-world factor across many coordinates. A supervised linear probe can recover useful combinations, but coordinate-level editing and sparse control work better when one axis corresponds to one factor. DSReg addresses that alignment problem by finding a rotation of the learned representation.

Why JEPA coordinates remain mixed

Joint-Embedding Predictive Architectures train an encoder to predict representations of related observations, such as different views or moments from the same scene. The cited LeJEPA result establishes linear identifiability in the Gaussian-latent setting: the learned representation contains the true latent state up to an invertible linear transformation.

Linear identifiability preserves information while leaving the coordinate system ambiguous. Position, color, velocity, and texture can remain mixed across dimensions, even when a linear map could recover them with sufficient labels. DSReg narrows that ambiguity to signed permutation, meaning coordinates may be reordered or flipped in direction but are no longer arbitrary mixtures.

Distinct footprints identify the factors

DSReg relies on an assumption the authors call Structural Diversity. Each latent factor must affect a distinct set of observed variables, such as pixels or sensor readings. That set forms the factor’s dependency footprint and corresponds locally to the nonzero entries in a Jacobian column.