Answers.Body: <p>The implementation of self-attention in the <a href="https://github.com/heykeetae/Self-Attention-N/blob/master/sagan_models.py" rel="nofollow noreferrer">source code</a> for the <a href="https://arxiv.org/abs/1805.08318" rel="nofollow noreferrer">&quot;Self-Attention Generative Adversarial Networks&quot;</a> (SAGAN) paper is somewhat easier to follow than that in the &quot;Denoising Diffusion Probabilistic Models&quot; paper. Assuming that both papers are using the same approach to self-attention on feature maps (which I think they are), the question can be answered by inspecting the definition of self-attention used in SAGAN:</p>
  
  <pre><code>class Self_Attn(nn.Module):
  &quot;&quot;&quot; Self attention Layer&quot;&quot;&quot;
  def __init__(self,in_dim,activation):
      super(Self_Attn,self).__init__()
      self.chanel_in = in_dim
      self.activation = activation
      
      self.query_conv = nn.Conv2d(in_channels = in_dim , out_channels = in_dim//8 , kernel_size= 1)
      self.key_conv = nn.Conv2d(in_channels = in_dim , out_channels = in_dim//8 , kernel_size= 1)
      self.value_conv = nn.Conv2d(in_channels = in_dim , out_channels = in_dim , kernel_size= 1)
      self.gamma = nn.Parameter(torch.zeros(1))
  
      self.softmax  = nn.Softmax(dim=-1) #
  def forward(self,x):
      &quot;&quot;&quot;
          inputs :
              x : input feature maps( B X C X W X H)
          returns :
              out : self attention value + input feature 
              attention: B X N X N (N is Width*Height)
      &quot;&quot;&quot;
      m_batchsize,C,width ,height = x.size()
      proj_query  = self.query_conv(x).view(m_batchsize,-1,width*height).permute(0,2,1) # B X CX(N)
      proj_key =  self.key_conv(x).view(m_batchsize,-1,width*height) # B X C x (*W*H)
      energy =  torch.bmm(proj_query,proj_key) # transpose check
      attention = self.softmax(energy) # BX (N) X (N) 
      proj_value = self.value_conv(x).view(m_batchsize,-1,width*height) # B X C X N
  
      out = torch.bmm(proj_value,attention.permute(0,2,1) )
      out = out.view(m_batchsize,C,width,height)
      
      out = self.gamma*out + x
      return out,attention
  </code></pre>
  <p>The calls of <code>.view(m_batchsize,-1, width*height)</code> in the forward pass are flattening the spatial image dimensions from <code>(W,H)</code> to <code>(N,)</code> where <code>N=H*W</code>, so the 2D array of pixels in the feature map is being converted into a 1D sequence of pixels, then  self-attention is being computed over these sequence elements in the usual manner. So the guess in the question is correct; self-attention is being computed in a pixel-wise manner, and the complaints following the guess are resolved when flattening the feature maps before computing self-attention.</p>
  
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2024-03-23T20:18:42.967Z
Answers.Text: This makes sense by simple analogy to the transformer: Each pixel is like a token, and the feature maps are like the embeddings of the token.  (The similarities between feature-maps for a pixel and token embeddings is strong.)  Attention maps tokens to tokens, so it feels correct for attention to map pixels to pixels.. That this is O(w^2 h^2) does make this approach seem fairly expensive.
Answers.UserDisplayName: Josiah Yoder
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2023-12-30T13:04:10.360Z
Answers.LastActivityDate: 2023-12-30T13:04:10.360Z
Answers.UserDisplayName: Rational Function
Body: <p>I am trying to understand the neural network architecture used by Ho et al. in &quot;Denoising Diffusion Probabilistic Models&quot; (<a href="http://arxiv.org/abs/2006.11239" rel="nofollow noreferrer">paper</a>, <a href="https://github.com/hojonathanho/diffusion/tree/master" rel="nofollow noreferrer">source code</a>). They include self-attention layers in the model, applying them to the feature maps output by previous convolutional layers (ResNet blocks). I understand self-attention in the context of sequential data, but here there is no sequence of vectors, just a single image to be processed by self-attention. I do not understand what the self-attention layer is doing to the feature maps.</p>
  <p><strong>Question: please could you explain the function of the self-attention layers in this CNN?</strong></p>
  <p>My guess is that the self-attention layer is treating pixels like sequence elements i.e. it uses keys/queries to find the pixels in the input feature map that are most relevant to the pixel at some spatial location in the input feature map, then forms a relevance weighted sum of these pixels (or rather their corresponding values) as the output at that location. But this perspective doesn't seem to be consistent with the mechanics of the operation. In sequential data, we have Q, K, V matrices which store embeddings pertaining to the different sequence elements as different rows in the matrices, and it seems that we're replacing these with Q, K, V matrices that are derived from a single image. So when we e.g. calculate <span class="math-container">$QK^\top$</span> we are not computing the dot-product similarity of the queries and keys of different elements=pixels, because a row of <span class="math-container">$Q$</span> and a row of <span class="math-container">$K$</span> does not correspond to one element=pixel in the image (contrast to sequential data where they would correspond to specific sequence elements).</p>
  <p>Note: from reading the paper/source code, I think that the self-attention operation is working according to the diagram below, but please correct me if this isn't the case.
  <a href="https://i.sstatic.net/dCjA4.png" rel="nofollow noreferrer"><img src="https://i.sstatic.net/dCjA4.png" alt="self attention operation on feature maps" /></a>
  (Figure reproduced from this <a href="https://www.mdpi.com/2072-4292/13/16/3316" rel="nofollow noreferrer">paper</a> and a similar set-up can be found in the self-attention GAN <a href="https://arxiv.org/abs/1805.08318" rel="nofollow noreferrer">paper</a>.)</p>
  
Comments.ContentLicense: CC BY-SA 4.0
Comments.CreationDate: 2023-12-30T11:43:40.647Z
Comments.Text: @DeepQZero Thank you very much for welcoming me and for the feedback on my question! I wanted to make sure that I didn't make the problem any more confusing than it needed to be since AI/ML is already fairly confusing intrinsically haha :)
Comments.UserDisplayName: Rational Function
Tags: <convolutional-neural-networks><attention><diffusion-models>
Title: Understanding the function of attention layers in a convolutional neural network (U-Net in a diffusion model)
UserDisplayName: Rational Function
ThreadText: # Understanding the function of attention layers in a convolutional neural network (U-Net in a diffusion model)
  
  (Asked by: username_0 on 2023-12-28)
  
  I am trying to understand the neural network architecture used by Ho et al. in "Denoising Diffusion Probabilistic Models" ([paper](http://arxiv.org/abs/2006.11239), [source code](https://github.com/hojonathanho/diffusion/tree/master)). They include self-attention layers in the model, applying them to the feature maps output by previous convolutional layers (ResNet blocks). I understand self-attention in the context of sequential data, but here there is no sequence of vectors, just a single image to be processed by self-attention. I do not understand what the self-attention layer is doing to the feature maps.
  
  **Question: please could you explain the function of the self-attention layers in this CNN?**
  
  My guess is that the self-attention layer is treating pixels like sequence elements i.e. it uses keys/queries to find the pixels in the input feature map that are most relevant to the pixel at some spatial location in the input feature map, then forms a relevance weighted sum of these pixels (or rather their corresponding values) as the output at that location. But this perspective doesn't seem to be consistent with the mechanics of the operation. In sequential data, we have Q, K, V matrices which store embeddings pertaining to the different sequence elements as different rows in the matrices, and it seems that we're replacing these with Q, K, V matrices that are derived from a single image. So when we e.g. calculate $QK^\top$ we are not computing the dot-product similarity of the queries and keys of different elements=pixels, because a row of $Q$ and a row of $K$ does not correspond to one element=pixel in the image (contrast to sequential data where they would correspond to specific sequence elements).
  
  Note: from reading the paper/source code, I think that the self-attention operation is working according to the diagram below, but please correct me if this isn't the case.
  [![self attention operation on feature maps](https://i.sstatic.net/dCjA4.png)](https://i.sstatic.net/dCjA4.png)
  (Figure reproduced from this [paper](https://www.mdpi.com/2072-4292/13/16/3316) and a similar set-up can be found in the self-attention GAN [paper](https://arxiv.org/abs/1805.08318).)
  
  **Tags:** convolutional-neural-networks, attention, diffusion-models
  > **Comments:**
  > **username_0**: @DeepQZero Thank you very much for welcoming me and for the feedback on my question! I wanted to make sure that I didn't make the problem any more confusing than it needed to be since AI/ML is already fairly confusing intrinsically haha :)  
  
  ---
  
  ## 1 answers
  
  ### Answer by username_0 on 2023-12-30. Score: 2 | Accepted answer
  
  The implementation of self-attention in the source code for the "Self-Attention Generative Adversarial Networks" (SAGAN) paper is somewhat easier to follow than that in the "Denoising Diffusion Probabilistic Models" paper. Assuming that both papers are using the same approach to self-attention on feature maps (which I think they are), the question can be answered by inspecting the definition of self-attention used in SAGAN:
  class Self_Attn(nn.Module):
  """ Self attention Layer"""
  def __init__(self,in_dim,activation):
      super(Self_Attn,self).__init__()
      self.chanel_in = in_dim
      self.activation = activation
      
      self.query_conv = nn.Conv2d(in_channels = in_dim , out_channels = in_dim//8 , kernel_size= 1)
      self.key_conv = nn.Conv2d(in_channels = in_dim , out_channels = in_dim//8 , kernel_size= 1)
      self.value_conv = nn.Conv2d(in_channels = in_dim , out_channels = in_dim , kernel_size= 1)
      self.gamma = nn.Parameter(torch.zeros(1))
  
      self.softmax  = nn.Softmax(dim=-1) #
  def forward(self,x):
      """
          inputs :
              x : input feature maps( B X C X W X H)
          returns :
              out : self attention value + input feature 
              attention: B X N X N (N is Width*Height)
      """
      m_batchsize,C,width ,height = x.size()
      proj_query  = self.query_conv(x).view(m_batchsize,-1,width*height).permute(0,2,1) # B X CX(N)
      proj_key =  self.key_conv(x).view(m_batchsize,-1,width*height) # B X C x (*W*H)
      energy =  torch.bmm(proj_query,proj_key) # transpose check
      attention = self.softmax(energy) # BX (N) X (N) 
      proj_value = self.value_conv(x).view(m_batchsize,-1,width*height) # B X C X N
  
      out = torch.bmm(proj_value,attention.permute(0,2,1) )
      out = out.view(m_batchsize,C,width,height)
      
      out = self.gamma*out + x
      return out,attention
  
  The calls of .view(m_batchsize,-1, width*height) in the forward pass are flattening the spatial image dimensions from (W,H) to (N,) where N=H*W, so the 2D array of pixels in the feature map is being converted into a 1D sequence of pixels, then  self-attention is being computed over these sequence elements in the usual manner. So the guess in the question is correct; self-attention is being computed in a pixel-wise manner, and the complaints following the guess are resolved when flattening the feature maps before computing self-attention.
  
  > **Comments:**
  > **username_1**: This makes sense by simple analogy to the transformer: Each pixel is like a token, and the feature maps are like the embeddings of the token.  (The similarities between feature-maps for a pixel and token embeddings is strong.)  Attention maps tokens to tokens, so it feels correct for attention to map pixels to pixels.. That this is O(w^2 h^2) does make this approach seem fairly expensive.  
  
  ---
  
  