Transformer-like Inference from Optimal Control
Provides a theoretical foundation for transformer architectures from first principles, offering a new perspective for understanding and potentially improving transformer models.
The paper derives transformer-like inference architectures from optimal control theory, showing that decoder-only transformer layers emerge as solutions to prediction problems. Numerical experiments reveal that transformers exploit non-Markovian structure when embedding dimension is insufficient.
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.