Causal Transformer 入门指南：从基础原理到实战应用

1次阅读

共计 2830 个字符，预计需要花费 8 分钟才能阅读完成。

传统 Transformer 在并行处理序列数据时表现优异，但在自回归生成任务（如文本生成）中存在一个关键缺陷：解码时可能访问未来信息（即信息泄露）。例如生成第 t 个词时，模型理论上能利用 t+1 时刻的信息，这违背了实际应用场景的因果性。

Causal Transformer 通过引入因果注意力掩码（Causal Attention Mask）解决该问题，确保模型仅能访问当前位置及之前的信息。这种特性使其在以下场景中成为必需品：

语言模型预训练（如 GPT 系列）
实时序列生成（语音合成、代码补全）
时间序列预测

因果掩码是一个下三角矩阵，其数学表示为：

M_{ij} = {0 if i ≥ j, -∞ otherwise}

通过将掩码应用于注意力分数的 Softmax 前，使得位置 i 只能关注位置 ≤i 的输入。下图展示了序列长度为 4 时的掩码效果：

[0 -∞ -∞ -∞]
[0  0 -∞ -∞]
[0  0  0 -∞]
[0  0  0  0]

给定输入序列 X ∈ R^(L×d)，其注意力计算流程为：

计算 Q = XW_q, K = XW_k, V = XW_v
注意力分数 S = QK^T/√d_k
应用掩码 S_masked = S + M
注意力权重 A = softmax(S_masked)
输出 O = AV

import torch
import torch.nn as nn
import torch.nn.functional as F

class CausalAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.d_head = d_model // n_heads
        self.n_heads = n_heads
        self.qkv_proj = nn.Linear(d_model, 3*d_model)
        self.out_proj = nn.Linear(d_model, d_model)

    def forward(self, x, mask=None):
        B, L, _ = x.shape
        qkv = self.qkv_proj(x)
        q, k, v = qkv.chunk(3, dim=-1)

        # Reshape for multi-head attention
        q = q.view(B, L, self.n_heads, self.d_head).transpose(1, 2)
        k = k.view(B, L, self.n_heads, self.d_head).transpose(1, 2)
        v = v.view(B, L, self.n_heads, self.d_head).transpose(1, 2)

        # Scaled dot-product attention
        attn_scores = (q @ k.transpose(-2, -1)) / (self.d_head ** 0.5)

        # Apply causal mask
        causal_mask = torch.triu(torch.ones(L, L) * -float('inf'), diagonal=1)
        causal_mask = causal_mask.to(x.device)
        attn_scores = attn_scores + causal_mask

        attn_weights = F.softmax(attn_scores, dim=-1)
        output = attn_weights @ v
        output = output.transpose(1, 2).contiguous().view(B, L, -1)
        return self.out_proj(output)

def train(model, dataloader, optimizer, epochs):
    model.train()
    for epoch in range(epochs):
        for batch in dataloader:
            inputs, targets = batch
            optimizer.zero_grad()

            # Forward pass with teacher forcing
            outputs = model(inputs)
            loss = F.cross_entropy(outputs.view(-1, vocab_size), targets.view(-1))

            # Backward pass
            loss.backward()
            optimizer.step()

def generate(model, prompt, max_len=50):
    model.eval()
    generated = prompt.clone()

    for _ in range(max_len):
        # Get predictions for last token only
        with torch.no_grad():
            logits = model(generated)[:, -1, :]

        # Sample next token
        probs = F.softmax(logits, dim=-1)
        next_token = torch.multinomial(probs, num_samples=1)
        generated = torch.cat([generated, next_token], dim=1)

    return generated