Averaging-Hard-Attention (AHAT) in Transformers

• Averaging-Hard-Attention (AHAT) is a transformer mechanism that hard-selects maximal inputs and uniformly averages them, encoding cardinality into its computations.
• AHAT bridges standard transformer architectures with uniform TC⁰ circuits, demonstrating its ability to simulate threshold-based counting operations.
• Explicit constructions in AHAT address recognition of languages like MAJORITY, DYCK₁, and PARITY, highlighting both practical applications and theoretical implications.

Averaging-Hard-Attention (AHAT) refers to a restricted form of the transformer attention mechanism in which, at each position, the model "hard-selects" all inputs that maximize a score function and computes their uniform average. Unlike the unique hard-attention model—which selects only a single (possibly arbitrarily chosen) maximizer—AHAT pools over all ties, encoding cardinalities into the computation and fundamentally increasing its counting expressivity. Recent results establish that fixed-width, constant-depth AHAT transformers with logarithmic precision are provably equivalent, in recognition power, to the classical circuit class uniform TC⁰: the languages recognized by polynomial-size, constant-depth, logspace-uniform threshold circuits with majority gates (Strobl, 2023, Barcelo et al., 2023, Hao et al., 2022).

1. Formal Description of Averaging-Hard-Attention

An AHAT transformer layer follows the standard encoder-layer template with the key alteration in its attention phase. For an input sequence $X = (x_1, \dots, x_n)$, where $x_i \in F_p^k$ for some finite-precision float space $F_p^k$ (typically $p = \Theta(\log n)$), an attention head computes a score matrix $S$ such that the $(i,j)$ entry is $s_{ij} = o(x_i, x_j)$, where $o : F_p^k \times F_p^k \to F_p$ is linear-space computable. For the $i$-th row, define $M(s^{(i)})$ as the set of $j$ achieving the maximum $s_{ij}$ over $j \in [n]$. The AHAT distribution $\alpha(s^{(i)})$ puts uniform weight $1/|M(s^{(i)})|$ on these maximizers and $0$ elsewhere. The attention output at position $i$ is then:

$$\text{Att}_o(X)_i = \left[ \frac{1}{|M(s^{(i)})|} \sum_{j \in M(s^{(i)})} x_j \right]_p$$

where $[\,\cdot\,]_p$ denotes truncation back into $F_p$. The overall layer output is the result of applying a feed-forward network $f$ to the concatenation of $x_i$ and the head outputs, followed by pointwise activation (typically ReLU), residual connection, and so forth (Strobl, 2023, Barcelo et al., 2023, Hao et al., 2022).

2. Circuit Complexity Characterization and Main Theorems

AHAT’s theoretical foundation is its rigorous connection to threshold circuits. The central result establishes that, under the constraint of fixed width, fixed depth, and logarithmic arithmetic precision, every AHAT transformer can be simulated by a uniform family of constant-depth, polynomial-size threshold circuits with majority gates (i.e., uniform TC⁰) (Strobl, 2023). More precisely, the following hold:

• Theorem 1 (Simulation): Each AHAT attention head, for each sequence index $i$, can be computed by a log-space uniform, depth-$O(1)$ threshold circuit of poly$(n)$ size. The circuit computes scores, finds maxima, tests equality, computes uniform averages, division, and truncation, all in parallel and in constant depth.
• Theorem 2 (Expressive Power): For any fixed constant number of layers, heads, and precision $p = \Theta(\log n)$, the language recognized (i.e., accepted or rejected) by the AHAT transformer is in the class of languages recognized by uniform TC⁰ circuits.

This establishes AHAT as exactly recognizing uniform TC⁰, matching the full expressive power of first-order logic with counting quantifiers and Presburger arithmetic over finite words (Barcelo et al., 2023).

3. AHAT Versus Unique Hard Attention and Other Transformer Models

A strict separation exists between models using unique hard attention (UHAT), generalized unique hard attention (GUHAT), and average-hard attention. UHAT and GUHAT restrict pooling to one maximizer (deterministic tie-breaking), thus cannot count and can only recognize languages in AC⁰—the class associated with AND, OR, NOT gates in constant-depth, uniform, poly-size circuits. AHAT, by averaging over all maximizers, encodes the cardinalities of maximal sets and thereby supports exact thresholding and counting operations, which are provably outside AC⁰ but inside TC⁰ (Hao et al., 2022, Barcelo et al., 2023). Thus, the following hierarchy holds:

$$\text{UHAT} \subsetneq \text{AHAT} \subseteq \text{TC}^0$$

Softmax-based attention transformers, under logarithmic precision, also fall within uniform TC⁰, but the relative power of AHAT and softmax attention in this regime is an open question (Strobl, 2023).

4. Explicit Constructions and Recognized Languages

AHAT’s recognition power includes, but is not limited to, the following concrete language classes:

• MAJORITY: The language of bitstrings with at least as many $1$s as $0$s. A single-layer, single-head AHAT can implement this by embedding $+1$ for $1$, $-1$ for $0$, assigning equal scores (i.e., all ties), and pooling to compute the normalized sum; thresholding determines acceptance (Hao et al., 2022).
• DYCK₁: Well-formed, balanced parenthesis sequences. Multiple layers and activation dimensions suffice: the network computes prefix sums and prefix minima via mask-restricted averaging and targeted scoring, and finally verifies net zero and non-negativity in the last layer (Hao et al., 2022).
• Parity: Languages recognizing even/odd number of occurrences of a symbol (e.g., $\mathrm{PARITY} = \{x : \#1(x) \bmod 2 = 0\}$) are in AHAT via linear thresholding over averaged maximal indices (Barcelo et al., 2023).
• Permutation-closed languages: Any permutation (Parikh image) of a regular language, corresponding to semilinear sets, is recognizable by AHAT since linear inequalities over symbol counts, i.e., Presburger-definable sets, are implementable in constant depth (Barcelo et al., 2023).

By contrast, languages in AC⁰, such as those requiring only bounded counting and no majority, are admitted by UHAT—however, UHAT cannot recognize MAJORITY or DYCK₁ (Hao et al., 2022, Barcelo et al., 2023).

5. Simulation by Uniform Threshold Circuits

The AHAT computation decomposes naturally into TC⁰-computable subroutines, as follows (Strobl, 2023):

Level	Description	TC⁰ Subroutine
Score Computation	Bilinear or linear-space computable scoring	Add/mult, $O(1)$-depth
Maximum Level	$\max_{j}s_{ij}$ computation	poly-size, depth-3 threshold
Equality Mask	Boolean $e_{ij} = (s_{ij} = m_i)$	Equality on $O(\log n)$ bits
Summation/Numerator	Select-and-sum $x_j$ where $e_{ij}=1$	Sum in TC⁰
Denominator	Count $|M(s^{(i)})| = \sum_j e_{ij}$	Majority + OR gating
Division and Trunc.	Compute $N_i/D_i$ in $F_p$, truncate if needed	Division/subsampling in TC⁰

Each phase is constant-depth, logspace-uniform, and polynomial size in $n$. Parallel evaluation across sequence positions is straightforward.

6. Logical and Structural Expressivity

AHAT implements first-order logic with counting quantifiers (FO[C]) and Presburger arithmetic over finite words. Via constant-depth, constant-width stacks of average-hard attention plus ReLU layers, temporal logic (LTL) with extension by counting terms can be exactly simulated. This yields, in logic-theoretic terms, the inclusion:

$$\mathsf{AHAT} = \mathsf{FO[C]} = \text{uniform } \mathrm{TC}^{0}$$

This equivalence is witnessed by explicit model constructions: Boolean combinations realized by ReLU + affine maps, temporal operators by constrained attention, and counting/Presburger tests by appropriate attention labeling and post-layer ReLU (Barcelo et al., 2023, Strobl, 2023).

7. Limitations, Open Questions, and Directions

Although AHAT achieves full uniform TC⁰ expressivity, important open problems persist. Notably, whether every uniform TC⁰ language can be realized by a fixed-width, fixed-depth AHAT network remains unresolved. There exist natural TC⁰-complete problems, but formal lower bounds or separation results for AHAT as a transformer architecture have not yet been proven (Strobl, 2023, Barcelo et al., 2023). The effect of positional embeddings, masking schemes, and possible Turing completeness with alternate embeddings or increasing depth/width is not fully understood. Furthermore, closure under operations such as concatenation and star is unclear for AHAT, mirroring open problems for uniform TC⁰ itself (Hao et al., 2022).

A plausible implication is that further relaxation of constraints—such as super-logarithmic precision, more flexible positional representations, or increased stacking—could potentially extend the class beyond TC⁰, but this is not established. In practical machine learning contexts, perfectly hard (max or average-hard) attention is rarely deployed due to gradients; however, these results inform the theoretical understanding of what can be efficiently captured by discrete transformer architectures.

---

References:

(Strobl, 2023, Barcelo et al., 2023, Hao et al., 2022)