Cross-Attention - PyTorch Basics

Cross-Attention

What you'll build: cross_attention(Q, K, V): scaled dot-product attention where the queries come from one sequence and the keys and values from another, possibly of a different length.

Background

In self-attention, Q, K and V are all projections of the same sequence. In cross-attention the queries come from one sequence (the decoder's tokens, or an image's latent pixels) and the keys and values from another (the encoder's output, or a text prompt). It is how a translation decoder looks at the source sentence and how Stable Diffusion conditions an image on text.

Formula

scores  = Q @ K.T / sqrt(d_k)          (seq_q, seq_kv)
weights = softmax(scores, axis=-1)     each row sums to 1: query i's distribution over the keys
output  = weights @ V                  (seq_q, d_v)

Dividing by sqrt(d_k) keeps the scores' variance around 1 for random unit-variance inputs; without it, large d_k pushes the softmax into a near one-hot regime with vanishing gradients.

Inputs

Name	Shape	Meaning
Q	(seq_q, d_k)	queries
K	(seq_kv, d_k)	keys (same feature size as Q)
V	(seq_kv, d_v)	values (same length as K)

Expected Return Format

An array of shape (seq_q, d_v): one output per query, a weighted average of the rows of V.

Edge Cases and Conventions

• No mask and no batch or head dimensions.
• Use a numerically stable softmax (subtract the row max).
• If all keys are identical, every query attends uniformly and gets the mean of V.

Constraints

• Vectorised matrix products; no loops over queries.
• PyTorch track: F.softmax is fine; don't call F.scaled_dot_product_attention or nn.MultiheadAttention.

Follow-up

• What is the time and memory complexity in seq_q and seq_kv?
• In an encoder-decoder transformer, which cross-attention tensors can be cached across decoding steps?

Input: cross_attention(np.array([[2.0, 0.0]]), np.array([[1.0, 0.0], [0.0, 1.0]]), np.array([[1.0, 0.0], [0.0, 1.0]]))
Output: array([[0.8044, 0.1956]])

Input: Q (4, 8), K (6, 8), V (6, 16)
Output: shape (4, 16)

import numpy as np

def cross_attention(Q, K, V):
    """
    Scaled dot-product cross-attention.

    Args:
        Q: Queries (seq_len_q, d_k)
        K: Keys (seq_len_kv, d_k)
        V: Values (seq_len_kv, d_v)

    Returns:
        output: (seq_len_q, d_v)
    """
    # Your code here
    pass