AI多任务学习踩坑记录

深度学习 多任务学习 PyTorch 模型优化
作者: 醉月思 发布日期: 2022-03-22 分类: 技术实践

如果只能用一句话说AI多任务学习：先把失败复现出来。

从单任务说起

先说说原来的单任务方案。情感分类用的是个标准的 BERT + Linear 层：

import torch
import torch.nn as nn
from transformers import BertModel, BertTokenizer

class SentimentClassifier(nn.Module):
    def __init__(self, num_labels=3):
        super().__init__()
        self.bert = BertModel.from_pretrained('bert-base-chinese')
        self.dropout = nn.Dropout(0.1)
        self.classifier = nn.Linear(self.bert.config.hidden_size, num_labels)

    def forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        pooled_output = outputs.pooler_output
        pooled_output = self.dropout(pooled_output)
        logits = self.classifier(pooled_output)
        return logits

这个模型训了三天，准确率 85%，上线后推理延迟 120ms，还算能接受。但加上主题识别和负面评论检测后，三个模型加起来延迟就到 350ms 了，客户开始投诉系统"卡顿"。

于是决定做多任务。一开始想得很简单，不就是多加几个分类头吗？

第一次尝试：朴素多任务

把三个任务揉到一个模型里，共享 BERT encoder，每个任务一个独立头：

class MultiTaskBERT(nn.Module):
    def __init__(self, num_sentiment=3, num_topic=10, num_negative=2):
        super().__init__()
        self.bert = BertModel.from_pretrained('bert-base-chinese')
        self.dropout = nn.Dropout(0.1)

        # 三个任务头
        self.sentiment_head = nn.Linear(self.bert.config.hidden_size, num_sentiment)
        self.topic_head = nn.Linear(self.bert.config.hidden_size, num_topic)
        self.negative_head = nn.Linear(self.bert.config.hidden_size, num_negative)

    def forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        pooled_output = outputs.pooler_output
        pooled_output = self.dropout(pooled_output)

        return {
            'sentiment': self.sentiment_head(pooled_output),
            'topic': self.topic_head(pooled_output),
            'negative': self.negative_head(pooled_output)
        }

损失函数也简单粗暴，三个 loss 直接相加：

criterion = nn.CrossEntropyLoss()

def compute_loss(predictions, labels):
    loss_sentiment = criterion(predictions['sentiment'], labels['sentiment'])
    loss_topic = criterion(predictions['topic'], labels['topic'])
    loss_negative = criterion(predictions['negative'], labels['negative'])

    return loss_sentiment + loss_topic + loss_negative

训练开始后，第一个 epoch 就出问题了。训练 loss 一直降，但验证集上的三个任务指标都惨不忍睹。情感分类准确率只有 60%，比单任务时掉了 25 个百分点。

问题定位：任务不平衡

打印训练日志后发现，三个任务的 loss 量级差异巨大。主题分类因为类别多（10 个），loss 经常在 2.0 以上；而负面评论检测只有 2 个类别，loss 只有 0.3 左右。模型被主题任务牵着走，另外两个任务基本没学到东西。

这时候才意识到，多任务学习的第一个坑就是任务平衡问题。

方案一：手动调权重

最直接的想法是给每个任务加个权重，手动平衡一下：

def compute_loss(predictions, labels, weights):
    loss_sentiment = criterion(predictions['sentiment'], labels['sentiment'])
    loss_topic = criterion(predictions['topic'], labels['topic'])
    loss_negative = criterion(predictions['negative'], labels['negative'])

    total_loss = (weights['sentiment'] * loss_sentiment +
                  weights['topic'] * loss_topic +
                  weights['negative'] * loss_negative)
    return total_loss

weights = {'sentiment': 1.0, 'topic': 0.3, 'negative': 2.0}

试了几轮调参，把情感和负面检测的权重调高，主题权重调低，三个任务的 loss 大概能走到一个水平线上。但准确率还是不稳定，有时候调完一组权重，情感上去了，主题又掉下来。

这种手动调参的方式就像是在玩一个永远玩不完的平衡游戏，而且每次数据分布一变又要重新调。

方案二：Uncertainty Weighting

后来查阅文献，找到了 Kendall 等人提出的 Uncertainty Weighting 方法。它的思想是让模型自己学习每个任务的权重，基于任务的不确定性来调整 loss 权重：

class UncertaintyWeightedLoss(nn.Module):
    def __init__(self, num_tasks):
        super().__init__()
        # 每个任务学一个 log σ^2 参数
        self.log_vars = nn.Parameter(torch.zeros(num_tasks))

    def forward(self, losses):
        weighted_losses = []
        for i, loss in enumerate(losses):
            # precision = 1 / σ^2
            precision = torch.exp(-self.log_vars[i])
            # loss / (2σ^2) + log σ
            weighted_loss = precision * loss + 0.5 * self.log_vars[i]
            weighted_losses.append(weighted_loss)
        return torch.stack(weighted_losses).sum()

# 使用方式
uncertainty_loss = UncertaintyWeightedLoss(num_tasks=3)

def compute_loss(predictions, labels):
    loss_sentiment = criterion(predictions['sentiment'], labels['sentiment'])
    loss_topic = criterion(predictions['topic'], labels['topic'])
    loss_negative = criterion(predictions['negative'], labels['negative'])

    return uncertainty_loss([loss_sentiment, loss_topic, loss_negative])

这个方法省去了手动调参的痛苦，模型会根据每个任务的难度自动调整权重。训练过程中，如果某个任务 loss 震荡比较大，对应的权重就会变小，避免主导训练。

实际使用后，三个任务的准确率都有提升，情感分类回到了 80% 左右，主题识别和负面检测也达到了预期水平。

第二个问题：负迁移

解决完平衡问题后，又碰到了新问题。有时候明明三个任务单训都很好，合在一起后某个任务反而变差了。这种现象叫负迁移，任务之间互相干扰。

排查数据后发现，有些样本的标签本身就矛盾。比如同一条评论，情感是正面，但被标记为负面评论，这种不一致数据会让模型学到冲突的表征。

解决方式也比较简单粗暴：数据清洗。写了个脚本检查标签的一致性，把明显矛盾的样本删掉或重新标注：

def clean_dataset(dataset):
    cleaned = []
    for item in dataset:
        text = item['text']
        sentiment = item['sentiment_label']
        negative = item['negative_label']

        # 如果情感是正面（0）但标记为负面（1），可能是标注错误
        if sentiment == 0 and negative == 1:
            # 简单策略：优先保留情感标签，重置负面标签
            item['negative_label'] = 0

        cleaned.append(item)
    return cleaned

数据质量上来后，负迁移问题缓解了不少。

模型架构优化

虽然基本能用了，但还想再压榨一下性能。开始尝试不同的模型架构。

多头注意力机制

原来的共享 encoder + 独立头虽然简单，但可能不够灵活。尝试给每个任务加一个轻量级的注意力层：

class TaskSpecificAttention(nn.Module):
    def __init__(self, hidden_size, num_heads=4):
        super().__init__()
        self.attention = nn.MultiheadAttention(hidden_size, num_heads, batch_first=True)
        self.layer_norm = nn.LayerNorm(hidden_size)

    def forward(self, x):
        # x: (batch, seq_len, hidden_size)
        attn_output, _ = self.attention(x, x, x)
        x = self.layer_norm(x + attn_output)
        return x[:, 0, :]  # 取 [CLS] 位置

class AdvancedMultiTaskBERT(nn.Module):
    def __init__(self, num_sentiment=3, num_topic=10, num_negative=2):
        super().__init__()
        self.bert = BertModel.from_pretrained('bert-base-chinese')

        # 任务特定的注意力层
        self.sentiment_attn = TaskSpecificAttention(self.bert.config.hidden_size)
        self.topic_attn = TaskSpecificAttention(self.bert.config.hidden_size)
        self.negative_attn = TaskSpecificAttention(self.bert.config.hidden_size)

        # 分类头
        self.sentiment_head = nn.Linear(self.bert.config.hidden_size, num_sentiment)
        self.topic_head = nn.Linear(self.bert.config.hidden_size, num_topic)
        self.negative_head = nn.Linear(self.bert.config.hidden_size, num_negative)

    def forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        sequence_output = outputs.last_hidden_state

        # 每个任务有自己的注意力表征
        sentiment_repr = self.sentiment_attn(sequence_output)
        topic_repr = self.topic_attn(sequence_output)
        negative_repr = self.negative_attn(sequence_output)

        return {
            'sentiment': self.sentiment_head(sentiment_repr),
            'topic': self.topic_head(topic_repr),
            'negative': self.negative_head(negative_repr)
        }

这个架构能让每个任务学到自己关注的文本特征，而不是完全依赖共享的 [CLS] 表征。实验下来，准确率提升 2-3 个百分点，但计算开销也增加了约 15%。

权衡后还是决定用回原来的简单架构，毕竟线上环境对延迟很敏感。

训练策略优化

除了模型架构，训练策略也很关键。试了几种不同的训练顺序：

渐进式训练

先训一个主任务（情感分类），等主任务稳定后，再加第二个任务，最后加第三个任务：

def progressive_training(model, train_loader, epochs_per_stage=5):
    # Stage 1: 只训情感分类
    optimizer = torch.optim.Adam(model.parameters(), lr=2e-5)
    for epoch in range(epochs_per_stage):
        for batch in train_loader:
            optimizer.zero_grad()
            preds = model(batch['input_ids'], batch['attention_mask'])
            loss = criterion(preds['sentiment'], batch['sentiment_label'])
            loss.backward()
            optimizer.step()

    # Stage 2: 加入主题识别
    for epoch in range(epochs_per_stage):
        for batch in train_loader:
            optimizer.zero_grad()
            preds = model(batch['input_ids'], batch['attention_mask'])
            loss = uncertainty_loss([
                criterion(preds['sentiment'], batch['sentiment_label']),
                criterion(preds['topic'], batch['topic_label'])
            ])
            loss.backward()
            optimizer.step()

    # Stage 3: 所有任务一起训
    for epoch in range(epochs_per_stage):
        for batch in train_loader:
            optimizer.zero_grad()
            preds = model(batch['input_ids'], batch['attention_mask'])
            loss = uncertainty_loss([
                criterion(preds['sentiment'], batch['sentiment_label']),
                criterion(preds['topic'], batch['topic_label']),
                criterion(preds['negative'], batch['negative_label'])
            ])
            loss.backward()
            optimizer.step()

这种方式确实能缓解模型"一开始就被多个任务拉扯"的问题，但训练时间会变长。对于数据量大的场景不太划算。

动态任务采样

另一个思路是动态调整每个任务的采样概率，让困难任务有更多训练机会：

class DynamicTaskSampler:
    def __init__(self, tasks, initial_weights=None):
        self.tasks = tasks
        self.weights = initial_weights or {task: 1.0 for task in tasks}
        self.task_losses = {task: [] for task in tasks}
        self.update_interval = 100  # 每 100 个 batch 更新一次权重
        self.step = 0

    def update_weights(self):
        self.step += 1
        if self.step % self.update_interval != 0:
            return

        # 基于最近 loss 的方差调整权重
        for task in self.tasks:
            if len(self.task_losses[task]) > 0:
                recent_losses = self.task_losses[task][-self.update_interval:]
                loss_std = np.std(recent_losses)
                # loss 越不稳定，权重越大
                self.weights[task] = loss_std

        # 归一化权重
        total_weight = sum(self.weights.values())
        for task in self.tasks:
            self.weights[task] /= total_weight

    def get_task(self):
        self.update_weights()
        tasks = list(self.tasks.keys())
        weights = [self.weights[task] for task in tasks]
        return np.random.choice(tasks, p=weights)

    def record_loss(self, task, loss):
        self.task_losses[task].append(loss)

这个方式能一定程度上缓解任务不平衡，但实现起来比较复杂，而且调参成本不低。最后还是用了最简单的 Uncertainty Weighting，稳定省心。

上线效果

折腾了一圈，最终上线的架构还是最初的那个：共享 BERT encoder + 三个独立分类头 + Uncertainty Weighting loss。

线上表现：

  • 推理延迟从 350ms 降到 180ms，相比三个独立模型快了差不多一倍
  • 情感分类准确率 82%，比单任务时低 3 个百分点，但换来了两个额外任务
  • 主题识别准确率 78%，负面检测 F1 0.85，都达到了业务预期
  • 模型大小从 3 × 110MB 降到 1 × 110MB，内存占用也降了

客户没再投诉卡顿，产品经理也满意多任务的能力。算是折腾出了点结果。

踩坑总结

回想整个过程，几个最深的感受：

多任务学习不是简单的"把 loss 加起来就行"。任务平衡、负迁移、架构设计、训练策略，每一个环节都可能踩坑。

数据质量比模型架构重要。一开始把大量时间花在调模型结构上，后来发现清洗数据反而收益更大。

不要追求完美的多任务学习。有时候几个独立模型加个轻量级的特征共享层，比搞个复杂的多任务架构更实用。

单任务模型像是只会干一活的工具人，多任务学习更像是在培养一个能同时搞定多个活儿的全栈工程师。前者专精但成本高，后者灵活但需要更多调教。

但技术选择从来不是二选一，而是看场景、看约束、看投入产出比。就像这次，如果客户对延迟没那么敏感，可能直接就上三个独立模型了，省得折腾这些坑。

所谓经验，大抵就是这些踩坑爬坑的过程攒出来的。

版权声明: 本文首发于 指尖魔法屋-AI多任务学习踩坑记录（https://blog.thinkmoon.cn/post/220-ai-multitask-learning-from-single-to-multi-practice/） 转载或引用必须申明原指尖魔法屋来源及源地址！

相关阅读

根据标签、分类和主题为你挑选的相近文章