Transformer架构中b,t,c向量多头注意力机制的高效实现与优化

1次阅读

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

在 Transformer 模型中，多头注意力机制是其核心组件之一。然而，传统的注意力计算在长序列场景下存在 O(t^2) 的复杂度问题，这会导致显存占用急剧增加和计算效率下降。具体来说，当序列长度 t 增加时，注意力矩阵的大小会呈平方级增长，这对 GPU 显存提出了极高的要求。

为了优化多头注意力的计算，我们可以从 batch(b)、sequence(t)、channel(c) 三个维度进行张量分割。这种分割策略能够有效减少单次计算的数据量，从而降低显存占用。

batch 维度分割 ：将大的 batch 分成多个小的子 batch，逐个处理。
sequence 维度分割 ：将长序列分成多个短的子序列，分别计算注意力。
channel 维度分割 ：将通道分成多个子通道，并行计算。

以下是一个分块计算的 PyTorch 实现示例，重点展示了如何通过分块减少显存占用：

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

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super(MultiHeadAttention, self).__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads

        self.w_q = nn.Linear(d_model, d_model)
        self.w_k = nn.Linear(d_model, d_model)
        self.w_v = nn.Linear(d_model, d_model)
        self.w_o = nn.Linear(d_model, d_model)

    def forward(self, x, mask=None):
        batch_size, seq_len, d_model = x.size()

        # Split into multiple heads
        q = self.w_q(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        k = self.w_k(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        v = self.w_v(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)

        # Compute attention scores
        scores = torch.matmul(q, k.transpose(-2, -1)) / (self.d_k ** 0.5)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)
        attn = F.softmax(scores, dim=-1)

        # Apply attention to values
        output = torch.matmul(attn, v)
        output = output.transpose(1, 2).contiguous().view(batch_size, seq_len, d_model)
        return self.w_o(output)

内存复用技术通过重复使用已分配的显存空间来减少显存占用。具体实现包括：