Multi-head attention

As explained in Natural Language Processing in Action, Second Edition

Multi-head attention computes multiple attention representations in parallel and combines them, potentially using distinct query and source sequences.

A word can matter in more than one way. Multi-head attention lets a transformer compute several attention representations in parallel, so different heads can attend to different positions and patterns. In self-attention, the queries, keys, and values all come from the input sequence’s embeddings. A decoder can also attend to an encoder’s output: its target sequence supplies the queries, while the encoder output supplies the keys and values. section 9.1.2 section 9.2.3

The computation builds on scaled dot-product attention:

• Project embeddings into queries, keys, and values. Query–key scores, scaled by 1/√dₖ and passed through softmax, weight the values. section 9.1.2
• Run this computation across multiple heads, concatenate their outputs, and apply an output projection. The number of heads is an architecture hyperparameter; the original transformer used eight heads with 64-dimensional keys and values. section 9.1.2
• Place the attention layer alongside a feedforward layer and normalization in a stackable transformer block. Encoder and decoder stacks use attention in different positions. section 9.2.2

Figure 9.4 Multi-head self-attention

The diagram shows the parallel heads and how their results become one projected output. section 9.1.2

Inspect attention during translation

For a translation experiment, load aligned German–English sentences, tokenize the two languages, and pass source and target sequences to a PyTorch transformer. A customized decoder layer calls self-attention with the target three times, then calls multi-head attention with the target and encoder memory. It returns the latter attention weights for inspection. In the supplied test example, a German sentence about a mother and her son produces an English translation; plotted weights put source and output tokens on the axes. The plots help inspect what the model attended to, but a translation still needs evaluation on held-out data. section 9.2.3

Listing 9.5 Inspecting attention weights with customized TransformerDecoderLayer

>>> from torch import Tensor
>>> from typing import Optional, Any

>>> class CustomDecoderLayer(nn.TransformerDecoderLayer):
...     def forward(self, tgt: Tensor, memory: Tensor,
...             tgt_mask: Optional[Tensor] = None,
...             memory_mask: Optional[Tensor] = None,
...             tgt_key_padding_mask: Optional[Tensor] = None
...             ) -> Tensor:
...         """Like decode but returns multi-head attention weights."""
...         tgt2 = self.self_attn(
...             tgt, tgt, tgt, attn_mask=tgt_mask,
...             key_padding_mask=tgt_key_padding_mask)[0]
...         tgt = tgt + self.dropout1(tgt2)
...         tgt = self.norm1(tgt)
...         tgt2, attention_weights = self.multihead_attn(
...             tgt, memory, memory,                           #1
...             attn_mask=memory_mask,
...             key_padding_mask=mem_key_padding_mask,
...             need_weights=True)
...         tgt = tgt + self.dropout2(tgt2)
...         tgt = self.norm2(tgt)
...         tgt2 = self.linear2(
...             self.dropout(self.activation(self.linear1(tgt))))
...         tgt = tgt + self.dropout3(tgt2)
...         tgt = self.norm3(tgt)
...         return tgt, attention_weights           #2

The decoder-layer code makes the distinction between target self-attention and attention over encoder memory concrete. section 9.2.3

How Natural Language Processing in Action, Second Edition covers it

How deep it goes

Works through scoring, head dimensions, concatenation, and projection, then implements weight inspection through PyTorch decoder customization rather than building attention from scratch.

Examples

Uses an eight-head transformer configuration, German–English translation samples, customized decoder code, and attention-weight plots.

What it assumes

Builds on token embeddings, query–key–value self-attention, BPE, and positional encoding before assembling transformer layers.

What you can do after

Apply it via PyTorch’s Transformer modules to train a translation model and inspect decoder attention weights.

How it connects

Links self-attention and cross-attention with embeddings, positional encoding, transformer blocks, layer normalization, and machine translation.

liveBookmanning.com

Read Natural Language Processing in Action, Second Edition on liveBook Get the book at manning.com

Also explained in

• 

  liveBookmanning.com
  Build a Text-to-Image Generator (from Scratch)An attention mechanism that computes attention through multiple separate query, key, and value subspaces to capture different relationships simultaneously.
• 

  liveBookmanning.com
  Transformers in ActionMultihead attention is an attention mechanism that computes multiple attention patterns in parallel so a model can represent different relationships among sequence positions.
• 

  liveBookmanning.com
  Learn Generative AI with PyTorchA self-attention variant that splits representations into multiple attention heads to capture different relationships in parallel.
• 

  liveBookmanning.com
  Inside Deep Learning: Math, Algorithms, ModelsAn attention mechanism that uses multiple attention heads to attend to different contexts or aspects of the input in parallel.
• 

  liveBookmanning.com
  Deep Learning with Python, Third EditionAn attention variant that runs several attention computations in parallel with different parameters to learn different relationships.
• 

  liveBookmanning.com
  Graph Neural Networks in ActionAn attention technique that computes multiple attention operations in parallel and combines their outputs.
• 

  liveBookmanning.com
  Build a Large Language Model (From Scratch)An attention technique that computes several independent attention heads in parallel and combines their outputs to capture different relationships in the same sequence.
• 

  liveBookmanning.com
  Deep Learning with R, Third EditionMultiHeadAttention is a Keras layer that computes attention through multiple parallel attention heads and combines their outputs.
• 

  liveBookmanning.com
  Time Series Forecasting Using Foundation ModelsAn attention mechanism with multiple parallel attention heads that combines information from different representation subspaces.
• 

  liveBookmanning.com
  Machine Learning Algorithms in DepthMultihead attention runs multiple attention mechanisms in parallel and combines their output representations.
• 

  liveBookmanning.com
  Deep Learning with PyTorch, Second EditionAn attention mechanism that runs several attention heads in parallel, each learning different projections and relationships.
• 

  liveBookmanning.com
  Optimization Algorithms: AI techniques for design, planning, and control problemsMulti-head self-attention computes several attention operations in parallel so a sequence representation can capture different relationships among its elements.
• 

  liveBookmanning.com
  The AI Pocket BookA transformer attention mechanism that splits representations into multiple heads to attend to different relationships in parallel.

Compare how 14 books explain multi-head attention →

Related concepts

• graph attention networkused by · A graph attention network is a graph neural network that uses attention weights to aggregate information from neighboring nodes, allowing their contributions to differ.
• grouped-query attentionequivalent to · Grouped-query attention partitions query heads into groups whose members share key and value projections, reducing attention computation and memory use.
• attentionkind of · Attention is a neural-network mechanism that computes input-dependent weights to combine information according to its relevance to a computation.
• attention headhas part · An attention head is an independently parameterized attention computation that evaluates relationships among input elements and combines information from them into a contextualized representation.
• transformer architectureused by · A transformer architecture is a neural-network architecture that processes sequences or other structured inputs using attention mechanisms and related transformations.
• encoderpart of · An encoder is a component that transforms input data into an encoded representation for subsequent processing, reconstruction, comparison, or transmission.

This explainer was generated by Manning's knowledge graph from the text of Natural Language Processing in Action, Second Edition. The section links point to the parts of the book it draws on.