√dₖ and softmax saturation (row entropy)

Background

Without the 1/√d_k scale, attention scores grow with the head dimension and softmax collapses to a near one-hot — "hard" attention with dead gradients. A clean way to measure softness is the average entropy of the attention rows: high entropy = broad/soft, near-zero = saturated.

For raw scores S and head dimension d_k:

$$A = \text{softmax}\!\left(\frac{S}{\sqrt{d_k}}\right), \qquad \bar H = \text{mean over rows of } \left(-\sum_j A_{ij}\ln A_{ij}\right)$$

Problem statement

Implement attention_entropy(scores, head_size) returning the mean per-row entropy of the scaled-softmax attention.

Input

• scores — raw QKᵀ scores, shape (T, T).
• head_size — int, the per-head dim d_k.

Output

Returns a float: the mean over rows of the row entropy (natural log).

Examples

Example 1 — uniform scores → maximal entropy

Input:  scores = [[0, 0, 0, 0]], head_size = 64
Output: 1.3862943611...   # ln(4)

Explanation: equal scores → uniform softmax over 4 keys → entropy ln 4.

Example 2 — large unscaled spread → saturated

Input:  scores = [[100, 0, 0, 0]], head_size = 1
Output: ~0.0

Constraints

• Scale: scores / sqrt(head_size).
• Row-wise (numerically stable) softmax, then -Σ A·ln A per row (use a small epsilon so ln 0 is safe).
• Return the mean row entropy as a float.

Notes

• Bigger head_size divides the scores down, raising entropy back toward ln T — which is exactly why √d_k keeps softmax in its learnable regime.