从 Decoder 到小型语言模型：标签移位与因果约束

import torch
from torch import nn
from torch.nn import functional as F
torch.manual_seed(22)
torch.set_num_threads(1)
sequence = (torch.arange(9)[None, :] + torch.arange(4)[:, None]) % 4
x, y = sequence[:, :-1], sequence[:, 1:]  # 外部只移位一次
class TinyLM(nn.Module):
    def __init__(self):
        super().__init__()
        self.embedding = nn.Embedding(4, 16)
        self.position = nn.Embedding(8, 16)
        self.block = nn.TransformerEncoderLayer(16, 2, 32, dropout=0.,
                                                batch_first=True, norm_first=True)
        self.head = nn.Linear(16, 4)
    def forward(self, ids):
        t = ids.shape[1]
        hidden = self.embedding(ids) + self.position(torch.arange(t))
        blocked = torch.ones(t, t, dtype=torch.bool).triu(1)  # 此接口True表示禁止
        return self.head(self.block(hidden, src_mask=blocked))
model = TinyLM()
optimizer = torch.optim.AdamW(model.parameters(), lr=.02)
assert x[0].tolist() == [0, 1, 2, 3, 0, 1, 2, 3]
assert y[0].tolist() == [1, 2, 3, 0, 1, 2, 3, 0]
with torch.no_grad():
    initial_logits = model(x)
    initial = F.cross_entropy(initial_logits.reshape(-1, 4), y.reshape(-1)).item()
    changed = x.clone()
    changed[:, -1] = (changed[:, -1] + 1) % 4
    assert torch.allclose(initial_logits[:, :-1], model(changed)[:, :-1], atol=1e-6)
for _ in range(60):
    optimizer.zero_grad(set_to_none=True)
    logits = model(x)
    loss = F.cross_entropy(logits.reshape(-1, 4), y.reshape(-1))
    assert torch.isfinite(loss)
    loss.backward()
    optimizer.step()
with torch.no_grad():
    logits = model(x)
    final = F.cross_entropy(logits.reshape(-1, 4), y.reshape(-1)).item()
assert logits.shape == (4, 8, 4)
assert final < initial * .2
assert (logits.argmax(-1) == y).float().mean() > .95
print(f"CPU toy LM checks passed; loss {initial:.4f} -> {final:.4f}")