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/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.