[Python] Why is attention scaled by sqrt(d_k) in Transformer architectures?

I have this code in transformer model:

keys = x @ W_key
queries = x @ W_query
values = x @ W_value
attention_scores = queries @ keys.T
# keys.shape[-1]**0.5: used to scale the attention scores before applying the softmax
attn_weights = torch.softmax(
attention_scores / keys.shape[-1]**0.5, dim=-1
)


I understand that keys.shape[-1]**0.5 helps stabilize training by preventing large dot products from pushing softmax into sharp regions. But what’s the statistical or mathematical reasoning for choosing this value? Instead of using a fixed parameter, why a learnable parameter is not used here, Wouldn’t that allow the model to adapt better?

Continue reading...