Introduction

What’s Wrong with Seq2Seq Models?

• The Seq2seq model originates from language modeling.
  • Normally has an encoder-decoder architecture:
    • Encoder processes the input sequence and compresses it into a fixed-length context vector, where this representation is expected to be a good summary of the meaning of the whole source sequence.
    • Decoder is initialized with the context vector to emit the transformed output.
    • Both encoder and decoder are typically RNNs (LSTMs/GRUs).
  • Transforms an input sequence (source) to new one (target) where both sequences can be of arbitrary lengths:

    

  • Limitation: incapable to remember long sentences because of the fixed-length context vector.
    • The context vector struggles to retain information from long sentences, often causing early parts of the sequence to be forgotten by the time processing completes.
      ⇒ The attention mechanism was born to resolve this problem.

Attention Cues in Biology

• When inspecting a visual scene, our optic nerve receives information at the order of  bits per second, far exceeding what our brain can fully process. Fortunately, our ancestors had learned from experience (or, data in our case) that not all sensory inputs are created equal.
  ⇒ Throughout human history, the capability of directing attention to only a fraction of information of interest has enabled our brain to allocate resources more smartly to survive, to grow, and to socialize, such as detecting predators, preys, and mates.

• To explain how our attention is deployed in the visual world, a two-component framework has emerged. Idea dates back to William James, who is considered the “father of American psychology” [James, 2007].
  ⇒ In this framework, subjects selectively direct the spotlight of attention using
  1. Nonvolitional cue
  2. Volitional cue
  ⇒ Nonvolitional cue is based on the saliency and conspicuity of objects in the environment:

  

  • While all the paper products are printed in black and white, the coffee cup is red.
    ⇒ This coffee is intrinsically salient and conspicuous in this visual environment.
    ⇒ Automatically and involuntarily draws attention.
    ⇒ You bring the fovea (the center of the macula where visual acuity is highest) onto the coffee.
  ⇒ After drinking coffee, you become caffeinated and want to read a book. You turn your head, refocus your eyes, and look at the book. In this task-dependent case you select the book under cognitive and volitional control:

  

  ⇒ Using the volitional cue based on variable selection criteria, this form of attention is more deliberate. It is also more powerful with the subject’s voluntary effort.

Queries, Keys, and Values

• Inspired by the nonvolitional and volitional attention cues that explain the attentional deployment.

• Consider the case where only nonvolitional cues are available. To bias selection over sensory inputs, we can simply use:
  • Parameterized fully-connected layer.
  • Non-parameterized max or average pooling.
  ⇒ What sets attention mechanisms apart from FC or pooling layers is inclusion of the volitional cues.
  • In the context of attention mechanisms, we refer to volitional cues as queries.
  • Given any query, attention mechanisms bias selection over sensory inputs via attention pooling.
    • These sensory inputs are called values in the context of attention mechanisms.
    • Every value is paired with a key, which can be thought of the nonvolitional cue of that sensory input.
  Design attention pooling such that the given query (volitional cue) can interact with keys (nonvolitional cues), which guides bias selection over values (sensory inputs):

  

Nadaraya-Watson Kernel Regression

• Example of ML algorithm with attention mechanisms.

• We generate a dataset according to the following non-linear function with the noise:
  Given this dataset, how to learn to predict output for any new input ?

Average Pooling

• Begin with the “dumbest” estimator for this regression problem.
  ⇒ Use average pooling to average over all the training outputs:

  

  ⇒ As we can see, this estimator is not so smart as average pooling omits the inputs .

Nonparametric Attention Pooling

• Better idea was proposed by Nadaraya and Watson to weight the outputs according to their input locations:
  where  is a kernel.
  ⇒ We can rewrite it in a more generalized form of attention pooling:
  where  is the query and  is the key-value pair, is an attention weight.
  • is assigned to the corresponding value .
  • Consider a Gaussian kernel defined as
    ⇒ Plugging the Gaussian kernel into the last two equations gives us
  • Key that is closer to the given query will get more attention via a larger attention weight assigned to the key’s corresponding value.
  ⇒ Predicted line is smooth and closer to the ground-truth than that produced by average pooling:

  

  • Now let us take a look at the attention weights. Here testing inputs are queries while training inputs are keys. Since both inputs are sorted, we can see that the closer the query-key pair is, the higher attention weight is in the attention pooling.

    

Parametric Attention Pooling

• We can easily integrate learnable parameters into attention pooling.
  ⇒ Distance between the query and the key can be multiplied by a learnable parameter :
  ⇒ After training the parametric attention model, we can plot its prediction:

  

  • Comparing with nonparametric attention pooling, the region with large attention weights becomes sharper in the learnable and parametric setting:

    

Attention Pooling and Attention Scoring Functions

• In the previous section, we obtained probability distribution over values that are paired with keys: output of the attention pooling is simply a weighted sum of the values based on the attention weights.

• At a high level, we can use the above algorithm to instantiate the framework of attention mechanisms.
  ⇒ Denoting an attention scoring function by , the following figure illustrates how the output of attention pooling can be computed as a weighted sum of values:

  

  • NB: attention weights are a probability distribution, weighted sum is a weighted average.
  • More formally, suppose that we have:
    • Query .
    • Key-value pairs , where any and any .
    ⇒ Attention pooling ( before) is instantiated as a weighted sum of the values:
    where the attention weight (scalar) for the query and key is computed by the softmax operation of an attention scoring function that maps two vectors to a scalar:

• Different choices of the attention scoring function lead to different behaviors of attention pooling.

Additive Attention

• When queries and keys are vectors of different lengths, we can use additive attention as the scoring function. Given a query and a key , the additive attention scoring function is
  where , and are learnable parameters.

  class AdditiveAttention(nn.Module):
  
      def __init__(self, key_size, query_size, num_hiddens, dropout):
          super().__init__()
          self.W_k = nn.Linear(key_size, num_hiddens, bias=False)
          self.W_q = nn.Linear(query_size, num_hiddens, bias=False)
          self.w_v = nn.Linear(num_hiddens, 1, bias=False)
          self.dropout = nn.Dropout(dropout)
  
      def forward(self, queries, keys, values, valid_lens):
          queries, keys = self.W_q(queries), self.W_k(keys)
          # After dimension expansion, shape of `queries`: (`batch_size`, no. of
          # queries, 1, `num_hiddens`) and shape of `keys`: (`batch_size`, 1,
          # no. of key-value pairs, `num_hiddens`). Sum them up with
          # broadcasting
          features = queries.unsqueeze(2) + keys.unsqueeze(1)
          features = torch.tanh(features)
          # There is only one output of `self.w_v`, so we remove the last
          # one-dimensional entry from the shape. Shape of `scores`:
          # (`batch_size`, no. of queries, no. of key-value pairs)
          scores = self.w_v(features).squeeze(-1)
          self.attention_weights = masked_softmax(scores, valid_lens)
          # Shape of `values`: (`batch_size`, no. of key-value pairs, value
          # dimension)
          return torch.bmm(self.dropout(self.attention_weights), values)

Scaled Dot-Product Attention

• More computationally efficient but requires both query and key having the same vector of length .

• Assume that all elements of query and key are i.i.d variables with zero mean and unit variance.
  ⇒ To ensure that the variance of the dot product still remains one regardless of vector length, the scaled dot-product attention scoring function is

• In practice, we often think in mini-batches for efficiency.
  ⇒ Scaled dot-product attention of queries , keys , and values is

  class DotProductAttention(nn.Module):
  
      def __init__(self, dropout):
          super().__init__()
          self.dropout = nn.Dropout(dropout)
  
      # Shape of `queries`: (`batch_size`, no. of queries, `d`)
      # Shape of `keys`: (`batch_size`, no. of key-value pairs, `d`)
      # Shape of `values`: (`batch_size`, no. of key-value pairs, value
      # dimension)
      # Shape of `valid_lens`: (`batch_size`,) or (`batch_size`, no. of queries)
      def forward(self, queries, keys, values, valid_lens=None):
          d = queries.shape[-1]
          # Set `transpose_b=True` to swap the last two dimensions of `keys`
          scores = torch.bmm(queries, keys.transpose(1,2)) / math.sqrt(d)
          self.attention_weights = masked_softmax(scores, valid_lens)
          return torch.bmm(self.dropout(self.attention_weights), values)

Masked Softmax Operation

• In some cases, not all the values should be fed into attention pooling.
  • E.g.: for efficient minibatch processing, some text sequences are padded with special tokens that do not carry meaning. To get an attention pooling over only meaningful tokens as values, we can specify a valid sequence length (in number of tokens) to filter out those beyond this specified range when computing softmax. In this way, we can implement such a masked softmax operation, where any value beyond the valid length is masked as zero.

Basic Attention Mechanisms

Bahdanau Attention

• As we get from the intro, attention mechanism was born to help memorize long source sentences in neural machine translation (NMT).
  ⇒ Rather than building a single context vector out of the encoder’s last hidden state, the secret sauce invented by attention is to create shortcuts between the context vector and the entire source input. The weights of these shortcut connections are customizable for each output element.
  ⇒ When predicting a token, not all the input tokens are relevant.
  ⇒ Model aligns only to parts of the input sequence that are relevant to the current prediction.
  ⇒ This is achieved by treating the context variable at the decoding time step as an output of attention pooling:

  

  where
  • is the number of tokens in the input sequence.
  • Decoder hidden state at time step is the query.
  • Encoder hidden states are both the keys and values.
  • Attention weight is computed using the additive attention scoring function.

• Matrix of alignment scores is a nice byproduct to explicitly show the correlation between source and target words:

  

Multi-Head Attention

• Given the same set of queries, keys, and values, we may want our model to combine knowledge from different behaviors of the same attention mechanism, such as capturing dependencies of various ranges (e.g., shorter-range vs. longer-range) within a sequence.
  ⇒ It may be beneficial to allow our attention mechanism to jointly use different representation subspaces of queries, keys, and values.
  • Usually, understanding the role of a word in a sentence requires understanding how it is related to different parts of the sentence. For example, in some languages, subjects define verb inflection (e.g., gender agreement), verbs define the case of their objects, and many more. In other words, each word is part of many relations.
  ⇒ Let model focus on different things from different representation subspaces at different positions.
  ⇒ Instead of performing a single attention pooling, queries, keys, and values can be transformed with independently learned linear projections. Then these  projected queries, keys, and values are fed into attention pooling in parallel. In the end,  attention pooling outputs (heads) are concatenated and transformed with another learned linear projection to produce the final output.
  ⇒ Given a query , a key , and a value , each attention head is computed as
  where
  • and are learnable parameters.
  • is attention pooling (e.g. additive attention, scaled dot-product attention, etc.).
  ⇒ Multi-head attention output is a linear transformation via of the concatenation of heads:

  

Self-Attention

• We often use CNNs or RNNs to encode a sequence.
  ⇒ Idea: with attention mechanisms in mind, try to feed a sequence of tokens into attention pooling so that the same set of tokens act as queries, keys, and values.
  ⇒ Since queries, keys, and values come from the same place, this performs self-attention.
  ⇒ Given a sequence of input tokens , , its self-attention outputs a sequence of the same length , where
  • E.g. it can be multi-head attention.

CNNs/RNNs vs. Self-Attention

• Let’s compare architectures for mapping a sequence of  tokens to another sequence of equal length, where each input or output token is represented by a -dimensional vector.
  • CNN (consider a convolutional layer whose kernel size is ):
    • Computational complexity of the convolutional layer is .
    • Since CNNs are hierarchical, there are sequential operations and maximum path length is .
  • RNN:
    • When updating the hidden state, multiplication of the weight matrix and the -dimensional hidden state has a computational complexity of . Since the sequence length is , the computational complexity of the recurrent layer is .
    • According to the figure above, there are sequential operations that cannot be parallelized and the maximum path length is also .
  • Self-attention:
    • Queries, keys, and values are all matrices. Consider the scaled dot-product attention, where a matrix is multiplied by a matrix, then the output matrix is multiplied by a matrix. As a result, the self-attention has a computational complexity.
    • As we can see in the figure, each token is directly connected to any other token via self-attention. Therefore, computation can be parallel with sequential operations and the maximum path length is also .

• However, the quadratic computational complexity with respect to a sequence length makes self-attention prohibitively slow for very long sequences.

Positional Encoding

• Self-attention ditches sequential operations in favor of parallel computation.
  ⇒ To use the sequence order information, we can inject absolute or relative positional information.
  ⇒ Add positional encoding to input representations where these encodings can be learned or fixed.

• Fixed positional encoding can be based on sine and cosine functions [Vaswani et al., 2017].
  ⇒ Called sinusoidal positional encoding therefore.
  ⇒ Input representation contains -dimensional embeddings for tokens of a sequence.
  ⇒ Positional encoding outputs using a positional embedding matrix of the same shape, whose element on the -th row and the -th or the -th column is
  • In the positional embedding matrix , rows correspond to positions within a sequence and columns represent different positional encoding dimensions.
    • E.g.: from the graph below, we can see that the -th and the -th columns of the positional embedding matrix have a higher frequency than the -th and the -th columns. The offset between the -th and the -th (same for the -th and the -th) columns is due to the alternation of sine and cosine functions.

      

• Let’s see how monotonically decreased frequency along encoding dimension relates to absolute positional info.
  ⇒ Let’s print out the binary representations of :

  for i in range(8):
      print(f'{i} in binary is {i:>03b}')
  
  0 in binary is 000
  1 in binary is 001
  2 in binary is 010
  3 in binary is 011
  4 in binary is 100
  5 in binary is 101
  6 in binary is 110
  7 in binary is 111

  ⇒ Lowest bit, 2nd-lowest bit, and 3rd-lowest bit alternate on every number, every 2 numbers, and every 4 numbers.
  ⇒ In binary representations, a higher bit has a lower frequency than a lower bit.
  ⇒ Similarly, positional encoding decreases frequencies along encoding dimension by using trigonometric functions:

  

  • Since the outputs are float numbers, such continuous representations are more space-efficient than binary representations.

• Besides capturing absolute positional information, the above positional encoding also allows a model to easily learn to attend by relative positions.
  • This is because for any fixed position offset , the positional encoding at position  can be represented by a linear projection of that at position :
    where and the projection matrix does not depend on any position index .
    ⇒ Any pair of can be linearly projected to for any fixed offset .

“Vanilla” Transformer

• We have compared CNNs, RNNs, and self-attention.
  ⇒ Self-attention enjoys both parallel computation and the shortest maximum path length.
  ⇒ It’s appealing to design deep architectures by using self-attention.
  ⇒ Transformer model is solely based on attention mechanisms.

Model

• Transformer is composed of an encoder and a decoder:

  

  • Input (source) and output (target) sequence embeddings are added with sinusoidal positional encoding before being fed into the encoder and the decoder that stack modules based on self-attention.
  • Encoder is a stack of multiple identical layers, where each layer has 2 sublayers.
    • The first is a multi-head self-attention pooling.
    • The second is a position-wise feed-forward network.
      • By position-wise (or point-wise), it means that it applies the same linear transformation (with the same weights) to each element in the sequence.
      • This can also be viewed as a convolutional layer with filter size of 1.
    • Inspired by the ResNet design, a residual connection is employed around both sublayers.
    • The addition from the residual connection is followed by layer normalization [Ba et al., 2016].
    ⇒ Encoder outputs a -dimensional vector representation for each position of the input sequence.
    ⇒ Generates an attention-based representation with capability to locate a specific piece of information from a large context.
  • Decoder is similar to the encoder, except that the decoder contains two multi-head attention sublayers instead of one in each identical repeating layers.
    • The first multi-head attention sublayer is masked to prevent positions from attending to the future. Thus, each position in decoder is allowed to only attend to all positions in the decoder up to that position.
      • This masked attention preserves the auto-regressive property, ensuring that the prediction only depends on those output tokens that have been generated.
    • The second multi-head attention sublayer is called encoder-decoder attention, where queries are from outputs of previous decoder layer, and keys and values are from transformer encoder outputs.
    ⇒ Retrieves information from the encoded representation.

Positionwise Feed-Forward Network

• Position-wise FFNN transforms representation at all the sequence positions using the same MLP.
  ⇒ This is why it is called position-wise:

  class PositionWiseFFN(nn.Module):
  """Positionwise feed-forward network."""
  
      def __init__(
  			self, 
  			ffn_num_input, 
  			ffn_num_hiddens, 
  			ffn_num_outputs,
  		):
          super().__init__()
          self.layer_1 = nn.Linear(ffn_num_input, ffn_num_hiddens)
          self.relu = nn.ReLU()
          self.layer_2 = nn.Linear(ffn_num_hiddens, ffn_num_outputs)
  
      def forward(self, X):
          return self.layer_2(self.relu(self.layer_1(X)))

  • Since the same MLP transforms at all the positions, when the inputs at all these positions are the same, their outputs are also identical:

    ffn = PositionWiseFFN(4, 4, 8)
    ffn.eval()
    ffn(torch.ones((2, 3, 4)))[0]
    ---
    tensor([[-0.3859, -0.2491, -0.3411, -0.0251,  0.4630, -0.4324, -0.1697,  0.2019],
            [-0.3859, -0.2491, -0.3411, -0.0251,  0.4630, -0.4324, -0.1697,  0.2019],
            [-0.3859, -0.2491, -0.3411, -0.0251,  0.4630, -0.4324, -0.1697,  0.2019]],
           grad_fn=<SelectBackward0>)

Transformers for Vision

• Transformer architecture was initially proposed for Seq2Seq learning, with a focus on machine translation.
  ⇒ Later on, it emerged as the model of choice in various natural language processing tasks.
  ⇒ In the field of computer vision the dominant architecture has remained the CNN.
  ⇒ Researchers started to wonder if it’s possible to do better by adapting Transformer to image data.
  ⇒ Transformers have also become a game-changer in computer vision.

• Vision Transformers (ViTs) extract patches from images and feed them into a Transformer encoder to obtain a global representation, which will finally be transformed into the output label:

  

  • Consider an input image with height , width , channels, and patch height and width both as .
    ⇒ Image is split into a sequence of patches (each patch is flattened to a vector of length ).
    ⇒ Image patches can be treated similarly to tokens in text sequences by Transformer encoders.

Patch Embedding

• Splitting an image into patches and linearly projecting these flattened patches can be simplified as a single convolution operation, where both the kernel size and the stride size are set to the patch size:

  class PatchEmbedding(nn.Module):
      def __init__(self, img_size=96, patch_size=16, num_hiddens=512):
          super().__init__()
          self._img_size = self._make_tuple(img_size)
          self._patch_size = self._make_tuple(patch_size)
          self.num_patches = (
  	        (self._img_size[0] // self._patch_size[0])
  	        * (self._img_size[1] // self._patch_size[1])
  	    )
          self.conv = nn.LazyConv2d(
  	        num_hiddens, 
  	        kernel_size=self._patch_size, 
  	        stride=self._patch_size,
  	    )
  
      def forward(self, X):
          # Output shape: (batch size, no. of patches, no. of channels)
          return self.conv(X).flatten(2).transpose(1, 2)
          
      @staticmethod
      def _make_tuple(x):
          return (x, x) if not isinstance(x, (tuple, list)) else x

ViT’s MLP

• The MLP of vision Transformer encoder is slightly different from the position-wise FFN of original Transformer encoder.
  • Activation function uses the GELU, which is smoother version of the ReLU.
  • Dropout is applied to the output of each fully connected layer in the MLP for regularization.

  class ViTMLP(nn.Module):
      def __init__(
          self, 
          mlp_num_hiddens, 
          mlp_num_outputs, 
          dropout=0.5,
      ):
          super().__init__()
          self.dense1 = nn.LazyLinear(mlp_num_hiddens)
          self.gelu = nn.GELU()
          self.dropout1 = nn.Dropout(dropout)
          self.dense2 = nn.LazyLinear(mlp_num_outputs)
          self.dropout2 = nn.Dropout(dropout)
  
      def forward(self, x):
          return self.dropout2(
              self.dense2(
                  self.dropout1(self.gelu(self.dense1(x)))
              )
          )

ViT’s “Add & Norm” Layer

• Normalization is applied right before the multi-head attention or MLP (not after as in vanilla Transformer).

  class ViTBlock(nn.Module):
      def __init__(
          self,
          num_hiddens,
          norm_shape,
          mlp_num_hiddens,
          num_heads,
          dropout,
          use_bias=False,
      ):
          super().__init__()
          self.ln1 = nn.LayerNorm(norm_shape)
          self.attention = nn.MultiheadAttention(
              embed_dim=num_hiddens,
              num_heads=num_heads,
              dropout=dropout,
              bias=use_bias,
          )
          self.ln2 = nn.LayerNorm(norm_shape)
          self.mlp = ViTMLP(
              mlp_num_hiddens, num_hiddens, dropout
          )
  
      def forward(self, X, attention_mask=None):
          # Apply LayerNorm and Self-Attention with residual connection
          # Transpose for MultiheadAttention
          # (sequence_length, batch_size, embed_dim)
          X_ln1 = self.ln1(X).transpose(0, 1)
  
          if attention_mask is not None:
              # Shape: (batch_size, sequence_length, sequence_length)
              attention_mask = attention_mask.unsqueeze(
                  1
              ).expand(-1, X_ln1.size(0), -1)
  
          attn_output, _ = self.attention(
              X_ln1,
              X_ln1,
              X_ln1,
              attn_mask=attention_mask,
          )
          # Residual connection and transpose back
          X = X + attn_output.transpose(0, 1)
  
          return X + self.mlp(self.ln2(X))

Summary

• Input image is fed into PatchEmbedding whose output is concatenated with <cls> token embedding.
  ⇒ It’s summed with learnable positional embeddings before dropout.
  ⇒ Then output is fed into Transformer encoder that stacks num_blks instances of ViTBlock class.
  ⇒ Finally, the representation of the <cls> token is projected by the network head.

  class ViT(nn.Module):
      """Vision Transformer."""
  
      def __init__(
          self,
          img_size,
          patch_size,
          num_hiddens,
          mlp_num_hiddens,
          num_heads,
          num_blks,
          emb_dropout,
          blk_dropout,
          lr=0.1,
          use_bias=False,
          num_classes=10,
      ):
          super().__init__()
          self.patch_embedding = PatchEmbedding(
              img_size,
              patch_size,
              num_hiddens,
          )
          self.cls_token = nn.Parameter(
              torch.zeros(1, 1, num_hiddens)
          )
          # Add the cls token
          num_steps = self.patch_embedding.num_patches + 1
          # Positional embeddings are learnable
          self.pos_embedding = nn.Parameter(
              torch.randn(1, num_steps, num_hiddens)
          )
          self.dropout = nn.Dropout(emb_dropout)
          self.blks = nn.Sequential(
              *[
                  ViTBlock(
                      num_hiddens,
                      num_hiddens,
                      mlp_num_hiddens,
                      num_heads,
                      blk_dropout,
                      use_bias,
                  )
                  for _ in range(num_blks)
              ]
          )
          self.head = nn.Sequential(
              nn.LayerNorm(num_hiddens),
              nn.Linear(num_hiddens, num_classes),
          )
  
      def forward(self, X):
          X = self.patch_embedding(X)
          X = torch.cat(
              (self.cls_token.expand(X.shape[0], -1, -1), X),
              1,
          )
          X = self.dropout(X + self.pos_embedding)
          X = self.blks(X)
          return self.head(X[:, 0])
  

• For small datasets like ImageNet (1.2M images), vision Transformer does not outperform the ResNet.
  • Transformers lack useful principles in convolution, such as translation invariance and locality.
  ⇒ However, the picture changes when training larger models on larger datasets (e.g., 300M images).
  ⇒ Vision Transformers outperform ResNets by a large margin in image classification.
  ⇒ Superiority of Transformers in scalability.

• It was shown to be effective with data-efficient training strategies of DeiT (Touvron et al., 2021).

• Quadratic complexity of self-attention makes the architecture less suitable for higher-resolution images.
  ⇒ Swin Transformers addressed the quadratic computational complexity with respect to image size.
  ⇒ Reinstated convolution-like priors, extending the applicability of Transformers to a range of computer vision tasks beyond image classification with state-of-the-art results (Liu et al., 2021).

Transformer Optimizations

Sources

• https://d2l.ai/chapter_attention-mechanisms-and-transformers/index.html

• https://lena-voita.github.io/nlp_course/seq2seq_and_attention.html

• https://lilianweng.github.io/posts/2018-06-24-attention/

• https://lilianweng.github.io/posts/2020-04-07-the-transformer-family/

• https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/

• https://huggingface.co/docs/text-generation-inference/conceptual/flash_attention

• https://medium.com/@sachinkalsi/flashattention-understanding-gpu-architecture-part-1-0a8a9a0bb725