chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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()
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user