Guaranteed behaviour of a small trained transformer at every input length: a machine-checked proof of concept

Abstract

The behaviour of a trained neural network is usually estimated by testing; this paper is a proof of concept for guaranteeing it by a machine-checked proof about the trained weights. The network is a one-layer, one-head, attention-only transformer without positional embeddings, trained at length 10 to output the largest of its input tokens. We prove in Lean 4, using only Lean's standard axioms, statements about what it computes on inputs of every length. The certified model has weights rounded to multiples of 2^-10 and exact real-valued softmax; the proofs work from its attention and value tables and never enumerate inputs. First, the model is correct whenever the other tokens are few enough and far enough below the maximum: for example, on any number of copies of a maximum of at least 23 together with at most 10^6 other tokens, each at least 16 below it. Second, it is provably wrong on explicit inputs, among them inputs of the training length, although it answers all of 2^19 random inputs of that length correctly. Third, the two kinds of statement meet in exact thresholds: on k copies of 20 followed by 21 the model answers 21 if and only if k <= 4, and in each of five training seeds, at every gap g <= 6 and for at least 39 of the 40 maxima m from 24 to 63, there is an explicit N such that the model is correct with up to N other tokens at most m-g and wrong on every larger number of copies of m-g. All theorems are checked by the Lean kernel from the weight matrices up, in under an hour per model on a desktop computer. Code, weights and proofs: https://github.com/Oliverds321/transformer-all-lengths-lean.

Keywords