The Expressive Power of Low Precision Softmax Transformers with (Summarized) Chain-of-Thought
Abstract
Existing expressivity results for transformers typically rely on hardmax attention, high precision, and other architectural modifications that disconnect them from the models used in practice. We bridge this gap by analyzing standard transformer decoders with softmax attention and rounding of activations and attention weights, while allowing depth and width to grow logarithmically with the context length. As an intermediate step, we construct hardmax transformers with ternary activations and well-separated attention scores that simulate Turing machines using Chain-of-Thought (CoT). This lets us convert the constructions to equivalent softmax transformers without the unrealistic parameter magnitudes or activation precision that prior approaches would require. Using the same technique, we analyze a recently proposed summarized CoT paradigm and show that it simulates Turing machines more efficiently, with model size scaling logarithmically in a space bound rather than a time bound. We empirically test predictions made by our results on a Sudoku reasoning task and find better alignment with learnability than for prior high-precision results. Our code is available at https://github.com/moritzbroe/transformer-expressivity.
Lay Summary
AI language models like ChatGPT are built on neural networks called transformers. Theorists study what such models can in principle compute, but their analyses typically rely on unrealistic simplifications, assuming much higher numerical precision and more idealized internal operations than real systems use. This leaves it unclear whether those theoretical results say anything about deployed models. We work towards closing this gap by proving that realistic transformers can solve any problem a computer can solve if they explicitly write out their steps. We further analyze a more efficient variant in which the model periodically summarizes its progress, keeping the amount of memory it needs small. Experiments training transformers to solve hard Sudoku puzzles confirm that our realistic theory aligns well with what these models actually learn.