| |
| |
| |
| |
| |
| |
|
|
| import math |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from einops import rearrange |
|
|
|
|
| def fft_conv( |
| u: torch.Tensor, |
| k: torch.Tensor, |
| dropout_mask: torch.Tensor | None, |
| gelu: bool = True, |
| k_rev: torch.Tensor | None = None, |
| ) -> torch.Tensor: |
| 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): |
| """ |
| Apply FFT convolution with a learned filter of length max_len. |
| |
| Args: |
| hidden_size (int): |
| The number of expected features in the input and output. |
| max_len (int): |
| The maximum sequence length. |
| |
| Returns: |
| y (torch.Tensor): |
| Output of shape [batch_size, seq_len, hidden_size]. |
| """ |
|
|
| def __init__(self, hidden_size: int, max_len: int, **kwargs) -> None: |
| 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) -> torch.Tensor: |
| """ |
| Args: |
| x (torch.Tensor): |
| Input of shape [batch_size, seq_len, hidden_size]. |
| |
| Returns: |
| y (torch.Tensor): |
| Output of shape [batch_size, seq_len, hidden_size]. |
| """ |
| x = x.transpose(1, 2) |
| y = fft_conv(u=x, k=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) -> None: |
| """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: int) -> torch.Tensor: |
| 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 (int, Optional): |
| The dimension of the positional embeddings. Must be odd and greater or equal to 3 (time, sine and cosine). |
| Default: 3. |
| d_hidden (int, Optional): |
| The number of features in the hidden layer of the MLP. Default: 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) -> None: |
| 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(emb_dim=d_emb, seq_len=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) -> torch.Tensor: |
| return self.mlp(self.pos_emb(seq_len)).transpose(1, 2) |
|
|
| def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor: |
| """ |
| Args: |
| x (torch.Tensor): |
| Input of shape [batch_size, seq_len, hidden_size]. |
| |
| Returns: |
| y (torch.Tensor): |
| Output of shape [batch_size, seq_len, hidden_size]. |
| """ |
| x = x.transpose(1, 2) |
| k = self.filter(x.shape[-1]) |
| y = fft_conv(u=x, k=k, dropout_mask=None, gelu=False) |
|
|
| y = y.transpose(1, 2) |
| return y.to(dtype=x.dtype) |
|
|
