Clean Representation Consistency Loss

• Clean representation consistency loss is a regularization technique that enforces invariant latent features across data augmentations and perturbations.
• It integrates prediction, consistency via KL divergence, and curvature penalties to maintain robust, interpretable representations.
• Empirical studies show that each component is critical to prevent error accumulation and ensure effective model-based control and planning.

A clean representation consistency loss is an objective function or regularization mechanism that enforces agreement between neural network representations or outputs under varying conditions—such as paired inputs, data augmentations, latent dynamics, or different data modalities—with the goal of learning representations that are robust, interpretable, and well-suited for downstream tasks such as control, clustering, detection, or generative modeling. The underlying principle is that "clean" representations are those that maintain essential semantic or structural invariants across transformations, reconstructions, or function compositions; consistency losses explicitly enforce such invariances.

1. Principles and Motivation

Clean representation consistency losses arise in representation learning where the key goal is to ensure that encoded latent variables or feature mappings preserve meaningful relationships between the original observations and their perturbed, predicted, or augmented counterparts. The motivation includes:

• Reducing error accumulation or model drift in latent spaces for sequential control or dynamical systems.
• Improving robustness to label noise, data augmentation, or out-of-distribution shifts.
• Aligning representations so that class, cluster, or state assignments are invariant under intra-sample or inter-sample perturbations.
• Enabling effective use of control or clustering algorithms that assume local linearity or consensus among multiple solution variants.

A paradigmatic example is the Prediction, Consistency, and Curvature (PCC) loss for locally linear control (Levine et al., 2019), where a combination of predictive reconstruction, latent-observation consistency, and curvature regularization is required to learn a latent space that supports iterative LQR (iLQR) or similar algorithms.

2. Key Components in Clean Representation Consistency Losses

Prediction (Reconstruction/Fidelity)

The prediction loss component enforces that, given an input observation $x_t$ and action $u_t$, the model can accurately forecast the next observation $x_{t+1}$ through an encoder-decoder framework mediated by latent dynamics: $$E: x_t \to z_t,\ F: (z_t, u_t) \to \hat{z}_{t+1},\ D: \hat{z}_{t+1} \to \hat{x}_{t+1}$$ The loss,

$$R_3'(P) = -\mathbb{E}_{x,u,x'} \left[ \log \hat{P}\left(x' \mid x, u\right) \right]$$

drives the model to reconstruct future states with high fidelity.

Consistency (Latent-Observation Correspondence)

The consistency term penalizes discrepancies between predicted latent states and those inferred by encoding actual future observations. Typically, this is measured using a Kullback-Leibler divergence: $$R_2''(\hat{P}) = \mathbb{E}_{x,u,x'} \left[ D_{KL}\left(E(\cdot \mid x') \| F \circ E(\cdot \mid x, u)\right) \right]$$ This ensures that latent space transitions are “clean,” i.e., their evolution matches what would be inferred from raw data, reducing compounding errors when the model is rolled out through multiple planning/control steps.

Curvature (Local Linearity)

In settings where the downstream algorithm (e.g., iLQR) assumes linearity, a curvature penalty regularizes the nonlinearity of latent transitions by encouraging local affine behavior: $$R_{LLC}(\hat{P}) = \mathbb{E}_{x,u} \left[ \mathbb{E}_{(\epsilon_z,\epsilon_u) \sim \mathcal{N}(0, \delta^2 I)} \left\| f_\mathcal{Z}(z+\epsilon_z,u+\epsilon_u) - f_\mathcal{Z}(z,u) - (\nabla_z f_\mathcal{Z}(z,u)\cdot\epsilon_z + \nabla_u f_\mathcal{Z}(z,u)\cdot\epsilon_u) \right\|_2^2 \right]$$ This term is essential to ensure that locally linear control is effective in the learned latent space.

3. Amortized Variational Approximation

Direct minimization of the negative log-likelihood or KL divergence terms is generally intractable due to latent variable marginalization. To this end, amortized variational inference is used to derive tractable evidence lower bounds (ELBOs), typically with a factorized recognition distribution: $$Q(z_t, \hat{z}_{t+1}\mid x_t, x_{t+1}, u_t) = Q(\hat{z}_{t+1}\mid x_{t+1})\cdot Q(z_t\mid \hat{z}_{t+1},x_t,u_t)$$ Monte Carlo estimates over this variational family enable computation of both prediction and consistency terms, allowing end-to-end gradient optimization.

4. Bi-level and Weighted Optimization

The total clean representation consistency loss is formulated as a weighted sum of the prediction, consistency, and curvature components: $$\hat{P}^*\in\operatorname{argmin}_{\hat{P}} (\lambda_p R_3'(\hat{P}) + \lambda_c R_2''(\hat{P}) + \lambda_{cur} R_{LLC}(\hat{P}))$$ where $\lambda_p, \lambda_c, \lambda_{cur}$ are hyperparameters controlling the importance of each term. This formulation can be embedded within a bi-level optimization, with the inner loop fitting the latent model and the outer loop addressing downstream planning, imitation, or reinforcement objectives.

5. Empirical Findings and Loss Ablations

Experimental studies demonstrate that:

• The full PCC loss yields smooth, robust, and consistent latent representations with strong control performance under local linear controllers.
• Ablating the consistency term leads to severe performance degradation (e.g., in the Planar task: goal occupancy dropped from $\sim35.7\%$ to $0\%$), confirming its role in preventing error accumulation.
• Omitting the curvature loss similarly diminishes control effectiveness, reflecting the necessity of local affine transitions.
• Replacing expensive curvature computations with amortized Jacobian predictors can deliver computational savings with comparable downstream performance.

These results substantiate that each loss component is necessary to achieve “clean” and usable representations for model-based control.

6. Practical Significance and Downstream Applications

Clean representation consistency losses are crucial for:

• Learning latent spaces that are compatible with local optimal control (iLQR) and planning algorithms.
• Ensuring that errors do not accumulate during rollouts by aligning predicted latent states to encodings of actual observations.
• Facilitating stable and reproducible training regimes for high-dimensional, partially observed optimal control domains.
• Potentially generalizing to other areas, such as sequential decision-making, unsupervised embedding, or multitask learning, where internal consistency under defined transformations or compositionality is critical.

While computationally more intensive, especially with non-amortized curvature regularization, the resulting benefits in terms of downstream task fitness, stability, and robustness substantiate the value of this approach.

7. Theoretical and Methodological Extensions

Amortized variational bounds render the computation of latent consistency losses tractable. Research directions include:

• Refinement of the variational family to better capture the posterior structure.
• Extension to stochastic latent dynamics and risk-sensitive objectives.
• Exploration of alternative regularizers for non-Euclidean or manifold-constrained latent spaces.
• Analysis of scalability and representation quality as observation and action dimensionality increases.

The principles of clean representation consistency, as originally formulated in the PCC framework, now inform a broad class of algorithms in representation learning, model-based reinforcement learning, and beyond (Levine et al., 2019).