How can the self-attention mechanism of a transformer account for contextual language contraints?
Try it by yourself.
Imagine a tiny toy language containing only the words
"we", "you", "are", "hear", "here", "why".
We have designed a simple generative pretrained transformer able to predict that
the next word of a sequence like "why are you",
is "here", while "you" alone would
probably most often be followed by "are".
Follow the modification of the word embedding of "you"
through a simplified self-attention mechanism and see how context changes the next-word
prediction from "are" to "here".
3-D Word embedding space
shaded side: decoder favors are
shaded side: decoder favors here.
What this step does
Prediction at this stage
2-D attention space
Comments
This demonstration uses a deliberately tiny vocabulary:
"we", "you", "are", "hear", "here", "why".
It follows the sequence "why are you"
and asks how context can change the model's prediction of the next word.
The model uses two distinct vector spaces. Word embeddings and all
context-dependent word representations live in a 3-D Word embedding
space, whereas attention-related features (Queries, Keys and Values) have dimension 2. The matrices
WQ,
WK and
WV
map from the 3-D Word embedding space into the 2-D attention space, while
WO
maps back from the 2-D attention space to the 3-D Word embedding space.
The representation of "you" initially lies in a
region of the 3-D Word embedding space where the decoder favors
"are". Attention gathers information from the
previous words and moves the representation across the fixed decoder
boundary into the region favoring "here".
Decoder regions in the 3-D Word embedding space
For any point in the 3-D embedding space
r=[x,y,z], the output decoder assigns one
score to each possible next word w in ("we", "you", "are", "hear", "here", "why")
score(w) =
Wout(w) · r
+
bout(w)
In this deliberately simplified example, all vocabulary items except
"are" and
"here" have very low scores, so the 3-D Word
embedding space is effectively partitioned into an
"are" region (in blue) and a
"here" region (in red).
Since this simplified decoder is linear - while real transformers use DNN-based non-linear classifiers - , the areas favoring "are" and "here" are separated by a plane where the two scores are equal:
score(w=are) = score(w=here)
The boundary does not move while the processing step changes.
The decoder is fixed; it is the word representation that moves through
the 3-D Word embedding space.
Step 0 — 3-D word embeddings
Each word begins as a fixed vector in the 3-D Word embedding space:
eᵢ ∈ ℝ³
No position or context is present yet.
Embedding table used by this example
Step 1 — Add 3-D positional information
A positional vector is added to every word embedding:
xᵢ = eᵢ + pᵢ
The result is still a 3-D vector in the same Word embedding space.
The coordinates are still described by the same three embedding
dimensions.
Positional vectors
Actual positioned representations
Step 2 — Project from 3-D into 2-D Query, Key and Value spaces
Three rectangular matrices map each 3-D word representation into
a 2-D attention vector:
The Query represents what the current word is looking for;
a Key represents what a word advertises as potentially relevant;
and a Value represents what information that word can contribute
if it receives attention.
For intance, to make it simple for tutorial purposes, the Key for "are"
could mean somthing like "I am a verb" and the key for "why" could mean "this is a question".
None of these Keys alone impose the next word to be "here", as
a sequence like "hear why you" would provide similar Keys while forbidding "here" as the next word.
It is the combination of these Keys which constrain the choice to "here".
WQ — 2 × 3
WK — 2 × 3
WV — 2 × 3
Actual 2-D Q, K and V vectors
Step 3 — Compare the final Query with the Keys
The Query of the final word
you is compared with the Key of each current and previous words,
via the computation of a 2-D scalar product:
attention_score(you,j) =
qyou · kⱼ / √2
The semantic interpretation is:
what is “you” looking for? is compared with
what can each available word provide or match on?
Actual attention scores
Step 4 — Softmax determines how much to listen to each word
αⱼ = softmax(score(you,j))
These weights answer:
how much should “you” listen to each available word?
Actual attention weights
Step 5 — Mix the 2-D Value vectors
Keys decide relevance; Values contain the contextual information that is actually
transferred:
a(you) = Σⱼ αⱼ vⱼ
Each word contributes its Value in proportion to its attention weight.
Actual weighted Value calculation for the final word
Step 6 — Translate the gathered information back into the 3-D Word embedding space
c(you) = WO a(you)
WO ∈ ℝ³ˣ²
WO translates the information
gathered by attention into a displacement in the 3-D Word embedding
space. The pink vector is therefore drawn from the end of the positioned
"you" vector rather than from the origin.
WO — 3 × 2
Actual 2-D attention output and 3-D correction
Step 7 — Add the context correction
h(you) = x(you) + c(you)
The final contextual representation is obtained by adding the context
correction to the positioned representation.
The embedding, positioned and contextual versions of
"you" keep the same color because they are three
successive versions of the same word representation. Dotted vectors show
the earlier embedding and positioned states, while the solid vector shows
the contextual representation. The pink
"+ context" displacement shows exactly how
attention moves the positioned representation to its contextual position.
Actual final representations
Output decoder
score(w) =
Wout(w) · h
+
bout(w)
The final contextual representation h(you) is sent to the
fixed output decoder. The decoder first produces one raw score for each candidate
next word. These raw scores are not probabilities yet.
Output weights and biases
Final decoder scores before softmax
Final probabilities after softmax
Summarizing
Prediction at the three representation checkpoints
About the two graphs
The graph on the left always represents the 3-D Word embedding space.
Its three axes are the three embedding dimensions. Its camera, axis
ranges and decoder boundary stay fixed.
The lightly shaded volumes indicate which of
are and
here receives the larger decoder score.
The graph on the right shows the separate 2-D attention space during
the internal attention steps. Its ranges are calculated once from the
complete forward pass.
This demonstration uses a deliberately tiny vocabulary: "we", "you", "are", "hear", "here", "why".
It follows the sequence "why are you" and asks how context can change the model's prediction of the next word.
The model uses two distinct vector spaces. Word embeddings and all context-dependent word representations live in a 3-D Word embedding space, whereas attention-related features (Queries, Keys and Values) have dimension 2. The matrices WQ, WK and WV map from the 3-D Word embedding space into the 2-D attention space, while WO maps back from the 2-D attention space to the 3-D Word embedding space.
The representation of "you" initially lies in a region of the 3-D Word embedding space where the decoder favors "are". Attention gathers information from the previous words and moves the representation across the fixed decoder boundary into the region favoring "here".
Decoder regions in the 3-D Word embedding space
For any point in the 3-D embedding space r=[x,y,z], the output decoder assigns one score to each possible next word w in ("we", "you", "are", "hear", "here", "why")
In this deliberately simplified example, all vocabulary items except "are" and "here" have very low scores, so the 3-D Word embedding space is effectively partitioned into an "are" region (in blue) and a "here" region (in red).
Since this simplified decoder is linear - while real transformers use DNN-based non-linear classifiers - , the areas favoring "are" and "here" are separated by a plane where the two scores are equal:
The boundary does not move while the processing step changes. The decoder is fixed; it is the word representation that moves through the 3-D Word embedding space.
Step 0 — 3-D word embeddings
Each word begins as a fixed vector in the 3-D Word embedding space:
No position or context is present yet.
Step 1 — Add 3-D positional information
A positional vector is added to every word embedding:
The result is still a 3-D vector in the same Word embedding space. The coordinates are still described by the same three embedding dimensions.
Step 2 — Project from 3-D into 2-D Query, Key and Value spaces
Three rectangular matrices map each 3-D word representation into a 2-D attention vector:
kᵢ = WK xᵢ WK ∈ ℝ²ˣ³
vᵢ = WV xᵢ WV ∈ ℝ²ˣ³
The Query represents what the current word is looking for; a Key represents what a word advertises as potentially relevant; and a Value represents what information that word can contribute if it receives attention.
For intance, to make it simple for tutorial purposes, the Key for "are" could mean somthing like "I am a verb" and the key for "why" could mean "this is a question". None of these Keys alone impose the next word to be "here", as a sequence like "hear why you" would provide similar Keys while forbidding "here" as the next word. It is the combination of these Keys which constrain the choice to "here".
Step 3 — Compare the final Query with the Keys
The Query of the final word you is compared with the Key of each current and previous words, via the computation of a 2-D scalar product:
The semantic interpretation is: what is “you” looking for? is compared with what can each available word provide or match on?
Step 4 — Softmax determines how much to listen to each word
These weights answer: how much should “you” listen to each available word?
Step 5 — Mix the 2-D Value vectors
Keys decide relevance; Values contain the contextual information that is actually transferred:
Each word contributes its Value in proportion to its attention weight.
Step 6 — Translate the gathered information back into the 3-D Word embedding space
WO translates the information gathered by attention into a displacement in the 3-D Word embedding space. The pink vector is therefore drawn from the end of the positioned "you" vector rather than from the origin.
Step 7 — Add the context correction
The final contextual representation is obtained by adding the context correction to the positioned representation.
The embedding, positioned and contextual versions of "you" keep the same color because they are three successive versions of the same word representation. Dotted vectors show the earlier embedding and positioned states, while the solid vector shows the contextual representation. The pink "+ context" displacement shows exactly how attention moves the positioned representation to its contextual position.
Output decoder
The final contextual representation h(you) is sent to the fixed output decoder. The decoder first produces one raw score for each candidate next word. These raw scores are not probabilities yet.
Summarizing
About the two graphs
The graph on the left always represents the 3-D Word embedding space. Its three axes are the three embedding dimensions. Its camera, axis ranges and decoder boundary stay fixed.
The lightly shaded volumes indicate which of are and here receives the larger decoder score.
The graph on the right shows the separate 2-D attention space during the internal attention steps. Its ranges are calculated once from the complete forward pass.