Kernels:
Trusted publisher
File size: 5,495 Bytes
e19323e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 | # Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
# For a list of all contributors, visit:
# https://github.com/fla-org/flash-linear-attention/graphs/contributors
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
# The time embedding fed to the filteres is normalized so that t_f = 1
t = torch.linspace(0, 1, self.seq_len)[None, :, None] # 1, L, 1
if emb_dim > 1:
bands = (emb_dim - 1) // 2
# To compute the right embeddings we use the "proper" linspace
t_rescaled = torch.linspace(0, seq_len - 1, seq_len)[None, :, None]
w = 2 * math.pi * t_rescaled / seq_len # 1, L, 1
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)
# final linear layer
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)
|