分类任务使用MAE损失函数遇无梯度提供错误的疑问

我在分类任务里想把回归常用的Mean Absolute Error（MAE）当损失函数用。已知y_pred和y_true是独热编码格式，但MAE需要实数输入。

第一种自定义损失函数实现：

def AgeAccuracyRegularity(y_pred,y_true):
    
    mae_func = tf.keras.losses.MeanAbsoluteError()

    y_pred_ages = K.cast(K.argmax(y_pred, axis=-1)+1,dtype='float32')
    y_true_ages = K.cast(K.argmax(y_true, axis=-1)+1,dtype='float32')

    res = mae_func(y_pred_ages, y_true_ages)
    
    return res

这段代码会抛出ValueError: No gradients provided for any variable错误；但加了无意义计算后的实现却能正常运行，且两种方式输出的尺寸和类型一致：

def AgeAccuracyRegularity(y_pred,y_true):
    
    mae_func = tf.keras.losses.MeanAbsoluteError()

    y_pred_ages = K.cast(K.argmax(y_pred, axis=-1)+1,dtype='float32')
    y_true_ages = K.cast(K.argmax(y_true, axis=-1)+1,dtype='float32')

    res = mae_func(y_pred_ages, y_true_ages)

    mae = mae_func(y_pred, y_true)

    return res-mae+mae

对此现象存在疑问。

• 第一个版本报错的核心：K.argmax是不可导操作。当你用它把独热编码的y_pred转成离散的类别索引后，直接切断了模型参数到损失值的梯度传播路径——离散索引和原始预测概率之间没有连续的导数关系，TensorFlow没法计算任何参数的梯度，自然抛出无梯度的错误。

• 第二个版本能运行的本质：看似无意义的res-mae+mae其实补了梯度链路。这里的mae = mae_func(y_pred, y_true)是直接对原始独热编码计算MAE，这个操作是可导的。TensorFlow构建计算图时会保留这条可导路径，哪怕最终返回值和res完全一样，梯度传播的链路已经被补上了。

• 隐藏的坑：这种写法只是绕过了梯度检查，但实际优化逻辑完全偏离你的需求——模型的梯度实际来自独热编码的MAE，而不是你想要的基于年龄数值的MAE。你的目标损失res因为argmax的不可导性，根本没参与参数更新，相当于模型在优化一个和你需求无关的损失函数。

要让模型基于年龄数值的MAE优化，得避开argmax这种不可导操作，改用概率加权计算年龄期望：

def AgeAccuracyRegularity(y_pred,y_true):
    # 生成年龄权重：假设类别索引0对应年龄1，依次递增
    age_weights = K.cast(K.arange(1, y_pred.shape[-1]+1), dtype='float32')
    # 用预测概率加权求和得到期望年龄（可导操作）
    y_pred_ages = K.sum(y_pred * age_weights, axis=-1)
    # 真实年龄同理（独热编码加权求和直接得到对应真实年龄）
    y_true_ages = K.sum(y_true * age_weights, axis=-1)
    
    mae_func = tf.keras.losses.MeanAbsoluteError()
    return mae_func(y_pred_ages, y_true_ages)

这种方式完整保留了梯度传播链路，损失值能正确指导模型更新，完全符合你的需求。

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