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

1次阅读

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

Transformer 模型在自然语言处理等领域取得了巨大成功，但其标准实现存在一个关键问题：在训练过程中，每个位置都能 ” 看到 ” 整个输入序列。这对于机器翻译等任务是有益的，但在自回归生成任务（如文本生成）中会导致信息泄露，因为模型会 ” 作弊 ” 地看到未来 token。

Causal Transformer 通过引入因果注意力掩码 (causal attention mask) 解决了这个问题。它确保在生成第 i 个 token 时，模型只能关注到 1 到 i - 1 的位置，这与人类写作时的思考过程类似。这种特性使得 Causal Transformer 特别适合：

文本生成
语音合成
时间序列预测
任何需要严格顺序建模的任务

Causal Transformer 的核心是因果注意力掩码。其数学表示为：

Attention(Q,K,V) = softmax((QK^T)/√d_k + M)V

其中 M 是掩码矩阵，定义如下：

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

图示说明：

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

这个上三角矩阵确保每个位置只能关注自身及之前的位置。

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

class CausalAttention(nn.Module):
    def __init__(self, embed_size, heads):
        super(CausalAttention, self).__init__()
        self.embed_size = embed_size
        self.heads = heads
        self.head_dim = embed_size // heads

        assert (self.head_dim * heads == embed_size), "Embedding size needs to be divisible by heads"

        self.values = nn.Linear(self.head_dim, self.head_dim, bias=False)
        self.keys = nn.Linear(self.head_dim, self.head_dim, bias=False)
        self.queries = nn.Linear(self.head_dim, self.head_dim, bias=False)
        self.fc_out = nn.Linear(heads * self.head_dim, embed_size)

    def forward(self, values, keys, query):
        N = query.shape[0]
        value_len, key_len, query_len = values.shape[1], keys.shape[1], query.shape[1]

        # Split embedding into self.heads pieces
        values = values.reshape(N, value_len, self.heads, self.head_dim)
        keys = keys.reshape(N, key_len, self.heads, self.head_dim)
        queries = query.reshape(N, query_len, self.heads, self.head_dim)

        values = self.values(values)
        keys = self.keys(keys)
        queries = self.queries(queries)

        energy = torch.einsum("nqhd,nkhd->nhqk", [queries, keys])

        # Create causal mask
        mask = torch.tril(torch.ones((query_len, key_len))).bool().to(query.device)
        energy = energy.masked_fill(~mask, float('-1e20'))

        attention = torch.softmax(energy / (self.embed_size ** (1/2)), dim=3)

        out = torch.einsum("nhql,nlhd->nqhd", [attention, values])
        out = out.reshape(N, query_len, self.heads * self.head_dim)

        return self.fc_out(out)

model = CausalTransformer()
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

for epoch in range(epochs):
    for batch in train_loader:
        optimizer.zero_grad()

        src = batch.src
        trg = batch.trg

        output = model(src, trg[:,:-1])  # Teacher forcing

        loss = criterion(output.reshape(-1, output.shape[-1]), trg[:,1:].reshape(-1))
        loss.backward()
        optimizer.step()

def generate(self, src, max_len):
    memory = self.encode(src)
    ys = torch.ones(1, 1).fill_(SOS_TOKEN).type_as(src).long()

    for i in range(max_len-1):
        out = self.decode(memory, ys)
        prob = self.generator(out[:, -1])
        _, next_word = torch.max(prob, dim=1)
        ys = torch.cat([ys, next_word.unsqueeze(0)], dim=1)

    return ys