Bi-LSTM与注意力机制实战：基于知识蒸馏的模型轻量化入门指南

1次阅读

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

在自然语言处理任务中，Bi-LSTM 因其能捕获双向上下文信息而被广泛使用。但当我们需要将模型部署到移动设备或嵌入式系统时，传统 Bi-LSTM 模型面临几个现实问题：

参数量庞大：一个标准的双层 Bi-LSTM 模型（hidden_size=256）参数量可达 3MB 以上
内存占用高：推理时需要同时存储前向和后向的隐藏状态
计算延迟明显：序列较长时，串行计算特性导致实时性差

这些限制使得原始模型难以在资源受限环境下落地，而知识蒸馏技术为我们提供了可行的优化路径。

模型类型	参数量	推理速度(ms)	准确率(%)
Bi-LSTM	3.2M	47	88.1
Bi-LSTM+Attention	3.4M	52	89.7
蒸馏后学生模型	0.9M	15	88.9

测试环境：Intel i7-9750H CPU，序列长度 128

首先实现带注意力机制的 Bi-LSTM 教师模型，关键代码如下：

import torch
import torch.nn as nn

class BiLSTMAttn(nn.Module):
    def __init__(self, vocab_size, embed_dim, hidden_dim, num_classes):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.lstm = nn.LSTM(embed_dim, hidden_dim, 
                           bidirectional=True, batch_first=True)
        # 注意力权重计算层
        self.attn = nn.Linear(2*hidden_dim, 1)
        self.fc = nn.Linear(2*hidden_dim, num_classes)

    def forward(self, x):
        embedded = self.embedding(x)  # (batch, seq, embed_dim)
        outputs, _ = self.lstm(embedded)  # (batch, seq, 2*hidden_dim)

        # 计算注意力得分
        attn_scores = self.attn(outputs)  # (batch, seq, 1)
        attn_weights = torch.softmax(attn_scores, dim=1)

        # 加权求和得到上下文向量
        context = torch.sum(attn_weights * outputs, dim=1)
        return self.fc(context)

知识蒸馏的核心是通过 KL 散度让学生模型学习教师模型的 ” 软标签 ”：

def distillation_loss(student_logits, teacher_logits, labels, 
                     temp=5.0, alpha=0.7):
    """
    student_logits: 学生模型输出
    teacher_logits: 教师模型输出
    labels: 真实标签
    temp: 温度参数
    alpha: 蒸馏损失权重
    """
    # 计算常规交叉熵损失
    ce_loss = nn.CrossEntropyLoss()(student_logits, labels)

    # 计算蒸馏损失（带温度系数的 KL 散度）soft_teacher = nn.functional.softmax(teacher_logits/temp, dim=-1)
    soft_student = nn.functional.log_softmax(student_logits/temp, dim=-1)
    kld_loss = nn.KLDivLoss(reduction='batchmean')(soft_student, soft_teacher)

    # 组合损失
    return alpha * (temp**2) * kld_loss + (1-alpha) * ce_loss

在 SST- 2 情感分析任务上的实验结果：