Dual Filter: A Transformer-like Inference Architecture for Hidden Markov Models

1. Problem Statement

The paper addresses the problem of causal nonlinear prediction for systems where observations are generated by an underlying Hidden Markov Model (HMM).

• Goal: Compute the conditional probability of the next token (observation), $P(Z_{T+1} = z \mid Z_1, \dots, Z_T)$, given a finite sequence of past tokens.
• Context: While modern Decoder-only Transformers excel at this task in generative AI, their internal mechanics are often viewed as "black boxes." The authors aim to derive a mathematically rigorous, transformer-like architecture from first principles (specifically optimal control theory) that solves this prediction problem for HMMs, rather than modeling the transformer itself.
• Distinction: The focus is strictly on inference (computing the prediction given a model), not on learning the model parameters from data.

2. Methodology

The authors propose a framework based on Optimal Control Theory and Duality, leading to a new algorithm called the Dual Filter.

A. Mathematical Formulation

1. Nonlinear Predictor Representation: The authors seek a representation for the conditional measure (posterior) $\pi_T$ in the form of a causal, nonlinear predictor:
   $$\pi_T(F) = (\text{constant}) - \sum_{t=1}^T U_{t-1}^T e(Z_t)$$
   Here, $U$ is a stochastic process (random weights) adapted to past observations, and $e(\cdot)$ is a mapping of tokens to vectors (analogous to one-hot encoding). Unlike linear predictors (Wiener filters) where weights are deterministic, these weights depend on the observation history, making the predictor nonlinear.

2. Optimal Control Reformulation:

   • The Minimum Mean-Squared Error (MMSE) estimation problem is reformulated as an Optimal Control Problem (OCP).
   • Dynamic Constraint: The system is modeled using a Backward Stochastic Difference Equation (BS$\Delta$E). This equation links the terminal condition (the target function to be estimated) backward in time to the initial state, driven by control inputs $U$ and martingale increments.
   • Duality Principle: A key theoretical contribution is establishing a duality between the nonlinear filtering problem and the optimal control problem. The authors prove that minimizing the control cost function is equivalent to minimizing the MMSE.

3. The Dual Filter Algorithm:

   • Solving the OCP yields an explicit formula for the optimal control input $U^{(opt)}_t$ as a function of the current conditional measure $\pi_t$.
   • This leads to a fixed-point equation on the space of probability measures.
   • The authors define a layer transformation $\mathcal{N}^{(\text{dfltr})}$ that maps a sequence of measures to a new sequence. The true conditional measure $\pi$ is the fixed point of this map: $\pi = \mathcal{N}^{(\text{dfltr})}(\pi)$.
   • Algorithms: Two algorithms are proposed to approximate this fixed point:
     1. Iterative Algorithm: Repeatedly applies the layer transformation (structurally similar to Transformer layers).
     2. Single-Shot Algorithm: Computes the solution in a single pass by leveraging the recursive structure of the backward equation, offering higher efficiency ($O(d^2 T)$ complexity, independent of vocabulary size $m$).

3. Key Contributions

1. New Representation: Derivation of a novel representation for causal nonlinear predictors (Eq. 1) based on optimal control inputs, generalizing the Wiener filter to the nonlinear HMM setting.
2. Duality Theory for HMMs: Establishment of a "Minimum Variance Duality" for discrete-time HMMs with discrete observations. This resolves a long-standing open problem in control theory, extending Kalman's duality (previously limited to Linear-Quadratic-Gaussian settings) to nonlinear, discrete systems.
3. Dual Filter Architecture: Introduction of the Dual Filter, an iterative algorithm that solves the fixed-point equation. It provides a mathematically grounded, transformer-like architecture derived from first principles.
4. Structural Correspondence: A detailed mapping between the Dual Filter and Decoder-only Transformers:
   • Input: Token embeddings correspond to initial measure estimates.
   • Layers: The Dual Filter's layer transformation $\mathcal{N}^{(\text{dfltr})}$ parallels the Transformer's self-attention layer $\mathcal{N}^{(\text{xfer})}$.
   • Mechanism: The paper argues that the "attention" mechanism in transformers can be interpreted as a specific type of nonlinear predictor derived from the control-theoretic duality, though the specific implementation differs (Dual Filter uses explicit Markov dynamics; Transformers learn attention weights).

4. Results

The paper validates the theory through numerical experiments using parameters mimicking large-scale transformer models (State dimension $d=384$, Vocabulary $m=65$, Sequence length $T=256$).

• Accuracy: The Dual Filter (both iterative and single-shot) successfully recovers the ground-truth conditional probabilities generated by the standard nonlinear filter (forward algorithm).
• Control Input Analysis:
  • When the HMM dynamics are highly mixing (rapid state changes), the optimal control input is sparse, relying only on recent observations.
  • When dynamics are deterministic or have long-range correlations, the control input spans a longer time horizon, effectively capturing long-term dependencies.
• Comparison with Transformers (nanoGPT):
  • The authors trained a nanoGPT model on data generated from the same HMM.
  • Favorable Case: When the emission matrix creates a "low-complexity" prediction task (sparse output distribution), the Transformer learns to approximate the nonlinear filter well.
  • Unfavorable Case: For more complex emission structures, the Transformer failed to converge to the ground truth, whereas the model-based Dual Filter remained accurate. This suggests Transformers may struggle with certain HMM structures without specific architectural inductive biases or sufficient data.

5. Significance and Future Directions

• Theoretical Bridge: The paper provides a rigorous mathematical link between classical stochastic control/filtering and modern deep learning architectures (Transformers). It suggests that Transformers can be viewed as learning approximations of optimal control laws for underlying latent processes.
• Interpretability: By deriving the architecture from first principles, the Dual Filter offers a transparent alternative to the "black box" nature of Transformers, explaining why certain structures (like causal masking and layer-wise transformations) are effective for prediction.
• Future Work:
  • Learning: Extending the framework to learn the model parameters $(\mu, A, C)$ or the layer transformation directly from data (history-based approaches).
  • Stability Analysis: Analyzing the contraction properties of the Dual Filter map to ensure stability in long sequences.
  • Generalization: Moving beyond HMMs to general non-Markovian processes and Partially Observed Markov Decision Processes (POMDPs).

In summary, this paper presents a foundational theory that reinterprets the Transformer's prediction capability through the lens of optimal control and duality, offering a new, mathematically tractable algorithm (the Dual Filter) that matches Transformer performance in specific settings while providing deeper theoretical insight.