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