TensorFlow在多任务学习中的应用

TensorFlow 多任务学习实战指南

一 核心概念与适用场景

• 多任务学习（MTL）通过让多个相关任务共享一部分网络参数，同时各自拥有任务特定的“头（Head）”，在数据效率、泛化能力与正则化方面通常优于单任务方案，广泛应用于计算机视觉、自然语言处理与医疗健康等场景。典型做法是“共享骨干 + 任务头”，如同时做人脸关键点定位、属性分类等联合任务。

二 建模与训练要点

• 损失合并与尺度对齐：直接相加多个任务的损失会被尺度更大的任务主导。可用不确定性加权（learnable noise parameter）自动对齐不同任务损失的尺度，避免手动调权重；也可采用损失加权求和作为简单基线。
• 学习率分治：共享层与不同任务头的学习速率往往不同。可在 TensorFlow 中按变量集合拆分梯度，分别用不同的 Adam 学习率更新共享参数与各任务头参数，提高收敛稳定性与效果。
• 梯度阻断与“估计作为特征”：当把任务 A 的输出作为任务 B 的输入时，若不希望 B 的梯度回传到 A，可用 tf.stop_gradient 阻断反向传播，稳定训练与解耦任务。
• 训练范式选择：常见有联合训练（所有任务一起优化）与交替训练（按批次或轮次在不同任务间切换）；前者实现简单、收敛稳定，后者在某些任务差异较大时更易平衡。

三 代码模板 Keras 版

• 使用不确定性加权自动平衡多任务损失，并支持为共享层与任务头设置不同学习率（TensorFlow 2.x / Keras）。

import tensorflow as tf
from tensorflow.keras.layers import Input, Dense, Lambda
from tensorflow.keras.models import Model
import tensorflow_probability as tfp

def build_mtl_model(input_dim, task1_dim, task2_dim,
                    shared_units=[128, 64],
                    lr_shared=1e-3, lr_head=1e-3):
    x = Input(shape=(input_dim,))
    # 共享骨干
    z = x
    for u in shared_units:
        z = Dense(u, activation='relu')(z)

    # 任务1头（分类示例）
    logits1 = Dense(task1_dim)(z)
    # 任务2头（回归示例）
    mu2    = Dense(task2_dim)(z)
    # 可学习的噪声参数（标量，任务2）
    log_sigma2 = Dense(1, activation='softplus')(z)  # > 0
    sigma2 = Lambda(lambda s: 1e-6 + s, name='sigma2')(log_sigma2)

    # 任务1损失：稀疏交叉熵
    loss1 = tf.reduce_mean(
        tf.nn.sparse_softmax_cross_entropy_with_logits(labels=y1, logits=logits1))

    # 任务2损失：异方差高斯负对数似然
    sq_err = tf.reduce_sum((y2 - mu2) ** 2, axis=-1)
    nll2  = 0.5 * (sq_err / sigma2 + tf.math.log(2 * np.pi * sigma2))
    loss2 = tf.reduce_mean(nll2)

    # 总损失 = 任务1损失 + 任务2损失（不确定性已内嵌到 loss2 的 sigma2）
    total_loss = loss1 + loss2

    model = Model(inputs=x, outputs=[logits1, mu2, sigma2])
    # 优化器：共享与任务头使用不同学习率
    opt_shared = tf.keras.optimizers.Adam(lr_shared)
    opt_heads  = tf.keras.optimizers.Adam(lr_head)
    # 变量分组
    shared_vars = model.layers[1].trainable_variables  # 共享层
    head1_vars  = model.layers[-4].trainable_variables  # 任务1头
    head2_vars  = (model.layers[-3].trainable_variables +
                  model.layers[-2].trainable_variables)  # 任务2头（mu + log_sigma2）
    # 分别最小化
    train_step = tf.function(lambda: [
        opt_shared.minimize(lambda: loss1, var_list=shared_vars),
        opt_heads.minimize(lambda: loss2, var_list=head1_vars + head2_vars)
    ])
    return model, train_step

# 使用示例（伪代码）
# model, train_step = build_mtl_model(input_dim=100, task1_dim=10, task2_dim=1)
# for xb, y1b, y2b in dataset:
#     loss = train_step(xb, y1b, y2b)

• 要点说明：
  • 任务1使用分类损失，任务2使用高斯 NLL（含可学习的噪声方差），两者直接相加即可，无需额外权重超参。
  • 通过分组优化器对共享层与任务头设置不同学习率，提升稳定性与效果。
  • 若需将任务 A 的输出作为任务 B 的输入且不回传梯度，可在连接处包裹 tf.stop_gradient。

四 数据流水线与实践建议

• 使用 tf.data 与 tf.image 构建高效输入流水线：并行读取、解析、增强、批处理与预取；在 tf.estimator 时代也可结合 Head API 将共享特征提取器与多个任务头解耦，便于快速搭建与调试 MTL 模型。
• 训练监控：为每个任务分别记录损失与指标（如分类准确率、回归 MAE），并在 TensorBoard 中对比收敛速度与尺度，必要时调整任务权重或学习率分组。