-
Notifications
You must be signed in to change notification settings - Fork 1
/
s4d.py
134 lines (105 loc) · 4.6 KB
/
s4d.py
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
"""Minimal version of S4D with extra options and features stripped out, for pedagogical purposes."""
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange, repeat
class DropoutNd(nn.Module):
def __init__(self, p: float = 0.5, tie=True, transposed=True):
"""
tie: tie dropout mask across sequence lengths (Dropout1d/2d/3d)
"""
super().__init__()
if p < 0 or p >= 1:
raise ValueError("dropout probability has to be in [0, 1), " "but got {}".format(p))
self.p = p
self.tie = tie
self.transposed = transposed
self.binomial = torch.distributions.binomial.Binomial(probs=1-self.p)
def forward(self, X):
"""X: (batch, dim, lengths...)."""
if self.training:
if not self.transposed: X = rearrange(X, 'b ... d -> b d ...')
mask_shape = X.shape[:2] + (1,)*(X.ndim-2) if self.tie else X.shape
mask = torch.rand(*mask_shape, device=X.device) < 1.-self.p
X = X * mask * (1.0/(1-self.p))
if not self.transposed: X = rearrange(X, 'b d ... -> b ... d')
return X
return X
class S4DKernel(nn.Module):
"""Generate convolution kernel from diagonal SSM parameters."""
def __init__(self, d_model, N=64, dt_min=0.001, dt_max=0.1, lr=None):
super().__init__()
# Generate dt
H = d_model
log_dt = torch.rand(H) * (
math.log(dt_max) - math.log(dt_min)
) + math.log(dt_min)
C = torch.randn(H, N // 2, dtype=torch.cfloat)
self.C = nn.Parameter(torch.view_as_real(C))
self.register("log_dt", log_dt, lr)
log_A_real = torch.log(0.5 * torch.ones(H, N//2))
A_imag = math.pi * repeat(torch.arange(N//2), 'n -> h n', h=H)
self.register("log_A_real", log_A_real, lr)
self.register("A_imag", A_imag, lr)
def forward(self, L):
"""
returns: (..., c, L) where c is number of channels (default 1)
"""
# Materialize parameters
dt = torch.exp(self.log_dt) # (H)
C = torch.view_as_complex(self.C) # (H N)
A = -torch.exp(self.log_A_real) + 1j * self.A_imag # (H N)
# Vandermonde multiplication
dtA = A * dt.unsqueeze(-1) # (H N)
K = dtA.unsqueeze(-1) * torch.arange(L, device=A.device) # (H N L)
C = C * (torch.exp(dtA)-1.) / A
K = 2 * torch.einsum('hn, hnl -> hl', C, torch.exp(K)).real
return K
def register(self, name, tensor, lr=None):
"""Register a tensor with a configurable learning rate and 0 weight decay"""
if lr == 0.0:
self.register_buffer(name, tensor)
else:
self.register_parameter(name, nn.Parameter(tensor))
optim = {"weight_decay": 0.0}
if lr is not None: optim["lr"] = lr
setattr(getattr(self, name), "_optim", optim)
class S4D(nn.Module):
def __init__(self, d_model, d_state=64, dropout=0.0, transposed=True, **kernel_args):
super().__init__()
self.h = d_model
self.n = d_state
self.d_output = self.h
self.transposed = transposed
self.D = nn.Parameter(torch.randn(self.h))
# SSM Kernel
self.kernel = S4DKernel(self.h, N=self.n, **kernel_args)
# Pointwise
self.activation = nn.GELU()
# dropout_fn = nn.Dropout2d # NOTE: bugged in PyTorch 1.11
dropout_fn = DropoutNd
self.dropout = dropout_fn(dropout) if dropout > 0.0 else nn.Identity()
# position-wise output transform to mix features
self.output_linear = nn.Sequential(
nn.Conv1d(self.h, 2*self.h, kernel_size=1),
nn.GLU(dim=-2),
)
def forward(self, u, prev_hidden=None, **kwargs): # absorbs return_output and transformer src mask
""" Input and output shape (B, H, L) """
if not self.transposed: u = u.transpose(-1, -2)
L = u.size(-1)
# Compute SSM Kernel
k = self.kernel(L=L) # (H L)
# Convolution
k_f = torch.fft.rfft(k, n=2*L) # (H L)
u_f = torch.fft.rfft(u, n=2*L) # (B H L)
y = torch.fft.irfft(u_f*k_f, n=2*L)[..., :L] # (B H L)
# Compute D term in state space equation - essentially a skip connection
y = y + u * self.D.unsqueeze(-1)
hidden = y[:, :, -1] # (B H)
y = self.dropout(self.activation(y))
y = self.output_linear(y)
if not self.transposed: y = y.transpose(-1, -2)
# return y, None # Return a dummy state to satisfy this repo's interface, but this can be modified
return y, hidden.unsqueeze(0)