如何平衡多输出模型的损失权重？分类与回归训练优化

我们有一个双输出模型，分类分支采用Categorical Cross Entropy(CCE)损失，回归分支采用MSE损失，训练时取两者损失均值。需要找到合适的损失权重策略，让两个分支能平等受益于梯度下降，避免先优化完一个分支后，再提升另一个分支权重时前者偏离理想区间。

初始训练损失与目标范围 #

• 初始损失值：
  • CCE = 0.6
  • MSE = 5e-3
• 理想损失范围：
  • CCE < 0.1
  • MSE < 5e-7

可行的损失权重策略 #

1. 基于损失比例的静态权重初始化

先计算初始损失的比值，让两个损失对总损失的初始贡献相等。初始时CCE与MSE的比值为 0.6 / 5e-3 = 120，因此设置权重：

• classification 权重 = 1
• regression 权重 = 120
  这样初始时两个损失项对总损失的贡献均为0.6（1*0.6 和 120*5e-3=0.6），保证梯度下降初期两个分支的影响力一致。

2. 动态权重调整（滑动窗口更新）

训练过程中定期计算当前分支损失，实时调整权重，维持两者加权损失的平衡。比如每N个epoch后：

1. 记录当前的CCE_loss和MSE_loss
2. 计算权重比例：weight_regression = CCE_loss / MSE_loss
3. 平滑更新loss_weights，保持classification_weight * CCE_loss ≈ regression_weight * MSE_loss
   这种方法能适应训练中损失的动态变化，避免某一分支损失主导总损失。

3. 损失标准化

对每个分支的损失做标准化处理，统一量级后再取均值。比如：

• CCE损失标准化：当前批次CCE损失 / 初始CCE值（0.6）
• MSE损失标准化：当前批次MSE损失 / 初始MSE值（5e-3）
  总损失取两个标准化损失的均值，天然保证分支间损失量级一致，无需手动调权重。

代码示例 #

静态权重初始化版本

model.compile(
    # 补充optimizer、metrics等其他参数
    loss={
        'classification': classification_loss,
        'regression': regression_loss,
    },
    loss_weights={
        'classification': 1,
        'regression': 120  # 0.6 / 5e-3 = 120
    }
)

动态权重调整（Keras回调实现）

from tensorflow.keras.callbacks import Callback

class DynamicLossWeightCallback(Callback):
    def __init__(self, alpha=0.1):
        self.alpha = alpha  # 平滑系数，避免权重突变

    def on_epoch_end(self, epoch, logs=None):
        current_cce = logs['classification_loss']
        current_mse = logs['regression_loss']
        # 计算目标权重比例
        target_weight_ratio = current_cce / current_mse
        # 平滑更新权重
        new_reg_weight = (1 - self.alpha) * self.model.loss_weights['regression'] + self.alpha * target_weight_ratio
        self.model.loss_weights = {
            'classification': 1,
            'regression': new_reg_weight
        }
        print(f"Updated loss weights: classification=1, regression={new_reg_weight:.2f}")

# 编译模型时初始权重设为1:120
model.compile(
    # 补充optimizer、metrics等其他参数
    loss={
        'classification': classification_loss,
        'regression': regression_loss,
    },
    loss_weights={
        'classification': 1,
        'regression': 120
    }
)

# 训练时加入回调
model.fit(
    # 补充训练数据、epochs等参数
    callbacks=[DynamicLossWeightCallback(alpha=0.1)]
)

内容的提问来源于stack exchange，提问作者Darren Rahnemoon