Teacher Forcing as Generalized Bayes: Optimization Geometry Mismatch in Switching Surrogates for Chaotic Dynamics

Abstract: Identity teacher forcing (ITF) enables stable training of deterministic recurrent surrogates for chaotic dynamical systems and has been highly effective for dynamical systems reconstruction (DSR) with recurrent neural networks (RNNs), including interpretable almost-linear RNNs (AL-RNNs). However, as an intervention-based prediction loss (and thus a generalized Bayes update), teacher forcing need not match the free-running model's marginal likelihood geometry. We compare the objective-induced curvatures of ITF and marginal likelihood in a probabilistic switching augmentation of AL-RNNs, estimating ambiguity-aware observed information via Louis' identity. In the switching setting studied here, conditioning on a single forced regime path (as ITF does) inflates curvature, while marginal likelihood curvature is reduced by a missing-information correction when multiple switching explanations remain plausible. In Lorenz-63 experiments, windowed evidence fine-tuning improves held-out evidence but can degrade dynamical quantities of interest (QoIs) relative to ITF-pretrained models.

• The paper demonstrates that teacher forcing induces a generalized Bayes update, leading to a significant geometry mismatch in model parameter space.
• Experiments with AL-RNN and PAL-RNN reveal that increased latent switching ambiguity under ITF can inflate curvature estimates by up to three orders of magnitude.
• Optimizing marginal likelihood may degrade long-term dynamical qualities, suggesting the need for QoI-aware, post-Bayesian objective functions.

Teacher Forcing as Generalized Bayes: Geometry Mismatch in Switching Surrogates for Chaotic Dynamics

Introduction and Motivation

The paper "Teacher Forcing as Generalized Bayes: Optimization Geometry Mismatch in Switching Surrogates for Chaotic Dynamics" (2604.25904) undertakes a rigorous investigation into the geometric implications of teacher forcing (TF), particularly its identity variant (ITF), when used to train recurrent surrogates for chaotic dynamical systems. The study operates within the context of dynamical systems reconstruction (DSR), where the aim is to model unknown chaotic systems from time series data in a way that preserves long-term invariant properties, specifically dynamical quantities of interest (QoIs) such as Lyapunov exponents and attractor geometry.

A fundamental challenge in DSR for chaotic systems is that short-horizon prediction losses, while stabilizing training, are not consistent with the maximum likelihood (ML) objective of the unforced, free-running model. The core claim of the paper is that teacher-forced objectives define a generalized Bayes update that induces a geometry in parameter space (curvature), which can substantially mismatch the observed information geometry associated with genuine ML training on probabilistic state-space models with latent switching variables.

Problem Formulation and Theoretical Foundations

The study builds on the almost-linear RNN (AL-RNN) family, which features a piecewise affine transition with gate-induced switching between linear regimes. Specifically, the AL-RNN is formulated as:

$$\bm{z}_{t+1} = F_{\bm\theta}(\bm{z}_t) = \bm{A}\bm{z}_t + \bm{W}\,\bm{\phi}^*(\bm{z}_t) + \bm{h},$$

where the nonlinearity is consolidated into $P$ ReLU-gated coordinates, endowing the system with explicit symbolic switching codes $\bm{c}_t$ that identify the active linear regime.

Identity Teacher Forcing (ITF) is utilized as a training intervention, where observed coordinates in the latent state are periodically overwritten with true data points, stabilizing gradients but breaking alignment with the target model’s marginal likelihood geometry.

The authors precise that the ITF-modified objective is a generalized Bayes update—manifesting a “Gibbs posterior”—where the induced loss function conditions on a particular intervention (forced regime path), rather than marginalizing over latent regime ambiguities as in probabilistic state-space models.

Figure 1: Directed graphical model for PAL-RNN, showing continuous latent states and discrete stochastic gates governed by a probit distribution.

The AL-RNN is then probabilistically augmented with a Probit-gated Switching SSM (PAL-RNN), in which gate variables are treated as latent random variables (Bernoulli distributed via a probit link), and the system follows a switching linear-Gaussian SSM specification. This enables estimation of marginal likelihood curvature via the observed information, implementable with Rao–Blackwellized particle (RBPF) methods and Louis’ identity. Crucially, Louis’ identity decomposes observed Fisher information into a “complete-data” contribution (as if the switching path was known) minus a “missing-information” term that captures the reduction in certainty due to latent regime ambiguity.

Empirical Characterization of Geometry Mismatch

Controlled Toy Model: Curvature-Entropy Relationship

The initial set of experiments utilizes a probit-gated switching AR(1) model with a tunable gate noise parameter $\sigma_g$ to induce varying levels of posterior switching ambiguity. As gate entropy increases, the missing-information ratio (MIR) escalates—reflecting a large reduction of curvature in the marginal likelihood geometry—while the observed curvature sharply decays.

Figure 2: Increasing gate noise $\sigma_g$ leads to higher posterior switching ambiguity, increased MIR, and diminished observed curvature in a probit-gated switching model.

AL-RNN/PAL-RNN Case Study: Lorenz-63 Dynamics

The main empirical study is conducted on AL-RNN surrogates trained with ITF on Lorenz-63 chaotic trajectories, subsequently extended to PAL-RNNs with stochastic switching. The authors estimate both the ITF-aligned curvature and the ambiguity-aware observed information for the drift parameters, varying process and observation noise levels.

Curvature Gap Quantification

The curvature gap

$$g_{\bm Q} = \log_{10}\left(T\,\mathrm{tr}(\bm{\mathcal I}_{\mathrm{ITF}})/\,\mathrm{tr}(\bm{\mathcal I}_{\text{obs}})\right)$$

is introduced to compare the “sharpness” of ITF-induced (intervention-based) curvature and ambiguity-aware observed information, with higher $g_{\bm Q}$ indicating increased overconfidence under ITF. The curvature inflation under ITF is shown to scale with the entropy (ambiguity) of the posterior switching code, exhibiting gaps up to three orders of magnitude.

Figure 3: Matrix-aware diagnostics show direction-dependent and anisotropic curvature mismatch, with ITF-aligned Fisher sometimes greatly exceeding ambiguity-aware curvature in selected subspaces.

Practical Implications: Evidence-QoI Misalignment

A key experimental result is that optimizing the windowed marginal likelihood via Particle-SAEM, starting from ITF-pretrained checkpoints, generally increases held-out evidence but can degrade performance on dynamical QoIs such as attractor geometry and Lyapunov exponents—often driving models away from chaos towards overly stable dynamics.

Figure 4: Attractor reconstructions demonstrate qualitative collapse of chaotic geometry following evidence-maximizing fine-tuning; the forced model preserves complexity, but likelihood-tuned models can yield degenerate, limit-cycle-like behavior.

Figure 5: RBPF diagnostics on a representative window highlight well-maintained effective sample size and stable particle diversity during inference.

Theoretical and Practical Implications

Geometry-Objective Dependence

Teacher forcing, as an intervention-based optimization, can induce substantially sharper (higher curvature) local parameter geometry than ambiguity-aware marginal likelihoods, especially in regimes of high switching uncertainty. This geometric mismatch has direct consequences for local identifiability, uncertainty quantification (e.g., Laplace approximations), and active experimental design. Importantly, there is no canonical posterior in DSR for chaotic systems; the geometry of uncertainty is inherently objective-dependent.

Evidence-Quality Divergence

The paper provides numerical evidence that maximizing windowed evidence can be anti-aligned with preserving long-horizon dynamical invariants. This raises limitations in classical (Bayesian or generalized Bayes) approaches that treat evidence or posterior precision as the sole criterion for learning or model selection in chaotic DSR.

Guidelines for Post-Bayesian DSR

The findings motivate a QoI-aware, post-Bayesian approach: learning objectives and geometry should be tailored to downstream DSR tasks, not solely determined by marginal likelihoods or prediction error. Intervention-aligned curvature proxies, like the ITF-aligned Fisher matrix, may better reflect practical sensitivities for long-horizon behavior in chaotic systems, and geometry-aware active data collection and experiment design should explicitly account for the “objective–geometry gap.”

Future Directions

• Extending geometry mismatch analysis beyond AL-RNNs to more general or highly nonlinear surrogates.
• Development of scalable, ambiguity-aware curvature summaries or approximations that can be integrated into uncertainty quantification, model calibration, and optimal experimental design.
• Augmentation of likelihood-based fine-tuning with explicit constraints on invariant measures or other physical statistics to avoid degradation of dynamical QoIs.
• Investigation of forcing strategies (e.g., GTF vs. STF/ITF) to interpolate between training stability and geometric fidelity with respect to the target marginal likelihood.
• Generalization to partial-observation settings, richer noise structures, and broader classes of chaotic and non-chaotic systems.

Conclusion

This paper rigorously demonstrates the geometric consequences of intervention-based training in chaotic dynamical surrogate modeling. The mismatch between teacher-forcing-induced and marginal-likelihood-induced curvatures, especially in the presence of latent switching ambiguity, has concrete effects on local identifiability, uncertainty estimation, and the relationship between prediction, evidence, and long-term dynamical invariants. These findings argue for principled, QoI-oriented approaches to optimization and inference in DSR, particularly when reconstructing chaotic systems where long-horizon behavior is of paramount interest.