深入解析AM-Softmax损失函数：从原理到PyTorch实战

1次阅读

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

在传统的分类任务中，Softmax 损失函数被广泛应用。然而，当我们需要学习具有判别性的特征（如人脸识别）时，传统的 Softmax 存在一些明显的局限性：

类内方差较大：同一类别的样本在特征空间中的分布可能比较分散
类间区分度不足：不同类别的决策边界不够明确
对角度度量不敏感：Softmax 主要关注分类正确性，而非特征的方向性

这些问题导致了在需要高判别性特征的任务中，传统 Softmax 的表现不尽如人意。AM-Softmax（Additive Margin Softmax）正是为了解决这些问题而提出的改进方案。

AM-Softmax 的核心思想是在角度空间中引入一个边际（margin），使得同类样本更加紧凑，不同类样本更加分离。其损失函数公式如下：

$$
L=-\frac{1}{N}\sum_{i}\log\frac{e^{s(\cos\theta_{y_i,i}-m)}}{e^{s(\cos\theta_{y_i,i}-m)}+\sum_{j\neq y_i}e^{s\cos\theta_{j,i}}}
$$

其中：

$s$ 是缩放因子（scale），控制决策边界的 ” 硬度 ”
$m$ 是边际（margin），控制类间距离
$\theta_{y_i,i}$ 是特征向量与对应类别权重向量的夹角

这个公式的关键改进在于在正类余弦值上减去一个边际 $m$，这使得模型需要学习更具判别性的特征才能达到相同的损失值。

下面是一个完整的 AM-Softmax PyTorch 实现，包含详细的注释：

import torch
import torch.nn as nn
import torch.nn.functional as F

class AMSoftmax(nn.Module):
    def __init__(self, feat_dim, num_classes, s=30.0, m=0.4):
        """
        AM-Softmax 损失函数实现
        :param feat_dim: 特征维度
        :param num_classes: 类别数量
        :param s: 缩放因子
        :param m: 边际值
        """
        super(AMSoftmax, self).__init__()
        self.feat_dim = feat_dim
        self.num_classes = num_classes
        self.s = s
        self.m = m
        self.W = nn.Parameter(torch.Tensor(num_classes, feat_dim))
        nn.init.xavier_normal_(self.W)  # 权重初始化

    def forward(self, feats, labels):
        """
        前向计算
        :param feats: 输入特征 [batch_size, feat_dim]
        :param labels: 真实标签 [batch_size]
        :return: 损失值
        """
        # 归一化特征和权重
        feats_norm = F.normalize(feats, p=2, dim=1)
        W_norm = F.normalize(self.W, p=2, dim=1)

        # 计算余弦相似度 [batch_size, num_classes]
        cos_theta = torch.mm(feats_norm, W_norm.t())

        # 获取对应类别的余弦值
        cos_theta_yi = cos_theta[torch.arange(cos_theta.size(0)), labels].view(-1, 1)

        # 应用边际（减去 m）cos_theta_yi_m = cos_theta_yi - self.m

        # 构建新的 logits
        new_logits = torch.cat((cos_theta_yi_m, cos_theta), dim=1)
        new_logits = new_logits * self.s

        # 计算交叉熵损失
        labels_onehot = torch.zeros_like(cos_theta)
        labels_onehot.scatter_(1, labels.view(-1, 1), 1.0)
        labels_onehot = torch.cat((labels_onehot, torch.zeros(labels.size(0), 1).to(feats.device)), dim=1)

        loss = F.cross_entropy(new_logits, torch.argmax(labels_onehot, dim=1))

        return loss

我们在 MNIST 数据集上设计了对照实验，比较了不同损失函数的表现：