File size: 4,187 Bytes
b66f552 | 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 | # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang, Yuqi Pan
import torch
import torch.nn.functional as F
def ttt_linear(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
w: torch.Tensor,
b: torch.Tensor,
eta: torch.Tensor,
scale: float,
eps: float,
mini_batch_size: int,
initial_state: torch.Tensor,
initial_state_bias: torch.Tensor,
output_final_state: bool,
):
B, H, T, D = q.shape
BT = mini_batch_size
NT = T // BT
# [NT, B, H, mini_batch_size, D]
_q = q.reshape(B, H, NT, BT, D).permute(2, 0, 1, 3, 4)
_k = k.reshape(B, H, NT, BT, D).permute(2, 0, 1, 3, 4)
_v = v.reshape(B, H, NT, BT, D).permute(2, 0, 1, 3, 4)
# [NT, B, H, BT, 1]
_eta = eta.reshape(B, H, NT, BT, 1).permute(2, 0, 1, 3, 4)
# [H, 1, D]
w = w.reshape(H, 1, D).to(torch.float32)
b = b.reshape(H, 1, D).to(torch.float32)
h = torch.zeros((B, H, D, D), device=v.device, dtype=torch.float32) if initial_state is None else initial_state
hb = torch.zeros((B, H, 1, D), device=v.device, dtype=torch.float32) if initial_state_bias is None else initial_state_bias
q *= scale
# [NT, B, H, BT, D]
o = torch.empty_like(_v)
for i in range(NT):
q_i, k_i, v_i, eta_i = [x[i] for x in [_q, _k, _v, _eta]]
kh = k_i @ h + hb
reconstruction_target = v_i - k_i
mean = kh.mean(-1, True)
var = kh.var(-1, unbiased=False, keepdim=True).to(torch.float32)
rstd = torch.sqrt(var + eps).to(torch.float32)
kh_hat = (kh - mean) / rstd
g = w * kh_hat + b - reconstruction_target
g *= w
v_new = (D * g - g.sum(-1, True) - kh_hat * (g * kh_hat).sum(-1, True)) / (rstd * D)
Attn = torch.tril(q_i @ k_i.transpose(-2, -1))
o_i = q_i @ h - (eta_i * Attn) @ v_new + hb - torch.tril(eta_i.expand_as(Attn)) @ v_new
h = h - (eta_i[:, :, -1, :, None] * k_i).transpose(-1, -2) @ v_new
hb = hb - torch.sum(eta_i[:, :, -1, :, None] * v_new, dim=-2, keepdim=True)
# layer norm with residuals
mean = o_i.mean(dim=-1, keepdim=True)
var = o_i.var(dim=-1, unbiased=False, keepdim=True).to(torch.float32)
rstd = torch.sqrt(var + eps).to(torch.float32)
o[i] = o_i + (o_i - mean) / rstd * w + b
# [B, H, T, D]
o = o.permute(1, 2, 0, 3, 4).reshape(B, H, T, D)
h = h if output_final_state else None
hb = hb if output_final_state else None
return o, h, hb
def chunk_ttt_linear_ref(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
w: torch.Tensor,
b: torch.Tensor,
eta: torch.Tensor,
scale: float = None,
eps: float = 1e-6,
mini_batch_size: int = 16,
initial_state: torch.Tensor = None,
initial_state_bias: torch.Tensor = None,
output_final_state: bool = False,
head_first: bool = False,
):
assert q.dtype == k.dtype == v.dtype
assert k.shape[-1] == v.shape[-1], "The key and value dimension must be the same."
if isinstance(eta, float):
eta = torch.full_like(q[:, :, :, :1], eta)
if scale is None:
scale = k.shape[-1] ** -0.5
if not head_first:
q = q.transpose(1, 2)
k = k.transpose(1, 2)
v = v.transpose(1, 2)
eta = eta.transpose(1, 2)
T = q.shape[-2]
padded = (mini_batch_size - (T % mini_batch_size)) % mini_batch_size
if padded > 0:
q = F.pad(q, (0, 0, 0, padded))
k = F.pad(k, (0, 0, 0, padded))
v = F.pad(v, (0, 0, 0, padded))
eta = F.pad(eta, (0, 0, 0, padded))
eta[:, :, -1, :] = eta[:, :, -(padded+1), :]
assert q.shape[-2] % mini_batch_size == 0, "Sequence length should be a multiple of mini_batch_size."
q, k, v, eta, w, b = map(lambda x: x.to(torch.float32), [q, k, v, eta, w, b])
o, final_state, final_state_bias = ttt_linear(
q,
k,
v,
w,
b,
eta,
scale,
eps,
mini_batch_size,
initial_state,
initial_state_bias,
output_final_state,
)
o = o[:, :, :T, :].contiguous()
if not head_first:
o = o.transpose(1, 2)
return o, final_state, final_state_bias
|