Bi-LSTM与多头自注意力机制入门指南：从理论到代码实现

1次阅读

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

在处理自然语言处理（NLP）任务时，传统的单向 RNN（如 LSTM）存在一个主要问题：梯度消失。尤其是当处理长文本时，单向 RNN 只能捕捉到从左到右的上下文信息，而忽略了从右到左的上下文。这在许多任务中（如机器翻译、文本分类）会导致性能下降，因为双向上下文信息往往对理解句子含义至关重要。

Bi-LSTM（双向长短期记忆网络）通过同时处理正向和反向的序列数据，能够更好地捕捉上下文信息。然而，Bi-LSTM 仍然存在计算效率低和难以并行化的问题。为了解决这些问题，多头自注意力机制（Multi-Head Self-Attention）应运而生，它能够高效地捕捉长距离依赖关系，并且天然支持并行计算。

模型	计算复杂度	并行性	位置感知
普通 LSTM	O(n)	低	隐式
Bi-LSTM	O(2n)	低	隐式
Transformer	O(n^2)	高	显式（位置编码）

从表格中可以看出，Bi-LSTM 在计算复杂度和并行性上不如 Transformer，但在某些任务中，Bi-LSTM 的表现仍然优于 Transformer，尤其是在数据量较小的情况下。

首先，我们需要安装必要的库：

pip install torch numpy matplotlib

接下来，我们实现一个结合 Bi-LSTM 和多头自注意力的模型：

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

class BiLSTMMultiHeadAttention(nn.Module):
    def __init__(self, vocab_size, embedding_dim, hidden_dim, num_heads):
        super(BiLSTMMultiHeadAttention, self).__init__()
        self.embedding = nn.Embedding(vocab_size, embedding_dim)
        self.bilstm = nn.LSTM(embedding_dim, hidden_dim, bidirectional=True, batch_first=True)
        self.multihead_attn = nn.MultiheadAttention(hidden_dim * 2, num_heads)
        self.fc = nn.Linear(hidden_dim * 2, 1)

    def forward(self, x, mask=None):
        # x shape: (batch_size, seq_len)
        embedded = self.embedding(x)  # (batch_size, seq_len, embedding_dim)
        lstm_out, _ = self.bilstm(embedded)  # (batch_size, seq_len, hidden_dim * 2)

        # Transpose for MultiheadAttention
        lstm_out = lstm_out.transpose(0, 1)  # (seq_len, batch_size, hidden_dim * 2)
        attn_output, attn_weights = self.multihead_attn(lstm_out, lstm_out, lstm_out, key_padding_mask=mask)
        attn_output = attn_output.transpose(0, 1)  # (batch_size, seq_len, hidden_dim * 2)

        # Global average pooling
        pooled = attn_output.mean(dim=1)  # (batch_size, hidden_dim * 2)
        output = self.fc(pooled)  # (batch_size, 1)
        return output.squeeze(), attn_weights

序列 padding 与 mask 处理 ：在 NLP 任务中，输入序列通常是变长的，因此需要进行 padding。我们可以使用torch.nn.utils.rnn.pad_sequence 来处理变长序列，并生成对应的 mask。

from torch.nn.utils.rnn import pad_sequence

# Example sequences
sequences = [torch.tensor([1, 2, 3]), torch.tensor([4, 5]), torch.tensor([6])]
padded_sequences = pad_sequence(sequences, batch_first=True, padding_value=0)

# Generate mask (0 for padding, 1 for actual tokens)
mask = (padded_sequences != 0).bool()

注意力权重的归一化处理：多头自注意力机制中的注意力权重是通过 softmax 函数进行归一化的，确保每个位置的权重之和为 1。
双向状态向量的拼接方式 ：Bi-LSTM 的输出是正向和反向 LSTM 的拼接，因此hidden_dim 需要乘以 2。

我们可以使用 matplotlib 来可视化不同 head 的注意力热力图：

import matplotlib.pyplot as plt
import numpy as np

def plot_attention_weights(attn_weights, sentence, head_idx=0):
    # attn_weights shape: (num_heads, seq_len, seq_len)
    plt.figure(figsize=(10, 10))
    plt.imshow(attn_weights[head_idx].detach().numpy(), cmap='hot', interpolation='nearest')
    plt.xticks(np.arange(len(sentence)), labels=sentence)
    plt.yticks(np.arange(len(sentence)), labels=sentence)
    plt.colorbar()
    plt.show()

# Example usage
sentence = ["I", "love", "NLP"]
model = BiLSTMMultiHeadAttention(vocab_size=100, embedding_dim=50, hidden_dim=64, num_heads=4)
output, attn_weights = model(torch.tensor([[1, 2, 3]]))
plot_attention_weights(attn_weights, sentence, head_idx=0)

当序列长度 >512 时的分块处理策略：对于长序列，可以将序列分成多个块分别处理，然后合并结果。
在 Tesla T4 显卡上的 batch_size 调优经验：Tesla T4 的显存为 16GB，建议 batch_size 设置为 32 或 64，具体取决于模型大小和序列长度。
注意力矩阵的内存占用计算公式：注意力矩阵的内存占用为(batch_size * num_heads * seq_len * seq_len * 4) bytes（假设 float32 精度）。