<p>I'm trying to visualize how the gradients change as we're increasing <span class="math-container">$d_{k}$</span> in the scaled dot-product attention and compare it to its non scaled version but I'm failing to produce a reasonable graph that transcribes the theory. The authors of the <a href="https://arxiv.org/abs/1706.03762" rel="nofollow noreferrer">Attention Is All You Need</a> paper explain that non scaling the dot-product leads to having large values which pushes the softmax into regions where its gradients are small, which is understandable because the values inside the softmax have such a high variance around a mean of 0 that the softmax will produce a mixture of values very close to 0 and 1.</p>
<p>I'm trying to visualize that using the following Python code:</p>
<pre><code>def scaled_softmax_dot_product(Q, K, scale):
    attention_logits = torch.matmul(Q, K.T)
    attention_logits /= scale
    return F.softmax(attention_logits, dim=-1)

mean_grads_with_scaling = []
mean_grads_without_scaling = []
dk_values = range(1, 321)

for d_k in dk_values:
    Q = torch.randn(16, d_k, requires_grad=True)
    K = torch.randn(16, d_k, requires_grad=True)

    # Compute attention with scaling
    output_scaled = scaled_softmax_dot_product(Q, K, scale=torch.sqrt(torch.tensor(d_k, dtype=torch.float)))
    output_scaled.backward(torch.ones_like(output_scaled))
    grad_with_scaling = Q.grad.mean().item()
    Q.grad.zero_()
    K.grad.zero_()

    # Compute attention without scaling
    output_unscaled = scaled_softmax_dot_product(Q, K, scale=1.0)
    output_unscaled.backward(torch.ones_like(output_unscaled))
    grad_without_scaling = Q.grad.mean().item()
    Q.grad.zero_()
    K.grad.zero_()

    mean_grads_with_scaling.append(grad_with_scaling)
    mean_grads_without_scaling.append(grad_without_scaling)

plt.figure(figsize=(10, 6))
plt.plot(dk_values, mean_grads_with_scaling, label='With Scaling', color='blue')
plt.plot(dk_values, mean_grads_without_scaling, label='Without Scaling', color='red')
plt.xlabel('Dimension of Keys (d_k)')
plt.ylabel('Mean Gradients')
plt.title('Evolution of Mean Gradients with and without Scaling')
plt.legend()
plt.show()
</code></pre>
<p>As you can see in the code I'm just computing the softmax of the non/scaled dot-product with respect to <span class="math-container">$Q$</span> and I was expecting the mean of the gradients for the non scaled dot-product to converge to <span class="math-container">$0$</span> as <span class="math-container">$d_k$</span> increases but that's not what the graph shows:</p>
<p><a href="https://i.sstatic.net/EjCOV.png" rel="nofollow noreferrer"><img src="https://i.sstatic.net/EjCOV.png" alt="enter image description here" /></a></p>
<p>I'd love to understand what is the error that I'm making to for getting this graph, or if I am misunderstanding the statement.</p>

CC BY-SA 4.0
2024-01-04T20:04:18.077Z
I think is due to numerical instability. When you do `attention_logits /= scale` you should add a small constant to avoid zero division. For example with `eps = 1e-7`, you could `attention_logits /= scale + eps`.
Luca Anzalone
CC BY-SA 4.0
2024-01-05T06:28:25.483Z
@LucaAnzalone I added the small constant and I'm still getting the same graph. I don't think the issue is due to division by zero because that's never happening in the above code, we're either dividing by $1.0$ or $\sqrt{d_{k}} >= 1.0$
Daviiid
<transformer><backpropagation><pytorch><attention><gradient>
What am I doing wrong that result in a graph indicating better gradients in non-scaled dot-product attention compared to the scaled version?
Daviiid
# What am I doing wrong that result in a graph indicating better gradients in non-scaled dot-product attention compared to the scaled version?

(Asked by: username_0 on 2024-01-04)

I'm trying to visualize how the gradients change as we're increasing $d_{k}$ in the scaled dot-product attention and compare it to its non scaled version but I'm failing to produce a reasonable graph that transcribes the theory. The authors of the Attention Is All You Need paper explain that non scaling the dot-product leads to having large values which pushes the softmax into regions where its gradients are small, which is understandable because the values inside the softmax have such a high variance around a mean of 0 that the softmax will produce a mixture of values very close to 0 and 1.
I'm trying to visualize that using the following Python code:
def scaled_softmax_dot_product(Q, K, scale):
    attention_logits = torch.matmul(Q, K.T)
    attention_logits /= scale
    return F.softmax(attention_logits, dim=-1)

mean_grads_with_scaling = []
mean_grads_without_scaling = []
dk_values = range(1, 321)

for d_k in dk_values:
    Q = torch.randn(16, d_k, requires_grad=True)
    K = torch.randn(16, d_k, requires_grad=True)

    # Compute attention with scaling
    output_scaled = scaled_softmax_dot_product(Q, K, scale=torch.sqrt(torch.tensor(d_k, dtype=torch.float)))
    output_scaled.backward(torch.ones_like(output_scaled))
    grad_with_scaling = Q.grad.mean().item()
    Q.grad.zero_()
    K.grad.zero_()

    # Compute attention without scaling
    output_unscaled = scaled_softmax_dot_product(Q, K, scale=1.0)
    output_unscaled.backward(torch.ones_like(output_unscaled))
    grad_without_scaling = Q.grad.mean().item()
    Q.grad.zero_()
    K.grad.zero_()

    mean_grads_with_scaling.append(grad_with_scaling)
    mean_grads_without_scaling.append(grad_without_scaling)

plt.figure(figsize=(10, 6))
plt.plot(dk_values, mean_grads_with_scaling, label='With Scaling', color='blue')
plt.plot(dk_values, mean_grads_without_scaling, label='Without Scaling', color='red')
plt.xlabel('Dimension of Keys (d_k)')
plt.ylabel('Mean Gradients')
plt.title('Evolution of Mean Gradients with and without Scaling')
plt.legend()
plt.show()

As you can see in the code I'm just computing the softmax of the non/scaled dot-product with respect to $Q$ and I was expecting the mean of the gradients for the non scaled dot-product to converge to $0$ as $d_k$ increases but that's not what the graph shows:

I'd love to understand what is the error that I'm making to for getting this graph, or if I am misunderstanding the statement.

**Tags:** transformer, backpropagation, pytorch, attention, gradient

> **Comments:**
> **username_1**: I think is due to numerical instability. When you do `attention_logits /= scale` you should add a small constant to avoid zero division. For example with `eps = 1e-7`, you could `attention_logits /= scale + eps`.  
> **username_0**: @LucaAnzalone I added the small constant and I'm still getting the same graph. I don't think the issue is due to division by zero because that's never happening in the above code, we're either dividing by $1.0$ or $\sqrt{d_{k}} >= 1.0$  

---

## 0 answers

