Skip to content

lacoco-lab/uhat-lean

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

23 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

Lean Formalization for Hard-Attention Transformers

This repository is a Lean formalization of key results from:

This repository was produced by GPT-5.5 with input from Michael Hahn.

The formalization focuses on unique hard-attention transformers (UHATs): fixed depth, fixed finite head count, deterministic hard argmax attention with a fixed tie rule, and unbounded input length. It also formalizes the Amiri et al. input-plus-chain-of-thought setting, where the input block is bidirectional and the CoT block is causal.

Source Versions

Line references in Lean comments and in this README refer to the local arXiv source files checked out next to this repository:

Paper External link Local source used for line numbers
Hahn 2020 https://arxiv.org/abs/1906.06755 ../multiplication-refreshed/background-hahn/arXiv-1906.06755v2/main.tex
Amiri et al. 2025 https://arxiv.org/abs/2502.02393 ../multiplication-refreshed/background-amiri-et-al/arxiv-clean.tex

The Hahn line numbers are specifically for the arXiv v2 source directory arXiv-1906.06755v2. The Amiri line numbers are for the local arXiv-clean source package whose 00README.json lists arxiv-clean.tex as the top-level file. The line numbers are intentionally source-file line numbers, not PDF page numbers.

Status

Result Status Main files
UHAT can compute Boolean AND with one layer Proved constructively HahnAmiri/And.lean
Fixed-depth UHAT cannot compute PARITY at all sufficiently large lengths Proved unconditionally in Lean HahnAmiri/HahnNoParity.lean, HahnAmiri/HahnStage3Distribution.lean, HahnAmiri/FiniteLLL.lean
Finite asymmetric Lovasz Local Lemma used by Hahn Stage 3 Proved internally, not assumed as an axiom HahnAmiri/FiniteLLL.lean
Fixed-depth UHAT cannot compute Hahn's Dyck language at all sufficiently large lengths Proved for 1Dyck, the one-bracket 2Dyck slice, and full two-bracket 2Dyck HahnAmiri/Dyck.lean
Amiri et al. PARITY CoT lower bound Proved: any fixed-depth UHAT+CoT realization of PARITY needs Omega(N) CoT length HahnAmiri/AmiriCoT.lean
PARITY CoT upper bound Proved: PARITY has a linear-length CoT in the same UHAT+CoT model HahnAmiri/AmiriCoT.lean

The project builds without sorry, admit, or project-specific axiom declarations. The usual Lean foundations such as Classical.choice remain available through Mathlib.

Paper-To-Lean Map

Paper passage Formalization
Hahn's hard-attention theorem statement, where the restricted computation depends on boundedly many inputs (main.tex:326-:330) Core restriction/dependency machinery in HahnAmiri/HardAttention.lean, HahnAmiri/ProbabilisticRestriction.lean, and the closed endpoints in HahnAmiri/HahnNoParity.lean
Hahn's PARITY and 2Dyck corollary (main.tex:331-:345) PARITY endpoint in HahnAmiri/Parity.lean; Dyck endpoint and full 2Dyck transfer in HahnAmiri/Dyck.lean
Iterated depth reduction from the final activation (main.tex:420-:425) Semantic collapse in HahnAmiri/HardAttention.lean and wrappers in HahnAmiri/HahnNoParity.lean
Stage 1 high-fanout restriction (main.tex:441-:447) Fanout/counting lemmas in HahnAmiri/HahnStages.lean
Stage 2 private-witness saturation and neighbor bound (main.tex:500-:514) Saturation construction and support-neighbor counting in HahnAmiri/HahnStages.lean and HahnAmiri/HahnStage3Distribution.lean
Stage 3 random restriction, Chernoff event, attention failure events, and LLL (main.tex:516-:553) Finite probability, Chernoff-style finite bounds, selected-support independence, and finite asymmetric LLL in HahnAmiri/FiniteProbability.lean, HahnAmiri/HahnStage3Distribution.lean, and HahnAmiri/FiniteLLL.lean
Stage 3 free-count arithmetic and layer removal (main.tex:573-:585) Closed finite-tail scalar arithmetic and depth-reduction wrappers in HahnAmiri/HahnNoParity.lean
Hahn's observation that AND-like witness languages remain easy (main.tex:347-:355) One-layer Boolean AND construction in HahnAmiri/And.lean
Amiri et al.'s CoT definition and arbitrary CoT vocabulary (arxiv-clean.tex:425-:432) CoTSpec, arbitrary Xi, and finite input alphabet interfaces in HahnAmiri/AmiriCoT.lean
Amiri et al.'s UHAT expressibility definition with uniformly bounded layers and heads (arxiv-clean.tex:453-:466) CoTExpressible and CoTRealizesAtLength in HahnAmiri/AmiriCoT.lean
Amiri et al.'s CoT fixing induction (arxiv-clean.tex:1535-:1588) Input-only restriction lift, per-token fixing step, induction state, and room arithmetic in HahnAmiri/AmiriCoT.lean
Amiri et al.'s PARITY CoT lower and upper bounds (arxiv-clean.tex:682-:696) amiri_parity_requires_OmegaN_cot, amiri_no_nonOmegaN_cot_for_parity, and parity_has_ThetaN_cot

Source Alignment

The README gives short excerpts only; the line ranges above identify the complete local source passages.

For Hahn, the public endpoints track the theorem and corollary around main.tex:326-:345, including the stated obstruction that hard attention "cannot model PARITY or 2Dyck" (main.tex:333). The Lean PARITY proof follows the bounded-dependency flip argument in main.tex:336-:338. The Lean Dyck proof follows the one-bracket reduction in main.tex:341-:344, including the "simpler variant 1Dyck" step, and then transfers it to full two-bracket 2Dyck.

The depth-reduction proof follows Hahn's three-stage organization. Stage 1 is the high-fanout pruning at main.tex:441-:447. Stage 2 is the iterative private-witness saturation loop at main.tex:500-:514, including the 2^c n / k iteration bound and the final k-neighbor estimate. Stage 3 is the random restriction and bad-event proof at main.tex:516-:553; the finite LLL is formalized where Hahn invokes the "Lovasz Local Lemma" (main.tex:544). The final collapse follows main.tex:582-:585, where the first layer is absorbed after the selected heads have bounded input dependency.

For Amiri et al., the CoT model keeps the finite input alphabet separate from the CoT alphabet, matching arxiv-clean.tex:425-:432; in particular, the CoT alphabet is "not restricted to be finite" (arxiv-clean.tex:431). The Lean definition uses one full-sequence UHAT at each length, with layer and head counts uniformly bounded as in arxiv-clean.tex:453-:466; it does not create a new transformer per CoT step. The CoT lower-bound proof follows the induction at arxiv-clean.tex:1535-:1588: first apply the Hahn-style restriction on the input block, then fix each generated token by spending at most a constant number of additional input positions per layer and head. The PARITY result formalizes the lower bound at arxiv-clean.tex:682-:687 and the matching prefix-parity upper bound at arxiv-clean.tex:696.

Main Theorems

AND

HahnAmiri.HardAttention.And.exists_one_layer_hard_attention_computes_and

This is a sanity check: hard attention can compute simple witness-style predicates by selecting an appropriate position.

Hahn PARITY Lower Bound

HahnAmiri.HardAttention.Parity.hahn_theorem_1_parity_lower_bound_all_lengths
HahnAmiri.HardAttention.Parity.HahnSequenceHardAttentionFamily.fails_for_all_sufficiently_large_lengths
HahnAmiri.HardAttention.Parity.HahnSequenceHardAttentionFamily.no_family_computes_parity_eventually

For fixed head count and fixed depth, no hard-attention transformer family indexed by input length computes PARITY at every sufficiently large length. The statement has no Stage-1, Stage-2, Stage-3, LLL, or scalar-asymptotic certificates in its assumptions.

Dyck Lower Bound

HahnAmiri.HardAttention.Restrictions.Dyck.hahn_theorem_dyck_lower_bound
HahnAmiri.HardAttention.Restrictions.Dyck.hahn_theorem_twoDyck_slice_lower_bound
HahnAmiri.HardAttention.Restrictions.Dyck.hahn_theorem_full_twoDyck_lower_bound

The Dyck file proves Hahn's final bounded-dependency obstruction. After the restriction and depth-reduction stages, the output depends on only boundedly many positions. In Hahn's prefix/suffix context, the formal proof constructs a compatible 1Dyck completion with a free close outside that bounded support, flips that close to an open, and obtains two compatible inputs that agree on the dependency support but differ on Dyck membership.

The closed theorem hahn_theorem_dyck_lower_bound constructs the required high-free scalar tower internally. On lengths n = 10m, it keeps the middle room needed by Hahn's prefix/suffix repair. The theorem hahn_theorem_twoDyck_slice_lower_bound transfers this to the one-bracket 2Dyck slice, and hahn_theorem_full_twoDyck_lower_bound proves the full two-bracket result by restriction to that slice.

Amiri et al. PARITY CoT Lower Bound

HahnAmiri.HardAttention.Amiri.amiri_parity_requires_OmegaN_cot
HahnAmiri.HardAttention.Amiri.amiri_no_nonOmegaN_cot_for_parity

The formalization uses one full-sequence UHAT for each input length. The input block is unmasked bidirectionally, while CoT positions can attend to the full input and only the causal CoT prefix. The lower bound is proved for Boolean inputs and arbitrary CoT alphabet.

The reusable obstruction interfaces are:

HahnAmiri.HardAttention.Amiri.PositiveFreeTargetObstruction
HahnAmiri.HardAttention.Amiri.LinearFreeTargetObstruction

These interfaces are intended as hooks for later MEDIAN, REACHABILITY, and related applications. Those task-specific instantiations are not completed in this repository.

PARITY CoT Upper Bound

HahnAmiri.HardAttention.Amiri.parityLinearCoTSpec
HahnAmiri.HardAttention.Amiri.parity_has_linear_cot_expressible
HahnAmiri.HardAttention.Amiri.parity_has_ThetaN_cot
HahnAmiri.HardAttention.Amiri.amiri_parity_cot_complexity_ThetaN

The construction uses CoT alphabet ZMod 2 and CoT length N + 1. The CoT tokens store the running prefix parity, so the final token is the parity of the input. This gives the matching linear upper bound in the same UHAT+CoT model.

File Guide

File Purpose
HahnAmiri/HardAttention.lean Core hard-attention layers, transformers, dependency supports, and depth collapse
HahnAmiri/ProbabilisticRestriction.lean Restrictions, compatible inputs, free counts, and fixing operations
HahnAmiri/FiniteProbability.lean Finite probability space utilities
HahnAmiri/FiniteLLL.lean Finite asymmetric Lovasz Local Lemma
HahnAmiri/HahnParameters.lean Hahn parameter constants and budget bookkeeping
HahnAmiri/HahnStages.lean Hahn Stage 1 and Stage 2 structures
HahnAmiri/HahnStage3Distribution.lean Hahn Stage 3 distribution, bad events, and selected-support independence
HahnAmiri/HahnNoParity.lean Closed fixed-depth Hahn PARITY lower bound
HahnAmiri/Parity.lean Bounded-dependency PARITY endpoint
HahnAmiri/Dyck.lean Dyck definitions and bounded-dependency obstruction
HahnAmiri/And.lean One-layer AND construction
HahnAmiri/AmiriCoT.lean Amiri et al. UHAT+CoT model, PARITY lower bound, and PARITY upper bound

What Is Not Formalized

  • Hahn's soft-attention lower bounds are not part of this development.
  • Amiri et al.'s MEDIAN, REACHABILITY, and MULTIPLICATION applications are not instantiated here, although the PARITY proof exposes reusable obstruction interfaces for such targets.

Build And Audit

This project uses Lake and Lean 4. The Lean toolchain is pinned in lean-toolchain.

lake build
rg -n "\b(sorry|admit|axiom)\b" HahnAmiri -g '*.lean'

The expected audit result is that the rg command reports no matches.

Citation

@article{hahn2020theoretical,
  title   = {Theoretical Limitations of Self-Attention in Neural Sequence Models},
  author  = {Hahn, Michael},
  journal = {Transactions of the Association for Computational Linguistics},
  volume  = {8},
  pages   = {156--171},
  year    = {2020}
}

@inproceedings{amiri2025lower,
  title     = {Lower Bounds for Chain-of-Thought Reasoning in Hard-Attention Transformers},
  author    = {Amiri, Alireza and Huang, Xinting and Rofin, Mark and Hahn, Michael},
  booktitle = {Proceedings of the 42nd International Conference on Machine Learning},
  year      = {2025}
}

About

No description, website, or topics provided.

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

No releases published

Packages

 
 
 

Contributors

Languages