---
title: Talking-Heads Attention in Transformers
url: https://www.emergentmind.com/topics/talking-heads-attention
type: topic
---

# Talking-Heads Attention in Transformers

Talking-Heads Attention is a generalization of standard multi-head attention in Transformer architectures wherein linear projections are inserted across the attention-heads dimension, immediately before and after the softmax operation. This modification creates information pathways among attention heads, enabling richer cross-head interactions. The construction, computational implications, empirical results, and theoretical motivations are distinct from standard multi-head formulations, yet involve a relatively minor increase in parametrization and compute burden. Talking-Heads Attention has demonstrated consistent improvements on large-scale masked language modeling and downstream transfer-learning tasks [2003.02436].

## 1. Standard Multi-Head Attention: Preliminaries

In multi-head self-attention, the input $X \in \mathbb{R}^{n \times d}$ (with sequence length $n$ and model dimension $d$) is projected via matrices $W^Q, W^K, W^V$ into $Q, K, V$ representations, which are reshaped into $h$ head-parallel blocks: $Q[i], K[i], V[i] \in \mathbb{R}^{n \times k}$ for each $i=1\dots h$. Each attention head computes the scaled softmax-dot-product as:
$$
A[i] = \text{softmax}(Q[i]K[i]^T/\sqrt{k}) \, V[i] \in \mathbb{R}^{n \times v}
$$
The outputs across heads are concatenated and linearly projected by $W^O$, yielding the final multi-head attention output. In this structure, heads' scoring and weighting operate in isolation, without interaction until the output merger.

## 2. Architectural Modifications: Head-Dimension Projections

Talking-Heads Attention introduces two learned matrices, $W_\text{pre} \in \mathbb{R}^{h_k \times h}$ and $W_\text{post} \in \mathbb{R}^{h \times h_v}$, which project across logit and value head dimensions, respectively. The typical configuration sets $h_k = h = h_v$, making both $W_\text{pre}$ and $W_\text{post}$ square matrices of dimension $h$.

The computation proceeds as follows:
1. Standard projections generate $Q \in \mathbb{R}^{n \times d_k \times h_k}$ and $K \in \mathbb{R}^{n \times d_k \times h_k}$.
2. Raw attention logits $L_0 \in \mathbb{R}^{n \times n \times h_k}$ are computed via pairwise dot products: $(L_0)_{p,q,r} = \sum_d Q_{p,d,r} K_{q,d,r}$.
3. Logit heads are mixed: $L = L_0 \times_3 W_\text{pre} \in \mathbb{R}^{n \times n \times h}$, where for each $(p,q)$, $L_{p,q,\cdot} = W_\text{pre}^T L_{p,q,\cdot}$.
4. Softmax is applied sequence-wise: $W = \text{softmax}(L) \in \mathbb{R}^{n \times n \times h}$.
5. Weight heads are further mixed: $U = W \times_3 W_\text{post} \in \mathbb{R}^{n \times n \times h_v}$.
6. $U$ is contracted with $V \in \mathbb{R}^{n \times d_v \times h_v}$ to yield per-head outputs, which are then projected as in the standard approach.

In tensor notation,
$$
L = (Q \,\overset{\text{head}}{\otimes}\, K) \cdot W_\text{pre}, \quad W = \text{softmax}(L), \quad U = W \cdot W_\text{post}, \quad \text{output} = (U \, \overset{\text{head}}{\otimes} \, V) W^O
$$

## 3. Computational and Parameter Complexity

Relative to standard multi-head attention, Talking-Heads Attention adds $O(h^2)$ parameters (two matrices of size $h \times h$ under $h_k = h_v = h$) and $O(n m h^2)$ arithmetic operations per layer. By contrast, conventional multi-head attention requires $O(h n m d_k)$ multiplies per layer and $O(h d d_k)$ parameters when $d_k = d_v$ is assumed for simplicity. The increase in operations and parameters is moderate when $h \ll d_k$ and is justified by empirical performance gains. The extra overhead arises specifically from the head-mixing tensor contractions and is proportional to $n$, $m$ (sequence dimensions), and $h^2$.

## 4. Empirical Evaluations and Ablation Analyses

Extensive experiments were conducted using T5 (12-layer encoder-decoder, $d_\text{model} = 768$), ALBERT (12-layer parameter-shared encoder), and BERT (12-layer, with relative positions, up to 768 heads).

Quantitative results on T5 (Table 1):
- At 12 heads, $d_k = 64$, standard multi-head: ln PPL = 1.678; Talking-Heads: ln PPL = 1.641.
- At 24 heads, $d_k = 32$, multi-head: 1.669; Talking-Heads: 1.624.
- SQuAD v1.1 F1 improved from 90.87 (multi-head) to 91.38 (Talking-Heads) with 12 heads, and to 91.83 with 24 heads.
- MNLI-m accuracy increased from 86.20 (multi-head) to 87.42 (Talking-Heads, 24 heads).

ALBERT (Table 7) revealed multi-head stagnation as head count increases, but Talking-Heads maintained or improved accuracy (e.g., avg accuracy rises from 80.78 to 81.44 as heads increase from 12 to 48).

BERT ablations (Table 9) using up to 768 heads ($d_k=1$) showed continuous gains: SQuAD1.1 F1 improved from 88.51 (12 heads) to 90.5 (768 heads), MNLI-m from 82.6 to 84.2. Ablation studies underscored the necessity of both pre- and post-softmax projections, and that most downstream gains were recoverable by applying Talking-Heads only to encoder self-attention.

### Summary Table: Perplexity and F1/Accuracy Gains

| Model  | Heads ($h$) | Multi-Head (PPL / F1 / Acc) | Talking-Heads (PPL / F1 / Acc) |
|--------|-------------|-----------------------------|---------------------------------|
| T5     | 12          | ln PPL = 1.678 / F1 = 90.87 | ln PPL = 1.641 / F1 = 91.38     |
| T5     | 24          | ln PPL = 1.669 / F1 = 91.83 | ln PPL = 1.624 / F1 = 91.83     |
| BERT   | 12          | F1 = 88.51 / Acc = 82.6     | F1 = 90.5 / Acc = 84.2          |

## 5. Mechanistic Interpretation and Information Flow

In standard multi-head attention, each head’s Q–K scoring and subsequent value weighting are isolated, which creates an information bottleneck, especially pronounced when $d_k$ is small relative to $h$. The insertion of $W_\text{pre}$ before softmax permits each head’s logits to incorporate information from every other head, forming cross-head compatibility patterns. $W_\text{post}$, applied after softmax, enables recombination of weight distributions prior to aggregation with $V$. Visualization of $W_\text{pre}$ and $W_\text{post}$ (see Figure 3 in [2003.02436]) demonstrates their dense, well-conditioned structure, with little evidence of dominance by any diagonal component. This suggests that head cross-talk is both active and nontrivial, mitigating bottlenecks observed in large-head, small-dimension settings.

## 6. Implementation Notes and Open Research Questions

The additional matrix multiplications required by head-mixing operations can be inefficient on hardware optimized for large GEMMs, raising the question of whether hardware or algorithmic innovations—such as locality-aware or memory-compressed attention—may further accelerate the method. A “dynamic” variant, in which $W_\text{pre}$ and $W_\text{post}$ receive small input-dependent offsets, was briefly explored; while it further improved pretraining perplexity, downstream accuracy benefits were not conclusively observed, suggesting the potential for future research into more sophisticated dynamic projections or data-dependent head mixing.

## 7. Broader Impact and Extensions

Talking-Heads Attention establishes a lightweight, extensible alternative to standard multi-head attention by introducing only $O(h^2)$ additional parameters and moderate computational overhead. The consistent improvements in pretraining log-perplexity and transfer learning performance across major Transformer variants highlight its practical relevance. While the gains are robust across a range of architectural hyperparameters, further work remains on optimizing small-matrix multiplication efficiency and exploring data-dependent projection schemes. The approach demonstrates that enhancing inter-head communication can compensate for limitations imposed by head dimension bottlenecks and offers a template for further architectural innovation within attention-based models [2003.02436].

Source: https://www.emergentmind.com/topics/talking-heads-attention