Cross-Attention › Implementation
Language AI HandbookAll chapters Cross-Attention › ImplementationOne section · Sections in this chapter7 / 14 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Let's implement a complete cross-attention module following the patterns used in production transformers. We'll build it step by step, starting with single-head attention and working up to the full multi-head version.
The implementation below follows the structure of the previous mathematical derivation closely, making it straightforward to verify that the code matches the formulas. In production code, you would use batched matrix operations via PyTorch or JAX, but the single-sequence NumPy implementation here makes the shapes and operations maximally transparent.
class CrossAttention:
"""
Cross-attention module for encoder-decoder transformers.
Queries come from the decoder, keys and values come from the encoder.
"""
def __init__(self, d_model, d_k, d_v):
"""
Initialize cross-attention with projection matrices.
Args:
d_model: Model dimension (size of input representations)
d_k: Query/key dimension
d_v: Value dimension
"""
self.d_k = d_k
self.scale = 1.0 / np.sqrt(d_k)
# Query projection (applied to decoder state)
self.W_Q = np.random.randn(d_model, d_k) * np.sqrt(
2.0 / (d_model + d_k)
)
# Key and value projections (applied to encoder output)
self.W_K = np.random.randn(d_model, d_k) * np.sqrt(
2.0 / (d_model + d_k)
)
self.W_V = np.random.randn(d_model, d_v) * np.sqrt(
2.0 / (d_model + d_v)
)
def __call__(self, decoder_state, encoder_output, encoder_mask=None):
"""
Apply cross-attention.
Args:
decoder_state: Current decoder representations (n_dec, d_model)
encoder_output: Encoder output representations (n_enc, d_model)
encoder_mask: Boolean mask for valid encoder positions (n_enc,)
Returns:
output: Contextualized decoder representations (n_dec, d_v)
attention_weights: Cross-attention weights (n_dec, n_enc)
"""
# Q from decoder, K and V from encoder
Q = decoder_state @ self.W_Q
K = encoder_output @ self.W_K
V = encoder_output @ self.W_V
# Compute attention scores
scores = Q @ K.T * self.scale
# Apply encoder padding mask if provided
if encoder_mask is not None:
mask = encoder_mask.reshape(1, -1)
scores = np.where(mask, scores, -1e9)
# Softmax and aggregate
attention_weights = softmax(scores)
output = attention_weights @ V
return output, attention_weightsThe Xavier initialization used for the projection matrices (scaling by ) ensures that activations have similar variance regardless of the network depth. This is especially important for cross-attention because it uses three different projection matrices, each of which could amplify or suppress activations if initialized poorly.
Let's test the module with a translation-like example:
# Simulate encoding "The quick brown fox"
source_tokens = ["The", "quick", "brown", "fox"]
n_source = len(source_tokens)
# Target tokens generated so far: "Le renard"
target_tokens = ["Le", "renard"]
n_target = len(target_tokens)
d_model = 16
d_k = d_v = 16
# Simulated representations
encoder_output = np.random.randn(n_source, d_model)
decoder_state = np.random.randn(n_target, d_model)
# Create and apply cross-attention
cross_attn = CrossAttention(d_model, d_k, d_v)
output, weights = cross_attn(decoder_state, encoder_output)Cross-attention example: English → French
Source: ['The', 'quick', 'brown', 'fox']
Target so far: ['Le', 'renard']
Encoder output shape: (4, 16)
Decoder state shape: (2, 16)
Cross-attention output shape: (2, 16)
Attention weights (target × source):
['The', 'quick', 'brown', 'fox']
Le [0.285 0.365 0.304 0.045]
renard [0.049 0.171 0.729 0.051]Each target token distributes its attention across all source tokens. The weights indicate which parts of the source are most relevant for each target position. In a trained model, "renard" (French for "fox") would attend strongly to "fox" in the English source, because the training signal would have taught the model to align these semantically equivalent words. With random projections, the weights are essentially uniform, which is expected: the model has not yet learned any useful alignments.
