Attention-based representations for multi-task computation

2026-08-04

AI Summary

This work investigates the minimal number of attention heads required in multi-head attention mechanisms for multi-task settings and characterizes their representational limitations. Focusing on extremum extraction and n-bit XOR tasks, the study establishes the first theoretical lower bounds relating the number of attention heads, embedding dimension, and computational precision, and presents tight constructions that match these bounds. By integrating linear predictors, polynomial threshold functions, and threshold-degree analysis of Boolean functions, the authors demonstrate that single-head attention necessitates exponentially more resources for extremum tasks. For XOR tasks, they achieve an optimal configuration where the product of the number of heads and polynomial degree equals n, thereby revealing the necessity and efficiency advantages of multi-head architectures in representing multiple tasks.

Abstract

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.

Problem

• multi-task computation
• attention heads
• vector representations
• Boolean functions
• threshold degree

Innovation

• multi-head attention
• multi-task representation
• theoretical bounds
• Boolean functions
• threshold degree