PaperlayerView exact revision on arXiv

Source-linked AI summary

Attention-based representations for multi-task computation

Daniel Hsu, Mingyue Xu

arXiv:2608.04243v1cs.LG

TL;DR

The paper asks when multiple attention heads are necessary for representations supporting multiple tasks. It proves bounds for simultaneous minimum/maximum prediction and for XOR, and constructs attention layers matching the XOR tradeoff. For symmetric Boolean functions, the tradeoff is expressed using threshold degree.

  • Problem

    It is unclear when multiple heads in one attention layer are necessary for representations that support multiple tasks, since a single head can approximate essentially any function if its dimension or precision is high enough.

    Source: §1 · p. 1

  • Method

    The paper proves lower bounds for simultaneous minimum/maximum computation and XOR, then constructs multi-head layers using rational-function representations and additive-basis arguments.

    Source: §1 · p. 2, §1.2 · p. 4

  • Results

    For n-bit XOR, the product of head count H and polynomial degree D must be at least n, and constructions attain this bound; for symmetric Boolean functions, n is replaced by threshold degree T.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Increasing the number of heads amplifies the threshold degree of symmetric Boolean functions computable by polynomial-threshold post-processing, while a large value-vector dimension cannot compensate for too few heads.

    Source: §6 · p. 20

    The minimum/maximum lower bound covers linear classifiers; its extension to nonlinear predictors is unclear.

    Source: §6 · p. 20

Abstract

from arXiv · show

Multi-head attention layers produce vector representations that support multiple downstream tasks. We establish bounds on the number of heads required in two simple and concrete multi-task scenarios. In the first scenario, a vector representation is sought so that linear predictors can compute both the smallest and largest numbers in a given list. In this case, it is known two attention heads with small embedding dimension and bit precision level suffice. We prove that a single attention head requires exponentially higher embedding dimension or precision level. In the second scenario, a vector representation is sought so that a polynomial threshold function can compute the XOR of a given string of $n$ bits. This scenario is analogous to the first one for $n=2$, since XOR is readily computed by a linear function using a vector representation that encodes both the AND and the OR of the two bits. We observe that $n$-bit XOR requires the product of the number of heads and the polynomial degree to be at least $n$, and we construct multi-head attention layers that match this lower bound. These results generalize to arbitrary (symmetric) Boolean functions, where the bound is given in terms of the threshold degree.

1 Introduction

The paper studies when multiple attention heads are needed for representations that support multiple tasks, establishing lower bounds for simultaneous minimum/maximum computation and XOR and matching XOR constructions. It situates these results among prior head-separation work and uses geometric and polynomial-based arguments to prove them.

1 Introduction

  • A single attention head supports linear prediction of both list minimum and maximum only with value-embedding dimension or precision that grows large with the input or domain size.Each task is separately supported by one head with a linear classifier; the lower bound concerns supporting both simultaneously.Source: §1 · p. 2
  • For n-bit XOR computed by a degree-D polynomial threshold function, the necessary and sufficient tradeoff is D × H ≥ n; symmetric Boolean functions use threshold degree T instead of n.The construction achieves value-embedding dimension O(min{D, log H}).Source: §1 · p. 2

1.1 Related works

  • Earlier work establishes head separations for other settings, including two-bit XOR, while the paper extends the XOR analysis to multiple heads and polynomial-threshold post-processing.The related-work discussion also distinguishes its minimum/maximum setting from Yu et al. [30] by task, exactness, head count, and classifier restrictions.Source: §1.1 · pp. 2–3

1.2 Proof techniques

  • For minimum/maximum, the lower bound uses the Erdős–Szekeres theorem to force many value-embedding constraints, then a volume argument to bound dimension or precision.Monotonic attention weights make one extreme receive relatively low weight, so its value embedding must stand out for linear prediction.Source: §1.2 · pp. 3–4
  • For XOR, rational-function representations of attention heads connect the lower bound to threshold degree, while rational-function constructions and additive bases yield matching upper bounds.The dimension-reduction construction chooses powers from a restricted additive D-basis and achieves dimension O(min{D, log H}).Source: §1.2 · p. 4

2 Preliminaries

The preliminaries specify the attention-layer model and define the prediction and polynomial concepts used in the results. They also introduce additive bases as a way to represent integers with bounded numbers of summands.

2 Preliminaries

  • The model uses standard softmax attention with one query vector; an H-head layer sums H head outputs sharing value-vector dimension d.The query/key-vector dimension is irrelevant to the analysis, and the constructions can use one-dimensional query and key vectors.Source: §2 · p. 5
  • A linear classifier selects the class with the largest weight-vector inner product; a polynomial sign-represents a Boolean function when its sign matches the function on every input.Threshold degree is the minimum degree of such a sign-representing polynomial, and symmetric Boolean functions depend only on input weight.Source: §2 · p. 5

Additive bases.

  • An additive D-basis with range T represents every non-negative integer up to T as a sum of at most D basis elements.The empty sum represents zero.Source: Additive bases · p. 5

3 Main results

For min and max, a single attention head with linear classifiers requires at least one resource—embedding dimension, precision, or value-vector norm—to grow roughly as the square root of the smaller of the input-domain and list sizes, while two heads suffice with exponentially smaller resource requirements when M = poly(n). For XOR and symmetric Boolean functions, the product of head count and polynomial degree must meet the function’s threshold degree, and matching constructions attain that tradeoff.

3.1 Attention-based minimum and maximum computation

  • At least one of d, p, and log2(ℓ) must grow roughly as the square root of min{M, n} for one head to support both min and max.When M ≳ n^2, the paper also gives a weaker, simpler lower-bound conclusion; two heads suffice with exponentially smaller (d, p) when M = poly(n).Source: §3.1 · p. 6

3.2 Attention-based XOR computation

  • For n-bit XOR, any degree-D polynomial threshold function composed with H attention heads must satisfy D × H ≥ n.The n-bit XOR function has threshold degree n.Source: §3.2 · p. 7
  • Matching constructions attain the D × H ≥ T lower bound with value-vector dimension at most O(min{D, log H}).The construction is stated for all positive D and H satisfying the threshold-degree condition.Source: §3.2 · p. 7
  • For a symmetric Boolean function with threshold degree T, H attention heads and a degree-D polynomial suffice whenever D × H ≥ T.The construction uses value-vector dimension at most min{D, 2 max{1, ⌈log2 H⌉} + 2}.Source: §3.2 · p. 7
  • For intermediate D, whether value-vector dimension can be improved below O(min{D, log H}) remains open.The proof gives constant dimension when D is constant or at least a constant positive power of H, such as D ≳ H^0.01.Source: §3.2 · p. 7

4 Proof of Theorem 1

The proof of Theorem 1 turns correct min and max classification into constraints on value vectors, then combines those constraints with a geometric bound on approximately triangular vector configurations. A monotone subsequence of attention weights supplies the relevant ordering, and bounded norm and integer precision limit the resulting volume growth.

4.1 Histogram representation

  • The input is represented by a histogram h whose entries count occurrences, with h_1 + ··· + h_M = n.The proof begins with this representation of an input in [M]^n.Source: §4.1 · p. 7

4.2 Monotonic subsequence of attention weights

  • When M ≥ (N − 1)^2 + 1, some increasing sequence x_1 < ··· < x_N has attention weights that are either non-decreasing or non-increasing.This follows from the Erdős–Szekeres theorem [8, 25].Source: §4.2 · p. 8

4.3 Constraints implied by a correct minimum classifier

  • For non-decreasing attention weights, correct minimum classification yields vectors λ_i with ⟨λ_i, v_xi⟩ = 1 and |⟨λ_i, v_xj⟩| < 1/(n − 1) for later j.These constraints make each selected value vector stand out from those that follow it.Source: §4.3 · pp. 8–9

4.4 Constraints implied by a correct maximum classifier

  • For non-increasing attention weights, correct maximum classification yields vectors μ_i with ⟨μ_i, v_xi⟩ = 1 and |⟨μ_i, v_xj⟩| < 1/(n − 1) for earlier j.Source: §4.4 · p. 9

4.5 Approximately triangular configurations

  • Lemma 6 bounds the length of an approximately triangular vector configuration using its span dimension, vector norm, and integer-coordinate conditions.The proof’s volume-growth argument limits rank-preserving steps under bounded precision and bounded norm.Source: §4.5 · p. 10
  • When a vector remains in the preceding span, its expression in the current basis must include a coefficient of large magnitude, producing volume growth.The argument relates this growth to bounded norm and integer precision.Source: §4.5 · p. 10
  • In the Lemma 6 process, vectors are considered in order and matrices track bases for their successive spans; steps either increase rank or replace a basis vector without increasing rank.The change-of-basis determinants are analyzed across these updates.Source: §4.5 · pp. 10–11

4.6 Finishing the proof

  • The proof finishes by applying the monotone-weight and classifier constraints to Lemma 6’s geometric bound to establish Theorem 1.The construction of the vectors for Lemma 6 depends on whether the selected attention weights are non-decreasing or non-increasing.Source: §4.6 · p. 13

5 Proofs of Theorems 2 and 3

The proofs establish the head-degree lower bound by translating attention layers into rational representations, then construct representations that attain the bound. For symmetric Boolean functions, the construction uses polynomial representations tailored to the degree and head-count regime.

5.2 Rational representations of attention layers

  • An H-head layer has a rational representation with numerator and denominator of degree at most H; composing it with a degree-D polynomial yields a degree-DH sign representation, so D × H ≥ T.The denominator is positive on Boolean inputs, preserving the sign in the conversion.Source: §5.2 · pp. 14–15

5.3 Realizing a polynomial threshold function with attention heads

  • H attention heads can realize d univariate polynomials of degree at most H, up to positive scaling that does not affect sign after homogeneous polynomial post-processing.This construction supplies the polynomial features used in the upper-bound argument.Source: §5.3 · pp. 16–17

5.4 Degree-restricted compositional sign-representations

  • When D × H ≥ T, factorization represents a degree-T univariate polynomial as a degree-at-most-D polynomial applied to at most D polynomials of degree at most H.The construction uses real roots and a positive factor from paired complex roots to obtain a sign representation.Source: §5.4 · pp. 17–18
  • For large D relative to a positive power of H, additive D-bases provide an alternative representation with smaller value-vector dimension.The proof gives a basis of cardinality at most 2p + 1 for range DH under the stated condition on D, H, and p.Source: §5.4 · pp. 18–19

5.5 Finishing the proof

  • Theorem 3 combines the two representation routes in separate D regimes to construct H attention heads and a post-processing polynomial for a symmetric Boolean function.The proof sets p := max{1, ⌈log2(H)⌉}; the additive-basis route gives dimension at most 2p + 2.Source: §5.5 · pp. 19–20

6 Discussion

The discussion limits Theorem 1 to linear post-processing and identifies unanswered questions about classifier scope and embedding dimension. The results also concern standard softmax attention, leaving positional embeddings, layer normalization, and multiple layers outside their scope.

6 Discussion

  • Theorem 1 establishes its lower bound only for linear classifiers; extending it to nonlinear classifiers remains unresolved.The authors also leave open whether nonlinear classifiers can avoid the lower bound without greatly increasing embedding dimension or precision.Source: §6 · p. 20
  • Increasing head count amplifies the threshold degree of symmetric Boolean functions computable with polynomial-threshold post-processing, and embedding dimension cannot compensate for too few heads.The constructions use value-vector dimension O(min{D, log H}); whether that dimension is necessary remains open.Source: §6 · p. 20
  • The results cover standard softmax attention without positional embeddings or layer normalization; the joint role of multiple heads and layers remains a future question.Source: §6 · p. 20

A Attention-based minimum computation

This appendix gives an attention-head construction for computing the minimum of a list of integers, using a linear classifier. The same construction computes the maximum by negating the query/key values.

A Attention-based minimum computation

  • A single attention head with value-vector dimension ⌈8 ln(M)⌉ and a linear classifier computes the minimum of n integers from [M].The construction uses sign vectors as value vectors and assigns a classifier weight vector to each integer.Source: A Attention-based minimum computation · p. 23
  • Negating all query/key values in the construction gives a corresponding attention-head construction for computing the maximum.Source: A Attention-based minimum computation · p. 23

B Deferred proofs

The deferred proofs establish auxiliary bounds for the minimum-computation argument and a polynomial basis used in the attention construction. They also show that the polynomials Q0 through QH form a basis for polynomials of degree at most H.

B Deferred proofs

  • For the minimum-computation proof, inputs with different histograms are used to derive constraints implied by correct minimum prediction.Source: B Deferred proofs · pp. 23–24
  • A matrix with columns of Euclidean norm at most L satisfies det(B^T B) ≤ L^(2r), while an integer matrix of full column rank satisfies det(A^T A) ≥ 1.The first bound follows from Hadamard’s inequality; the second uses positive definiteness and integrality.Source: B Deferred proofs · p. 25
  • The polynomials Q0, Q1, …, QH form a basis for the vector space of polynomials of degree at most H.Source: B Deferred proofs · p. 25

C Additive bases

The appendix constructs additive D-bases with range DH and cardinality at most 2p + 1 under a condition on D, H, and p. A power-of-two example illustrates the construction, while a counting argument limits what is possible for arbitrary parameters.

C Additive bases

  • For power-of-two H, the candidate basis combines H with powers of two and their complements to represent values using at most D elements.When using q copies of H and remainder terms would exceed D, the construction reduces the number of H terms and uses B2 to make up the difference.Source: C Additive bases · p. 26
  • When D + 2 ≥ 2p(b − 1), a basis formed from H and the sets B1 and B2 has range DH and cardinality at most 2p + 1.Here b = ⌈H^(1/p)⌉, B1 contains powers b^i below H, and B2 contains their complements H − β.Source: C Additive bases · p. 26
  • Remainders below H use at most p(b − 1) elements of B1; when that plus the quotient exceeds D, B2 terms replace one H while keeping the total at most D.The proof handles the cases q + p(b − 1) ≤ D and q + p(b − 1) ≥ D + 1 separately.Source: C Additive bases · p. 27
  • Counting possible sums from k basis elements against DH + 1 target integers limits the range attainable by an additive D-basis.The appendix uses this counting argument to show that constant-size bases cannot work for arbitrary D and H.Source: C Additive bases · p. 27

D Lower bounds against post-processing by Lipschitz predictors

The appendix bounds attention heads followed by Lipschitz predictors: supporting both min and max requires the Lipschitz constant or some pairwise value-vector distance to grow at least roughly as the square root of the input size. Its proof compares carefully chosen input sequences and their attention outputs.

D Lower bounds against post-processing by Lipschitz predictors

  • Supporting both min and max with Lipschitz post-processing requires the Lipschitz constant or some pairwise value-vector distance to grow at least roughly as the square root of the input size.The proposition assumes pairwise value-vector distances are bounded by D and both predictors are L-Lipschitz.Source: D Lower bounds against post-processing… · pp. 27–28
  • The proof invokes a lemma guaranteeing input pairs under bounded pairwise value-vector distances, then uses the Lipschitz property of gmin and an analogous argument for max.The proof notes that its approach is similar in spirit to Yu et al. [30].Source: D Lower bounds against post-processing… · pp. 28–29
  • For i<j with αi≤αj, where αi=exp(⟨q,ki⟩) and αj=exp(⟨q,kj⟩), the proof compares (i,j,…,j) with (j,j,…,j), giving difference (vi−vj)/[1+(n−1)αj/αi].It also compares (i,…,i,j) with (i,…,i,i), giving difference (vj−vi)/[(n−1)αi/αj+1].Source: D Lower bounds against post-processing… · pp. 28–29
Loading 2608.04243v1…