Dual Filter: A Mathematical Framework for Inference using Transformer-like Architectures

TL;DR

The paper develops a mathematically principled framework for causal nonlinear prediction in finite-HMM settings and connects it to decoder-only transformers. By reformulating MMSE prediction as an optimal-control problem with a backward dynamic constraint (BS$\Delta$E), it derives an explicit optimal-control law and a corresponding fixed-point operator on probability measures, called the dual filter. This yields a transformer-like iterative algorithm and a direct correspondence between the dual-filter updates and transformer layers, offering a principled view of attention as a data-dependent control mechanism. Numerical experiments with transformer-scale parameters illustrate the algorithm’s predictive performance and illuminate how spectral properties of the hidden dynamics shape the optimal control and predictive behavior.

Abstract

This paper presents a mathematical framework for causal nonlinear prediction in settings where observations are generated from an underlying hidden Markov model (HMM). Both the problem formulation and the proposed solution are motivated by the decoder-only transformer architecture, in which a finite sequence of observations (tokens) is mapped to the conditional probability of the next token. Our objective is not to construct a mathematical model of a transformer. Rather, our interest lies in deriving, from first principles, transformer-like architectures that solve the prediction problem for which the transformer is designed. The proposed framework is based on an original optimal control approach, where the prediction objective (MMSE) is reformulated as an optimal control problem. An analysis of the optimal control problem is presented leading to a fixed-point equation on the space of probability measures. To solve the fixed-point equation, we introduce the dual filter, an iterative algorithm that closely parallels the architecture of decoder-only transformers. These parallels are discussed in detail along with the relationship to prior work on mathematical modeling of transformers as transport on the space of probability measures. Numerical experiments are provided to illustrate the performance of the algorithm using parameter values used in researchscale transformer models.

Figures (6)

• Figure 1: Graphical model for the HMM for $T=1$.
• Figure 2: Eigenvalues of the transition matrix $A$ as a function of the homotopy parameter $\alpha$. (Middle): Plot of the second eigenvalue magnitude $|\lambda_2|$ as a function of $\alpha$. (Left): Eigenvalue spectrum for $A = A^{\text{(stoch)}}$ ($\alpha = 0$). (Right): Eigenvalue spectrum for $A = A^{\text{(circ)}}$ ($\alpha = 1$). In the middle plot, three representative values of $\alpha$---0.3 (green), 0.9 (orange), and 1.0 (blue)---are highlighted. Insets show the corresponding eigenvalue spectra and matrix structures; darker shades indicate higher matrix entry values.
• Figure 3: Comparison with the single-shot algorithm. (Top): Time traces of the error $\{\varepsilon_t^{(\ell)} : 1 \leq t \leq T\}$ for $\ell = 0, 1$. The dashed line corresponds to the initial error ($\ell = 0$), while the solid line shows the error after one iteration ($\ell = 1$). (Bottom): Top-ten conditional probabilities $\{p_t^{(\ell)}(z_i) : i = 1, \ldots, 10\}$ for five representative time points. For each $t$, these are obtained by sorting the conditional probability vector and selecting the top ten. The gray shading indicates the ground truth (from the nonlinear filter), and the solid line shows the result from the dual filter.
• Figure 4: Comparison with the iterative algorithm: (Top): Time traces of the error $\{\varepsilon_t^{(\ell)} : 1 \leq t \leq T\}$ for $\ell=1,2,\hdots,L$ (with $L=6$). The dashed line corresponds to the initial error ($\ell = 0$), while the solid lines show the error after subsequent iterations, with darker shades used as $\ell$ increases. (Bottom): Top-ten conditional probabilities $\{p_t^{(\ell)}(z_i) : i = 1, \ldots, 10\}$ for five representative time points. For each $t$, these are obtained by sorting the conditional probability vector and selecting the top ten. The gray shading indicates the ground truth (from the nonlinear filter), and the solid line shows the result from the dual filter.
• Figure 5: Optimal control inputs as a function of $|\lambda_2|$. (Top): $|\lambda_2| = 0.3$ (green), (Middle): $|\lambda_2| = 0.9$ (orange), (Bottom): $|\lambda_2| = 1.0$ (blue). These three cases correspond to different values of the homotopy parameter $\alpha$. Note that the x-axis (time $t$) is plotted on a non-linear scale. The right-hand panels show the top ten conditional probabilities $\{p_T^{(\ell)}(z_i): i = 1, \ldots, 10\}$ at the terminal time $t = T$. These are obtained by sorting the conditional probability vector and selecting the top ten.

Theorems & Definitions (35)

• Example 1: m=1
• Remark 1: Linear vs nonlinear predictors
• Remark 2: Input (tokens)
• Remark 3: Output (prediction at time $T$)
• Remark 4: Predictions for intermediate times
• Remark 5: Inference and learning
• Proposition 6
• Remark 7
• Remark 8
• Remark 9