Provably Overwhelming Transformer Models with Designed Inputs

TL;DR

This work introduces the notion of \'overwhelming\' for trained transformers, wherein a fixed input prefix $s$ and final token $q$ yield outputs invariant to free suffixes $t$ of length up to $n_{free}$. It develops a two-stage, rigorous framework based on worst-case deviation $\mathcal{W}$ and peak-to-peak difference $PTP$, plus practical algorithms to certify overwhelm for a given model, input prefix, and context length; it further refines bounds under permutation-invariant inputs using LP relaxations. The authors demonstrate the approach on a concrete single-layer, decoder-style transformer with RoPE, LN, and an MLP, providing both theoretical bounds and empirical validation on the AG News dataset with various input-restriction scenarios; they show that overcoming $W$ relative to $PTP$ certifies that the next-token output is fixed regardless of allowable free tokens. This work is a stepping stone toward provable guarantees for trained transformers, with potential applications in safety evaluations, prompt-engineering no-go results, and the study of model complexity in constrained contexts.

Abstract

We develop an algorithm which, given a trained transformer model $\mathcal{M}$ as input, as well as a string of tokens $s$ of length $n_{fix}$ and an integer $n_{free}$, can generate a mathematical proof that $\mathcal{M}$ is ``overwhelmed'' by $s$, in time and space $\widetilde{O}(n_{fix}^2 + n_{free}^3)$. We say that $\mathcal{M}$ is ``overwhelmed'' by $s$ when the output of the model evaluated on this string plus any additional string $t$, $\mathcal{M}(s + t)$, is completely insensitive to the value of the string $t$ whenever length($t$) $\leq n_{free}$. Along the way, we prove a particularly strong worst-case form of ``over-squashing'', which we use to bound the model's behavior. Our technique uses computer-aided proofs to establish this type of operationally relevant guarantee about transformer models. We empirically test our algorithm on a single layer transformer complete with an attention head, layer-norm, MLP/ReLU layers, and RoPE positional encoding. We believe that this work is a stepping stone towards the difficult task of obtaining useful guarantees for trained transformer models.

Figures (4)

• Figure 1: Two snippets of code that are identical except for the order of the assignment statements.
• Figure 2: The worst-case deviation for the model for the three different input restrictions and two different permutation classes. The $x$-axis is the number of tokens and the $y$-axis is the upper-bound on worst-case deviation.
• Figure 3: The peak-to-peak difference for the model for the three different input restrictions and two different permutation classes. Red "x"s indicate that the worst-case deviation is less than half of the peak-to-peak difference and, thus, the model output is provably invariant over the permutation class. The $x$-axis is the number of tokens and the $y$-axis is the peak-to-peak deviation.
• Figure 4: The peak-to-peak difference and the worst-case deviation when we continue generation through multiple tokens. The $x$-axis connotes the total number of tokens ($n_{ctx}$) throughout a generation. When the line for $PTP / 2$ is above the worst-case deviation line, then our algorithm provably guarantees that the output is fixed.

Theorems & Definitions (49)

• Definition 1.1: Overwhelm, Overwhelmed, Overwhelming
• Theorem 1.2: Formally restated in \ref{['thm:InpRes']}
• Theorem 1.3: Formally restated in \ref{['thm:InpResPermInv']}
• Definition 2.1: $1$-Layer Transformer
• Definition 2.2: One Hot Space
• Definition 3.1: Lipschitz Constant
• Definition 3.2: Worst-Case Deviation
• Definition 3.3: ${F_{\infty}}$-norm
• Definition 4.1: Input Restriction, o2014analysis definition 3.18
• Definition 4.2: Designed Space