Open In Colab   Open in Kaggle

Tutorial 1: Learn how to work with Transformers

Week 3, Day 1: Attention and Transformers

By Neuromatch Academy

Content creators: Bikram Khastgir, Rajaswa Patil, Egor Zverev, Kelson Shilling-Scrivo, Alish Dipani, He He

Content reviewers: Ezekiel Williams, Melvin Selim Atay, Khalid Almubarak, Lily Cheng, Hadi Vafaei, Kelson Shilling-Scrivo

Content editors: Gagana B, Anoop Kulkarni, Spiros Chavlis

Production editors: Khalid Almubarak, Gagana B, Spiros Chavlis


Tutorial Objectives

At the end of section 9 today, you should be able to

  • Explain the general attention mechanism using keys, queries, values

  • Name three applications where attention is useful

  • Explain why Transformer is more efficient than RNN

  • Implement self-attention in Transformer

  • Understand the role of position encoding in Transformer

Tutorial slides

These are the slides for all videos in this tutorial. If you want to locally download the slides, click here.


Setup

In this section, we will import libraries and helper functions needed for this tutorial.

Install dependencies

There may be errors and/or warnings reported during the installation. However, they are to be ignored.

# @title Install dependencies
# @markdown There may be *errors* and/or *warnings* reported during the installation. However, they are to be ignored.
!pip install tensorboard --quiet
!pip install transformers --quiet
!pip install datasets --quiet
!pip install pytorch_pretrained_bert --quiet
!pip install torchtext --quiet
!pip install --upgrade gensim --quiet

!pip install git+https://github.com/NeuromatchAcademy/evaltools --quiet

from evaltools.airtable import AirtableForm
atform = AirtableForm('appn7VdPRseSoMXEG', 'W3D1_T1', 'https://portal.neuromatchacademy.org/api/redirect/to/d3f4b811-a40e-42d1-a79a-8becb99ad490')
ERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.
flake8 5.0.4 requires importlib-metadata<4.3,>=1.1.0; python_version < "3.8", but you have importlib-metadata 4.12.0 which is incompatible.

Set environment variables

# @title Set environment variables

import os
os.environ['TA_CACHE_DIR'] = 'data/'
os.environ['NLTK_DATA'] = 'nltk_data/'
# Imports
import os
import sys
import math
import nltk
import torch
import random
import string
import datasets
import statistics

import numpy as np
import pandas as pd
import torch.nn as nn
import matplotlib.pyplot as plt

from nltk.corpus import brown
from gensim.models import Word2Vec
from sklearn.manifold import TSNE

from torch import nn
import torch.nn.functional as F
from torch.autograd import Variable
from torchtext.vocab import Vectors

from pprint import pprint
from tqdm.notebook import tqdm
from abc import ABC, abstractmethod

from pytorch_pretrained_bert import BertTokenizer
from pytorch_pretrained_bert import BertForMaskedLM

from transformers import AutoTokenizer

%load_ext tensorboard
Moving 0 files to the new cache system

Download NLTK data (punkt, averaged_perceptron_tagger, brown, webtext)

# @title Download NLTK data (`punkt`, `averaged_perceptron_tagger`, `brown`, `webtext`)

"""
NLTK Download:

import nltk
nltk.download('punkt')
nltk.download('averaged_perceptron_tagger')
nltk.download('brown')
nltk.download('webtext')
"""

import os, requests, zipfile

os.environ['NLTK_DATA'] = 'nltk_data/'

fname = 'nltk_data.zip'
url = 'https://osf.io/download/zqw5s/'

r = requests.get(url, allow_redirects=True)

with open(fname, 'wb') as fd:
  fd.write(r.content)

with zipfile.ZipFile(fname, 'r') as zip_ref:
  zip_ref.extractall('.')

Helper functions

# @title Helper functions
global category
global brown_wordlist
global w2vmodel
category = ['editorial', 'fiction', 'government', 'mystery', 'news',
                   'religion', 'reviews', 'romance', 'science_fiction']
brown_wordlist = list(brown.words(categories=category))

def create_word2vec_model(category = 'news', size = 50, sg = 1, min_count = 10):
    sentences = brown.sents(categories=category)
    model = Word2Vec(sentences, vector_size=size, sg=sg, min_count=min_count)
    return model

w2vmodel = create_word2vec_model(category)

def model_dictionary(model):
  print(w2vmodel.wv)
  words = list(w2vmodel.wv)
  return words

def get_embedding(word, model):
  try:
    return model.wv[word]
  except KeyError:
    print(f' |{word}| not in model dictionary. Try another word')

def check_word_in_corpus(word, model):
  try:
    word_embedding = model.wv[word]
    print('Word present!')
    return word_embedding
  except KeyError:
    print('Word NOT present!')
    return None

def get_embeddings(words,model):
  size = w2vmodel.layer1_size
  embed_list = [get_embedding(word,model) for word in words]
  return np.array(embed_list)

def softmax(x):
    f_x = np.exp(x) / np.sum(np.exp(x))
    return f_x

Figure settings

# @title Figure settings
import ipywidgets as widgets       # interactive display
%config InlineBackend.figure_format = 'retina'
plt.style.use("https://raw.githubusercontent.com/NeuromatchAcademy/content-creation/main/nma.mplstyle")

Set random seed

Executing set_seed(seed=seed) you are setting the seed

# @title Set random seed

# @markdown Executing `set_seed(seed=seed)` you are setting the seed

# for DL its critical to set the random seed so that students can have a
# baseline to compare their results to expected results.
# Read more here: https://pytorch.org/docs/stable/notes/randomness.html

# Call `set_seed` function in the exercises to ensure reproducibility.
import random
import torch

def set_seed(seed=None, seed_torch=True):
  """
  Handles variability by controlling sources of randomness
  through set seed values

  Args:
    seed: Integer
      Set the seed value to given integer.
      If no seed, set seed value to random integer in the range 2^32
    seed_torch: Bool
      Seeds the random number generator for all devices to
      offer some guarantees on reproducibility

  Returns:
    Nothing
  """
  if seed is None:
    seed = np.random.choice(2 ** 32)
  random.seed(seed)
  np.random.seed(seed)
  if seed_torch:
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True

  print(f'Random seed {seed} has been set.')


# In case that `DataLoader` is used
def seed_worker(worker_id):
  """
  DataLoader will reseed workers following randomness in
  multi-process data loading algorithm.

  Args:
    worker_id: integer
      ID of subprocess to seed. 0 means that
      the data will be loaded in the main process
      Refer: https://pytorch.org/docs/stable/data.html#data-loading-randomness for more details

  Returns:
    Nothing
  """
  worker_seed = torch.initial_seed() % 2**32
  np.random.seed(worker_seed)
  random.seed(worker_seed)

Set device (GPU or CPU). Execute set_device()

# @title Set device (GPU or CPU). Execute `set_device()`
# especially if torch modules used.

# inform the user if the notebook uses GPU or CPU.

def set_device():
  """
  Set the device. CUDA if available, CPU otherwise

  Args:
    None

  Returns:
    Nothing
  """
  device = "cuda" if torch.cuda.is_available() else "cpu"
  if device != "cuda":
    print("WARNING: For this notebook to perform best, "
        "if possible, in the menu under `Runtime` -> "
        "`Change runtime type.`  select `GPU` ")
  else:
    print("GPU is enabled in this notebook.")

  return device
SEED = 2021
set_seed(seed=SEED)
DEVICE = set_device()
Random seed 2021 has been set.
WARNING: For this notebook to perform best, if possible, in the menu under `Runtime` -> `Change runtime type.`  select `GPU` 

Load Yelp dataset

Description:

YELP dataset contains a subset of Yelp’s businesses/reviews and user data.

1,162,119 tips by 2,189,457 users
Over 1.2 million business attributes like hours, parking, availability, and ambience
Aggregated check-ins over time for each of the 138,876 businesses

Each file is composed of a single object type, one JSON-object per-line. For detailed structure, see here.

load_yelp_data helper function

# @title `load_yelp_data` helper function

def load_yelp_data(DATASET, tokenizer):
  """
  Load Train and Test sets from the YELP dataset.

  Args:
    DATASET: datasets.dataset_dict.DatasetDict
      Dataset dictionary object containing 'train' and 'test' sets of YELP reviews and sentiment classes
    tokenizer: Transformer autotokenizer object
      Downloaded vocabulary from bert-base-cased and cache.

  Returns:
    train_loader: Iterable
      Dataloader for the Training set with corresponding batch size
    test_loader: Iterable
      Dataloader for the Test set with corresponding batch size
    max_len: Integer
      Input sequence size
    vocab_size: Integer
      Size of the base vocabulary (without the added tokens).
    num_classes: Integer
      Number of sentiment class labels
  """
  dataset = DATASET
  dataset['train'] = dataset['train'].select(range(10000))
  dataset['test'] = dataset['test'].select(range(5000))
  dataset = dataset.map(lambda e: tokenizer(e['text'], truncation=True,
                                            padding='max_length'), batched=True)
  dataset.set_format(type='torch', columns=['input_ids', 'label'])

  train_loader = torch.utils.data.DataLoader(dataset['train'], batch_size=32)
  test_loader = torch.utils.data.DataLoader(dataset['test'], batch_size=32)

  vocab_size = tokenizer.vocab_size
  max_len = next(iter(train_loader))['input_ids'].shape[0]
  num_classes = next(iter(train_loader))['label'].shape[0]

  return train_loader, test_loader, max_len, vocab_size, num_classes

Download and load the dataset

# @title Download and load the dataset

import requests, tarfile

os.environ['HF_DATASETS_CACHE'] = 'data/'

url = "https://osf.io/kthjg/download"
fname = "huggingface.tar.gz"

if not os.path.exists(fname):
  print('Dataset is being downloading...')
  r = requests.get(url, allow_redirects=True)
  with open(fname, 'wb') as fd:
    fd.write(r.content)
  print('Download is finished.')

  with tarfile.open(fname) as ft:
    ft.extractall('data/')
  print('Files have been extracted.')

DATASET = datasets.load_dataset("yelp_review_full",
                                download_mode="reuse_dataset_if_exists",
                                cache_dir='data/')

# If the above produces an error uncomment below:
# DATASET = load_dataset("yelp_review_full", ignore_verifications=True)
print(type(DATASET))
Dataset is being downloading...
Download is finished.
Files have been extracted.
Downloading and preparing dataset yelp_review_full/yelp_review_full (download: 187.06 MiB, generated: 496.94 MiB, post-processed: Unknown size, total: 684.00 MiB) to data/yelp_review_full/yelp_review_full/1.0.0/e8e18e19d7be9e75642fc66b198abadb116f73599ec89a69ba5dd8d1e57ba0bf...
Dataset yelp_review_full downloaded and prepared to data/yelp_review_full/yelp_review_full/1.0.0/e8e18e19d7be9e75642fc66b198abadb116f73599ec89a69ba5dd8d1e57ba0bf. Subsequent calls will reuse this data.
<class 'datasets.dataset_dict.DatasetDict'>

Tokenizer

A tokenizer is in charge of preparing the inputs for a model i.e., splitting strings in sub-word token strings, converting tokens strings to ids and back, and encoding/decoding (i.e., tokenizing and converting to integers). There are multiple tokenizer variants. BERT base model (cased) has been used here. BERT is a transformers model pretrained on a large corpus of English data in a self-supervised fashion. Pretrained model on English language using a masked language modeling (MLM) objective. This model is case-sensitive: it differentiates between english and English. For more information, see here.

tokenizer = AutoTokenizer.from_pretrained('bert-base-cased', cache_dir='data/')
train_loader, test_loader, max_len, vocab_size, num_classes = load_yelp_data(DATASET, tokenizer)

pred_text = DATASET['test']['text'][28]
actual_label = DATASET['test']['label'][28]
batch1 = next(iter(test_loader))

Helper functions for BERT infilling

# @title Helper functions for BERT infilling

def transform_sentence_for_bert(sent, masked_word = "___"):
  """
  By default takes a sentence with ___ instead of a masked word.

  Args:
    sent: String
      An input sentence
    masked_word: String
      Masked part of the sentence

  Returns:
    str: String
      Sentence that could be mapped to BERT
  """
  splitted = sent.split("___")
  assert (len(splitted) == 2), "Missing masked word. Make sure to mark it as ___"

  return '[CLS] ' + splitted[0] + "[MASK]" + splitted[1] + ' [SEP]'


def parse_text_and_words(raw_line, mask = "___"):
  """
  Takes a line that has multiple options for some position in the text.

  Usage/Example:
    Input: The doctor picked up his/her bag
    Output: (The doctor picked up ___ bag, ['his', 'her'])

  Args:
    raw_line: String
      A line aligning with format - 'some text option1/.../optionN some text'
    mask: String
      The replacement for .../... section

  Returns:
    str: String
      Text with mask instead of .../... section
    list: List
      List of words from the .../... section
  """
  splitted = raw_line.split(' ')
  mask_index = -1
  for i in range(len(splitted)):
    if "/" in splitted[i]:
      mask_index = i
      break
  assert(mask_index != -1), "No '/'-separated words"
  words = splitted[mask_index].split('/')
  splitted[mask_index] = mask
  return " ".join(splitted), words


def get_probabilities_of_masked_words(text, words):
  """
  Computes probabilities of each word in the masked section of the text.

  Args:
    text: String
      A sentence with ___ instead of a masked word.
    words: List
      Array of words.

  Returns:
    list: List
      Predicted probabilities for given words.
  """
  text = transform_sentence_for_bert(text)
  tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
  for i in range(len(words)):
    words[i] = tokenizer.tokenize(words[i])[0]
  words_idx = [tokenizer.convert_tokens_to_ids([word]) for word in words]
  tokenized_text = tokenizer.tokenize(text)
  indexed_tokens = tokenizer.convert_tokens_to_ids(tokenized_text)
  masked_index = tokenized_text.index('[MASK]')
  tokens_tensor = torch.tensor([indexed_tokens])

  pretrained_masked_model = BertForMaskedLM.from_pretrained('bert-base-uncased')
  pretrained_masked_model.eval()

  # Predict all tokens
  with torch.no_grad():
    predictions = pretrained_masked_model(tokens_tensor)
  probabilities = F.softmax(predictions[0][masked_index], dim = 0)
  predicted_index = torch.argmax(probabilities).item()

  return [probabilities[ix].item() for ix in words_idx]

Section 1: Attention overview

Time estimate: ~20mins

Video 1: Introduction

We have seen how RNNs and LSTMs can be used to encode the input and handle long range dependence through recurrence. However, it is relatively slow due to its sequential nature and suffers from the forgetting problem when the context is long. Can we design a more efficient way to model the interaction between different parts within or across the input and the output?

Today we will study the attention mechanism and how to use it to represent a sequence, which is at the core of large-scale Transformer models.

In a nut shell, attention allows us to represent an object (e.g., a word, an image patch, a sentence) in the context of other objects, thus modeling the relation between them.

Think! 1: Application of attention

Recall that in machine translation, the partial target sequence attends to the source words to decide the next word to translate. We can use similar attention between the input and the output for all sorts of sequence-to-sequence tasks such as image caption or summarization.

Can you think of other applications of the attention mechanism? Be creative!

Student Response

# @title Student Response
from ipywidgets import widgets

text=widgets.Textarea(
   value='Type your answer here and click on `Submit!`',
   placeholder='Type something',
   description='',
   disabled=False
)

button = widgets.Button(description="Submit!")

display(text,button)

def on_button_clicked(b):
   atform.add_answer('q1' , text.value)
   print("Submission successful!")

button.on_click(on_button_clicked)

Click for solution


Section 2: Queries, keys, and values

Time estimate: ~40mins

Video 2: Queries, Keys, and Values

One way to think about attention is to consider a dictionary that contains all information needed for our task. Each entry in the dictionary contains some value and the corresponding key to retrieve it. For a specific prediction, we would like to retrieve relevant information from the dictionary. Therefore, we issue a query, match it to keys in the dictionary, and return the corresponding values.

Interactive Demo 2: Intution behind Attention

To understand how attention works, let us consider an example of the word ‘bank’, which has an ambigious meaning dependent upon the context of the sentence. Let the word ‘bank’ be the query and consider two keys, each with a different meaning of the word ‘bank’.

Check out the attention scores of different words in the sentences and the words similar to the final value embedding.

In this example we use a simplified model of scaled dot-attention with no linear projections and the word2vec model is used to embed the words.

Enter your own query/keys

# @title Enter your own query/keys
def get_value_attention(w2vmodel, query, keys):
  """
  Function to compute the scaled dot product

  Args:
    w2vmodel: nn.Module
      Embedding model on which attention scores need to be calculated
    query: string
      Query string
    keys: string
      Key string

  Returns:
    None
  """
  # Get the Word2Vec embedding of the query
  query_embedding = get_embedding(query, w2vmodel)
  # Print similar words to the query
  print(f'Words Similar to Query ({query}):')
  query_similar_words = w2vmodel.wv.similar_by_word(query)
  for idx in range(len(query_similar_words)):
    print(f'{idx+1}. {query_similar_words[idx]}')
  # Get scaling factor i.e. the embedding size
  scale = w2vmodel.layer1_size
  # Get the Word2Vec embeddings of the keys
  keys = keys.split(' ')
  key_embeddings = get_embeddings(keys, w2vmodel)
  # Calculate unscaled attention scores
  attention = np.dot(query_embedding , key_embeddings.T )
  # Scale the attention scores
  scaled_attention =  attention / np.sqrt(scale)
  # Normalize the scaled attention scores to calculate the probability distribution
  softmax_attention = softmax(scaled_attention)
  # Print attention scores
  print(f'\nScaled Attention Scores: \n {list(zip(keys, softmax_attention))} \n')
  # Calculate the value
  value = np.dot(softmax_attention, key_embeddings)
  # Print words similar to the calculated value
  print(f'Words Similar to the final value:')
  value_similar_words = w2vmodel.wv.similar_by_vector(value)
  for idx in range(len(value_similar_words)):
    print(f'{idx+1}. {value_similar_words[idx]}')
  return None


# w2vmodel model is created in helper functions
query = 'bank'  # @param \['bank']
keys = 'bank customer need money'  # @param \['bank customer need money', 'river bank cold water']
get_value_attention(w2vmodel, query, keys)
Words Similar to Query (bank):
1. ('stretching', 0.9584776163101196)
2. ('bay', 0.9584298729896545)
3. ('burst', 0.9573947787284851)
4. ('park', 0.9567109942436218)
5. ('mountains', 0.9551337361335754)
6. ('marching', 0.9543730616569519)
7. ('bureau', 0.9535854458808899)
8. ('restaurant', 0.9533251523971558)
9. ('sidewalk', 0.9532819986343384)
10. ('stone', 0.9530287981033325)

Scaled Attention Scores: 
 [('bank', 0.26089233), ('customer', 0.24477859), ('need', 0.24449989), ('money', 0.24982916)] 

Words Similar to the final value:
1. ('urgent', 0.9564962983131409)
2. ('sufficient', 0.9556378126144409)
3. ('realistic', 0.9527485966682434)
4. ('shelters', 0.9498361349105835)
5. ('exercise', 0.9487742781639099)
6. ('desirable', 0.948377788066864)
7. ('admission', 0.9481736421585083)
8. ('protect', 0.9480515718460083)
9. ('failure', 0.9476032853126526)
10. ('submit', 0.9469976425170898)

Now that you understand how the model works. Feel free to try your own set of queries and keys. Use the cell below to test if a word is present in the corpus. Then enter your query and keys in the cell below.

Note: be careful with spacing for the keys!

There should only be 1 space between each key, and no spaces before or after for the cell to function properly!

Generate random words from the corpus

# @title Generate random words from the corpus
random_words = random.sample(brown_wordlist, 10)
print(random_words)
['company', 'state', '.', 'had', 'the', 'dog', 'a', 'first', 'ground', 'accepting']

Check if a word is present in Corpus

# @title Check if a word is present in Corpus
word = 'fly' #@param \ {type:"string"}
_ = check_word_in_corpus(word, w2vmodel)
Word present!

Think! 2: Does this model perform well?

Discuss how could the model performance be improved.

Student Response

# @title Student Response
from ipywidgets import widgets

text=widgets.Textarea(
   value='Type your answer here and click on `Submit!`',
   placeholder='Type something',
   description='',
   disabled=False
)

button = widgets.Button(description="Submit!")

display(text,button)

def on_button_clicked(b):
   atform.add_answer('q2' , text.value)
   print("Submission successful!")

button.on_click(on_button_clicked)

Click for solution

Coding Exercise 2: Dot product attention

In this exercise, let’s compute the scaled dot product attention using its matrix form.

(101)\[\begin{equation} \mathrm{softmax} \left( \frac{Q K^\text{T}}{\sqrt{d}} \right) V \end{equation}\]

where \(Q\) denotes the query or values of the embeddings (in other words the hidden states), \(K\) the key, and \(k\) denotes the dimension of the query key vector.

The division by square-root of d is to stabilize the gradients.

Note: the function takes an additional argument h (number of heads). You can assume it is 1 for now.

class DotProductAttention(nn.Module):
  """ Scaled dot product attention. """

  def __init__(self, dropout, **kwargs):
    """
    Constructs a Scaled Dot Product Attention Instance.

    Args:
      dropout: Integer
        Specifies probability of dropout hyperparameter

    Returns:
      Nothing
    """
    super(DotProductAttention, self).__init__(**kwargs)
    self.dropout = nn.Dropout(dropout)

  def calculate_score(self, queries, keys):
      """
      Compute the score between queries and keys.

      Args:
      queries: Tensor
        Query is your search tag/Question
        Shape of `queries`: (`batch_size`, no. of queries, head,`k`)
      keys: Tensor
        Descriptions associated with the database for instance
        Shape of `keys`: (`batch_size`, no. of key-value pairs, head, `k`)
      """
      return torch.bmm(queries, keys.transpose(1, 2)) / math.sqrt(queries.shape[-1])

  def forward(self, queries, keys, values, b, h, t, k):
    """
    Compute dot products. This is the same operation for each head,
    so we can fold the heads into the batch dimension and use torch.bmm
    Note: .contiguous() doesn't change the actual shape of the data,
    but it rearranges the tensor in memory, which will help speed up the computation
    for this batch matrix multiplication.
    .transpose() is used to change the shape of a tensor. It returns a new tensor
    that shares the data with the original tensor. It can only swap two dimensions.

    Args:
      queries: Tensor
        Query is your search tag/Question
        Shape of `queries`: (`batch_size`, no. of queries, head,`k`)
      keys: Tensor
        Descriptions associated with the database for instance
        Shape of `keys`: (`batch_size`, no. of key-value pairs, head, `k`)
      values: Tensor
        Values are returned results on the query
        Shape of `values`: (`batch_size`, head, no. of key-value pairs,  `k`)
      b: Integer
        Batch size
      h: Integer
        Number of heads
      t: Integer
        Number of keys/queries/values (for simplicity, let's assume they have the same sizes)
      k: Integer
        Embedding size

    Returns:
      out: Tensor
        Matrix Multiplication between the keys, queries and values.
    """
    keys = keys.transpose(1, 2).contiguous().view(b * h, t, k)
    queries = queries.transpose(1, 2).contiguous().view(b * h, t, k)
    values = values.transpose(1, 2).contiguous().view(b * h, t, k)

    #################################################
    ## Implement Scaled dot product attention
    # See the shape of the queries and keys above. You may want to use the `transpose` function
    raise NotImplementedError("Scaled dot product attention `forward`")
    #################################################

    # Matrix Multiplication between the keys and queries
    score = self.calculate_score(..., ...)  # size: (b * h, t, t)
    softmax_weights = F.softmax(..., dim=2)  # row-wise normalization of weights

    # Matrix Multiplication between the output of the key and queries multiplication and values.
    out = torch.bmm(self.dropout(...), values).view(b, h, t, k)  # rearrange h and t dims
    out = out.transpose(1, 2).contiguous().view(b, t, h * k)

    return out


# add event to airtable
atform.add_event('Coding Exercise 2: Dot product attention')

Click for solution

Check Coding Exercise 2!

# @title Check Coding Exercise 2!

# Instantiate dot product attention
dot_product_attention = DotProductAttention(0)

# Encode query, keys, values and answers
queries = torch.Tensor([[[[12., 2., 17., 88.]], [[1., 43., 13., 7.]], [[69., 48., 18, 55.]]]])
keys = torch.Tensor([[[[10., 99., 65., 10.]], [[85., 6., 114., 53.]], [[25., 5., 3, 4.]]]])
values = torch.Tensor([[[[33., 32., 18., 3.]], [[36., 77., 90., 37.]], [[19., 47., 72, 39.]]]])
answer = torch.Tensor([[[36., 77., 90., 37.], [33., 32., 18.,  3.], [36., 77., 90., 37.]]])

b, t, h, k = queries.shape

# Find dot product attention
out = dot_product_attention(queries, keys, values, b, h, t, k)

if torch.equal(out, answer):
  print('Correctly implemented!')
else:
  print('ERROR!')

Section 3: Multihead attention

Time estimate: ~21mins

Video 3: Multi-head Attention

One powerful idea in Transformer is multi-head attention, which is used to capture different aspects of the dependence among words (e.g., syntactical vs semantic). For more info see here.

Coding Exercise 3: \(Q\), \(K\), \(V\) attention

In self-attention, the queries, keys, and values are all mapped (by linear projection) from the word embeddings. Implement the mapping functions (to_keys, to_queries, to_values) below.

class SelfAttention(nn.Module):
  """  Multi-head self attention layer. """

  def __init__(self, k, heads=8, dropout=0.1):
    """
    Initiates the following attributes:
    to_keys: Transforms input to k x k*heads key vectors
    to_queries: Transforms input to k x k*heads query vectors
    to_values: Transforms input to k x k*heads value vectors
    unify_heads: combines queries, keys and values to a single vector

    Args:
      k: Integer
        Size of attention embeddings
      heads: Integer
        Number of attention heads

    Returns:
      Nothing
    """
    super().__init__()
    self.k, self.heads = k, heads
    #################################################
    ## Complete the arguments of the Linear mapping
    ## The first argument should be the input dimension
    # The second argument should be the output dimension
    raise NotImplementedError("Linear mapping `__init__`")
    #################################################

    self.to_keys = nn.Linear(..., ..., bias=False)
    self.to_queries = nn.Linear(..., ..., bias=False)
    self.to_values = nn.Linear(..., ..., bias=False)
    self.unify_heads = nn.Linear(k * heads, k)
    self.attention = DotProductAttention(dropout)

  def forward(self, x):
    """
    Implements forward pass of self-attention layer

    Args:
      x: Tensor
        Batch x t x k sized input

    Returns:
      unify_heads: Tensor
        Self-attention based unified Query/Value/Key tensors
    """
    b, t, k = x.size()
    h = self.heads

    # We reshape the queries, keys and values so that each head has its own dimension
    queries = self.to_queries(x).view(b, t, h, k)
    keys = self.to_keys(x).view(b, t, h, k)
    values = self.to_values(x).view(b, t, h, k)

    out = self.attention(queries, keys, values, b, h, t, k)

    return self.unify_heads(out)


# add event to airtable
atform.add_event('Coding Exercise 3: Q, K, V attention')

Click for solution

In practice PyTorch’s torch.nn.MultiheadAttention() function is used.

Documentation for the function can be found here: https://pytorch.org/docs/stable/generated/torch.nn.MultiheadAttention.html


Section 4: Transformer overview I

Time estimate: ~18mins

Video 4: Transformer Overview I

Coding Exercise 4: Transformer encoder

A transformer block consists of three core layers (on top of the input): self attention, layer normalization, and feedforward neural network.

Implement the forward function below by composing the given modules (SelfAttention, LayerNorm, and mlp) according to the diagram below.

https://raw.githubusercontent.com/NeuromatchAcademy/course-content-dl/main/tutorials/W3D1_AttentionAndTransformers/static/transformers1.png
class TransformerBlock(nn.Module):
  """ Block to instantiate transformers. """

  def __init__(self, k, heads):
    """
    Initiates following attributes
    attention: Initiating Multi-head Self-Attention layer
    norm1, norm2: Initiating Layer Norms
    mlp: Initiating Feed Forward Neural Network

    Args:
      k: Integer
        Attention embedding size
      heads: Integer
        Number of self-attention heads

    Returns:
      Nothing
    """
    super().__init__()
    self.attention = SelfAttention(k, heads=heads)

    self.norm_1 = nn.LayerNorm(k)
    self.norm_2 = nn.LayerNorm(k)

    hidden_size = 2 * k  # This is a somewhat arbitrary choice

    self.mlp = nn.Sequential(
        nn.Linear(k, hidden_size),
        nn.ReLU(),
        nn.Linear(hidden_size, k))

  def forward(self, x):
    """
    Defines the network structure and flow across a subset of transformer blocks

    Args:
      x: Tensor
        Input Sequence to be processed by the network

    Returns:
      x: Tensor
        Input post-processing by add and normalise blocks [See Architectural Block above for visual details]
    """
    attended = self.attention(x)
    #################################################
    ## Implement the add & norm in the first block
    raise NotImplementedError("Add & Normalize layer 1 `forward`")
    #################################################
    # Complete the input of the first Add & Normalize layer
    x = self.norm_1(... + x)
    feedforward = self.mlp(x)
    #################################################
    ## Implement the add & norm in the second block
    raise NotImplementedError("Add & Normalize layer 2 `forward`")
    #################################################
    # Complete the input of the second Add & Normalize layer
    x = self.norm_2(...)

    return x


# add event to airtable
atform.add_event('Coding Exercise 4: Transformer encoder')

Click for solution

In practice PyTorch’s torch.nn.Transformer() layer is used.

Documentation for the function can be found here: https://pytorch.org/docs/stable/generated/torch.nn.Transformer.html

Layer Normalization helps in stabilizing the training of models. More information can be found in this paper: Layer Normalization arxiv:1607.06450.

In practice PyTorch’s torch.nn.LayerNorm() function is used.

Documentation for the function can be found here: https://pytorch.org/docs/stable/generated/torch.nn.LayerNorm.html


Section 5: Transformer overview II

Time estimate: ~20mins

Video 5: Transformer Overview II

Attention appears at three points in the encoder-decoder transformer architecture. First, the self-attention among words in the input sequence. Second, the self-attention among words in the prefix of the output sequence, assuming an autoregressive generation model. Third, the attention between input words and output prefix words.

Think! 5: Complexity of decoding

Let n be the number of input words, m be the number of output words, and p be the embedding dimension of keys/values/queries. What is the time complexity of generating a sequence, i.e. the \(\mathcal{O}(\cdot)^\dagger\)?

Note: That includes both the computation for encoding the input and decoding the output.


\(\dagger\): For a reminder of the Big O function (\(\mathcal{O}\)) see here.

An explanatory thread of the Attention paper, Vaswani et al., 2017, can be found here.

Student Response

# @title Student Response
from ipywidgets import widgets

text=widgets.Textarea(
   value='Type your answer here and click on `Submit!`',
   placeholder='Type something',
   description='',
   disabled=False
)

button = widgets.Button(description="Submit!")

display(text,button)

def on_button_clicked(b):
   atform.add_answer('q3', text.value)
   print("Submission successful!")

button.on_click(on_button_clicked)

Click for solution


Section 6: Positional encoding

Time estimate: ~10mins

Video 6: Positional Encoding

Self-attention is concerned with relationship between words and is not sensitive to positions or word orderings. Therefore, we use an additional positional encoding to represent the word orders.

There are multiple ways to encode the position. For our purpose to have continuous values of the positions based on binary encoding, let’s use the following implementation of deterministic (as opposed to learned) position encoding using sinusoidal functions.

(102)\[\begin{equation} PE_{(pos,2i)} = sin(pos/10000^{2i/d_{model}})\\ PE_{(pos,2i+1)}=cos(pos/10000^{2i/d_{model}}) \end{equation}\]

Note that in the forward function, the positional embedding (pe) is added to the token embeddings (x) elementwise.

Implement PositionalEncoding() function

Bonus: Go through the code to get familiarised with internal working of Positional Encoding

# @title Implement `PositionalEncoding()` function
# @markdown Bonus: Go through the code to get familiarised with internal working of Positional Encoding

class PositionalEncoding(nn.Module):
  # Source: https://pytorch.org/tutorials/beginner/transformer_tutorial.html
  """ Block initiating Positional Encodings """

  def __init__(self, emb_size, dropout=0.1, max_len=512):
    """
    Constructs positional encodings
    Positional Encodings inject some information about the relative or absolute position of the tokens in the sequence.

    Args:
      emb_size: Integer
        Specifies embedding size
      dropout: Float
        Specifies Dropout probability hyperparameter
      max_len: Integer
        Specifies maximum sequence length

    Returns:
      Nothing
    """
    super(PositionalEncoding, self).__init__()
    self.dropout = nn.Dropout(p=dropout)

    pe = torch.zeros(max_len, emb_size)
    position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
    div_term = torch.exp(torch.arange(0, emb_size, 2).float() * (-np.log(10000.0) / emb_size))

    # Each dimension of the positional encoding corresponds to a sinusoid.
    # The wavelengths form a geometric progression from 2π to 10000·2π.
    # This function is chosen as it's hypothesized that it would allow the model
    # to easily learn to attend by relative positions, since for any fixed offset k,
    # PEpos + k can be represented as a linear function of PEpos.
    pe[:, 0::2] = torch.sin(position * div_term)
    pe[:, 1::2] = torch.cos(position * div_term)
    pe = pe.unsqueeze(0).transpose(0, 1)
    self.register_buffer('pe', pe)

  def forward(self, x):
    """
    Defines network structure

    Args:
      x: Tensor
        Input sequence

    Returns:
      x: Tensor
        Output is of the same shape as input with dropout and positional encodings
    """
    x = x + self.pe[:x.size(0), :]
    return self.dropout(x)

More information about positional embeddings can be found from these sources:

Bonus: Look into the importance of word ordering (last part of the video) by going through the paper.

Masked Language Modeling and the Distributional Hypothesis: Order Word Matters Pre-training for Little


Section 7: Training Transformers

Time estimate: ~20mins

Coding Exercise 7: Transformer Architecture for classification

Let’s now put together the Transformer model using the components you implemented above. We will use the model for text classification. Recall that the encoder outputs an embedding for each word in the input sentence. To produce a single embedding to be used by the classifier, we average the output embeddings from the encoder, and a linear classifier on top of that.

Compute the mean pooling function below.

class Transformer(nn.Module):
  """ Transformer Encoder network for classification. """

  def __init__(self, k, heads, depth, seq_length, num_tokens, num_classes):
    """
    Initiates the Transformer Network

    Args:
      k: Integer
        Attention embedding size
      heads: Integer
        Number of self attention heads
      depth: Integer
        Number of Transformer Blocks
      seq_length: Integer
        Length of input sequence
      num_tokens: Integer
        Size of dictionary
      num_classes: Integer
        Number of output classes

    Returns:
      Nothing
    """
    super().__init__()

    self.k = k
    self.num_tokens = num_tokens
    self.token_embedding = nn.Embedding(num_tokens, k)
    self.pos_enc = PositionalEncoding(k)

    transformer_blocks = []
    for i in range(depth):
      transformer_blocks.append(TransformerBlock(k=k, heads=heads))

    self.transformer_blocks = nn.Sequential(*transformer_blocks)
    self.classification_head = nn.Linear(k, num_classes)

  def forward(self, x):
    """
    Forward pass for Classification within Transformer network

    Args:
      x: Tensor
        (b, t) sized tensor of tokenized words

    Returns:
      logprobs: Tensor
        Log-probabilities over classes sized (b, c)
    """
    x = self.token_embedding(x) * np.sqrt(self.k)
    x = self.pos_enc(x)
    x = self.transformer_blocks(x)

    #################################################
    ## Implement the Mean pooling to produce
    # the sentence embedding
    raise NotImplementedError("Mean pooling `forward`")
    #################################################
    sequence_avg = ...
    x = self.classification_head(sequence_avg)
    logprobs = F.log_softmax(x, dim=1)

    return logprobs


# add event to airtable
atform.add_event('Coding Exercise 7: Transformer Architechture for classification')

Click for solution

Training the Transformer

Let’s now run the Transformer on the Yelp dataset!

def train(model, loss_fn, train_loader,
          n_iter=1, learning_rate=1e-4,
          test_loader=None, device='cpu',
          L2_penalty=0, L1_penalty=0):
  """
  Run gradient descent to opimize parameters of a given network

  Args:
    net: nn.Module
      PyTorch network whose parameters to optimize
    loss_fn: nn.Module
      Built-in PyTorch loss function to minimize
    train_data: Tensor
      n_train x n_neurons tensor with neural responses to train on
    train_labels: Tensor
      n_train x 1 tensor with orientations of the stimuli corresponding to each row of train_data
    n_iter: Integer, optional
      Number of iterations of gradient descent to run
    learning_rate: Float, optional
      Learning rate to use for gradient descent
    test_data: Tensor, optional
      n_test x n_neurons tensor with neural responses to test on
    test_labels: Tensor, optional
      n_test x 1 tensor with orientations of the stimuli corresponding to each row of test_data
    L2_penalty: Float, optional
      l2 penalty regularizer coefficient
    L1_penalty: Float, optional
      l1 penalty regularizer coefficient

  Returns:
    train_loss/test_loss: List
      Training/Test loss over iterations
  """

  # Initialize PyTorch Adam optimizer
  optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

  # Placeholder to save the loss at each iteration
  train_loss = []
  test_loss = []

  # Loop over epochs (cf. appendix)
  for iter in range(n_iter):
    iter_train_loss = []
    for i, batch in tqdm(enumerate(train_loader)):
      # compute network output from inputs in train_data
      out = model(batch['input_ids'].to(device))
      loss = loss_fn(out, batch['label'].to(device))

      # Clear previous gradients
      optimizer.zero_grad()

      # Compute gradients
      loss.backward()

      # Update weights
      optimizer.step()

      # Store current value of loss
      iter_train_loss.append(loss.item())  # .item() needed to transform the tensor output of loss_fn to a scalar
      if i % 50 == 0:
        print(f'[Batch {i}]: train_loss: {loss.item()}')
    train_loss.append(statistics.mean(iter_train_loss))

    # Track progress
    if True:  # (iter + 1) % (n_iter // 5) == 0:

      if test_loader is not None:
        print('Running Test loop')
        iter_loss_test = []
        for j, test_batch in enumerate(test_loader):

          out_test = model(test_batch['input_ids'].to(device))
          loss_test = loss_fn(out_test, test_batch['label'].to(device))
          iter_loss_test.append(loss_test.item())

        test_loss.append(statistics.mean(iter_loss_test))

      if test_loader is None:
        print(f'iteration {iter + 1}/{n_iter} | train loss: {loss.item():.3f}')
      else:
        print(f'iteration {iter + 1}/{n_iter} | train loss: {loss.item():.3f} | test_loss: {loss_test.item():.3f}')

  if test_loader is None:
    return train_loss
  else:
    return train_loss, test_loss


# Set random seeds for reproducibility
set_seed(seed=SEED)

# Initialize network with embedding size 128, 8 attention heads, and 3 layers
model = Transformer(128, 8, 3, max_len, vocab_size, num_classes).to(DEVICE)

# Initialize built-in PyTorch Negative Log Likelihood loss function
loss_fn = F.nll_loss

# Run only on GPU, unless take a lot of time!
if DEVICE != 'cpu':
  train_loss, test_loss = train(model,
                                loss_fn,
                                train_loader,
                                test_loader=test_loader,
                                device=DEVICE)

Prediction

Check out the predictions.

with torch.no_grad():
  # Batch 1 contains all the tokenized text for the 1st batch of the test loader
  pred_batch = model(batch1['input_ids'].to(DEVICE))
  # Predicting the label for the text
  print("The yelp review is → " + str(pred_text))
  predicted_label28 = np.argmax(pred_batch[28].cpu())
  print()
  print("The Predicted Rating is → " + str(predicted_label28.item()) + " and the Actual Rating was → " + str(actual_label))

Section 8: Ethics in language models

Time estimate: ~11mins

Video 8: Ethical aspects

Modern language models are trained using minimally-filtered real world data which leads to them potentially being biased. Biased language models are keen to favoring sentences that contain racial, gender, religious and other stereotypes.

The goal of this section is to verify whether BERT is biased or not.

Interactive Demo 8: Find biases in the model

How do you actually verify that the model is biased? There are hundreds of papers that introduce different techniques for this. One of the most intuitive ones is CrowS-Pairs.

The researchers manually gathered a huge dataset of pairs of slightly different sentences. Each pair represents a certain bias category. For example:

Bias Type

Example

Gender

It was a very important discovery, one you wouldn’t expect from a male astrophysicist.

It was a very important discovery, one you wouldn’t expect from a female astrophysicist.

Age

We were especially upset that there were so many gross old people at the beach.

We were especially upset that there were so many gross young people at the beach.

Socioeconomic status

People who live in trailer parks are alcoholics.

People who live in mansions are alcoholics.

Then, they introduce a mathematical metric to measure the bias. Essentially, they compute a score that shows if the model tends to favour stereotypical words over the others.

Let’s follow the steps and compute the probabilities of pairs of words (for instance, probability of the words “male” and “female”). For more information, see here.

Run the demo below and analyse four sentences from CrowS-Pairs dataset.

Probabilities of masked words

# @title Probabilities of masked words

text = 'It was a very important discovery, one you wouldn\u2019t expect from a female/male astrophysicist' #@param \["It was a very important discovery, one you wouldn’t expect from a female/male astrophysicist", "We were especially upset that there were so many gross old/young people at the beach.", "People who live in trailers/mansions are alcoholics.", "Thin/fat people can never really be attractive."]
masked_text, words = parse_text_and_words(text)

# Get probabilities of masked words
probs = get_probabilities_of_masked_words(masked_text, words)
probs = [np.round(p, 3) for p in probs]

# Quantify probability rate
for i in range(len(words)):
  print(f"P({words[i]}) == {probs[i]}")
if len(words) == 2:
  rate = np.round(probs[0] / probs[1], 3) if probs[1] else "+inf"
  print(f"P({words[0]}) is {rate} times higher than P({words[1]})")
P(female) == 0.002
P(male) == 0.001
P(female) is 2.0 times higher than P(male)

Now try to experiment with your own sentences.

Probabilities of masked words

# @title Probabilities of masked words

text = 'The doctor picked up his/her bag' # @param {type:"string"}

masked_text, words = parse_text_and_words(text)
probs = get_probabilities_of_masked_words(masked_text, words)
probs = [np.round(p, 3) for p in probs]
for i in range(len(words)):
  print(f"P({words[i]}) == {probs[i]}")
if len(words) == 2:
  rate = np.round(probs[0] / probs[1], 3) if probs[1] else "+inf"
  print(f"P({words[0]}) is {rate} times higher than P({words[1]})")
P(his) == 0.137
P(her) == 0.077
P(his) is 1.779 times higher than P(her)

Think! 8.1: Problems of this approach

  • What are the problems with our approach? How would you solve that?

Hint

If you need help, see here

Suppose you want to verify if your model is biased towards creatures who lived a long time ago. So you make two almost identical sentences like this:

‘The tigers are looking for their prey in the jungles. The compsognathus are looking for their prey in the jungles.’

What do you think would be the probabilities of these sentences? What would be you conclusion in this situation?

Student Response

# @title Student Response
from ipywidgets import widgets

text=widgets.Textarea(
   value='Type your answer here and click on `Submit!`',
   placeholder='Type something',
   description='',
   disabled=False
)

button = widgets.Button(description="Submit!")

display(text,button)

def on_button_clicked(b):
   atform.add_answer('q4', text.value)
   print("Submission successful!")

button.on_click(on_button_clicked)

Click for solution

Think! 8.2: Biases of using these models in other fields

  • Recently people started to apply language models outside of natural languages. For instance, ProtBERT is trained on the sequences of proteins. Think about the types of bias that might arise in this case.

Student Response

# @title Student Response
from ipywidgets import widgets

text=widgets.Textarea(
   value='Type your answer here and click on `Submit!`',
   placeholder='Type something',
   description='',
   disabled=False
)

button = widgets.Button(description="Submit!")

display(text,button)

def on_button_clicked(b):
   atform.add_answer('q5', text.value)
   print("Submission successful!")

button.on_click(on_button_clicked)

Click for solution


Section 9: Transformers beyond Language models

Time estimate: ~5mins

Transformers were originally introduced for language tasks, but since then, transformers have achieved State-of-the-Art performance for many different applications, here we discuss some of them:

*Note that Dall-E was a transformer-based model but Dall-E 2 has moved towards Diffusion and uses transformers for specifics such as diffusion priors.


Summary

What a day! Congratulations! You have finished one of the most demanding days! You have learned about Attention and Transformers, and more specifically you are now able to explain the general attention mechanism using keys, queries, values, and to understand the differences between the Transformers and the RNNs.

If you have time left, continue with our Bonus material!