One Attention Reading Is Not Yet a Transformer Layer

An attention reading answers what one position retrieves, but a reusable layer also needs a stable input-output width and a batch-safe interface. Add the output projection and identify which parts of a complete Transformer are still absent.

Suppose a query produces the attention reading

h_i=\sum_{j=1}^{T}a_{ij}v_j.

This equation answers a local question: what value vector should position $i$ read from the positions it is allowed to inspect? A neural-network layer must answer a larger question: what input does it accept, and what output can the next part of the network rely on?

Recall the Central Operation

For one sequence and one head,

Q=XW_Q,\qquad K=XW_K,\qquad V=XW_V,

A=\operatorname{softmax}_{last}\!\left( \frac{QK^\top}{\sqrt{d_k}}+M \right),\qquad H=AV.

Each row of $A$ belongs to one query position. Each column refers to a position whose value may be read. In causal attention, $M$ prevents a query from reading future positions.

Language Modeling derived this operation and calculated a complete four-token example. We retain it as the starting point. The new question is how $H$ becomes the output of a stable layer.

A Layer Needs a Stable Width

The input matrix has shape

X:T\times d_{model}.

The value projection is allowed to choose a different width:

W_V:d_{model}\times d_v,\qquad H:T\times d_v.

If $d_v\ne d_{model}$ , then $H$ cannot be added to $X$ , and the next layer cannot assume that the model width stayed fixed. An output projection solves this interface problem:

W_O:d_v\times d_{model},\qquad Y=HW_O:T\times d_{model}.

Even when $d_v=d_{model}$ , $W_O$ is not redundant. It learns how the features produced by attention are written back into model-width coordinates. Chapter 2 will also use it to mix the concatenated outputs of several heads.

The Interface Extends to Batches

Training normally processes several sequences together. Adding a batch axis changes the shapes but not the operation:

Quantity	Single sequence	Batch
Input $X$	$(T,d_{model})$	$(B,T,d_{model})$
Queries and keys	$(T,d_k)$	$(B,T,d_k)$
Values	$(T,d_v)$	$(B,T,d_v)$
Scores and weights	$(T,T)$	$(B,T,T)$
Layer output $Y$	$(T,d_{model})$	$(B,T,d_{model})$

The computation is independent for each batch item. A token in sequence 1 must never acquire an attention edge to sequence 2 merely because both are stored in one tensor.

Attention Predates the Transformer

Bahdanau, Cho, and Bengio proposed a learned soft alignment for neural machine translation in 2014. At each decoder step, that model searched over source annotations instead of relying only on one fixed-length source summary. The source and target sides had different roles.

The 2017 Transformer used attention throughout an encoder-decoder architecture and distinguished self-attention, where Q, K, and V come from the same sequence, from encoder-decoder attention, where queries and key-value records come from different sequences. The Transformer combined attention with multiple heads, position information, residual paths, normalization, and feed-forward layers. No single attention reading is the whole architecture.

Q1. Restore the layer width

A single head produces $H$ with shape $(12,6)$ . The model width is 10. What shape must $W_O$ have so that $Y=HW_O$ has model width 10?

Choose the output-projection shape

6 × 10

10 × 6

12 × 10

Select one choice, then check.

Hint

The input width of

W_O

must equal the final width of

H

Solution

W_O

has shape

(6,10)

. Then

(12,6)(6,10)=(12,10)

One attention reading is only the central operation inside a reusable self-attention layer. The output projection restores the model width and the batch interface keeps sequences independent. A complete Transformer block still needs residual updates, normalization, and a position-wise nonlinear sublayer.

References

Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio, Neural Machine Translation by Jointly Learning to Align and Translate (2014). The paper replaces exclusive reliance on one fixed source vector with a learned soft search over source annotations.
Ashish Vaswani et al., Attention Is All You Need (2017). Sections 3.1 and 3.2 define the original Transformer architecture and its attention mechanisms.