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