| |
| |
| |
| |
| |
| |
|
|
| import math |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from einops import rearrange |
|
|
|
|
| def fft_conv(u, k, dropout_mask, gelu=True, k_rev=None): |
| seqlen = u.shape[-1] |
| fft_size = 2 * seqlen |
| k_f = torch.fft.rfft(k, n=fft_size) / fft_size |
| if k_rev is not None: |
| k_rev_f = torch.fft.rfft(k_rev, n=fft_size) / fft_size |
| k_f = k_f + k_rev_f.conj() |
| u_f = torch.fft.rfft(u.to(dtype=k.dtype), n=fft_size) |
|
|
| if len(u.shape) > 3: |
| k_f = k_f.unsqueeze(1) |
| y = torch.fft.irfft(u_f * k_f, n=fft_size, norm="forward")[..., :seqlen] |
|
|
| out = y + u |
| if gelu: |
| out = F.gelu(out) |
| if dropout_mask is not None: |
| return (out * rearrange(dropout_mask, "b H -> b H 1")).to(dtype=u.dtype) |
| else: |
| return out.to(dtype=u.dtype) |
|
|
|
|
| class LongConvolution(nn.Module): |
| """ |
| LongConvolution applies a convolution operation on the input tensor using a fixed |
| filter of length max_len. |
| The filter is learned during training and is applied using FFT convolution. |
| |
| Args: |
| hidden_size (int): The number of expected features in the input and output. |
| max_len (int): The maximum sequence length. |
| |
| Returns: |
| y: [batch_size, seq_len, hidden_size] tensor |
| """ |
|
|
| def __init__( |
| self, |
| hidden_size: int, |
| max_len: int, |
| **kwargs, |
| ): |
| """ |
| Initializes the LongConvolution module. |
| Args: |
| hidden_size (int): The number of expected features in the input and output. |
| max_len (int): The maximum sequence length. |
| """ |
| super().__init__() |
| self.hidden_size = hidden_size |
| self.filter = nn.Parameter(torch.randn(self.hidden_size, max_len), requires_grad=True) |
|
|
| def forward(self, x: torch.Tensor, *args, **kwargs): |
| """ |
| Applies the LongConvolution operation on the input tensor. |
| Args: |
| x: [batch_size, seq_len, hidden_size] tensor |
| Returns: |
| y: [batch_size, seq_len, hidden_size] tensor |
| """ |
| x = x.transpose(1, 2) |
| y = fft_conv(x, self.filter, dropout_mask=None, gelu=False) |
| y = y.transpose(1, 2) |
| return y.to(dtype=x.dtype) |
|
|
|
|
| class PositionalEmbedding(nn.Module): |
| def __init__(self, emb_dim: int, seq_len: int, **kwargs): |
| """Complex exponential positional embeddings for implicit long convolution filters.""" |
| super().__init__() |
|
|
| self.seq_len = seq_len |
| |
| t = torch.linspace(0, 1, self.seq_len)[None, :, None] |
|
|
| if emb_dim > 1: |
| bands = (emb_dim - 1) // 2 |
| |
| t_rescaled = torch.linspace(0, seq_len - 1, seq_len)[None, :, None] |
| w = 2 * math.pi * t_rescaled / seq_len |
|
|
| f = torch.linspace(1e-4, bands - 1, bands)[None, None] |
| z = torch.exp(-1j * f * w) |
| z = torch.cat([t, z.real, z.imag], dim=-1) |
| self.z = nn.Parameter(z, requires_grad=False) |
|
|
| def forward(self, L): |
| return self.z[:, :L] |
|
|
|
|
| class ImplicitLongConvolution(nn.Module): |
| """ |
| Long convolution with implicit filter parameterized by an MLP. |
| |
| Args: |
| hidden_size (int): |
| The number of expected features in the input and output. |
| max_len (int): |
| The maximum sequence length. |
| d_emb (Optional[int]): |
| The dimension of the positional embeddings. Must be odd and greater or equal to 3 (time, sine and cosine). |
| Defaults to 3. |
| d_hidden (Optional[int]): |
| The number of features in the hidden layer of the MLP. Defaults to 16. |
| |
| Attributes: |
| pos_emb (`PositionalEmbedding`): The positional embedding layer. |
| mlp (`nn.Sequential`): The MLP that parameterizes the implicit filter. |
| |
| """ |
|
|
| def __init__( |
| self, |
| hidden_size: int, |
| max_len: int, |
| d_emb: int = 3, |
| d_hidden: int = 16, |
| **kwargs, |
| ): |
| """ |
| Long convolution with implicit filter parameterized by an MLP. |
| |
| |
| """ |
| super().__init__() |
| self.hidden_size = hidden_size |
| self.d_emb = d_emb |
|
|
| assert ( |
| d_emb % 2 != 0 and d_emb >= 3 |
| ), "d_emb must be odd and greater or equal to 3 (time, sine and cosine)" |
| self.pos_emb = PositionalEmbedding(d_emb, max_len) |
|
|
| |
| self.mlp = nn.Sequential( |
| nn.Linear(d_emb, d_hidden), |
| torch.nn.ReLU(), |
| nn.Linear(d_hidden, hidden_size), |
| ) |
|
|
| def filter(self, seq_len: int, *args, **kwargs): |
| return self.mlp(self.pos_emb(seq_len)).transpose(1, 2) |
|
|
| def forward(self, x: torch.Tensor, *args, **kwargs): |
| """ |
| Args: |
| x: [batch_size, seq_len, hidden_size] tensor |
| |
| Returns: |
| y: [batch_size, seq_len, hidden_size] tensor |
| """ |
| x = x.transpose(1, 2) |
| k = self.filter(x.shape[-1]) |
| y = fft_conv(x, k, dropout_mask=None, gelu=False) |
|
|
| y = y.transpose(1, 2) |
| return y.to(dtype=x.dtype) |
|
|
