Demo: Why Scale Attention Scores?

Dot-product variance grows with head width. We will widen a random attention head while holding sequence length and input variance fixed, then measure both score scale and attention concentration.

The setup defines a stable row-wise softmax and normalized entropy. Entropy near one means attention is spread across keys. Entropy near zero means it is concentrated on very few keys.

import numpy as np
import matplotlib.pyplot as plt

plt.style.use("../figure_style.mplstyle")
rng = np.random.default_rng(20)

def row_softmax(scores):
    shifted = scores - scores.max(axis=1, keepdims=True)
    weights = np.exp(shifted)
    return weights / weights.sum(axis=1, keepdims=True)

def normalized_entropy(weights):
    row_entropy = -(weights * np.log(weights + 1e-15)).sum(axis=1)
    maximum_entropy = np.log(weights.shape[1])
    return row_entropy / maximum_entropy

sequence_length = 16
dimensions = np.array([4, 16, 64, 256, 1024])
trials = 300
print(f"Averaging {trials * sequence_length:,} attention rows at each width.")

Averaging 4,800 attention rows at each width.

Question: Does Dividing by \(\sqrt{d_k}\) Stabilize Attention?

For each width, we compute unscaled scores \(\mathbf Q\mathbf K^\top\) and scaled scores \(\mathbf Q\mathbf K^\top/\sqrt{d_k}\). Predict how their standard deviations and softmax entropies change with \(d_k\).

score_std = {"Unscaled": [], "Scaled": []}
entropy = {"Unscaled": [], "Scaled": []}
for dk in dimensions:
    variances = {"Unscaled": [], "Scaled": []}
    entropies = {"Unscaled": [], "Scaled": []}
    for trial in range(trials):
        Q = rng.standard_normal((sequence_length, dk))
        K = rng.standard_normal((sequence_length, dk))
        scores = Q @ K.T
        for name, values in [("Unscaled", scores), ("Scaled", scores / np.sqrt(dk))]:
            weights = row_softmax(values)
            variances[name].append(np.mean(values ** 2))
            entropies[name].append(normalized_entropy(weights).mean())

    for name in score_std:
        # Scores have population mean zero, so E[score^2] is their variance.
        score_std[name].append(np.sqrt(np.mean(variances[name])))
        entropy[name].append(np.mean(entropies[name]))

Each trial uses an ordinary query matrix and key matrix, with one row per token. The paired calculations use the same scores, so scaling is the only difference.

fig, axes = plt.subplots(1, 2, figsize=(8, 3.5), layout="constrained")
for label, marker in [("Unscaled", "o"), ("Scaled", "s")]:
    axes[0].plot(dimensions, score_std[label], marker=marker, linewidth=1.5, markersize=4, label=label)
    axes[1].plot(dimensions, entropy[label], marker=marker, linewidth=1.5, markersize=4, label=label)
axes[0].plot(dimensions, np.sqrt(dimensions), linestyle=":",
             linewidth=1, label="$\\sqrt{d_k}$")
for ax in axes:
    ax.set_xscale("log", base=2)
    ax.set_xlabel("Head Width $d_k$")
    ax.legend(frameon=False)
axes[0].set_yscale("log")
axes[0].set_ylabel("Score Standard Deviation")
axes[1].set_ylabel("Mean Normalized Attention Entropy")
axes[1].set_ylim(0, 1.02)
plt.show()

print("Scaled score standard deviations:", np.round(score_std["Scaled"], 3))
print("Scaled attention entropies:       ", np.round(entropy["Scaled"], 3))

Scaled score standard deviations: [1.009 0.997 1.003 0.999 0.999]
Scaled attention entropies:        [0.855 0.854 0.85  0.851 0.851]

Answer. Without scaling, score magnitude grows like \(\sqrt{d_k}\) and the softmax becomes increasingly concentrated. Scaling keeps the score variance near one and makes the entropy nearly independent of width.

The factor removes concentration caused by width alone. Learned query-key alignment can still produce concentrated attention.