|
|
import math |
|
|
import torch |
|
|
|
|
|
import torch.nn.functional as F |
|
|
|
|
|
from torch import nn, einsum |
|
|
from functools import partial |
|
|
from einops import rearrange, repeat, pack, unpack |
|
|
|
|
|
def exists(val): |
|
|
return val is not None |
|
|
|
|
|
def default(value, d): |
|
|
return value if exists(value) else d |
|
|
|
|
|
def empty(tensor): |
|
|
return tensor.numel() == 0 |
|
|
|
|
|
def pad_to_multiple(tensor, multiple, dim=-1, value=0): |
|
|
seqlen = tensor.shape[dim] |
|
|
m = seqlen / multiple |
|
|
|
|
|
if m.is_integer(): return False, tensor |
|
|
return True, F.pad(tensor, (*((0,) * (-1 - dim) * 2), 0, (math.ceil(m) * multiple - seqlen)), value = value) |
|
|
|
|
|
def look_around(x, backward = 1, forward = 0, pad_value = -1, dim = 2): |
|
|
t = x.shape[1] |
|
|
dims = (len(x.shape) - dim) * (0, 0) |
|
|
padded_x = F.pad(x, (*dims, backward, forward), value = pad_value) |
|
|
|
|
|
return torch.cat([padded_x[:, ind:(ind + t), ...] for ind in range(forward + backward + 1)], dim = dim) |
|
|
|
|
|
def rotate_half(x): |
|
|
x1, x2 = rearrange(x, 'b ... (r d) -> b ... r d', r = 2).unbind(dim = -2) |
|
|
return torch.cat((-x2, x1), dim = -1) |
|
|
|
|
|
def apply_rotary_pos_emb(q, k, freqs, scale = 1): |
|
|
q_len = q.shape[-2] |
|
|
q_freqs = freqs[..., -q_len:, :] |
|
|
|
|
|
inv_scale = scale ** -1 |
|
|
if scale.ndim == 2: scale = scale[-q_len:, :] |
|
|
|
|
|
q = (q * q_freqs.cos() * scale) + (rotate_half(q) * q_freqs.sin() * scale) |
|
|
k = (k * freqs.cos() * inv_scale) + (rotate_half(k) * freqs.sin() * inv_scale) |
|
|
|
|
|
return q, k |
|
|
|
|
|
def orthogonal_matrix_chunk(cols, qr_uniform_q=False, device=None): |
|
|
unstructured_block = torch.randn((cols, cols), device=device) |
|
|
|
|
|
q, r = torch.linalg.qr(unstructured_block.cpu(), mode="reduced") |
|
|
q, r = map(lambda t: t.to(device), (q, r)) |
|
|
|
|
|
if qr_uniform_q: |
|
|
d = r.diag(0) |
|
|
q *= d.sign() |
|
|
|
|
|
return q.t() |
|
|
|
|
|
def gaussian_orthogonal_random_matrix(nb_rows, nb_columns, scaling=0, qr_uniform_q=False, device=None): |
|
|
nb_full_blocks = int(nb_rows / nb_columns) |
|
|
block_list = [] |
|
|
|
|
|
for _ in range(nb_full_blocks): |
|
|
block_list.append(orthogonal_matrix_chunk(nb_columns, qr_uniform_q=qr_uniform_q, device=device)) |
|
|
|
|
|
remaining_rows = nb_rows - nb_full_blocks * nb_columns |
|
|
|
|
|
if remaining_rows > 0: |
|
|
block_list.append( |
|
|
orthogonal_matrix_chunk( |
|
|
nb_columns, |
|
|
qr_uniform_q=qr_uniform_q, |
|
|
device=device |
|
|
)[:remaining_rows] |
|
|
) |
|
|
if scaling == 0: |
|
|
multiplier = torch.randn( |
|
|
(nb_rows, nb_columns), |
|
|
device=device |
|
|
).norm(dim=1) |
|
|
elif scaling == 1: |
|
|
multiplier = math.sqrt( |
|
|
(float(nb_columns)) |
|
|
) * torch.ones( |
|
|
(nb_rows,), |
|
|
device=device |
|
|
) |
|
|
else: raise ValueError(f"{scaling} != 0, 1") |
|
|
|
|
|
return multiplier.diag() @ torch.cat(block_list) |
|
|
|
|
|
def linear_attention(q, k, v): |
|
|
return einsum( |
|
|
"...ed,...nd->...ne", |
|
|
k, |
|
|
q |
|
|
) if v is None else einsum( |
|
|
"...de,...nd,...n->...ne", |
|
|
einsum( |
|
|
"...nd,...ne->...de", |
|
|
k, |
|
|
v |
|
|
), |
|
|
q, |
|
|
1.0 / (einsum( |
|
|
"...nd,...d->...n", |
|
|
q, |
|
|
k.sum(dim=-2).type_as(q) |
|
|
) + 1e-8) |
|
|
) |
|
|
|
|
|
def softmax_kernel(data, *, projection_matrix, is_query, normalize_data=True, eps=1e-4, device=None): |
|
|
b, h, *_ = data.shape |
|
|
|
|
|
data_normalizer = (data.shape[-1] ** -0.25) if normalize_data else 1.0 |
|
|
ratio = projection_matrix.shape[0] ** -0.5 |
|
|
|
|
|
data_dash = torch.einsum("...id,...jd->...ij", (data_normalizer * data), repeat(projection_matrix, "j d -> b h j d", b=b, h=h).type_as(data)) |
|
|
diag_data = (((data**2).sum(dim=-1) / 2.0) * (data_normalizer**2)).unsqueeze(dim=-1) |
|
|
|
|
|
return (ratio * ((data_dash - diag_data - data_dash.max(dim=-1, keepdim=True).values).exp() + eps) if is_query else ratio * ((data_dash - diag_data + eps).exp())).type_as(data) |
|
|
|
|
|
class SinusoidalEmbeddings(nn.Module): |
|
|
def __init__( |
|
|
self, |
|
|
dim, |
|
|
scale_base = None, |
|
|
use_xpos = False, |
|
|
theta = 10000 |
|
|
): |
|
|
super().__init__() |
|
|
inv_freq = 1. / (theta ** (torch.arange(0, dim, 2).float() / dim)) |
|
|
self.register_buffer('inv_freq', inv_freq) |
|
|
self.use_xpos = use_xpos |
|
|
self.scale_base = scale_base |
|
|
assert not (use_xpos and not exists(scale_base)) |
|
|
scale = (torch.arange(0, dim, 2) + 0.4 * dim) / (1.4 * dim) |
|
|
self.register_buffer('scale', scale, persistent = False) |
|
|
|
|
|
def forward(self, x): |
|
|
seq_len, device = x.shape[-2], x.device |
|
|
t = torch.arange(seq_len, device = x.device).type_as(self.inv_freq) |
|
|
|
|
|
freqs = torch.einsum('i , j -> i j', t, self.inv_freq) |
|
|
freqs = torch.cat((freqs, freqs), dim = -1) |
|
|
|
|
|
if not self.use_xpos: return freqs, torch.ones(1, device = device) |
|
|
|
|
|
power = (t - (seq_len // 2)) / self.scale_base |
|
|
scale = self.scale ** rearrange(power, 'n -> n 1') |
|
|
|
|
|
return freqs, torch.cat((scale, scale), dim = -1) |
|
|
|
|
|
class LocalAttention(nn.Module): |
|
|
def __init__( |
|
|
self, |
|
|
window_size, |
|
|
causal = False, |
|
|
look_backward = 1, |
|
|
look_forward = None, |
|
|
dropout = 0., |
|
|
shared_qk = False, |
|
|
rel_pos_emb_config = None, |
|
|
dim = None, |
|
|
autopad = False, |
|
|
exact_windowsize = False, |
|
|
scale = None, |
|
|
use_rotary_pos_emb = True, |
|
|
use_xpos = False, |
|
|
xpos_scale_base = None |
|
|
): |
|
|
super().__init__() |
|
|
look_forward = default(look_forward, 0 if causal else 1) |
|
|
assert not (causal and look_forward > 0) |
|
|
self.scale = scale |
|
|
self.window_size = window_size |
|
|
self.autopad = autopad |
|
|
self.exact_windowsize = exact_windowsize |
|
|
self.causal = causal |
|
|
self.look_backward = look_backward |
|
|
self.look_forward = look_forward |
|
|
self.dropout = nn.Dropout(dropout) |
|
|
self.shared_qk = shared_qk |
|
|
self.rel_pos = None |
|
|
self.use_xpos = use_xpos |
|
|
if use_rotary_pos_emb and (exists(rel_pos_emb_config) or exists(dim)): |
|
|
if exists(rel_pos_emb_config): dim = rel_pos_emb_config[0] |
|
|
self.rel_pos = SinusoidalEmbeddings(dim, use_xpos = use_xpos, scale_base = default(xpos_scale_base, window_size // 2)) |
|
|
|
|
|
def forward(self, q, k, v, mask = None, input_mask = None, attn_bias = None, window_size = None): |
|
|
mask = default(mask, input_mask) |
|
|
assert not (exists(window_size) and not self.use_xpos) |
|
|
|
|
|
( |
|
|
_, |
|
|
autopad, |
|
|
pad_value, |
|
|
window_size, |
|
|
causal, |
|
|
look_backward, |
|
|
look_forward, |
|
|
shared_qk |
|
|
) = ( |
|
|
q.shape, |
|
|
self.autopad, |
|
|
-1, |
|
|
default( |
|
|
window_size, |
|
|
self.window_size |
|
|
), |
|
|
self.causal, |
|
|
self.look_backward, |
|
|
self.look_forward, |
|
|
self.shared_qk |
|
|
) |
|
|
|
|
|
(q, packed_shape), (k, _), (v, _) = map(lambda t: pack([t], '* n d'), (q, k, v)) |
|
|
|
|
|
if autopad: |
|
|
orig_seq_len = q.shape[1] |
|
|
(_, q), (_, k), (_, v) = map(lambda t: pad_to_multiple(t, self.window_size, dim = -2), (q, k, v)) |
|
|
|
|
|
b, n, dim_head, device, dtype = *q.shape, q.device, q.dtype |
|
|
scale = default(self.scale, dim_head ** -0.5) |
|
|
|
|
|
assert (n % window_size) == 0 |
|
|
windows = n // window_size |
|
|
|
|
|
if shared_qk: k = F.normalize(k, dim = -1).type(k.dtype) |
|
|
|
|
|
seq = torch.arange(n, device = device) |
|
|
b_t = rearrange(seq, '(w n) -> 1 w n', w = windows, n = window_size) |
|
|
bq, bk, bv = map(lambda t: rearrange(t, 'b (w n) d -> b w n d', w = windows), (q, k, v)) |
|
|
|
|
|
bq = bq * scale |
|
|
look_around_kwargs = dict(backward = look_backward, forward = look_forward, pad_value = pad_value) |
|
|
|
|
|
bk = look_around(bk, **look_around_kwargs) |
|
|
bv = look_around(bv, **look_around_kwargs) |
|
|
|
|
|
if exists(self.rel_pos): |
|
|
pos_emb, xpos_scale = self.rel_pos(bk) |
|
|
bq, bk = apply_rotary_pos_emb(bq, bk, pos_emb, scale = xpos_scale) |
|
|
|
|
|
bq_t = b_t |
|
|
bq_k = look_around(b_t, **look_around_kwargs) |
|
|
bq_t = rearrange(bq_t, '... i -> ... i 1') |
|
|
bq_k = rearrange(bq_k, '... j -> ... 1 j') |
|
|
|
|
|
pad_mask = bq_k == pad_value |
|
|
sim = einsum('b h i e, b h j e -> b h i j', bq, bk) |
|
|
|
|
|
if exists(attn_bias): |
|
|
heads = attn_bias.shape[0] |
|
|
assert (b % heads) == 0 |
|
|
|
|
|
attn_bias = repeat(attn_bias, 'h i j -> (b h) 1 i j', b = b // heads) |
|
|
sim = sim + attn_bias |
|
|
|
|
|
mask_value = -torch.finfo(sim.dtype).max |
|
|
if shared_qk: |
|
|
self_mask = bq_t == bq_k |
|
|
sim = sim.masked_fill(self_mask, -5e4) |
|
|
del self_mask |
|
|
|
|
|
if causal: |
|
|
causal_mask = bq_t < bq_k |
|
|
if self.exact_windowsize: causal_mask = causal_mask | (bq_t > (bq_k + (self.window_size * self.look_backward))) |
|
|
sim = sim.masked_fill(causal_mask, mask_value) |
|
|
del causal_mask |
|
|
|
|
|
sim = sim.masked_fill( |
|
|
((bq_k - (self.window_size * self.look_forward)) > bq_t) | (bq_t > (bq_k + (self.window_size * self.look_backward))) | pad_mask, |
|
|
mask_value |
|
|
) if not causal and self.exact_windowsize else sim.masked_fill( |
|
|
pad_mask, |
|
|
mask_value |
|
|
) |
|
|
|
|
|
if exists(mask): |
|
|
batch = mask.shape[0] |
|
|
assert (b % batch) == 0 |
|
|
|
|
|
h = b // mask.shape[0] |
|
|
if autopad: _, mask = pad_to_multiple(mask, window_size, dim = -1, value = False) |
|
|
|
|
|
mask = repeat( |
|
|
rearrange( |
|
|
look_around( |
|
|
rearrange( |
|
|
mask, |
|
|
'... (w n) -> (...) w n', |
|
|
w = windows, |
|
|
n = window_size |
|
|
), |
|
|
**{ |
|
|
**look_around_kwargs, |
|
|
'pad_value': False |
|
|
} |
|
|
), |
|
|
'... j -> ... 1 j' |
|
|
), |
|
|
'b ... -> (b h) ...', |
|
|
h = h |
|
|
) |
|
|
|
|
|
sim = sim.masked_fill(~mask, mask_value) |
|
|
del mask |
|
|
|
|
|
out = rearrange( |
|
|
einsum( |
|
|
'b h i j, b h j e -> b h i e', |
|
|
self.dropout(sim.softmax(dim = -1)), |
|
|
bv |
|
|
), |
|
|
'b w n d -> b (w n) d' |
|
|
) |
|
|
|
|
|
if autopad: out = out[:, :orig_seq_len, :] |
|
|
out, *_ = unpack(out, packed_shape, '* n d') |
|
|
|
|
|
return out |
|
|
|
|
|
class FastAttention(nn.Module): |
|
|
def __init__( |
|
|
self, |
|
|
dim_heads, |
|
|
nb_features=None, |
|
|
ortho_scaling=0, |
|
|
causal=False, |
|
|
generalized_attention=False, |
|
|
kernel_fn=nn.ReLU(), |
|
|
qr_uniform_q=False, |
|
|
no_projection=False |
|
|
): |
|
|
super().__init__() |
|
|
nb_features = default(nb_features, int(dim_heads * math.log(dim_heads))) |
|
|
self.dim_heads = dim_heads |
|
|
self.nb_features = nb_features |
|
|
self.ortho_scaling = ortho_scaling |
|
|
self.create_projection = partial( |
|
|
gaussian_orthogonal_random_matrix, |
|
|
nb_rows=self.nb_features, |
|
|
nb_columns=dim_heads, |
|
|
scaling=ortho_scaling, |
|
|
qr_uniform_q=qr_uniform_q |
|
|
) |
|
|
projection_matrix = self.create_projection() |
|
|
self.register_buffer("projection_matrix", projection_matrix) |
|
|
self.generalized_attention = generalized_attention |
|
|
self.kernel_fn = kernel_fn |
|
|
self.no_projection = no_projection |
|
|
self.causal = causal |
|
|
|
|
|
@torch.no_grad() |
|
|
def redraw_projection_matrix(self): |
|
|
projections = self.create_projection() |
|
|
self.projection_matrix.copy_(projections) |
|
|
del projections |
|
|
|
|
|
def forward(self, q, k, v): |
|
|
if self.no_projection: q, k = q.softmax(dim=-1), (k.exp() if self.causal else k.softmax(dim=-2)) |
|
|
else: |
|
|
create_kernel = partial(softmax_kernel, projection_matrix=self.projection_matrix, device=q.device) |
|
|
q, k = create_kernel(q, is_query=True), create_kernel(k, is_query=False) |
|
|
|
|
|
attn_fn = linear_attention if not self.causal else self.causal_linear_fn |
|
|
return attn_fn(q, k, None) if v is None else attn_fn(q, k, v) |
|
|
|
|
|
class SelfAttention(nn.Module): |
|
|
def __init__( |
|
|
self, |
|
|
dim, |
|
|
causal=False, |
|
|
heads=8, |
|
|
dim_head=64, |
|
|
local_heads=0, |
|
|
local_window_size=256, |
|
|
nb_features=None, |
|
|
feature_redraw_interval=1000, |
|
|
generalized_attention=False, |
|
|
kernel_fn=nn.ReLU(), |
|
|
qr_uniform_q=False, |
|
|
dropout=0.0, |
|
|
no_projection=False |
|
|
): |
|
|
super().__init__() |
|
|
assert dim % heads == 0 |
|
|
dim_head = default(dim_head, dim // heads) |
|
|
inner_dim = dim_head * heads |
|
|
self.fast_attention = FastAttention( |
|
|
dim_head, |
|
|
nb_features, |
|
|
causal=causal, |
|
|
generalized_attention=generalized_attention, |
|
|
kernel_fn=kernel_fn, |
|
|
qr_uniform_q=qr_uniform_q, |
|
|
no_projection=no_projection |
|
|
) |
|
|
self.heads = heads |
|
|
self.global_heads = heads - local_heads |
|
|
self.local_attn = ( |
|
|
LocalAttention( |
|
|
window_size=local_window_size, |
|
|
causal=causal, |
|
|
autopad=True, |
|
|
dropout=dropout, |
|
|
look_forward=int(not causal), |
|
|
rel_pos_emb_config=(dim_head, local_heads) |
|
|
) if local_heads > 0 else None |
|
|
) |
|
|
self.to_q = nn.Linear(dim, inner_dim) |
|
|
self.to_k = nn.Linear(dim, inner_dim) |
|
|
self.to_v = nn.Linear(dim, inner_dim) |
|
|
self.to_out = nn.Linear(inner_dim, dim) |
|
|
self.dropout = nn.Dropout(dropout) |
|
|
|
|
|
@torch.no_grad() |
|
|
def redraw_projection_matrix(self): |
|
|
self.fast_attention.redraw_projection_matrix() |
|
|
|
|
|
def forward(self, x, context=None, mask=None, context_mask=None, name=None, inference=False, **kwargs): |
|
|
_, _, _, h, gh = *x.shape, self.heads, self.global_heads |
|
|
cross_attend = exists(context) |
|
|
context = default(context, x) |
|
|
context_mask = default(context_mask, mask) if not cross_attend else context_mask |
|
|
|
|
|
q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h=h), (self.to_q(x), self.to_k(context), self.to_v(context))) |
|
|
(q, lq), (k, lk), (v, lv) = map(lambda t: (t[:, :gh], t[:, gh:]), (q, k, v)) |
|
|
|
|
|
attn_outs = [] |
|
|
|
|
|
if not empty(q): |
|
|
if exists(context_mask): v.masked_fill_(~context_mask[:, None, :, None], 0.0) |
|
|
if cross_attend: pass |
|
|
else: out = self.fast_attention(q, k, v) |
|
|
|
|
|
attn_outs.append(out) |
|
|
|
|
|
if not empty(lq): |
|
|
assert (not cross_attend), "not cross_attend" |
|
|
|
|
|
out = self.local_attn(lq, lk, lv, input_mask=mask) |
|
|
attn_outs.append(out) |
|
|
|
|
|
return self.dropout(self.to_out(rearrange(torch.cat(attn_outs, dim=1), "b h n d -> b n (h d)"))) |
