Vertex-Softmax: Tight Transformer Verification via Exact Softmax Optimization

TL;DR

Vertex-Softmax introduces an exact softmax optimization method for transformer verification, improving certification tightness and efficiency by leveraging vertex-based solutions and a novel threshold structure theorem.

Contribution

It provides the first tight, exact softmax bound for transformer verification, with a log-linear complexity and formal optimality guarantees, enhancing existing verification methods.

Findings

Significantly improves certified rates on MNIST, Fashion-MNIST, CIFAR-10.

Tightly bounds softmax over score intervals with formal optimality.

Outperforms alpha-CROWN and branch-and-bound methods at lower computational cost.

Abstract

Certified verification of transformer attention requires bounding the softmax function over interval constraints on the pre-softmax scores. Existing verifiers relax softmax ndependently of the downstream objective, leaving avoidable slack. We prove that the exact optimum of this score-box problem is attained at a vertex of the constraint box, and establish a threshold structure theorem showing that, after sorting the objective coefficients, the optimum lies among only linearly many candidates, yielding the Vertex-Softmax primitive with log-linear complexity in the sequence length. We further prove a formal optimality result showing that Vertex-Softmax is the tightest sound bound obtainable from score intervals alone, characterizing precisely what additional structure (score correlations, score-value coupling) is needed for further improvement. Integrated into a CROWN Convex Relaxation based Optimization for Worst-case Neurons)-style verifier with a formal soundness guarantee, Vertex-Softmax significantly improves certified rates and substantially tightens lower bounds across MNIST, Fashion-MNIST, and CIFAR-10 attention models, while consistently matching or outperforming alpha-CROWN and branch-and-bound baselines at a fraction of their cost.