class_docstr: Compositional Attention, as proposed in
  "Compositional Attention: Disentangling search and retrieval"_, S. Mittal et al.
  
  A key insight from this proposal is that the attention mechanism can be conceived as two steps:
  a search and a retrieval operation. When queried, the model can search for the most relevant information
  (Softmax(QKt)), then retrieve information given the Value.
  
  Contrary to the original attention proposal, which does not consider interactions in between heads,
  the compositional attention will consider all possible interactions and softmax over that dimension,
  so that the information retrieved covers the most relevant dimensions. The number of heads and rules to
  use is thus typically smaller than for a comparable traditional Transformer, and asking for the same number of heads
  may not fit in memory.
  
  Args:
      dim_model: dimension of the incoming latent space
      num_heads: number of heads *for the search operation*
      dim_attn: dimension (embedding) of the attention
      num_rules: number of rules to consider *for the retrieval operation*
      dim_selection: dimension of the scoring/selection space for the retrievals
      dim_key, dim_value: dimensions of K and V, if different from Q
      dropout: attention dropout probability
      qk_rule: QK product will drive the retrieval process
      nonlinear: use a non linear method to score the retrievals
      bias: use bias in the initial projection step
      causal: causal computations (attend to the past only)
  
  _"Compositional Attention: Disentangling search and retrieval": https://arxiv.org/pdf/2110.09419v1.pdf
class_docstr_tok: ["Compositional","Attention",",","as","proposed","in","''","Compositional","Attention",":","Disentangling","search","and","retrieval","''","_",",","S.","Mittal","et","al",".","A","key","insight","from","this","proposal","is","that","the","attention","mechanism","can","be","conceived","as","two","steps",":","a","search","and","a","retrieval","operation",".","When","queried",",","the","model","can","search","for","the","most","relevant","information","(","Softmax","(","QKt",")",")",",","then","retrieve","information","given","the","Value",".","Contrary","to","the","original","attention","proposal",",","which","does","not","consider","interactions","in","between","heads",",","the","compositional","attention","will","consider","all","possible","interactions","and","softmax","over","that","dimension",",","so","that","the","information","retrieved","covers","the","most","relevant","dimensions",".","The","number","of","heads","and","rules","to","use","is","thus","typically","smaller","than","for","a","comparable","traditional","Transformer",",","and","asking","for","the","same","number","of","heads","may","not","fit","in","memory",".","Args",":","dim_model",":","dimension","of","the","incoming","latent","space","num_heads",":","number","of","heads","*","for","the","search","operation","*","dim_attn",":","dimension","(","embedding",")","of","the","attention","num_rules",":","number","of","rules","to","consider","*","for","the","retrieval","operation","*","dim_selection",":","dimension","of","the","scoring\/selection","space","for","the","retrievals","dim_key",",","dim_value",":","dimensions","of","K","and","V",",","if","different","from","Q","dropout",":","attention","dropout","probability","qk_rule",":","QK","product","will","drive","the","retrieval","process","nonlinear",":","use","a","non","linear","method","to","score","the","retrievals","bias",":","use","bias","in","the","initial","projection","step","causal",":","causal","computations","(","attend","to","the","past","only",")","_","''","Compositional","Attention",":","Disentangling","search","and","retrieval","''",":","https",":","\/\/arxiv.org\/pdf\/2110.09419v1.pdf"]
code: def _reset_parameters(self):
      # NOTE: in_proj_container is already initialized
  
      if self.qk_rule:
          nn.init.xavier_uniform_(
              self.value_k.weight, gain=1 / math.sqrt(2)
          )
          nn.init.xavier_uniform_(
              self.value_q.weight, gain=1 / math.sqrt(2)
          )
      else:
          nn.init.xavier_uniform_(self.value_q.weight)
          if self.nonlinear:
              nn.init.xavier_uniform_(self.score_network[0].weight)
              nn.init.xavier_uniform_(self.score_network[2].weight)
          else:
              nn.init.xavier_uniform_(self.score_network.weight)
  
      nn.init.xavier_uniform_(self.out_proj.weight)
      if self.out_proj.bias is not None:
          nn.init.constant_(self.out_proj.bias, 0.0)
  
file_path: xformers/xformers/components/attention/compositional.py
imports: import math
  from dataclasses import dataclass
  from typing import Optional
  import torch
  import torch.nn.functional
  from torch import Tensor, nn
  from xformers.components.attention import Attention, AttentionConfig, AttentionMask, register_attention
  from xformers.components.attention.core import _softmax
  from xformers.components.input_projection import InputProjection, InputProjectionConfig
label_desc: Use image node_id 2 for calling the CompositionalAttention obj's underlying member method code with example usage: obj._reset_parameters() without return types