wehub/labmlai--annotated_deep_learning_paper_implementations
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:19:01 +08:00
commit 3b90d1192f
2172 changed files with 594509 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
labml_nn/transformers/__init__.py
+115
View File
@@ -0,0 +1,115 @@
"""
---
title: Transformers
summary: >
This is a collection of PyTorch implementations/tutorials of
transformers and related techniques.
---
# Transformers
This module contains [PyTorch](https://pytorch.org/)
implementations and explanations of original transformer
from paper [Attention Is All You Need](https://arxiv.org/abs/1706.03762),
and derivatives and enhancements of it.
* [Multi-head attention](mha.html)
* [Transformer Encoder and Decoder Models](models.html)
* [Position-wise Feed Forward Network (FFN)](feed_forward.html)
* [Fixed positional encoding](positional_encoding.html)
## [Transformer XL](xl/index.html)
This implements Transformer XL model using
[relative multi-head attention](xl/relative_mha.html)
## [Rotary Positional Embeddings](rope/index.html)
This implements Rotary Positional Embeddings (RoPE)
## [Attention with Linear Biases](alibi/index.html)
This implements Attention with Linear Biases (ALiBi).
## [RETRO](retro/index.html)
This implements the Retrieval-Enhanced Transformer (RETRO).
## [Compressive Transformer](compressive/index.html)
This is an implementation of compressive transformer
that extends upon [Transformer XL](xl/index.html) by compressing
the oldest memories to give a longer attention span.
## [GPT Architecture](gpt/index.html)
This is an implementation of GPT-2 architecture.
## [GLU Variants](glu_variants/simple.html)
This is an implementation of the paper
[GLU Variants Improve Transformer](https://arxiv.org/abs/2002.05202).
## [kNN-LM](knn/index.html)
This is an implementation of the paper
[Generalization through Memorization: Nearest Neighbor Language Models](https://arxiv.org/abs/1911.00172).
## [Feedback Transformer](feedback/index.html)
This is an implementation of the paper
[Accessing Higher-level Representations in Sequential Transformers with Feedback Memory](https://arxiv.org/abs/2002.09402).
## [Switch Transformer](switch/index.html)
This is a miniature implementation of the paper
[Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://arxiv.org/abs/2101.03961).
Our implementation only has a few million parameters and doesn't do model parallel distributed training.
It does single GPU training but we implement the concept of switching as described in the paper.
## [Fast Weights Transformer](fast_weights/index.html)
This is an implementation of the paper
[Linear Transformers Are Secretly Fast Weight Memory Systems in PyTorch](https://arxiv.org/abs/2102.11174).
## [FNet: Mixing Tokens with Fourier Transforms](fnet/index.html)
This is an implementation of the paper
[FNet: Mixing Tokens with Fourier Transforms](https://arxiv.org/abs/2105.03824).
## [Attention Free Transformer](aft/index.html)
This is an implementation of the paper
[An Attention Free Transformer](https://arxiv.org/abs/2105.14103).
## [Masked Language Model](mlm/index.html)
This is an implementation of Masked Language Model used for pre-training in paper
[BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://arxiv.org/abs/1810.04805).
## [MLP-Mixer: An all-MLP Architecture for Vision](mlp_mixer/index.html)
This is an implementation of the paper
[MLP-Mixer: An all-MLP Architecture for Vision](https://arxiv.org/abs/2105.01601).
## [Pay Attention to MLPs (gMLP)](gmlp/index.html)
This is an implementation of the paper
[Pay Attention to MLPs](https://arxiv.org/abs/2105.08050).
## [Vision Transformer (ViT)](vit/index.html)
This is an implementation of the paper
[An Image Is Worth 16x16 Words: Transformers For Image Recognition At Scale](https://arxiv.org/abs/2010.11929).
## [Primer EZ](primer_ez/index.html)
This is an implementation of the paper
[Primer: Searching for Efficient Transformers for Language Modeling](https://arxiv.org/abs/2109.08668).
## [Hourglass](hour_glass/index.html)
This is an implementation of the paper
[Hierarchical Transformers Are More Efficient Language Models](https://arxiv.org/abs/2110.13711)
"""
from .configs import TransformerConfigs
from .models import TransformerLayer, Encoder, Decoder, Generator, EncoderDecoder
from .mha import MultiHeadAttention
from labml_nn.transformers.xl.relative_mha import RelativeMultiHeadAttention
labml_nn/transformers/aft/__init__.py
+234
View File
@@ -0,0 +1,234 @@
"""
---
title: An Attention Free Transformer
summary: >
This is an annotated implementation/tutorial of the AFT (Attention Free Transformer) in PyTorch.
---
# An Attention Free Transformer
This is a [PyTorch](https://pytorch.org) implementation of the paper
[An Attention Free Transformer](https://arxiv.org/abs/2105.14103).
This paper replaces the [self-attention layer](../mha.html) with a new efficient operation,
that has memory complexity of $\mathcal{O}(Td)$, where $T$ is the sequence length
and $d$ is the dimensionality of embeddings.
The paper introduces AFT along with AFT-local and AFT-conv.
Here we have implemented AFT-local which pays attention to closeby tokens
in an autoregressive model.
## Attention Free Transformer
AFT (similar to [MHA](../mha.html)) first transforms the embeddings $X$ into
query $Q = XW^Q$, key $K = XW^K$ and value $V = XW^V$ tensors with learned weights.
The output for each position $t \in [1, T]$ is calculated with the following operation.
$$Y_t = \sigma(Q_t) \odot
\frac{\sum_{t'=1}^T \exp(K_{t'} + w_{t,t'}) \odot V_{t'}}
{\sum_{t'=1}^T \exp(K_{t'} + w_{t,t'})}$$
, where $\odot$ is element-wise product, $\sigma$ is a non-linearity (sigmoid) and
$w \in \mathbb{R}^{T \times T}$ is a learned matrix of pair-wise position biases.
This means that we take the weighted average of values
and multiply them by the query. This eliminates the need to calculate the $T \times T$ attention
matrix that [MHA](../mha.html) requires, and therefore reduce the memory requirement.
## AFT Local
AFT Local only apply learned pair-wise position biases locally:
\begin{align}
w'_{t,t'} =
\begin{cases}
w_{t,t'}, & {\text{for } \lvert t-t' \rvert \lt s} \\
0, & \text{otherwise}
\end{cases}
\end{align}
, where $s \le T$ is the local window size.
Although $w'_{t,t'}$ is $0$ outside the local window the AFT operation still uses key-value pairs from
other areas. This is different from local transformers where embeddings outside the local window are
completely not visible.
Here is [the training code](experiment.html) for a AFT Local model.
"""
from typing import Optional
import torch
from torch import nn
class AFTLocal(nn.Module):
"""
### AFT Local Operation
$$Y_t = \sigma(Q_t) \odot
\frac{\sum_{t'=1}^T \exp(K_{t'} + w_{t,t'}) \odot V_{t'}}
{\sum_{t'=1}^T \exp(K_{t'} + w_{t,t'})}$$
where,
\begin{align}
w'_{t,t'} =
\begin{cases}
w_{t,t'}, & {\text{for } \lvert t-t' \rvert \lt s} \\
0, & \text{otherwise}
\end{cases}
\end{align}
"""
def __init__(self, d_model: int, seq_len: int, local_window_size: int, bias: bool = True):
"""
* `d_model` is the number of features in the `query`, `key` and `value` vectors.
* `seq_len` is $T$
* `local_window_size` is the local window size $s$
* `bias` is whether to have a bias parameter for transformations for $Q$, $K$ and $V$.
"""
super().__init__()
# Local window size $s$
self.local_window_size = local_window_size
# These transform the `query`, `key` and `value` vectors.
self.query = nn.Linear(d_model, d_model, bias=bias)
self.key = nn.Linear(d_model, d_model, bias=bias)
self.value = nn.Linear(d_model, d_model, bias=bias)
# Pair-wise positional biases $w \in \mathbb{R}^{T \times T}$
self.pos_bias = nn.Parameter(torch.zeros(seq_len, seq_len), requires_grad=True)
# Mask for $w_{t,t'}$
self.local_mask = nn.Parameter(self.create_local_mask(seq_len, local_window_size), requires_grad=False)
# Activation $\sigma$
self.activation = nn.Sigmoid()
# Output layer
self.output = nn.Linear(d_model, d_model)
@staticmethod
def create_local_mask(seq_len, local_window_size):
"""
#### Create local mask
This creates a mask for
\begin{align}
m_{t,t'} =
\begin{cases}
1, & {\text{for } \lvert t-t' \rvert \lt s} \\
0, & \text{otherwise}
\end{cases}
\end{align}
"""
# Initialize to ones
local_mask = torch.ones(seq_len, seq_len, dtype=torch.bool)
# Make $t' - t \ge s$ zero
local_mask = torch.tril(local_mask, local_window_size - 1)
# Make $t - t' \ge s$ zero
local_mask = torch.triu(local_mask, -(local_window_size - 1))
#
return local_mask
def forward(self, *,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
mask: Optional[torch.Tensor] = None):
"""
`query`, `key` and `value` are the tensors that store
collection of token embeddings for *query*, *key* and *value*.
They have shape `[seq_len, batch_size, d_model]`.
`mask` has shape `[seq_len, seq_len, batch_size]` and
`mask[i, j, b]` indicates whether for batch `b`,
query at position `i` has access to key-value at position `j`.
"""
# `query`, `key` and `value` have shape `[seq_len, batch_size, d_model]`
seq_len, _, _ = query.shape
if mask is not None:
# `mask` has shape `[seq_len_q, seq_len_k, batch_size]`,
# where first dimension is the query dimension.
# If the query dimension is equal to $1$ it will be broadcasted.
assert mask.shape[0] == 1 or mask.shape[0] == query.shape[0]
assert mask.shape[1] == key.shape[0]
assert mask.shape[2] == 1 or mask.shape[2] == query.shape[1]
# Transform query, key and value embeddings
query = self.query(query)
key = self.key(key)
value = self.value(value)
# Get
#
# \begin{align}
# w'_{t,t'} =
# \begin{cases}
# w_{t,t'}, & {\text{for }\lvert t-t' \rvert \lt s} \\
# 0, & \text{otherwise}
# \end{cases}
# \end{align}
#
# using the mask
pos_bias = self.pos_bias[:seq_len, :seq_len] * self.local_mask[:seq_len, :seq_len]
pos_bias = pos_bias.unsqueeze(-1)
pos_bias.masked_fill_(~mask, float('-inf'))
# \begin{align}
# Y_t &= \sigma(Q_t) \odot
# \frac{\sum_{t'=1}^T \exp(K_{t'} + w_{t,t'}) \odot V_{t'}}
# {\sum_{t'=1}^T \exp(K_{t'} + w_{t,t'})} \\
# &= \sigma(Q_t) \odot
# \frac{\sum_{t'=1}^T \exp(w_{t,t'}) \odot \exp(K_{t'}) \odot V_{t'}}
# {\sum_{t'=1}^T \exp(w_{t,t'}) \odot \exp(K_{t'})}
# \end{align}
#
# We compute $\exp(w_{t,t'})$, $\exp(K_{t'}) \odot V_{t'}$ and $\exp(K_{t'})$
# separately and do a matrix multiplication. We use einsum for clarity.
# We subtract $\max_{t'}(K_{t'})$ and $\max_{t'}(w_{t,t'})$ before calculating the exponents to stabilize
# the softmax calculation.
#
# If $x_i$ is large $\exp(x_i)$ becomes huge and the computation of
# $\frac{\sum\exp(x_i)y_i}{\sum\exp(x_i)}$becomes unstable.
# Subtracting a constant before calculating the exponent from numerator and denominator will cancel out.
# and can help stabilize the computation.
# So we subtract $\max(x_i)$ to stabilize the computation.
max_key = key.max(dim=0, keepdims=True)[0]
max_pos_bias = pos_bias.max(dim=1, keepdims=True)[0]
# $\exp \big(K_{t'}- \max_{t'}(K_{t'})\big)$
exp_key = torch.exp(key - max_key)
# $\exp \big(w_{t,t'} - \max_{t'}(w_{t,t'})\big)$
exp_pos_bias = torch.exp(pos_bias - max_pos_bias)
# The numerator part $\sum_{t'=1}^T \exp(w_{t,t'}) \odot \exp(K_{t'}) \odot V_{t'}$
num = torch.einsum('ijb,jbd->ibd', exp_pos_bias, exp_key * value)
# The denominator part $\sum_{t'=1}^T \exp(w_{t,t'}) \odot \exp(K_{t'})$
den = torch.einsum('ijb,jbd->ibd', exp_pos_bias, exp_key)
# Output $$Y_t = \sigma(Q_t) \odot
# \frac{\sum_{t'=1}^T \exp(w_{t,t'}) \odot \exp(K_{t'}) \odot V_{t'}}
# {\sum_{t'=1}^T \exp(w_{t,t'}) \odot \exp(K_{t'})}$$
y = self.activation(query) * num / den
# Output layer
return self.output(y)
def _test_local_mask():
"""
Test local mask
"""
from labml.logger import inspect
inspect(AFTLocal.create_local_mask(10, 4))
#
if __name__ == '__main__':
_test_local_mask()
labml_nn/transformers/aft/experiment.py
+162
View File
@@ -0,0 +1,162 @@
"""
---
title: Attention Free Transformer (AFT) Experiment
summary: This experiment trains an Attention Free Transformer (AFT) based model on Tiny Shakespeare dataset.
---
# [Attention Free Transformer (AFT)](index.html) Experiment
This is an annotated PyTorch experiment to train a [AFT model](index.html).
This is based on
[general training loop and configurations for auto-regressive NLP task](../../experiments/nlp_autoregression.html).
"""
import torch
from labml import experiment
from labml.configs import option
from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs
from labml_nn.transformers import TransformerConfigs, Encoder
from labml_nn.transformers.utils import subsequent_mask
from torch import nn
class AutoregressiveTransformer(nn.Module):
"""
## Simple autoregressive model
This consists of a token embedding layer, transformer encoder, and
a final linear layer that gives token logits.
"""
def __init__(self, encoder: Encoder, src_embed: nn.Module, generator: nn.Module):
"""
* `encoder` is the transformer [Encoder](../models.html#Encoder)
* `src_embed` is the token
[embedding module (with positional encodings)](../models.html#EmbeddingsWithLearnedPositionalEncoding)
* `generator` is the [final fully connected layer](../models.html#Generator) that gives the logits.
"""
super().__init__()
self.src_embed = src_embed
self.encoder = encoder
self.generator = generator
# The mask will be initialized on the first call
self.mask = None
def forward(self, x: torch.Tensor):
# Create subsequent mask if mask is not initialized
# or if the size of the mask is different
if self.mask is None or self.mask.size(0) != len(x):
# Subsequent mask, will mask out tokens from seeing future tokens
self.mask = subsequent_mask(len(x)).to(x.device)
# Get the token embeddings with positional encodings
x = self.src_embed(x)
# Transformer encoder
x = self.encoder(x, self.mask)
# Get logits
x = self.generator(x)
# Return results
# (second value is for state, since our trainer is used with RNNs also)
return x, None
class Configs(NLPAutoRegressionConfigs):
"""
## Configurations
This inherits from
[`NLPAutoRegressionConfigs`](../../experiments/nlp_autoregression.html#NLPAutoRegressionConfigs)
"""
# GPT model
model: AutoregressiveTransformer
# Transformer
transformer: TransformerConfigs
local_window_size: int = 32
@option(Configs.transformer, 'Transformer')
def _transformer_configs(c: Configs):
"""
### Transformer configurations
"""
# We use our
# [configurable transformer implementation](../configs.html#TransformerConfigs)
conf = TransformerConfigs()
# Set the vocabulary sizes for embeddings and generating logits
conf.n_src_vocab = c.n_tokens
conf.n_tgt_vocab = c.n_tokens
# Set the embedding size
conf.d_model = c.d_model
# Replace self-attention with an [AFT Local Module](index.html)
from labml_nn.transformers.aft import AFTLocal
conf.encoder_attn = AFTLocal(c.d_model, c.seq_len, c.local_window_size)
#
return conf
@option(Configs.model)
def _model(c: Configs):
"""
Create an auto-regressive model
"""
m = AutoregressiveTransformer(c.transformer.encoder,
c.transformer.src_embed,
c.transformer.generator).to(c.device)
return m
def main():
# Create experiment
experiment.create(name="aft")
# Create configs
conf = Configs()
# Override configurations
experiment.configs(conf, {
# Use character level tokenizer
'tokenizer': 'character',
# Prompt separator is blank
'prompt_separator': '',
# Starting prompt for sampling
'prompt': 'It is ',
# Use Tiny Shakespeare dataset
'text': 'tiny_shakespeare',
# Use a context size of $128$
'seq_len': 256,
# Train for $32$ epochs
'epochs': 128,
# Batch size $128$
'batch_size': 32,
# Switch between training and validation for $10$ times
# per epoch
'inner_iterations': 10,
# Embedding size
'd_model': 128,
# FFN hidden dimension size
'transformer.ffn.d_ff': 256,
# Optimizer
'optimizer.optimizer': 'Noam',
'optimizer.learning_rate': 1.,
})
# Set models for saving and loading
experiment.add_pytorch_models({'model': conf.model})
# Start the experiment
with experiment.start():
# Run training
conf.run()
#
if __name__ == '__main__':
main()
labml_nn/transformers/aft/readme.md
+13
View File
@@ -0,0 +1,13 @@
# [An Attention Free Transformer](https://nn.labml.ai/transformers/aft/index.html)
This is a [PyTorch](https://pytorch.org) implementation of the paper
[An Attention Free Transformer](https://arxiv.org/abs/2105.14103).
This paper replaces the [self-attention layer](https://nn.labml.ai/transformers/mha.html)
with a new efficient operation,
that has memory complexity of O(Td), where T is the sequence length
and $d$ is the dimensionality of embeddings.
The paper introduces AFT along with AFT-local and AFT-conv.
Here we have implemented AFT-local which pays attention to closeby tokens
in an autoregressive model.
labml_nn/transformers/alibi/__init__.py
+205
View File
@@ -0,0 +1,205 @@
"""
---
title: Attention with Linear Biases (ALiBi)
summary: >
Documented implementation with explanations of Attention with Linear Biases (ALiBi)
---
# Attention with Linear Biases (ALiBi)
This is an implementation of Attention with Linear Biases (ALiBi) from the paper
[Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation](https://arxiv.org/abs/2108.12409).
This replaces positional encodings with biases added to attention scores (attention logits, before the softmax).
This is a relative scheme tested on autoregressive tasks, and the bias is higher for closeby tokens
and lower for far-away tokens.
The biases decrease linearly in the log scale (because it's before the softmax) and each head has a different slope.
Here's the attention formula for $i$-th token,
\begin{align}
\mathbf{a}_i
&= \text{softmax} \bigg( \mathbf{q}_i \mathbf{K}^\top + m \cdot \big[-(i-1), \dots, -1, 0 \big] \bigg) \\
&= \text{softmax} \bigg( \mathbf{q}_i \mathbf{K}^\top + m \cdot \big[0, 1, \dots, (i - 1) \big] \bigg)
\end{align}
where $\mathbf{q}_i \in \mathbb{R}^d$ is the query of the $i$-th token, $K \in \mathbb{R}^{i \times d}$ are the keys
up to $i$, and $d$ the number of features per head.
Note that the above equality halts because $\text{softmax}$ is invariant to translations
(you can add any constant to all elements without changing the result).
Here is [the training code](experiment.html) for a ALiBi model.
"""
import math
from typing import Optional
import torch
from torch import nn
from labml.logger import inspect
from labml_nn.transformers.mha import MultiHeadAttention
def get_slopes(n_heads: int):
"""
## Get head-specific slope $m$ for each head
* `n_heads` is the number of heads in the attention layer $n$
The slope for first head is
$$\frac{1}{2^{\frac{8}{n}}} = 2^{-\frac{8}{n}}$$
The slopes for the rest of the heads are in a geometric series with a ratio same as above.
For instance when the number of heads is $8$ the slopes are
$$\frac{1}{2^1}, \frac{1}{2^2}, \dots, \frac{1}{2^8}$$
"""
# Get the closest power of 2 to `n_heads`.
# If `n_heads` is not a power of 2, then we first calculate slopes to the closest (smaller) power of 2,
# and then add the remaining slopes.
n = 2 ** math.floor(math.log2(n_heads))
# $2^{-\frac{8}{n}}$
m_0 = 2.0 ** (-8.0 / n)
# $2^{-1\frac{8}{n}}, 2^{-2 \frac{8}{n}}, 2^{-3 \frac{8}{n}}, \dots$
m = torch.pow(m_0, torch.arange(1, 1 + n))
# If `n_heads` is not a power of 2, then we add the remaining slopes.
# We calculate the remaining slopes for $n * 2$ (avoiding slopes added previously).
# And pick the slopes upto `n_heads`.
if n < n_heads:
# $2^{-\frac{8}{2n}}$
m_hat_0 = 2.0 ** (-4.0 / n)
# $2^{-1\frac{8}{2n}}, 2^{-3 \frac{8}{2n}}, 2^{-5 \frac{8}{2n}}, \dots$
# Note that we take steps by $2$ to avoid slopes added previously.
m_hat = torch.pow(m_hat_0, torch.arange(1, 1 + 2 * (n_heads - n), 2))
# Concatenate the slopes with the remaining slopes.
m = torch.cat([m, m_hat])
return m
@torch.no_grad()
def get_alibi_biases(n_heads: int, mask: torch.Tensor):
"""
## Calculate the attention biases matrix
* `n_heads` is the number of heads in the attention layer
* `mask` is the attention mask of shape `[seq_len_q, seq_len_k]`
This returns a matrix of shape `[seq_len_q, seq_len_k, n_heads, ]` with ALiBi attention biases.
"""
# Get slopes $m$ for each head
m = get_slopes(n_heads).to(mask.device)
# Calculate distances $[0, 1, \dots, N]$
# Here we calculate the distances using the mask.
#
# Since it's causal mask we can just use $[0, 1, \dots, N]$ too.
# `distance = torch.arange(mask.shape[1], dtype=torch.long, device=mask.device)[None, :]`
distance = mask.cumsum(dim=-1)
# Multiply them pair-wise to get the AliBi bias matrix
return distance[:, :, None] * m[None, None, :]
class AlibiMultiHeadAttention(MultiHeadAttention):
"""
## Attention with Linear Biases (ALiBi)
We override [Multi-Head Attention](../mha.html).
"""
def __init__(self, heads: int, d_model: int, dropout_prob: float = 0.1):
super().__init__(heads, d_model, dropout_prob)
# To cache AliBi the biases
self.alibi_biases = None
def forward(self, *,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
mask: Optional[torch.Tensor] = None):
"""
`query`, `key` and `value` are the tensors that store
collection of *query*, *key* and *value* vectors.
They have shape `[seq_len, batch_size, d_model]`.
`mask` has shape `[seq_len, seq_len, batch_size]` and
`mask[i, j, b]` indicates whether for batch `b`,
query at position `i` has access to key-value at position `j`.
"""
# ALiBi only works with causal masks.
assert mask is not None
assert mask.shape[0] == mask.shape[1] and mask.shape[2] == 1
# `query`, `key` and `value` have shape `[seq_len, batch_size, d_model]`
seq_len, batch_size, _ = query.shape
# Add head dimension to mask and check its shape.
mask = self.prepare_mask(mask, query.shape, key.shape)
# Prepare `query`, `key` and `value` for attention computation.
# These will then have shape `[seq_len, batch_size, heads, d_k]`.
query = self.query(query)
key = self.key(key)
value = self.value(value)
# Compute attention scores $Q K^\top$.
# This gives a tensor of shape `[seq_len, seq_len, batch_size, heads]`.
scores = self.get_scores(query, key)
# Scale scores $\frac{Q K^\top}{\sqrt{d_k}}$
scores *= self.scale
# Create AliBi biases if it's not cached
if self.alibi_biases is None or self.alibi_biases.shape[1] < seq_len:
# `mask` has shape `[seq_len, seq_len, 1, 1]`
self.alibi_biases = get_alibi_biases(scores.shape[-1], mask[:, :, 0, 0])
# Add AliBi biases to attention scores.
# ALiBi biases has shape `[seq_len, seq_len, n_heads]`
# and `scores` has shape `[seq_len, seq_len, batch_size, n_heads]`
scores += self.alibi_biases[:seq_len, :seq_len, None, :]
# Apply mask
scores = scores.masked_fill(mask == 0, float('-inf'))
# $softmax$ attention along the key sequence dimension
# $\underset{seq}{softmax}\Bigg(\frac{Q K^\top}{\sqrt{d_k}}\Bigg)$
attn = self.softmax(scores)
# Apply dropout
attn = self.dropout(attn)
# Multiply by values
# $$\underset{seq}{softmax}\Bigg(\frac{Q K^\top}{\sqrt{d_k}}\Bigg)V$$
x = torch.einsum("ijbh,jbhd->ibhd", attn, value)
# Concatenate multiple heads
x = x.reshape(seq_len, batch_size, -1)
# Output layer
return self.output(x)
def _test_alibi():
"""
Simple test function to see the slopes.
"""
inspect(get_slopes(12).tolist(), _n=-1)
from labml_nn.transformers.utils import subsequent_mask
mask = subsequent_mask(8)[:, :, 0]
inspect(mask)
inspect(get_alibi_biases(12, mask)[:, :, 3], _n=-1)
#
if __name__ == '__main__':
_test_alibi()
labml_nn/transformers/alibi/experiment.py
+155
View File
@@ -0,0 +1,155 @@
"""
---
title: Attention with Linear Biases (ALiBi) Experiment
summary: This experiment trains an Attention with Linear Biases (ALiBi) based model on Tiny Shakespeare dataset.
---
# [Attention with Linear Biases (ALiBi)](index.html) Experiment
This is an annotated PyTorch experiment to train a [ALiBi model](index.html).
This is based on [our GPT model](../gpt/index.html).
"""
import torch
from torch.utils.data import DataLoader
from labml import experiment, tracker
from labml.configs import option, calculate
from labml_nn.helpers.datasets import SequentialUnBatchedDataset
from labml_nn.transformers.alibi import AlibiMultiHeadAttention
from labml_nn.experiments.nlp_autoregression import transpose_batch
from labml_nn.transformers import TransformerConfigs
from labml_nn.transformers.gpt import Configs as GPTConfigs
class Configs(GPTConfigs):
"""
## Configurations
We extend [GPT configurations](../gpt/index.html) and change the attention mechanism.
"""
# ALiBi based transformer (defined below)
transformer: TransformerConfigs = 'GPT_ALiBi'
# Longer validation set
valid_seq_len: int = 128
valid_loader = 'shuffled_longer_valid_loader'
def other_metrics(self, output: torch.Tensor, target: torch.Tensor):
"""
Log losses at the initial and final tokens
"""
# If there are more tokens that the training sequence length (during validation),
if self.seq_len < output.shape[0]:
# Log the loss at training sequence length
tracker.add(f'loss.{self.seq_len - 1}.', self.loss_func(output[self.seq_len - 1], target[self.seq_len - 1]))
# Log the loss at the first token
tracker.add(f'loss.0.', self.loss_func(output[0], target[0]))
# Log the loss at the final token
tracker.add(f'loss.{int(output.shape[0]) - 1}.', self.loss_func(output[-1], target[-1]))
def _alibi_mha(c: TransformerConfigs):
"""
Create an ALiBi attention module
"""
return AlibiMultiHeadAttention(c.n_heads, c.d_model, dropout_prob=c.dropout)
# Set all attention mechanisms to ALiBi
calculate(TransformerConfigs.encoder_attn, 'alibi_mha', _alibi_mha)
calculate(TransformerConfigs.decoder_attn, 'alibi_mha', _alibi_mha)
calculate(TransformerConfigs.decoder_mem_attn, 'alibi_mha', _alibi_mha)
@option(Configs.valid_loader)
def shuffled_longer_valid_loader(c: Configs):
"""
Shuffled validation data loader with `valid_seq_len` sequence length
"""
return DataLoader(SequentialUnBatchedDataset(text=c.text.valid,
dataset=c.text,
seq_len=c.valid_seq_len),
batch_size=c.batch_size,
collate_fn=transpose_batch,
shuffle=True)
@option(Configs.transformer, 'GPT_ALiBi')
def _transformer_configs(c: Configs):
"""
### ALiBi based Transformer configurations
"""
# We use our
# [configurable transformer implementation](../configs.html#TransformerConfigs)
conf = TransformerConfigs()
# Set the vocabulary sizes for embeddings and generating logits
conf.n_src_vocab = c.n_tokens
conf.n_tgt_vocab = c.n_tokens
# GPT uses GELU activation for position wise feedforward
conf.ffn.activation = 'GELU'
# ALiBi doesn't use positional embeddings
conf.src_embed = 'no_pos'
conf.tgt_embed = 'no_pos'
# Set all attention mechanisms to ALiBi
conf.encoder_attn = 'alibi_mha'
conf.decoder_attn = 'alibi_mha'
conf.decoder_mem_attn = 'alibi_mha'
#
return conf
def main():
# Create experiment
experiment.create(name="gpt_alibi")
# Create configs
conf = Configs()
# Override configurations
experiment.configs(conf, {
# Use character level tokenizer
'tokenizer': 'character',
# Prompt separator is blank
'prompt_separator': '',
# Starting prompt for sampling
'prompt': 'It is ',
# Use Tiny Shakespeare dataset
'text': 'tiny_shakespeare',
# 'text': 'tiny_shakespeare_no_split',
# Use a context size of $128$
'seq_len': 64,
# Use a context size of $128$
'valid_seq_len': 80,
# Train for $32$ epochs
'epochs': 128,
# Batch size $128$
'batch_size': 128,
# Switch between training and validation for $10$ times
# per epoch
'inner_iterations': 10,
# Transformer configurations
'transformer.d_model': 128,
'transformer.ffn.d_ff': 512,
'transformer.n_heads': 8,
'transformer.n_layers': 4,
'transformer.dropout': 0.1,
})
# Set models for saving and loading
experiment.add_pytorch_models({'model': conf.model})
# Start the experiment
with experiment.start():
# Run training
conf.run()
#
if __name__ == '__main__':
main()
labml_nn/transformers/basic/__init__.py
View File
labml_nn/transformers/basic/autoregressive_experiment.ipynb
+295
View File
@@ -0,0 +1,295 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "AYV_dMVDxyc2",
"pycharm": {
"name": "#%% md\n"
}
},
"source": [
"[![Github](https://img.shields.io/github/stars/labmlai/annotated_deep_learning_paper_implementations?style=social)](https://github.com/labmlai/annotated_deep_learning_paper_implementations)\n",
"[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/transformers/basic/autoregressive_experiment.ipynb)\n",
"\n",
"## Transformer Experiment\n",
"\n",
"This trains a simple transformer with\n",
"[multi headed attention](https://nn.labml.ai/transformers/mha.html)\n",
"introduced in [Attention Is All You Need](https://arxiv.org/abs/1706.03762)\n",
"on an NLP auto-regression task (with Tiny Shakespeare dataset)."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "AahG_i2y5tY9",
"pycharm": {
"name": "#%% md\n"
}
},
"source": [
"### Install the packages"
]
},
{
"cell_type": "code",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "ZCzmCrAIVg0L",
"outputId": "cf107fb2-4d50-4c67-af34-367624553421",
"pycharm": {
"name": "#%%\n"
}
},
"source": [
"!pip install labml-nn --quiet"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "SE2VUQ6L5zxI",
"pycharm": {
"name": "#%% md\n"
}
},
"source": [
"### Imports"
]
},
{
"cell_type": "code",
"metadata": {
"id": "0hJXx_g0wS2C",
"pycharm": {
"name": "#%%\n"
}
},
"source": [
"from labml import experiment\n",
"from labml_nn.transformers.basic.autoregressive_experiment import Configs"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "Lpggo0wM6qb-",
"pycharm": {
"name": "#%% md\n"
}
},
"source": [
"### Create an experiment"
]
},
{
"cell_type": "code",
"metadata": {
"id": "bFcr9k-l4cAg",
"pycharm": {
"name": "#%%\n"
}
},
"source": [
"experiment.create(name=\"transformer\", writers={'screen'})"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "-OnHLi626tJt",
"pycharm": {
"name": "#%% md\n"
}
},
"source": [
"### Configurations"
]
},
{
"cell_type": "code",
"metadata": {
"id": "Piz0c5f44hRo",
"pycharm": {
"name": "#%%\n"
}
},
"source": [
"conf = Configs()"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "wwMzCqpD6vkL",
"pycharm": {
"name": "#%% md\n"
}
},
"source": [
"Set experiment configurations and assign a configurations dictionary to override configurations"
]
},
{
"cell_type": "code",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 17
},
"id": "e6hmQhTw4nks",
"outputId": "29634715-42f4-4405-fb11-fc9522608627",
"pycharm": {
"name": "#%%\n"
}
},
"source": [
"experiment.configs(conf, {\n",
" # Use character level tokenizer\n",
" 'tokenizer': 'character',\n",
" # Prompt separator is blank\n",
" 'prompt_separator': '',\n",
" # Starting prompt for sampling\n",
" 'prompt': 'It is ',\n",
" # Use Tiny Shakespeare dataset\n",
" 'text': 'tiny_shakespeare',\n",
"\n",
" # Use a context size of $256$\n",
" 'seq_len': 512,\n",
" # Train for 32 epochs\n",
" 'epochs': 32,\n",
" # Batch size $32$\n",
" 'batch_size': 16,\n",
" # Switch between training and validation for $10$ times\n",
" # per epoch\n",
" 'inner_iterations': 10,\n",
"\n",
" # Model size\n",
" 'd_model': 256,\n",
" 'transformer.n_heads': 16,\n",
" 'transformer.ffn.d_ff': 1024,\n",
"\n",
" # Use [Noam optimizer](../../optimizers/noam.html)\n",
" 'optimizer.optimizer': 'Noam',\n",
" 'optimizer.learning_rate': 1.,\n",
"})"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "EvI7MtgJ61w5",
"pycharm": {
"name": "#%% md\n"
}
},
"source": [
"Set PyTorch models for loading and saving"
]
},
{
"cell_type": "code",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 255
},
"id": "GDlt7dp-5ALt",
"outputId": "e7548e8f-c541-4618-dc5a-1597cae42003",
"pycharm": {
"name": "#%%\n"
}
},
"source": [
"experiment.add_pytorch_models({'model': conf.model})"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "KJZRf8527GxL",
"pycharm": {
"name": "#%% md\n"
}
},
"source": [
"### Start the experiment and run the training loop."
]
},
{
"cell_type": "code",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 1000
},
"id": "aIAWo7Fw5DR8",
"outputId": "db979785-bfe3-4eda-d3eb-8ccbe61053e5",
"pycharm": {
"name": "#%%\n"
}
},
"source": [
"# Start the experiment\n",
"with experiment.start():\n",
" conf.run()"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "code",
"metadata": {
"id": "oBXXlP2b7XZO",
"pycharm": {
"name": "#%%\n"
}
},
"source": [],
"outputs": [],
"execution_count": null
}
],
"metadata": {
"accelerator": "GPU",
"colab": {
"collapsed_sections": [],
"name": "Transformer",
"provenance": []
},
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.7.11"
}
},
"nbformat": 4,
"nbformat_minor": 4
}
labml_nn/transformers/basic/autoregressive_experiment.py
+155
View File
@@ -0,0 +1,155 @@
"""
---
title: Transformer Auto-Regression Experiment
summary: >
This trains a simple transformer model on NLP auto-regression.
---
# Transformer Auto-Regression Experiment
[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/transformers/basic/autoregressive_experiment.ipynb)
This trains a simple transformer introduced in [Attention Is All You Need](https://arxiv.org/abs/1706.03762)
on an NLP auto-regression task (with Tiny Shakespeare dataset).
"""
import torch
from torch import nn
from labml import experiment
from labml.configs import option
from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs
from labml_nn.transformers import TransformerConfigs, Encoder
from labml_nn.transformers.utils import subsequent_mask
class AutoregressiveTransformer(nn.Module):
"""
## Auto-Regressive model
"""
def __init__(self, encoder: Encoder, src_embed: nn.Module, generator: nn.Module):
"""
* `encoder` is the transformer [Encoder](../models.html#Encoder)
* `src_embed` is the token
[embedding module (with positional encodings)](../models.html#EmbeddingsWithLearnedPositionalEncoding)
* `generator` is the [final fully connected layer](../models.html#Generator) that gives the logits.
"""
super().__init__()
self.src_embed = src_embed
self.encoder = encoder
self.generator = generator
# The mask will be initialized on the first call
self.mask = None
def forward(self, x: torch.Tensor):
# Create subsequent mask if mask is not initialized
# or if the size of the mask is different
if self.mask is None or self.mask.size(0) != len(x):
# Subsequent mask, will mask out tokens from seeing future tokens
self.mask = subsequent_mask(len(x)).to(x.device)
# Get the token embeddings with positional encodings
x = self.src_embed(x)
# Transformer encoder
x = self.encoder(x, self.mask)
# Get logits
x = self.generator(x)
# Return results
# (second value is for state, since our trainer is used with RNNs also)
return x, None
class Configs(NLPAutoRegressionConfigs):
"""
## Configurations
This inherits from
[`NLPAutoRegressionConfigs`](../../experiments/nlp_autoregression.html#NLPAutoRegressionConfigs)
"""
# GPT model
model: AutoregressiveTransformer
# Transformer
transformer: TransformerConfigs
@option(Configs.transformer, 'Transformer')
def _transformer_configs(c: Configs):
"""
### Transformer configurations
"""
# We use our
# [configurable transformer implementation](../configs.html#TransformerConfigs)
conf = TransformerConfigs()
# Set the vocabulary sizes for embeddings and generating logits
conf.n_src_vocab = c.n_tokens
conf.n_tgt_vocab = c.n_tokens
#
conf.d_model = c.d_model
#
return conf
@option(Configs.model)
def _model(c: Configs):
"""
Create GPT model and initialize weights
"""
m = AutoregressiveTransformer(c.transformer.encoder,
c.transformer.src_embed,
c.transformer.generator).to(c.device)
return m
def main():
# Create experiment
experiment.create(name="transformer")
# Create configs
conf = Configs()
# Override configurations
experiment.configs(conf, {
# Use character level tokenizer
'tokenizer': 'character',
# Prompt separator is blank
'prompt_separator': '',
# Starting prompt for sampling
'prompt': 'It is ',
# Use Tiny Shakespeare dataset
'text': 'tiny_shakespeare',
# Use a context size of $256$
'seq_len': 512,
# Train for 32 epochs
'epochs': 32,
# Batch size $32$
'batch_size': 16,
# Switch between training and validation for $10$ times
# per epoch
'inner_iterations': 10,
# Model size
'd_model': 256,
'transformer.n_heads': 16,
'transformer.ffn.d_ff': 1024,
# Use [Noam optimizer](../../optimizers/noam.html)
'optimizer.optimizer': 'Noam',
'optimizer.learning_rate': 1.,
})
# Set models for saving and loading
experiment.add_pytorch_models({'model': conf.model})
# Start the experiment
with experiment.start():
# Run training
conf.run()
#
if __name__ == '__main__':
main()
labml_nn/transformers/basic/with_sophia.py
+159
View File
@@ -0,0 +1,159 @@
"""
---
title: Transformer Auto-Regression Experiment with [Sophia-G optimizer](../../optimizers/sophia.html)
summary: >
This trains a simple transformer model on NLP auto-regression with Sophia-G optimizer.
---
# Transformer Auto-Regression Experiment with [Sophia-G optimizer](../../optimizers/sophia.html)
This trains a simple transformer introduced in [Attention Is All You Need](https://arxiv.org/abs/1706.03762)
on an NLP auto-regression task (with Tiny Shakespeare dataset) with [Sophia-G optimizer](../../optimizers/sophia.html).
"""
import torch
from labml import experiment, tracker
from labml_nn.helpers.trainer import BatchIndex
from labml_nn.optimizers.sophia import Sophia
from labml_nn.transformers.basic.autoregressive_experiment import Configs as TransformerAutoRegressionConfigs
class Configs(TransformerAutoRegressionConfigs):
"""
## Configurations
This inherits from [`Configs`](autoregressive_experiment.html)
"""
hess_interval: int = 10
optimizer: Sophia
def step(self, batch: any, batch_idx: BatchIndex):
"""
### Training or validation step with Gauss-Newton-Bartlett (GNB) Hessian diagonal estimator
"""
# Set training/eval mode
self.model.train(self.mode.is_train)
# Move data to the device
data, target = batch[0].to(self.device), batch[1].to(self.device)
# Estimate the Hessian diagonal every $k$ steps
if isinstance(self.optimizer, Sophia) and self.mode.is_train and batch_idx.idx % self.hess_interval == 0:
# Get model outputs
output, *_ = self.model(data)
# Create a categorical distribution from logits
samp_dist = torch.distributions.Categorical(logits=output)
# Sample $\hat{y}$
y_sample = samp_dist.sample()
# Calculate and log loss
loss = self.loss_func(output, y_sample)
tracker.add("loss.hess.", loss)
# Calculate gradients
loss.backward()
# Clip gradients
torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=self.grad_norm_clip)
# Update EMA Hessian diagonal
#
# \begin{align}
# \hat{h}_t &= B \cdot \nabla_\theta \hat{L} (\theta) \odot \nabla_\theta \hat{L} (\theta) \\
# h_t &= \beta_2 h_{t-k} + (1 - \beta_2) \hat{h}_t
# \end{align}
self.optimizer.update_hessian(data.numel())
# Clear the gradients
self.optimizer.zero_grad()
else:
# Move data to the device
data, target = batch[0].to(self.device), batch[1].to(self.device)
# Update global step (number of tokens processed) when in training mode
if self.mode.is_train:
tracker.add_global_step(data.shape[0] * data.shape[1])
# Get model outputs.
# It's returning a tuple for states when using RNNs.
# This is not implemented yet. 😜
output, *_ = self.model(data)
# Calculate and log loss
loss = self.loss_func(output, target)
tracker.add("loss.", loss)
# Calculate and log accuracy
self.accuracy(output, target)
self.accuracy.track()
self.other_metrics(output, target)
# Train the model
if self.mode.is_train:
# Calculate gradients
loss.backward()
# Clip gradients
torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=self.grad_norm_clip)
# Take optimizer step
self.optimizer.step()
# Log the model parameters and gradients on last batch of every epoch
if batch_idx.is_last and self.is_log_model_params_grads:
tracker.add('model', self.model)
# Clear the gradients
self.optimizer.zero_grad()
# Save the tracked metrics
tracker.save()
def main():
# Create experiment
experiment.create(name="transformer")
# Create configs
conf = Configs()
# Override configurations
experiment.configs(conf, {
# Use character level tokenizer
'tokenizer': 'character',
# Prompt separator is blank
'prompt_separator': '',
# Starting prompt for sampling
'prompt': 'It is ',
# Use Tiny Shakespeare dataset
'text': 'tiny_shakespeare',
# Use a context size of $256$
'seq_len': 512,
# Train for 32 epochs
'epochs': 32,
# Batch size $32$
'batch_size': 16,
# Switch between training and validation for $10$ times
# per epoch
'inner_iterations': 10,
# Model size
'd_model': 256,
'transformer.n_heads': 16,
'transformer.ffn.d_ff': 1024,
# Use [Sophia optimizer](../../optimizers/sophia.html)
'optimizer.optimizer': 'Sophia',
'optimizer.learning_rate': 3e-4,
'optimizer.rho': 0.03,
})
# Set models for saving and loading
experiment.add_pytorch_models({'model': conf.model})
# Start the experiment
with experiment.start():
# Run training
conf.run()
#
if __name__ == '__main__':
main()
labml_nn/transformers/compressive/__init__.py
+334
View File
@@ -0,0 +1,334 @@
"""
---
title: Compressive Transformer
summary: >
Documented implementation with explanations of a
Compressive Transformer model.
---
# Compressive Transformer
This is an implementation of
[Compressive Transformers for Long-Range Sequence Modelling](https://arxiv.org/abs/1911.05507)
in [PyTorch](https://pytorch.org).
This is an extension of [Transformer XL](../xl/index.html) where past memories
are compressed to give a longer attention range.
That is, the furthest $n_{cm} c$ memories are compressed into
$n_{cm}$ memories, where $c$ is the compression rate.
## Compression operation
The compression operation is defined as
$f_c: \mathbb{R}^{nc \times d} \rightarrow \mathbb{R}^{n \times d}$.
The paper introduces multiple choices for $f_c$ and we have only implemented
1D convolution which seems to give the best results.
Each layer has a separate compression operation $f_c^{(i)}$ where
$i$ is the layer number.
## Training compression operation
Since training compression with BPTT requires maintaining
a very large computational graph (many time steps), the paper proposes
an *auto-encoding loss* and an *attention reconstruction loss*.
The auto-encoding loss decodes the original memories from the compressed memories
and calculates the loss.
Attention reconstruction loss computes the multi-headed attention results
on the compressed memory and on uncompressed memory and gets a mean squared error
between them.
We have implemented the latter here since it gives better results.
This implementation uses pre-layer normalization
while the paper uses post-layer normalization.
Pre-layer norm does the layer norm before [FFN](../feedforward.html) and
self-attention, and the pass-through in the residual connection is not normalized.
This is supposed to be more stable in standard transformer setups.
Here are [the training code](experiment.html) and a notebook for training a compressive transformer
model on the Tiny Shakespeare dataset.
[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/transformers/compressive/experiment.ipynb)
"""
from typing import Optional, List
import torch
import torch.nn.functional as F
from torch import nn
from labml_nn.transformers.feed_forward import FeedForward
from labml_nn.transformers.mha import PrepareForMultiHeadAttention
from labml_nn.transformers.xl.relative_mha import RelativeMultiHeadAttention
from labml_nn.utils import clone_module_list
class Conv1dCompression(nn.Module):
"""
## 1D Convolution Compression $f_c$
This is a simple wrapper around
[`nn.Conv1d`](https://pytorch.org/docs/stable/generated/torch.nn.Conv1d.html)
with some tensor dimension permutations.
"""
def __init__(self, compression_rate: int, d_model: int):
"""
* `compression_rate` $c$
* `d_model` is the embedding size
"""
super().__init__()
self.conv = nn.Conv1d(d_model, d_model, kernel_size=compression_rate, stride=compression_rate)
def forward(self, mem: torch.Tensor):
"""
`mem` has shape `[seq_len, batch, d_model]`
"""
# Permute the dimensions of `mem` so that we can run it through the convolution layer.
# The convolution layer accepts in the form `[batch, features, sequence]`
mem = mem.permute(1, 2, 0)
# Get compressed memory by running it through the convolution layer
c_mem = self.conv(mem)
# Permute back to form `[seq_len, batch, d_model]`
return c_mem.permute(2, 0, 1)
class CompressiveTransformerLayer(nn.Module):
"""
## Compressive Transformer Layer
This is the implementation of a single compressive transformer layer
"""
def __init__(self, *,
d_model: int,
self_attn: RelativeMultiHeadAttention,
feed_forward: FeedForward,
dropout_prob: float,
compress: Conv1dCompression):
"""
* `d_model` is the token embedding size
* `self_attn` is the [self attention module](../xl/relative_mha.html)
* `feed_forward` is the [feed forward module](../feed_forward.html)
* `dropout_prob` is the probability of dropping out after self attention and FFN
* `compress` is the compression function $f_c$
"""
super().__init__()
self.compress = compress
self.size = d_model
self.self_attn = self_attn
self.feed_forward = feed_forward
self.dropout = nn.Dropout(dropout_prob)
self.norm_self_attn = nn.LayerNorm([d_model])
self.norm_ff = nn.LayerNorm([d_model])
def concat_memory(self, z: torch.Tensor, mem: Optional[torch.Tensor], c_mem: Optional[torch.Tensor]):
"""
Concatenate the normalized token embeddings with memory and compressed memory.
* `z` is layer normalized token embeddings.
* `mem` and `c_mem` are memory and compressed memory (not normalized).
"""
# If there is no memory just return the token embeddings
if mem is None:
return z
# If there are compressed memory concatenate that with memory
if c_mem is not None:
mem = torch.cat((c_mem, mem), dim=0)
# Run the memory through the normalization layer
mem = self.norm_self_attn(mem)
# Concatenate normalized memory and normalized token embeddings
return torch.cat((mem, z), dim=0)
def forward(self, *,
x: torch.Tensor,
mem: Optional[torch.Tensor],
c_mem: Optional[torch.Tensor],
mask: torch.Tensor):
"""
* `x` is a tensor of token level feature vectors of shape `[seq_len, batch_size, d_model]`
* `mem` is a tensor of the past token level feature vectors (memory) of shape `[mem_len, batch_size, d_model]`
* `c_mem` is a tensor of the compressed memory `[c_mem_len, batch_size, d_model]`
* `mask` is a matrix of shape `[seq_len, c_mem_len + mem_len + seq_len, batch_size]` or `[seq_len, c_mem_len + mem_len + seq_len, 1]`.
`mask[i, j]` is true if token at `i` can see token at `j`.
"""
# Normalize the vectors before doing self attention
z = self.norm_self_attn(x)
# Normalize and concatenate memory and compressed memory
m_z = self.concat_memory(z, mem, c_mem)
# Attention
self_attn = self.self_attn(query=z, key=m_z, value=m_z, mask=mask)
# Add the attention results
x = x + self.dropout(self_attn)
# Normalize for feed-forward
z = self.norm_ff(x)
# Pass through the feed-forward network
ff = self.feed_forward(z)
# Add the feed-forward results back
x = x + self.dropout(ff)
#
return x
class CompressiveTransformer(nn.Module):
"""
## Compressive Transformer Model
This consists of multiple compressive transformer layers
"""
def __init__(self, layer: CompressiveTransformerLayer, n_layers: int):
super().__init__()
# Make copies of the transformer layer
self.layers = clone_module_list(layer, n_layers)
# Final normalization layer
self.norm = nn.LayerNorm([layer.size])
def forward(self, x: torch.Tensor, mem: List[torch.Tensor], c_mem: List[torch.Tensor], mask: torch.Tensor):
"""
* `x` is a tensor of the token embeddings vectors of shape `[seq_len, batch_size, d_model]`
* `mem` is a list of tensors of the past token level feature vectors of shape
`[mem_len, batch_size, d_model]` for each layer
* `c_mem` is a list of tensors of the compressed memory
`[c_mem_len, batch_size, d_model]` for each layer
* `mask` is the masking matrix
"""
# List to store token level feature vectors,
# which will become the memories for the next sequential batch.
new_mem = []
# Run through each transformer layer
for i, layer in enumerate(self.layers):
# Add to the list of feature vectors
new_mem.append(x.detach())
# Memory
m = mem[i] if mem else None
# Compressed Memory
cm = c_mem[i] if c_mem else None
# Run through the transformer XL layer
x = layer(x=x, mem=m, c_mem=cm, mask=mask)
# Finally, normalize the vectors
return self.norm(x), new_mem
class AttentionReconstructionLoss:
"""
## Attention Reconstruction Loss
Attention reconstruction loss recreates the self-attention output with
uncompressed memory and with compressed memory and calculates the mean squared error
between the two. It does this without positional encoding.
When calculating and training the compression function $f_c$ with attention
reconstruction loss, all parameters but $f_c$ are frozen.
This includes key/value projections and bias/scaling after normalization.
Since this loss can be computed independently of the cross-entropy-loss of the model
you can have a separate optimizer that only updates $f_c$.
However, we use the same optimizer to update $f_c$ so when calculating
attention reconstruction loss, we detach all other parameters except $f_c$
from the gradient computation.
"""
def __init__(self, layers: nn.ModuleList):
"""
`layers` is the list of Compressive Transformer layers
"""
self.layers = layers
self.loss_func = nn.MSELoss()
def prepare_for_attn(self, pmha: PrepareForMultiHeadAttention, x: torch.Tensor):
"""
This is a reimplementation of ['PrepareForMultiHeadAttention'](../mha.html#PrepareMHA)
where the projections are done with the parameters detached from gradient computation.
* `pmha` is the ['PrepareForMultiHeadAttention'](../mha.html#PrepareMHA) module
* `x` is tensor with the token embeddings
"""
# Shape of the input except embedding dimension; `[seq_len, batch_size]`.
head_shape = x.shape[:-1]
# Detach projection weights and bias
weight = pmha.linear.weight.detach()
bias = pmha.linear.bias.detach() if pmha.linear.bias is not None else None
# Linear transform
x = F.linear(x, weight, bias)
# Split last dimension into heads
x = x.view(*head_shape, pmha.heads, pmha.d_k)
# Output has shape `[seq_len, batch_size, heads, d_k]` or `[batch_size, d_model]`
return x
def attn(self, layer: RelativeMultiHeadAttention, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor):
"""
This is a reimplementation of ['Multi-Head Attention'](../mha.html#MHA) which calls
`prepare_for_attn` instead of ['PrepareForMultiHeadAttention'](../mha.html#PrepareMHA)
to detach projection parameters.
"""
# Calculate query, key and value projections
query = self.prepare_for_attn(layer.query, query)
key = self.prepare_for_attn(layer.key, key)
value = self.prepare_for_attn(layer.value, value)
# Compute attention scores $Q K^\top$.
# This gives a tensor of shape `[seq_len, seq_len, batch_size, heads]`.
scores = torch.einsum('ibhd,jbhd->ijbh', query, key)
# Scale scores $\frac{Q K^\top}{\sqrt{d_k}}$
scores *= layer.scale
# $softmax$ attention along the key sequence dimension
# $\underset{seq}{softmax}\Bigg(\frac{Q K^\top}{\sqrt{d_k}}\Bigg)$
attn = layer.softmax(scores)
# Multiply by values
# $$\underset{seq}{softmax}\Bigg(\frac{Q K^\top}{\sqrt{d_k}}\Bigg)V$$
return torch.einsum("ijbh,jbhd->ibhd", attn, value)
def norm(self, ln: nn.LayerNorm, x: torch.Tensor):
"""
Perform layer normalization with shift and scale parameters detached.
"""
# Detach shift(`bias`) and scaling(`weight`) parameters
weight = ln.weight.detach() if ln.weight is not None else None
bias = ln.bias.detach() if ln.bias is not None else None
# Layer normalization
return F.layer_norm(x, ln.normalized_shape, weight, bias, ln.eps)
def calc_loss(self, layer: CompressiveTransformerLayer, h: torch.Tensor, mem: torch.Tensor):
"""
This calculates the loss for a layer
"""
# Detach the token embeddings and memory.
h = h.detach()
mem = mem.detach()
# Compress the memory with $f_c^{(i)}$.
# The parameters of $f_c^{(i)}$ are the only parameters not detached from gradient computation.
c_mem = layer.compress(mem)
# Normalize the embeddings and memories
h = self.norm(layer.norm_self_attn, h)
mem = self.norm(layer.norm_self_attn, mem)
c_mem = self.norm(layer.norm_self_attn, c_mem)
# Calculate the attention with uncompressed memory
attn_mem = self.attn(layer.self_attn, h, mem, mem)
# Calculate the attention with compressed memory
attn_cmem = self.attn(layer.self_attn, h, c_mem, c_mem)
# Calculate the mean square error
return self.loss_func(attn_cmem, attn_mem)
def __call__(self, h: List[torch.Tensor], mem: List[torch.Tensor]):
# Calculate the losses for each layer
losses = [self.calc_loss(layer, h[n], mem[n]) for n, layer in enumerate(self.layers)]
# Sum of the losses
return sum(losses)
labml_nn/transformers/compressive/experiment.ipynb
+227
View File
@@ -0,0 +1,227 @@
{
"nbformat": 4,
"nbformat_minor": 0,
"metadata": {
"colab": {
"name": "Compressive Transformer",
"provenance": [],
"collapsed_sections": []
},
"kernelspec": {
"name": "python3",
"display_name": "Python 3"
},
"accelerator": "GPU"
},
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "AYV_dMVDxyc2"
},
"source": [
"[![Github](https://img.shields.io/github/stars/labmlai/annotated_deep_learning_paper_implementations?style=social)](https://github.com/labmlai/annotated_deep_learning_paper_implementations)\n",
"[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/transformers/compressive/experiment.ipynb) \n",
"\n",
"## Compressive Transformer\n",
"\n",
"This is an experiment training Shakespeare dataset with a Compressive Transformer model."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "AahG_i2y5tY9"
},
"source": [
"Install the `labml-nn` package"
]
},
{
"cell_type": "code",
"metadata": {
"id": "ZCzmCrAIVg0L",
"colab": {
"base_uri": "https://localhost:8080/"
},
"outputId": "cf107fb2-4d50-4c67-af34-367624553421"
},
"source": [
"!pip install labml-nn"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "SE2VUQ6L5zxI"
},
"source": [
"Imports"
]
},
{
"cell_type": "code",
"metadata": {
"id": "0hJXx_g0wS2C"
},
"source": [
"import torch\n",
"import torch.nn as nn\n",
"\n",
"from labml import experiment\n",
"from labml.configs import option\n",
"from labml_nn.transformers.compressive.experiment import Configs"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "Lpggo0wM6qb-"
},
"source": [
"Create an experiment"
]
},
{
"cell_type": "code",
"metadata": {
"id": "bFcr9k-l4cAg"
},
"source": [
"experiment.create(name=\"compressive_transformer\")"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "-OnHLi626tJt"
},
"source": [
"Initialize configurations"
]
},
{
"cell_type": "code",
"metadata": {
"id": "Piz0c5f44hRo"
},
"source": [
"conf = Configs()"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "wwMzCqpD6vkL"
},
"source": [
"Set experiment configurations and assign a configurations dictionary to override configurations"
]
},
{
"cell_type": "code",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 17
},
"id": "e6hmQhTw4nks",
"outputId": "29634715-42f4-4405-fb11-fc9522608627"
},
"source": [
"experiment.configs(conf,\n",
" # A dictionary of configurations to override\n",
" {'tokenizer': 'character',\n",
" 'text': 'tiny_shakespeare',\n",
" 'optimizer.learning_rate': 2.5e-4,\n",
" 'optimizer.optimizer': 'AdamW',\n",
" 'prompt': 'It is',\n",
" 'prompt_separator': '',\n",
"\n",
" 'train_loader': 'sequential_train_loader',\n",
" 'valid_loader': 'sequential_valid_loader',\n",
"\n",
" 'seq_len': 8,\n",
" 'mem_len': 8,\n",
" 'epochs': 128,\n",
" 'batch_size': 32,\n",
" 'inner_iterations': 25,\n",
" 'compression_rate': 2,\n",
" })"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "EvI7MtgJ61w5"
},
"source": [
"Set PyTorch models for loading and saving"
]
},
{
"cell_type": "code",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 255
},
"id": "GDlt7dp-5ALt",
"outputId": "e7548e8f-c541-4618-dc5a-1597cae42003"
},
"source": [
"experiment.add_pytorch_models({'model': conf.model})"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "KJZRf8527GxL"
},
"source": [
"Start the experiment and run the training loop."
]
},
{
"cell_type": "code",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 1000
},
"id": "aIAWo7Fw5DR8",
"outputId": "db979785-bfe3-4eda-d3eb-8ccbe61053e5"
},
"source": [
"# Start the experiment\n",
"with experiment.start():\n",
" conf.run()"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "code",
"metadata": {
"id": "oBXXlP2b7XZO"
},
"source": [
""
],
"outputs": [],
"execution_count": null
}
]
}
labml_nn/transformers/compressive/experiment.py
+346
View File
@@ -0,0 +1,346 @@
"""
---
title: Compressive Transformer Experiment
summary: This experiment trains a compressive transformer model on tiny Shakespeare dataset.
---
# Compressive Transformer Experiment
This is an annotated PyTorch experiment to train a compressive transformer model.
"""
from typing import List, Tuple, NamedTuple
import torch
import torch.nn as nn
from labml import experiment, tracker, monit, logger
from labml.configs import option
from labml.logger import Text
from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs
from labml_nn.helpers.metrics import SimpleStateModule
from labml_nn.helpers.trainer import BatchIndex
from labml_nn.transformers.compressive import CompressiveTransformer, AttentionReconstructionLoss, \
CompressiveTransformerLayer, Conv1dCompression
class CompressedMemory(NamedTuple):
mem: List[torch.Tensor]
c_mem: List[torch.Tensor]
class AutoregressiveModel(nn.Module):
"""
## Auto regressive model
"""
def __init__(self, n_vocab: int, d_model: int, transformer: CompressiveTransformer):
super().__init__()
# Token embedding module
self.src_embed = nn.Embedding(n_vocab, d_model)
# Transformer
self.transformer = transformer
# Final layer
self.generator = nn.Linear(d_model, n_vocab)
# Masks
self.mask_x = None
self.mask_mem = None
def forward(self, x: torch.Tensor, mem: CompressedMemory):
# Get memory and compressed memory
if mem is not None:
mem, c_mem = mem.mem, mem.c_mem
else:
mem = []
c_mem = []
# Total length of the memory and compressed memory (for masks)
m_len = len(mem[0]) if mem else 0
if c_mem:
m_len += len(c_mem[0])
# Create a subsequent mask for tokens
if self.mask_x is None or self.mask_x.shape[0] < len(x):
from labml_nn.transformers.utils import subsequent_mask
self.mask_x = subsequent_mask(len(x)).to(x.device)
# Create an all ones (full visibility) mask for memory
if self.mask_mem is None or self.mask_mem.shape[1] < m_len or self.mask_mem.shape[0] < len(x):
self.mask_mem = self.mask_x.new_ones(len(x), m_len, 1)
# Concatenate the masks if there is memory
if m_len:
mask = torch.cat((self.mask_mem[:len(x), :m_len], self.mask_x[:len(x), :len(x)]), dim=1)
# Use only the subsequent mask otherwise
else:
mask = self.mask_x[:len(x), :len(x)]
# Token embeddings
x = self.src_embed(x)
# Run it through the transformer
res, mem = self.transformer(x, mem, c_mem, mask)
# Generate logits of the next token
res = self.generator(res)
#
return res, mem
class Configs(NLPAutoRegressionConfigs):
"""
## Configurations
The default configurations can and will be overridden when we start the experiment.
"""
model: AutoregressiveModel
# Token embedding size
d_model: int = 128
# Number of attention heads
heads: int = 4
# Dropout probability
dropout: float = 0.0
# Number of features in FFN hidden layer
d_ff: int = 256
# Number of transformer layers
n_layers: int = 6
# Number of memories to keep
mem_len: int = 8
# State module to maintain memories when switching between training and validation
memory = SimpleStateModule()
# Attention Reconstruction Loss
attention_reconstruction_loss: AttentionReconstructionLoss
# Compression rate
compression_rate: int = 4
# Compressed memory length
c_mem_len: int = 128
def init(self):
# Set tracker configurations
tracker.set_scalar("accuracy.*", True)
tracker.set_scalar("loss.*", True)
# Do not print the attention reconstruction loss in the terminal
tracker.set_scalar("ar_loss.*", False)
# This will keep the accuracy metric stats and memories separate for training and validation.
self.state_modules = [self.accuracy, self.memory]
@torch.no_grad()
def merge_compress_memory(self, mem: CompressedMemory, new_mem: List[torch.Tensor]) \
-> Tuple[CompressedMemory, List[torch.Tensor]]:
"""
Concatenate new memories and compress the oldest memories.
"""
# If the configurations specify not to use memory
if self.mem_len == 0 and self.c_mem_len == 0:
return CompressedMemory([], []), []
# Get memory and compressed memory
if mem is not None:
mem, c_mem = mem.mem, mem.c_mem
else:
mem, c_mem = [], []
# Concatenate new memories with old memory
if mem:
mem = [torch.cat((m, x), dim=0) for m, x in zip(mem, new_mem)]
else:
mem = new_mem
# Compress the oldest memories if there are more memories than `mem_len`
if len(mem[0]) > self.mem_len:
# Calculate the number of compressed memories to make $n_{cm} = \bigg\lceil\frac{n'_m - N_m}{c}\bigg\rceil$,
# where $n'_m$ is the number of memories we have
# and $N_m$ is the maximum number of memories we maintain (`mem_len`).
n_c_mem = (len(mem[0]) - self.mem_len + self.compression_rate - 1) // self.compression_rate
# Number of memories to compress $c n_{cm}$
n_old = n_c_mem * self.compression_rate
# A list to keep memories that need to be compressed for each layer.
mem_to_compress = []
# A list to keep the memories that do not get compressed for each layer.
uncompressed_mem = []
# Iterate through memories of each layer.
for m in mem:
# Split the memories at $c n_{cm}$
cm, m = torch.split(m, [n_old, len(m) - n_old])
# Collect memories to compress
mem_to_compress.append(cm)
# Collect remaining memories
uncompressed_mem.append(m)
# Update the memories
mem = uncompressed_mem
# Compress the memories
new_c_mem = []
for i, layer in enumerate(self.model.transformer.layers):
new_c_mem.append(layer.compress(mem_to_compress[i]))
# Concatenate newly compressed memories with old compressed memories
if c_mem:
c_mem = [torch.cat((m, nm), dim=0) for m, nm in zip(c_mem, new_c_mem)]
# If there are no old compressed memories
else:
c_mem = new_c_mem
# Truncate old memories
if len(c_mem[0]) > self.c_mem_len:
c_mem = [m[-self.c_mem_len:] for m in c_mem]
# No memories are compressed if the number of memories is less than `mem_len`
else:
mem_to_compress = []
# Return memories and the memories that were compressed.
# Memories that were compressed are needed for the reconstruction loss computation.
return CompressedMemory(mem, c_mem), mem_to_compress
def step(self, batch: any, batch_idx: BatchIndex):
"""
### Training/validation step
"""
# Move data to the device
data, target = batch[0].to(self.device), batch[1].to(self.device)
# Update global step (number of tokens processed) when in training mode
if self.mode.is_train:
tracker.add_global_step(data.shape[0] * data.shape[1])
# Get memories
mem = self.memory.get()
# Run the model
output, new_mem = self.model(data, mem)
# Merge and compress memory
mem, mem_to_compress = self.merge_compress_memory(mem, new_mem)
# Update memories
self.memory.set(mem)
# Calculate and log cross entropy loss
loss = self.loss_func(output, target)
tracker.add("loss.", loss)
# Calculate attention reconstruction loss if memories were compressed in this step
if mem_to_compress:
# Get attention reconstruction loss
ar_loss = self.attention_reconstruction_loss(new_mem, mem_to_compress)
# Track attention reconstruction loss
tracker.add("ar_loss.", ar_loss)
# Add attention reconstruction loss to loss
loss = loss + ar_loss
# Calculate and log accuracy
self.accuracy(output, target)
self.accuracy.track()
# Train the model
if self.mode.is_train:
# Calculate gradients
loss.backward()
# Clip gradients
torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=self.grad_norm_clip)
# Take optimizer step
self.optimizer.step()
# Log the model parameters and gradients on last batch of every epoch
if batch_idx.is_last:
tracker.add('model', self.model)
# Clear the gradients
self.optimizer.zero_grad()
# Save the tracked metrics
tracker.save()
def sample(self):
"""
### Sampling function to generate samples periodically while training
"""
# Starting prompt
prompt = self.prompt
# Collect output for printing
log = [(prompt, Text.subtle)]
# memory
mem = CompressedMemory([], [])
# Sample 25 tokens
for i in monit.iterate('Sample', 25):
# Tokenize the prompt
data = self.text.text_to_i(prompt).unsqueeze(-1)
# Move to device
data = data.to(self.device)
# Get the model output
output, new_mem = self.model(data, mem)
# Get the model prediction (greedy)
output = output.argmax(dim=-1).squeeze(1)
# Add the prediction to prompt
prompt += self.prompt_separator + self.text.itos[output[-1]]
# Only feed the last character to model in next iteration, rest will go in as memories
prompt = prompt[-1:]
# Add the prediction for logging
log += [(self.prompt_separator + self.text.itos[output[-1]], Text.value)]
# Update and compress memory
mem, _ = self.merge_compress_memory(mem, new_mem)
# Print the sampled output
logger.log(log)
@option(Configs.model)
def autoregressive_model(c: Configs):
"""
### Initialize the auto-regressive model
"""
from labml_nn.transformers.xl import RelativeMultiHeadAttention
from labml_nn.transformers.feed_forward import FeedForward
m = AutoregressiveModel(c.n_tokens, c.d_model, CompressiveTransformer(
CompressiveTransformerLayer(d_model=c.d_model,
self_attn=RelativeMultiHeadAttention(c.heads, c.d_model, c.dropout),
feed_forward=FeedForward(c.d_model, c.d_ff, c.dropout),
dropout_prob=c.dropout,
compress=Conv1dCompression(c.compression_rate, c.d_model)), c.n_layers))
return m.to(c.device)
@option(Configs.attention_reconstruction_loss)
def attention_reconstruction_loss(c: Configs):
"""
### Initialize the attention reconstruction loss
"""
return AttentionReconstructionLoss(c.model.transformer.layers)
def main():
"""
### Run the experiment
"""
# Create experiment
experiment.create(name="compressive_transformer", comment='')
# Create configs
conf = Configs()
# Load configurations
experiment.configs(conf,
# A dictionary of configurations to override
{'tokenizer': 'character',
'text': 'tiny_shakespeare',
'optimizer.learning_rate': 2.5e-4,
'optimizer.optimizer': 'AdamW',
'prompt': 'It is',
'prompt_separator': '',
'train_loader': 'sequential_train_loader',
'valid_loader': 'sequential_valid_loader',
'seq_len': 8,
'mem_len': 8,
'epochs': 128,
'batch_size': 32,
'inner_iterations': 25,
'compression_rate': 2,
})
# Set models for saving and loading
experiment.add_pytorch_models({'model': conf.model})
# Start the experiment
with experiment.start():
# `TrainValidConfigs.run`
conf.run()
#
if __name__ == '__main__':
main()
labml_nn/transformers/compressive/readme.md
+42
View File
@@ -0,0 +1,42 @@
# [Compressive Transformer](https://nn.labml.ai/transformers/compressive/index.html)
This is an implementation of
[Compressive Transformers for Long-Range Sequence Modelling](https://arxiv.org/abs/1911.05507)
in [PyTorch](https://pytorch.org).
This is an extension of [Transformer XL](https://nn.labml.ai/transformers/xl/index.html) where past memories
are compressed to give a longer attention range.
That is, the furthest $n_{cm} c$ memories are compressed into
$n_{cm}$ memories, where $c$ is the compression rate.
## Compression operation
The compression operation is defined as
$f_c: \mathbb{R}^{nc \times d} \rightarrow \mathbb{R}^{n \times d}$.
The paper introduces multiple choices for $f_c$ and we have only implemented
1D convolution which seems to give the best results.
Each layer has a separate compression operation $f_c^{(i)}$ where
$i$ is the layer number.
## Training compression operation
Since training compression with BPTT requires maintaining
a very large computational graph (many time steps), the paper proposes
an *auto-encoding loss* and an *attention reconstruction loss*.
The auto-encoding loss decodes the original memories from the compressed memories
and calculates the loss.
Attention reconstruction loss computes the multi-headed attention results
on the compressed memory and on uncompressed memory and gets a mean squared error
between them.
We have implemented the latter here since it gives better results.
This implementation uses pre-layer normalization
while the paper uses post-layer normalization.
Pre-layer norm does the layer norm before [FFN](../feedforward.html) and
self-attention, and the pass-through in the residual connection is not normalized.
This is supposed to be more stable in standard transformer setups.
Here are [the training code](https://nn.labml.ai/transformers/compressive/experiment.html) and a notebook for training a compressive transformer
model on the Tiny Shakespeare dataset.
[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/transformers/compressive/experiment.ipynb)
labml_nn/transformers/configs.py
+325
View File
@@ -0,0 +1,325 @@
"""
---
title: Configurable Transformer Components
summary: These are configurable components that can be re-used quite easily.
---
# Configurable Transformer Components
"""
import copy
import torch.nn as nn
from labml.configs import BaseConfigs, option, calculate, aggregate
from .feed_forward import FeedForward
from .mha import MultiHeadAttention
from .models import EmbeddingsWithPositionalEncoding, EmbeddingsWithLearnedPositionalEncoding, TransformerLayer, \
Encoder, Decoder, Generator, EncoderDecoder
class FeedForwardConfigs(BaseConfigs):
"""
<a id="FFN"></a>
## FFN Configurations
Creates a Position-wise FeedForward Network defined in
[`feed_forward.py`](feed_forward.html).
"""
# Position-wise feedforward layer
ffn: FeedForward
# Number of features in the embedding
d_model: int
# Number of features in in the hidden layer
d_ff: int = 2048
# Dropout probability
dropout: float = 0.1
# Activation in position-wise feedforward layer
activation: nn.Module = 'ReLU'
# Whether the FFN layer should be gated
is_gated: bool = False
# Whether the first fully connected layer should have a learnable bias
bias1: bool = True
# Whether the second fully connected layer should have a learnable bias
bias2: bool = True
# Whether the fully connected layer for the gate should have a learnable bias
bias_gate: bool = False
# Predefined GLU variants
glu_variant: str = 'none'
@option(FeedForwardConfigs.activation, 'ReLU')
def _ffn_activation_relu():
"""
### ReLU activation
$$\max(0, x)$$
"""
return nn.ReLU()
@option(FeedForwardConfigs.activation, 'GELU')
def _ffn_activation_gelu():
"""
### GELU activation
$$x \Phi(x)$$ where $\Phi(x) = P(X \le x), X \sim \mathcal{N}(0,1)$
It was introduced in paper [Gaussian Error Linear Units](https://arxiv.org/abs/1606.08415).
"""
return nn.GELU()
@option(FeedForwardConfigs.ffn, 'default')
def _feed_forward(c: FeedForwardConfigs):
"""
Initialize a [feed forward network](feed_forward.html)
"""
return FeedForward(c.d_model, c.d_ff,
dropout=c.dropout,
activation=c.activation,
is_gated=c.is_gated,
bias1=c.bias1,
bias2=c.bias2,
bias_gate=c.bias_gate)
# ## GLU Variants
# These are variants with gated hidden layers for the FFN
# as introduced in paper [GLU Variants Improve Transformer](https://arxiv.org/abs/2002.05202).
# We have omitted the bias terms as specified in the paper.
# ### FFN with Gated Linear Units
#
# $$FFN_{GLU}(x)(x, W_1, V, W_2) = (\sigma(x W_1) \otimes x V) W_2$$
aggregate(FeedForwardConfigs.glu_variant, 'GLU',
(FeedForwardConfigs.is_gated, True),
(FeedForwardConfigs.bias1, False),
(FeedForwardConfigs.bias2, False),
(FeedForwardConfigs.bias_gate, False),
(FeedForwardConfigs.activation, nn.Sigmoid()))
# ### FFN with Bilinear hidden layer
#
# $$FFN_{Bilinear}(x)(x, W_1, V, W_2) = (x W_1 \otimes x V) W_2$$
aggregate(FeedForwardConfigs.glu_variant, 'Bilinear',
(FeedForwardConfigs.is_gated, True),
(FeedForwardConfigs.bias1, False),
(FeedForwardConfigs.bias2, False),
(FeedForwardConfigs.bias_gate, False),
(FeedForwardConfigs.activation, nn.Identity()))
# ### FFN with ReLU gate
#
# $$FFN_{ReGLU}(x)(x, W_1, V, W_2) = (\max(0, x W_1) \otimes x V) W_2$$
aggregate(FeedForwardConfigs.glu_variant, 'ReGLU',
(FeedForwardConfigs.is_gated, True),
(FeedForwardConfigs.bias1, False),
(FeedForwardConfigs.bias2, False),
(FeedForwardConfigs.bias_gate, False),
(FeedForwardConfigs.activation, nn.ReLU()))
# ### FFN with GELU gate
#
# $$FFN_{GEGLU}(x)(x, W_1, V, W_2) = (\text{GELU}(x W_1) \otimes x V) W_2$$
aggregate(FeedForwardConfigs.glu_variant, 'GEGLU',
(FeedForwardConfigs.is_gated, True),
(FeedForwardConfigs.bias1, False),
(FeedForwardConfigs.bias2, False),
(FeedForwardConfigs.bias_gate, False),
(FeedForwardConfigs.activation, nn.GELU()))
# ### FFN with Swish gate
#
# $$FFN_{SwiGLU}(x)(x, W_1, V, W_2) = (\text{Swish}_1(x W_1) \otimes x V) W_2$$
# where $\text{Swish}_\beta(x) = x \sigma(\beta x)$
aggregate(FeedForwardConfigs.glu_variant, 'SwiGLU',
(FeedForwardConfigs.is_gated, True),
(FeedForwardConfigs.bias1, False),
(FeedForwardConfigs.bias2, False),
(FeedForwardConfigs.bias_gate, False),
(FeedForwardConfigs.activation, nn.SiLU()))
class TransformerConfigs(BaseConfigs):
"""
<a id="TransformerConfigs"></a>
## Transformer Configurations
This defines configurations for a transformer.
The configurations are calculate using option functions.
These are lazy loaded and therefore only the necessary modules
are calculated.
"""
# Number of attention heads
n_heads: int = 8
# Transformer embedding size
d_model: int = 512
# Number of layers
n_layers: int = 6
# Dropout probability
dropout: float = 0.1
# Number of tokens in the source vocabulary (for token embeddings)
n_src_vocab: int
# Number of tokens in the target vocabulary (to generate logits for prediction)
n_tgt_vocab: int
# The encoder self attention
encoder_attn: MultiHeadAttention = 'mha'
# The decoder self attention
decoder_attn: MultiHeadAttention = 'mha'
# The decoder memory attention
decoder_mem_attn: MultiHeadAttention = 'mha'
# Configurable Feedforward Layer
ffn: FeedForwardConfigs
# Encoder layer
encoder_layer: TransformerLayer = 'default'
# Decoder layer
decoder_layer: TransformerLayer = 'default'
# Encoder consisting of multiple encoder layers
encoder: Encoder = 'default'
# Encoder consisting of multiple decoder layers
decoder: Decoder = 'default'
# Embedding layer for source
src_embed: nn.Module = 'fixed_pos'
# Embedding layer for target (for decoder)
tgt_embed: nn.Module = 'fixed_pos'
# Logit generator for prediction
generator: Generator = 'default'
# Encoder-decoder
encoder_decoder: EncoderDecoder
# ### Multi-head Attention
def _mha(c: TransformerConfigs):
return MultiHeadAttention(c.n_heads, c.d_model, dropout_prob=c.dropout)
calculate(TransformerConfigs.encoder_attn, 'mha', _mha)
calculate(TransformerConfigs.decoder_attn, 'mha', _mha)
calculate(TransformerConfigs.decoder_mem_attn, 'mha', _mha)
# ### Relative Multi-head Attention
def _relative_mha(c: TransformerConfigs):
from labml_nn.transformers.xl.relative_mha import RelativeMultiHeadAttention
return RelativeMultiHeadAttention(c.n_heads, c.d_model)
calculate(TransformerConfigs.encoder_attn, 'relative', _relative_mha)
calculate(TransformerConfigs.decoder_attn, 'relative', _relative_mha)
calculate(TransformerConfigs.decoder_mem_attn, 'relative', _relative_mha)
@option(TransformerConfigs.ffn, 'default')
def _feed_forward(c: TransformerConfigs):
"""
Create feedforward layer configurations
"""
conf = FeedForwardConfigs()
conf.set_default(FeedForwardConfigs.d_model, func=lambda: c.d_model)
conf.set_default(FeedForwardConfigs.dropout, func=lambda: c.dropout)
return conf
@option(TransformerConfigs.encoder_layer, 'default')
def _encoder_layer(c: TransformerConfigs):
"""
Encoder layer
"""
return TransformerLayer(d_model=c.d_model, self_attn=c.encoder_attn,
src_attn=None, feed_forward=copy.deepcopy(c.ffn.ffn),
dropout_prob=c.dropout)
@option(TransformerConfigs.decoder_layer, 'default')
def _decoder_layer(c: TransformerConfigs):
"""
Decoder layer
"""
return TransformerLayer(d_model=c.d_model, self_attn=c.decoder_attn,
src_attn=c.decoder_mem_attn, feed_forward=copy.deepcopy(c.ffn.ffn),
dropout_prob=c.dropout)
@option(TransformerConfigs.encoder, 'default')
def _encoder(c: TransformerConfigs):
"""
Encoder
"""
return Encoder(c.encoder_layer, c.n_layers)
@option(TransformerConfigs.decoder, 'default')
def _decoder(c: TransformerConfigs):
"""
Decoder
"""
return Decoder(c.decoder_layer, c.n_layers)
@option(TransformerConfigs.generator, 'default')
def _generator(c: TransformerConfigs):
"""
Logit generator
"""
return Generator(c.n_tgt_vocab, c.d_model)
# ### Fixed Positional Embeddings
@option(TransformerConfigs.src_embed, 'fixed_pos')
def _src_embed_with_positional(c: TransformerConfigs):
"""
Source embedding with fixed positional encodings
"""
return EmbeddingsWithPositionalEncoding(c.d_model, c.n_src_vocab)
@option(TransformerConfigs.tgt_embed, 'fixed_pos')
def _tgt_embed_with_positional(c: TransformerConfigs):
"""
Target embedding with fixed positional encodings
"""
return EmbeddingsWithPositionalEncoding(c.d_model, c.n_tgt_vocab)
# ### Learned Positional Embeddings
@option(TransformerConfigs.src_embed, 'learned_pos')
def _src_embed_with_learned_positional(c: TransformerConfigs):
"""
Source embedding with learned positional encodings
"""
return EmbeddingsWithLearnedPositionalEncoding(c.d_model, c.n_src_vocab)
@option(TransformerConfigs.tgt_embed, 'learned_pos')
def _tgt_embed_with_learned_positional(c: TransformerConfigs):
"""
Target embedding with learned positional encodings
"""
return EmbeddingsWithLearnedPositionalEncoding(c.d_model, c.n_tgt_vocab)
# ### No Positional Embeddings
@option(TransformerConfigs.src_embed, 'no_pos')
def _src_embed_without_positional(c: TransformerConfigs):
"""
Source embedding without positional encodings
"""
return nn.Embedding(c.n_src_vocab, c.d_model)
@option(TransformerConfigs.tgt_embed, 'no_pos')
def _tgt_embed_without_positional(c: TransformerConfigs):
return nn.Embedding(c.n_tgt_vocab, c.d_model)
@option(TransformerConfigs.encoder_decoder, 'default')
def _encoder_decoder(c: TransformerConfigs):
return EncoderDecoder(c.encoder, c.decoder, c.src_embed, c.tgt_embed, c.generator)
labml_nn/transformers/fast_weights/__init__.py
+327
View File
@@ -0,0 +1,327 @@
"""
---
title: Linear Transformers Are Secretly Fast Weight Memory Systems
summary: >
This is an annotated implementation/tutorial of
Linear Transformers Are Secretly Fast Weight Memory Systems in PyTorch.
---
# Fast weights transformer
The paper
[Linear Transformers Are Secretly Fast Weight Memory Systems in PyTorch](https://arxiv.org/abs/2102.11174)
finds similarities between linear self-attention and fast weight systems
and makes modifications to self-attention update rule based on that.
It also introduces a simpler, yet effective kernel function.
*The authors have provided an [official implementation](https://github.com/ischlag/fast-weight-transformers)
of the paper including other variants they compare with in the paper.*
## Fast weights
Consider a sequence of inputs $\big\{x^{(i)}\big\}^L_{i=1}$ or length $L$
and each step is a vector of size $d_{in}$; i.e. $x \in \mathbb{R}^{d_{in}}$.
The fast weight model generates a weight matrix at each step to produce output
$\big\{y^{(i)}\big\}^L_{i=1}$, $y \in \mathbb{R}^{d_{out}}$
\begin{align}
a^{(i)}, b^{(i)} &= \textcolor{orange}{W_a} x^{(i)}, \textcolor{orange}{W_b} x^{(i)} \\
\textcolor{cyan}{W^{(i)}} &= \sigma \Big( \textcolor{cyan}{W^{(i-1)}} + a^{(i)} \otimes b^{(i)} \Big) \\
y^{(i)} &= \textcolor{cyan}{W^{(i)}} x^{(i)}
\end{align}
$\otimes$ is the outer product ($a \otimes b = a b^\top$), where elements of the two vectors are multiplied with each other
to give a matrix.
$\sigma$ is an activation function.
$\textcolor{orange}{W_a}$ and $\textcolor{orange}{W_b}$ are trainable weights (parameters).
$\textcolor{cyan}{W^{(i)}}$ are the fast weights that are generated at each step.
## Linear self-attention
Original transformer self-attention is, (omitting $\frac{1}{d_k}$ for clarity)
\begin{align}
y^{(i)} &= \Big[v^{(1)}, v^{(2)}, ..., v^{(i)}\Big] \text{softmax}
\bigg(
\Big[k^{(1)}, k^{(2)}, ..., k^{(i)}\Big] ^\top
q^{(i)}
\bigg) \\
&= \sum^i_{j=1} \frac
{ v^{(j)} \kappa(k^{(j)}, q^{(i)}) }
{ \sum^i_{j'=1} \kappa(k^{(j')}, q^{(i)}) } \\
\end{align}
where $\kappa(k, q) = \text{exp}(k \cdot q)$
The idea behind linearizing self attention is to replace softmax
kernel $\kappa$ with a different kernel $\kappa '$ so that we can calculate the
denominator of the self attention function faster:
$$\kappa '(k, q) = \textcolor{lightgreen}{\phi(k)}^\top \textcolor{lightgreen}{\phi(q)}$$
This gives
\begin{align}
y^{(i)} &= \frac
{\Big( \sum^i_{j=1} v^{(j)} \otimes \textcolor{lightgreen}{\phi(k^{(j)})} \Big)
\textcolor{lightgreen}{\phi(q^{(i)})} }
{ \Big( \sum^i_{j'=1}
\textcolor{lightgreen}{\phi(k^{(j')})} \Big)
\textcolor{lightgreen}{\phi(q^{(i)})} }
\end{align}
With $\textcolor{cyan}{W^{(i)}} = \sum^i_{j=1} v^{(j)} \otimes \phi(k^{(j)})$ and
$z^{(i)} = \sum^i_{j=1} \textcolor{lightgreen}{\phi(k^{(j)})}$, we can calculate them efficiently:
\begin{align}
\textcolor{cyan}{W^{(i)}} &= \textcolor{cyan}{W^{(i-1)}} + v^{(i)} \otimes \textcolor{lightgreen}{\phi(k^{(i)})} \\
z^{(i)} &= z{(i)} + \textcolor{lightgreen}{\phi(k^{(i)})} \\
y^{(i)} &= \frac{1}{z^{(i)} \cdot \textcolor{lightgreen}{\phi(q^{(i)})}}
W^{(i)} \textcolor{lightgreen}{\phi(q^{(i)})}
\end{align}
This is quite similar to fast weights.
The paper introduces a new linear attention projection function $\textcolor{lightgreen}{\phi}$
a new update rule for $\textcolor{cyan}{W^{(i)}} = f(\textcolor{cyan}{W^{(i-1)}})$ and change the normalization
$\frac{1}{z^{(i)} \cdot \textcolor{lightgreen}{\phi(q^{(i)})}}$
Here are [the training code](experiment.html) and a notebook for training a fast weights
transformer on the Tiny Shakespeare dataset.
[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/transformers/fast_weights/experiment.ipynb)
"""
import torch
from torch import nn
from labml_nn.transformers.feed_forward import FeedForward
from labml_nn.transformers.mha import PrepareForMultiHeadAttention
from labml_nn.utils import clone_module_list
class DPFP(nn.Module):
"""
## Deterministic Parameter Free Project (DPFP)
This is the new projection function $\textcolor{lightgreen}{\phi}$ introduced in the paper.
DPFP projects $k$ of dimensionality $d_{key}$ to dimensionality $d_{dot} = 2 d_{key} \nu$,
where $\nu \in \\{1, 2, ..., 2 d_{key} - 1 \\}$ is a hyper-parameter.
$$\textcolor{lightgreen}{\phi_{2 d_{key} (i - 1) + j}(k)}
= \text{ReLU}\Big(\big[k, -k\big]\Big)_{j}
\text{ReLU}\Big(\big[k, -k\big]\Big)_{i + j}$$
where $\big[k, -k\big]$ is the concatenation of $k$ and $-k$ to give a vector of
size $2 d_{key}$, $i \in \\{1, 2, ..., \nu \\}$, and $j \in \\{1, 2, ..., 2 d_{key}\\}$.
$x_i$ is the $i$-th element of vector $x$ and is rolled around if
$i$ is larger than the number of elements in $x$.
Basically, it creates a new vector by multiplying elements of $[k, -k]$ shifted by $i$.
This produces projections that are sparse (only a few elements of $phi$ are non-zero) and
orthogonal ($\textcolor{lightgreen}{\phi(k^{(i)})} \cdot \textcolor{lightgreen}{\phi(k^{(j)})}
\approx 0$ for most $i, j$
unless $k^{(i)}$ and $k^{(j)}$ are very similar.
### Normalization
Paper introduces a simple normalization for $\textcolor{lightgreen}{\phi}$,
$$\textcolor{lightgreen}{\phi '(k)} =
\frac{\textcolor{lightgreen}{\phi(k)}}{\sum^{d_{dot}}_{j=1} \textcolor{lightgreen}{\phi(k)_j}}$$
*Check the paper for derivation.*
"""
def __init__(self, nu: int = 1, eps: float = 1e-6):
"""
* `nu` is the hyper-parameter $\nu$.
* `eps` is the small value used to make sure there is no division-by-zero when normalizing.
"""
super().__init__()
self.nu = nu
self.relu = nn.ReLU()
self.eps = eps
def forward(self, k: torch.Tensor):
# Get $\textcolor{lightgreen}{\phi(k)}$
k = self.dpfp(k)
# Normalize by $\sum^{d_{dot}}_{j=1} \textcolor{lightgreen}{\phi(k)_j}$
return k / (torch.sum(k, dim=-1, keepdim=True) + self.eps)
def dpfp(self, k: torch.Tensor):
"""
$$\textcolor{lightgreen}{\phi(k)}$$
"""
# $x = \text{ReLU}\Big(\big[k, -k\big]\Big)$
x = self.relu(torch.cat([k, -k], dim=-1))
# Shift and roll by $i \in \\{1, 2, ..., \nu \\}$,
# to get $$x'_{i,j} = \text{ReLU}\Big(\big[k, -k\big]\Big)_{i+j}$$
x_rolled = [x.roll(shifts=i, dims=-1) for i in range(1, self.nu + 1)]
# Concatenate to get
# $$x'_{2 d_{key} (i - 1) + j} = \text{ReLU}\Big(\big[k, -k\big]\Big)_{i+j}$$
x_rolled = torch.cat(x_rolled, dim=-1)
# Concatenate copies of $x$
x_repeat = torch.cat([x] * self.nu, dim=-1)
# Multiply them,
# $$\textcolor{lightgreen}{\phi_{2 d_{key} (i - 1) + j}(k)}
# = \text{ReLU}\Big(\big[k, -k\big]\Big)_{j}
# \text{ReLU}\Big(\big[k, -k\big]\Big)_{i + j}$$
return x_repeat * x_rolled
class FastWeightsAttention(nn.Module):
"""
## Fast Weights Attention
The paper introduces a new update rule for calculating $\textcolor{cyan}{W^{(i)}}$.
The model first retrieves the current value
$\bar{v}^{(i)}$ paired with the key $k^{(i)}$.
Then stores a combination $v^{(i)}_{new}$
of the retrieved value $\bar{v}^{(i)}$ and the input $v^{(i)}$.
\begin{align}
k^{(i)}, v^{(i)}, q^{(i)} &=
\textcolor{orange}{W_k} x^{(i)}, \textcolor{orange}{W_v} x^{(i)}, \textcolor{orange}{W_q} x^{(i)} \\
\bar{v}^{(i)} &= \textcolor{cyan}{W^{(i-1)}} \textcolor{lightgreen}{\phi'(k^{(i)})} \\
\beta^{(i)} &= \sigma \Big(\textcolor{orange}{W_\beta} x^{(i)} \Big) \\
v^{(i)}_{new} &= \beta^{(i)} v^{(i)} + \Big(1 - \beta^{(i)} \Big) \bar{v}^{(i)} \\
\textcolor{cyan}{W^{(i)}}
&= \textcolor{cyan}{W^{(i-1)}} + v^{(i)}_{new} \otimes \textcolor{lightgreen}{\phi'(k^{(i)})} \\
&= \textcolor{cyan}{W^{(i-1)}} +
\beta^{(i)} \Big( v^{(i)} - \bar{v}^{(i)} \Big ) \otimes \textcolor{lightgreen}{\phi'(k^{(i)})} \\
y^{(i)} &= \textcolor{cyan}{W^{(i)}} \textcolor{lightgreen}{\phi'(q^{(i)})}
\end{align}
where $\textcolor{orange}{W_\beta}$ is a trainable parameter and $\sigma$ is the sigmoid function.
Note that we don't need the normalization term $z$ because $\textcolor{lightgreen}{\phi'}$ is normalized.
"""
def __init__(self, heads: int, d_model: int, dropout_prob: float, phi: DPFP):
super().__init__()
# Number of features per head $d_k$
self.d_k = d_model // heads
# Number of heads
self.heads = heads
# These transform the `query`, `key` and `value` multi-headed attention.
self.query = PrepareForMultiHeadAttention(d_model, heads, self.d_k, bias=False)
self.key = PrepareForMultiHeadAttention(d_model, heads, self.d_k, bias=False)
self.value = PrepareForMultiHeadAttention(d_model, heads, self.d_k, bias=False)
# Interpolation weight function $\sigma \Big(\textcolor{orange}{W_\beta} x^{(i)} \Big)$ for each head
self.interpolation_weight = nn.Sequential(
PrepareForMultiHeadAttention(d_model, heads, 1, bias=False),
nn.Sigmoid()
)
# $\textcolor{lightgreen}{\phi'}$
self.phi = phi
# Output layer
self.output = nn.Linear(d_model, d_model)
# Dropout
self.dropout = nn.Dropout(dropout_prob)
def forward(self, x: torch.Tensor):
# Get the number of steps $L$
seq_len = x.shape[0]
# $\textcolor{lightgreen}{\phi'(q^{(i)})}$ for all steps and heads
query = self.phi(self.query(x))
# $\textcolor{lightgreen}{\phi'(k^{(i)})}$ for all steps and heads
key = self.phi(self.key(x))
# $v^{(i)}$ for all steps and heads
value = self.value(x)
# $\beta^{(i)}$ for all steps and heads
beta = self.interpolation_weight(x)
# $\textcolor{cyan}{W^{(0)}}$
weights = key.new_zeros((key.shape[1], key.shape[2], value.shape[3], key.shape[3]))
# List to store outputs $y^{(i)}$
outputs = []
# Iterate through steps
for i in range(seq_len):
# $$\bar{v}^{(i)} = \textcolor{cyan}{W^{(i-1)}} \textcolor{lightgreen}{\phi'(k^{(i)})}$$
value_existing = torch.einsum('bhvk,bhk->bhv', weights, key[i])
# $$\textcolor{cyan}{W^{(i)}}
# = \textcolor{cyan}{W^{(i-1)}} +
# \beta^{(i)} \Big( v^{(i)} - \bar{v}^{(i)} \Big ) \otimes \textcolor{lightgreen}{\phi'(k^{(i)})}$$
weights = weights + torch.einsum('bhv,bhk->bhvk', beta[i] * (value[i] - value_existing), key[i])
# $$y^{(i)} = \textcolor{cyan}{W^{(i)}} \textcolor{lightgreen}{\phi'(q^{(i)})}$$
y = torch.einsum('bhvk,bhk->bhv', weights, query[i])
# Merge multiple heads and append to `outputs`
outputs.append(y.reshape(y.shape[0], -1))
# Stack outputs at each step into a single tensor
x = torch.stack(outputs)
# Output layer
return self.output(x)
class FastWeightsAttentionTransformerLayer(nn.Module):
"""
This is a general transformer layer that combines self attention and feedforward network.
"""
def __init__(self, *,
d_model: int,
attn: FastWeightsAttention,
feed_forward: FeedForward,
dropout_prob: float):
super().__init__()
# Transformer size $d_{model}$
self.size = d_model
# Fast weights attention module
self.attn = attn
# Feed-forward network
self.feed_forward = feed_forward
# Dropout layer
self.dropout = nn.Dropout(dropout_prob)
# Normalization layers
self.norm_self_attn = nn.LayerNorm([d_model])
self.norm_ff = nn.LayerNorm([d_model])
def forward(self, x: torch.Tensor):
# Calculate fast weights self attention
attn = self.attn(x)
# Add the self attention results
x = x + self.dropout(attn)
# Normalize for feed-forward
z = self.norm_ff(x)
# Pass through the feed-forward network
ff = self.feed_forward(z)
# Add the feed-forward results back
x = x + self.dropout(ff)
#
return x
class FastWeightsAttentionTransformer(nn.Module):
"""
This is a general transformer module with multiple transformer layers
"""
def __init__(self, layer: FastWeightsAttentionTransformerLayer, n_layers: int):
super().__init__()
# Make copies of the transformer layer
self.layers = clone_module_list(layer, n_layers)
# Final normalization layer
self.norm = nn.LayerNorm([layer.size])
def forward(self, x: torch.Tensor):
for i, layer in enumerate(self.layers):
# Get layer output
x = layer(x)
# Normalize the output
return self.norm(x)
labml_nn/transformers/fast_weights/experiment.ipynb
+219
View File
@@ -0,0 +1,219 @@
{
"nbformat": 4,
"nbformat_minor": 0,
"metadata": {
"colab": {
"name": "Fast Weights Transformer",
"provenance": [],
"collapsed_sections": []
},
"kernelspec": {
"name": "python3",
"display_name": "Python 3"
},
"accelerator": "GPU"
},
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "AYV_dMVDxyc2"
},
"source": [
"[![Github](https://img.shields.io/github/stars/labmlai/annotated_deep_learning_paper_implementations?style=social)](https://github.com/labmlai/annotated_deep_learning_paper_implementations)\n",
"[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/transformers/fast_weights/experiment.ipynb) \n",
"\n",
"## Fast Weights Transformer\n",
"\n",
"This is an experiment training Shakespeare dataset with a Compressive Transformer model."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "AahG_i2y5tY9"
},
"source": [
"Install the `labml-nn` package"
]
},
{
"cell_type": "code",
"metadata": {
"id": "ZCzmCrAIVg0L",
"colab": {
"base_uri": "https://localhost:8080/"
},
"outputId": "b73a038d-62dc-48e0-c2f1-9454224e9137"
},
"source": [
"!pip install labml-nn"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "SE2VUQ6L5zxI"
},
"source": [
"Imports"
]
},
{
"cell_type": "code",
"metadata": {
"id": "0hJXx_g0wS2C"
},
"source": [
"from labml import experiment\n",
"from labml_nn.transformers.fast_weights.experiment import Configs"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "Lpggo0wM6qb-"
},
"source": [
"Create an experiment"
]
},
{
"cell_type": "code",
"metadata": {
"id": "bFcr9k-l4cAg"
},
"source": [
"experiment.create(name=\"fast_weights_transformer\")"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "-OnHLi626tJt"
},
"source": [
"Initialize configurations"
]
},
{
"cell_type": "code",
"metadata": {
"id": "Piz0c5f44hRo"
},
"source": [
"conf = Configs()"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "wwMzCqpD6vkL"
},
"source": [
"Set experiment configurations and assign a configurations dictionary to override configurations"
]
},
{
"cell_type": "code",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 17
},
"id": "e6hmQhTw4nks",
"outputId": "6660f5c6-347c-4370-f19a-ff7bd9b96535"
},
"source": [
"experiment.configs(conf,\n",
" # A dictionary of configurations to override\n",
" {'tokenizer': 'character',\n",
" 'text': 'tiny_shakespeare',\n",
" 'optimizer.learning_rate': 1.0,\n",
" 'optimizer.optimizer': 'Noam',\n",
" 'prompt': 'It is',\n",
" 'prompt_separator': '',\n",
"\n",
" 'train_loader': 'shuffled_train_loader',\n",
" 'valid_loader': 'shuffled_valid_loader',\n",
"\n",
" 'seq_len': 128,\n",
" 'epochs': 128,\n",
" 'batch_size': 16,\n",
" 'inner_iterations': 25})"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "EvI7MtgJ61w5"
},
"source": [
"Set PyTorch models for loading and saving"
]
},
{
"cell_type": "code",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 255
},
"id": "GDlt7dp-5ALt",
"outputId": "5154be3b-dbfc-4002-82f6-1b6d58970e21"
},
"source": [
"experiment.add_pytorch_models({'model': conf.model})"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "KJZRf8527GxL"
},
"source": [
"Start the experiment and run the training loop."
]
},
{
"cell_type": "code",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 1000
},
"id": "aIAWo7Fw5DR8",
"outputId": "583dc28a-338e-4089-a415-e21d8b65fa3e"
},
"source": [
"with experiment.start():\n",
" conf.run()"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "code",
"metadata": {
"id": "oBXXlP2b7XZO"
},
"source": [
""
],
"outputs": [],
"execution_count": null
}
]
}
labml_nn/transformers/fast_weights/experiment.py
+115
View File
@@ -0,0 +1,115 @@
"""
---
title: Train Fast Weights Transformer
summary: This is training code with notes for a Fast Weights Transformer.
---
# Train Fast Weights Transformer
This trains a fast weights transformer model for auto-regression.
Here’s a Colab notebook for training a fast weights transformer on Tiny Shakespeare dataset.
[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/transformers/fast_weights/experiment.ipynb)
"""
import torch
from torch import nn
from labml import experiment
from labml.configs import option
from labml.utils.pytorch import get_modules
from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs
class AutoregressiveModel(nn.Module):
"""
## Auto regressive model
"""
def __init__(self, n_vocab: int, d_model: int, transformer: nn.Module):
super().__init__()
# Token embedding module
self.src_embed = nn.Embedding(n_vocab, d_model)
self.transformer = transformer
self.generator = nn.Linear(d_model, n_vocab)
def forward(self, x: torch.Tensor):
# Embed the tokens
x = self.src_embed(x)
# Run it through the the transformer
res = self.transformer(x)
# Generate logits of the next token
return self.generator(res), None
class Configs(NLPAutoRegressionConfigs):
"""
## Configurations
The default configs can and will be over-ridden when we start the experiment
"""
model: AutoregressiveModel
d_model: int = 512
nu: int = 1
heads: int = 8
dropout: float = 0.0
d_ff: int = 2048
n_layers: int = 6
@option(Configs.model)
def fast_weights_transformer(c: Configs):
"""
Create [fast weights transformer](index.html).
"""
from labml_nn.transformers.fast_weights import FastWeightsAttentionTransformer, \
FastWeightsAttentionTransformerLayer, FastWeightsAttention, FeedForward
from labml_nn.transformers.fast_weights import DPFP
return AutoregressiveModel(
c.n_tokens, c.d_model,
FastWeightsAttentionTransformer(
FastWeightsAttentionTransformerLayer(d_model=c.d_model,
attn=FastWeightsAttention(c.heads, c.d_model, c.dropout, DPFP(nu=c.nu)),
feed_forward=FeedForward(c.d_model, c.d_ff, c.dropout),
dropout_prob=c.dropout),
c.n_layers)).to(c.device)
def main():
# Create experiment
experiment.create(name="fast_weights_transformer")
# Create configs
conf = Configs()
# Load configurations
experiment.configs(conf,
# A dictionary of configurations to override
{'tokenizer': 'character',
'text': 'tiny_shakespeare',
'optimizer.learning_rate': 1.0,
'optimizer.optimizer': 'Noam',
'prompt': 'It is',
'prompt_separator': '',
'train_loader': 'shuffled_train_loader',
'valid_loader': 'shuffled_valid_loader',
'seq_len': 128,
'epochs': 128,
'batch_size': 16,
'inner_iterations': 25})
# Set models for saving and loading
experiment.add_pytorch_models(get_modules(conf))
# Start the experiment
with experiment.start():
# Run the training loop
conf.run()
if __name__ == '__main__':
main()
labml_nn/transformers/fast_weights/readme.md
+10
View File
@@ -0,0 +1,10 @@
# [Fast weights transformer](https://nn.labml.ai/transformers/fast_weights/index.html)
This is an annotated implementation of the paper
[Linear Transformers Are Secretly Fast Weight Memory Systems in PyTorch](https://arxiv.org/abs/2102.11174).
Here is the [annotated implementation](https://nn.labml.ai/transformers/fast_weights/index.html).
Here are [the training code](https://nn.labml.ai/transformers/fast_weights/experiment.html)
and a notebook for training a fast weights transformer on the Tiny Shakespeare dataset.
[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/transformers/fast_weights/experiment.ipynb)
labml_nn/transformers/fast_weights/token_wise.py
+129
View File
@@ -0,0 +1,129 @@
"""
---
title: Fast Weight Systems
summary: >
This is an annotated implementation/tutorial of
Linear Transformers Are Secretly Fast Weight Memory Systems in PyTorch.
---
"""
from typing import Optional
import torch
from torch import nn
from labml_nn.transformers.fast_weights import DPFP
from labml_nn.transformers.feed_forward import FeedForward
from labml_nn.transformers.mha import PrepareForMultiHeadAttention
from labml_nn.utils import clone_module_list
class FastWeightsAttention(nn.Module):
def __init__(self, heads: int, d_model: int, dropout_prob: float, phi: DPFP):
super().__init__()
# Number of features per head
self.d_k = d_model // heads
#
self.heads = heads
# These transform the `query` multi-headed attention.
self.query = PrepareForMultiHeadAttention(d_model, heads, self.d_k, bias=False)
# These transform the `key` and `value` for multi-headed attention.
self.key = PrepareForMultiHeadAttention(d_model, heads, self.d_k, bias=False)
self.value = PrepareForMultiHeadAttention(d_model, heads, self.d_k, bias=False)
self.gate = nn.Sequential(PrepareForMultiHeadAttention(d_model, heads, 1, bias=False),
nn.Sigmoid())
self.phi = phi
# Output layer
self.output = nn.Linear(d_model, d_model)
# Dropout
self.dropout = nn.Dropout(dropout_prob)
def forward(self, x: torch.Tensor, weights: Optional[torch.Tensor]):
query = self.phi(self.query(x))
key = self.phi(self.key(x))
value = self.value(x)
if weights is None:
weights = key.new_zeros((key.shape[0], key.shape[1], value.shape[2], key.shape[2]))
value_existing = torch.einsum('bhvk,bhk->bhv', weights, key)
beta = self.gate(x)
weights = weights + torch.einsum('bhv,bhk->bhvk', beta * (value - value_existing), key)
x = torch.einsum('bhvk,bhk->bhv', weights, query)
# Concatenate multiple heads
x = x.reshape(x.shape[0], -1)
# Output layer
return self.output(x), weights
class FastWeightsAttentionTransformerLayer(nn.Module):
def __init__(self, *,
d_model: int,
attn: FastWeightsAttention,
feed_forward: FeedForward,
dropout_prob: float):
super().__init__()
# Transformer size $d_{model}$
self.size = d_model
#
self.attn = attn
self.feed_forward = feed_forward
self.dropout = nn.Dropout(dropout_prob)
# Normalization layers
self.norm_self_attn = nn.LayerNorm([d_model])
self.norm_ff = nn.LayerNorm([d_model])
def forward(self, x: torch.Tensor, weights: Optional[torch.Tensor]):
attn, weights = self.attn(x, weights)
# Add the self attention results
x = x + self.dropout(attn)
# Normalize for feed-forward
z = self.norm_ff(x)
# Pass through the feed-forward network
ff = self.feed_forward(z)
# Add the feed-forward results back
x = x + self.dropout(ff)
#
return x, weights
class FastWeightsAttentionTransformer(nn.Module):
def __init__(self, layer: FastWeightsAttentionTransformerLayer, n_layers: int):
super().__init__()
# Make copies of the transformer layer
self.layers = clone_module_list(layer, n_layers)
# Final normalization layer
self.norm = nn.LayerNorm([layer.size])
def forward(self, x_seq: torch.Tensor):
# Split the input to a list along the sequence axis
x_seq = torch.unbind(x_seq, dim=0)
# List to store the outputs
res = []
# For each input step
weights = [None for _ in range(len(self.layers))]
for x in x_seq:
# Run through each layer
for i, layer in enumerate(self.layers):
# Get layer output
x, weights[i] = layer(x, weights[i])
res.append(x)
# Stack the output tensors
res = torch.stack(res)
# Normalize the output
return self.norm(res)
labml_nn/transformers/feed_forward.py
+93
View File
@@ -0,0 +1,93 @@
"""
---
title: Position-wise Feed-Forward Network (FFN)
summary: Documented reusable implementation of the position wise feedforward network.
---
# Position-wise Feed-Forward Network (FFN)
This is a [PyTorch](https://pytorch.org) implementation
of position-wise feedforward network used in transformer.
FFN consists of two fully connected layers.
Number of dimensions in the hidden layer $d_{ff}$, is generally set to around
four times that of the token embedding $d_{model}$.
So it is sometime also called the expand-and-contract network.
There is an activation at the hidden layer, which is
usually set to ReLU (Rectified Linear Unit) activation, $$\max(0, x)$$
That is, the FFN function is,
$$FFN(x, W_1, W_2, b_1, b_2) = \max(0, x W_1 + b_1) W_2 + b_2$$
where $W_1$, $W_2$, $b_1$ and $b_2$ are learnable parameters.
Sometimes the
GELU (Gaussian Error Linear Unit) activation is also used instead of ReLU.
$$x \Phi(x)$$ where $\Phi(x) = P(X \le x), X \sim \mathcal{N}(0,1)$
### Gated Linear Units
This is a generic implementation that supports different variants including
[Gated Linear Units](https://arxiv.org/abs/2002.05202) (GLU).
We have also implemented experiments on these:
* [experiment that uses `labml.configs`](glu_variants/experiment.html)
* [simpler version from scratch](glu_variants/simple.html)
"""
import torch
from torch import nn
class FeedForward(nn.Module):
"""
## FFN module
"""
def __init__(self, d_model: int, d_ff: int,
dropout: float = 0.1,
activation=nn.ReLU(),
is_gated: bool = False,
bias1: bool = True,
bias2: bool = True,
bias_gate: bool = True):
"""
* `d_model` is the number of features in a token embedding
* `d_ff` is the number of features in the hidden layer of the FFN
* `dropout` is dropout probability for the hidden layer
* `is_gated` specifies whether the hidden layer is gated
* `bias1` specified whether the first fully connected layer should have a learnable bias
* `bias2` specified whether the second fully connected layer should have a learnable bias
* `bias_gate` specified whether the fully connected layer for the gate should have a learnable bias
"""
super().__init__()
# Layer one parameterized by weight $W_1$ and bias $b_1$
self.layer1 = nn.Linear(d_model, d_ff, bias=bias1)
# Layer one parameterized by weight $W_1$ and bias $b_1$
self.layer2 = nn.Linear(d_ff, d_model, bias=bias2)
# Hidden layer dropout
self.dropout = nn.Dropout(dropout)
# Activation function $f$
self.activation = activation
# Whether there is a gate
self.is_gated = is_gated
if is_gated:
# If there is a gate the linear layer to transform inputs to
# be multiplied by the gate, parameterized by weight $V$ and bias $c$
self.linear_v = nn.Linear(d_model, d_ff, bias=bias_gate)
def forward(self, x: torch.Tensor):
# $f(x W_1 + b_1)$
g = self.activation(self.layer1(x))
# If gated, $f(x W_1 + b_1) \otimes (x V + b) $
if self.is_gated:
x = g * self.linear_v(x)
# Otherwise
else:
x = g
# Apply dropout
x = self.dropout(x)
# $(f(x W_1 + b_1) \otimes (x V + b)) W_2 + b_2$ or $f(x W_1 + b_1) W_2 + b_2$
# depending on whether it is gated
return self.layer2(x)
labml_nn/transformers/feedback/__init__.py
+528
View File
@@ -0,0 +1,528 @@
"""
---
title: Feedback Transformer
summary: >
This is an annotated implementation/tutorial the Feedback Transformer in PyTorch.
---
# Feedback Transformer
This is a [PyTorch](https://pytorch.org) implementation of the paper
[Accessing Higher-level Representations in Sequential Transformers with Feedback Memory](https://arxiv.org/abs/2002.09402).
Normal transformers process tokens in parallel. Each transformer layer pays attention
to the outputs of the previous layer.
Feedback transformer pays attention to the output of all layers in previous steps.
So this adds recurrence, and we need to process token-by-token.
This slows down the training significantly (about 5X - 10X depending on the sequence length).
However, when predicting Feedback Transformer is faster because you can predict the next token
if you cache the memory vectors.
In order to speed up the training, the paper discusses starting with a short sequence length and
gradually increasing it.
They also discuss using a pretrained parallel transformer as the starting point.
The original feedback transformer doesn't keep the outputs of all layers.
Instead it keeps weighted sum of the output of all layers.
This reduces the memory used for caching during prediction.
The first half of this file implements this.
The updated feedback transformer shares weights $W^l_k$ and $W^l_v$ used
to calculate keys and values among the layers.
We then calculate the keys and values for each step only once and keep
them cached.
The [second half](#shared_kv) of this file implements this.
We implemented a custom PyTorch function to improve performance.
Here's [the training code](experiment.html) and a notebook for training a feedback transformer on Tiny Shakespeare dataset.
[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/transformers/feedback/experiment.ipynb)
"""
import math
from typing import Optional
import torch
from torch import nn
from labml_nn.transformers.feed_forward import FeedForward
from labml_nn.transformers.mha import PrepareForMultiHeadAttention
from labml_nn.utils import clone_module_list
class FeedbackAttention(nn.Module):
r"""
## Feedback Attention
This module computes recurrent attention similar to attention from original transformers
paper.
$$\mathop{Attention}(Q, K, V) = \underset{seq}{\mathop{softmax}}\Bigg(\frac{Q^\top K}{\sqrt{d_k}}\Bigg)V$$
"""
def __init__(self, heads: int, d_model: int, dropout_prob: float = 0.1, *,
is_kv_precomputed: bool = False):
"""
* 'heads' is the number of attention heads
* `d_model` is the number of features in the transformer
* `dropout_prob` is the attention dropout probability
* `is_kv_precomputed` is whether key, value tensors are already calculated
"""
super().__init__()
# Number of features per head
self.d_k = d_model // heads
#
self.heads = heads
# These transform the `query` multi-headed attention.
self.query = PrepareForMultiHeadAttention(d_model, heads, self.d_k, bias=False)
# These transform the `key` and `value` for multi-headed attention.
if not is_kv_precomputed:
self.key = PrepareForMultiHeadAttention(d_model, heads, self.d_k, bias=False)
self.value = PrepareForMultiHeadAttention(d_model, heads, self.d_k, bias=True)
# Keys and values are already calculated
else:
self.key = None
self.value = None
# Output layer
self.output = nn.Linear(d_model, d_model)
# Dropout
self.dropout = nn.Dropout(dropout_prob)
# Scaling factor before the softmax
self.scale = 1 / math.sqrt(self.d_k)
# Softmax for attention along the time dimension of `key`
self.softmax = nn.Softmax(dim=0)
# Number of relative positions
self.P = 2 ** 12
# Relative positional embeddings for key relative to the query.
self.key_pos_embeddings = nn.Parameter(torch.zeros((self.P, heads, self.d_k)), requires_grad=True)
# Relative positional embedding bias for key relative to the query.
self.key_pos_bias = nn.Parameter(torch.zeros((self.P, heads)), requires_grad=True)
# Positional embeddings for the query is independent of the position of the query
self.query_pos_bias = nn.Parameter(torch.zeros((heads, self.d_k)), requires_grad=True)
# We store attentions so that it can be used for logging, or other computations if needed
self.attn = None
def get_scores(self, query: torch.Tensor, key: torch.Tensor):
r"""
### Get attention scores
We use relative positional encodings for attention, similar
to [relative multi-head attention form Transformer-XL paper](../relative_mha.html).
Attention from current step's query to key in step $j$ (relative to current step) is,
\begin{align}
A_{j} &= Q^\top K_j \\
&= lin_q(X^q + P_q)^\top lin_k(X^k_j + P_j) \\
&= (Q + U^Q)^\top(K_j + U^K_j) \\
&= \underset{\textcolor{lightgreen}{A}}{Q^\top K_j} +
\underset{\textcolor{lightgreen}{B}}{Q^\top U^K_j} +
\underset{\textcolor{lightgreen}{C}}{{U^Q}^\top K_j} +
\underset{\textcolor{lightgreen}{D}}{{U^Q}^\top U^K_j}
\end{align}
where $Q, K_j$, are linear transformations of
original embeddings $X^q, X^k_j$
and $U^Q, U^K_j$ are linear transformations of
positional encodings $P_q, P_j$.
We replace term $\textcolor{lightgreen}{D}$ with $S_j$.
"""
# $U^K_j$
key_pos_emb = self.key_pos_embeddings[-key.shape[0]:]
# $U^Q$
query_pos_bias = self.query_pos_bias[None, :, :]
# $S_j$
key_pos_bias = self.key_pos_bias[-key.shape[0]:]
# $\underset{\textcolor{lightgreen}{A}}{Q^\top K_j} + \underset{\textcolor{lightgreen}{C}}{{U^Q}^\top K_j}$
ac = torch.einsum('bhd,jbhd->jbh', query + query_pos_bias, key)
# $\underset{\textcolor{lightgreen}{B}}{Q^\top U^K_j} + \underset{\textcolor{lightgreen}{D}}{S_j}$
bd = torch.einsum('bhd,jhd->jbh', query, key_pos_emb) + key_pos_bias[:, None, :]
# $A_j$
return ac + bd
def forward(self, *,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor):
"""
* `query` has shape `[batch_size, d_model]`
* `key` and `value` has shape `[seq_len, batch_size, d_model]`
"""
# Prepare `query`, `key` and `value` for attention computation
# `key` and `value` will then have shape `[seq_len, batch_size, heads, d_k]`
# and `query` will have shape `[batch_size, heads, d_k]`
query = self.query(query)
if self.key:
key = self.key(key)
if self.value:
value = self.value(value)
# Compute attention scores.
# Results in a tensor of shape `[seq_len, batch_size, heads]`
scores = self.get_scores(query, key)
# Scale scores $\frac{1}{\sqrt{d_k}}$
scores *= self.scale
# Softmax
attn = self.softmax(scores)
# Apply dropout
attn = self.dropout(attn)
# Multiply by the values
x = torch.einsum("jbh,jbhd->bhd", attn, value)
# Concatenate multiple heads
x = x.reshape(x.shape[0], -1)
# Output layer
return self.output(x)
class FeedbackTransformerLayer(nn.Module):
"""
## Feedback Transformer Layer
This implements a single transformer layer in the feedback transformer.
"""
def __init__(self, *,
d_model: int,
attn: FeedbackAttention,
feed_forward: FeedForward,
dropout_prob: float):
"""
* `d_model` is the number of features in the transformer
* `attn` is the feedback attention module
* `feed_forward` is the position-wise feed forward layer
* `dropout_prob` is the dropout probability for dropout layers after attention and feed-forward
"""
super().__init__()
# Transformer size $d_{model}$
self.size = d_model
#
self.attn = attn
self.feed_forward = feed_forward
self.dropout = nn.Dropout(dropout_prob)
# Normalization layers
self.norm_self_attn = nn.LayerNorm([d_model])
self.norm_ff = nn.LayerNorm([d_model])
def forward(self, *,
x: torch.Tensor,
key: Optional[torch.Tensor],
value: Optional[torch.Tensor]):
# If there is memory
if key is not None:
# Normalize the vectors before doing self attention
z = self.norm_self_attn(x)
# Run through self attention, i.e. keys and values are from self
self_attn = self.attn(query=z, key=key, value=value)
# Add the self attention results
x = x + self.dropout(self_attn)
# Normalize for feed-forward
z = self.norm_ff(x)
# Pass through the feed-forward network
ff = self.feed_forward(z)
# Add the feed-forward results back
x = x + self.dropout(ff)
#
return x
class FeedbackTransformer(nn.Module):
"""
## Feedback Transformer Module
"""
def __init__(self, layer: FeedbackTransformerLayer, n_layers: int):
"""
* `layer` is the feedback transformer layer, which we clone for each layer
* `n_layers` is the number of layers in the transformer
"""
super().__init__()
# Make copies of the transformer layer
self.layers = clone_module_list(layer, n_layers)
# Final normalization layer
self.norm = nn.LayerNorm([layer.size])
# Memory vectors are computed as a weighted sum of representations of each layer.
# This is the weights parameter for that.
self.weights = nn.Parameter(torch.ones(n_layers + 1), requires_grad=True)
# Softmax for weights before taking the weighted sum
self.softmax = nn.Softmax(0)
def forward(self, x_seq: torch.Tensor):
"""
* `x_seq` is the input with shape `[seq_len, batch_size, d_model]`
"""
# Split the input to a list along the sequence axis
x_seq = torch.unbind(x_seq, dim=0)
# List to store the outputs
res = []
# List to store the memory vectors
mem = []
# For each input step
for x in x_seq:
# List to store layer outputs
layer_outputs = [x]
# If there is memory, stack them into a vector
mem_tensor = torch.stack(mem) if mem else None
# Run through each layer
for layer in self.layers:
# Get layer output
x = layer(x=x, key=mem_tensor, value=mem_tensor)
# Append them to the list of layer outputs
layer_outputs.append(x)
# Stack the layer outputs to a tensor
layer_outputs = torch.stack(layer_outputs)
# Calculate the memory vector as a weighted sum of layer outputs
mem.append(torch.einsum('lbd,l->bd', layer_outputs, self.softmax(self.weights)))
# Append the output to results
res.append(x)
# Stack the output tensors
res = torch.stack(res)
# Normalize the output
return self.norm(res)
# <a id="shared_kv"></a>
#
# # Shared keys and values among layers
class StackFunction(torch.autograd.Function):
"""
### Stack Function implementation
We implement a custom function instead of appending to a python list
and then doing `torch.stack`.
This greatly improves the performance over calling `torch.stack` at
each step along the sequence.
Everytime `torch.stack` is called, it creates a new tensor, while
this method and the accompanying class `Stack` share memory for each step.
"""
@staticmethod
def forward(ctx, memory, memory_grad, last, n):
"""
* `ctx` is the context of the function (which lets us cache stuff)
* `memory` is the shared memory tensor where we stack and store the values of each step (keys & values)
* `memory_grad` is the shared memory tensor to store and accumulate gradients of each step
* `last` is the last value stacked
* `n` is the number of steps (i.e. size of the stack)
This returns the stacked tensor for steps upto `n`.
"""
# Cache accumulated gradients
ctx._mem_grad = memory_grad
# Cache the size of the stack
ctx._n = n
# Return the stack
return memory[:n + 1]
@staticmethod
def backward(ctx, grad_output):
"""
* `grad_output` is the gradient with respect to the output of about `forward` function
This accumulates the gradients in the shared memory tensor and return the
gradients with respect to the `last` result in the stack.
"""
# Get the current size of the stack
n = ctx._n
# Get the accumulated gradients
memory_grad = ctx._mem_grad
# Add the gradients
memory_grad[:n + 1] += grad_output
# Return the gradients w.r.t to last value in the stack
return None, None, memory_grad[n], None
class Stack:
"""
### Stack Module
This uses the stack function defined above, and does the necessary initializations.
"""
def __init__(self, max_len: int):
"""
* `max_len` is the maximum size of the stack
"""
self.max_len = max_len
self.memory = None
self.memory_grad = None
self.last = None
self.n = -1
self.last_get_n = -1
def append(self, n: int, value: torch.Tensor):
"""
* `n` is the size of the stack
* `value` is the tensor that needs to be added to the stack
"""
# You need to get (use) the stack after adding a value.
# Otherwise this implementation fails
assert n == 0 or self.last_get_n == n - 1, f"{n}, {self.last_get_n}"
# Do this without gradients
with torch.no_grad():
# Initialize the shared memory tensor to keep the stack
if self.memory is None or self.memory.shape[1:] != value.shape:
# This should only happen when the stack is empty
assert n == 0
# Create a tensor for the stack
self.memory = value.new_zeros(self.max_len, *value.shape, requires_grad=False)
# Create a tensor to accumulate the gradients
self.memory_grad = value.new_zeros(self.memory.shape, requires_grad=False)
# The memory is already initialized but we are resetting the stack.
#
# This could have been another function like `reset`, but
# we found this easier to use.
elif n == 0:
# Reset accumulated gradients
self.memory_grad.fill_(0.)
# Set the value in the correct position of the stack
self.memory.data[n] = value.detach()
# Keep track of the stack (for debugging)
self.n = n
# Keep track of the last value added to the stack.
# We need this to be passed on to `StackFunction` in order
# to get the gradients propagated backwards.
self.last = value
def get(self):
"""
Returns the stack
"""
# Keep track of the size of the stack when it was used.
# This is used for a sanity check in `append`.
self.last_get_n = self.n
# Take it all through `StackFunction` so that `StackFunction.backwards`
# is called by PyTorch during backpropagation.
return StackFunction.apply(self.memory, self.memory_grad, self.last, self.n)
def free(self):
"""
To release memory
"""
self.memory = None
self.memory_grad = None
self.last = None
class FeedbackTransformerKV(nn.Module):
"""
## Updated Feedback Transformer Module
This is the updated feedback transformer module that caches the keys and values.
"""
def __init__(self, layer: FeedbackTransformerLayer, n_layers: int, d_model: int, heads: int):
"""
* `layer` is the feedback transformer layer, which we clone for each layer
* `n_layers` is the number of layers in the transformer
* `d_model` is the number of features in the transformer
* 'heads' is the number of attention heads
"""
super().__init__()
# Make copies of the transformer layer
self.layers = clone_module_list(layer, n_layers)
# Final normalization layer
self.norm = nn.LayerNorm([layer.size])
# Memory vectors are computed as a weighted sum of representations of each layer.
# This is the weights parameter for that.
self.weights = nn.Parameter(torch.ones(n_layers + 1), requires_grad=True)
# Softmax for weights before taking the weighted sum
self.softmax = nn.Softmax(0)
# Number of features in a head
d_k = d_model // heads
# Module to transform embeddings (memory) to get keys
self.key = PrepareForMultiHeadAttention(d_model, heads, d_k, bias=False)
# Module to transform embeddings (memory) to get keys
self.value = PrepareForMultiHeadAttention(d_model, heads, d_k, bias=False)
# Memory for stacked keys
self.mem_key = Stack(512)
# Memory for stacked values
self.mem_value = Stack(512)
def forward(self, x_seq: torch.Tensor):
"""
* `x_seq` is the input with shape `[seq_len, batch_size, d_model]`
"""
# Split the input to a list along the sequence axis
x_seq = torch.unbind(x_seq, dim=0)
# List to store the outputs
res = []
# For each input step
for step, x in enumerate(x_seq):
# List to store layer outputs
layer_outputs = [x]
# Stack of keys and values
key_tensor = None
value_tensor = None
# Get the keys and values tensors if we are beyond the initial step
if step > 0:
key_tensor = self.mem_key.get()
value_tensor = self.mem_value.get()
# Run through each layer
for layer in self.layers:
# Get layer output
x = layer(x=x, key=key_tensor, value=value_tensor)
# Append them to the list of layer outputs
layer_outputs.append(x)
# Stack the layer outputs to a tensor
layer_outputs = torch.stack(layer_outputs)
# Calculate the memory vector as a weighted sum of layer outputs
mem = torch.einsum('lbd,l->bd', layer_outputs, self.softmax(self.weights))
# Calculate the keys from memory and add it to the stack
self.mem_key.append(step, self.key(mem))
# Calculate the values from memory and add it to the stack
self.mem_value.append(step, self.value(mem))
# Append the output to results
res.append(x)
# Stack the output tensors
res = torch.stack(res)
# Normalize the output
return self.norm(res)
def free(self):
self.mem_key.free()
self.mem_value.free()
labml_nn/transformers/feedback/experiment.ipynb
+232
View File
@@ -0,0 +1,232 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "AYV_dMVDxyc2"
},
"source": [
"[![Github](https://img.shields.io/github/stars/labmlai/annotated_deep_learning_paper_implementations?style=social)](https://github.com/labmlai/annotated_deep_learning_paper_implementations)\n",
"[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/transformers/feedback/experiment.ipynb) \n",
"\n",
"## Feedback Transformer\n",
"\n",
"This is an experiment training Shakespeare dataset with Feedback Transformer."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "AahG_i2y5tY9"
},
"source": [
"Install the `labml-nn` package"
]
},
{
"cell_type": "code",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "ZCzmCrAIVg0L",
"outputId": "aa1fe63d-1755-4394-dcdf-9897ac6c1ee4"
},
"source": [
"!pip install labml-nn"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "SE2VUQ6L5zxI"
},
"source": [
"Imports"
]
},
{
"cell_type": "code",
"metadata": {
"id": "0hJXx_g0wS2C"
},
"source": [
"from labml import experiment\n",
"from labml_nn.transformers.feedback.experiment import Configs"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "Lpggo0wM6qb-"
},
"source": [
"Create an experiment"
]
},
{
"cell_type": "code",
"metadata": {
"id": "bFcr9k-l4cAg"
},
"source": [
"experiment.create(name=\"feedback_transformer\")"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "-OnHLi626tJt"
},
"source": [
"Initialize configurations"
]
},
{
"cell_type": "code",
"metadata": {
"id": "Piz0c5f44hRo"
},
"source": [
"conf = Configs()"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "wwMzCqpD6vkL"
},
"source": [
"Set experiment configurations and assign a configurations dictionary to override configurations"
]
},
{
"cell_type": "code",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 17
},
"id": "e6hmQhTw4nks",
"outputId": "91d99011-7a61-48fa-ee5c-2fa845883cec"
},
"source": [
"experiment.configs(conf,\n",
" {'tokenizer': 'character',\n",
" 'text': 'tiny_shakespeare',\n",
" 'optimizer.learning_rate': 1.0,\n",
" 'optimizer.optimizer': 'Noam',\n",
" 'prompt': 'It is',\n",
" 'prompt_separator': '',\n",
"\n",
" 'model': 'feedback_transformer',\n",
"\n",
" 'train_loader': 'shuffled_train_loader',\n",
" 'valid_loader': 'shuffled_valid_loader',\n",
"\n",
" 'seq_len': 64,\n",
" 'epochs': 128,\n",
" 'batch_size': 80,\n",
" 'inner_iterations': 25})"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "EvI7MtgJ61w5"
},
"source": [
"Set PyTorch models for loading and saving"
]
},
{
"cell_type": "code",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 255
},
"id": "GDlt7dp-5ALt",
"outputId": "bee57c09-4e71-4329-debb-b23ed309a3ef"
},
"source": [
"experiment.add_pytorch_models({'model': conf.model})"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "KJZRf8527GxL"
},
"source": [
"Start the experiment and run the training loop."
]
},
{
"cell_type": "code",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 1000
},
"id": "aIAWo7Fw5DR8",
"outputId": "bfb52e21-1913-4bd0-b67f-bef8a14f0f05"
},
"source": [
"# Start the experiment\n",
"with experiment.start():\n",
" conf.run()"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "code",
"metadata": {
"id": "oBXXlP2b7XZO"
},
"source": [],
"outputs": [],
"execution_count": null
}
],
"metadata": {
"accelerator": "GPU",
"colab": {
"collapsed_sections": [],
"name": "Feedback Transformer",
"provenance": []
},
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.7.5"
}
},
"nbformat": 4,
"nbformat_minor": 4
}
labml_nn/transformers/feedback/experiment.py
+136
View File
@@ -0,0 +1,136 @@
"""
---
title: Train Feedback Transformer
summary: This is training code with notes for a feedback transformer.
---
# Train Feedback Transformer
This trains a [feedback transformer](index.html) model for auto-regression.
You can pick the original feedback transformer or the new version
where the keys and values are precalculated.
Here's a Colab notebook for training a feedback transformer on Tiny Shakespeare dataset.
[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/transformers/feedback/experiment.ipynb)
"""
import torch
from labml import experiment
from labml.configs import option
from labml.utils.pytorch import get_modules
from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs
from torch import nn
class AutoregressiveModel(nn.Module):
"""
## Auto regressive model
"""
def __init__(self, n_vocab: int, d_model: int, transformer: nn.Module):
super().__init__()
# Token embedding module
self.src_embed = nn.Embedding(n_vocab, d_model)
self.transformer = transformer
self.generator = nn.Linear(d_model, n_vocab)
def forward(self, x: torch.Tensor):
# Embed the tokens
x = self.src_embed(x)
# Run it through the the transformer
res = self.transformer(x)
# Generate logits of the next token
return self.generator(res), None
class Configs(NLPAutoRegressionConfigs):
"""
## Configurations
The default configs can and will be over-ridden when we start the experiment
"""
model: AutoregressiveModel
d_model: int = 512
heads: int = 8
dropout: float = 0.0
d_ff: int = 2048
n_layers: int = 6
@option(Configs.model)
def feedback_transformer(c: Configs):
"""
Create [original feedback transformer](index.html).
"""
from labml_nn.transformers.feedback import FeedbackTransformer, FeedbackTransformerLayer, \
FeedbackAttention, FeedForward
return AutoregressiveModel(
c.n_tokens, c.d_model,
FeedbackTransformer(
FeedbackTransformerLayer(d_model=c.d_model,
attn=FeedbackAttention(c.heads, c.d_model, c.dropout),
feed_forward=FeedForward(c.d_model, c.d_ff, c.dropout),
dropout_prob=c.dropout),
c.n_layers)).to(c.device)
@option(Configs.model)
def feedback_transformer_kv(c: Configs):
"""
Create [updated feedback transformer](index.html#kv_shared), with precalculated keys and values.
"""
from labml_nn.transformers.feedback import FeedbackTransformerKV, FeedbackTransformerLayer, \
FeedbackAttention, FeedForward
return AutoregressiveModel(
c.n_tokens, c.d_model,
FeedbackTransformerKV(
FeedbackTransformerLayer(d_model=c.d_model,
attn=FeedbackAttention(c.heads, c.d_model, c.dropout,
is_kv_precomputed=True),
feed_forward=FeedForward(c.d_model, c.d_ff, c.dropout),
dropout_prob=c.dropout),
c.n_layers, c.d_model, c.heads)).to(c.device)
def main():
# Create experiment
experiment.create(name="feedback_transformer")
# Create configs
conf = Configs()
# Load configurations
experiment.configs(conf,
# A dictionary of configurations to override
{'tokenizer': 'character',
'text': 'tiny_shakespeare',
'optimizer.learning_rate': 1.0,
'optimizer.optimizer': 'Noam',
'prompt': 'It is',
'prompt_separator': '',
# Use `feedback_transformer` for original feedback transformer
'model': 'feedback_transformer_kv',
'train_loader': 'shuffled_train_loader',
'valid_loader': 'shuffled_valid_loader',
'seq_len': 128,
'epochs': 128,
'batch_size': 64,
'inner_iterations': 25})
# Set models for saving and loading
experiment.add_pytorch_models(get_modules(conf))
# Start the experiment
with experiment.start():
# Run the training loop
conf.run()
if __name__ == '__main__':
main()
labml_nn/transformers/feedback/readme.md
+34
View File
@@ -0,0 +1,34 @@
# [Feedback Transformer](https://nn.labml.ai/transformers/feedback/index.html)
This is a [PyTorch](https://pytorch.org) implementation of the paper
[Accessing Higher-level Representations in Sequential Transformers with Feedback Memory](https://arxiv.org/abs/2002.09402).
Normal transformers process tokens in parallel. Each transformer layer pays attention
to the outputs of the previous layer.
Feedback transformer pays attention to the output of all layers in previous steps.
So this adds recurrence, and we need to process token-by-token.
This slows down the training significantly (about 5X - 10X depending on the sequence length).
However, when predicting Feedback Transformer is faster because you can predict the next token
if you cache the memory vectors.
In order to speed up the training the paper discusses starting with a short sequence length and
gradually increasing it.
They also discuss using a pretrained parallel transformer as the starting point.
The original feedback transformer doesn't keep the outputs of all layers.
Instead it keeps weighted sum of the output of all layers.
This reduces the memory used for caching during prediction.
The first half of this file implements this.
The updated feedback transformer shares weights used
to calculate keys and values among the layers.
We then calculate the keys and values for each step only once and keep
them cached.
The [second half](#shared_kv) of this file implements this.
We implemented a custom PyTorch function to improve performance.
Here's [the training code](experiment.html) and a notebook for training a feedback transformer on Tiny Shakespeare dataset.
[Colab Notebook](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/transformers/feedback/experiment.ipynb)
[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/transformers/feedback/experiment.ipynb)
labml_nn/transformers/flash/__init__.py
+1000
View File
File diff suppressed because it is too large Load Diff
labml_nn/transformers/flash/test.py
+179
View File
@@ -0,0 +1,179 @@
"""
### Test Flash Attention Implementation
This is the code to test and measure performance of our flash attention implementation
"""
import torch
import triton
from labml import logger, monit
from labml_nn.transformers.flash import attention
HI_PRES_TORCH = torch.float32
@torch.no_grad()
def _calc_abs_rel_error(a: torch.Tensor, b: torch.Tensor, atol=1e-2):
"""
#### Calculate absolute and relative error for reporting
"""
d = (a - b).abs()
max_abs = d.max()
d = (d - atol).clamp(min=0)
d = d / b.abs()
max_rel = d.max()
return max_abs.cpu().item(), max_rel.cpu().item()
def test_fwd_bwd(batch_size, n_heads, k_heads, q_seq_len, kv_seq_len, d_head, causal, dtype, device):
"""
#### Compare our implementation with naive PyTorch attention
"""
with monit.section(f'Init {q_seq_len} {kv_seq_len} {d_head}'):
torch.manual_seed(20)
q = (torch.empty((batch_size, n_heads, q_seq_len, d_head),
dtype=dtype, device=device).normal_(mean=0.0, std=0.5).requires_grad_())
k = (torch.empty((batch_size, k_heads, kv_seq_len, d_head),
dtype=dtype, device=device).normal_(mean=0.0, std=0.5).requires_grad_())
v = (torch.empty((batch_size, k_heads, kv_seq_len, d_head),
dtype=dtype, device=device).normal_(mean=0.0, std=0.5).requires_grad_())
sm_scale = d_head ** -0.5
d_out = torch.randn_like(q)
# reference implementation
mask = torch.tril(torch.ones((q_seq_len, kv_seq_len), device=device, dtype=torch.bool))
torch.cuda.synchronize()
with monit.section('Pytorch'):
p = torch.matmul(q.view(batch_size, k_heads, -1, q_seq_len, d_head),
k.transpose(2, 3)[:, :, None, :, :]) * sm_scale
if causal:
p[:, :, :, ~mask] = float("-inf")
p = torch.softmax(p.to(HI_PRES_TORCH), dim=-1).to(dtype)
ref_out = torch.matmul(p, v[:, :, None, :, :])
ref_out = ref_out.view(q.shape)
ref_out.backward(d_out)
ref_dv, v.grad = v.grad.clone(), None
ref_dk, k.grad = k.grad.clone(), None
ref_dq, q.grad = q.grad.clone(), None
torch.cuda.synchronize()
with monit.section('Triton'):
assert q.dtype == dtype
tri_out = attention(q, k, v, causal, sm_scale).to(dtype)
monit.progress(0.5)
tri_out.backward(d_out)
monit.progress(0.9)
tri_dv, v.grad = v.grad.clone(), None # type: ignore
tri_dk, k.grad = k.grad.clone(), None # type: ignore
tri_dq, q.grad = q.grad.clone(), None # type: ignore
torch.cuda.synchronize()
with monit.section('Test') as s:
# compare
passed = True
if not torch.allclose(tri_out, ref_out, atol=1e-2, rtol=0.):
abs_err, rel_err = _calc_abs_rel_error(ref_out, tri_out)
logger.log(('[FAILED]', logger.Text.danger), f' Out mismatch {abs_err} {rel_err}')
passed = False
rtol = 1e-1
if not torch.allclose(tri_dq, ref_dq, atol=1e-2, rtol=rtol):
abs_err, rel_err = _calc_abs_rel_error(ref_dq, tri_dq)
logger.log(('[FAILED]', logger.Text.danger), f' dQ mismatch {abs_err} {rel_err}')
passed = False
if not torch.allclose(tri_dv, ref_dv, atol=1e-2, rtol=rtol):
abs_err, rel_err = _calc_abs_rel_error(ref_dv, tri_dv)
logger.log(('[FAILED]', logger.Text.danger), f' dV mismatch {abs_err} {rel_err}')
passed = False
if not torch.allclose(tri_dk, ref_dk, atol=1e-2, rtol=rtol):
abs_err, rel_err = _calc_abs_rel_error(ref_dk, tri_dk)
logger.log(('[FAILED]', logger.Text.danger), f' dK mismatch {abs_err} {rel_err}')
passed = False
if passed:
logger.log('[PASSED]', logger.Text.success)
s.success = True
else:
s.success = False
torch.cuda.synchronize()
def _perf_triton_fn(*, device, dtype, batch_size, k_heads, n_groups, seq_len, d_head, causal):
"""
Get a partial function to test performance of our implementation
"""
q = torch.randn((batch_size, k_heads * n_groups, seq_len, d_head), dtype=dtype, device=device, requires_grad=True)
k = torch.randn((batch_size, k_heads, seq_len, d_head), dtype=dtype, device=device, requires_grad=True)
v = torch.randn((batch_size, k_heads, seq_len, d_head), dtype=dtype, device=device, requires_grad=True)
sm_scale = d_head ** -0.5
return lambda: attention(q, k, v, causal, sm_scale)
def _perf_flash(*, batch_size, k_heads, n_groups, seq_len, d_head, causal, device, dtype):
"""
Get a partial function to test performance of original flash implementation
"""
q = torch.randn((batch_size, seq_len, k_heads * n_groups, d_head), dtype=dtype, device=device, requires_grad=True)
k = torch.randn((batch_size, seq_len, k_heads, d_head), dtype=dtype, device=device, requires_grad=True)
v = torch.randn((batch_size, seq_len, k_heads, d_head), dtype=dtype, device=device, requires_grad=True)
from flash_attn import flash_attn_func
return lambda: flash_attn_func(q, k, v, causal=causal)
def measure_performance(name, fn, *, batch_size, k_heads, n_groups, seq_len, d_head, causal, is_bwd: bool):
"""
### Measure the speed
"""
if is_bwd:
o = fn()
do = torch.randn_like(o)
fn = lambda: o.backward(do, retain_graph=True)
ms = triton.testing.do_bench(fn)
flops_per_matmul = 2.0 * batch_size * k_heads * n_groups * seq_len * seq_len * d_head
total_flops = 2 * flops_per_matmul
if causal:
total_flops *= 0.5
if is_bwd:
total_flops *= 2.5 # 2.0(bwd) + 0.5(recompute)
tf_ps = total_flops * 1e-12 / (ms * 1e-3)
logger.log((f'{name}', logger.Text.key), ': ', f'{ms :,.1f}ms', ' ', f'{tf_ps :,.2f}TFps')
def main():
device = torch.device('cuda:0')
torch.cuda.set_device(device)
dtype = torch.float16
# only works on post-Ampere GPUs right now
test_fwd_bwd(1, 4, 1, 2048, 2048, 128, True, dtype=dtype, device=device)
test_fwd_bwd(16, 32, 8, 2001, 4001, 128, False, dtype=dtype, device=device)
test_fwd_bwd(4, 32, 8, 2048, 1024, 128, False, dtype=dtype, device=device)
test_fwd_bwd(4, 32, 8, 2001, 4001, 128, True, dtype=dtype, device=device)
_conf = {
'batch_size': 16,
'k_heads': 8,
'n_groups': 4,
'seq_len': 2048,
'd_head': 128,
}
for _causal in [False, True]:
for is_bwd in [False, True]:
logger.log(f'{"Causal" if _causal else "Non-causal"} {" Backward" if is_bwd else ""}', logger.Text.title)
measure_performance(f'flash', _perf_flash(causal=_causal, device=device, dtype=dtype, **_conf),
is_bwd=is_bwd,
causal=_causal, **_conf)
measure_performance(f'triton', _perf_triton_fn(causal=_causal, device=device, dtype=dtype, **_conf),
is_bwd=is_bwd,
causal=_causal, **_conf)
if __name__ == "__main__":
main()
labml_nn/transformers/fnet/__init__.py
+91
View File
@@ -0,0 +1,91 @@
"""
---
title: "FNet: Mixing Tokens with Fourier Transforms"
summary: >
This is an annotated implementation/tutorial of FNet in PyTorch.
---
# FNet: Mixing Tokens with Fourier Transforms
This is a [PyTorch](https://pytorch.org) implementation of the paper
[FNet: Mixing Tokens with Fourier Transforms](https://arxiv.org/abs/2105.03824).
This paper replaces the [self-attention layer](../mha.html) with two
[Fourier transforms](https://en.wikipedia.org/wiki/Discrete_Fourier_transform) to
*mix* tokens.
This is a $7 \times$ more efficient than self-attention.
The accuracy loss of using this over self-attention is about 92% for
[BERT](https://paperswithcode.com/method/bert) on
[GLUE benchmark](https://paperswithcode.com/dataset/glue).
## Mixing tokens with two Fourier transforms
We apply Fourier transform along the hidden dimension (embedding dimension)
and then along the sequence dimension.
$$
\mathcal{R}\big(\mathcal{F}_\text{seq} \big(\mathcal{F}_\text{hidden} (x) \big) \big)
$$
where $x$ is the embedding input, $\mathcal{F}$ stands for the fourier transform and
$\mathcal{R}$ stands for the real component in complex numbers.
This is very simple to implement on PyTorch - just 1 line of code.
The paper suggests using a precomputed DFT matrix and doing matrix multiplication to get the
Fourier transformation.
Here is [the training code](experiment.html) for using a FNet based model for classifying
[AG News](https://paperswithcode.com/dataset/ag-news).
"""
from typing import Optional
import torch
from torch import nn
class FNetMix(nn.Module):
"""
## FNet - Mix tokens
This module simply implements
$$
\mathcal{R}\big(\mathcal{F}_\text{seq} \big(\mathcal{F}_\text{hidden} (x) \big) \big)
$$
The structure of this module is made similar to a [standard attention module](../mha.html) so that we can simply
replace it.
"""
def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, mask: Optional[torch.Tensor] = None):
"""
The [normal attention module](../mha.html) can be fed with different token embeddings for
$\text{query}$,$\text{key}$, and $\text{value}$ and a mask.
We follow the same function signature so that we can replace it directly.
For FNet mixing, $$x = \text{query} = \text{key} = \text{value}$$ and masking is not possible.
Shape of `query` (and `key` and `value`) is `[seq_len, batch_size, d_model]`.
"""
# $\text{query}$,$\text{key}$, and $\text{value}$ all should be equal to $x$ for token mixing
assert query is key and key is value
# Token mixing doesn't support masking. i.e. all tokens will see all other token embeddings.
assert mask is None
# Assign to `x` for clarity
x = query
# Apply the Fourier transform along the hidden (embedding) dimension
# $$\mathcal{F}_\text{hidden} (x)$$
#
# The output of the Fourier transform is a tensor of
# [complex numbers](https://pytorch.org/docs/stable/complex_numbers.html).
fft_hidden = torch.fft.fft(x, dim=2)
# Apply the Fourier transform along the sequence dimension
# $$\mathcal{F}_\text{seq} \big(\mathcal{F}_\text{hidden} (x) \big)$$
fft_seq = torch.fft.fft(fft_hidden, dim=0)
# Get the real component
# $$\mathcal{R}\big(\mathcal{F}_\text{seq} \big(\mathcal{F}_\text{hidden} (x) \big) \big)$$
return torch.real(fft_seq)
labml_nn/transformers/fnet/experiment.py
+154
View File
@@ -0,0 +1,154 @@
"""
---
title: FNet Experiment
summary: This experiment trains a FNet based model on AG News dataset.
---
# [FNet](index.html) Experiment
This is an annotated PyTorch experiment to train a [FNet model](index.html).
This is based on
[general training loop and configurations for AG News classification task](../../experiments/nlp_classification.html).
"""
import torch
from torch import nn
from labml import experiment
from labml.configs import option
from labml_nn.experiments.nlp_classification import NLPClassificationConfigs
from labml_nn.transformers import Encoder
from labml_nn.transformers import TransformerConfigs
class TransformerClassifier(nn.Module):
"""
# Transformer based classifier model
"""
def __init__(self, encoder: Encoder, src_embed: nn.Module, generator: nn.Linear):
"""
* `encoder` is the transformer [Encoder](../models.html#Encoder)
* `src_embed` is the token
[embedding module (with positional encodings)](../models.html#EmbeddingsWithLearnedPositionalEncoding)
* `generator` is the [final fully connected layer](../models.html#Generator) that gives the logits.
"""
super().__init__()
self.src_embed = src_embed
self.encoder = encoder
self.generator = generator
def forward(self, x: torch.Tensor):
# Get the token embeddings with positional encodings
x = self.src_embed(x)
# Transformer encoder
x = self.encoder(x, None)
# Get logits for classification.
#
# We set the `[CLS]` token at the last position of the sequence.
# This is extracted by `x[-1]`, where `x` is of
# shape `[seq_len, batch_size, d_model]`
x = self.generator(x[-1])
# Return results
# (second value is for state, since our trainer is used with RNNs also)
return x, None
class Configs(NLPClassificationConfigs):
"""
## Configurations
This inherits from
[`NLPClassificationConfigs`](../../experiments/nlp_classification.html)
"""
# Classification model
model: TransformerClassifier
# Transformer
transformer: TransformerConfigs
@option(Configs.transformer)
def _transformer_configs(c: Configs):
"""
### Transformer configurations
"""
# We use our
# [configurable transformer implementation](../configs.html#TransformerConfigs)
conf = TransformerConfigs()
# Set the vocabulary sizes for embeddings and generating logits
conf.n_src_vocab = c.n_tokens
conf.n_tgt_vocab = c.n_tokens
#
return conf
@option(TransformerConfigs.encoder_attn)
def fnet_mix():
"""
Create `FNetMix` module that can replace the self-attention in
[transformer encoder layer](../models.html#TransformerLayer)
.
"""
from labml_nn.transformers.fnet import FNetMix
return FNetMix()
@option(Configs.model)
def _model(c: Configs):
"""
Create classification model
"""
m = TransformerClassifier(c.transformer.encoder,
c.transformer.src_embed,
nn.Linear(c.d_model, c.n_classes)).to(c.device)
return m
def main():
# Create experiment
experiment.create(name="fnet")
# Create configs
conf = Configs()
# Override configurations
experiment.configs(conf, {
# Use world level tokenizer
'tokenizer': 'basic_english',
# Train for $32$ epochs
'epochs': 32,
# Switch between training and validation for $10$ times
# per epoch
'inner_iterations': 10,
# Transformer configurations (same as defaults)
'transformer.d_model': 512,
'transformer.ffn.d_ff': 2048,
'transformer.n_heads': 8,
'transformer.n_layers': 6,
# Use [FNet](index.html) instead of self-a
# ttention
'transformer.encoder_attn': 'fnet_mix',
# Use [Noam optimizer](../../optimizers/noam.html)
'optimizer.optimizer': 'Noam',
'optimizer.learning_rate': 1.,
})
# Set models for saving and loading
experiment.add_pytorch_models({'model': conf.model})
# Start the experiment
with experiment.start():
# Run training
conf.run()
#
if __name__ == '__main__':
main()
labml_nn/transformers/fnet/readme.md
+12
View File
@@ -0,0 +1,12 @@
# [FNet: Mixing Tokens with Fourier Transforms](https://nn.labml.ai/transformers/fnet/index.html)
This is a [PyTorch](https://pytorch.org) implementation of the paper
[FNet: Mixing Tokens with Fourier Transforms](https://arxiv.org/abs/2105.03824).
This paper replaces the [self-attention layer](https://nn.labml.ai/transformers//mha.html) with two
[Fourier transforms](https://en.wikipedia.org/wiki/Discrete_Fourier_transform) to
*mix* tokens.
This is a 7X more efficient than self-attention.
The accuracy loss of using this over self-attention is about 92% for
[BERT](https://paperswithcode.com/method/bert) on
[GLUE benchmark](https://paperswithcode.com/dataset/glue).
labml_nn/transformers/glu_variants/__init__.py
+13
View File
@@ -0,0 +1,13 @@
"""
---
title: Gated Linear Units and Variants
summary: >
Train an auto-regressive transformer with Gated Linear Units and variants
for the position-wise feedforward network (FFN).
---
# Gated Linear Units and Variants
* [Experiment that uses `labml.configs`](experiment.html)
* [Simpler version from scratch](simple.html)
"""
labml_nn/transformers/glu_variants/experiment.py
+132
View File
@@ -0,0 +1,132 @@
"""
---
title: Gated Linear Units and Variants
summary: >
Train an auto-regressive transformer with Gated Linear Units and variants
for the position-wise feedforward network (FFN).
---
# Gated Linear Units and Variants
This trains a simple [transformer](../../) model for auto-regression.
We try different variants for the [position-wise feedforward network](../feed_forward).
The reusable & configurable are defined in [`configs.py`](configs.html).
"""
import torch
from labml import experiment
from labml.configs import option
from labml.utils.pytorch import get_modules
from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs
from labml_nn.transformers import Encoder, Generator, TransformerConfigs
from labml_nn.transformers.utils import subsequent_mask
from torch import nn
class AutoregressiveModel(nn.Module):
"""
## Auto regressive model
"""
def __init__(self, src_embed: nn.Module, encoder: Encoder, generator: Generator):
super().__init__()
# Token embedding module
self.src_embed = src_embed
# Transformer based encoder
self.encoder = encoder
# Next token generation layer;
# this give logits of the the next token
self.generator = generator
# This will be initialized on the first call
self.src_mask = None
def forward(self, src: torch.Tensor):
# Create subsequent mask, so that the transformer can only pay attention to past tokens.
if self.src_mask is None or self.src_mask.size(0) != len(src):
self.src_mask = subsequent_mask(len(src)).to(src.device)
# Embed the tokens (`src`) and run it through the the transformer
res = self.encoder(self.src_embed(src), self.src_mask)
# Generate logits of the next token
return self.generator(res), None
class Configs(NLPAutoRegressionConfigs):
"""
## Configurations
The default configs can and will be over-ridden when we start the experiment
"""
transformer: TransformerConfigs
model: AutoregressiveModel
@option(Configs.model)
def autoregressive_model(c: Configs):
"""
Initialize the auto-regressive model
"""
m = AutoregressiveModel(c.transformer.src_embed, c.transformer.encoder, c.transformer.generator)
return m.to(c.device)
@option(Configs.transformer)
def transformer_c(c: Configs):
"""
Initialize the [configurable transformer](../configs.html) encoder for our autoregressive model.
"""
tc = TransformerConfigs()
tc.n_src_vocab = c.n_tokens
tc.n_tgt_vocab = c.n_tokens
return tc
def main():
# Create experiment
experiment.create(name="glu_variants")
# Create configs
conf = Configs()
# Load configurations
experiment.configs(conf,
# A dictionary of configurations to override
{'tokenizer': 'character',
'prompt_separator': '',
'prompt': 'It is ',
'text': 'tiny_shakespeare',
'optimizer.optimizer': 'Noam',
'optimizer.learning_rate': 1.,
'optimizer.d_model': 256,
'seq_len': 1024,
'epochs': 128,
'batch_size': 6,
'inner_iterations': 10,
# GLU Variant, one of GLU, Bilinear, ReGLU, GEGLU, SwiGLU
#
# These are defined in the [configurable FFN](../configs.html#FFN)
# implementation
'transformer.ffn.glu_variant': 'Bilinear',
# Transformer configurations
'transformer.d_model': 256,
'transformer.ffn.d_ff': 1024,
'transformer.n_heads': 8,
'transformer.n_layers': 6})
# This is needed to initialize models
conf.n_tokens = conf.text.n_tokens
# Set models for saving and loading
experiment.add_pytorch_models(get_modules(conf))
# Start the experiment
with experiment.start():
# `TrainValidConfigs.run`
conf.run()
if __name__ == '__main__':
main()
labml_nn/transformers/glu_variants/simple.ipynb
+216
View File
@@ -0,0 +1,216 @@
{
"nbformat": 4,
"nbformat_minor": 0,
"metadata": {
"colab": {
"name": "Gated Linear Units and Variants",
"provenance": [],
"collapsed_sections": [],
"toc_visible": true
},
"kernelspec": {
"name": "python3",
"display_name": "Python 3"
},
"accelerator": "GPU"
},
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "AYV_dMVDxyc2"
},
"source": [
"[![Github](https://img.shields.io/github/stars/labmlai/annotated_deep_learning_paper_implementations?style=social)](https://github.com/labmlai/annotated_deep_learning_paper_implementations)\n",
"[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/transformers/glu_variants/simple.ipynb) \n",
"\n",
"## Gated Linear Units and Variants\n",
"\n",
"This trains a simple [transformer](https://nn.labml.ai/transformers/) model for auto-regression.\n",
"We try different variants for the [position-wise feedforward network](https://nn.labml.ai/transformers/feed_forward.html).\n",
"\n",
"Annotated trainer code is at [`simple.py`](https://nn.labml.ai/transformers/glu_variants/simple.html)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "AahG_i2y5tY9"
},
"source": [
"Install the `labml-nn` package"
]
},
{
"cell_type": "code",
"metadata": {
"id": "ZCzmCrAIVg0L",
"colab": {
"base_uri": "https://localhost:8080/"
},
"outputId": "2de76edb-9911-496d-9f8c-281dad6f5680"
},
"source": [
"!pip install labml-nn"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "SE2VUQ6L5zxI"
},
"source": [
"Imports"
]
},
{
"cell_type": "code",
"metadata": {
"id": "0hJXx_g0wS2C"
},
"source": [
"import dataclasses\n",
"\n",
"import torch\n",
"import torch.nn as nn\n",
"from labml import experiment\n",
"from labml_nn.transformers.glu_variants.simple import Configs, Trainer"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "Lpggo0wM6qb-"
},
"source": [
"Create an experiment"
]
},
{
"cell_type": "code",
"metadata": {
"id": "bFcr9k-l4cAg"
},
"source": [
"experiment.create(name=\"glu_variants\")"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "-OnHLi626tJt"
},
"source": [
"Initialize configurations"
]
},
{
"cell_type": "code",
"metadata": {
"id": "Piz0c5f44hRo"
},
"source": [
"conf = Configs()"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "wwMzCqpD6vkL"
},
"source": [
"Set experiment configurations and assign a configurations dictionary to override configurations"
]
},
{
"cell_type": "code",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 17
},
"id": "e6hmQhTw4nks",
"outputId": "77eca625-7205-49ea-f275-23f2710c4d84"
},
"source": [
"experiment.configs(dataclasses.asdict(conf))"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "DHyNvXfnzeWQ"
},
"source": [
"Create [`Trainer`](https://nn.labml.ai/transformers/glu_variants/simple.html)"
]
},
{
"cell_type": "code",
"metadata": {
"id": "59ZeTv5SzcVe"
},
"source": [
"trainer = Trainer(conf)"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "EvI7MtgJ61w5"
},
"source": [
"Set PyTorch models for loading and saving"
]
},
{
"cell_type": "code",
"metadata": {
"id": "GDlt7dp-5ALt"
},
"source": [
"experiment.add_pytorch_models({'model': trainer.model})"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "KJZRf8527GxL"
},
"source": [
"Start the experiment and run the training loop."
]
},
{
"cell_type": "code",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 255
},
"id": "aIAWo7Fw5DR8",
"outputId": "18b8b334-f9e7-458b-f900-5828b4f9a5c8"
},
"source": [
"with experiment.start():\n",
" trainer.train()"
],
"outputs": [],
"execution_count": null
}
]
}
labml_nn/transformers/glu_variants/simple.py
+303
View File
@@ -0,0 +1,303 @@
"""
---
title: Gated Linear Units and Variants
summary: >
Train an auto-regressive transformer with Gated Linear Units and variants
for the position-wise feedforward network (FFN).
---
# Gated Linear Units and Variants
This trains a simple [transformer](../../) model for auto-regression.
We try different variants for the [position-wise feedforward network](../feed_forward).
*This is a simpler implementation that doesn't use [`labml.configs`](experiment.html) module.
We decided to write a simpler implementation to make it easier for readers who are not familiar.*
[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/transformers/glu_variants/simple.ipynb)
"""
import dataclasses
import torch
from labml import experiment, lab, tracker, monit, logger
from labml.logger import Text
from labml.utils.download import download_file
from labml_nn.experiments.nlp_autoregression import transpose_batch
from labml_nn.optimizers.noam import Noam
from labml_nn.transformers import Encoder, MultiHeadAttention
from labml_nn.transformers.feed_forward import FeedForward
from labml_nn.transformers.models import EmbeddingsWithPositionalEncoding, TransformerLayer
from labml_nn.transformers.utils import subsequent_mask
from torch import nn
from torch.utils.data import Dataset, DataLoader
class AutoregressiveModel(nn.Module):
"""
## Auto regressive model
"""
def __init__(self, src_embed: nn.Module, encoder: Encoder, generator: nn.Module):
super().__init__()
# Token embedding module
self.src_embed = src_embed
# Transformer based encoder
self.encoder = encoder
# Next token generation layer;
# this gives logits of the the next token
self.generator = generator
# This will be initialized on the first call
self.src_mask = None
def forward(self, src: torch.Tensor):
# Create subsequent mask, so that the transformer can only pay attention to past tokens.
if self.src_mask is None or self.src_mask.size(0) != len(src):
self.src_mask = subsequent_mask(len(src)).to(src.device)
# Embed the tokens (`src`) and run it through the the transformer
res = self.encoder(self.src_embed(src), self.src_mask)
# Generate logits of the next token
return self.generator(res)
@dataclasses.dataclass
class Configs:
"""
### Configurations
"""
d_model: int = 512
seq_len: int = 128
batch_size: int = 32
n_layers: int = 6
n_heads: int = 8
dropout: float = 0.1
d_ff: int = 2048
glu_variant: str = 'GLU'
epochs: int = 5
grad_norm_clip: float = 0.5
class TinyShakespeareDataset(Dataset):
"""
### Tiny Shakespeare Dataset
"""
def __init__(self, seq_len: int):
# Location of the text file
path = lab.get_data_path() / 'tiny_shakespeare.txt'
# Download the file
download_file('https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt', path)
# Read the downloaded file
with open(str(path), 'r') as f:
text = f.read()
# Extract the characters
chars = list(set(text))
# Character to id (integer) map
self.stoi = {c: i for i, c in enumerate(chars)}
# Id to character map
self.itos = {i: c for i, c in enumerate(chars)}
# Length of a training sample
self.seq_len = seq_len
# Data in the form of a tensor of ids
self.data = self.text_to_i(text)
def text_to_i(self, text: str):
"""
Transform the text into a tensor of ids
"""
return torch.tensor([self.stoi[c] for c in text], dtype=torch.long)
def __len__(self):
"""
Number of samples in the dataset.
*This will read the dataset `seq_len` times in a single epoch.*
"""
return len(self.data) - self.seq_len - 1
def __getitem__(self, idx):
"""
Return a sample
"""
return self.data[idx:idx + self.seq_len], self.data[idx + 1:idx + self.seq_len + 1]
class Trainer:
"""
## Trainer
"""
def __init__(self, configs: Configs):
# Get the device
self.device = torch.device('cpu')
if torch.cuda.is_available():
self.device = torch.device('cuda:0')
# Initialize the dataset
self.dataset = TinyShakespeareDataset(configs.seq_len)
# Initialize the dataloader
self.dataloader = DataLoader(self.dataset,
batch_size=configs.batch_size,
collate_fn=transpose_batch,
shuffle=True)
# FFN with Gated Linear Unit
# $$FFN_{GLU}(x)(x, W_1, V, W_2) = (\sigma(x W_1) \otimes x V) W_2$$
if configs.glu_variant == 'GLU':
ffn = FeedForward(configs.d_model, configs.d_ff, configs.dropout, nn.Sigmoid(), True, False, False, False)
# FFN with Bilinear hidden layer
# $$FFN_{Bilinear}(x)(x, W_1, V, W_2) = (x W_1 \otimes x V) W_2$$
elif configs.glu_variant == 'Bilinear':
ffn = FeedForward(configs.d_model, configs.d_ff, configs.dropout, nn.Identity(), True, False, False, False)
# FFN with ReLU gate
# $$FFN_{ReGLU}(x)(x, W_1, V, W_2) = (\max(0, x W_1) \otimes x V) W_2$$
elif configs.glu_variant == 'ReGLU':
ffn = FeedForward(configs.d_model, configs.d_ff, configs.dropout, nn.ReLU(), True, False, False, False)
# FFN with GELU gate
# $$FFN_{GEGLU}(x)(x, W_1, V, W_2) = (\text{GELU}(x W_1) \otimes x V) W_2$$
elif configs.glu_variant == 'GEGLU':
ffn = FeedForward(configs.d_model, configs.d_ff, configs.dropout, nn.GELU(), True, False, False, False)
# FFN with Swish gate
# $$FFN_{SwiGLU}(x)(x, W_1, V, W_2) = (\text{Swish}_1(x W_1) \otimes x V) W_2$$
# where $\text{Swish}_\beta(x) = x \sigma(\beta x)$
elif configs.glu_variant == 'SwiGLU':
ffn = FeedForward(configs.d_model, configs.d_ff, configs.dropout, nn.SiLU(), True, False, False, False)
# FFN with ReLU activation
# $$FFN_{ReLU}(x)(x, W_1, W_2, b_1, b_2) = \text{ReLU}_1(x W_1 + b_1) W_2 + b_2$$
elif configs.glu_variant == 'ReLU':
ffn = FeedForward(configs.d_model, configs.d_ff, configs.dropout, nn.ReLU())
# FFN with ReLU activation
# $$FFN_{GELU}(x)(x, W_1, W_2, b_1, b_2) = \text{GELU}_1(x W_1 + b_1) W_2 + b_2$$
elif configs.glu_variant == 'GELU':
ffn = FeedForward(configs.d_model, configs.d_ff, configs.dropout, nn.GELU())
else:
raise ValueError(f'Unknown variant {configs.glu_variant}')
# Number of different characters
n_chars = len(self.dataset.stoi)
# Initialize [Multi-Head Attention module](../mha.html)
mha = MultiHeadAttention(configs.n_heads, configs.d_model, configs.dropout)
# Initialize the [Transformer Block](../models.html#TransformerLayer)
transformer_layer = TransformerLayer(d_model=configs.d_model, self_attn=mha, src_attn=None,
feed_forward=ffn, dropout_prob=configs.dropout)
# Initialize the model with an
# [embedding layer](../models.html#EmbeddingsWithPositionalEncoding)
# (with fixed positional encoding)
# [transformer encoder](../models.html#Encoder) and
# a linear layer to generate logits.
self.model = AutoregressiveModel(EmbeddingsWithPositionalEncoding(configs.d_model, n_chars),
Encoder(transformer_layer, configs.n_layers),
nn.Linear(configs.d_model, n_chars))
# Move the model to the current device
self.model.to(self.device)
# Initialize [Noam optimizer](../../optimizers/noam.html)
self.optimizer = Noam(self.model.parameters(), lr=1.0, warmup=2_000, d_model=configs.d_model)
# Cross-entropy loss
self.loss_func = nn.CrossEntropyLoss()
# Number of training epochs;
# *note that our dataset definition repeats the data `seq_len` times in a single epoch*
self.epochs = configs.epochs
# Gradient clipping norm
self.grad_norm_clip = configs.grad_norm_clip
# Set tracker configurations
tracker.set_scalar("loss.*", True)
def sample(self):
"""
### Sampling function to generate samples periodically while training
"""
# Starting prompt
prompt = 'It is'
# Collect output for printing
log = [(prompt, Text.subtle)]
# Sample 25 tokens
for i in monit.iterate('Sample', 25):
# Tokenize the prompt
data = self.dataset.text_to_i(prompt).unsqueeze(-1)
data = data.to(self.device)
# Get the model output
output = self.model(data)
# Get the model prediction (greedy)
output = output.argmax(dim=-1).squeeze()
# Add the prediction to prompt
prompt += self.dataset.itos[output[-1].item()]
# Add the prediction for logging
log += [(self.dataset.itos[output[-1].item()], Text.value)]
# Print the sampled output
logger.log(log)
def train(self):
"""
### Train the model
"""
# Loop for the given number of epochs
for _ in monit.loop(self.epochs):
# Iterate over the minibatches
for i, batch in monit.enum('Train', self.dataloader):
# Move data to the device
data, target = batch[0].to(self.device), batch[1].to(self.device)
# Set tracker step, as the number of characters trained on
tracker.add_global_step(data.shape[0] * data.shape[1])
# Set model state to training
self.model.train()
# Evaluate the model
output = self.model(data)
# Calculate loss
loss = self.loss_func(output.view(-1, output.shape[-1]), target.view(-1))
# Log the loss
tracker.add("loss.train", loss)
# Calculate gradients
loss.backward()
# Clip gradients
torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=self.grad_norm_clip)
# Take optimizer step
self.optimizer.step()
# Log the model parameters and gradients
if (i + 1) % 100 == 0:
tracker.add('model', self.model)
# Clear the gradients
self.optimizer.zero_grad()
# Generate a sample
if (i + 1) % 100 == 0:
self.model.eval()
with torch.no_grad():
self.sample()
# Save the tracked metrics
if (i + 1) % 10 == 0:
tracker.save()
def main():
# Create experiment
experiment.create(name="glu_variants")
# Create configs
configs = Configs()
# Load configurations
experiment.configs(dataclasses.asdict(configs))
# Create trainer
trainer = Trainer(configs)
# Set models for training and loading
experiment.add_pytorch_models({'model': trainer.model})
# Start the experiment
with experiment.start():
# Train the model
trainer.train()
if __name__ == '__main__':
main()
labml_nn/transformers/gmlp/__init__.py
+157
View File
@@ -0,0 +1,157 @@
"""
---
title: Pay Attention to MLPs (gMLP)
summary: >
This is an annotated implementation/tutorial of Pay Attention to MLPs (gMLP) in PyTorch.
---
# Pay Attention to MLPs (gMLP)
This is a [PyTorch](https://pytorch.org) implementation of the paper
[Pay Attention to MLPs](https://arxiv.org/abs/2105.08050).
This paper introduces a Multilayer Perceptron (MLP) based architecture with gating,
which they name **gMLP**. It consists of a stack of $L$ *gMLP* blocks.
Here is [the training code](experiment.html) for a gMLP model based autoregressive model.
"""
from typing import Optional
import torch
from torch import nn
class GMLPBlock(nn.Module):
"""
## gMLP Block
Each block does the following transformations to input embeddings
$X \in \mathbb{R}^{n \times d}$ where $n$ is the sequence length
and $d$ is the dimensionality of the embeddings:
\begin{align}
Z &= \sigma(XU) \\
\tilde{Z} &= s(Z) \\
Y &= \tilde{Z}V \\
\end{align}
where $V$ and $U$ are learnable projection weights.
$s(\cdot)$ is the Spacial Gating Unit defined below.
Output dimensionality of $s(\cdot)$ will be half of $Z$.
$\sigma$ is an activation function such as
[GeLU](https://pytorch.org/docs/stable/generated/torch.nn.GELU.html).
"""
def __init__(self, d_model: int, d_ffn: int, seq_len: int):
"""
* `d_model` is the dimensionality ($d$) of $X$
* `d_ffn` is the dimensionality of $Z$
* `seq_len` is the length of the token sequence ($n$)
"""
super().__init__()
# Normalization layer fro Pre-Norm
self.norm = nn.LayerNorm([d_model])
# Activation function $\sigma$
self.activation = nn.GELU()
# Projection layer for $Z = \sigma(XU)$
self.proj1 = nn.Linear(d_model, d_ffn)
# Spacial Gating Unit $s(\cdot)$
self.sgu = SpacialGatingUnit(d_ffn, seq_len)
# Projection layer for $Y = \tilde{Z}V$
self.proj2 = nn.Linear(d_ffn // 2, d_model)
# Embedding size (required by [Encoder](../models.html#Encoder).
# We use the encoder module from transformer architecture and plug
# *gMLP* block as a replacement for the [Transformer Layer](../models.html#Encoder).
self.size = d_model
def forward(self, *, x: torch.Tensor, mask: Optional[torch.Tensor] = None):
"""
* `x` is the input embedding tensor $X$ of shape `[seq_len, batch_size, d_model]`
* `mask` is a boolean mask of shape `[seq_len, seq_len, 1]` that controls the visibility of tokens
among each other.
"""
# Keep a copy for shortcut connection
shortcut = x
# Normalize $X$
x = self.norm(x)
# Projection and activation $Z = \sigma(XU)$
z = self.activation(self.proj1(x))
# Spacial Gating Unit $\tilde{Z} = s(Z)$
z = self.sgu(z, mask)
# Final projection $Y = \tilde{Z}V$
z = self.proj2(z)
# Add the shortcut connection
return z + shortcut
class SpacialGatingUnit(nn.Module):
"""
## Spatial Gating Unit
$$s(Z) = Z_1 \odot f_{W,b}(Z_2)$$
where $f_{W,b}(Z) = W Z + b$ is a linear transformation along the sequence dimension,
and $\odot$ is element-wise multiplication.
$Z$ is split into to parts of equal size $Z_1$ and $Z_2$ along the channel dimension (embedding dimension).
"""
def __init__(self, d_z: int, seq_len: int):
"""
* `d_z` is the dimensionality of $Z$
* `seq_len` is the sequence length
"""
super().__init__()
# Normalization layer before applying $f_{W,b}(\cdot)$
self.norm = nn.LayerNorm([d_z // 2])
# Weight $W$ in $f_{W,b}(\cdot)$.
#
# The paper notes that it's important to initialize weights to small values and the bias to $1$,
# so that during the initial training $s(\cdot)$ is close to identity (apart from the split).
self.weight = nn.Parameter(torch.zeros(seq_len, seq_len).uniform_(-0.01, 0.01), requires_grad=True)
# Weight $b$ in $f_{W,b}(\cdot)$
#
# The paper notes that it's important to initialize bias to $1$.
self.bias = nn.Parameter(torch.ones(seq_len), requires_grad=True)
def forward(self, z: torch.Tensor, mask: Optional[torch.Tensor] = None):
"""
* `z` is the input $Z$ of shape `[seq_len, batch_size, d_z]`
* `mask` is is a boolean mask of shape `[seq_len, seq_len, 1]` that controls the visibility of tokens
among each other. The last dimension of size `1` is the batch, which we have in other transformer
implementations and was left for compatibility.
"""
# Get sequence length
seq_len = z.shape[0]
# Split $Z$ into $Z_1$ and $Z_2$
z1, z2 = torch.chunk(z, 2, dim=-1)
# Check mask
if mask is not None:
# `mask` has shape `[seq_len_q, seq_len_k, batch_size]`.
# The batch dimension should be of size `1` because this implementation supports
# only same mask for all samples in the batch.
assert mask.shape[0] == 1 or mask.shape[0] == seq_len
assert mask.shape[1] == seq_len
# Here we only support the same mask for all samples
assert mask.shape[2] == 1
# Remove the batch dimension
mask = mask[:, :, 0]
# Normalize $Z_2$ before $f_{W,b}(\cdot)$
z2 = self.norm(z2)
# Get the weight matrix; truncate if larger than `seq_len`
weight = self.weight[:seq_len, :seq_len]
# Apply mask to the weights.
#
# If $W_{i,j}$ is $0$ then $f_{W,b}(Z_2)_i$ will not get any information
# from token $j$.
if mask is not None:
weight = weight * mask
# $f_{W,b}(Z_2) = W Z_2 + b$
z2 = torch.einsum('ij,jbd->ibd', weight, z2) + self.bias[:seq_len, None, None]
# $Z_1 \odot f_{W,b}(Z_2)$
return z1 * z2
labml_nn/transformers/gmlp/experiment.py
+113
View File
@@ -0,0 +1,113 @@
"""
---
title: Pay Attention to MLPs (gMLP) Experiment
summary: This experiment trains a gMLP based model on Tiny Shakespeare dataset.
---
# [Pay Attention to MLPs (gMLP)](index.html) Experiment
This is an annotated PyTorch experiment to train a [gMLP model](index.html).
The paper also applies a Stochastic Depth regularization where some layers are removed randomly during training.
We have not implemented that here.
This is based on
[training loop and configurations for a simple transformer auto-regressive NLP task](../basic/autoregressive_experiment.html).
"""
from labml import experiment
from labml.configs import option
from labml_nn.transformers import TransformerConfigs
from labml_nn.transformers.basic.autoregressive_experiment import Configs as BasicAutoRegressionConfigs
from labml_nn.transformers.gmlp import GMLPBlock
class Configs(BasicAutoRegressionConfigs):
"""
## Configurations
This inherits from
[training loop and configurations for a simple transformer auto-regressive NLP task](../basic/autoregressive_transformer.html).
"""
# Transformer
transformer: TransformerConfigs = 'gMLP'
# gMLP Block
gmlp: GMLPBlock
# `d_ffn` for gMLP projection layer
d_ffn: int = 2048
@option(Configs.gmlp, 'gMLP')
def _gmlp_configs(c: Configs):
"""
### Create a gMLP block
"""
return GMLPBlock(c.d_model, c.d_ffn, c.seq_len)
@option(Configs.transformer, 'gMLP')
def _transformer_configs(c: Configs):
"""
### Transformer configurations
"""
# We use our
# [configurable transformer implementation](../configs.html#TransformerConfigs)
conf = TransformerConfigs()
# Set the vocabulary sizes for embeddings and generating logits
conf.n_src_vocab = c.n_tokens
conf.n_tgt_vocab = c.n_tokens
# Set model size
conf.d_model = c.d_model
# Replace the encoder layer with a gMLP layer
conf.encoder_layer = c.gmlp
return conf
def main():
# Create experiment
experiment.create(name="gMLP")
# Create configs
conf = Configs()
# Override configurations
experiment.configs(conf, {
# Use character level tokenizer
'tokenizer': 'character',
# Prompt separator is blank
'prompt_separator': '',
# Starting prompt for sampling
'prompt': 'It is ',
# Use Tiny Shakespeare dataset
'text': 'tiny_shakespeare',
# Use a context size of $256$
'seq_len': 256,
# Train for $128$ epochs
'epochs': 128,
# Batch size $32$
'batch_size': 32,
# Switch between training and validation for $10$ times
# per epoch
'inner_iterations': 10,
# Model size
'd_model': 512,
'd_ffn': 2048,
# Use [Noam optimizer](../../optimizers/noam.html)
'optimizer.optimizer': 'Noam',
'optimizer.learning_rate': 1.,
})
# Set models for saving and loading
experiment.add_pytorch_models({'model': conf.model})
# Start the experiment
with experiment.start():
# Run training
conf.run()
#
if __name__ == '__main__':
main()
labml_nn/transformers/gmlp/readme.md
+9
View File
@@ -0,0 +1,9 @@
# [Pay Attention to MLPs (gMLP)](https://nn.labml.ai/transformers/gmlp/index.html)
This is a [PyTorch](https://pytorch.org) implementation of the paper
[Pay Attention to MLPs](https://arxiv.org/abs/2105.08050).
This paper introduces a Multilayer Perceptron (MLP) based architecture with gating,
which they name **gMLP**. It consists of a stack of $L$ *gMLP* blocks.
Here is [the training code](https://nn.labml.ai/transformers/gmlp/experiment.html) for a gMLP model based autoregressive model.
labml_nn/transformers/gpt/__init__.py
+264
View File
@@ -0,0 +1,264 @@
"""
---
title: GPT
summary: >
Implementation/tutorial of GPT model and training code.
---
# GPT
This is a tutorial/implementation of
[OpenAI GPT architecture](https://openai.com/blog/better-language-models/)
in [PyTorch](https://pytorch.org).
We got a bunch of implementation details from
[minGPT](https://github.com/karpathy/minGPT)
by [@karpathy](https://twitter.com/karpathy).
This implementation also uses character tiny shakespeare dataset.
GPT model is essentially a standard transformer with a few tweaks.
GPT-2 and especially GPT-3 models are quite large and won't fit on a
single GPU and will need model parallelism.
This implementation doesn't even use data parallelism and is intended to be
more of a tutorial.
Main differences of this compared to a simple autoregressive transformer
are the parameter initialization, weight decay, and learning rate schedule.
For the transformer we reuse the
[existing labml/nn transformer implementation](../transformers/index.html).
Here's a notebook for training a GPT model on Tiny Shakespeare dataset.
[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/transformers/gpt/experiment.ipynb)
"""
import torch
from torch import nn
from labml import experiment
from labml.configs import option
from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs
from labml_nn.optimizers.configs import OptimizerConfigs
from labml_nn.transformers import TransformerConfigs, Encoder
from labml_nn.transformers.utils import subsequent_mask
class GPT(nn.Module):
"""
## GPT model
This consists of a token embedding layer, transformer encoder, and
a final linear layer that gives token logits.
"""
def __init__(self, encoder: Encoder, src_embed: nn.Module, generator: nn.Module):
"""
* `encoder` is the transformer [Encoder](../models.html#Encoder)
* `src_embed` is the token
[embedding module (with positional encodings)](../models.html#EmbeddingsWithLearnedPositionalEncoding)
* `generator` is the [final fully connected layer](../models.html#Generator) that gives the logits.
"""
super().__init__()
self.src_embed = src_embed
self.encoder = encoder
self.generator = generator
# The mask will be initialized on the first call
self.mask = None
def forward(self, x: torch.Tensor):
# Create subsequent mask if mask is not initialized
# or if the size of the mask is different
if self.mask is None or self.mask.size(0) != len(x):
# Subsequent mask, will mask out tokens from seeing future tokens
self.mask = subsequent_mask(len(x)).to(x.device)
# Get the token embeddings with positional encodings
x = self.src_embed(x)
# Transformer encoder
x = self.encoder(x, self.mask)
# Get logits
x = self.generator(x)
# Return results
# (second value is for state, since our trainer is used with RNNs also)
return x, None
class Configs(NLPAutoRegressionConfigs):
"""
## Configurations
This inherits from
[`NLPAutoRegressionConfigs`](../../experiments/nlp_autoregression.html#NLPAutoRegressionConfigs)
"""
# GPT model
model: GPT
# Transformer
transformer: TransformerConfigs
# Weight decay
weight_decay: float = 0.1
# Number of tokens for wamup
warmup_steps: int = 128 * 128 * 20
# Custom optimizer
optimizer = 'transformer_optimizer'
@option(Configs.transformer, 'GPT')
def _transformer_configs(c: Configs):
"""
### Transformer configurations
"""
# We use our
# [configurable transformer implementation](../configs.html#TransformerConfigs)
conf = TransformerConfigs()
# Set the vocabulary sizes for embeddings and generating logits
conf.n_src_vocab = c.n_tokens
conf.n_tgt_vocab = c.n_tokens
# GPT uses GELU activation for position wise feedforward
conf.ffn.activation = 'GELU'
#
return conf
def _init_weights(module):
"""
### Initialize weights
Weights of linear layers and embedding layers are initialized
to $\mathcal{N}(0, 0.02)$
instead of the default Xavier initialzation.
"""
if not isinstance(module, (nn.Linear, nn.Embedding)):
return
module.weight.data.normal_(mean=0.0, std=0.02)
# Initialize biases to $0$
if isinstance(module, nn.Linear) and module.bias is not None:
module.bias.data.zero_()
@option(Configs.model)
def _model(c: Configs):
"""
Create GPT model and initialize weights
"""
m = GPT(c.transformer.encoder,
c.transformer.src_embed,
c.transformer.generator).to(c.device)
# Apply custom weight initialization
m.apply(_init_weights)
return m
@option(NLPAutoRegressionConfigs.optimizer)
def transformer_optimizer(c: NLPAutoRegressionConfigs):
"""
### Create custom optimizer with weight decay
This code is taken from [minGPT](https://github.com/karpathy/minGPT).
This applies weight decay only to weights of linear layers.
"""
# Collect names of parameters to apply weight decay
decay = set()
for mn, m in c.model.named_modules():
for pn, p in m.named_parameters():
fpn = f'{mn}.{pn}' if mn else pn # full param name
if fpn.endswith('weight') and isinstance(m, nn.Linear):
decay.add(fpn)
# Get all the parameters
param_dict = {pn: p for pn, p in c.model.named_parameters()}
# Parameters that are not decayed
no_decay = set(param_dict.keys()) - decay
# create the pytorch optimizer object
opt_groups = [
{"params": [param_dict[pn] for pn in sorted(list(decay))], "weight_decay": c.weight_decay},
{"params": [param_dict[pn] for pn in sorted(list(no_decay))], "weight_decay": 0.0},
]
# Create a [configurable optimizer](../optimizers/configs.html#OptimizerConfigs),
# so that we can change these simply by passing
# a config dictionary.
optimizer = OptimizerConfigs()
# Set parameter groups for optimization.
optimizer.parameters = opt_groups
# Use [cosine decay optimizer](../optimizers/adam_warmup_cosine_decay.html).
# This is what GPT uses.
optimizer.optimizer = 'AdamWarmupCosineDecay'
# Set model embedding size, required if we use [Noam optimizer](../optimizers/noam.html)
# which has an exponential decay.
optimizer.d_model = c.d_model
# Set default weight decay.
# This is not required since we set the weight decay in the parameter groups.
optimizer.weight_decay = c.weight_decay
# GPT uses a maximum learning rate of $6 \times 10^{-4}$.
optimizer.learning_rate = 6e-4
# $\beta_1 = 0.9, \beta_2 = 0.95$
optimizer.betas = (0.9, 0.95)
# $\epsilon = 10^{-8}$
optimizer.eps = 1e-8
# Weight decay is decoupled from gradients
optimizer.weight_decouple = True
# Total number of optimization steps for learning rate cosine decay
optimizer.total_steps = c.epochs * len(c.text.train) // (c.batch_size * c.seq_len)
# Number of warmup optimization steps
optimizer.warmup = c.warmup_steps // (c.batch_size * c.seq_len)
return optimizer
def main():
# Create experiment
experiment.create(name="gpt")
# Create configs
conf = Configs()
# Override configurations
experiment.configs(conf, {
# Use character level tokenizer
'tokenizer': 'character',
# Prompt separator is blank
'prompt_separator': '',
# Starting prompt for sampling
'prompt': 'It is ',
# Use Tiny Shakespeare dataset
'text': 'tiny_shakespeare',
# Use a context size of $128$
'seq_len': 128,
# Train for $32$ epochs
'epochs': 32,
# Batch size $128$
'batch_size': 128,
# Switch between training and validation for $10$ times
# per epoch
'inner_iterations': 10,
# Transformer configurations
'transformer.d_model': 512,
'transformer.ffn.d_ff': 2048,
'transformer.n_heads': 8,
'transformer.n_layers': 6
})
# Set models for saving and loading
experiment.add_pytorch_models({'model': conf.model})
# Start the experiment
with experiment.start():
# Run training
conf.run()
#
if __name__ == '__main__':
main()
labml_nn/transformers/gpt/experiment.ipynb
+230
View File
@@ -0,0 +1,230 @@
{
"nbformat": 4,
"nbformat_minor": 0,
"metadata": {
"colab": {
"name": "GPT",
"provenance": [],
"collapsed_sections": []
},
"kernelspec": {
"name": "python3",
"display_name": "Python 3"
},
"accelerator": "GPU"
},
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "AYV_dMVDxyc2"
},
"source": [
"[![Github](https://img.shields.io/github/stars/labmlai/annotated_deep_learning_paper_implementations?style=social)](https://github.com/labmlai/annotated_deep_learning_paper_implementations)\n",
"[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/transformers/gpt/experiment.ipynb) \n",
"\n",
"## Training a model with GPT architecture\n",
"\n",
"This is an experiment training Tiny Shakespeare dataset with GPT architecture model."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "AahG_i2y5tY9"
},
"source": [
"Install the `labml-nn` package"
]
},
{
"cell_type": "code",
"metadata": {
"id": "ZCzmCrAIVg0L",
"colab": {
"base_uri": "https://localhost:8080/"
},
"outputId": "9c544df4-3fc7-4152-b50d-07919dbfb9de"
},
"source": [
"!pip install labml-nn"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "SE2VUQ6L5zxI"
},
"source": [
"Imports"
]
},
{
"cell_type": "code",
"metadata": {
"id": "0hJXx_g0wS2C"
},
"source": [
"from labml import experiment\n",
"from labml_nn.transformers.gpt import Configs"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "Lpggo0wM6qb-"
},
"source": [
"Create an experiment"
]
},
{
"cell_type": "code",
"metadata": {
"id": "bFcr9k-l4cAg"
},
"source": [
"experiment.create(name=\"gpt\")"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "-OnHLi626tJt"
},
"source": [
"Initialize [GPT configurations](https://nn.labml.ai/transformers/gpt/)"
]
},
{
"cell_type": "code",
"metadata": {
"id": "Piz0c5f44hRo"
},
"source": [
"conf = Configs()"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "wwMzCqpD6vkL"
},
"source": [
"Set experiment configurations and assign a configurations dictionary to override configurations"
]
},
{
"cell_type": "code",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 17
},
"id": "e6hmQhTw4nks",
"outputId": "018b39d7-7d84-4651-8b33-80de77d58ace"
},
"source": [
"experiment.configs(conf, {\n",
" # Use character level tokenizer\n",
" 'tokenizer': 'character',\n",
" # Prompt separator is blank\n",
" 'prompt_separator': '',\n",
" # Starting prompt for sampling\n",
" 'prompt': 'It is ',\n",
" # Use Tiny Shakespeare dataset\n",
" 'text': 'tiny_shakespeare',\n",
"\n",
" # Use a context size of $128$\n",
" 'seq_len': 128,\n",
" # Train for $32$ epochs\n",
" 'epochs': 32,\n",
" # Batch size $128$\n",
" 'batch_size': 128,\n",
" # Switch between training and validation for $10$ times\n",
" # per epoch\n",
" 'inner_iterations': 10,\n",
"\n",
" # Transformer configurations\n",
" 'transformer.d_model': 512,\n",
" 'transformer.ffn.d_ff': 2048,\n",
" 'transformer.n_heads': 8,\n",
" 'transformer.n_layers': 6\n",
"})"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "EvI7MtgJ61w5"
},
"source": [
"Set PyTorch models for loading and saving"
]
},
{
"cell_type": "code",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 459
},
"id": "GDlt7dp-5ALt",
"outputId": "e7768185-4a3c-4fec-f688-73c0463db015"
},
"source": [
"experiment.add_pytorch_models({'model': conf.model})"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "KJZRf8527GxL"
},
"source": [
"Start the experiment and run the training loop."
]
},
{
"cell_type": "code",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 1000
},
"id": "aIAWo7Fw5DR8",
"outputId": "d4ca5ae0-6a97-439e-a318-010cb3f288cd"
},
"source": [
"# Start the experiment\n",
"with experiment.start():\n",
" conf.run()"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "code",
"metadata": {
"id": "oBXXlP2b7XZO"
},
"source": [
""
],
"outputs": [],
"execution_count": null
}
]
}
labml_nn/transformers/hour_glass/__init__.py
+291
View File
@@ -0,0 +1,291 @@
"""
---
title: Hierarchical Transformers Are More Efficient Language Models
summary: >
This is an annotated implementation/tutorial of hourglass model in PyTorch.
---
# Hierarchical Transformers Are More Efficient Language Models
This is a [PyTorch](https://pytorch.org) implementation of the paper
[Hierarchical Transformers Are More Efficient Language Models](https://arxiv.org/abs/2110.13711).
This paper introduces a hierarchical transformer architecture to handle long sequences
efficiently. The first half of the transformer layers down-sample tokens and the second
half up-samples with direct skip connections between layers of the same resolution.
This is a little similar to [U-Net](../../diffusion/ddpm/unet.html) for vision tasks.
They try different up-sampling and down-sampling techniques and build a model
with the best performing up and down-sampling techniques which they call the
hourglass model.
Here we have implemented the simplest up-sampling and down-sampling techniques for simplicity.
We will consider adding more complex (and better performing) implementations later.
Here is [the training code](experiment.html) for the hourglass model.
"""
from typing import List
import torch
from torch import nn
from labml_nn.transformers import MultiHeadAttention, TransformerLayer
from labml_nn.transformers.feed_forward import FeedForward
from labml_nn.transformers.utils import subsequent_mask
class HourGlass(nn.Module):
"""
## Hourglass model
This model recursively adds layers to the middle while shortening the sequence by down-sampling.
The shortened sequence processed by another hourglass model is sandwiched between two normal transformer
layers. (A transformer layer has a [self-attention layer](../mha.html)
and a [position-wise feed-forward layer](../feed_forward.html)).
"""
def __init__(self, n_heads: int, d_model: int, dropout: float, d_ff: int, shortening_factors: List[int]):
"""
* `n_heads` is the number of heads in [multi-head attention layers](../mha.html)
* `d_model` is the size of the token embeddings
* `dropout` is the dropout probability
* `d_ff` is the dimensionality of the hidden layer in [position-wise feed-forward layers](../feed_forward.html)
* `shortening_factors` is the list of shortening factors
"""
super().__init__()
# The transformer layer before down-sampling
self.pre = TransformerLayer(d_model=d_model,
# [Multi-head attention layer](../mha.html)
self_attn=MultiHeadAttention(n_heads, d_model, dropout),
# [Position wise feed-forward layers](.. / feed_forward.html)
feed_forward=FeedForward(d_model, d_ff, dropout),
#
dropout_prob=dropout)
# Auto-regressive mask
self.mask = AutoregressiveMask()
# The shortening factor $k$ (or the down-sampling rate)
k = shortening_factors[0]
# We shift the tokens to the right by $k - 1$ steps to make sure
# information doesn't leak from the future tokens to past tokens
# as a result of down-sampling and up-sampling
self.shift_right = ShiftRight(k - 1)
# Shortening or the down-sampling layer. We use the simplest form - average pooling.
# The paper shows that attention based down sampling works best, which we haven't implemented yet.
self.shortening = AvgPoolShortening(k)
# If there are no more shortening (middle of the hourglass)
if len(shortening_factors) == 1:
# The center layer is another transformer layer
self.shortened = TransformerLayer(d_model=d_model,
self_attn=MultiHeadAttention(n_heads, d_model, dropout),
feed_forward=FeedForward(d_model, d_ff, dropout),
dropout_prob=dropout)
# Autoregressive mask
self.mask_short = AutoregressiveMask()
self.hour_glass = None
else:
# Insert another hourglass model recursively
self.hour_glass = HourGlass(n_heads, d_model, dropout, d_ff, shortening_factors[1:])
# Up-sampling layer. We use naive up-sampling for simplicity and the paper shows attention based up sampling
# works better.
self.up_sampling = NaiveUpSampling(k)
# The final transformer layer after up-sampling
self.post = TransformerLayer(d_model=d_model,
self_attn=MultiHeadAttention(n_heads, d_model, dropout),
feed_forward=FeedForward(d_model, d_ff, dropout),
dropout_prob=dropout)
def forward(self, x: torch.Tensor):
# Initial transformer layer
# $$x \leftarrow PreVanillaLayers(x)$$
x = self.pre(x=x, mask=self.mask(x))
# Shifting and shortening
# $$x' \leftarrow Shortening(ShiftRight(x,k−1),k)$$
x_short = self.shortening(self.shift_right(x))
# If we are at the center of the hourglass,
# $$\textbf{\small if } \text{\small E\scriptsize MPTY}(shorten\_factors) \textbf{\small then}$$
if self.hour_glass is None:
# Center transformer layer
# $$x' \leftarrow ShortenedLayers(x')$$
x_short = self.shortened(x=x_short, mask=self.mask_short(x_short))
# $$\textbf{else}$$
else:
# $$x' \leftarrow \text{\small H\scriptsize OURGLASS}(x, shorten\_factors)$$
x_short = self.hour_glass(x_short)
# Up-sample the shortened sequence and add a skip connection
# $$x \leftarrow x + Upsampling(x, x', k)$$
x = x + self.up_sampling(x, x_short)
# Final transformer layer
# $$x \leftarrow PostVanillaLayers(x)$$
x = self.post(x=x, mask=self.mask(x))
#
return x
class ShiftRight(nn.Module):
"""
### Shift right operation
This shifts the sequence to the right by the given number of steps
"""
def __init__(self, shift: int):
"""
* `shift` is the number of steps to shift by
"""
super().__init__()
# cannot be negative
assert shift >= 0
#
self.shift = shift
def forward(self, x: torch.Tensor):
"""
* `x` is a tensor of shape `[seq_len, ...]`
"""
# If the shift is $0$ return the original
if self.shift == 0:
return x
# Zeros to be appended to the left
prefix = x.new_zeros([self.shift, *x.shape[1:]])
# Concatenate the zeros and truncate the right
return torch.cat([prefix, x[:-self.shift]])
class AvgPoolShortening(nn.Module):
"""
### Average pool shortening
This down-samples by a given factor with average pooling
"""
def __init__(self, k: int):
"""
* `k` is the shortening factor
"""
super().__init__()
# Average pooling layer
self.pool = nn.AvgPool1d(k, ceil_mode=True)
def forward(self, x: torch.Tensor):
"""
* `x` is of shape `[seq_len, batch_size, d_model]`
"""
# Pooling layer accepts shape `[batch_size, d_model, seq_len]` so we
# permute axes.
return self.pool(x.permute(1, 2, 0)).permute(2, 0, 1)
class NaiveUpSampling(nn.Module):
"""
### Naive up-sampling
This up-samples by repeating
"""
def __init__(self, k: int):
"""
* `k` is the shortening factor
"""
super().__init__()
self.k = k
def forward(self, x: torch.Tensor, x_short: torch.Tensor):
"""
* `x` is the tensor with embeddings before down-sampling
* `x_short` is the tensor of higher density (to be up-sampled) representations
"""
# Repeat across the sequence dimension
expanded = torch.repeat_interleave(x_short, self.k, dim=0)
# Truncate the extra embeddings at the end
expanded = expanded[:x.shape[0]]
#
return expanded
class AutoregressiveMask(nn.Module):
"""
### Generate auto-regressive mask
"""
def __init__(self):
super().__init__()
self.mask = None
def forward(self, x: torch.Tensor):
# Create a mask if we haven't created or sizes have changed
if self.mask is None or self.mask.size(0) != len(x):
# [Subsequent mask](../utils.html), will mask out tokens from seeing future tokens
self.mask = subsequent_mask(len(x)).to(x.device)
#
return self.mask
class LinearPoolingShortening(nn.Module):
"""
### 🚧 Linear pooling for down-sampling
This concatenates the consecutive tokens embeddings that need to be merged and do a linear
transformation to map it to the size of a single token embedding.
"""
def __init__(self):
super().__init__()
raise NotImplementedError
class AttentionBasedShortening(nn.Module):
"""
### 🚧 Down-sampling with attention
\begin{align}
x' &= S(x) + Attention \Big(Q=S(x),K = x, V =x \Big) \\
x' &= x' + FFN(x')
\end{align}
where $S(x)$ is average pooling or linear pooling.
"""
def __init__(self):
super().__init__()
raise NotImplementedError
class LinearUpSampling(nn.Module):
"""
### 🚧 Linear projection for up-sampling
Make a linear projection of dense token embeddings to a size of $d_{\text{model}} k$.
"""
def __init__(self):
super().__init__()
raise NotImplementedError
class AttentionBasedUpSampling(nn.Module):
"""
### 🚧 Attention based up-sampling
\begin{align}
x &= U(x,x') + Attention \Big(Q=U(x,x'),K = x', V = x' \Big) \\
x &= x + FFN(x)
\end{align}
where $U(x,x') = x + LinearUpsampling(x')$
"""
def __init__(self):
super().__init__()
raise NotImplementedError
labml_nn/transformers/hour_glass/experiment.py
+157
View File
@@ -0,0 +1,157 @@
"""
---
title: Hierarchical Transformers Are More Efficient Language Models Experiment
summary: This experiment trains a hourglass model on Tiny Shakespeare dataset.
---
# [Hierarchical Transformers Are More Efficient Language Models](index.html) Experiment
This is an annotated PyTorch experiment to train a [hourglass](index.html).
This is based on
[training loop and configurations for a simple transformer auto-regressive NLP task](../basic/autoregressive_experiment.html).
"""
import math
from typing import List
import torch
from torch import nn
from labml import experiment
from labml.configs import option
from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs
from labml_nn.transformers.hour_glass import HourGlass
from labml_nn.transformers.positional_encoding import PositionalEncoding
class AutoregressiveTransformer(nn.Module):
"""
## Autoregressive language model
"""
def __init__(self, n_tokens: int, d_model: int, dropout: float, hour_glass: HourGlass):
"""
* `n_tokens` is the vocabulary size
* `d_model` is the size of the token embeddings
* `dropout` is the dropout probability
* `hour_glass` is the [hourglass model](index.html)
"""
super().__init__()
# Token embeddings
self.embedding = nn.Embedding(n_tokens, d_model)
# [Fixed positional embeddings](../positional_encoding.html).
#
# 📝 The
# [official paper implementation](https://github.com/google/trax/blob/master/trax/models/research/hourglass.py)
# use [relative attention](../xl/relative_mha.html)
self.pos_embedding = PositionalEncoding(d_model, dropout)
# [hourglass model](index.html)
self.hour_glass = hour_glass
# To normalize the final embeddings
self.norm = nn.LayerNorm([d_model])
# Embedding size
self.d_model = d_model
# Final linear layer to predict the logits
self.output = nn.Linear(d_model, n_tokens)
def __call__(self, x: torch.Tensor):
"""
* `x` is the tensor with token indexes of shape `[seq_len, batch_size]`
"""
# Get embeddings
x = self.embedding(x)
# Add [positional embeddings](../positional_encoding.html)
if self.pos_embedding is not None:
x = self.pos_embedding(x * math.sqrt(self.d_model))
# Hourglass
x = self.hour_glass(x)
# Get logits
output = self.output(self.norm(x))
# Return the logits
return output, None
class Configs(NLPAutoRegressionConfigs):
"""
## Configurations
This inherits from
[training loop and configurations for a simple transformer auto-regressive NLP task](../basic/autoregressive_transformer.html).
"""
# Model
model: AutoregressiveTransformer
# Number of attention heads
n_heads: int = 8
# Dropout probability
dropout: float = 0.1
# Size of feed-forward hidden layer
d_ff: int = 512
# Token embedding size
d_model: int = 256
# Shortening factors
shortening_factors: List[int] = [8, 4]
@option(Configs.model)
def _model(c: Configs):
"""
Create the model
"""
# Create hourglass model
hour_glass = HourGlass(c.n_heads, c.d_model, c.dropout, c.d_ff, c.shortening_factors)
# Create the auto-regressive wrapper
m = AutoregressiveTransformer(c.n_tokens, c.d_model, c.dropout, hour_glass).to(c.device)
#
return m
def main():
# Create experiment
experiment.create(name="hour_glass")
# Create configs
conf = Configs()
# Override configurations
experiment.configs(conf, {
# Use character level tokenizer
'tokenizer': 'character',
# Prompt separator is blank
'prompt_separator': '',
# Starting prompt for sampling
'prompt': 'It is ',
# Use Tiny Shakespeare dataset
'text': 'tiny_shakespeare',
# Use a context size of $256$
'seq_len': 256,
# Train for $128$ epochs
'epochs': 128,
# Batch size $32$
'batch_size': 32,
# Switch between training and validation for $10$ times
# per epoch
'inner_iterations': 10,
# Use [Noam optimizer](../../optimizers/noam.html)
'optimizer.optimizer': 'Noam',
'optimizer.learning_rate': 1.,
#
})
# Set models for saving and loading
experiment.add_pytorch_models({'model': conf.model})
# Start the experiment
with experiment.start():
# Run training
conf.run()
#
if __name__ == '__main__':
main()
labml_nn/transformers/jax_transformer/__init__.py
+1010
View File
File diff suppressed because it is too large Load Diff
labml_nn/transformers/knn/__init__.py
+43
View File
@@ -0,0 +1,43 @@
"""
---
title: k-Nearest Neighbor Language Models
summary: >
This is a simple PyTorch implementation/tutorial of the paper
Generalization through Memorization: Nearest Neighbor Language Models using FAISS.
It runs a kNN model on the final transformer layer embeddings to improve the
loss of transformer based language models.
It's also great for domain adaptation without pre-training.
---
# k-Nearest Neighbor Language Models
This is a [PyTorch](https://pytorch.org) implementation of the paper
[Generalization through Memorization: Nearest Neighbor Language Models](https://arxiv.org/abs/1911.00172).
It uses k-nearest neighbors to improve perplexity of autoregressive transformer models.
An autoregressive language model estimates $p(w_t | \textcolor{yellowgreen}{c_t})$,
where $w_t$ is the token at step $t$
and $c_t$ is the context, $\textcolor{yellowgreen}{c_t} = (w_1, w_2, ..., w_{t-1})$.
This paper, improves $p(w_t | \textcolor{yellowgreen}{c_t})$ using a k-nearest neighbor search
on key-value pairs $\big(f(c_i), w_i\big)$, with search key $f(\textcolor{yellowgreen}{c_t})$.
Here $f(\textcolor{yellowgreen}{c_t})$ is an embedding of the context $\textcolor{yellowgreen}{c_t}$.
The paper (and this implementation) uses the **input to the feed-forward layer of the
final layer of the transformer** as $f(\textcolor{yellowgreen}{c_t})$.
We use [FAISS](https://github.com/facebookresearch/faiss) to index $f(c_i)$.
### Implementation
So to run $k$NN-LM we need to:
* [Train a transformer model](train_model.html)
* [Build an index](build_index.html) of $\big(f(c_i), w_i\big)$
* [Evaluate kNN-ML](eval_knn.html) using $k$NN seach on $\big(f(c_i), w_i\big)$
with $f(\textcolor{yellowgreen}{c_t})$
This experiment uses a small dataset so that we can run this without using up a few hundred giga-bytes
of disk space for the index.
The official implementation of $k$NN-LM can be found [here](https://github.com/urvashik/knnlm).
"""
labml_nn/transformers/knn/build_index.py
+156
View File
@@ -0,0 +1,156 @@
"""
---
title: Build FAISS index for k-NN search
summary: This builds the FAISS index with the transformer embeddings.
---
# Build FAISS index for k-NN search
We want to build the index of $\big(f(c_i), w_i\big)$.
We store $f(c_i)$ and $w_i$ in memory mapped numpy arrays.
We find $f(c_i)$ nearest to $f(c_t)$ using [FAISS](https://github.com/facebookresearch/faiss).
FAISS indexes $\big(f(c_i), i\big)$ and we query it with $f(c_t)$.
"""
from typing import Optional
import faiss
import numpy as np
import torch
from labml import experiment, monit, lab
from labml.utils.pytorch import get_modules
from labml_nn.transformers.knn.train_model import Configs
def load_experiment(run_uuid: str, checkpoint: Optional[int] = None):
"""
Load a saved experiment from [train model](train_model.html).
"""
# Create configurations object
conf = Configs()
# Load custom configurations used in the experiment
conf_dict = experiment.load_configs(run_uuid)
# We need to get inputs to the feed forward layer, $f(c_i)$
conf_dict['is_save_ff_input'] = True
# This experiment is just an evaluation; i.e. nothing is tracked or saved
experiment.evaluate()
# Initialize configurations
experiment.configs(conf, conf_dict)
# Set models for saving/loading
experiment.add_pytorch_models(get_modules(conf))
# Specify the experiment to load from
experiment.load(run_uuid, checkpoint)
# Start the experiment; this is when it actually loads models
experiment.start()
return conf
def gather_keys(conf: Configs):
"""
## Gather $\big(f(c_i), w_i\big)$ and save them in numpy arrays
*Note that these numpy arrays will take up a lot of space (even few hundred gigabytes)
depending on the size of your dataset*.
"""
# Dimensions of $f(c_i)$
d_model = conf.transformer.d_model
# Training data loader
data_loader = conf.trainer.data_loader
# Number of contexts; i.e. number of tokens in the training data minus one.
# $\big(f(c_i), w_i\big)$ for $i \in [2, T]$
n_keys = data_loader.data.shape[0] * data_loader.data.shape[1] - 1
# Numpy array for $f(c_i)$
keys_store = np.memmap(str(lab.get_data_path() / 'keys.npy'), dtype=np.float32, mode='w+', shape=(n_keys, d_model))
# Numpy array for $w_i$
vals_store = np.memmap(str(lab.get_data_path() / 'vals.npy'), dtype=np.int, mode='w+', shape=(n_keys, 1))
# Number of keys $f(c_i)$ collected
added = 0
with torch.no_grad():
# Loop through data
for i, batch in monit.enum("Collect data", data_loader, is_children_silent=True):
# $w_i$ the target labels
vals = batch[1].view(-1, 1)
# Input data moved to the device of the model
data = batch[0].to(conf.device)
# Run the model
_ = conf.model(data)
# Get $f(c_i)$
keys = conf.model.ff_input.view(-1, d_model)
# Save keys, $f(c_i)$ in the memory mapped numpy array
keys_store[added: added + keys.shape[0]] = keys.cpu()
# Save values, $w_i$ in the memory mapped numpy array
vals_store[added: added + keys.shape[0]] = vals
# Increment the number of collected keys
added += keys.shape[0]
def build_index(conf: Configs, n_centeroids: int = 2048, code_size: int = 64, n_probe: int = 8, n_train: int = 200_000):
"""
## Build FAISS index
[Getting started](https://github.com/facebookresearch/faiss/wiki/Getting-started),
[faster search](https://github.com/facebookresearch/faiss/wiki/Faster-search),
and [lower memory footprint](https://github.com/facebookresearch/faiss/wiki/Lower-memory-footprint)
tutorials on FAISS will help you learn more about FAISS usage.
"""
# Dimensions of $f(c_i)$
d_model = conf.transformer.d_model
# Training data loader
data_loader = conf.trainer.data_loader
# Number of contexts; i.e. number of tokens in the training data minus one.
# $\big(f(c_i), w_i\big)$ for $i \in [2, T]$
n_keys = data_loader.data.shape[0] * data_loader.data.shape[1] - 1
# Build an index with Verenoi cell based faster search with compression that
# doesn't store full vectors.
quantizer = faiss.IndexFlatL2(d_model)
index = faiss.IndexIVFPQ(quantizer, d_model, n_centeroids, code_size, 8)
index.nprobe = n_probe
# Load the memory mapped numpy array of keys
keys_store = np.memmap(str(lab.get_data_path() / 'keys.npy'), dtype=np.float32, mode='r', shape=(n_keys, d_model))
# Pick a random sample of keys to train the index with
random_sample = np.random.choice(np.arange(n_keys), size=[min(n_train, n_keys)], replace=False)
with monit.section('Train index'):
# Train the index to store the keys
index.train(keys_store[random_sample])
# Add keys to the index; $\big(f(c_i), i\big)$
for s in monit.iterate('Index', range(0, n_keys, 1024)):
e = min(s + 1024, n_keys)
# $f(c_i)$
keys = keys_store[s:e]
# $i$
idx = np.arange(s, e)
# Add to index
index.add_with_ids(keys, idx)
with monit.section('Save'):
# Save the index
faiss.write_index(index, str(lab.get_data_path() / 'faiss.index'))
def main():
# Load the experiment. Replace the run uuid with you run uuid from
# [training the model](train_model.html).
conf = load_experiment('4984b85c20bf11eb877a69c1a03717cd')
# Set model to evaluation mode
conf.model.eval()
# Collect $\big(f(c_i), w_i\big)$
gather_keys(conf)
# Add them to the index for fast search
build_index(conf)
if __name__ == '__main__':
main()
labml_nn/transformers/knn/eval_knn.py
+157
View File
@@ -0,0 +1,157 @@
"""
---
title: Evaluate k-nearest neighbor language model
summary: >
This runs the kNN model and merges the kNN results with transformer output to
achieve better results than just using the transformer.
---
# Evaluate k-nearest neighbor language model
"""
from typing import Optional, List
import faiss
import numpy as np
import torch
from labml import monit, lab
from labml.logger import inspect
from labml_nn.transformers.knn.train_model import Configs
def knn(queries: torch.Tensor, index: faiss.IndexFlatL2, keys_store: np.ndarray, vals_store: np.ndarray, n_tokens: int):
"""
## $k$-NN to get $p(w_t, c_t)$
Here we refer to $f(\textcolor{yellowgreen}{c_t})$ as queries,
$f(c_i)$ as keys and $w_i$ as values.
"""
# Save shape of queries to reshape results
queries_shape = queries.shape
# Flatten the `batch` and `sequence` dimensions of queries
queries = queries.view(-1, queries_shape[-1])
# Find 10 nearest neighbors of $f(\textcolor{yellowgreen}{c_t})$ among $f(c_i)$.
# `distance` is the distance given by FAISS and `idx`, $i$ is the index of it in `keys_store`.
distance, idx = index.search(queries.numpy(), 10)
# Get $f(c_i)$
keys_found = queries.new_tensor(keys_store[idx])
# Get $w_i$
vals_found = torch.tensor(vals_store[idx]).squeeze(-1)
# We are going to calculate the cosine similarity between normalized vectors
# Normalize $f(c_i)$
keys_found_n = keys_found / torch.sqrt((keys_found ** 2).sum(-1, keepdims=True) + 1e-10)
# Normalize $f(\textcolor{yellowgreen}{c_t})$
queries_n = queries / torch.sqrt((queries ** 2).sum(-1, keepdims=True) + 1e-10)
# Get the dot-product, or cosine similarity
dot_prod = (keys_found_n * queries_n.unsqueeze(1)).sum(-1)
# Token-wise logits
logits_token = dot_prod.new_zeros(queries.shape[0], n_tokens)
# Scatter and accumulate token logits based on the nearest neighbors
_ = logits_token.scatter_(dim=1, index=vals_found, src=dot_prod, reduce='add')
# Reshape the logits
logits_token = logits_token.reshape(queries_shape[0], queries_shape[1], -1)
return logits_token
def validation_loss(knn_weights: List[float], last_n: Optional[int], conf: Configs, index: faiss.IndexFlatL2,
keys_store: np.ndarray, vals_store: np.ndarray):
"""
## Calculate validation loss
We calculate the validation loss of the combined on $k$-NN prediction and transformer prediction.
The weight given to the $k$-NN model is given by `knn_weight`.
It's a list of weights and we calculate the validation loss for each.
"""
# List of losses for each `knn_weights`
losses = [[] for _ in knn_weights]
# Number of samples in each batch
n_samples = []
with torch.no_grad():
# Iterate through validation data
for i, batch in monit.enum("Validation", conf.validator.data_loader, is_children_silent=True):
# Get data and target labels
data, target = batch[0].to(conf.device), batch[1].to(conf.device)
# Run the model and get predictions $p(w_t, c_t)$
res = conf.model(data)
# Get $k$-NN predictions
res_knn = knn(conf.model.ff_input.cpu(), index, keys_store, vals_store, conf.n_tokens)
res_knn = res_knn.to(conf.device)
# This is to calculate only the loss for `last_n` tokens.
# This is important because the first predictions (along the sequence)
# of transformer model has very few past tokens to look at.
if last_n:
res = res[-last_n:]
res_knn = res_knn[-last_n:]
target = target[-last_n:]
# Number of samples
n_s = res.shape[0] * data.shape[1]
n_samples.append(n_s)
# Calculate scores for each of `knn_weights`.
for i, c in enumerate(knn_weights):
# Calculate the loss
loss = conf.loss_func(res_knn * c + (1 - c) * res, target)
losses[i].append(loss * n_s)
return losses, n_samples
def load_index(conf: Configs, n_probe: int = 8):
"""
## Load the index
"""
# Dimensions of $f(c_i)$
d_model = conf.transformer.d_model
# Training data loader
data_loader = conf.trainer.data_loader
# Number of contexts; i.e. number of tokens in the training data minus one.
# $\big(f(c_i), w_i\big)$ for $i \in [2, T]$
n_keys = data_loader.data.shape[0] * data_loader.data.shape[1] - 1
# Load FAISS index
with monit.section('Load index'):
index = faiss.read_index(str(lab.get_data_path() / 'faiss.index'))
# Set number of cells to probe
index.nprobe = n_probe
# Load memory mapped numpy arrays
keys_store = np.memmap(str(lab.get_data_path() / 'keys.npy'), dtype=np.float32, mode='r', shape=(n_keys, d_model))
vals_store = np.memmap(str(lab.get_data_path() / 'vals.npy'), dtype=np.int, mode='r', shape=(n_keys, 1))
return index, keys_store, vals_store
def main():
from labml_nn.transformers.knn.build_index import load_experiment
# Load the experiment. Replace the run uuid with you run uuid from
# [training the model](train_model.html).
conf = load_experiment('4984b85c20bf11eb877a69c1a03717cd')
# Set model to evaluation mode
conf.model.eval()
# Load index
index, keys_store, vals_store = load_index(conf)
# List of weights given to $k$-NN prediction. We will evaluate the validation loss for
# each of the weights
knn_weights = [i / 20 for i in range(10)]
# Evaluate validation loss
losses, n_samples = validation_loss(knn_weights, None, conf, index, keys_store, vals_store)
# Output the losses for each of `knn_weights`.
inspect({c: np.sum(losses[i]) / np.sum(n_samples) for i, c in enumerate(knn_weights)})
if __name__ == '__main__':
main()
labml_nn/transformers/knn/train_model.py
+144
View File
@@ -0,0 +1,144 @@
"""
---
title: Train Autoregressive Transformer
summary: This is training code with notes for a basic auto-regressive transformer.
---
# Train Autoregressive Transformer
This trains a simple [transformer](../../) model for auto-regression.
"""
import torch
from torch import nn
from labml import experiment
from labml.configs import option
from labml.utils.pytorch import get_modules
from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs
from labml_nn.transformers import Encoder, Generator, TransformerConfigs
from labml_nn.transformers.utils import subsequent_mask
class AutoregressiveModel(nn.Module):
"""
## Auto regressive model
"""
def __init__(self, src_embed: nn.Module, encoder: Encoder, generator: Generator, *,
is_save_ff_input: bool = False):
super().__init__()
# Token embedding module
self.src_embed = src_embed
# Transformer based encoder
self.encoder = encoder
# Whether the last layer of the encoder should
# save the input to the feed-forward layer.
# This is out $f(c_t)$, the embedding of the context.
self.encoder.layers[-1].is_save_ff_input = is_save_ff_input
# Next token generation layer;
# this give logits of the the next token
self.generator = generator
# This will be initialized on the first call
self.src_mask = None
@property
def ff_input(self) -> torch.Tensor:
"""
Retrieve saved $f(c_t)$
"""
return self.encoder.layers[-1].ff_input
def forward(self, src: torch.Tensor):
# Create subsequent mask, so that the transformer can only pay attention to past tokens.
if self.src_mask is None or self.src_mask.size(0) != len(src):
self.src_mask = subsequent_mask(len(src)).to(src.device)
# Embed the tokens (`src`) and run it through the the transformer
res = self.encoder(self.src_embed(src), self.src_mask)
# Generate logits of the next token
return self.generator(res), None
class Configs(NLPAutoRegressionConfigs):
"""
## Configurations
The default configs can and will be over-ridden when we start the experiment
"""
transformer: TransformerConfigs
model: AutoregressiveModel
is_save_ff_input = False
@option(Configs.model)
def autoregressive_model(c: Configs):
"""
Initialize the auto-regressive model
"""
m = AutoregressiveModel(
# Get the source token embedding layer, encoder and
# final token generator from configurable transformer
src_embed=c.transformer.src_embed,
encoder=c.transformer.encoder,
generator=c.transformer.generator,
# Whether to save $f(c_t)$
is_save_ff_input=c.is_save_ff_input)
return m.to(c.device)
@option(Configs.transformer)
def transformer_c(c: Configs):
"""
Initialize the configurable transformer encoder for our autoregressive model
"""
tc = TransformerConfigs()
tc.n_src_vocab = c.n_tokens
tc.n_tgt_vocab = c.n_tokens
return tc
def main():
# Create experiment
experiment.create(name="knn_lm")
# Create configs
conf = Configs()
# Load configurations
experiment.configs(conf,
# A dictionary of configurations to override
{'tokenizer': 'character',
'prompt_separator': '',
'prompt': 'It is ',
'text': 'tiny_shakespeare',
'optimizer.optimizer': 'Noam',
'optimizer.learning_rate': 1.,
'optimizer.d_model': 256,
'seq_len': 1024,
'epochs': 128,
'batch_size': 6,
'inner_iterations': 10,
# Transformer configurations
'transformer.d_model': 256,
'transformer.ffn.d_ff': 1024,
'transformer.n_heads': 8,
'transformer.n_layers': 6})
# This is needed to initialize models
conf.n_tokens = conf.text.n_tokens
# Set models for saving and loading
experiment.add_pytorch_models(get_modules(conf))
# Start the experiment
with experiment.start():
# `TrainValidConfigs.run`
conf.run()
if __name__ == '__main__':
main()
labml_nn/transformers/label_smoothing_loss.py
+68
View File
@@ -0,0 +1,68 @@
"""
---
title: Label Smoothing Loss
summary: >
This is an implementation of label smoothing loss, that can be used as
an alternative to cross entropy loss for improved accuracy.
---
# Label Smoothing Loss
"""
import matplotlib.pyplot as plt
import numpy as np
import torch
from torch import nn
class LabelSmoothingLoss(nn.Module):
def __init__(self, size: int, padding_idx: int, smoothing: float = 0.0):
super().__init__()
self.loss = nn.KLDivLoss(reduction='sum')
self.padding_idx = padding_idx
self.confidence = 1.0 - smoothing
self.smoothing = smoothing
self.size = size
self.true_dist = None
def forward(self, x: torch.Tensor, target: torch.Tensor):
assert x.shape[1] == self.size
true_dist = x.clone()
true_dist.fill_(self.smoothing / (self.size - 2))
true_dist.scatter_(1, target.unsqueeze(1), self.confidence)
true_dist[:, self.padding_idx] = 0
mask = torch.nonzero(target == self.padding_idx, as_tuple=False)
if mask.dim() > 0:
true_dist.index_fill_(0, mask.squeeze(), 0.0)
self.true_dist = true_dist
return self.loss(x, true_dist.detach())
def _test_label_smoothing():
smooth_loss = LabelSmoothingLoss(5, 0, 0.4)
predict = torch.tensor([[0, 0.2, 0.7, 0.1, 0],
[0, 0.2, 0.7, 0.1, 0],
[0, 0.2, 0.7, 0.1, 0]], dtype=torch.float)
_ = smooth_loss(predict.log(),
torch.tensor([2, 1, 0], dtype=torch.long))
# Show the target distributions expected by the system.
plt.imshow(smooth_loss.true_dist)
plt.show()
smooth_loss = LabelSmoothingLoss(5, 0, 0.1)
def loss_sample(x):
d = x + 3 * 1
predict2 = torch.tensor([[0, x / d, 1 / d, 1 / d, 1 / d],
], dtype=torch.float)
# print(predict)
return smooth_loss(predict2.log(),
torch.tensor([1], dtype=torch.long)).item()
plt.plot(np.arange(1, 100), [loss_sample(x) for x in range(1, 100)])
plt.show()
if __name__ == '__main__':
_test_label_smoothing()
labml_nn/transformers/mha.py
+206
View File
@@ -0,0 +1,206 @@
"""
---
title: Multi-Headed Attention (MHA)
summary: >
This implements the Multi-Headed Attention used in transformers
using PyTorch with explanations.
---
# Multi-Headed Attention (MHA)
[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/transformers/basic/autoregressive_experiment.ipynb)
This is a tutorial/implementation of multi-headed attention
from paper [Attention Is All You Need](https://arxiv.org/abs/1706.03762)
in [PyTorch](https://pytorch.org/).
The implementation is inspired from [Annotated Transformer](https://nlp.seas.harvard.edu/2018/04/03/attention.html).
Here is the [training code](basic/autoregressive_experiment.html) that uses a basic transformer
with MHA for NLP auto-regression.
[Here is an experiment implementation](basic/autoregressive_experiment.html) that trains a simple transformer.
"""
import math
from typing import Optional, List
import torch
from torch import nn
from labml import tracker
class PrepareForMultiHeadAttention(nn.Module):
"""
<a id="PrepareMHA"></a>
## Prepare for multi-head attention
This module does a linear transformation and splits the vector into given
number of heads for multi-head attention.
This is used to transform **key**, **query**, and **value** vectors.
"""
def __init__(self, d_model: int, heads: int, d_k: int, bias: bool):
super().__init__()
# Linear layer for linear transform
self.linear = nn.Linear(d_model, heads * d_k, bias=bias)
# Number of heads
self.heads = heads
# Number of dimensions in vectors in each head
self.d_k = d_k
def forward(self, x: torch.Tensor):
# Input has shape `[seq_len, batch_size, d_model]` or `[batch_size, d_model]`.
# We apply the linear transformation to the last dimension and split that into
# the heads.
head_shape = x.shape[:-1]
# Linear transform
x = self.linear(x)
# Split last dimension into heads
x = x.view(*head_shape, self.heads, self.d_k)
# Output has shape `[seq_len, batch_size, heads, d_k]` or `[batch_size, heads, d_model]`
return x
class MultiHeadAttention(nn.Module):
r"""
<a id="MHA"></a>
## Multi-Head Attention Module
This computes scaled multi-headed attention for given `query`, `key` and `value` vectors.
$$\mathop{Attention}(Q, K, V) = \underset{seq}{\mathop{softmax}}\Bigg(\frac{Q K^\top}{\sqrt{d_k}}\Bigg)V$$
In simple terms, it finds keys that matches the query, and gets the values of
those keys.
It uses dot-product of query and key as the indicator of how matching they are.
Before taking the $softmax$ the dot-products are scaled by $\frac{1}{\sqrt{d_k}}$.
This is done to avoid large dot-product values causing softmax to
give very small gradients when $d_k$ is large.
Softmax is calculated along the axis of of the sequence (or time).
"""
def __init__(self, heads: int, d_model: int, dropout_prob: float = 0.1, bias: bool = True):
"""
* `heads` is the number of heads.
* `d_model` is the number of features in the `query`, `key` and `value` vectors.
"""
super().__init__()
# Number of features per head
self.d_k = d_model // heads
# Number of heads
self.heads = heads
# These transform the `query`, `key` and `value` vectors for multi-headed attention.
self.query = PrepareForMultiHeadAttention(d_model, heads, self.d_k, bias=bias)
self.key = PrepareForMultiHeadAttention(d_model, heads, self.d_k, bias=bias)
self.value = PrepareForMultiHeadAttention(d_model, heads, self.d_k, bias=True)
# Softmax for attention along the time dimension of `key`
self.softmax = nn.Softmax(dim=1)
# Output layer
self.output = nn.Linear(d_model, d_model)
# Dropout
self.dropout = nn.Dropout(dropout_prob)
# Scaling factor before the softmax
self.scale = 1 / math.sqrt(self.d_k)
# We store attentions so that it can be used for logging, or other computations if needed
self.attn = None
def get_scores(self, query: torch.Tensor, key: torch.Tensor):
"""
### Calculate scores between queries and keys
This method can be overridden for other variations like relative attention.
"""
# Calculate $Q K^\top$ or $S_{ijbh} = \sum_d Q_{ibhd} K_{jbhd}$
return torch.einsum('ibhd,jbhd->ijbh', query, key)
def prepare_mask(self, mask: torch.Tensor, query_shape: List[int], key_shape: List[int]):
"""
`mask` has shape `[seq_len_q, seq_len_k, batch_size]`, where first dimension is the query dimension.
If the query dimension is equal to $1$ it will be broadcasted.
"""
assert mask.shape[0] == 1 or mask.shape[0] == query_shape[0]
assert mask.shape[1] == key_shape[0]
assert mask.shape[2] == 1 or mask.shape[2] == query_shape[1]
# Same mask applied to all heads.
mask = mask.unsqueeze(-1)
# resulting mask has shape `[seq_len_q, seq_len_k, batch_size, heads]`
return mask
def forward(self, *,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
mask: Optional[torch.Tensor] = None):
"""
`query`, `key` and `value` are the tensors that store
collection of *query*, *key* and *value* vectors.
They have shape `[seq_len, batch_size, d_model]`.
`mask` has shape `[seq_len, seq_len, batch_size]` and
`mask[i, j, b]` indicates whether for batch `b`,
query at position `i` has access to key-value at position `j`.
"""
# `query`, `key` and `value` have shape `[seq_len, batch_size, d_model]`
seq_len, batch_size, _ = query.shape
if mask is not None:
mask = self.prepare_mask(mask, query.shape, key.shape)
# Prepare `query`, `key` and `value` for attention computation.
# These will then have shape `[seq_len, batch_size, heads, d_k]`.
query = self.query(query)
key = self.key(key)
value = self.value(value)
# Compute attention scores $Q K^\top$.
# This gives a tensor of shape `[seq_len, seq_len, batch_size, heads]`.
scores = self.get_scores(query, key)
# Scale scores $\frac{Q K^\top}{\sqrt{d_k}}$
scores *= self.scale
# Apply mask
if mask is not None:
scores = scores.masked_fill(mask == 0, float('-inf'))
# $softmax$ attention along the key sequence dimension
# $\underset{seq}{softmax}\Bigg(\frac{Q K^\top}{\sqrt{d_k}}\Bigg)$
attn = self.softmax(scores)
# Save attentions if debugging
tracker.debug('attn', attn)
# Apply dropout
attn = self.dropout(attn)
# Multiply by values
# $$\underset{seq}{softmax}\Bigg(\frac{Q K^\top}{\sqrt{d_k}}\Bigg)V$$
x = torch.einsum("ijbh,jbhd->ibhd", attn, value)
# Save attentions for any other calculations
self.attn = attn.detach()
# Concatenate multiple heads
x = x.reshape(seq_len, batch_size, -1)
# Output layer
return self.output(x)
labml_nn/transformers/mlm/__init__.py
+139
View File
@@ -0,0 +1,139 @@
"""
---
title: Masked Language Model
summary: >
This is an annotated implementation/tutorial of the Masked Language Model in PyTorch.
---
# Masked Language Model (MLM)
This is a [PyTorch](https://pytorch.org) implementation of the Masked Language Model (MLM)
used to pre-train the BERT model introduced in the paper
[BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://arxiv.org/abs/1810.04805).
## BERT Pretraining
BERT model is a transformer model.
The paper pre-trains the model using MLM and with next sentence prediction.
We have only implemented MLM here.
### Next sentence prediction
In *next sentence prediction*, the model is given two sentences `A` and `B` and the model
makes a binary prediction whether `B` is the sentence that follows `A` in the actual text.
The model is fed with actual sentence pairs 50% of the time and random pairs 50% of the time.
This classification is done while applying MLM. *We haven't implemented this here.*
## Masked LM
This masks a percentage of tokens at random and trains the model to predict
the masked tokens.
They **mask 15% of the tokens** by replacing them with a special `[MASK]` token.
The loss is computed on predicting the masked tokens only.
This causes a problem during fine-tuning and actual usage since there are no `[MASK]` tokens
at that time.
Therefore we might not get any meaningful representations.
To overcome this **10% of the masked tokens are replaced with the original token**,
and another **10% of the masked tokens are replaced with a random token**.
This trains the model to give representations about the actual token whether or not the
input token at that position is a `[MASK]`.
And replacing with a random token causes it to
give a representation that has information from the context as well;
because it has to use the context to fix randomly replaced tokens.
## Training
MLMs are harder to train than autoregressive models because they have a smaller training signal.
i.e. only a small percentage of predictions are trained per sample.
Another problem is since the model is bidirectional, any token can see any other token.
This makes the "credit assignment" harder.
Let's say you have the character level model trying to predict `home *s where i want to be`.
At least during the early stages of the training, it'll be super hard to figure out why the
replacement for `*` should be `i`, it could be anything from the whole sentence.
Whilst, in an autoregressive setting the model will only have to use `h` to predict `o` and
`hom` to predict `e` and so on. So the model will initially start predicting with a shorter context first
and then learn to use longer contexts later.
Since MLMs have this problem it's a lot faster to train if you start with a smaller sequence length
initially and then use a longer sequence length later.
Here is [the training code](experiment.html) for a simple MLM model.
"""
from typing import List
import torch
class MLM:
"""
## Masked LM (MLM)
This class implements the masking procedure for a given batch of token sequences.
"""
def __init__(self, *,
padding_token: int, mask_token: int, no_mask_tokens: List[int], n_tokens: int,
masking_prob: float = 0.15, randomize_prob: float = 0.1, no_change_prob: float = 0.1,
):
"""
* `padding_token` is the padding token `[PAD]`.
We will use this to mark the labels that shouldn't be used for loss calculation.
* `mask_token` is the masking token `[MASK]`.
* `no_mask_tokens` is a list of tokens that should not be masked.
This is useful if we are training the MLM with another task like classification at the same time,
and we have tokens such as `[CLS]` that shouldn't be masked.
* `n_tokens` total number of tokens (used for generating random tokens)
* `masking_prob` is the masking probability
* `randomize_prob` is the probability of replacing with a random token
* `no_change_prob` is the probability of replacing with original token
"""
self.n_tokens = n_tokens
self.no_change_prob = no_change_prob
self.randomize_prob = randomize_prob
self.masking_prob = masking_prob
self.no_mask_tokens = no_mask_tokens + [padding_token, mask_token]
self.padding_token = padding_token
self.mask_token = mask_token
def __call__(self, x: torch.Tensor):
"""
* `x` is the batch of input token sequences.
It's a tensor of type `long` with shape `[seq_len, batch_size]`.
"""
# Mask `masking_prob` of tokens
full_mask = torch.rand(x.shape, device=x.device) < self.masking_prob
# Unmask `no_mask_tokens`
for t in self.no_mask_tokens:
full_mask &= x != t
# A mask for tokens to be replaced with original tokens
unchanged = full_mask & (torch.rand(x.shape, device=x.device) < self.no_change_prob)
# A mask for tokens to be replaced with a random token
random_token_mask = full_mask & (torch.rand(x.shape, device=x.device) < self.randomize_prob)
# Indexes of tokens to be replaced with random tokens
random_token_idx = torch.nonzero(random_token_mask, as_tuple=True)
# Random tokens for each of the locations
random_tokens = torch.randint(0, self.n_tokens, (len(random_token_idx[0]),), device=x.device)
# The final set of tokens that are going to be replaced by `[MASK]`
mask = full_mask & ~random_token_mask & ~unchanged
# Make a clone of the input for the labels
y = x.clone()
# Replace with `[MASK]` tokens;
# note that this doesn't include the tokens that will have the original token unchanged and
# those that get replace with a random token.
x.masked_fill_(mask, self.mask_token)
# Assign random tokens
x[random_token_idx] = random_tokens
# Assign token `[PAD]` to all the other locations in the labels.
# The labels equal to `[PAD]` will not be used in the loss.
y.masked_fill_(~full_mask, self.padding_token)
# Return the masked input and the labels
return x, y
labml_nn/transformers/mlm/experiment.py
+308
View File
@@ -0,0 +1,308 @@
"""
---
title: Masked Language Model Experiment
summary: This experiment trains Masked Language Model (MLM) on Tiny Shakespeare dataset.
---
# [Masked Language Model (MLM)](index.html) Experiment
This is an annotated PyTorch experiment to train a [Masked Language Model](index.html).
"""
from typing import List
import torch
from torch import nn
from labml import experiment, tracker, logger
from labml.configs import option
from labml.logger import Text
from labml_nn.helpers.metrics import Accuracy
from labml_nn.helpers.trainer import BatchIndex
from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs
from labml_nn.transformers import Encoder, Generator
from labml_nn.transformers import TransformerConfigs
from labml_nn.transformers.mlm import MLM
class TransformerMLM(nn.Module):
"""
# Transformer based model for MLM
"""
def __init__(self, *, encoder: Encoder, src_embed: nn.Module, generator: Generator):
"""
* `encoder` is the transformer [Encoder](../models.html#Encoder)
* `src_embed` is the token
[embedding module (with positional encodings)](../models.html#EmbeddingsWithLearnedPositionalEncoding)
* `generator` is the [final fully connected layer](../models.html#Generator) that gives the logits.
"""
super().__init__()
self.generator = generator
self.src_embed = src_embed
self.encoder = encoder
def forward(self, x: torch.Tensor):
# Get the token embeddings with positional encodings
x = self.src_embed(x)
# Transformer encoder
x = self.encoder(x, None)
# Logits for the output
y = self.generator(x)
# Return results
# (second value is for state, since our trainer is used with RNNs also)
return y, None
class Configs(NLPAutoRegressionConfigs):
"""
## Configurations
This inherits from
[`NLPAutoRegressionConfigs`](../../experiments/nlp_autoregression.html)
because it has the data pipeline implementations that we reuse here.
We have implemented a custom training step form MLM.
"""
# MLM model
model: TransformerMLM
# Transformer
transformer: TransformerConfigs
# Number of tokens
n_tokens: int = 'n_tokens_mlm'
# Tokens that shouldn't be masked
no_mask_tokens: List[int] = []
# Probability of masking a token
masking_prob: float = 0.15
# Probability of replacing the mask with a random token
randomize_prob: float = 0.1
# Probability of replacing the mask with original token
no_change_prob: float = 0.1
# [Masked Language Model (MLM) class](index.html) to generate the mask
mlm: MLM
# `[MASK]` token
mask_token: int
# `[PADDING]` token
padding_token: int
# Prompt to sample
prompt: str = [
"We are accounted poor citizens, the patricians good.",
"What authority surfeits on would relieve us: if they",
"would yield us but the superfluity, while it were",
"wholesome, we might guess they relieved us humanely;",
"but they think we are too dear: the leanness that",
"afflicts us, the object of our misery, is as an",
"inventory to particularise their abundance; our",
"sufferance is a gain to them Let us revenge this with",
"our pikes, ere we become rakes: for the gods know I",
"speak this in hunger for bread, not in thirst for revenge.",
]
def init(self):
"""
### Initialization
"""
# `[MASK]` token
self.mask_token = self.n_tokens - 1
# `[PAD]` token
self.padding_token = self.n_tokens - 2
# [Masked Language Model (MLM) class](index.html) to generate the mask
self.mlm = MLM(padding_token=self.padding_token,
mask_token=self.mask_token,
no_mask_tokens=self.no_mask_tokens,
n_tokens=self.n_tokens,
masking_prob=self.masking_prob,
randomize_prob=self.randomize_prob,
no_change_prob=self.no_change_prob)
# Accuracy metric (ignore the labels equal to `[PAD]`)
self.accuracy = Accuracy(ignore_index=self.padding_token)
# Cross entropy loss (ignore the labels equal to `[PAD]`)
self.loss_func = nn.CrossEntropyLoss(ignore_index=self.padding_token)
#
super().init()
def step(self, batch: any, batch_idx: BatchIndex):
"""
### Training or validation step
"""
# Move the input to the device
data = batch[0].to(self.device)
# Update global step (number of tokens processed) when in training mode
if self.mode.is_train:
tracker.add_global_step(data.shape[0] * data.shape[1])
# Get the masked input and labels
with torch.no_grad():
data, labels = self.mlm(data)
# Get model outputs.
# It's returning a tuple for states when using RNNs.
# This is not implemented yet.
output, *_ = self.model(data)
# Calculate and log the loss
loss = self.loss_func(output.view(-1, output.shape[-1]), labels.view(-1))
tracker.add("loss.", loss)
# Calculate and log accuracy
self.accuracy(output, labels)
self.accuracy.track()
# Train the model
if self.mode.is_train:
# Calculate gradients
loss.backward()
# Clip gradients
torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=self.grad_norm_clip)
# Take optimizer step
self.optimizer.step()
# Log the model parameters and gradients on last batch of every epoch
if batch_idx.is_last:
tracker.add('model', self.model)
# Clear the gradients
self.optimizer.zero_grad()
# Save the tracked metrics
tracker.save()
@torch.no_grad()
def sample(self):
"""
### Sampling function to generate samples periodically while training
"""
# Empty tensor for data filled with `[PAD]`.
data = torch.full((self.seq_len, len(self.prompt)), self.padding_token, dtype=torch.long)
# Add the prompts one by one
for i, p in enumerate(self.prompt):
# Get token indexes
d = self.text.text_to_i(p)
# Add to the tensor
s = min(self.seq_len, len(d))
data[:s, i] = d[:s]
# Move the tensor to current device
data = data.to(self.device)
# Get masked input and labels
data, labels = self.mlm(data)
# Get model outputs
output, *_ = self.model(data)
# Print the samples generated
for j in range(data.shape[1]):
# Collect output from printing
log = []
# For each token
for i in range(len(data)):
# If the label is not `[PAD]`
if labels[i, j] != self.padding_token:
# Get the prediction
t = output[i, j].argmax().item()
# If it's a printable character
if t < len(self.text.itos):
# Correct prediction
if t == labels[i, j]:
log.append((self.text.itos[t], Text.value))
# Incorrect prediction
else:
log.append((self.text.itos[t], Text.danger))
# If it's not a printable character
else:
log.append(('*', Text.danger))
# If the label is `[PAD]` (unmasked) print the original.
elif data[i, j] < len(self.text.itos):
log.append((self.text.itos[data[i, j]], Text.subtle))
# Print
logger.log(log)
@option(Configs.n_tokens)
def n_tokens_mlm(c: Configs):
"""
Number of tokens including `[PAD]` and `[MASK]`
"""
return c.text.n_tokens + 2
@option(Configs.transformer)
def _transformer_configs(c: Configs):
"""
### Transformer configurations
"""
# We use our
# [configurable transformer implementation](../configs.html#TransformerConfigs)
conf = TransformerConfigs()
# Set the vocabulary sizes for embeddings and generating logits
conf.n_src_vocab = c.n_tokens
conf.n_tgt_vocab = c.n_tokens
# Embedding size
conf.d_model = c.d_model
#
return conf
@option(Configs.model)
def _model(c: Configs):
"""
Create classification model
"""
m = TransformerMLM(encoder=c.transformer.encoder,
src_embed=c.transformer.src_embed,
generator=c.transformer.generator).to(c.device)
return m
def main():
# Create experiment
experiment.create(name="mlm")
# Create configs
conf = Configs()
# Override configurations
experiment.configs(conf, {
# Batch size
'batch_size': 64,
# Sequence length of $32$. We use a short sequence length to train faster.
# Otherwise it takes forever to train.
'seq_len': 32,
# Train for 1024 epochs.
'epochs': 1024,
# Switch between training and validation for $1$ times
# per epoch
'inner_iterations': 1,
# Transformer configurations (same as defaults)
'd_model': 128,
'transformer.ffn.d_ff': 256,
'transformer.n_heads': 8,
'transformer.n_layers': 6,
# Use [Noam optimizer](../../optimizers/noam.html)
'optimizer.optimizer': 'Noam',
'optimizer.learning_rate': 1.,
})
# Set models for saving and loading
experiment.add_pytorch_models({'model': conf.model})
# Start the experiment
with experiment.start():
# Run training
conf.run()
#
if __name__ == '__main__':
main()
labml_nn/transformers/mlm/readme.md
+55
View File
@@ -0,0 +1,55 @@
# [Masked Language Model (MLM)](https://nn.labml.ai/transformers/mlm/index.html)
This is a [PyTorch](https://pytorch.org) implementation of Masked Language Model (MLM)
used to pre-train the BERT model introduced in the paper
[BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://arxiv.org/abs/1810.04805).
## BERT Pretraining
BERT model is a transformer model.
The paper pre-trains the model using MLM and with next sentence prediction.
We have only implemented MLM here.
### Next sentence prediction
In *next sentence prediction*, the model is given two sentences `A` and `B` and the model
makes a binary prediction whether `B` is the sentence that follows `A` in the actual text.
The model is fed with actual sentence pairs 50% of the time and random pairs 50% of the time.
This classification is done while applying MLM. *We haven't implemented this here.*
## Masked LM
This masks a percentage of tokens at random and trains the model to predict
the masked tokens.
They **mask 15% of the tokens** by replacing them with a special `[MASK]` token.
The loss is computed on predicting the masked tokens only.
This causes a problem during fine-tuning and actual usage since there are no `[MASK]` tokens
at that time.
Therefore we might not get any meaningful representations.
To overcome this **10% of the masked tokens are replaced with the original token**,
and another **10% of the masked tokens are replaced with a random token**.
This trains the model to give representations about the actual token whether or not the
input token at that position is a `[MASK]`.
And replacing with a random token causes it to
give a representation that has information from the context as well;
because it has to use the context to fix randomly replaced tokens.
## Training
MLMs are harder to train than autoregressive models because they have a smaller training signal.
i.e. only a small percentage of predictions are trained per sample.
Another problem is since the model is bidirectional, any token can see any other token.
This makes the "credit assignment" harder.
Let's say you have the character level model trying to predict `home *s where i want to be`.
At least during the early stages of the training, it'll be super hard to figure out why the
replacement for `*` should be `i`, it could be anything from the whole sentence.
Whilst, in an autoregressive setting the model will only have to use `h` to predict `o` and
`hom` to predict `e` and so on. So the model will initially start predicting with a shorter context first
and then learn to use longer contexts later.
Since MLMs have this problem it's a lot faster to train if you start with a smaller sequence length
initially and then use a longer sequence length later.
Here is [the training code](https://nn.labml.ai/transformers/mlm/experiment.html) for a simple MLM model.
labml_nn/transformers/mlp_mixer/__init__.py
+78
View File
@@ -0,0 +1,78 @@
"""
---
title: "MLP-Mixer: An all-MLP Architecture for Vision"
summary: >
This is an annotated implementation/tutorial of MLP-Mixer: An all-MLP Architecture for Vision in PyTorch.
---
# MLP-Mixer: An all-MLP Architecture for Vision
This is a [PyTorch](https://pytorch.org) implementation of the paper
[MLP-Mixer: An all-MLP Architecture for Vision](https://arxiv.org/abs/2105.01601).
This paper applies the model on vision tasks.
The model is similar to a transformer with attention layer being replaced by a MLP
that is applied across the patches (or tokens in case of a NLP task).
Our implementation of MLP Mixer is a drop in replacement for the [self-attention layer](../mha.html)
in [our transformer implementation](../models.html).
So it's just a couple of lines of code, transposing the tensor to apply the MLP
across the sequence dimension.
Although the paper applied MLP Mixer on vision tasks,
we tried it on a [masked language model](../mlm/index.html).
[Here is the experiment code](experiment.html).
"""
from typing import Optional
import torch
from torch import nn
class MLPMixer(nn.Module):
"""
## MLP Mixer
This module is a drop-in replacement for [self-attention layer](../mha.html).
It transposes the input tensor before feeding it to the MLP and transposes back,
so that the MLP is applied across the sequence dimension (across tokens or image patches) instead
of the feature dimension.
"""
def __init__(self, mlp: nn.Module):
"""
* `ffn` is the MLP module.
"""
super().__init__()
self.mlp = mlp
def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, mask: Optional[torch.Tensor] = None):
"""
The [normal attention module](../mha.html) can be fed with different token embeddings for
$\text{query}$,$\text{key}$, and $\text{value}$ and a mask.
We follow the same function signature so that we can replace it directly.
For MLP mixing, $$x = \text{query} = \text{key} = \text{value}$$ and masking is not possible.
Shape of `query` (and `key` and `value`) is `[seq_len, batch_size, d_model]`.
"""
# $\text{query}$,$\text{key}$, and $\text{value}$ all should be the same
assert query is key and key is value
# MLP mixer doesn't support masking. i.e. all tokens will see all other token embeddings.
assert mask is None
# Assign to `x` for clarity
x = query
# Transpose so that the last dimension is the sequence dimension.
# New shape is `[d_model, batch_size, seq_len]`
x = x.transpose(0, 2)
# Apply the MLP across tokens
x = self.mlp(x)
# Transpose back into original form
x = x.transpose(0, 2)
#
return x
labml_nn/transformers/mlp_mixer/experiment.py
+115
View File
@@ -0,0 +1,115 @@
"""
---
title: MLP Mixer experiment
summary: This experiment trains MLP Mixer on Tiny Shakespeare dataset.
---
# [MLP Mixer](index.html) Experiment
This is an annotated PyTorch experiment to train a [MLP Mixer Model](index.html).
"""
from labml import experiment
from labml.configs import option
from labml_nn.transformers import TransformerConfigs
from labml_nn.transformers.configs import FeedForwardConfigs
from labml_nn.transformers.mlm.experiment import TransformerMLM, Configs as MLMConfigs
class Configs(MLMConfigs):
"""
## Configurations
This inherits from
[`MLMConfigs`](../mlm/experiment.html) where we define an experiment for
[Masked Language Models](../mlm.index.html).
"""
# Configurable [Feed-Forward Network](../feed_forward.html) for the MLP
mix_mlp: FeedForwardConfigs
@option(Configs.mix_mlp)
def _mix_mlp_configs(c: Configs):
"""
The mixing MLP configurations
"""
conf = FeedForwardConfigs()
# Size of the MLP is the sequence length, because it is applied across tokens
conf.d_model = c.seq_len
# The paper suggests $GELU$ activation
conf.activation = 'GELU'
#
return conf
@option(Configs.transformer)
def _transformer_configs(c: Configs):
"""
### Transformer configurations
"""
# We use our
# [configurable transformer implementation](../configs.html#TransformerConfigs)
conf = TransformerConfigs()
# Set the vocabulary sizes for embeddings and generating logits
conf.n_src_vocab = c.n_tokens
conf.n_tgt_vocab = c.n_tokens
# Embedding size
conf.d_model = c.d_model
# Change attention module to [MLPMixer](index.html)
from labml_nn.transformers.mlp_mixer import MLPMixer
conf.encoder_attn = MLPMixer(c.mix_mlp.ffn)
#
return conf
def main():
# Create experiment
experiment.create(name="mlp_mixer_mlm")
# Create configs
conf = Configs()
# Override configurations
experiment.configs(conf, {
# Batch size
'batch_size': 64,
# Sequence length of $32$. We use a short sequence length to train faster.
# Otherwise MLM models take forever to train.
'seq_len': 32,
# Train for 1024 epochs.
'epochs': 1024,
# Switch between training and validation for $1$ times
# per epoch
'inner_iterations': 1,
# Transformer configurations
'd_model': 128,
'transformer.ffn.d_ff': 256,
'transformer.n_heads': 8,
'transformer.n_layers': 6,
'transformer.ffn.activation': 'GELU',
# Mixer MLP hidden layer size
'mix_mlp.d_ff': 128,
# Use [Noam optimizer](../../optimizers/noam.html)
'optimizer.optimizer': 'Noam',
'optimizer.learning_rate': 1.,
})
# Set models for saving and loading
experiment.add_pytorch_models({'model': conf.model})
# Start the experiment
with experiment.start():
# Run training
conf.run()
#
if __name__ == '__main__':
main()
labml_nn/transformers/mlp_mixer/readme.md
+17
View File
@@ -0,0 +1,17 @@
# [MLP-Mixer: An all-MLP Architecture for Vision](https://nn.labml.ai/transformers/mlp_mixer/index.html)
This is a [PyTorch](https://pytorch.org) implementation of the paper
[MLP-Mixer: An all-MLP Architecture for Vision](https://arxiv.org/abs/2105.01601).
This paper applies the model on vision tasks.
The model is similar to a transformer with attention layer being replaced by a MLP
that is applied across the patches (or tokens in case of a NLP task).
Our implementation of MLP Mixer is a drop in replacement for the [self-attention layer](https://nn.labml.ai/transformers/mha.html)
in [our transformer implementation](https://nn.labml.ai/transformers/models.html).
So it's just a couple of lines of code, transposing the tensor to apply the MLP
across the sequence dimension.
Although the paper applied MLP Mixer on vision tasks,
we tried it on a [masked language model](https://nn.labml.ai/transformers/mlm/index.html).
[Here is the experiment code](https://nn.labml.ai/transformers/mlp_mixer/experiment.html).
labml_nn/transformers/models.py
+224
View File
@@ -0,0 +1,224 @@
"""
---
title: Transformer Encoder and Decoder Models
summary: >
These are PyTorch implementations of Transformer based encoder and decoder models,
as well as other related modules.
---
# Transformer Encoder and Decoder Models
[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/transformers/basic/autoregressive_experiment.ipynb)
"""
import math
import torch
import torch.nn as nn
from labml_nn.utils import clone_module_list
from .feed_forward import FeedForward
from .mha import MultiHeadAttention
from .positional_encoding import get_positional_encoding
class EmbeddingsWithPositionalEncoding(nn.Module):
"""
<a id="EmbeddingsWithPositionalEncoding"></a>
## Embed tokens and add [fixed positional encoding](positional_encoding.html)
"""
def __init__(self, d_model: int, n_vocab: int, max_len: int = 5000):
super().__init__()
self.linear = nn.Embedding(n_vocab, d_model)
self.d_model = d_model
self.register_buffer('positional_encodings', get_positional_encoding(d_model, max_len))
def forward(self, x: torch.Tensor):
pe = self.positional_encodings[:x.shape[0]].requires_grad_(False)
return self.linear(x) * math.sqrt(self.d_model) + pe
class EmbeddingsWithLearnedPositionalEncoding(nn.Module):
"""
<a id="EmbeddingsWithLearnedPositionalEncoding"></a>
## Embed tokens and add parameterized positional encodings
"""
def __init__(self, d_model: int, n_vocab: int, max_len: int = 5000):
super().__init__()
self.linear = nn.Embedding(n_vocab, d_model)
self.d_model = d_model
self.positional_encodings = nn.Parameter(torch.zeros(max_len, 1, d_model), requires_grad=True)
def forward(self, x: torch.Tensor):
pe = self.positional_encodings[:x.shape[0]]
return self.linear(x) * math.sqrt(self.d_model) + pe
class TransformerLayer(nn.Module):
"""
<a id="TransformerLayer"></a>
## Transformer Layer
This can act as an encoder layer or a decoder layer. We use pre-norm.
"""
def __init__(self, *,
d_model: int,
self_attn: MultiHeadAttention,
src_attn: MultiHeadAttention = None,
feed_forward: FeedForward,
dropout_prob: float):
"""
* `d_model` is the token embedding size
* `self_attn` is the self attention module
* `src_attn` is the source attention module (when this is used in a decoder)
* `feed_forward` is the feed forward module
* `dropout_prob` is the probability of dropping out after self attention and FFN
"""
super().__init__()
self.size = d_model
self.self_attn = self_attn
self.src_attn = src_attn
self.feed_forward = feed_forward
self.dropout = nn.Dropout(dropout_prob)
self.norm_self_attn = nn.LayerNorm([d_model])
if self.src_attn is not None:
self.norm_src_attn = nn.LayerNorm([d_model])
self.norm_ff = nn.LayerNorm([d_model])
# Whether to save input to the feed forward layer
self.is_save_ff_input = False
def forward(self, *,
x: torch.Tensor,
mask: torch.Tensor,
src: torch.Tensor = None,
src_mask: torch.Tensor = None):
# Normalize the vectors before doing self attention
z = self.norm_self_attn(x)
# Run through self attention, i.e. keys and values are from self
self_attn = self.self_attn(query=z, key=z, value=z, mask=mask)
# Add the self attention results
x = x + self.dropout(self_attn)
# If a source is provided, get results from attention to source.
# This is when you have a decoder layer that pays attention to
# encoder outputs
if src is not None:
# Normalize vectors
z = self.norm_src_attn(x)
# Attention to source. i.e. keys and values are from source
attn_src = self.src_attn(query=z, key=src, value=src, mask=src_mask)
# Add the source attention results
x = x + self.dropout(attn_src)
# Normalize for feed-forward
z = self.norm_ff(x)
# Save the input to the feed forward layer if specified
if self.is_save_ff_input:
self.ff_input = z.clone()
# Pass through the feed-forward network
ff = self.feed_forward(z)
# Add the feed-forward results back
x = x + self.dropout(ff)
return x
class Encoder(nn.Module):
"""
<a id="Encoder"></a>
## Transformer Encoder
"""
def __init__(self, layer: TransformerLayer, n_layers: int):
super().__init__()
# Make copies of the transformer layer
self.layers = clone_module_list(layer, n_layers)
# Final normalization layer
self.norm = nn.LayerNorm([layer.size])
def forward(self, x: torch.Tensor, mask: torch.Tensor):
# Run through each transformer layer
for layer in self.layers:
x = layer(x=x, mask=mask)
# Finally, normalize the vectors
return self.norm(x)
class Decoder(nn.Module):
"""
<a id="Decoder"></a>
## Transformer Decoder
"""
def __init__(self, layer: TransformerLayer, n_layers: int):
super().__init__()
# Make copies of the transformer layer
self.layers = clone_module_list(layer, n_layers)
# Final normalization layer
self.norm = nn.LayerNorm([layer.size])
def forward(self, x: torch.Tensor, memory: torch.Tensor, src_mask: torch.Tensor, tgt_mask: torch.Tensor):
# Run through each transformer layer
for layer in self.layers:
x = layer(x=x, mask=tgt_mask, src=memory, src_mask=src_mask)
# Finally, normalize the vectors
return self.norm(x)
class Generator(nn.Module):
"""
<a id="Generator"></a>
## Generator
This predicts the tokens and gives the lof softmax of those.
You don't need this if you are using `nn.CrossEntropyLoss`.
"""
def __init__(self, n_vocab: int, d_model: int):
super().__init__()
self.projection = nn.Linear(d_model, n_vocab)
def forward(self, x):
return self.projection(x)
class EncoderDecoder(nn.Module):
"""
<a id="EncoderDecoder"></a>
## Combined Encoder-Decoder
"""
def __init__(self, encoder: Encoder, decoder: Decoder, src_embed: nn.Module, tgt_embed: nn.Module, generator: nn.Module):
super().__init__()
self.encoder = encoder
self.decoder = decoder
self.src_embed = src_embed
self.tgt_embed = tgt_embed
self.generator = generator
# This was important from their code.
# Initialize parameters with Glorot / fan_avg.
for p in self.parameters():
if p.dim() > 1:
nn.init.xavier_uniform_(p)
def forward(self, src: torch.Tensor, tgt: torch.Tensor, src_mask: torch.Tensor, tgt_mask: torch.Tensor):
# Run the source through encoder
enc = self.encode(src, src_mask)
# Run encodings and targets through decoder
return self.decode(enc, src_mask, tgt, tgt_mask)
def encode(self, src: torch.Tensor, src_mask: torch.Tensor):
return self.encoder(self.src_embed(src), src_mask)
def decode(self, memory: torch.Tensor, src_mask: torch.Tensor, tgt: torch.Tensor, tgt_mask: torch.Tensor):
return self.decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask)
labml_nn/transformers/positional_encoding.py
+76
View File
@@ -0,0 +1,76 @@
"""
---
title: Fixed Positional Encodings
summary: >
Implementation with explanation of fixed positional encodings as
described in paper Attention is All You Need.
---
# Fixed Positional Encodings
The positional encoding encodes the position along the sequence into
a vector of size `d_model`.
\begin{align}
PE_{p,2i} &= sin\Bigg(\frac{p}{10000^{\frac{2i}{d_{model}}}}\Bigg) \\
PE_{p,2i + 1} &= cos\Bigg(\frac{p}{10000^{\frac{2i}{d_{model}}}}\Bigg)
\end{align}
Where $1 \leq 2i, 2i + 1 \leq d_{model}$
are the feature indexes in the encoding, and $p$ is the position.
"""
import math
import numpy as np
import torch
import torch.nn as nn
class PositionalEncoding(nn.Module):
def __init__(self, d_model: int, dropout_prob: float, max_len: int = 5000):
super().__init__()
self.dropout = nn.Dropout(dropout_prob)
self.register_buffer('positional_encodings', get_positional_encoding(d_model, max_len), False)
def forward(self, x: torch.Tensor):
pe = self.positional_encodings[:x.shape[0]].detach().requires_grad_(False)
x = x + pe
x = self.dropout(x)
return x
def get_positional_encoding(d_model: int, max_len: int = 5000):
# Empty encodings vectors
encodings = torch.zeros(max_len, d_model)
# Position indexes
position = torch.arange(0, max_len, dtype=torch.float32).unsqueeze(1)
# $2 * i$
two_i = torch.arange(0, d_model, 2, dtype=torch.float32)
# $10000^{\frac{2i}{d_{model}}}$
div_term = torch.exp(two_i * -(math.log(10000.0) / d_model))
# $PE_{p,2i} = sin\Bigg(\frac{p}{10000^{\frac{2i}{d_{model}}}}\Bigg)$
encodings[:, 0::2] = torch.sin(position * div_term)
# $PE_{p,2i + 1} = cos\Bigg(\frac{p}{10000^{\frac{2i}{d_{model}}}}\Bigg)$
encodings[:, 1::2] = torch.cos(position * div_term)
# Add batch dimension
encodings = encodings.unsqueeze(1).requires_grad_(False)
return encodings
def _test_positional_encoding():
import matplotlib.pyplot as plt
plt.figure(figsize=(15, 5))
pe = get_positional_encoding(20, 100)
plt.plot(np.arange(100), pe[:, 0, 4:8].numpy())
plt.legend(["dim %d" % p for p in [4, 5, 6, 7]])
plt.title("Positional encoding")
plt.show()
if __name__ == '__main__':
_test_positional_encoding()
labml_nn/transformers/primer_ez/__init__.py
+129
View File
@@ -0,0 +1,129 @@
"""
---
title: "Primer: Searching for Efficient Transformers for Language Modeling"
summary: >
This is an annotated implementation/tutorial of
Primer: Searching for Efficient Transformers for Language Modeling for Vision in PyTorch.
---
# Primer: Searching for Efficient Transformers for Language Modeling
This is a [PyTorch](https://pytorch.org) implementation of the paper
[Primer: Searching for Efficient Transformers for Language Modeling](https://arxiv.org/abs/2109.08668).
The authors do an evolutionary search for transformer architectures.
They name the architecture found using the search Primer (PRIMitives searched transformER).
**Primer EZ** is the architecture with the two most robust modifications in Primer compared to
the original transformer.
Primer EZ trains a lot faster than the vanilla transformer.
### Squared ReLU
The most effective modification found by the search is using a square ReLU instead of ReLU in
the [position-wise feedforward module](../feed_forward.html).
$$y = {\max(x, 0)}^2$$
### Multi-DConv-Head Attention (MDHA)
The next effective modification is a depth-wise $3 \times 1$ convolution after multi-head projection
for queries, keys, and values.
The convolution is along the sequence dimension and per channel (depth-wise).
To be clear, if the number of channels in each head is $d_k$ the convolution will have $1 \times 3$
kernels for each of the $d_k$ channels.
[Here is the experiment code](experiment.html), for Primer EZ.
"""
import torch
from torch import nn
from labml_nn.transformers import MultiHeadAttention
class SquaredReLU(nn.Module):
"""
## Squared ReLU activation
$$y = {\max(x, 0)}^2$$
Squared ReLU is used as the activation function in the
[position wise feedforward module](../feed_forward.html).
"""
def __init__(self):
super().__init__()
self.relu = nn.ReLU()
def forward(self, x: torch.Tensor):
# Apply ReLU
x = self.relu(x)
# Square it
return x * x
class SpatialDepthWiseConvolution(nn.Module):
"""
## Spatial Depth Wise Convolution
"""
def __init__(self, d_k: int, kernel_size: int = 3):
"""
* `d_k` is the number of channels in each head
"""
super().__init__()
self.kernel_size = kernel_size
# We use PyTorch's `Conv1d` module.
# We set the number of groups to be equal to the number of channels so that it does a separate convolution
# (with different kernels) for each channel.
# We add padding to both sides and later crop the right most `kernel_size - 1` results
self.conv = nn.Conv1d(in_channels=d_k, out_channels=d_k,
kernel_size=(kernel_size,), padding=(kernel_size - 1,), groups=d_k)
def forward(self, x: torch.Tensor):
"""
`x` has shape `[seq_len, batch_size, heads, d_k]`
"""
# Get the shape
seq_len, batch_size, heads, d_k = x.shape
# Permute to `[batch_size, heads, d_k, seq_len]`
x = x.permute(1, 2, 3, 0)
# Change the shape to `[batch_size * heads, d_k, seq_len]`
x = x.view(batch_size * heads, d_k, seq_len)
# 1D convolution accepts input of the form `[N, channels, sequence]`
x = self.conv(x)
# Crop the right most `kernel_size - 1` results since we padded both sides
x = x[:, :, :-(self.kernel_size - 1)]
# Reshape to `[batch_size, heads, d_k, seq_len]`
x = x.view(batch_size, heads, d_k, seq_len)
# Permute to `[seq_len, batch_size, heads, d_k]`
x = x.permute(3, 0, 1, 2)
#
return x
class MultiDConvHeadAttention(MultiHeadAttention):
"""
## Multi-DConv-Head Attention (MDHA)
We extend our original implementation of [Multi-Head Attention](../mha.html#MHA)
and add the spatial depth-wise convolution to query, key and value projections.
"""
def __init__(self, heads: int, d_model: int, dropout_prob: float = 0.1):
super().__init__(heads, d_model, dropout_prob)
# [Multi-Head Attention](../mha.html#MHA) will create query, key and value projection modules
# `self.query`, `self.key`, and `self.value`.
#
# We combine a spatial depth-wise convolution layer to each of them and replace
# `self.query`, `self.key`, and `self.value`.
#
# 📝 *We feel this cleaner implementation is easier to understand since it clearly shows the difference
# between this and vanilla transformer multi-head attention*.
self.query = nn.Sequential(self.query, SpatialDepthWiseConvolution(self.d_k))
self.key = nn.Sequential(self.key, SpatialDepthWiseConvolution(self.d_k))
self.value = nn.Sequential(self.value, SpatialDepthWiseConvolution(self.d_k))
labml_nn/transformers/primer_ez/efficient.py
+60
View File
@@ -0,0 +1,60 @@
import math
import torch
from torch import nn
from labml_nn.transformers import MultiHeadAttention
class SpatialDepthWiseConvolution(nn.Module):
"""
## Spatial Depth Wise Convolution
This is actually slower
"""
def __init__(self, d_k: int, kernel_size: int = 3):
"""
* `d_k` is the number of channels in each head
"""
super().__init__()
self.kernel_size = kernel_size
# We use PyTorch's `Conv1d` module.
# We set the number of groups to be equal to the number of channels so that it does a separate convolution
# (with different kernels) for each channel.
# We add padding to both sides and later crop the right most `kernel_size - 1` results
rng = 1 / math.sqrt(kernel_size)
self.kernels = nn.Parameter(torch.zeros((kernel_size, d_k)).uniform_(-rng, rng))
def forward(self, x: torch.Tensor):
"""
`x` has shape `[seq_len, batch_size, heads, d_k]`
"""
res = x * self.kernels[0].view(1, 1, 1, -1)
for i in range(1, len(self.kernels)):
res[i:] += x[:-i] * self.kernels[i].view(1, 1, 1, -1)
return res
class MultiDConvHeadAttention(MultiHeadAttention):
"""
## Multi-DConv-Head Attention (MDHA)
We extend our original implementation of [Multi-Head Attention](../mha.html#MHA)
and add the spatial depth-wise convolution to query, key and value projections.
"""
def __init__(self, heads: int, d_model: int, dropout_prob: float = 0.1):
super().__init__(heads, d_model, dropout_prob)
# [Multi-Head Attention](../mha.html#MHA) will create query, key and value projection modules
# `self.query`, `self.key`, and `self.value`.
#
# We combine a spatial depth-wise convolution layer to each of them and replace
# `self.query`, `self.key`, and `self.value`.
self.query = nn.Sequential(self.query, SpatialDepthWiseConvolution(self.d_k))
self.key = nn.Sequential(self.key, SpatialDepthWiseConvolution(self.d_k))
self.value = nn.Sequential(self.value, SpatialDepthWiseConvolution(self.d_k))
labml_nn/transformers/primer_ez/experiment.py
+126
View File
@@ -0,0 +1,126 @@
"""
---
title: Primer EZ experiment
summary: This experiment trains Primer EZ on Tiny Shakespeare dataset.
---
# [Primer EZ](index.html) Experiment
This is an annotated PyTorch experiment to train a [Primer EZ transformer](index.html).
This is based on our [vanilla transformer experiment](../basic/experiment.html).
We use the same experiment and add the Primer EZ modifications.
"""
from labml import experiment
from labml.configs import option
from labml_nn.transformers import TransformerConfigs
from labml_nn.transformers.basic.autoregressive_experiment import Configs
from labml_nn.transformers.configs import FeedForwardConfigs
from labml_nn.transformers.primer_ez import SquaredReLU
@option(FeedForwardConfigs.activation, 'SquaredReLU')
def _squared_relu():
"""
Add the [option](https://docs.labml.ai/api/configs.html#labml.configs.option)
of [**squared ReLU**](index.html) to [configurable](../configs.html#FFN)
[feed forward module](../feed_forward.html).
"""
return SquaredReLU()
@option(TransformerConfigs.encoder_attn, 'MultiDConvHeadAttention')
def _d_conv_mha(c: TransformerConfigs):
"""
Add the [option](https://docs.labml.ai/api/configs.html#labml.configs.option)
of [**Multi-DConv-Head Attention**](index.html) to
[configurable transformer](../configs.html#TransformerConfigs)
"""
from labml_nn.transformers.primer_ez import MultiDConvHeadAttention
return MultiDConvHeadAttention(c.n_heads, c.d_model, dropout_prob=c.dropout)
@option(TransformerConfigs.encoder_attn, 'MultiDSharedConvHeadAttention')
def _d_shared_conv_mha(c: TransformerConfigs):
"""
Add the [option](https://docs.labml.ai/api/configs.html#labml.configs.option)
of [**Multi Depth-wise Shared Conv Head Attention**](variations.html) to
[configurable transformer](../configs.html#TransformerConfigs)
📝 *This is a variation we tried*
"""
from labml_nn.transformers.primer_ez.variations import MultiDSharedConvHeadAttention
return MultiDSharedConvHeadAttention(c.n_heads, c.d_model, dropout_prob=c.dropout)
@option(TransformerConfigs.encoder_attn, 'MultiDPHConvHeadAttention')
def _d_per_head_conv_mha(c: TransformerConfigs):
"""
Add the [option](https://docs.labml.ai/api/configs.html#labml.configs.option)
of [**Multi Depth-wise Per Head Conv Head Attention**](variation.html) to
[configurable transformer](../configs.html#TransformerConfigs)
📝 *This is a variation we tried*
"""
from labml_nn.transformers.primer_ez.variations import MultiDPHConvHeadAttention
return MultiDPHConvHeadAttention(c.n_heads, c.d_model, dropout_prob=c.dropout)
def main():
# Create experiment
experiment.create(name="primer_ez")
# Create configs
conf = Configs()
# Override configurations
experiment.configs(conf, {
# Use character level tokenizer
'tokenizer': 'character',
# Prompt separator is blank
'prompt_separator': '',
# Starting prompt for sampling
'prompt': 'It is ',
# Use Tiny Shakespeare dataset
'text': 'tiny_shakespeare',
# Use a context size of $256$
'seq_len': 256,
# Train for $128$ epochs
'epochs': 128,
# Batch size $32$
'batch_size': 32,
# Switch between training and validation for $10$ times
# per epoch
'inner_iterations': 10,
# Model size
'd_model': 512,
'transformer.ffn.d_ff': 2048,
# Use Adam optimizer
'optimizer.optimizer': 'Adam',
'optimizer.learning_rate': 2.5e-4,
# ⭐️ Use [**squared ReLU**](index.html) activation in the feed forward network.
#
# *Replace this with `ReLU` for $ReLU$.*
'transformer.ffn.activation': 'SquaredReLU',
# ⭐️ Use [**Multi-DConv-Head Attention**](index.html) for encoder attention.
#
# *Replace this with `mha` for original multi-head attention.*
'transformer.encoder_attn': 'MultiDConvHeadAttention',
})
# Set models for saving and loading
experiment.add_pytorch_models({'model': conf.model})
# Start the experiment
with experiment.start():
# Run training
conf.run()
#
if __name__ == '__main__':
main()
labml_nn/transformers/primer_ez/readme.md
+25
View File
@@ -0,0 +1,25 @@
# [Primer: Searching for Efficient Transformers for Language Modeling](https://nn.labml.ai/transformers/primer_ez/index.html)
This is a [PyTorch](https://pytorch.org) implementation of the paper
[Primer: Searching for Efficient Transformers for Language Modeling](https://arxiv.org/abs/2109.08668).
The authors do an evolutionary search for transformer architectures.
They name the architecture found using the search as Primer (PRIMitives searched transformER).
**Primer EZ** is the architecture with the two most robust modifications in Primer compared to
the original transformer.
Primer EZ trains a lot faster than the vanilla transformer.
### Squared ReLU
The most effective modification found by the search is using a square ReLU instead of ReLU in
the [position-wise feedforward module](https://nn.labml.ai/transformers/feed_forward.html).
### Multi-DConv-Head Attention (MDHA)
The next effective modification is a depth-wise 3 X 1 convolution after multi-head projection
for queries, keys, and values.
The convolution is along the sequence dimension and per channel (depth-wise).
To be clear, if the number of channels in each head is d_k the convolution will have 1 X 3
kernels for each of the d_k channels.
[Here is the experiment code](https://nn.labml.ai/transformers/primer_ez/experiment.html), for Primer EZ.
labml_nn/transformers/primer_ez/variations.py
+143
View File
@@ -0,0 +1,143 @@
"""
---
title: Primer EZ variations
summary: We tried some variations to Primer EZ.
---
# [Primer EZ](index.html) Variations
We tried some variations to see which changes in Primer EZ has most benefits.
"""
import torch
from torch import nn
from labml_nn.transformers import MultiHeadAttention
class SpatialDepthWiseSharedConvolution(nn.Module):
"""
## Spatial Depth Wise Shared Convolution
We share the same kernel across all channels.
"""
def __init__(self, kernel_size: int = 3):
"""
"""
super().__init__()
self.kernel_size = kernel_size
# We use PyTorch's `Conv1d` module.
# We add padding to both sides and later crop the right most `kernel_size - 1` results
self.conv = nn.Conv1d(in_channels=1, out_channels=1, kernel_size=(kernel_size,), padding=(kernel_size - 1,))
def forward(self, x: torch.Tensor):
"""
`x` has shape `[seq_len, batch_size, heads, d_k]`
"""
# Get the shape
seq_len, batch_size, heads, d_k = x.shape
# Permute to `[batch_size, heads, d_k, seq_len]`
x = x.permute(1, 2, 3, 0)
# Change the shape to `[batch_size * heads * d_k, seq_len]`
x = x.view(batch_size * heads * d_k, 1, seq_len)
# 1D convolution accepts input of the form `[N, channels, sequence]`
x = self.conv(x)
# Crop the right most `kernel_size - 1` results since we padded both sides
x = x[:, :, :-(self.kernel_size - 1)]
# Reshape to `[batch_size, heads, d_k, seq_len]`
x = x.view(batch_size, heads, d_k, seq_len)
# Permute to `[seq_len, batch_size, heads, d_k]`
x = x.permute(3, 0, 1, 2)
#
return x
class MultiDSharedConvHeadAttention(MultiHeadAttention):
"""
## Multi-Depth-wise-Shared-Conv-Head Attention
We extend our original implementation of [Multi-Head Attention](../mha.html#MHA)
and add the spatial depth-wise shared convolution to query, key and value projections.
"""
def __init__(self, heads: int, d_model: int, dropout_prob: float = 0.1):
super().__init__(heads, d_model, dropout_prob)
# [Multi-Head Attention](../mha.html#MHA) will create query, key and value projection modules
# `self.query`, `self.key`, and `self.value`.
#
# We combine a spatial depth-wise shared convolution layer to each of them and replace
# `self.query`, `self.key`, and `self.value`.
self.query = nn.Sequential(self.query, SpatialDepthWiseSharedConvolution())
self.key = nn.Sequential(self.key, SpatialDepthWiseSharedConvolution())
self.value = nn.Sequential(self.value, SpatialDepthWiseSharedConvolution())
class SpatialDepthWisePerHeadConvolution(nn.Module):
"""
## Spatial Depth Wise Per Head Convolution
"""
def __init__(self, heads: int, d_k: int, kernel_size: int = 3):
"""
* `heads` is the number of heads
* `d_k` is the number of channels in each head
"""
super().__init__()
self.kernel_size = kernel_size
# We use PyTorch's `Conv1d` module.
# We set the number of groups to be equal to the number of channels from each head
# so that it does a separate convolution
# (with different kernels) for each channel and head.
# We add padding to both sides and later crop the right most `kernel_size - 1` results
self.conv = nn.Conv1d(in_channels=d_k * heads, out_channels=d_k * heads,
kernel_size=(kernel_size,), padding=(kernel_size - 1,), groups=d_k * heads)
def forward(self, x: torch.Tensor):
"""
`x` has shape `[seq_len, batch_size, heads, d_k]`
"""
# Get the shape
seq_len, batch_size, heads, d_k = x.shape
# Permute to `[batch_size, heads, d_k, seq_len]`
x = x.permute(1, 2, 3, 0)
# Change the shape to `[batch_size heads * d_k, seq_len]`
x = x.view(batch_size, heads * d_k, seq_len)
# 1D convolution accepts input of the form `[N, channels, sequence]`
x = self.conv(x)
# Crop the right most `kernel_size - 1` results since we padded both sides
x = x[:, :, :-(self.kernel_size - 1)]
# Reshape to `[batch_size, heads, d_k, seq_len]`
x = x.view(batch_size, heads, d_k, seq_len)
# Permute to `[seq_len, batch_size, heads, d_k]`
x = x.permute(3, 0, 1, 2)
#
return x
class MultiDPHConvHeadAttention(MultiHeadAttention):
"""
## Multi-per-Head-Depth-wise-Conv-Head Attention
We extend our original implementation of [Multi-Head Attention](../mha.html#MHA)
and add the spatial depth-wise convolution to query, key and value projections.
"""
def __init__(self, heads: int, d_model: int, dropout_prob: float = 0.1):
super().__init__(heads, d_model, dropout_prob)
# [Multi-Head Attention](../mha.html#MHA) will create query, key and value projection modules
# `self.query`, `self.key`, and `self.value`.
#
# We combine a spatial per-head depth-wise convolution layer to each of them and replace
# `self.query`, `self.key`, and `self.value`.
self.query = nn.Sequential(self.query, SpatialDepthWisePerHeadConvolution(heads, self.d_k))
self.key = nn.Sequential(self.key, SpatialDepthWisePerHeadConvolution(heads, self.d_k))
self.value = nn.Sequential(self.value, SpatialDepthWisePerHeadConvolution(heads, self.d_k))
labml_nn/transformers/relative_mha.py
+7
View File
@@ -0,0 +1,7 @@
"""
---
title: Relative Multi-Headed Attention
summary: Relative Multi-Headed Attention from paper Transformer-XL.
redirect: https://nn.labml.ai/transformers/xl/relative_mha.html
---
"""
labml_nn/transformers/retro/__init__.py
+35
View File
@@ -0,0 +1,35 @@
"""
---
title: Retrieval-Enhanced Transformer (Retro)
summary: >
This is a PyTorch implementation/tutorial of the paper
Improving language models by retrieving from trillions of tokens.
It builds a key-value database of chunks of text and retrieves and uses them when
making predictions.
---
# Retrieval-Enhanced Transformer (Retro)
This is a [PyTorch](https://pytorch.org) implementation of the paper
[Improving language models by retrieving from trillions of tokens](https://arxiv.org/abs/2112.04426).
It builds a database of chunks of text.
It is a key-value database where the keys are indexed by the BERT embeddings of the chunks.
They use a frozen pre-trained BERT model to calculate these embeddings.
The values are the corresponding chunks and an equal length of text proceeding that chunk.
Then the model retrieves text similar (nearest neighbors) to the input to the model from this database.
These retrieved texts are used to predict the output.
Since we use a frozen BERT model for retrieval we can pre-calculate all the nearest neighbors for the
training dataset.
This speeds up the training process.
Components:
* [BERT embeddings](bert_embeddings.html): Code to get BERT embeddings of chunks of text.
* [Key-value database](database.html): Build and retrieve chunks
* [Model](model.html)
* [Dataset](dataset.html): Pre-calculate the nearest neighbors
* [Training code](train.html)
"""
labml_nn/transformers/retro/bert_embeddings.py
+148
View File
@@ -0,0 +1,148 @@
"""
---
title: BERT Embeddings of chunks of text
summary: >
Generate BERT embeddings for chunks using a frozen BERT model
---
# BERT Embeddings of chunks of text
This is the code to get BERT embeddings of chunks for [RETRO model](index.html).
"""
from typing import List
import torch
from transformers import BertTokenizer, BertModel
from labml import lab, monit
class BERTChunkEmbeddings:
"""
## BERT Embeddings
For a given chunk of text $N$ this class generates BERT embeddings $\text{B\small{ERT}}(N)$.
$\text{B\small{ERT}}(N)$ is the average of BERT embeddings of all the tokens in $N$.
"""
def __init__(self, device: torch.device):
self.device = device
# Load the BERT tokenizer from [HuggingFace](https://huggingface.co/bert-base-uncased)
with monit.section('Load BERT tokenizer'):
self.tokenizer = BertTokenizer.from_pretrained('bert-base-uncased',
cache_dir=str(
lab.get_data_path() / 'cache' / 'bert-tokenizer'))
# Load the BERT model from [HuggingFace](https://huggingface.co/bert-base-uncased)
with monit.section('Load BERT model'):
self.model = BertModel.from_pretrained("bert-base-uncased",
cache_dir=str(lab.get_data_path() / 'cache' / 'bert-model'))
# Move the model to `device`
self.model.to(device)
@staticmethod
def _trim_chunk(chunk: str):
"""
In this implementation, we do not make chunks with a fixed number of tokens.
One of the reasons is that this implementation uses character-level tokens and BERT
uses its sub-word tokenizer.
So this method will truncate the text to make sure there are no partial tokens.
For instance, a chunk could be like `s a popular programming la`, with partial
words (partial sub-word tokens) on the ends.
We strip them off to get better BERT embeddings.
As mentioned earlier this is not necessary if we broke chunks after tokenizing.
"""
# Strip whitespace
stripped = chunk.strip()
# Break words
parts = stripped.split()
# Remove first and last pieces
stripped = stripped[len(parts[0]):-len(parts[-1])]
# Remove whitespace
stripped = stripped.strip()
# If empty return original string
if not stripped:
return chunk
# Otherwise, return the stripped string
else:
return stripped
def __call__(self, chunks: List[str]):
"""
### Get $\text{B\small{ERT}}(N)$ for a list of chunks.
"""
# We don't need to compute gradients
with torch.no_grad():
# Trim the chunks
trimmed_chunks = [self._trim_chunk(c) for c in chunks]
# Tokenize the chunks with BERT tokenizer
tokens = self.tokenizer(trimmed_chunks, return_tensors='pt', add_special_tokens=False, padding=True)
# Move token ids, attention mask and token types to the device
input_ids = tokens['input_ids'].to(self.device)
attention_mask = tokens['attention_mask'].to(self.device)
token_type_ids = tokens['token_type_ids'].to(self.device)
# Evaluate the model
output = self.model(input_ids=input_ids,
attention_mask=attention_mask,
token_type_ids=token_type_ids)
# Get the token embeddings
state = output['last_hidden_state']
# Calculate the average token embeddings.
# Note that the attention mask is `0` if the token is empty padded.
# We get empty tokens because the chunks are of different lengths.
emb = (state * attention_mask[:, :, None]).sum(dim=1) / attention_mask[:, :, None].sum(dim=1)
#
return emb
def _test():
"""
### Code to test BERT embeddings
"""
from labml.logger import inspect
# Initialize
device = torch.device('cuda:0')
bert = BERTChunkEmbeddings(device)
# Sample
text = ["Replace me by any text you'd like.",
"Second sentence"]
# Check BERT tokenizer
encoded_input = bert.tokenizer(text, return_tensors='pt', add_special_tokens=False, padding=True)
inspect(encoded_input, _expand=True)
# Check BERT model outputs
output = bert.model(input_ids=encoded_input['input_ids'].to(device),
attention_mask=encoded_input['attention_mask'].to(device),
token_type_ids=encoded_input['token_type_ids'].to(device))
inspect({'last_hidden_state': output['last_hidden_state'],
'pooler_output': output['pooler_output']},
_expand=True)
# Check recreating text from token ids
inspect(bert.tokenizer.convert_ids_to_tokens(encoded_input['input_ids'][0]), _n=-1)
inspect(bert.tokenizer.convert_ids_to_tokens(encoded_input['input_ids'][1]), _n=-1)
# Get chunk embeddings
inspect(bert(text))
#
if __name__ == '__main__':
_test()
labml_nn/transformers/retro/database.py
+156
View File
@@ -0,0 +1,156 @@
"""
---
title: Database for nearest neighbor retrieval
summary: >
Nearest neighbor retrieval and creation of the database
---
# Database for nearest neighbor retrieval
This is the build the database and retrieves nearest neighbors for
[RETRO model](index.html).
We use [FAISS library](https://faiss.ai/) for the database whilst the paper had used the SCaNN library.
"""
from typing import List, Optional
import faiss
import numpy as np
import torch
from labml import lab, monit
from labml_nn.helpers.datasets import TextFileDataset
from labml_nn.transformers.retro.bert_embeddings import BERTChunkEmbeddings
def build_database(chunk_len: int = 16, batch_size: int = 64, d_emb: int = 768, n_centeroids: int = 256,
code_size: int = 64, n_probe: int = 8, n_train: int = 50_000):
"""
## Build Database
* `chunk_len` is the length of a chunk (number of characters)
* `batch_size` is the batch size to use when calculating $\text{B\small{ERT}}(N)$
* `d_emb` is the number of features in $\text{B\small{ERT}}(N)$ embeddings
[lists to select in FAISS index](https://faiss.ai/cpp_api/struct/structfaiss_1_1IndexIVFPQ.html)
* `n_centeroids` is the number of lists in the index
* `code_size` encoded vector size in the index
* `n_probe` is the number of lists to probe
* `n_train' is the number of keys to train the index on
"""
# Load the dataset text file
dataset = TextFileDataset(
lab.get_data_path() / 'tiny_shakespeare.txt',
list,
url='https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt')
# Get training data (a string)
text = dataset.train
# Split the text into chunks of `chunk_length`
chunks = [text[i:i + chunk_len] for i in range(0, len(text), chunk_len) if i + chunk_len * 2 < len(text)]
# Get the offsets of each of the chunks
chunk_offsets = np.array([i for i in range(0, len(text), chunk_len) if i + chunk_len * 2 < len(text)])
# Number of chunks
n_chunks = len(chunks)
# Initialize BERT to get $\text{B\small{ERT}}(N)$
bert = BERTChunkEmbeddings(torch.device('cuda:0'))
# Get chunk embeddings by processing `batch_size` number of chunks on each iteration
chunk_emb = []
for i in monit.iterate('Get embeddings', range(0, n_chunks, batch_size)):
chunk_emb.append(bert(chunks[i: i + batch_size]).cpu())
# Merge them into a single tensor
chunk_emb = torch.cat(chunk_emb, dim=0).numpy()
# Create the [FAISS index](https://faiss.ai/cpp_api/struct/structfaiss_1_1IndexIVFPQ.html)
quantizer = faiss.IndexFlatL2(d_emb)
index = faiss.IndexIVFPQ(quantizer, d_emb, n_centeroids, code_size, 8)
index.nprobe = n_probe
# Get a random sample of the the chunk indexes
random_sample = np.random.choice(np.arange(n_chunks), size=[min(n_train, n_chunks)], replace=False)
# Train the index to store the keys
with monit.section('Train index'):
index.train(chunk_emb[random_sample])
# Add the chunks to the index in batches of size `1024`
for s in monit.iterate('Index', range(0, n_chunks, 1024)):
e = min(s + 1024, n_chunks)
# Add to index
index.add_with_ids(chunk_emb[s:e], chunk_offsets[s: e])
# Save the index
with monit.section('Save'):
faiss.write_index(index, str(lab.get_data_path() / 'retro.index'))
class RetroIndex:
"""
## Index for retrieving nearest neighbors
"""
def __init__(self, chunk_len: int = 16, n_probe: int = 8,
n_neighbors: int = 2, n_extra: int = 2,
exclude_neighbor_span: int = 8):
"""
* `chunk_len` is the chunk length
* `n_probe` is the number of lists to probe
* `n_neighbors` is the number of neighbors to retrieve
* `n_extra` is the number of extra neighbors to retrieve since we will be
removing neighbors overlapping with the query chunk
* `exclude_neighbor_span` is the extra text length to avoid when checking for overlaps
"""
self.n_neighbors = n_neighbors
self.chunk_len = chunk_len
self.exclude_neighbor_span = exclude_neighbor_span
self.n_extra = n_extra
# Initialize BERT to get $\text{B\small{ERT}}(N)$
self.bert = BERTChunkEmbeddings(torch.device('cuda:0'))
# Load the database
with monit.section('Load index'):
self.index = faiss.read_index(str(lab.get_data_path() / 'retro.index'))
self.index.nprobe = n_probe
def filter_neighbors(self, offset: int, neighbor_offsets: List[int]):
"""
#### Filter neighbors that overlap with the query
The positions of the neighbors are given by `neighbor_offsets` and the position
of the query chunk is `offset`.
"""
return [n for n in neighbor_offsets
if n < offset - (self.chunk_len + self.exclude_neighbor_span)
or n > offset + (self.chunk_len + self.exclude_neighbor_span)]
def __call__(self, query_chunks: List[str], offsets: Optional[List[int]]):
"""
### Retrieve nearest neighbors
"""
# Get $\text{B\small{ERT}}(N)$ of query chunks
emb = self.bert(query_chunks).cpu()
# Get `n_neighbors + n_extra` nearest neighbors from the database
distance, neighbor_offsets = self.index.search(emb.numpy(), self.n_neighbors + self.n_extra)
# If the query chunk offsets are given filter out overlapping chunks
if offsets is not None:
neighbor_offsets = [self.filter_neighbors(off, n_off)
for off, n_off in zip(offsets, neighbor_offsets)]
# Get the closest `n_neighbors` after filtering
neighbor_offsets = [n_off[:self.n_neighbors] for n_off in neighbor_offsets]
#
return neighbor_offsets
#
if __name__ == '__main__':
build_database()
labml_nn/transformers/retro/dataset.py
+137
View File
@@ -0,0 +1,137 @@
"""
---
title: Training dataset for RETRO
summary: >
Create a dataset for RETRO model training
---
# RETRO training dataset
We pre-retrieve nearest neighbors from the [key-value database](database.html)
and create the dataset to train the [RETRO](index.html)
[model](model.html).
"""
import json
from pathlib import Path
import numpy as np
import torch
from torch.utils.data import Dataset as PyTorchDataset
from labml import lab, monit
from labml_nn.helpers.datasets import TextFileDataset, TextDataset
from labml_nn.transformers.retro.database import RetroIndex
def build_dataset(chunk_len: int = 16, chunks_per_sample: int = 32, skip_range: int = 8):
"""
## Build the dataset
* `chunk_len` is the chunk length
* `chunks_per_sample` is the number of chunks per training sample
* `skip_range` is the maximum number of characters to skip between two samples.
We skip a few characters between samples to make sure the samples
aren't aligned perfectly with the chunks in the [database](database.html)
"""
# Load the text file
dataset = TextFileDataset(
lab.get_data_path() / 'tiny_shakespeare.txt',
list,
url='https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt')
# Training portion of it
text = dataset.train
# Load the index for retrieving neighbors
index = RetroIndex()
# The input sample offsets
sample_offsets = []
# Cursor for the text
i = 0
while i < len(text):
# Skip a few characters to make sure it's not aligned with the neighbors
skip = np.random.randint(skip_range)
i += skip
# Stop if we've reached the end of the text
if i + chunks_per_sample * chunk_len > len(text):
break
# Collect the offset
sample_offsets.append(i)
# Increment the cursor
i += chunks_per_sample * chunk_len
# For samples
samples = []
# Iterate through sample offsets
for i in monit.iterate('Gather Neighbors', sample_offsets):
# Get the sample including an extra character (for prediction)
sample = text[i: i + chunks_per_sample * chunk_len + 1]
# The input
src = sample[:-1]
# Break it into chunks
chunks = [src[j:j + chunk_len] for j in range(0, len(src), chunk_len)]
# The chunk offsets
chunk_offsets = [j + i for j in range(0, len(src), chunk_len)]
# Retrieve nearest neighbors
neighbor_offsets = index(chunks, chunk_offsets)
# Get neighbor texts. The neighbor length is twice the `chunk_len`
neighbors = [[text[j: j + chunk_len * 2] for j in n_off] for n_off in neighbor_offsets]
# Add to list of samples
samples.append((sample[:-1], sample[1:], neighbors))
# Save the samples in JSON.
# We don't need to use complex dataset storage mechanisms or pre-tokenize
# since our dataset is small.
with open(str(lab.get_data_path() / 'retro_train_dataset.json'), 'w') as f:
f.write(json.dumps(samples))
class Dataset(PyTorchDataset):
"""
## Dataset
This is the PyTorch dataset that loads the dataset created
by `build_dataset`.
"""
def __init__(self, file_path: Path, tds: TextDataset):
"""
* `file_path` is the path of the saved JSON file
* `tds` is the `TextDataset`
"""
self.tds = tds
# Load the samples
with open(str(file_path), 'r') as f:
self.samples = json.loads(f.read())
def __len__(self):
"""
Number of samples
"""
return len(self.samples)
def __getitem__(self, idx: int):
"""
Get a sample
"""
# Get the sample
s = self.samples[idx]
# Tokenize
src = self.tds.text_to_i(s[0])
tgt = self.tds.text_to_i(s[1])
neighbors = torch.stack([torch.stack([self.tds.text_to_i(n) for n in chunks]) for chunks in s[2]])
#
return src, tgt, neighbors
#
if __name__ == '__main__':
build_dataset()
labml_nn/transformers/retro/model.py
+619
View File
@@ -0,0 +1,619 @@
"""
---
title: RETRO model
summary: >
RETRO model with encoder for neighbors and autoregressive decoder
---
# RETRO model
This is the model definition for
[RETRO](index.html).
"""
import math
from typing import Set
import torch
from torch import nn
from labml.logger import inspect
class RotaryPositionalEmbeddings(nn.Module):
"""
## [RoPE embeddings](../rope/index.html)
*We use rotary position embeddings in self-attention layers.
We assume the positional information gets embedded in embeddings
and therefore not use them in causal attention.
[Non-causal self-attention needs explicit positional information
because it cannot infer it](https://arxiv.org/abs/3999902edc8511eba3db37f65e372566).*
"""
def __init__(self, d: int, base: int = 10_000):
"""
* `d` is the number of features $d$
* `base` is the constant used for calculating $\Theta$
"""
super().__init__()
# $\Theta = {\theta_i = 10000^{\frac{2(i-1)}{d}}, i \in [1, 2, ..., \frac{d}{2}]}$
self.theta = nn.Parameter(1. / (base ** (torch.arange(0, d, 2).float() / d)), requires_grad=False)
def forward(self, x: torch.Tensor):
"""
* `x` is the Tensor at the head of a key or a query with shape `[ batch_size, seq_len, n_heads, d]`
"""
# Extract the shape
batch_size, seq_len, n_heads, d = x.shape
# $\frac{d}{2}$
d_2 = d // 2
# Create position indexes `[0, 1, ..., seq_len - 1]`
seq_idx = torch.arange(seq_len, device=x.device).type_as(self.theta)
# Calculate the product of position index and $\theta_i$
idx_theta = torch.einsum('n,d->nd', seq_idx, self.theta)
# Concatenate so that for row $m$ we have
# $[m \theta_0, m \theta_1, ..., m \theta_{\frac{d}{2}}, m \theta 0, m \theta 1, ..., m \theta_{\frac{d}{2}}]$
idx_theta2 = torch.cat([idx_theta, idx_theta], dim=1)
# Calculate
# $[-x^{(\frac{d}{2} + 1)}, -x^{(\frac{d}{2} + 2)}, ..., -x^{(d)}, x^{(1)}, x^{(2)}, ..., -x^{(\frac{d}{2})}]$
neg_half_x = torch.cat([-x[:, :, :, d_2:], x[:, :, :, :d_2]], dim=-1)
# Calculate
#
# \begin{align}
# \begin{pmatrix}
# x^{(i)}_m \cos m \theta_i - x^{(i + \frac{d}{2})}_m \sin m \theta_i \\
# x^{(i + \frac{d}{2})}_m \cos m\theta_i + x^{(i)}_m \sin m \theta_i \\
# \end{pmatrix} \\
# \end{align}
#
# for $i \in {1, 2, ..., \frac{d}{2}}$
rx = (x * idx_theta2.cos()[None, :, None, :]) + (neg_half_x * idx_theta2.sin()[None, :, None, :])
#
return rx
class SelfAttention(nn.Module):
"""
## Self-Attention Layer $\text{A\small{TTN}}$
This applies causal and non-causal [multi-headed self-attention](../mha.html).
"""
def __init__(self, d_model: int, n_heads: int, d_k: int, is_causal: bool):
"""
* `d_model` is the number of features in transformer embeddings
* `n_heads` is the number of attention heads
* `d_k` is the number of features per head
* `is_causal` indicates whether this is causal attention (masked)
"""
super().__init__()
self.is_causal = is_causal
self.n_heads = n_heads
self.d_k = d_k
# To scale attentions before softmax by $\frac{1}{\sqrt{d_k}}$
self.scale = 1 / math.sqrt(self.d_k)
# Linear layers for query, key and value heads.
self.query = nn.Linear(d_model, n_heads * d_k)
self.key = nn.Linear(d_model, n_heads * d_k)
self.value = nn.Linear(d_model, n_heads * d_k)
# Pre-norm layer. The paper uses RMSNorm instead.
self.norm = nn.LayerNorm(d_model)
# Softmax for attention probabilities
self.softmax = nn.Softmax(dim=-1)
# Rotary positional embeddings
self.rotary_pe = RotaryPositionalEmbeddings(self.d_k)
# Final linear layer
self.output = nn.Linear(n_heads * d_k, d_model)
def mask_attention(self, attn: torch.Tensor):
"""
### Mask the attention layer for causal attention
* `attn` is the attention matrix of shape `[batch_size, n_heads, seq_len, seq_len]`
"""
# No masking for non-causal attention
if not self.is_causal:
return attn
# Create a triangular mask
mask = torch.tril(attn.new_ones(attn.shape[-2:]))
# Filter by the mask
return attn.masked_fill(mask == 0, float('-inf'))
def forward(self, h: torch.Tensor):
"""
* `h` is the transformer embeddings of shape `[batch_size, seq_len, d_model]`
"""
# Residual connection
h_res = h
# Pre-normalization
h = self.norm(h)
# Get query, key, and values and split them in to heads.
# These will have shapes `[batch_size, seq_len, n_heads, d_k]`
mh_shape = (*h.shape[:-1], self.n_heads, self.d_k)
q = self.query(h).view(mh_shape)
k = self.key(h).view(mh_shape)
v = self.value(h).view(mh_shape)
# Apply rotary positional embeddings
q = self.rotary_pe(q)
k = self.rotary_pe(k)
# Calculate attentions
attn = torch.einsum('bihd,bjhd->bhij', q, k)
# Scale it by $\frac{1}{\sqrt{d_k}}$
attn = attn * self.scale
# Apply masks if it's causal attention
attn = self.mask_attention(attn)
# Calculate attention probabilities
attn = self.softmax(attn)
# Get values
h = torch.einsum("bhij,bjhd->bihd", attn, v)
# Change from shape `[batch_size, seq_len, n_heads, d_k]`
# to `[batch_size, seq_len, n_heads * d_k]`
h = h.reshape(*h.shape[:-2], -1)
# Apply final linear layer.
# The result will have shape `[batch_size, seq_len, d_model]`
h = self.output(h)
# Add the residual connection
return h + h_res
class CrossAttention(nn.Module):
"""
## Cross-Attention Layer $\text{C\small{A}}$
This is similar to the self-attention layer defined above, except that
it gets keys and values from a different set of embeddings than the queries.
This is used in the encoder to encode the retrieved chunks based on the
input chunks.
*We do not use any explicit positional embeddings here.
We assume that the model can represent positional information in the embeddings implicitly.*
"""
def __init__(self, d_model: int, n_heads: int, d_k: int):
"""
* `d_model` is the number of features in transformer embeddings
* `n_heads` is the number of attention heads
* `d_k` is the number of features per head
"""
super().__init__()
self.n_heads = n_heads
self.d_k = d_k
# To scale attentions before softmax by $\frac{1}{\sqrt{d_k}}$
self.scale = 1 / math.sqrt(self.d_k)
# Linear layers for query, key and value heads.
self.query = nn.Linear(d_model, n_heads * d_k)
self.key = nn.Linear(d_model, n_heads * d_k)
self.value = nn.Linear(d_model, n_heads * d_k)
# Pre-norm layer for the query embeddings. The paper uses RMSNorm instead.
self.norm = nn.LayerNorm(d_model)
# Softmax for attention probabilities
self.softmax = nn.Softmax(dim=-1)
# Final linear layer
self.output = nn.Linear(n_heads * d_k, d_model)
def forward(self, e: torch.Tensor, h: torch.Tensor):
"""
* `e` are the retrieved nearest neighbor chunk embeddings with shape
`[batch_size, chunks, neighbors, neighbor_len, d_model]`
* `h` are the input chunks from which the nearest neighbors were retrieved with shape
`[batch_size, chunks, chunk_len, d_model]`. This is already normalized.
"""
# Residual connection
e_res = e
# Normalize retrieved chunks
e = self.norm(e)
# Get query from the retrieved chunks
q = self.query(e).view(*e.shape[:-1], self.n_heads, self.d_k)
# Get keys and values from the input chunks
k = self.key(h).view(*h.shape[:-1], self.n_heads, self.d_k)
v = self.value(h).view(*h.shape[:-1], self.n_heads, self.d_k)
# Calculate attention scores for all chunks.
# Each retrieved neighbor will pay attention to the original chunk that retrieved it.
# This will have shape `[batch_size, chunks, neighbors, n_heads, neighbor_len, chunk_len]`
attn = torch.einsum('bcnihd,bcjhd->bcnhij', q, k)
# Scale attention scores
attn = attn * self.scale
# Calculate softmax across the last dimension
attn = self.softmax(attn)
# Gather values
e = torch.einsum("bcnhij,bcjhd->bcnihd", attn, v)
# Change from shape `[batch_size, chunks, neighbors, neighbor_len, n_heads, d_k]`
# to `[batch_size, chunks, neighbors, neighbor_len, n_heads * d_k]`
e = e.reshape(*e.shape[:-2], -1)
# Apply final linear layer.
# The result will have shape `[batch_size, chunks, neighbors, neighbor_len, d_model]`
e = self.output(e)
# Add residual connection
return e + e_res
class ChunkedCrossAttention(nn.Module):
"""
## Chunked Cross-Attention Layer $\text{C\small{CA}}$
This is similar to the cross-attention layer defined above.
This is used in the decoder to pay attention to the retrieved neighbor chunks.
*We do not use any explicit positional embeddings here.
We assume that the model can represent positional information in the embeddings implicitly.*
"""
def __init__(self, d_model: int, n_heads: int, d_k: int, chunk_len: int):
"""
* `d_model` is the number of features in transformer embeddings
* `n_heads` is the number of attention heads
* `d_k` is the number of features per head
* `chunk_len` is the length of a chunk
"""
super().__init__()
self.chunk_len = chunk_len
self.n_heads = n_heads
self.d_k = d_k
# To scale attentions before softmax by $\frac{1}{\sqrt{d_k}}$
self.scale = 1 / math.sqrt(self.d_k)
# Linear layers for query, key and value heads.
self.query = nn.Linear(d_model, n_heads * d_k)
self.key = nn.Linear(d_model, n_heads * d_k)
self.value = nn.Linear(d_model, n_heads * d_k)
# Pre-norm layer for the query embeddings. The paper uses RMSNorm instead.
self.norm = nn.LayerNorm(d_model)
# Softmax for attention probabilities
self.softmax = nn.Softmax(dim=-1)
# Final linear layer
self.output = nn.Linear(n_heads * d_k, d_model)
def forward(self, h: torch.Tensor, e: torch.Tensor):
"""
`h` are the input embeddings of shape `[batch_size, seq_len, d_model]`
`e` are the retrieved nearest neighbors of shape `[batch_size, chunks, neighbors, neighbor_len, d_model]`
"""
# Get shape
batch_size, chunks, neighbors, neighbor_len, d_model = e.shape
# No attention if there are no chunks (for short inputs when sampling)
if chunks == 0:
return h
# Residual connection
h_res = h
# Remove the first `chunk_len - 1` embeddings.
# The input pays attention to neighbors retrieved and encoded using the past tokens only;
# so that there is no information leakage.
# That is the retrieved neighbors from the first chunks will have information from the first chunk.
# So by shifting the sequence to the left by `chunk_len - 1` we make sure that information only flows
# to the right.
h = h[:, self.chunk_len - 1:]
# Pre-norm
h = self.norm(h)
# Append empty embeddings to the end to be able to split the input into chunks
if h.shape[1] < chunks * self.chunk_len:
h = torch.cat((h, h.new_zeros(batch_size, chunks * self.chunk_len - h.shape[1], d_model)), dim=1)
# Reshape the input into chunks.
h = h.reshape(batch_size, chunks, self.chunk_len, d_model)
# Get query from the input
q = self.query(h).view(*h.shape[:-1], self.n_heads, self.d_k)
# Get keys and values from the retrieved neighbors
k = self.key(e).view(*e.shape[:-1], self.n_heads, self.d_k)
v = self.value(e).view(*e.shape[:-1], self.n_heads, self.d_k)
# Calculate attention scores for input chunks.
# Each chunk will pay attention to neighbors retrieved by the previous chunk.
# This will have shape `[batch_size, chunks, heads, chunk_len, neighbors, neighbor_len]`
attn = torch.einsum('bcihd,bcnjhd->bchinj', q, k)
# Scale attention scores
attn = attn * self.scale
# Apply softmax over the last two dimensions `neighbors, neighbor_len`
attn = self.softmax(attn.view(*attn.shape[:-2], -1)).view(attn.shape)
# Gather values
h = torch.einsum("bchinj,bcnjhd->bcihd", attn, v)
# Change from shape `[batch_size, chunks, chunk_len, n_heads, d_k]`
# to `[batch_size, chunks * chunk_len, n_heads * d_k]`
h = h.reshape(batch_size, chunks * self.chunk_len, -1)
# Apply final linear layer.
# The result will have shape `[batch_size, chunks * chunk_len, d_model]`
h = self.output(h)
# Append `chunk_len - 1` zero embedding to the left; i.e. right shift it back
h = torch.cat((h.new_zeros(batch_size, self.chunk_len - 1, d_model), h), dim=1)
# Truncate and add the residual connection
return h[:, :h_res.shape[1]] + h_res
class FeedForward(nn.Module):
"""
### Position-wise Feed Forward Layer $\text{F\small{FW}}$
This consists of two linear layers and an activation in the middle.
"""
def __init__(self, d_model: int, d_ff: int):
"""
* `d_model` is the number of features in transformer embeddings
* `d_ff` is the number features in the hidden layer
"""
super().__init__()
# The two linear layers
self.lin1 = nn.Linear(d_model, d_ff)
self.lin2 = nn.Linear(d_ff, d_model)
# ReLU Activation
self.act = nn.ReLU()
# Pre-norm layer
self.norm = nn.LayerNorm(d_model)
def forward(self, h: torch.Tensor):
"""
`h` are the embeddings of shape `[batch_size, seq_len, d_model]`
"""
# Residual
h_res = h
# Pre-norm
h = self.norm(h)
# First linear layer
h = self.lin1(h)
# Activation
h = self.act(h)
# Second linear layer
h = self.lin2(h)
# Add the residual connection
return h + h_res
class NearestNeighborEncoder(nn.Module):
"""
## Nearest Neighbor Encoder $\text{E\small{NCODER}}(\text{R\small{ET}}(C_u)_{1 \le u \le l}, H)$
This module encodes the retrieved nearest neighbors
"""
def __init__(self, chunk_len: int, n_layers: int, ca_layers: Set[int],
d_model: int, n_heads: int, d_k: int, d_ff: int):
"""
* `chunk_len` is the length of a chunk
* `n_layer` is the number of layers in the encoder $L_{\text{enc}}$
* `ca_layers` are the layers with cross attention $P_{\text{enc}}$
* `d_model` is the number of features in embeddings
* `n_heads` is the number of heads in attention layers
* `d_k` is the size of attention heads
* `d_ff` is the size of the feed-forward networks hidden layers
"""
super().__init__()
self.ca_layers = ca_layers
self.chunk_len = chunk_len
# Cross-attention layers
self.ca = nn.ModuleList([CrossAttention(d_model, n_heads, d_k) for _ in range(len(ca_layers))])
# Bi-directional self attention layers
self.attn = nn.ModuleList([SelfAttention(d_model, n_heads, d_k, is_causal=False) for _ in range(n_layers)])
# Feed forward layers
self.ffw = nn.ModuleList([FeedForward(d_model, d_ff) for _ in range(n_layers)])
# Pre-normalization layer for $H$
self.norm_h = nn.LayerNorm(d_model)
def forward(self, e: torch.Tensor, h: torch.Tensor):
"""
* `e` are token embeddings of the retrieved nearest neighbors,
$\text{E\small{MB}}\big(\text{R\small{ET}}(C_u)_{1 \le u \le l}\big)$
of shape `[batch_size, chunks, neighbors, neighbor_len, d_model]`
* `h` is are the input token embeddings, $H$
of shape `[batch_size, seq_len, d_model]`
*The chunks $u \in [1, l]$ and neighbors $j \in [1, k]$ are processed in parallel.*
"""
# Get shape
batch_size, chunks, neighbors, neighbor_len, d_model = e.shape
# $(H_u)_{u \in [1, l]} \leftarrow \text{S\small{PLIT}}(H)$
h_split = h[:, :self.chunk_len * chunks, :].reshape(batch_size, chunks, self.chunk_len, d_model)
# Pre-norm
h_split = self.norm_h(h_split)
# Keep the index of the cross attention layer
p_ca = 0
# For all layers $p' \in [1, L_{\text{enc}}]$
for p in range(len(self.attn)):
# Bi-directional self attention
# $E^j_u \leftarrow \text{A\small{TTN}}_{\text{enc}}(E^j_u)$
e = self.attn[p](e.view(-1, neighbor_len, d_model)).view(e.shape)
# Cross attention if $p' \in P_{\text{enc}}$
if p in self.ca_layers:
# $E^j_u \leftarrow \text{C\small{A}}_{\text{enc}}(E^j_u, H_u)$
e = self.ca[p_ca](e, h_split)
# Incremnt the cross attention index
p_ca += 1
# Feed forward layer $E^j_u \leftarrow \text{F\small{FW}}_{\text{enc}}(E^j_u)$
e = self.ffw[p](e)
# return $E$
return e
class RetroModel(nn.Module):
"""
## Retro Model
This is the Retro decoder
"""
def __init__(self, n_vocab: int, d_model: int, n_layers: int, ca_layers: Set[int], chunk_len: int,
n_heads: int, d_k: int, d_ff: int, encoder: NearestNeighborEncoder):
"""
* `v_vocab` is the number of tokens in the vocabulary
* `d_model` is the number of features in embeddings
* `n_layers` is the number of layers in the decoder $L$
* `ca_layers` are the layers with cross attention $P$
* `chunk_len` is the length of a chunk
* `n_heads` is the number of heads in attention layers
* `d_k` is the size of attention heads
* `d_ff` is the size of the feed-forward networks hidden layers
* `encoder` is the nearest neighbor encoder
"""
super().__init__()
self.ca_layers = ca_layers
self.encoder = encoder
# Token embedding layer
self.emb = nn.Embedding(n_vocab, d_model)
# Chunked cross attention layers $\text{C\small{CA}}$
self.cca = nn.ModuleList(
[ChunkedCrossAttention(d_model, n_heads, d_k, chunk_len) for _ in range(len(ca_layers))])
# Attention layers $\text{A\small{TTN}}$
self.attn = nn.ModuleList([SelfAttention(d_model, n_heads, d_k, is_causal=True) for _ in range(n_layers)])
# Feed forward layers $\text{F\small{FW}}$
self.ffw = nn.ModuleList([FeedForward(d_model, d_ff) for _ in range(n_layers)])
# Readout layer $\text{R\small{EAD}}$
self.read = nn.Linear(d_model, n_vocab)
# Pre-normalization layer for nearest neighbor embeddings from
# $\text{E\small{NCODER}}(\text{R\small{ET}}(C_u)_{1 \le u \le l}, H)$
self.norm_e = nn.LayerNorm(d_model)
def forward(self, x: torch.Tensor, ret: torch.Tensor):
"""
* `x` is the input sequence, $X$ of shape `[batch_size, seq_len]`
* `ret` are the retrieved neighbors
$\text{R\small{ET}}(C_u)_{1 \le u \le l}$
of shape `[batch_size, chunks, neighbors, neighbor_len]`
"""
# Get input embeddings $H \leftarrow \text{E\small{MB}}(X)$
h = self.emb(x)
# Embeddings of the retrieved neighbors
# $E^j_u = \text{E\small{MB}}_{\text{enc}}\big(\text{R\small{ET}}(C_u)^j\big)$.
#
# We use same embeddings for both input and neighbors
ret_emb = self.emb(ret)
# Keep index of the chunked cross attention layer
p_ca = 0
# For all layers $p \in [1, L]$
for p in range(len(self.attn)):
# Causal self attention $H \leftarrow \text{A\small{TTN}}(H)$
h = self.attn[p](h)
# Get encoder embeddings before the first $\text{C\small{CA}}$ layer,
# when $p = \min(P)$
if self.ca_layers and p == min(self.ca_layers):
# $E = \text{E\small{NCODER}}(\text{R\small{ET}}(C_u)_{1 \le u \le l}, H)$
#
# We passed the embeddings of $\text{R\small{ET}}(C_u)_{1 \le u \le l}$ to encoder.
e = self.encoder(ret_emb, h)
# Normalize encoder embeddings
e = self.norm_e(e)
# Chunked-cross attention if $p \in P$
if p in self.ca_layers:
# $H \leftarrow \text{C\small{CA}}(H, E)$
h = self.cca[p_ca](h, e)
# Increment chunked cross-attention index
p_ca += 1
# $H \leftarrow \text{F\small{FW}}(H)$
h = self.ffw[p](h)
# $O \leftarrow \text{R\small{EAD}}(H)$
return self.read(h)
def _test():
"""
### Test the model with fake data
"""
chunk_len = 4
d_model = 8
d_ff = 32
n_heads = 2
d_k = 4
device = torch.device('cuda:0')
m = RetroModel(5, d_model, 6, {2, 5}, chunk_len, n_heads, d_k, d_ff,
encoder=NearestNeighborEncoder(chunk_len, 2, {1}, d_model, n_heads, d_k, d_ff))
m.to(device)
x = [1, 2, 4, 4, 0, 1, 2, 3, 4, 3]
ret = [
[[0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1]],
[[0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1]],
]
res = m(torch.tensor([x] * 10).to(device), torch.tensor([ret] * 10).to(device))
inspect(res)
#
if __name__ == '__main__':
_test()
labml_nn/transformers/retro/train.py
+223
View File
@@ -0,0 +1,223 @@
"""
---
title: RETRO training
summary: >
Training RETRO model with Tiny Shakespeare dataset
---
# RETRO training
This is the training code for
[RETRO](index.html).
"""
import torch
from labml import monit, lab, tracker, experiment, logger
from labml.logger import Text
from labml_nn.helpers.datasets import TextFileDataset
from labml_nn.optimizers.noam import Noam
from labml_nn.transformers.retro import model as retro
from labml_nn.transformers.retro.dataset import Dataset, RetroIndex
from labml_nn.transformers.retro.model import RetroModel, NearestNeighborEncoder
from torch import nn
from torch.utils.data import DataLoader, RandomSampler
class Sampler:
"""
## Sampler
This class greedily samples from a model.
"""
def __init__(self, device: torch.device, model: retro.RetroModel, tds: TextFileDataset, chunk_len: int):
"""
* `device` is the device of the model
* `model` is the [Retro mode](retro.html)
* `tds` is the text dataset (used to get neighbor chunks)
* `chunk_len` is the length of a chunk
"""
self.chunk_len = chunk_len
self.tds = tds
self.model = model
self.device = device
# [Retro index](database.html)
self.index = RetroIndex()
def retrieve_nearest_neighbours(self, chunk: str):
"""
### Retrieve nearest neighbors of a given chunk
"""
# Retrieve the offsets of the nearest neighbors
neighbor_offsets = self.index([chunk], None)
# Get the neighbors (with neighbor length equal to `chunk_len * 2`)
text = self.tds.train
neighbors = [text[j: j + self.chunk_len * 2] for j in neighbor_offsets[0]]
#
return neighbors
def sample(self, prompt: str, sample_len: int):
"""
### Sample text from the given prompt
"""
# To store nearest neighbors as strings
neighbors_str = []
# Sampled text
sampled = ''
# Sample `sample_len` tokens
for i in range(sample_len):
# We need to retrieve neighbors,
# if there are more sampled chunks than we have already retrieved for
while len(neighbors_str) < len(prompt) // self.chunk_len:
# Get the last chunk for which we haven't retrieved neighbors
off = len(neighbors_str) * self.chunk_len
chunk = prompt[off: off + self.chunk_len]
# Retrieve nearest neighbors
neighbors_str.append(self.retrieve_nearest_neighbours(chunk))
# Tokenize the input
src = self.tds.text_to_i(prompt)
# Tokenize the retrieved neighbors
neighbors = torch.stack([torch.stack([self.tds.text_to_i(n) for n in chunk]) for chunk in neighbors_str])
# Move them to the same device as the model
src = src.to(self.device)
neighbors = neighbors.to(self.device)
# Get model output
res = self.model(src[None, :], neighbors[None, :, :, :])
# Greedily sample the last token
token = res[0, -1, :].argmax(dim=-1)
# Add the sampled token text to the prompt and sample text
prompt += self.tds.itos[token.item()]
sampled += self.tds.itos[token.item()]
#
return sampled
class Trainer:
"""
## Retro trainer
"""
def __init__(self, device: torch.device, model: retro.RetroModel,
dataloader: DataLoader, optimizer: torch.optim.Optimizer):
"""
* `device` is the device of the model
* `model` is the [Retro mode](retro.html)
* `dataloader` is the dataloader for the [dataset with pre-retrieved neighbors](dataset.html)
* `optimizer` is the optimizer
"""
self.optimizer = optimizer
self.device = device
self.dataloader = dataloader
self.model = model
self.loss_func = nn.CrossEntropyLoss()
def __call__(self):
"""
### Train the model for an epoch
"""
# Iterate through training data
for i, (src, tgt, neighbors) in monit.enum('Train', self.dataloader):
# Move data to the device
src, tgt, neighbors = src.to(self.device), tgt.to(self.device), neighbors.to(self.device)
# Forward pass
res = self.model(src, neighbors)
# Calculate loss
loss = self.loss_func(res.view(-1, res.shape[-1]), tgt.view(-1))
# Clear the gradients
self.optimizer.zero_grad()
# Backward pass
loss.backward()
# Optimize the model
self.optimizer.step()
# Save training statistics and increment the global step counter
tracker.save({'loss.train': loss})
tracker.add_global_step(len(src))
def train():
"""
## Create and train a small model
"""
# Create an experiment
experiment.create(name='retro_small')
# GPU device
device = torch.device('cuda:0')
# Load Tiny Shakespeare dataset
tds = TextFileDataset(
lab.get_data_path() / 'tiny_shakespeare.txt',
list,
url='https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt')
# Load [Retro dataset](dataset.html)
train_dataset = Dataset(lab.get_data_path() / 'retro_train_dataset.json', tds)
# Create dataloader
train_dl = DataLoader(train_dataset,
batch_size=4,
sampler=RandomSampler(train_dataset, replacement=True))
# Hyper-parameters
chunk_len = 16
d_model = 128
d_ff = 512
n_heads = 16
d_k = 16
# Create the nearest neighbor encoder
nearest_neighbor_encoder = NearestNeighborEncoder(chunk_len, 6, {3}, d_model, n_heads, d_k, d_ff)
# Create the model
model = RetroModel(tds.n_tokens, d_model, 6,
{3, 5},
chunk_len, n_heads, d_k, d_ff,
encoder=nearest_neighbor_encoder)
# Move the model to the device
model = model.to(device)
# Create the optimizer
optimizer = Noam(model.parameters(), lr=1., d_model=d_model, warmup=2_000)
# Create the `Trainer`
trainer = Trainer(device, model, train_dl, optimizer)
# Create the `Sampler`
sampler = Sampler(device, model, tds, chunk_len)
#
prompt = '''Second Citizen:\nOne word, good citizens.\n\nFirst Citizen:'''
# Set models for saving and loading
experiment.add_pytorch_models(model=model)
# Start the experiment
with experiment.start():
# Train for `32` epochs
for epoch in monit.loop(32):
# Train
trainer()
# Print a new line
tracker.new_line()
# Sample from the `prompt`
logger.log([(prompt.replace('\n', '\\n\n'), Text.subtle),
(sampler.sample(prompt, 128).replace('\n', '\\n\n'), Text.none)])
# Save models
#
if __name__ == '__main__':
train()
labml_nn/transformers/rope/__init__.py
+235
View File
@@ -0,0 +1,235 @@
"""
---
title: Rotary Positional Embeddings (RoPE)
summary: >
Annotated implementation of RoPE from paper
RoFormer: Enhanced Transformer with Rotary Position Embedding
---
# Rotary Positional Embeddings (RoPE)
This is an implementation of
[Rotary Positional Embeddings (RoPE)](https://arxiv.org/abs/2104.09864)
in [PyTorch](https://pytorch.org).
Rotary Positional Embeddings (RoPE) encode position information of tokens
with a rotation matrix that naturally incorporates explicit relative position
dependency.
Here's [the training code](experiment.html) for training a transformer model with RoPE
on Tiny Shakespeare dataset.
"""
import torch
from torch import nn
from labml.logger import inspect
from labml_nn.transformers.mha import MultiHeadAttention
class RotaryPositionalEmbeddings(nn.Module):
"""
## RoPE module
Rotary encoding transforms pairs of features by rotating in the 2D plane.
That is, it organizes the $d$ features as $\frac{d}{2}$ pairs.
Each pair can be considered a coordinate in a 2D plane, and the encoding will rotate it
by an angle depending on the position of the token.
### For a pair of features
Let $x^{(1)}_m$ and $x^{(2)}_m$ be two features of the
key or query of any head at position $m$.
Or for simplicity assume $x$ has only two features.
Then the transformation is,
\begin{align}
RoPE\big(x^{(1)}_m, x^{(2)}_m, m\big) &=
\begin{pmatrix}
\cos m \theta & - \sin m \theta \\
\sin m \theta & \cos m \theta
\end{pmatrix}
\begin{pmatrix}
x^{(1)}_m \\
x^{(2)}_m \\
\end{pmatrix} \\
&=
\begin{pmatrix}
x^{(1)}_m \cos m\theta - x^{(2)}_m \sin m \theta \\
x^{(2)}_m \cos m\theta + x^{(1)}_m \sin m \theta \\
\end{pmatrix} \\
\end{align}
where $\theta$ is a constant angle. The other pairs of features are transformed similarly.
### Attention is relative
For a pair of features, dot-product attention score between two positions $m$ and $n$ would be
\begin{align}
\Big \langle RoPE\big(x^{(1)}_m, x^{(2)}_m, m\big), RoPE\big(x^{(1)}_n, x^{(2)}_n, n\big) \Big \rangle &= \\
(x^{(1)}_m \cos m\theta - x^{(2)}_m \sin m \theta)(x^{(1)}_n \cos n\theta - x^{(2)}_n \sin n \theta) &+ \\
(x^{(2)}_m \cos m\theta + x^{(1)}_m \sin m \theta)(x^{(2)}_n \cos n\theta + x^{(1)}_n \sin n \theta) &= \\
x^{(1)}_m x^{(1)}_n (\cos m\theta \cos n\theta + \sin m \theta \sin n \theta) &+ \\
x^{(1)}_m x^{(2)}_n (-\cos m\theta \sin n\theta + \sin m \theta \cos n \theta) &+ \\
x^{(2)}_m x^{(1)}_n (-\sin m\theta \cos n\theta + \cos m \theta \sin n \theta) &+ \\
x^{(2)}_m x^{(2)}_n (\sin m\theta \sin n\theta + \cos m \theta \cos n \theta) &= \\
x^{(1)}_m x^{(1)}_n \cos (m - n) \theta +
x^{(1)}_m x^{(2)}_n \sin(m - n) \theta &+ \\
- x^{(2)}_m x^{(1)}_n \sin (m - n) \theta +
x^{(2)}_m x^{(2)}_n \cos (m - n) \theta &= \\
\big(x^{(1)}_m \cos (m - n)\theta - x^{(2)}_m \sin (m - n) \theta\big) x^{(1)}_n &+ \\
\big(x^{(2)}_m \cos (m - n)\theta + x^{(1)}_m \sin (m - n) \theta\big) x^{(2)}_n &= \\
\Big \langle RoPE\big(x^{(1)}_m, x^{(2)}_m, m - n\big), RoPE\big(x^{(1)}_n, x^{(2)}_n, 0\big) \Big \rangle
\end{align}
This shows that for dot-production attention the rotary encodings gives relative attention.
### For all features
The features are grouped into pairs and handled as above. They use a different $\theta$ for each pair.
The paper suggests using $\Theta = {\theta_i = 10000^{\frac{2(i-1)}{d}}, i \in [1, 2, ..., \frac{d}{2}]}$
for the $\frac{d}{2}$ pairs of features.
The original implementation of RoPE divide the $d$-dimension features into $\frac{d}{2}$ pairs of features ($i$, $i + 1$).
In this implementation we pair feature $i$ with feature $i + \frac{d}{2}$. So for position $m$ we transform
\begin{align}
\begin{pmatrix}
x^{(i)}_m \\
x^{(i + \frac{d}{2})}_m
\end{pmatrix}
\end{align}
to
\begin{align}
\begin{pmatrix}
x^{(i)}_m \cos m \theta_i - x^{(i + \frac{d}{2})}_m \sin m \theta_i \\
x^{(i + \frac{d}{2})}_m \cos m\theta_i + x^{(i)}_m \sin m \theta_i \\
\end{pmatrix} \\
\end{align}
"""
def __init__(self, d: int, base: int = 10_000):
"""
* `d` is the number of features $d$
* `base` is the constant used for calculating $\Theta$
"""
super().__init__()
self.base = base
self.d = d
self.cos_cached = None
self.sin_cached = None
def _build_cache(self, x: torch.Tensor):
"""
Cache $\cos$ and $\sin$ values
"""
# Return if cache is already built
if self.cos_cached is not None and x.shape[0] <= self.cos_cached.shape[0]:
return
# Get sequence length
seq_len = x.shape[0]
# $\Theta = {\theta_i = 10000^{-\frac{2(i-1)}{d}}, i \in [1, 2, ..., \frac{d}{2}]}$
theta = 1. / (self.base ** (torch.arange(0, self.d, 2).float() / self.d)).to(x.device)
# Create position indexes `[0, 1, ..., seq_len - 1]`
seq_idx = torch.arange(seq_len, device=x.device).float().to(x.device)
# Calculate the product of position index and $\theta_i$
idx_theta = torch.einsum('n,d->nd', seq_idx, theta)
# Concatenate so that for row $m$ we have
# $[m \theta_0, m \theta_1, ..., m \theta_{\frac{d}{2}}, m \theta_0, m \theta_1, ..., m \theta_{\frac{d}{2}}]$
idx_theta2 = torch.cat([idx_theta, idx_theta], dim=1)
# Cache them
self.cos_cached = idx_theta2.cos()[:, None, None, :]
self.sin_cached = idx_theta2.sin()[:, None, None, :]
def _neg_half(self, x: torch.Tensor):
# $\frac{d}{2}$
d_2 = self.d // 2
# Calculate $[-x^{(\frac{d}{2} + 1)}, -x^{(\frac{d}{2} + 2)}, ..., -x^{(d)}, x^{(1)}, x^{(2)}, ..., x^{(\frac{d}{2})}]$
return torch.cat([-x[:, :, :, d_2:], x[:, :, :, :d_2]], dim=-1)
def forward(self, x: torch.Tensor):
"""
* `x` is the Tensor at the head of a key or a query with shape `[seq_len, batch_size, n_heads, d]`
"""
# Cache $\cos$ and $\sin$ values
self._build_cache(x)
# Sequence length
seq_len = x.shape[0]
# Split the features, we can choose to apply rotary embeddings only to a partial set of features.
x_rope, x_pass = x[..., :self.d], x[..., self.d:]
# Calculate
# $[-x^{(\frac{d}{2} + 1)}, -x^{(\frac{d}{2} + 2)}, ..., -x^{(d)}, x^{(1)}, x^{(2)}, ..., x^{(\frac{d}{2})}]$
neg_half_x = self._neg_half(x_rope)
# Calculate
#
# \begin{align}
# \begin{pmatrix}
# x^{(i)}_m \cos m \theta_i - x^{(i + \frac{d}{2})}_m \sin m \theta_i \\
# x^{(i + \frac{d}{2})}_m \cos m\theta_i + x^{(i)}_m \sin m \theta_i \\
# \end{pmatrix} \\
# \end{align}
#
# for $i \in {1, 2, ..., \frac{d}{2}}$
x_rope = (x_rope * self.cos_cached[:seq_len]) + (neg_half_x * self.sin_cached[:seq_len])
#
return torch.cat((x_rope, x_pass), dim=-1)
class RotaryPEMultiHeadAttention(MultiHeadAttention):
"""
## Multi-head attention with rotary positional embeddings
We override [multi-head attention from original transformer](../mha.html).
"""
def __init__(self, heads: int, d_model: int, rope_percentage: float = 0.5, dropout_prob: float = 0.0):
super().__init__(heads, d_model, dropout_prob)
# Rotary positional embedding layers
d_rope = int(self.d_k * rope_percentage)
self.query_rotary_pe = RotaryPositionalEmbeddings(d_rope)
self.key_rotary_pe = RotaryPositionalEmbeddings(d_rope)
def get_scores(self, query: torch.Tensor, key: torch.Tensor):
"""
### Calculate scores between queries and keys
"""
# Calculate dot-product with RoPE
return torch.einsum('ibhd,jbhd->ijbh', self.query_rotary_pe(query), self.key_rotary_pe(key))
def _test_rotary():
"""
Testing RoPE with a simple example
"""
x = torch.tensor([[1, 2, 3, 4], [4, 5, 6, 7], [7, 8, 9, 10]], dtype=torch.float)
x = x[:, None, None, :]
inspect(x)
rotary_pe = RotaryPositionalEmbeddings(4)
inspect(rotary_pe(x))
if __name__ == '__main__':
_test_rotary()
labml_nn/transformers/rope/experiment.py
+102
View File
@@ -0,0 +1,102 @@
"""
---
title: Rotary Positional Embeddings (RoPE) Experiment
summary: This experiment trains a transformer model with Rotary Positional Embeddings (RoPE) on tiny Shakespeare dataset.
---
# Rotary Positional Embeddings (RoPE) Experiment
This is an annotated PyTorch experiment to train a transformer model with Rotary Positional Embeddings (RoPE).
"""
from labml import experiment
from labml.configs import option, calculate
from labml_nn.transformers import TransformerConfigs
from labml_nn.transformers.basic.autoregressive_experiment import AutoregressiveTransformer, Configs
# ### Rotary PE attention
def _rotary_pe_mha(c: TransformerConfigs):
from labml_nn.transformers.rope import RotaryPEMultiHeadAttention
return RotaryPEMultiHeadAttention(c.n_heads, c.d_model, 1.)
# Configuration options
calculate(TransformerConfigs.encoder_attn, 'rotary', _rotary_pe_mha)
calculate(TransformerConfigs.decoder_attn, 'rotary', _rotary_pe_mha)
calculate(TransformerConfigs.decoder_mem_attn, 'rotary', _rotary_pe_mha)
@option(Configs.model, 'rotary_pe_transformer')
def _model(c: Configs):
"""
Create an autoregressive model and initialize weights
"""
m = AutoregressiveTransformer(c.transformer.encoder,
c.transformer.src_embed,
c.transformer.generator).to(c.device)
return m
def main():
# Create experiment
experiment.create(name="rotary_pe_transformer", writers={'screen'})
# Create configs
conf = Configs()
# Override configurations
experiment.configs(conf, {
# No fixed positional embeddings
'transformer.src_embed': 'no_pos',
'transformer.tgt_embed': 'no_pos',
# Encoder with RoPE
'transformer.encoder_attn': 'rotary',
#
'model': 'rotary_pe_transformer',
# Use character level tokenizer
'tokenizer': 'character',
# Prompt separator is blank
'prompt_separator': '',
# Starting prompt for sampling
'prompt': 'It is ',
# Use Tiny Shakespeare dataset
'text': 'tiny_shakespeare',
# Use a context size of $256$
'seq_len': 512,
# Train for 32 epochs
'epochs': 32,
# Batch size $4$
'batch_size': 4,
# Switch between training and validation for $10$ times
# per epoch
'inner_iterations': 10,
# Model size
'd_model': 128,
'transformer.ffn.d_ff': 512,
'transformer.n_heads': 16,
'transformer.dropout': 0.0,
# Use [Noam optimizer](../../optimizers/noam.html)
'optimizer.optimizer': 'Noam',
'optimizer.learning_rate': 1.,
'dataloader_shuffle_with_replacement': True
})
# Set models for saving and loading
experiment.add_pytorch_models({'model': conf.model})
# Start the experiment
with experiment.start():
# Run training
conf.run()
#
if __name__ == '__main__':
main()
labml_nn/transformers/rope/value_pe/__init__.py
+246
View File
@@ -0,0 +1,246 @@
"""
---
title: Rotary Positional Embeddings with Relative distance (RoPER)
summary: >
This is an implementation of RoPER which adds relative distance information to embeddings on
top of RoPE introduced in RoFormer: Enhanced Transformer with Rotary Position Embedding
---
*RoPER is work by [Georges Harik (@gharik)](https://twitter.com/gharik),
and this implementation is based on his original code.*
# Rotary Positional Embeddings with Relative distance (RoPER)
[Rotary Positional Embeddings (RoPE)](https://arxiv.org/abs/2104.09864) includes
relative positions in attention score calculation.
However, the embeddings themselves do not get any positional information
, [except what it can get implicitly from causal attention](https://arxiv.org/abs/2c364684b15b11ecac827bce58715ee7).
RoPER adds relative positional information explicitly to value embeddings.
Specifically, it adds the relative positions of the tokens it paid attention to.
We use same rotary positional embeddings to rotate the values in attention,
Then, after taking the weighted sum,
we rotate the final in the opposite direction.
Which is equivalent to rotating each of the values (before attention) relative to the current position.
Here's [the training code](experiment.html) for training a transformer model with RoPER
on an arithmetic addition where we can see significant improvement over RoPE.
### Relative distances in embeddings
For any head, let $a_{n,i}$ be the attention from position $n$ to position $i$,
and $v_i$ be the value embeddings at position $i$. Let's denote individual features
as $v^{(1)}_i, v^{(2)}_i, \dots$.
Normally, we would take the weight sum of value embeddings
$$o^{(j)}_n = \sum_i a_{n,i} v^{(j)}_i$$
This doesn't explicitly add any distance information about the positions $i$ to final
result $o^{(j)}_n$.
RoPER pairs features like RoPE and transform them.
For a pair $v^{(1)}_m$ and $v^{(2)}_m$ it transforms them by
$RoPE\big(v^{(1)}_m, v^{(2)}_m, m\big)$.
Let us donate the transformed features with $\hat{v}^{(1)}_m, \hat{v}^{(2)}_m$.
Then it rotates the weighted sum $\hat{o}^{(j)}_n$ in the the reverse direction with
$RoPE\big(\hat{o}^{(1)}_n, \hat{o}^{(2)}_n, -n\big)$.
*Note the *$-n$.
Note that,
\begin{align}
RoPE\big(x^{(1)}_m, x^{(2)}_m, m\big) &=
\begin{pmatrix}
\cos m \theta & - \sin m \theta \\
\sin m \theta & \cos m \theta
\end{pmatrix}
\begin{pmatrix}
x^{(1)}_m \\
x^{(2)}_m \\
\end{pmatrix} \\
&=
\begin{pmatrix}
x^{(1)}_m \cos m\theta - x^{(2)}_m \sin m \theta \\
x^{(2)}_m \cos m\theta + x^{(1)}_m \sin m \theta \\
\end{pmatrix} \\
\end{align}
Final output after with the transformations is,
\begin{align}
RoPE\big(\hat{o}^{(1)}_n, \hat{o}^{(2)}_n, -n\big) &= \\
\begin{pmatrix}
\hat{o}^{(1)}_n \cos n\theta + \hat{o}^{(2)}_n \sin n \theta \\
\hat{o}^{(2)}_n \cos n\theta - \hat{o}^{(1)}_n \sin n \theta \\
\end{pmatrix} \\
\end{align}
*Note that *$\sin (-n \theta) = -\sin n \theta$.
Let's expand the first term $\hat{o}^{(1)}_n \cos n\theta + \hat{o}^{(2)}_n \sin n \theta$,
\begin{align}
\hat{o}^{(1)}_n \cos n\theta + \hat{o}^{(2)}_n \sin n \theta &= \\
\sum_i a_{n,i} \hat{v}^{(1)}_i \cos n\theta + \sum_i a_{n,i} \hat{v}^{(2)}_i \sin n \theta &= \\
\sum_i a_{n,i} \Big( v^{(1)}_i \cos i\theta - v^{(2)}_i \sin i \theta \Big) \cos n\theta &+ \\
\sum_i a_{n,i} \Big( v^{(2)}_i \cos i\theta + v^{(1)}_i \sin i \theta \Big) \sin m \theta &= \\
\sum_i a_{n,i} v^{(1)}_i \Big( \cos i\theta \cos n\theta + \sin i \theta \sin n \theta \Big) &+ \\
\sum_i a_{n,i} v^{(2)}_i \Big( \cos i\theta \sin n\theta - \sin i \theta \cos n \theta \Big) &= \\
\sum_i a_{n,i} v^{(1)}_i \cos (i - n) \theta - \sum_i a_{n,i} v^{(2)}_i \sin (i - n) \theta &= \\
\sum_i a_{n,i} v^{(1)}_i \cos (i - n) \theta - \sum_i a_{n,i} v^{(2)}_i \sin (i - n) \theta
\end{align}
Simiarly we can show the second term is equal to,
$$\sum_i a_{n,i} v^{(1)}_i \cos (i - n) \theta + \sum_i a_{n,i} v^{(2)}_i \sin (i - n) \theta$$
Which gives,
\begin{align}
RoPE\big(\hat{o}^{(1)}_n, \hat{o}^{(2)}_n, -n\big) &= \\
\begin{pmatrix}
\sum_i a_{n,i} v^{(1)}_i \cos (i - n) \theta - \sum_i a_{n,i} v^{(2)}_i \sin (i - n) \theta \\
\sum_i a_{n,i} v^{(1)}_i \cos (i - n) \theta + \sum_i a_{n,i} v^{(2)}_i \sin (i - n) \theta \\
\end{pmatrix} &= \\
\sum_i a_{n,i} RoPE \big (v^{(1)}_i, v^{(1)}_i, (i - n) \theta \big)
\end{align}
That is, the weighted average of values rotated relative to current position.
[Here's an experiment](arithmetic_experiment.html) that uses RoPER on an arthmetic addition task.
"""
from typing import Optional
import torch
from labml_nn.transformers.rope import RotaryPositionalEmbeddings, RotaryPEMultiHeadAttention
class ReverseRotaryPositionalEmbeddings(RotaryPositionalEmbeddings):
"""
## RoPE module that rotates in the opposite direction
This inherits from [RoPE rotation implementation](../index.html) and changes the direction.
"""
def forward(self, x: torch.Tensor):
"""
* `x` is the Tensor at the head of a key or a query with shape `[seq_len, batch_size, n_heads, d]`
"""
# Cache $\cos$ and $\sin$ values
self._build_cache(x)
# Split the features, we can choose to apply rotary embeddings only to a partial set of features.
x_rope, x_pass = x[..., :self.d], x[..., self.d:]
# Calculate
# $[-x^{(\frac{d}{2} + 1)}, -x^{(\frac{d}{2} + 2)}, ..., -x^{(d)}, x^{(1)}, x^{(2)}, ..., x^{(\frac{d}{2})}]$
neg_half_x = self._neg_half(x_rope)
# Calculate
#
# \begin{align}
# \begin{pmatrix}
# x^{(i)}_m \cos -m \theta_i - x^{(i + \frac{d}{2})}_m \sin -m \theta_i \\
# x^{(i + \frac{d}{2})}_m \cos -m\theta_i + x^{(i)}_m \sin -m \theta_i \\
# \end{pmatrix} = \\
# \begin{pmatrix}
# x^{(i)}_m \cos m \theta_i + x^{(i + \frac{d}{2})}_m \sin m \theta_i \\
# x^{(i + \frac{d}{2})}_m \cos m\theta_i - x^{(i)}_m \sin m \theta_i \\
# \end{pmatrix} \\
# \end{align}
#
# for $i \in {1, 2, ..., \frac{d}{2}}$
x_rope = (x_rope * self.cos_cached[:x.shape[0]]) - (neg_half_x * self.sin_cached[:x.shape[0]])
#
return torch.cat((x_rope, x_pass), dim=-1)
class RotaryValuePEMultiHeadAttention(RotaryPEMultiHeadAttention):
"""
## Multi-head attention with rotary positional embeddings
We override [multi-head attention from original transformer](../mha.html).
"""
def __init__(self, heads: int, d_model: int,
rope_percentage: float = 0.5, rope_value_percentage: float = 0.5,
dropout_prob: float = 0.0):
super().__init__(heads, d_model, rope_percentage, dropout_prob)
# Rotary positional embedding layers
d_rope_value = int(self.d_k * rope_value_percentage)
self.value_rotary_pe = RotaryPositionalEmbeddings(d_rope_value)
self.value_reverse_rotary_pe = ReverseRotaryPositionalEmbeddings(d_rope_value)
def forward(self, *,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
mask: Optional[torch.Tensor] = None):
"""
`query`, `key` and `value` are the tensors that store
collection of *query*, *key* and *value* vectors.
They have shape `[seq_len, batch_size, d_model]`.
`mask` has shape `[seq_len, seq_len, batch_size]` and
`mask[i, j, b]` indicates whether for batch `b`,
query at position `i` has access to key-value at position `j`.
"""
# `query`, `key` and `value` have shape `[seq_len, batch_size, d_model]`
seq_len, batch_size, _ = query.shape
if mask is not None:
mask = self.prepare_mask(mask, query.shape, key.shape)
# Prepare `query`, `key` and `value` for attention computation.
# These will then have shape `[seq_len, batch_size, heads, d_k]`.
query = self.query(query)
key = self.key(key)
value = self.value(value)
# Compute attention scores $Q K^\top$.
# This gives a tensor of shape `[seq_len, seq_len, batch_size, heads]`.
scores = self.get_scores(query, key)
# Scale scores $\frac{Q K^\top}{\sqrt{d_k}}$
scores *= self.scale
# Apply mask
if mask is not None:
scores = scores.masked_fill(mask == 0, float('-inf'))
# $softmax$ attention along the key sequence dimension
# $\underset{seq}{softmax}\Bigg(\frac{Q K^\top}{\sqrt{d_k}}\Bigg)$
attn = self.softmax(scores)
# Apply dropout
attn = self.dropout(attn)
# Rotate value embeddings before taking the weighted sum so that they contain positional information
value = self.value_rotary_pe(value)
# Multiply by values
# $$\underset{seq}{softmax}\Bigg(\frac{Q K^\top}{\sqrt{d_k}}\Bigg)V$$
x = torch.einsum("ijbh,jbhd->ibhd", attn, value)
# Rotate in the opposite direction so that each embedding hold the relative positions
x = self.value_reverse_rotary_pe(x)
# Save attentions for any other calculations
self.attn = attn.detach()
# Concatenate multiple heads
x = x.reshape(seq_len, batch_size, -1)
# Output layer
return self.output(x)
labml_nn/transformers/rope/value_pe/arithmetic_experiment.py
+93
View File
@@ -0,0 +1,93 @@
"""
---
title: Rotary Positional Embeddings with Relative distance (RoPER) Experiment
summary: This experiment trains a transformer model with Rotary Positional Embeddings with
Relative Distance (RoPER) on the arithmetic addition task.
---
# Rotary Positional Embeddings with Relative distance ([RoPER](index.html)) Experiment
"""
from labml import experiment
from labml.configs import calculate
from labml_nn.experiments.arithmetic_dataset import ArithmeticAutoregression
from labml_nn.transformers import TransformerConfigs
from labml_nn.transformers.rope.experiment import Configs as RoPEConfigs
class Configs(RoPEConfigs, ArithmeticAutoregression):
"""
We inherit [RoPE experiment](../experiment.html) and use it for
[arithmetic addition task](../../experiments/arithmetic_dataset.html).
We add the option to change attention to use Rotary Positional Embeddings with Relative distance (RoPER)
below.
"""
pass
def _rotary_value_pe_mha(c: TransformerConfigs):
"""
Use Rotary Positional Embeddings with Relative distance ([RoPER](index.html)) in attention.
"""
from labml_nn.transformers.rope.value_pe import RotaryValuePEMultiHeadAttention
return RotaryValuePEMultiHeadAttention(c.n_heads, c.d_model, 1., 1.)
# Configuration options
calculate(TransformerConfigs.encoder_attn, 'rotary_value', _rotary_value_pe_mha)
calculate(TransformerConfigs.decoder_attn, 'rotary_value', _rotary_value_pe_mha)
calculate(TransformerConfigs.decoder_mem_attn, 'rotary_value', _rotary_value_pe_mha)
def main():
# Create experiment
experiment.create(name="roper_addition", comment="rotary value 7", writers={'screen', 'labml'})
# Create configs
conf = Configs()
# Override configurations
experiment.configs(conf, {
'max_digits': 7,
# No fixed positional embeddings
'transformer.src_embed': 'no_pos',
'transformer.tgt_embed': 'no_pos',
# Encoder with RoPER attention
'transformer.encoder_attn': 'rotary_value',
# Encoder with RoPE attention
# 'transformer.encoder_attn': 'rotary',
#
'model': 'rotary_pe_transformer',
# Use a context size of $256$
'seq_len': 512,
# Train for 32 epochs
'epochs': 20,
# Batch size $4$
'batch_size': 16,
# Model size
'd_model': 128,
'transformer.ffn.d_ff': 512,
'transformer.n_heads': 4,
'transformer.dropout': 0.0,
# Use [Adam optimizer](../../optimizers/noam.html)
'optimizer.optimizer': 'Adam',
'optimizer.learning_rate': 2.5e-4,
})
# Set models for saving and loading
experiment.add_pytorch_models({'model': conf.model})
# Start the experiment
with experiment.start():
# Run training
conf.run()
#
if __name__ == '__main__':
main()
labml_nn/transformers/rope/value_pe/experiment.py
+96
View File
@@ -0,0 +1,96 @@
"""
---
title: Rotary Positional Embeddings (RoPE) Experiment
summary: This experiment trains a transformer model with Rotary Positional Embeddings (RoPE) on tiny Shakespeare dataset.
---
# Rotary Positional Embeddings (RoPE) Experiment
This is an annotated PyTorch experiment to train a transformer model with Rotary Positional Embeddings (RoPE).
"""
from labml import experiment
from labml.configs import calculate
from labml_nn.transformers import TransformerConfigs
from labml_nn.transformers.rope.experiment import Configs as RoPEConfigs
# ### Rotary PE attention
class Configs(RoPEConfigs): # , ArithmeticAutoregression):
pass
def _rotary_value_pe_mha(c: TransformerConfigs):
from labml_nn.transformers.rope.value_pe import RotaryValuePEMultiHeadAttention
return RotaryValuePEMultiHeadAttention(c.n_heads, c.d_model, 1., 1.)
# Configuration options
calculate(TransformerConfigs.encoder_attn, 'rotary_value', _rotary_value_pe_mha)
calculate(TransformerConfigs.decoder_attn, 'rotary_value', _rotary_value_pe_mha)
calculate(TransformerConfigs.decoder_mem_attn, 'rotary_value', _rotary_value_pe_mha)
def main():
# Create experiment
experiment.create(name="rotary_shakespeare", comment="rotary value", writers={'screen', 'labml'})
# Create configs
conf = Configs()
# Override configurations
experiment.configs(conf, {
# No fixed positional embeddings
'transformer.src_embed': 'no_pos',
'transformer.tgt_embed': 'no_pos',
# Encoder with RoPE
'transformer.encoder_attn': 'rotary_value',
# 'transformer.encoder_attn': 'rotary',
#
'model': 'rotary_pe_transformer',
# Use character level tokenizer
'tokenizer': 'character',
# Prompt separator is blank
'prompt_separator': '',
# Starting prompt for sampling
'prompt': 'It is ',
# Use Tiny Shakespeare dataset
'text': 'tiny_shakespeare',
# Use a context size of $256$
'seq_len': 512,
# Train for 32 epochs
'epochs': 24,
# Batch size $4$
'batch_size': 16,
# Switch between training and validation for $10$ times
# per epoch
'inner_iterations': 4,
# Model size
'd_model': 128,
'transformer.ffn.d_ff': 512,
'transformer.n_heads': 4,
'transformer.dropout': 0.0,
# Use [Adam optimizer](../../optimizers/noam.html)
'optimizer.optimizer': 'Adam',
'optimizer.learning_rate': 2.5e-4,
'dataloader_shuffle_with_replacement': True
})
# Set models for saving and loading
experiment.add_pytorch_models({'model': conf.model})
# Start the experiment
with experiment.start():
# Run training
conf.run()
#
if __name__ == '__main__':
main()
labml_nn/transformers/switch/__init__.py
+238
View File
@@ -0,0 +1,238 @@
"""
---
title: Switch Transformer
summary: >
This is an annotated implementation/tutorial a miniature version of Switch Transformer in PyTorch.
---
# Switch Transformer
This is a miniature [PyTorch](https://pytorch.org) implementation of the paper
[Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://arxiv.org/abs/2101.03961).
Our implementation only has a few million parameters and doesn't do model parallel distributed training.
It does single GPU training, but we implement the concept of switching as described in the paper.
The Switch Transformer uses different parameters for each token by switching among parameters
based on the token.
Therefore, only a fraction of parameters are chosen for each token.
So you can have more parameters but less computational cost.
The switching happens at the Position-wise Feedforward network (FFN) of each transformer block.
Position-wise feedforward network consists of two sequentially fully connected layers.
In switch transformer we have multiple FFNs (multiple experts),
and we chose which one to use based on a router.
The output is a set of probabilities for picking a FFN,
and we pick the one with the highest probability and only evaluate that.
So essentially the computational cost is the same as having a single FFN.
In our implementation this doesn't parallelize well when you have many or large FFNs since it's all
happening on a single GPU.
In a distributed setup you would have each FFN (each very large) on a different device.
The paper introduces another loss term to balance load among the experts (FFNs) and
discusses dropping tokens when routing is not balanced.
Here's [the training code](experiment.html) and a notebook for training a switch transformer on Tiny Shakespeare dataset.
[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/transformers/switch/experiment.ipynb)
"""
import torch
from torch import nn
from labml_nn.transformers.feed_forward import FeedForward
from labml_nn.transformers.mha import MultiHeadAttention
from labml_nn.utils import clone_module_list
class SwitchFeedForward(nn.Module):
"""
## Routing among multiple FFNs
"""
def __init__(self, *,
capacity_factor: float,
drop_tokens: bool,
is_scale_prob: bool,
n_experts: int,
expert: FeedForward,
d_model: int):
"""
* `capacity_factor` is the capacity of each expert as a factor relative to ideally balanced load
* `drop_tokens` specifies whether to drop tokens if more tokens are routed to an expert than the capacity
* `is_scale_prob` specifies whether to multiply the input to the FFN by the routing probability
* `n_experts` is the number of experts
* `expert` is the expert layer, a [FFN module](../feed_forward.html)
* `d_model` is the number of features in a token embedding
* `d_ff` is the number of features in the hidden layer of the FFN
* `dropout` is dropout probability in the FFN
"""
super().__init__()
self.capacity_factor = capacity_factor
self.is_scale_prob = is_scale_prob
self.n_experts = n_experts
self.drop_tokens = drop_tokens
# make copies of the FFNs
self.experts = clone_module_list(expert, n_experts)
# Routing layer and softmax
self.switch = nn.Linear(d_model, n_experts)
self.softmax = nn.Softmax(dim=-1)
def forward(self, x: torch.Tensor):
"""
* `x` is the input to the switching module with shape `[seq_len, batch_size, d_model]`
"""
# Capture the shape to change shapes later
seq_len, batch_size, d_model = x.shape
# Flatten the sequence and batch dimensions
x = x.view(-1, d_model)
# Get routing probabilities for each of the tokens.
# $$p_i(x) = \frac{e^{h(x)_i}}{\sum^N_j e^{h(x)_j}}$$
# where $N$ is the number of experts `n_experts` and
# $h(\cdot)$ is the linear transformation of token embeddings.
route_prob = self.softmax(self.switch(x))
# Get the maximum routing probabilities and the routes.
# We route to the expert with highest probability
route_prob_max, routes = torch.max(route_prob, dim=-1)
# Get indexes of tokens going to each expert
indexes_list = [torch.eq(routes, i).nonzero(as_tuple=True)[0] for i in range(self.n_experts)]
# Initialize an empty tensor to store outputs
final_output = x.new_zeros(x.shape)
# Capacity of each expert.
# $$\mathrm{expert\;capacity} =
# \frac{\mathrm{tokens\;per\;batch}}{\mathrm{number\;of\;experts}}
# \times \mathrm{capacity\;factor}$$
capacity = int(self.capacity_factor * len(x) / self.n_experts)
# Number of tokens routed to each expert.
counts = x.new_tensor([len(indexes_list[i]) for i in range(self.n_experts)])
# Initialize an empty list of dropped tokens
dropped = []
# Only drop tokens if `drop_tokens` is `True`.
if self.drop_tokens:
# Drop tokens in each of the experts
for i in range(self.n_experts):
# Ignore if the expert is not over capacity
if len(indexes_list[i]) <= capacity:
continue
# Shuffle indexes before dropping
indexes_list[i] = indexes_list[i][torch.randperm(len(indexes_list[i]))]
# Collect the tokens over capacity as dropped tokens
dropped.append(indexes_list[i][capacity:])
# Keep only the tokens upto the capacity of the expert
indexes_list[i] = indexes_list[i][:capacity]
# Get outputs of the expert FFNs
expert_output = [self.experts[i](x[indexes_list[i], :]) for i in range(self.n_experts)]
# Assign to final output
for i in range(self.n_experts):
final_output[indexes_list[i], :] = expert_output[i]
# Pass through the dropped tokens
if dropped:
dropped = torch.cat(dropped)
final_output[dropped, :] = x[dropped, :]
if self.is_scale_prob:
# Multiply by the expert outputs by the probabilities $y = p_i(x) E_i(x)$
final_output = final_output * route_prob_max.view(-1, 1)
else:
# Don't scale the values but multiply by $\frac{p}{\hat{p}} = 1$ so that the gradients flow
# (this is something we experimented with).
final_output = final_output * (route_prob_max / route_prob_max.detach()).view(-1, 1)
# Change the shape of the final output back to `[seq_len, batch_size, d_model]`
final_output = final_output.view(seq_len, batch_size, d_model)
# Return
#
# * the final output
# * number of tokens routed to each expert
# * sum of probabilities for each expert
# * number of tokens dropped.
# * routing probabilities of the selected experts
#
# These are used for the load balancing loss and logging
return final_output, counts, route_prob.sum(0), len(dropped), route_prob_max
class SwitchTransformerLayer(nn.Module):
"""
# Switch Transformer Block
This is the same as [normal transformer block](../models.html#TransformerLayer)
with handling extra outputs of switch feedforward module.
"""
def __init__(self, *,
d_model: int,
attn: MultiHeadAttention,
feed_forward: SwitchFeedForward,
dropout_prob: float):
"""
* `d_model` is the token embedding size
* `attn` is the attention module
* `feed_forward` is the feed forward module (which is the switching module in this case)
* `dropout_prob` is the probability of dropping out after self attention and FFN
"""
super().__init__()
self.size = d_model
self.attn = attn
self.feed_forward = feed_forward
self.dropout = nn.Dropout(dropout_prob)
self.norm_self_attn = nn.LayerNorm([d_model])
self.norm_ff = nn.LayerNorm([d_model])
def forward(self, *,
x: torch.Tensor,
mask: torch.Tensor):
# Normalize the vectors before doing self attention
z = self.norm_self_attn(x)
# Run through self attention, i.e. keys and values are from self
self_attn = self.attn(query=z, key=z, value=z, mask=mask)
# Add the self attention results
x = x + self.dropout(self_attn)
# Normalize for feed-forward
z = self.norm_ff(x)
# Pass through the switching feed-forward network
ff, counts, route_prob, n_dropped, route_prob_max = self.feed_forward(z)
# Add the feed-forward results back
x = x + self.dropout(ff)
return x, counts, route_prob, n_dropped, route_prob_max
class SwitchTransformer(nn.Module):
"""
## Switch Transformer
"""
def __init__(self, layer: SwitchTransformerLayer, n_layers: int):
super().__init__()
# Make copies of the transformer layer
self.layers = clone_module_list(layer, n_layers)
# Final normalization layer
self.norm = nn.LayerNorm([layer.size])
def forward(self, x: torch.Tensor, mask: torch.Tensor):
# Run through each transformer layer
counts, route_prob, n_dropped, route_prob_max = [], [], [], []
for layer in self.layers:
x, f, p, n_d, p_max = layer(x=x, mask=mask)
counts.append(f)
route_prob.append(p)
n_dropped.append(n_d)
route_prob_max.append(p_max)
# Finally, normalize the vectors
x = self.norm(x)
#
return x, torch.stack(counts), torch.stack(route_prob), n_dropped, torch.stack(route_prob_max)
labml_nn/transformers/switch/experiment.ipynb
+228
View File
@@ -0,0 +1,228 @@
{
"nbformat": 4,
"nbformat_minor": 0,
"metadata": {
"colab": {
"name": "Switch Transformer",
"provenance": [],
"collapsed_sections": []
},
"kernelspec": {
"name": "python3",
"display_name": "Python 3"
},
"accelerator": "GPU"
},
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "AYV_dMVDxyc2"
},
"source": [
"[![Github](https://img.shields.io/github/stars/labmlai/annotated_deep_learning_paper_implementations?style=social)](https://github.com/labmlai/annotated_deep_learning_paper_implementations)\n",
"[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/transformers/switch/experiment.ipynb) \n",
"\n",
"## Switch Transformer\n",
"\n",
"This is an experiment training Shakespeare dataset with a small Switch Transformer."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "AahG_i2y5tY9"
},
"source": [
"Install the `labml-nn` package"
]
},
{
"cell_type": "code",
"metadata": {
"id": "ZCzmCrAIVg0L",
"colab": {
"base_uri": "https://localhost:8080/"
},
"outputId": "41bb262e-d7e4-4dd9-cf8c-b2a1724889b7"
},
"source": [
"!pip install labml-nn"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "SE2VUQ6L5zxI"
},
"source": [
"Imports"
]
},
{
"cell_type": "code",
"metadata": {
"id": "0hJXx_g0wS2C"
},
"source": [
"from labml import experiment\n",
"from labml_nn.transformers.switch.experiment import Configs"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "Lpggo0wM6qb-"
},
"source": [
"Create an experiment"
]
},
{
"cell_type": "code",
"metadata": {
"id": "bFcr9k-l4cAg"
},
"source": [
"experiment.create(name=\"switch_transformer\")"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "-OnHLi626tJt"
},
"source": [
"Initialize configurations"
]
},
{
"cell_type": "code",
"metadata": {
"id": "Piz0c5f44hRo"
},
"source": [
"conf = Configs()"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "wwMzCqpD6vkL"
},
"source": [
"Set experiment configurations and assign a configurations dictionary to override configurations"
]
},
{
"cell_type": "code",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 17
},
"id": "e6hmQhTw4nks",
"outputId": "0bc4e738-adc7-4003-a030-4080df882bbb"
},
"source": [
"experiment.configs(conf,\n",
" # A dictionary of configurations to override\n",
" {'tokenizer': 'character',\n",
" 'text': 'tiny_shakespeare',\n",
" 'optimizer.learning_rate': 1.,\n",
" 'optimizer.optimizer': 'Noam',\n",
" 'prompt': 'It is',\n",
" 'prompt_separator': '',\n",
"\n",
" 'transformer': 'switch_transformer',\n",
" 'is_scale_prob': False,\n",
" 'n_experts': 4,\n",
"\n",
" 'drop_tokens': True,\n",
" 'capacity_factor': 1.2,\n",
"\n",
" 'train_loader': 'shuffled_train_loader',\n",
" 'valid_loader': 'shuffled_valid_loader',\n",
"\n",
" 'seq_len': 64,\n",
" 'epochs': 128,\n",
" 'batch_size': 32,\n",
" 'inner_iterations': 25,\n",
" })"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "EvI7MtgJ61w5"
},
"source": [
"Set PyTorch models for loading and saving"
]
},
{
"cell_type": "code",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 272
},
"id": "GDlt7dp-5ALt",
"outputId": "93e0f3b1-d0fe-4525-d9f6-9ffab9ea7f9b"
},
"source": [
"experiment.add_pytorch_models({'model': conf.model})"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "KJZRf8527GxL"
},
"source": [
"Start the experiment and run the training loop."
]
},
{
"cell_type": "code",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 1000
},
"id": "aIAWo7Fw5DR8",
"outputId": "12a92c2e-d248-436b-a6f1-7cf92b5289e9"
},
"source": [
"# Start the experiment\n",
"with experiment.start():\n",
" conf.run()"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "code",
"metadata": {
"id": "oBXXlP2b7XZO"
},
"source": [
""
],
"outputs": [],
"execution_count": null
}
]
}
labml_nn/transformers/switch/experiment.py
+236
View File
@@ -0,0 +1,236 @@
"""
---
title: Switch Transformer Experiment
summary: This experiment trains a small switch transformer on tiny Shakespeare dataset.
---
# Switch Transformer Experiment
This is an annotated PyTorch experiment to train a switch transformer.
[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/transformers/switch/experiment.ipynb)
"""
import torch
import torch.nn as nn
from labml import experiment, tracker
from labml.configs import option
from labml_nn.helpers.trainer import BatchIndex
from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs
class AutoregressiveModel(nn.Module):
"""
## Auto regressive model
"""
def __init__(self, n_vocab: int, d_model: int, transformer: nn.Module):
super().__init__()
# Token embedding module
self.src_embed = nn.Embedding(n_vocab, d_model)
# Transformer
self.transformer = transformer
# Final layer
self.generator = nn.Linear(d_model, n_vocab)
self.mask = None
def forward(self, x: torch.Tensor):
# Initialize the subsequent mask
if self.mask is None or self.mask.size(0) != len(x):
from labml_nn.transformers.utils import subsequent_mask
self.mask = subsequent_mask(len(x)).to(x.device)
# Token embeddings
x = self.src_embed(x)
# Run it through the transformer
res, counts, route_prob, n_dropped, route_prob_max = self.transformer(x, self.mask)
# Generate logits of the next token
res = self.generator(res)
#
return res, counts, route_prob, n_dropped, route_prob_max
class Configs(NLPAutoRegressionConfigs):
"""
## Configurations
This extends [`NLPAutoRegressionConfigs`](../../experiments/nlp_autoregression.html).
The default configs can and will be over-ridden when we start the experiment
"""
model: AutoregressiveModel
transformer: nn.Module
# Token embedding size
d_model: int = 128
# Number of attention heads
heads: int = 4
# Dropout probability
dropout: float = 0.0
# Number of features in FFN hidden layer
d_ff: int = 256
# Number of transformer layers
n_layers: int = 6
# Number of experts
n_experts: int = 4
# Load balancing coefficient
load_balancing_loss_ceof = 0.01
# Whether to scale the chosen expert outputs by the routing probability
is_scale_prob: bool = True
# Whether to drop tokens
drop_tokens: bool = False
# Capacity factor to determine capacity of each model
capacity_factor: float = 1.0
def init(self):
super().init()
# Initialize tracking indicators
tracker.set_scalar("lb_loss.*", False)
tracker.set_scalar("route.*", False)
tracker.set_scalar("dropped.*", False)
def step(self, batch: any, batch_idx: BatchIndex):
"""
### Training or validation step
"""
# Move data to the device
data, target = batch[0].to(self.device), batch[1].to(self.device)
# Update global step (number of tokens processed) when in training mode
if self.mode.is_train:
tracker.add_global_step(data.shape[0] * data.shape[1])
# Get model outputs.
output, counts, route_prob, n_dropped, route_prob_max = self.model(data)
# Calculate and cross entropy loss
cross_entropy_loss = self.loss_func(output, target)
# Total number of tokens processed, $T$, in the current batch $\mathscr{B}$
total = counts.sum(dim=-1, keepdims=True)
# Fraction of tokens routed to each expert
# $$f_i = \frac{1}{T} \sum_{x \in \mathscr{B}} \mathbf{1} \{ \mathop{argmax} p(x), i \}$$
# $f_i$ is the count of tokens where the argmax of $p(x)$ is equal to $i$.
route_frac = counts / total
# Mean routing probability
# $$P_i = \frac{1}{T} \sum_{x \in \mathscr{B}} p_i (x)$$
route_prob = route_prob / total
# Load balancing loss
# $$\mathscr{L} = N \sum_{i=1}^N f_i \cdot P_i$$
# $\mathscr{L}$ is the loss for a single layer and here we are
# taking the sum of losses across all layers.
load_balancing_loss = self.n_experts * (route_frac * route_prob).sum()
# Track stats
tracker.add('dropped.', total.new_tensor(n_dropped) / total)
tracker.add('route.min.', route_frac.min())
tracker.add('route.max.', route_frac.max())
tracker.add('route.std.', route_frac.std())
tracker.add('route.max_prob.', route_prob_max)
tracker.add("loss.", cross_entropy_loss)
tracker.add("lb_loss.", load_balancing_loss)
# Combined loss.
# The load balancing loss is multiplied by a coefficient $\alpha$ which is
# set to something small like $\alpha = 0.01$.
loss = cross_entropy_loss + self.load_balancing_loss_ceof * load_balancing_loss
# Calculate and log accuracy
self.accuracy(output, target)
self.accuracy.track()
# Train the model
if self.mode.is_train:
# Calculate gradients
loss.backward()
# Clip gradients
torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=self.grad_norm_clip)
# Take optimizer step
self.optimizer.step()
# Log the model parameters and gradients on last batch of every epoch
if batch_idx.is_last:
tracker.add('model', self.model)
# Clear the gradients
self.optimizer.zero_grad()
# Save the tracked metrics
tracker.save()
@option(Configs.model)
def autoregressive_model(c: Configs):
"""
### Initialize the auto-regressive model
"""
m = AutoregressiveModel(c.n_tokens, c.d_model, c.transformer)
return m.to(c.device)
@option(Configs.transformer)
def switch_transformer(c: Configs):
"""
### Initialize the switch transformer
"""
from labml_nn.transformers.switch import SwitchTransformer, SwitchTransformerLayer, SwitchFeedForward
from labml_nn.transformers import MultiHeadAttention
from labml_nn.transformers.feed_forward import FeedForward
return SwitchTransformer(
SwitchTransformerLayer(d_model=c.d_model,
attn=MultiHeadAttention(c.heads, c.d_model, c.dropout),
feed_forward=SwitchFeedForward(capacity_factor=c.capacity_factor,
drop_tokens=c.drop_tokens,
is_scale_prob=c.is_scale_prob,
n_experts=c.n_experts,
expert=FeedForward(c.d_model, c.d_ff, c.dropout),
d_model=c.d_model),
dropout_prob=c.dropout),
c.n_layers)
def main():
"""
### Run the experiment
"""
# Create experiment
experiment.create(name="switch_transformer", comment='')
# Create configs
conf = Configs()
# Load configurations
experiment.configs(conf,
# A dictionary of configurations to override
{'tokenizer': 'character',
'text': 'tiny_shakespeare',
'optimizer.learning_rate': 1.,
'optimizer.optimizer': 'Noam',
'prompt': 'It is',
'prompt_separator': '',
'transformer': 'switch_transformer',
'n_experts': 4,
'drop_tokens': True,
'capacity_factor': 1.2,
'train_loader': 'shuffled_train_loader',
'valid_loader': 'shuffled_valid_loader',
'seq_len': 64,
'epochs': 128,
'batch_size': 32,
'inner_iterations': 25,
})
# Set models for saving and loading
experiment.add_pytorch_models({'model': conf.model})
# Start the experiment
with experiment.start():
# `TrainValidConfigs.run`
conf.run()
#
if __name__ == '__main__':
main()
labml_nn/transformers/switch/readme.md
+27
View File
@@ -0,0 +1,27 @@
# [Switch Transformer](https://nn.labml.ai/transformers/switch/index.html)
This is a miniature [PyTorch](https://pytorch.org) implementation of the paper
[Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://arxiv.org/abs/2101.03961).
Our implementation only has a few million parameters and doesn't do model parallel distributed training.
It does single GPU training, but we implement the concept of switching as described in the paper.
The Switch Transformer uses different parameters for each token by switching among parameters
based on the token.
Therefore, only a fraction of parameters are chosen for each token.
So you can have more parameters but less computational cost.
The switching happens at the Position-wise Feedforward network (FFN) of each transformer block.
Position-wise feedforward network consists of two sequentially fully connected layers.
In switch transformer we have multiple FFNs (multiple experts),
and we chose which one to use based on a router.
The output is a set of probabilities for picking a FFN,
and we pick the one with the highest probability and only evaluate that.
So essentially the computational cost is the same as having a single FFN.
In our implementation this doesn't parallelize well when you have many or large FFNs since it's all
happening on a single GPU.
In a distributed setup you would have each FFN (each very large) on a different device.
The paper introduces another loss term to balance load among the experts (FFNs) and
discusses dropping tokens when routing is not balanced.
Here's [the training code](experiment.html) and a notebook for training a switch transformer on Tiny Shakespeare dataset.
labml_nn/transformers/utils.py
+27
View File
@@ -0,0 +1,27 @@
"""
---
title: Utilities for Transformer
summary: A bunch of utility functions and classes for transformers.
---
# Utilities for Transformer
"""
import torch
def subsequent_mask(seq_len):
"""
## Subsequent mask to mask out data from future (subsequent) time steps
"""
mask = torch.tril(torch.ones(seq_len, seq_len)).to(torch.bool).unsqueeze(-1)
return mask
def _subsequent_mask():
from labml.logger import inspect
inspect(subsequent_mask(10)[:, :, 0])
if __name__ == '__main__':
_subsequent_mask()
labml_nn/transformers/vit/__init__.py
+213
View File
@@ -0,0 +1,213 @@
"""
---
title: Vision Transformer (ViT)
summary: >
A PyTorch implementation/tutorial of the paper
"An Image Is Worth 16x16 Words: Transformers For Image Recognition At Scale"
---
# Vision Transformer (ViT)
This is a [PyTorch](https://pytorch.org) implementation of the paper
[An Image Is Worth 16x16 Words: Transformers For Image Recognition At Scale](https://arxiv.org/abs/2010.11929).
Vision transformer applies a pure transformer to images
without any convolution layers.
They split the image into patches and apply a transformer on patch embeddings.
[Patch embeddings](#PathEmbeddings) are generated by applying a simple linear transformation
to the flattened pixel values of the patch.
Then a standard transformer encoder is fed with the patch embeddings, along with a
classification token `[CLS]`.
The encoding on the `[CLS]` token is used to classify the image with an MLP.
When feeding the transformer with the patches, learned positional embeddings are
added to the patch embeddings, because the patch embeddings do not have any information
about where that patch is from.
The positional embeddings are a set of vectors for each patch location that get trained
with gradient descent along with other parameters.
ViTs perform well when they are pre-trained on large datasets.
The paper suggests pre-training them with an MLP classification head and
then using a single linear layer when fine-tuning.
The paper beats SOTA with a ViT pre-trained on a 300 million image dataset.
They also use higher resolution images during inference while keeping the
patch size the same.
The positional embeddings for new patch locations are calculated by interpolating
learning positional embeddings.
Here's [an experiment](experiment.html) that trains ViT on CIFAR-10.
This doesn't do very well because it's trained on a small dataset.
It's a simple experiment that anyone can run and play with ViTs.
"""
import torch
from torch import nn
from labml_nn.transformers import TransformerLayer
from labml_nn.utils import clone_module_list
class PatchEmbeddings(nn.Module):
"""
<a id="PatchEmbeddings"></a>
## Get patch embeddings
The paper splits the image into patches of equal size and do a linear transformation
on the flattened pixels for each patch.
We implement the same thing through a convolution layer, because it's simpler to implement.
"""
def __init__(self, d_model: int, patch_size: int, in_channels: int):
"""
* `d_model` is the transformer embeddings size
* `patch_size` is the size of the patch
* `in_channels` is the number of channels in the input image (3 for rgb)
"""
super().__init__()
# We create a convolution layer with a kernel size and and stride length equal to patch size.
# This is equivalent to splitting the image into patches and doing a linear
# transformation on each patch.
self.conv = nn.Conv2d(in_channels, d_model, patch_size, stride=patch_size)
def forward(self, x: torch.Tensor):
"""
* `x` is the input image of shape `[batch_size, channels, height, width]`
"""
# Apply convolution layer
x = self.conv(x)
# Get the shape.
bs, c, h, w = x.shape
# Rearrange to shape `[patches, batch_size, d_model]`
x = x.permute(2, 3, 0, 1)
x = x.view(h * w, bs, c)
# Return the patch embeddings
return x
class LearnedPositionalEmbeddings(nn.Module):
"""
<a id="LearnedPositionalEmbeddings"></a>
## Add parameterized positional encodings
This adds learned positional embeddings to patch embeddings.
"""
def __init__(self, d_model: int, max_len: int = 5_000):
"""
* `d_model` is the transformer embeddings size
* `max_len` is the maximum number of patches
"""
super().__init__()
# Positional embeddings for each location
self.positional_encodings = nn.Parameter(torch.zeros(max_len, 1, d_model), requires_grad=True)
def forward(self, x: torch.Tensor):
"""
* `x` is the patch embeddings of shape `[patches, batch_size, d_model]`
"""
# Get the positional embeddings for the given patches
pe = self.positional_encodings[:x.shape[0]]
# Add to patch embeddings and return
return x + pe
class ClassificationHead(nn.Module):
"""
<a id="ClassificationHead"></a>
## MLP Classification Head
This is the two layer MLP head to classify the image based on `[CLS]` token embedding.
"""
def __init__(self, d_model: int, n_hidden: int, n_classes: int):
"""
* `d_model` is the transformer embedding size
* `n_hidden` is the size of the hidden layer
* `n_classes` is the number of classes in the classification task
"""
super().__init__()
# First layer
self.linear1 = nn.Linear(d_model, n_hidden)
# Activation
self.act = nn.ReLU()
# Second layer
self.linear2 = nn.Linear(n_hidden, n_classes)
def forward(self, x: torch.Tensor):
"""
* `x` is the transformer encoding for `[CLS]` token
"""
# First layer and activation
x = self.act(self.linear1(x))
# Second layer
x = self.linear2(x)
#
return x
class VisionTransformer(nn.Module):
"""
## Vision Transformer
This combines the [patch embeddings](#PatchEmbeddings),
[positional embeddings](#LearnedPositionalEmbeddings),
transformer and the [classification head](#ClassificationHead).
"""
def __init__(self, transformer_layer: TransformerLayer, n_layers: int,
patch_emb: PatchEmbeddings, pos_emb: LearnedPositionalEmbeddings,
classification: ClassificationHead):
"""
* `transformer_layer` is a copy of a single [transformer layer](../models.html#TransformerLayer).
We make copies of it to make the transformer with `n_layers`.
* `n_layers` is the number of [transformer layers](../models.html#TransformerLayer).
* `patch_emb` is the [patch embeddings layer](#PatchEmbeddings).
* `pos_emb` is the [positional embeddings layer](#LearnedPositionalEmbeddings).
* `classification` is the [classification head](#ClassificationHead).
"""
super().__init__()
# Patch embeddings
self.patch_emb = patch_emb
self.pos_emb = pos_emb
# Classification head
self.classification = classification
# Make copies of the transformer layer
self.transformer_layers = clone_module_list(transformer_layer, n_layers)
# `[CLS]` token embedding
self.cls_token_emb = nn.Parameter(torch.randn(1, 1, transformer_layer.size), requires_grad=True)
# Final normalization layer
self.ln = nn.LayerNorm([transformer_layer.size])
def forward(self, x: torch.Tensor):
"""
* `x` is the input image of shape `[batch_size, channels, height, width]`
"""
# Get patch embeddings. This gives a tensor of shape `[patches, batch_size, d_model]`
x = self.patch_emb(x)
# Concatenate the `[CLS]` token embeddings before feeding the transformer
cls_token_emb = self.cls_token_emb.expand(-1, x.shape[1], -1)
x = torch.cat([cls_token_emb, x])
# Add positional embeddings
x = self.pos_emb(x)
# Pass through transformer layers with no attention masking
for layer in self.transformer_layers:
x = layer(x=x, mask=None)
# Get the transformer output of the `[CLS]` token (which is the first in the sequence).
x = x[0]
# Layer normalization
x = self.ln(x)
# Classification head, to get logits
x = self.classification(x)
#
return x
labml_nn/transformers/vit/experiment.py
+94
View File
@@ -0,0 +1,94 @@
"""
---
title: Train a Vision Transformer (ViT) on CIFAR 10
summary: >
Train a Vision Transformer (ViT) on CIFAR 10
---
# Train a [Vision Transformer (ViT)](index.html) on CIFAR 10
"""
from labml import experiment
from labml.configs import option
from labml_nn.experiments.cifar10 import CIFAR10Configs
from labml_nn.transformers import TransformerConfigs
class Configs(CIFAR10Configs):
"""
## Configurations
We use [`CIFAR10Configs`](../../experiments/cifar10.html) which defines all the
dataset related configurations, optimizer, and a training loop.
"""
# [Transformer configurations](../configs.html#TransformerConfigs)
# to get [transformer layer](../models.html#TransformerLayer)
transformer: TransformerConfigs
# Size of a patch
patch_size: int = 4
# Size of the hidden layer in classification head
n_hidden_classification: int = 2048
# Number of classes in the task
n_classes: int = 10
@option(Configs.transformer)
def _transformer():
"""
Create transformer configs
"""
return TransformerConfigs()
@option(Configs.model)
def _vit(c: Configs):
"""
### Create model
"""
from labml_nn.transformers.vit import VisionTransformer, LearnedPositionalEmbeddings, ClassificationHead, \
PatchEmbeddings
# Transformer size from [Transformer configurations](../configs.html#TransformerConfigs)
d_model = c.transformer.d_model
# Create a vision transformer
return VisionTransformer(c.transformer.encoder_layer, c.transformer.n_layers,
PatchEmbeddings(d_model, c.patch_size, 3),
LearnedPositionalEmbeddings(d_model),
ClassificationHead(d_model, c.n_hidden_classification, c.n_classes)).to(c.device)
def main():
# Create experiment
experiment.create(name='ViT', comment='cifar10')
# Create configurations
conf = Configs()
# Load configurations
experiment.configs(conf, {
# Optimizer
'optimizer.optimizer': 'Adam',
'optimizer.learning_rate': 2.5e-4,
# Transformer embedding size
'transformer.d_model': 512,
# Training epochs and batch size
'epochs': 32,
'train_batch_size': 64,
# Augment CIFAR 10 images for training
'train_dataset': 'cifar10_train_augmented',
# Do not augment CIFAR 10 images for validation
'valid_dataset': 'cifar10_valid_no_augment',
})
# Set model for saving/loading
experiment.add_pytorch_models({'model': conf.model})
# Start the experiment and run the training loop
with experiment.start():
conf.run()
#
if __name__ == '__main__':
main()
labml_nn/transformers/vit/readme.md
+32
View File
@@ -0,0 +1,32 @@
# [Vision Transformer (ViT)](https://nn.labml.ai/transformer/vit/index.html)
This is a [PyTorch](https://pytorch.org) implementation of the paper
[An Image Is Worth 16x16 Words: Transformers For Image Recognition At Scale](https://arxiv.org/abs/2010.11929).
Vision transformer applies a pure transformer to images
without any convolution layers.
They split the image into patches and apply a transformer on patch embeddings.
[Patch embeddings](https://nn.labml.ai/transformer/vit/index.html#PathEmbeddings) are generated by applying a simple linear transformation
to the flattened pixel values of the patch.
Then a standard transformer encoder is fed with the patch embeddings, along with a
classification token `[CLS]`.
The encoding on the `[CLS]` token is used to classify the image with an MLP.
When feeding the transformer with the patches, learned positional embeddings are
added to the patch embeddings, because the patch embeddings do not have any information
about where that patch is from.
The positional embeddings are a set of vectors for each patch location that get trained
with gradient descent along with other parameters.
ViTs perform well when they are pre-trained on large datasets.
The paper suggests pre-training them with an MLP classification head and
then using a single linear layer when fine-tuning.
The paper beats SOTA with a ViT pre-trained on a 300 million image dataset.
They also use higher resolution images during inference while keeping the
patch size the same.
The positional embeddings for new patch locations are calculated by interpolating
learning positional embeddings.
Here's [an experiment](https://nn.labml.ai/transformer/vit/experiment.html) that trains ViT on CIFAR-10.
This doesn't do very well because it's trained on a small dataset.
It's a simple experiment that anyone can run and play with ViTs.
labml_nn/transformers/xl/__init__.py
+140
View File
@@ -0,0 +1,140 @@
"""
---
title: Transformer XL
summary: >
Documented implementation with explanations of a
Transformer-XL model.
---
# Transformer XL
This is an implementation of
[Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context](https://arxiv.org/abs/1901.02860)
in [PyTorch](https://pytorch.org).
Transformer has a limited attention span,
equal to the length of the sequence trained in parallel.
All these positions have a fixed positional encoding.
Transformer XL increases this attention span by letting
each of the positions pay attention to precalculated past embeddings.
For instance if the context length is $l$, it will keep the embeddings of
all layers for previous batch of length $l$ and feed them to current step.
If we use fixed-positional encodings these pre-calculated embeddings will have
the same positions as the current context.
They introduce relative positional encoding, where the positional encodings
are introduced at the attention calculation.
Annotated implementation of relative multi-headed attention is in [`relative_mha.py`](relative_mha.html).
Here's [the training code](experiment.html) and a notebook for training a transformer XL model on Tiny Shakespeare dataset.
[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/transformers/xl/experiment.ipynb)
"""
from typing import List, Optional
import torch
import torch.nn as nn
from labml_nn.utils import clone_module_list
from .relative_mha import RelativeMultiHeadAttention
from ..feed_forward import FeedForward
class TransformerXLLayer(nn.Module):
"""
## Transformer XL Layer
The transformer XL model comprises of a number of these layers.
"""
def __init__(self, *,
d_model: int,
self_attn: RelativeMultiHeadAttention,
feed_forward: FeedForward,
dropout_prob: float):
"""
* `d_model` is the token embedding size
* `self_attn` is the [self attention module](relative_mha.html)
* `feed_forward` is the feed forward module
* `dropout_prob` is the probability of dropping out after self attention and FFN
"""
super().__init__()
self.size = d_model
self.self_attn = self_attn
self.feed_forward = feed_forward
self.dropout = nn.Dropout(dropout_prob)
self.norm_self_attn = nn.LayerNorm([d_model])
self.norm_ff = nn.LayerNorm([d_model])
def forward(self, *,
x: torch.Tensor,
mem: Optional[torch.Tensor],
mask: torch.Tensor):
"""
* `x` is a tensor of the token level feature vectors of shape `[seq_len, batch_size, d_model]`
* `mem` is a tensor of the past token level feature vectors of shape `[mem_len, batch_size, d_model]`
* `mask` is a matrix of shape `[seq_len, mem_len + seq_len, batch_size]` or `[seq_len, mem_len + seq_len, 1]`.
`mask[i, j]` is true if token at `i` can see token at `j`.
"""
# Normalize the vectors before doing self attention
z = self.norm_self_attn(x)
# If there is memory
if mem is not None:
# Normalize it
mem = self.norm_self_attn(mem)
# Concatenate with `z`
m_z = torch.cat((mem, z), dim=0)
# Ignore if there is no memory
else:
m_z = z
# Attention
self_attn = self.self_attn(query=z, key=m_z, value=m_z, mask=mask)
# Add the attention results
x = x + self.dropout(self_attn)
# Normalize for feed-forward
z = self.norm_ff(x)
# Pass through the feed-forward network
ff = self.feed_forward(z)
# Add the feed-forward results back
x = x + self.dropout(ff)
#
return x
class TransformerXL(nn.Module):
"""
## Transformer XL Model
This consists of multiple transformer XL layers
"""
def __init__(self, layer: TransformerXLLayer, n_layers: int):
super().__init__()
# Make copies of the transformer layer
self.layers = clone_module_list(layer, n_layers)
# Final normalization layer
self.norm = nn.LayerNorm([layer.size])
def forward(self, x: torch.Tensor, mem: List[torch.Tensor], mask: torch.Tensor):
"""
* `x` is a tensor of the token embeddings vectors of shape `[seq_len, batch_size, d_model]`
* `mem` is a list of tensors of the past token level feature vectors of shape
`[mem_len, batch_size, d_model]` for each layer
* `mask` is the masking matrix
"""
# List to store token level feature vectors,
# which will become the memories for the next sequential batch.
new_mem = []
# Run through each transformer layer
for i, layer in enumerate(self.layers):
# Add to the list of feature vectors
new_mem.append(x.detach())
# Memory
m = mem[i] if mem else None
# Run through the transformer XL layer
x = layer(x=x, mem=m, mask=mask)
# Finally, normalize the vectors
return self.norm(x), new_mem
labml_nn/transformers/xl/experiment.ipynb
+222
View File
@@ -0,0 +1,222 @@
{
"nbformat": 4,
"nbformat_minor": 0,
"metadata": {
"colab": {
"name": "Transformer XL",
"provenance": [],
"collapsed_sections": []
},
"kernelspec": {
"name": "python3",
"display_name": "Python 3"
},
"accelerator": "GPU"
},
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "AYV_dMVDxyc2"
},
"source": [
"[![Github](https://img.shields.io/github/stars/labmlai/annotated_deep_learning_paper_implementations?style=social)](https://github.com/labmlai/annotated_deep_learning_paper_implementations)\n",
"[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/transformers/xl/experiment.ipynb) \n",
"\n",
"## Transformer XL\n",
"\n",
"This is an experiment training Shakespeare dataset with a Transformer XL model."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "AahG_i2y5tY9"
},
"source": [
"Install the `labml-nn` package"
]
},
{
"cell_type": "code",
"metadata": {
"id": "ZCzmCrAIVg0L",
"colab": {
"base_uri": "https://localhost:8080/"
},
"outputId": "1a4a59ce-b300-4d9f-baee-15720d696773"
},
"source": [
"!pip install labml-nn"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "SE2VUQ6L5zxI"
},
"source": [
"Imports"
]
},
{
"cell_type": "code",
"metadata": {
"id": "0hJXx_g0wS2C"
},
"source": [
"from labml import experiment\n",
"from labml_nn.transformers.xl.experiment import Configs"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "Lpggo0wM6qb-"
},
"source": [
"Create an experiment"
]
},
{
"cell_type": "code",
"metadata": {
"id": "bFcr9k-l4cAg"
},
"source": [
"experiment.create(name=\"transformer_xl\")"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "-OnHLi626tJt"
},
"source": [
"Initialize configurations"
]
},
{
"cell_type": "code",
"metadata": {
"id": "Piz0c5f44hRo"
},
"source": [
"conf = Configs()"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "wwMzCqpD6vkL"
},
"source": [
"Set experiment configurations and assign a configurations dictionary to override configurations"
]
},
{
"cell_type": "code",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 17
},
"id": "e6hmQhTw4nks",
"outputId": "839820be-f5a9-476d-d458-7d2ff3278e89"
},
"source": [
"experiment.configs(conf,\n",
" # A dictionary of configurations to override\n",
" {'tokenizer': 'character',\n",
" 'text': 'tiny_shakespeare',\n",
" 'optimizer.learning_rate': 1.,\n",
" 'optimizer.optimizer': 'Noam',\n",
" 'prompt': 'It is',\n",
" 'prompt_separator': '',\n",
"\n",
" 'train_loader': 'sequential_train_loader',\n",
" 'valid_loader': 'sequential_valid_loader',\n",
"\n",
" 'seq_len': 2,\n",
" 'mem_len': 32,\n",
" 'epochs': 128,\n",
" 'batch_size': 32,\n",
" 'inner_iterations': 25,\n",
" })"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "EvI7MtgJ61w5"
},
"source": [
"Set PyTorch models for loading and saving"
]
},
{
"cell_type": "code",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 255
},
"id": "GDlt7dp-5ALt",
"outputId": "0543a726-fcf4-4493-dbe9-ab62c5bea94f"
},
"source": [
"experiment.add_pytorch_models({'model': conf.model})"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "KJZRf8527GxL"
},
"source": [
"Start the experiment and run the training loop."
]
},
{
"cell_type": "code",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 493
},
"id": "aIAWo7Fw5DR8",
"outputId": "64764a3f-e8d7-4c06-b803-0da6ff756163"
},
"source": [
"# Start the experiment\n",
"with experiment.start():\n",
" conf.run()"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "code",
"metadata": {
"id": "oBXXlP2b7XZO"
},
"source": [
""
],
"outputs": [],
"execution_count": null
}
]
}
labml_nn/transformers/xl/experiment.py
+257
View File
@@ -0,0 +1,257 @@
"""
---
title: Transformer XL Experiment
summary: This experiment trains a transformer XL model on tiny Shakespeare dataset.
---
# Transformer XL Experiment
This is an annotated PyTorch experiment to train a transformer xl model.
"""
from typing import List
import torch
import torch.nn as nn
from labml import experiment, tracker, monit, logger
from labml.configs import option
from labml.logger import Text
from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs
from labml_nn.helpers.metrics import SimpleStateModule
from labml_nn.helpers.trainer import BatchIndex
from labml_nn.transformers.xl import TransformerXL, TransformerXLLayer
class AutoregressiveModel(nn.Module):
"""
## Auto regressive model
"""
def __init__(self, n_vocab: int, d_model: int, transformer: TransformerXL):
super().__init__()
# Token embedding module
self.src_embed = nn.Embedding(n_vocab, d_model)
# Transformer
self.transformer = transformer
# Final layer
self.generator = nn.Linear(d_model, n_vocab)
# Masks
self.mask_x = None
self.mask_mem = None
def forward(self, x: torch.Tensor, mem: List[torch.Tensor]):
# Length of the memory
m_len = len(mem[0]) if mem else 0
# Create a subsequent mask for tokens
if self.mask_x is None or self.mask_x.shape[0] < len(x):
from labml_nn.transformers.utils import subsequent_mask
self.mask_x = subsequent_mask(len(x)).to(x.device)
# Create an all ones (full visibility) mask for memory
if self.mask_mem is None or self.mask_mem.shape[1] < m_len or self.mask_mem.shape[0] < len(x):
self.mask_mem = self.mask_x.new_ones(len(x), m_len, 1)
# Concatenate the masks if there is memory
if m_len:
mask = torch.cat((self.mask_mem[:len(x), :m_len], self.mask_x[:len(x), :len(x)]), dim=1)
# Use the subsequent mask otherwise
else:
mask = self.mask_x[:len(x), :len(x)]
# Token embeddings
x = self.src_embed(x)
# Run it through the transformer
res, mem = self.transformer(x, mem, mask)
# Generate logits of the next token
res = self.generator(res)
#
return res, mem
class Configs(NLPAutoRegressionConfigs):
"""
## Configurations
The default configs can and will be over-ridden when we start the experiment
"""
model: AutoregressiveModel
# Token embedding size
d_model: int = 128
# Number of attention heads
heads: int = 4
# Dropout probability
dropout: float = 0.0
# Number of features in FFN hidden layer
d_ff: int = 256
# Number of transformer layers
n_layers: int = 6
# Number of memories to keep
mem_len: int = 128
# State module to maintain memories when switching between training and validation
memory = SimpleStateModule()
def init(self):
# Set tracker configurations
tracker.set_scalar("accuracy.*", True)
tracker.set_scalar("loss.*", True)
# This will keep the accuracy metric stats and memories separate for training and validation.
self.state_modules = [self.accuracy, self.memory]
def merge_memory(self, old_mem, new_mem):
"""
Concatenate memories and remove old memories to keep a maximum of
`mem_len` memories.
"""
# If it's configured not to use memory
if self.mem_len == 0:
return []
# Concatenate with old memory
if old_mem:
mem = [torch.cat((m, x), dim=0) for m, x in zip(old_mem, new_mem)]
else:
mem = new_mem
# Truncate old memories
if len(mem[0]) > self.mem_len:
mem = [m[-self.mem_len:] for m in mem]
#
return mem
def step(self, batch: any, batch_idx: BatchIndex):
"""
### Training/validation step
"""
# Move data to the device
data, target = batch[0].to(self.device), batch[1].to(self.device)
# Update global step (number of tokens processed) when in training mode
if self.mode.is_train:
tracker.add_global_step(data.shape[0] * data.shape[1])
# Get memories
mem = self.memory.get()
# Run the model
output, new_mem = self.model(data, mem)
# Merge memory
mem = self.merge_memory(mem, new_mem)
# Update memories
self.memory.set(mem)
# Calculate and log cross entropy loss
loss = self.loss_func(output, target)
tracker.add("loss.", loss)
# Calculate and log accuracy
self.accuracy(output, target)
self.accuracy.track()
# Train the model
if self.mode.is_train:
# Calculate gradients
loss.backward()
# Clip gradients
torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=self.grad_norm_clip)
# Take optimizer step
self.optimizer.step()
# Log the model parameters and gradients on last batch of every epoch
if batch_idx.is_last:
tracker.add('model', self.model)
# Clear the gradients
self.optimizer.zero_grad()
# Save the tracked metrics
tracker.save()
def sample(self):
"""
### Sampling function to generate samples periodically while training
"""
# Starting prompt
prompt = self.prompt
# Collect output for printing
log = [(prompt, Text.subtle)]
# memory
mem = []
# Sample 25 tokens
for i in monit.iterate('Sample', 25):
# Tokenize the prompt
data = self.text.text_to_i(prompt).unsqueeze(-1)
# Move to device
data = data.to(self.device)
# Get the model output
output, new_mem = self.model(data, mem)
# Get the model prediction (greedy)
output = output.argmax(dim=-1).squeeze(1)
# Add the prediction to prompt
prompt += self.prompt_separator + self.text.itos[output[-1]]
# Only feed the last character to model in next iteration, rest will go in as memories
prompt = prompt[-1:]
# Add the prediction for logging
log += [(self.prompt_separator + self.text.itos[output[-1]], Text.value)]
# Update memory
mem = self.merge_memory(mem, new_mem)
# Print the sampled output
logger.log(log)
@option(Configs.model)
def autoregressive_model(c: Configs):
"""
### Initialize the auto-regressive model
"""
from labml_nn.transformers.xl import RelativeMultiHeadAttention
from labml_nn.transformers.feed_forward import FeedForward
m = AutoregressiveModel(c.n_tokens, c.d_model, TransformerXL(
TransformerXLLayer(d_model=c.d_model,
self_attn=RelativeMultiHeadAttention(c.heads, c.d_model, c.dropout),
feed_forward=FeedForward(c.d_model, c.d_ff, c.dropout),
dropout_prob=c.dropout), c.n_layers))
return m.to(c.device)
def main():
"""
### Run the experiment
"""
# Create experiment
experiment.create(name="transformer_xl", comment='')
# Create configs
conf = Configs()
# Load configurations
experiment.configs(conf,
# A dictionary of configurations to override
{'tokenizer': 'character',
'text': 'tiny_shakespeare',
'optimizer.learning_rate': 1.,
'optimizer.optimizer': 'Noam',
'prompt': 'It is',
'prompt_separator': '',
'train_loader': 'sequential_train_loader',
'valid_loader': 'sequential_valid_loader',
'seq_len': 2,
'mem_len': 32,
'epochs': 128,
'batch_size': 32,
'inner_iterations': 25,
})
# Set models for saving and loading
experiment.add_pytorch_models({'model': conf.model})
# Start the experiment
with experiment.start():
# `TrainValidConfigs.run`
conf.run()
#
if __name__ == '__main__':
main()
labml_nn/transformers/xl/readme.md
+23
View File
@@ -0,0 +1,23 @@
# [Transformer XL](https://nn.labml.ai/transformers/xl/index.html)
This is an implementation of
[Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context](https://arxiv.org/abs/1901.02860)
in [PyTorch](https://pytorch.org).
Transformer has a limited attention span,
equal to the length of the sequence trained in parallel.
All these positions have a fixed positional encoding.
Transformer XL increases this attention span by letting
each of the positions pay attention to precalculated past embeddings.
For instance if the context length is $l$, it will keep the embeddings of
all layers for previous batch of length $l$ and feed them to current step.
If we use fixed-positional encodings these pre-calculated embeddings will have
the same positions as the current context.
They introduce relative positional encoding, where the positional encodings
are introduced at the attention calculation.
Annotated implementation of relative multi-headed attention is in [`relative_mha.py`](https://nn.labml.ai/transformers/xl/relative_mha.html).
Here's [the training code](https://nn.labml.ai/transformers/xl/experiment.html) and a notebook for training a transformer XL model on Tiny Shakespeare dataset.
[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/transformers/xl/experiment.ipynb)
labml_nn/transformers/xl/relative_mha.py
+152
View File
@@ -0,0 +1,152 @@
"""
---
title: Relative Multi-Headed Attention
summary: >
Documented implementation with explanations of
Relative Multi-Headed Attention from paper Transformer-XL.
---
# Relative Multi-Headed Attention
This is an implementation of relative multi-headed attention from paper
[Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context](https://arxiv.org/abs/1901.02860)
in [PyTorch](https://pytorch.org).
"""
import torch
from torch import nn
from labml.logger import inspect
from labml_nn.transformers.mha import MultiHeadAttention
def shift_right(x: torch.Tensor):
"""
This method shifts $i^{th}$ row of a matrix by $i$ columns.
If the input is `[[1, 2 ,3], [4, 5 ,6], [7, 8, 9]]`, the shifted
result would be `[[1, 2 ,3], [0, 4, 5], [6, 0, 7]]`.
*Ideally we should mask out the lower triangle but it's ok for our purpose*.
"""
# Concatenate a column of zeros
zero_pad = x.new_zeros(x.shape[0], 1, *x.shape[2:])
x_padded = torch.cat([x, zero_pad], dim=1)
# Reshape and remove excess elements from the end
x_padded = x_padded.view(x.shape[1] + 1, x.shape[0], *x.shape[2:])
x = x_padded[:-1].view_as(x)
#
return x
class RelativeMultiHeadAttention(MultiHeadAttention):
"""
## Relative Multi-Head Attention Module
We override [Multi-Head Attention](mha.html) module so we only need to
write the `get_scores` method.
"""
def __init__(self, heads: int, d_model: int, dropout_prob: float = 0.1):
# The linear transformations do not need a bias since we
# explicitly include it when calculating scores.
# However having a bias for `value` might make sense.
super().__init__(heads, d_model, dropout_prob, bias=False)
# Number of relative positions
self.P = 2 ** 12
# Relative positional embeddings for key relative to the query.
# We need $2P$ embeddings because the keys can be before or after the query.
self.key_pos_embeddings = nn.Parameter(torch.zeros((self.P * 2, heads, self.d_k)), requires_grad=True)
# Relative positional embedding bias for key relative to the query.
self.key_pos_bias = nn.Parameter(torch.zeros((self.P * 2, heads)), requires_grad=True)
# Positional embeddings for the query is independent of the position of the query
self.query_pos_bias = nn.Parameter(torch.zeros((heads, self.d_k)), requires_grad=True)
def get_scores(self, query: torch.Tensor, key: torch.Tensor):
r"""
### Get relative attention scores
With absolute attention
\begin{align}
A^{abs}_{j} &= lin_q(X^q_i + P_i)^\top lin_k(X^k_j + P_j) \\
&= \underset{\textcolor{lightgreen}{A}}{Q_i^\top K_j} +
\underset{\textcolor{lightgreen}{B}}{Q_i^\top U^K_j} +
\underset{\textcolor{lightgreen}{C}}{{U^Q_i}^\top K_j} +
\underset{\textcolor{lightgreen}{D}}{{U^Q_i}^\top U^K_j}
\end{align}
where $Q_i, K_j$, are linear transformations of
original embeddings $X^q_i, X^k_j$
and $U^Q_i, U^K_j$ are linear transformations of
absolute positional encodings $P_i, P_j$.
They reason out that the attention to a given key should be the same regardless of
the position of query.
Hence replace $\underset{\textcolor{lightgreen}{C}}{{U^Q_i}^\top K_j}$
with a constant $\underset{\textcolor{lightgreen}{C}}{\textcolor{orange}{v^\top} K_j}$.
For the second and third terms relative positional encodings are introduced.
So $\underset{\textcolor{lightgreen}{B}}{Q_i^\top U^K_j}$ is
replaced with $\underset{\textcolor{lightgreen}{B}}{Q_i^\top \textcolor{orange}{R_{i - j}}}$
and $\underset{\textcolor{lightgreen}{D}}{{U^Q_i}^\top U^K_j}$
with $\underset{\textcolor{lightgreen}{D}}{\textcolor{orange}{S_{i-j}}}$.
\begin{align}
A^{rel}_{i,j} &= \underset{\mathbf{\textcolor{lightgreen}{A}}}{Q_i^\top K_j} +
\underset{\mathbf{\textcolor{lightgreen}{B}}}{Q_i^\top \textcolor{orange}{R_{i - j}}} +
\underset{\mathbf{\textcolor{lightgreen}{C}}}{\textcolor{orange}{v^\top} K_j} +
\underset{\mathbf{\textcolor{lightgreen}{D}}}{\textcolor{orange}{S_{i-j}}}
\end{align}
"""
# $\textcolor{orange}{R_k}$
key_pos_emb = self.key_pos_embeddings[self.P - key.shape[0]:self.P + query.shape[0]]
# $\textcolor{orange}{S_k}$
key_pos_bias = self.key_pos_bias[self.P - key.shape[0]:self.P + query.shape[0]]
# $\textcolor{orange}{v^\top}$
query_pos_bias = self.query_pos_bias[None, None, :, :]
# ${(\textcolor{lightgreen}{\mathbf{A + C}})}_{i,j} =
# Q_i^\top K_j +
# \textcolor{orange}{v^\top} K_j$
ac = torch.einsum('ibhd,jbhd->ijbh', query + query_pos_bias, key)
# $\textcolor{lightgreen}{\mathbf{B'}_{i,k}} = Q_i^\top \textcolor{orange}{R_k}$
b = torch.einsum('ibhd,jhd->ijbh', query, key_pos_emb)
# $\textcolor{lightgreen}{\mathbf{D'}_{i,k}} = \textcolor{orange}{S_k}$
d = key_pos_bias[None, :, None, :]
# Shift the rows of $\textcolor{lightgreen}{\mathbf{(B' + D')}_{i,k}}$
# to get $$\textcolor{lightgreen}{\mathbf{(B + D)}_{i,j} = \mathbf{(B' + D')}_{i,i - j}}$$
bd = shift_right(b + d)
# Remove extra positions
bd = bd[:, -key.shape[0]:]
# Return the sum $$
# \underset{\mathbf{\textcolor{lightgreen}{A}}}{Q_i^\top K_j} +
# \underset{\mathbf{\textcolor{lightgreen}{B}}}{Q_i^\top \textcolor{orange}{R_{i - j}}} +
# \underset{\mathbf{\textcolor{lightgreen}{C}}}{\textcolor{orange}{v^\top} K_j} +
# \underset{\mathbf{\textcolor{lightgreen}{D}}}{\textcolor{orange}{S_{i-j}}}
# $$
return ac + bd
def _test_shift_right():
x = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
inspect(x)
inspect(shift_right(x))
x = torch.arange(1, 6)[None, :, None, None].repeat(5, 1, 1, 1)
inspect(x[:, :, 0, 0])
inspect(shift_right(x)[:, :, 0, 0])
x = torch.arange(1, 6)[None, :, None, None].repeat(3, 1, 1, 1)
inspect(x[:, :, 0, 0])
inspect(shift_right(x)[:, :, 0, 0])
if __name__ == '__main__':
_test_shift_right()