Logit scaling before the final softmax

The problem, from first principles

[04-seq-modeling/04-attention/01-scaled-dot-product-attention] divided raw attention scores by 1/sqrt(d_k) for a concrete numerical reason: a dot product between two d_k-dimensional vectors has variance proportional to d_k, so without scaling, larger dimensions produce larger-magnitude scores, pushing softmax toward its SATURATED regime (extremely close to one-hot, near-zero gradient almost everywhere). [09-untied-embeddings]'s final output projection, hidden_states @ output_weight.T, is ALSO a dot product, between a d_model-dimensional hidden state and each vocabulary word's d_model-dimensional row, and it has the EXACT SAME statistical problem: the resulting logits' magnitude naturally grows with d_model, and without any correction, a model with a larger hidden dimension produces disproportionately SHARPER, more overconfident next-token probability distributions, purely as a side effect of its width, unrelated to how genuinely confident the underlying prediction actually should be.

Dividing the final logits by 1/sqrt(d_model) before softmax is the exact same fix, applied at the exact same conceptual spot in the computation, just one layer later in the model, at the very END rather than inside every attention call.

From theory to code

Implement scale_logits_before_softmax(logits, d_model), a direct logits / sqrt(d_model).

Constraints

• Scales by 1 / sqrt(d_model), not 1 / d_model (an easy, meaningfully different mistake, [01-scaled-dot-product-attention]'s own scaling used the same square root for the same reason).
• Applied to the LOGITS (the RAW output of [09-untied-embeddings]'s output projection), strictly BEFORE softmax is computed, never after.
• A purely elementwise, linear operation: scaling every logit by the same constant factor.

Hints