基于Bi-LSTM与多头自注意力机制的文本分类实战与性能优化

1次阅读

没有评论

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

在自然语言处理任务中，循环神经网络（RNN）曾经是处理序列数据的首选架构。但随着研究的深入和实际应用的验证，传统 RNN 在处理文本分类任务时暴露出几个明显缺陷：

梯度消失 / 爆炸问题：当处理长文本时，RNN 难以有效捕捉远距离的依赖关系，这导致模型难以理解那些需要长期记忆支持的语义信息。
单向信息流限制：标准 RNN 只能单向（通常是从左到右）处理文本，无法同时利用上下文信息，这在很多需要双向理解的场景中表现不佳。
固定长度上下文窗口：传统 RNN 的隐藏状态实际上形成了一个固定大小的上下文表示，这对于长度变化大的文本处理不够灵活。

为了克服传统 RNN 的局限性，研究者们提出了两种主要的改进方案：双向长短期记忆网络（Bi-LSTM）和 Transformer 架构。让我们先看看它们各自的优缺点：

Bi-LSTM 的优势：
– 通过双向处理能捕捉前后文信息
– LSTM 单元设计缓解了梯度消失问题
– 对序列位置信息敏感，适合处理有序数据

Bi-LSTM 的不足：
– 仍然存在一定程度的长期依赖问题
– 计算无法并行，训练速度较慢
– 对全局关系的建模能力有限

Transformer 的优势：
– 多头注意力机制能捕捉全局依赖关系
– 高度并行化的计算结构
– 对长距离关系建模能力强

Transformer 的不足：
– 对位置信息依赖显式编码
– 在小数据集上容易过拟合
– 计算复杂度随序列长度平方增长

结合 Bi-LSTM 和 Transformer 的优势，我们设计了一个混合架构，其核心思想是：

使用 Bi-LSTM 作为底层特征提取器，捕获序列的局部和顺序特征
在 Bi-LSTM 之上应用多头自注意力机制，建模全局依赖关系
通过残差连接保持信息的流动

关键设计细节：

注意力权重计算：
对于 Bi-LSTM 输出的隐藏状态序列 H∈ℝ^(batch×seq_len×hidden_dim)，我们计算查询 Q、键 K 和值 V：

Q = HW_Q, K = HW_K, V = HW_V

其中 W_Q, W_K, W_V 是可学习参数矩阵。注意力得分计算为：

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

信息流整合：
将多头注意力的输出与原始 Bi-LSTM 输出进行残差连接，然后通过层归一化：

Output = LayerNorm(H + MultiHeadAttention(H))

以下是基于 PyTorch 的实现代码，包含数据预处理、模型定义和训练循环：

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torch.nn.utils.rnn import pad_sequence

# 数据预处理
class TextDataset(Dataset):
    def __init__(self, texts, labels, tokenizer, max_len):
        self.texts = texts
        self.labels = labels
        self.tokenizer = tokenizer
        self.max_len = max_len

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        text = self.texts[idx]
        label = self.labels[idx]
        tokens = self.tokenizer(text)[:self.max_len]
        return torch.tensor(tokens, dtype=torch.long), torch.tensor(label, dtype=torch.long)

# 模型定义
class BiLSTMAttention(nn.Module):
    def __init__(self, vocab_size, embed_dim, hidden_dim, num_heads, num_layers, num_classes):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.bilstm = nn.LSTM(embed_dim, hidden_dim//2, 
                             num_layers=num_layers, 
                             bidirectional=True, 
                             batch_first=True)

        # 多头注意力层
        self.attention = nn.MultiheadAttention(hidden_dim, num_heads)
        self.fc = nn.Linear(hidden_dim, num_classes)
        self.layer_norm = nn.LayerNorm(hidden_dim)

    def forward(self, x):
        # 输入 x 维度: (batch_size, seq_len)
        embedded = self.embedding(x)  # (batch_size, seq_len, embed_dim)

        # Bi-LSTM 层
        lstm_out, _ = self.bilstm(embedded)  # (batch_size, seq_len, hidden_dim)

        # 调整维度用于多头注意力
        lstm_out = lstm_out.transpose(0, 1)  # (seq_len, batch_size, hidden_dim)

        # 多头注意力层
        attn_out, _ = self.attention(lstm_out, lstm_out, lstm_out)

        # 残差连接和层归一化
        out = self.layer_norm(lstm_out + attn_out)
        out = out.transpose(0, 1)  # 恢复维度 (batch_size, seq_len, hidden_dim)

        # 取最后一个时间步作为分类特征
        out = out[:, -1, :]
        return self.fc(out)

# 训练循环
def train(model, dataloader, criterion, optimizer, device):
    model.train()
    total_loss = 0

    for batch in dataloader:
        inputs, labels = batch
        inputs, labels = inputs.to(device), labels.to(device)

        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()

        # 梯度裁剪
        nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

        optimizer.step()
        total_loss += loss.item()

    return total_loss / len(dataloader)

我们在 AG News 数据集上进行了对比实验，结果如下：