Transformer-like Inference from Optimal Control

TL;DR

This paper derives transformer-like inference architectures from optimal control principles, providing a theoretical foundation and explicit algorithms that mirror transformer layer structures.

Contribution

It introduces a novel framework connecting transformers to optimal control theory, deriving inference algorithms with layer structures similar to transformers from first principles.

Findings

Optimal control-based inference algorithms match transformer layer structures.

Transformers implicitly exploit non-Markovian structure when embedding dimension is limited.

Numerical experiments compare optimal control solutions with trained transformer attention weights.

Abstract

Decoder-only transformers compute the conditional probability of the next token from a sequence of past observations. This paper derives, from first principles, inference architectures that solve the same prediction problem - and in doing so, recovers transformer-like layer operations as a consequence of optimal control theory. The framework is developed for two model classes: a nonlinear model of discrete-valued processes, directly motivated by the transformer, and a linear Gaussian model as a tractable baseline. For both model classes, the prediction objective is reformulated as an optimal control problem whose solution yields an explicit inference algorithm, the dual filter, with a layer structure that mirrors the layer structure of a decoder-only transformer. Numerical experiments provide a comparison of the optimal control to attention weights from a trained transformer. These experiments reveal that when the embedding dimension is insufficient, the transformer implicitly exploits non-Markovian structure.