Why scaled dot product

Attention Scoring Functions

Scoring Functions

Attention pooling needs a scoring function a(\mathbf{q}, \mathbf{k}) that softmax turns into weights:

\alpha(\mathbf{q}, \mathbf{k}_i) = \frac{\exp\, a(\mathbf{q}, \mathbf{k}_i)}{\sum_j \exp\, a(\mathbf{q}, \mathbf{k}_j)}.

Output = weighted sum of values; weights = softmax of scoring function a.

Two scorings dominate practice

• Scaled dot product — a = \mathbf{q}^\top \mathbf{k}/\sqrt{d}. Cheap, parameter-free; query and key share a dimension. The Transformer choice.
• Additive (Bahdanau) — a tiny MLP over [\mathbf{q}; \mathbf{k}]. More expressive, learns the metric, allows different \mathbf{q}/\mathbf{k} shapes.

Both feed into the same softmax + value-pooling pipeline.

Setup

from d2l import tensorflow as d2l
import tensorflow as tf

For d-dimensional queries and keys with independent, zero-mean, unit-variance coordinates,

\operatorname{Var}(\mathbf{q}^\top\mathbf{k}) = \operatorname{Var}\left(\sum_{\ell=1}^d q_\ell k_\ell\right) = d.

As d grows, raw dot products become large in magnitude, softmax saturates, and gradients shrink. Scaling by 1/\sqrt d keeps the logit variance approximately constant:

a(\mathbf{q}, \mathbf{k}_i) = \mathbf{q}^\top \mathbf{k}_i / \sqrt{d}.

A Gaussian-kernel view gives useful geometric intuition, but the variance argument is the operational reason used in Transformers.

Masked softmax

Padded sequences in a minibatch — we don’t want <pad> keys to receive attention mass. Set their pre-softmax scores to a large negative number so \exp flushes them to zero:

def masked_softmax(X, valid_lens):
    """Perform softmax operation by masking elements on the last axis."""
    # X: 3D tensor, valid_lens: 1D or 2D tensor
    def _sequence_mask(X, valid_len, value=0):
        maxlen = tf.shape(X)[1]
        mask = tf.range(start=0, limit=maxlen, dtype=tf.float32)[
            None, :] < tf.cast(valid_len[:, None], dtype=tf.float32)
        return tf.where(mask, X, value)

    if valid_lens is None:
        return tf.nn.softmax(X, axis=-1)
    else:
        shape = tf.shape(X)
        if len(valid_lens.shape) == 1:
            valid_lens = tf.repeat(valid_lens, repeats=shape[1])
        else:
            valid_lens = tf.reshape(valid_lens, shape=(-1,))
        # On the last axis, replace masked elements with a very large negative
        # value, whose exponentiation outputs 0
        X = _sequence_mask(tf.reshape(X, (-1, shape[-1])), valid_lens,
                           value=-1e6)
        return tf.nn.softmax(tf.reshape(X, shape), axis=-1)

Masked softmax in action

Random scores; specify a valid length per row:

masked_softmax(tf.random.uniform(shape=(2, 2, 4)), tf.constant([2, 3]))

<tf.Tensor: shape=(2, 2, 4), dtype=float32, numpy=
array([[[0.6731507 , 0.32684928, 0.        , 0.        ],
        [0.37237114, 0.6276289 , 0.        , 0.        ]],

       [[0.35941443, 0.35476714, 0.28581837, 0.        ],
        [0.406084  , 0.30974853, 0.28416747, 0.        ]]], dtype=float32)>

Per-row mask vectors work too:

masked_softmax(tf.random.uniform((2, 2, 4)), tf.constant([[1, 3], [2, 4]]))

<tf.Tensor: shape=(2, 2, 4), dtype=float32, numpy=
array([[[1.        , 0.        , 0.        , 0.        ],
        [0.3033781 , 0.2762389 , 0.42038298, 0.        ]],

       [[0.55870575, 0.44129428, 0.        , 0.        ],
        [0.19670235, 0.2990804 , 0.14326371, 0.36095354]]], dtype=float32)>

Batched matmul

Attention runs in batches; weights × values is a batched matmul. bmm does the right thing — confirm shapes:

Q = d2l.ones((2, 3, 4))
K = d2l.ones((2, 4, 6))
d2l.check_shape(tf.matmul(Q, K).numpy(), (2, 3, 6))

DotProductAttention class

Stateless layer — no parameters, just \mathbf{Q}\mathbf{K}^\top/\sqrt d, masked softmax, then weighted sum of values:

class DotProductAttention(tf.keras.layers.Layer):
    """Scaled dot product attention."""
    def __init__(self, dropout):
        super().__init__()
        self.dropout = tf.keras.layers.Dropout(dropout)
        
    # Shape of queries: (batch_size, no. of queries, d)
    # Shape of keys: (batch_size, no. of key-value pairs, d)
    # Shape of values: (batch_size, no. of key-value pairs, value dimension)
    # Shape of valid_lens: (batch_size,) or (batch_size, no. of queries)
    def call(self, queries, keys, values, valid_lens=None, training=False,
             **kwargs):
        d = tf.cast(tf.shape(queries)[-1], dtype=tf.float32)
        scores = tf.matmul(queries, keys, transpose_b=True)/tf.math.sqrt(d)
        self.attention_weights = masked_softmax(scores, valid_lens)
        return tf.matmul(self.dropout(
            self.attention_weights, training=training), values)

DotProduct demo

2 queries, 10 keys/values, valid lengths (2, 6) — only the first 2 / first 6 keys per batch get nonzero weight:

queries = tf.random.normal(shape=(2, 1, 2))
keys = tf.random.normal(shape=(2, 10, 2))
values = tf.random.normal(shape=(2, 10, 4))
valid_lens = tf.constant([2, 6])

attention = DotProductAttention(dropout=0.5)
d2l.check_shape(attention(queries, keys, values, valid_lens, training=False),
                (2, 1, 4))

Visualize the resulting attention matrix:

d2l.show_heatmaps(d2l.reshape(attention.attention_weights, (1, 1, 2, 10)),
                  xlabel='Keys', ylabel='Queries')

AdditiveAttention class

a(\mathbf{q}, \mathbf{k}) = \mathbf{w}_v^\top \tanh(\mathbf{W}_q\mathbf{q} + \mathbf{W}_k\mathbf{k}). Learnable \mathbf{W}_q, \mathbf{W}_k, \mathbf{w}_v. Lets queries and keys live in different feature spaces.

class AdditiveAttention(tf.keras.layers.Layer):
    """Additive attention."""
    def __init__(self, key_size, query_size, num_hiddens, dropout, **kwargs):
        super().__init__(**kwargs)
        self.W_k = tf.keras.layers.Dense(num_hiddens, use_bias=False)
        self.W_q = tf.keras.layers.Dense(num_hiddens, use_bias=False)
        self.w_v = tf.keras.layers.Dense(1, use_bias=False)
        self.dropout = tf.keras.layers.Dropout(dropout)
        
    def call(self, queries, keys, values, valid_lens, training=False, **kwargs):
        queries, keys = self.W_q(queries), self.W_k(keys)
        # After dimension expansion, shape of queries: (batch_size, no. of
        # queries, 1, num_hiddens) and shape of keys: (batch_size, 1, no. of
        # key-value pairs, num_hiddens). Sum them up with broadcasting
        features = tf.expand_dims(queries, axis=2) + tf.expand_dims(
            keys, axis=1)
        features = tf.nn.tanh(features)
        # There is only one output of self.w_v, so we remove the last
        # one-dimensional entry from the shape. Shape of scores: (batch_size,
        # no. of queries, no. of key-value pairs)
        scores = tf.squeeze(self.w_v(features), axis=-1)
        self.attention_weights = masked_softmax(scores, valid_lens)
        # Shape of values: (batch_size, no. of key-value pairs, value
        # dimension)
        return tf.matmul(self.dropout(
            self.attention_weights, training=training), values)

Additive demo

Same shapes as before, with mismatched query/key dims allowed:

queries = tf.random.normal(shape=(2, 1, 20))

attention = AdditiveAttention(key_size=2, query_size=20, num_hiddens=8,
                              dropout=0.1)
d2l.check_shape(attention(queries, keys, values, valid_lens, training=False),
                (2, 1, 4))

Visualize:

d2l.show_heatmaps(d2l.reshape(attention.attention_weights, (1, 1, 2, 10)),
                  xlabel='Keys', ylabel='Queries')

Recap

• Scoring function a + softmax \Rightarrow attention weights; pool values with those weights.
• Scaled dot product is the default — cheap, parameter-free, scales by 1/\sqrt d to control softmax saturation.
• Additive attention is more flexible (separate \mathbf{q}/ \mathbf{k} shapes, learned metric) but slower and less used at modern scale.
• Masked softmax is the workhorse for handling padded sequences in batched inference.