Leakage-Free Predictive Coding

• Leakage-free predictive coding is an information-theoretic framework that removes decoders to prevent leakage of crucial dynamics in latent representations.
• It employs a contrastive predictive coding objective to maximize latent-space mutual information, ensuring predictive sufficiency for efficient control.
• The method demonstrates state-of-the-art control performance with faster training compared to traditional high-dimensional reconstruction approaches.

Leakage-free predictive coding is an information-theoretic control framework for high-dimensional optimal control that removes the observation decoder from the learning pipeline, thereby eliminating information leakage often encountered in prior latent dynamics approaches. This paradigm is designed to construct a latent, controllable representation of observations for decision-making, where the dynamical system can be efficiently controlled directly in the latent space. Leakage-free predictive coding maintains predictive sufficiency of the latent variables without reconstructing high-dimensional observations, which traditionally introduces nuisance parameters and may hinder control performance. The method is instantiated through a specific training objective combining predictive coding with controllability regularizers, and achieves superior empirical results on standard control benchmarks (Shu et al., 2020).

1. Information-Theoretic Leakage-Free Objective

Leakage-free predictive coding discards the conventional VAE-style objective that relies on an explicit decoder $D$ for reconstructing the next observation $x_{t+1}$ from the latent state $z_{t+1}$. Instead, the framework optimizes the encoder $E$ and latent dynamics $F$ such that only the information relevant to predicting the next latent state is maintained.

The predictive-suboptimality of a representation $E$ is defined as

$$\ell^*(E) = \min_{q} \mathbb{E}_{p(x_t, x_{t+1}, u_t)} \Big[ D_{\mathrm{KL}} ( p(x_{t+1}|x_t,u_t) \;\|\,q(x_{t+1}|x_t,u_t) ) \Big],$$

where $q$ can only access the history via $(Z_t, U_t)$ (where $Z_t = E(X_t)$). Lemma 1 states that

$x_{t+1}$0

As $x_{t+1}$1 is environment-dependent and fixed, maximizing the mutual information $x_{t+1}$2 pushes the suboptimality to zero. This is operationalized via a contrastive predictive coding bound, with the objective

$x_{t+1}$3

over a batch of size $x_{t+1}$4, with no need for a decoder (Shu et al., 2020).

2. Decoder-Free Principle and Information Retention

The elimination of the decoder $x_{t+1}$5 is justified by the result that predictive sufficiency can be achieved by maximizing latent-space mutual information, without reconstructing observations. In standard dynamical-VAE settings, the decoder introduces a large parameter set used solely for the reconstruction objective and ignored at deployment. Leakage-free predictive coding leverages the mutual information upper bound to guarantee that, by sufficiently maximizing $x_{t+1}$6, the latent embedding loses at most $x_{t+1}$7 in KL-risk relative to modeling $x_{t+1}$8. In the limit, maximizing this quantity recovers the full predictive power of the transition model. Thus, enforcing a high mutual information in latent space is information-theoretically equivalent to maximizing next-observation likelihood, but avoids the intractable task of high-dimensional reconstruction.

3. Structure of Latent Dynamics and Local Linearization

Latent dynamics within leakage-free predictive coding are parameterized as a conditional Gaussian,

$x_{t+1}$9

where $z_{t+1}$0 is a compact MLP. This design ensures that the latent dynamics remain amenable to first-order local linearization, as required for efficient second-order optimal control strategies such as iLQR. At control time, iterative LQR operates in latent space by locally linearizing the learned dynamics via the Jacobians $z_{t+1}$1 and $z_{t+1}$2, which is enabled by the gently curved, low-dimensional nature of $z_{t+1}$3. The control problem is thus reduced to solving

$z_{t+1}$4

via repeated local linearizations (Shu et al., 2020).

4. Assumptions, Regularization, and Controllability Constraints

The approach presumes Markovian state transitions ($z_{t+1}$5) and bounded, Lipschitz continuous cost $z_{t+1}$6. Multiple regularizers are adopted to ensure that the latent space is both predictive and suitable for control:

• Consistency Regularizer: Latent-dynamics MLE, maximizing

$z_{t+1}$7

aligns the learned dynamics with the actual transitions in the latent space.

• Curvature Penalty: To guarantee reliable local linearization, a Taylor-remainder penalty,

$z_{t+1}$8

is enforced for Gaussian perturbations $z_{t+1}$9.

• Latent-Noise Stabilization: Fixed Gaussian noise $E$0 is injected into $E$1 to prevent the encoder from collapsing and to ensure proper spread in the latent space.

The full training objective takes the form

$E$2

with hyperparameters $E$3 tuned empirically on small domains (Shu et al., 2020).

5. Empirical Performance Relative to LCE Baselines

Leakage-free predictive coding, exemplified by the PC³ model, demonstrates superior control performance on several standard benchmarks when compared to latent control approaches that reconstruct pixels (PCC) and globally linear latent models with online updates (SOLAR). The average fraction of time spent in a "goal" region, aggregated over multiple seeds and tasks, is summarized below:

Task	PC³ (all models)	PCC (all)	SOLAR (all)
Planar	74.35 ± 0.76 %	56.60 ± 3.15 %	68.00 ± 3.80 %
Balance	99.12 ± 0.66 %	91.90 ± 1.72 %	67.00 ± 2.60 %
Swing-Up	58.40 ± 3.53 %	26.41 ± 2.64 %	35.40 ± 1.90 %
Cartpole	96.26 ± 0.95 %	94.44 ± 1.34 %	91.20 ± 5.40 %
3-Link	42.40 ± 3.23 %	14.17 ± 2.20 %	0.00 ± 0.00 %

In addition to higher task completion rates, PC³ is approximately 1.8 times faster to train than PCC and roughly 50 times faster than SOLAR, attributed to the absence of a high-dimensional pixel decoder (Shu et al., 2020).

6. Implications and Significance

Leakage-free predictive coding resolves the critical issue of information leakage typically present when using decoders in learning latent models for control. By eliminating the decoder and maximizing mutual information in the latent space via contrastive predictive coding, the framework preserves all predictive information relevant for next-step dynamics. This approach yields a well-structured, controllable latent space compatible with locally linear control techniques and achieves state-of-the-art empirical results on nontrivial control domains. This suggests a fundamental shift in latent dynamical model design, prioritizing information retention for control rather than high-fidelity reconstruction of observations (Shu et al., 2020).