File size: 10,621 Bytes
98c193b 55003bd 98c193b 55003bd 98c193b 55003bd 98c193b 55003bd 98c193b 55003bd 98c193b ae02605 8c85090 ae02605 8c85090 ae02605 8c85090 ae02605 | 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 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 | """Transformer denoiser for masked diffusion.
Bidirectional by default (a masked diffusion denoiser must see both sides);
`causal=True` gives the matched autoregressive baseline. Rotary is the default
positional encoding - it was the strongest option in our previous study.
"""
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
def sinusoidal(max_len, d):
pe = torch.zeros(max_len, d)
pos = torch.arange(max_len).unsqueeze(1).float()
div = torch.exp(torch.arange(0, d, 2).float() * (-math.log(10000.0) / d))
pe[:, 0::2] = torch.sin(pos * div)
pe[:, 1::2] = torch.cos(pos * div)
return pe
def alibi_slopes(n_heads):
start = 2 ** (-(2 ** -(math.log2(n_heads) - 3))) if math.log2(n_heads).is_integer() else 0.5
return torch.tensor([start * (start ** i) for i in range(n_heads)])
def apply_rope(x, cos, sin):
# x: (B, H, L, Dh)
x1, x2 = x[..., 0::2], x[..., 1::2]
rx = torch.stack([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
return rx.flatten(-2)
class Attention(nn.Module):
def __init__(self, d, n_heads, pe, causal, max_len):
super().__init__()
self.h, self.dh, self.pe, self.causal = n_heads, d // n_heads, pe, causal
self.qkv = nn.Linear(d, 3 * d, bias=False)
self.proj = nn.Linear(d, d, bias=False)
if pe == "alibi":
self.register_buffer("slopes", alibi_slopes(n_heads), persistent=False)
if pe == "rope":
inv = 1.0 / (10000 ** (torch.arange(0, self.dh, 2).float() / self.dh))
ang = torch.arange(max_len).float()[:, None] * inv[None, :]
self.register_buffer("cos", ang.cos(), persistent=False)
self.register_buffer("sin", ang.sin(), persistent=False)
def forward(self, x, pad_mask, pos_ids=None):
B, L, D = x.shape
q, k, v = self.qkv(x).chunk(3, dim=-1)
q = q.view(B, L, self.h, self.dh).transpose(1, 2)
k = k.view(B, L, self.h, self.dh).transpose(1, 2)
v = v.view(B, L, self.h, self.dh).transpose(1, 2)
if self.pe == "rope":
if pos_ids is None:
cos, sin = self.cos[:L][None, None], self.sin[:L][None, None]
else:
cos, sin = self.cos[pos_ids][:, None], self.sin[pos_ids][:, None]
q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
att = (q @ k.transpose(-2, -1)) / math.sqrt(self.dh)
if self.pe == "alibi":
if pos_ids is None:
pos = torch.arange(L, device=x.device)
rel = (pos[None, :] - pos[:, None]).abs().float()[None]
else:
rel = (pos_ids[:, None, :] - pos_ids[:, :, None]).abs().float()
att = att - self.slopes.view(1, -1, 1, 1) * rel[:, None]
if self.causal:
cm = torch.ones(L, L, device=x.device, dtype=torch.bool).tril()
att = att.masked_fill(~cm[None, None], float("-inf"))
if pad_mask is not None:
att = att.masked_fill(~pad_mask[:, None, None, :], float("-inf"))
att = att.softmax(-1)
out = (att @ v).transpose(1, 2).reshape(B, L, D)
return self.proj(out)
class Block(nn.Module):
def __init__(self, d, n_heads, pe, causal, max_len):
super().__init__()
self.n1, self.n2 = nn.LayerNorm(d), nn.LayerNorm(d)
self.attn = Attention(d, n_heads, pe, causal, max_len)
self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))
def forward(self, x, pad_mask, pos_ids=None):
x = x + self.attn(self.n1(x), pad_mask, pos_ids)
return x + self.mlp(self.n2(x))
class Transformer(nn.Module):
def __init__(self, vocab, d=384, n_layers=6, n_heads=6, pe="ape", causal=False,
max_len=256, budget_bins=0):
super().__init__()
self.pe_kind, self.causal = pe, causal
self.budget_bins = budget_bins
self.emb = nn.Embedding(vocab, d)
# Budget conditioning. A K-pass decode only visits mask ratios
# {1, (K-1)/K, ..., 1/K}, so a model trained for one K is mis-specified
# at every other K - measured directly: our 20-digit model scores 99.8%
# at one pass and 85.7% at twenty. Telling the model which budget it is
# being decoded at lets ONE network serve every budget instead of a
# family of networks each good at a single point.
# budget_bins=0 keeps the parameter absent entirely, so checkpoints
# trained before this existed still load.
if budget_bins:
self.budget_emb = nn.Embedding(budget_bins, d)
# Segment embedding: which part of the equation a token belongs to.
# Place-value ids deliberately collide across operands and answer, so a
# bidirectional model needs this to tell those tokens apart; a causal
# model gets the same information free from the attention mask.
if pe == "ape":
self.pos = nn.Embedding(max_len, d)
elif pe == "sin":
self.register_buffer("pos_sin", sinusoidal(max_len, d), persistent=False)
self.blocks = nn.ModuleList([Block(d, n_heads, pe, causal, max_len) for _ in range(n_layers)])
self.norm = nn.LayerNorm(d)
self.head = nn.Linear(d, vocab, bias=False)
self.apply(self._init)
@staticmethod
def _init(m):
if isinstance(m, (nn.Linear, nn.Embedding)):
nn.init.normal_(m.weight, std=0.02)
if isinstance(m, nn.Linear) and m.bias is not None:
nn.init.zeros_(m.bias)
@staticmethod
def budget_bin(K, n_bins, device):
"""Map a decoding budget K to a bin index by log2, clamped.
Budgets are used on a log scale (1, 2, 4, ... L), so log2 is the natural
parameterisation and keeps the table small.
"""
if not torch.is_tensor(K):
K = torch.as_tensor(K, device=device)
K = K.to(device).float().clamp_min(1.0)
return torch.log2(K).round().long().clamp_(0, n_bins - 1)
def forward(self, idx, pad_mask=None, pos_ids=None, budget=None):
B, L = idx.shape
x = self.emb(idx)
if self.budget_bins and budget is not None:
b = self.budget_bin(budget, self.budget_bins, idx.device)
if b.ndim == 0:
b = b.expand(B)
x = x + self.budget_emb(b)[:, None, :]
if self.pe_kind == "ape":
p = torch.arange(L, device=idx.device)[None] if pos_ids is None else pos_ids
x = x + self.pos(p)
elif self.pe_kind == "sin":
p = torch.arange(L, device=idx.device)[None] if pos_ids is None else pos_ids
x = x + self.pos_sin[p]
for b in self.blocks:
x = b(x, pad_mask, pos_ids)
return self.head(self.norm(x))
def n_params(self):
return sum(p.numel() for p in self.parameters() if p.requires_grad)
class RecurrentDenoiser(nn.Module):
"""A weight-shared block applied R times, supervised at every step.
Motivated by a measurement, not a hunch. On hard Sudoku a 37.86M-parameter
12-layer denoiser reaches 89.4%, while the Recurrent Transformer of Yang et
al. (2023) reaches 99.5% on a comparable split with 211k parameters - 180x
smaller. The difference is not capacity, it is that constraint propagation is
an ITERATIVE algorithm: solving a hard Sudoku takes tens of rounds of
"eliminate, propagate, repeat". A fixed 12-layer feedforward network cannot
express thirty rounds of propagation at any width; one block applied thirty
times can.
This is the same argument THEORY.md makes for addition - that a K-pass decode
has effective depth L*K and "borrows depth from the decoding loop" - applied
to the architecture instead of to the decoder.
Two details carry the method:
INPUT INJECTION. The token embedding is re-added at every recurrence.
Without it the input signal decays through thirty applications of the same
block and the model forgets the clues it is solving for.
DEEP SUPERVISION. Training loss is taken at EVERY recurrence, not only the
last. That forces each application to make progress on its own rather than
letting the stack learn one entangled thirty-step function, and it is what
lets inference run MORE recurrences than training - the per-step objective
is the same at every step, so the map is iterable beyond where it was fit.
"""
def __init__(self, vocab, d=128, n_layers=1, n_heads=4, pe="ape",
max_len=128, recurrences=32, hidden_mult=4, inject=True):
super().__init__()
self.pe_kind, self.causal = pe, False
self.recurrences = recurrences
# `inject` exists to separate two things the headline comparison
# otherwise conflates: weight sharing, and re-supplying the input at
# every application. Turning it off keeps the architecture and the depth
# identical and changes only the injection, which is the clean isolation.
self.inject = inject
self.emb = nn.Embedding(vocab, d)
if pe == "ape":
self.pos = nn.Embedding(max_len, d)
elif pe == "sin":
self.register_buffer("pos_sin", sinusoidal(max_len, d), persistent=False)
self.blocks = nn.ModuleList([Block(d, n_heads, pe, False, max_len)
for _ in range(n_layers)])
self.norm = nn.LayerNorm(d)
self.head = nn.Linear(d, vocab, bias=False)
self.apply(Transformer._init)
def _inject(self, idx):
x = self.emb(idx)
if self.pe_kind == "ape":
x = x + self.pos(torch.arange(idx.shape[1], device=idx.device))[None]
elif self.pe_kind == "sin":
x = x + self.pos_sin[torch.arange(idx.shape[1], device=idx.device)][None]
return x
def forward(self, idx, pad_mask=None, pos_ids=None, budget=None,
recurrences=None, return_all=False):
"""Returns final logits, or the list of logits from every recurrence."""
R = recurrences or self.recurrences
inp = self._inject(idx)
h = inp
outs = []
for _ in range(R):
if self.inject:
h = h + inp # input injection: keep the clues alive
for b in self.blocks:
h = b(h, pad_mask, pos_ids)
outs.append(self.head(self.norm(h)))
return outs if return_all else outs[-1]
def n_params(self):
return sum(p.numel() for p in self.parameters() if p.requires_grad)
|