Body: <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>
  
Comments.ContentLicense: CC BY-SA 4.0
Comments.CreationDate: 2024-01-04T20:04:18.077Z
Comments.Text: 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`.
Comments.UserDisplayName: Luca Anzalone
Comments.ContentLicense: CC BY-SA 4.0
Comments.CreationDate: 2024-01-05T06:28:25.483Z
Comments.Text: @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$
Comments.UserDisplayName: Daviiid
Tags: <transformer><backpropagation><pytorch><attention><gradient>
Title: What am I doing wrong that result in a graph indicating better gradients in non-scaled dot-product attention compared to the scaled version?
UserDisplayName: Daviiid
ThreadText: # 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
  
  