ArcFace损失函数公式详解与实战优化：从理论到高效实现

1次阅读

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

在人脸识别任务中，最大的挑战是 类别间相似度高 （不同人的面部特征可能很接近）和 类内差异大（同一个人在不同光照、角度下差异显著）。传统 Softmax 损失函数在处理这种细粒度分类时存在明显缺陷：

Softmax 仅追求分类正确，缺乏对特征判别性的显式约束
学到的特征在嵌入空间呈 ” 放射状 ” 分布，类间边界模糊
容易受到噪声干扰，测试时泛化能力不足

数学表达上，Softmax 损失可写为：
$$L_{softmax} = -\log\frac{e^{W_{y_i}^T x_i + b_{y_i}}}{\sum_{j=1}^n e^{W_j^T x_i + b_j}}$$
其中 $W_j$ 是全连接层权重，$x_i$ 是输入特征。这个公式没有显式优化特征间的角度关系。

决策边界由权重向量决定：
$$W_1^Tx + b_1 = W_2^Tx + b_2$$

引入余弦 margin：
$$L_{cos} = -\log\frac{e^{s(\cosθ_{y_i} – m)}}{e^{s(\cosθ_{y_i} – m)} + \sum_{j≠y_i} e^{s\cosθ_j}}$$

将 margin 添加到角度空间而非余弦空间：
$$L_{arc} = -\log\frac{e^{s\cos(θ_{y_i} + m)}}{e^{s\cos(θ_{y_i} + m)} + \sum_{j≠y_i} e^{s\cosθ_j}}$$

关键区别在于：
1. ArcFace 的 margin 是在角度 $θ$ 上直接叠加（angular margin）
2. 几何解释相当于在超球面上压缩类内分布
3. 决策边界变为：$\cos(θ_1 + m) = \cosθ_2$

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

class ArcFace(nn.Module):
    def __init__(self, feat_dim, num_classes, margin=0.3, scale=64):
        super().__init__()
        self.W = nn.Parameter(torch.Tensor(feat_dim, num_classes))
        nn.init.xavier_normal_(self.W)  # 权重初始化
        self.margin = margin  # angular margin 值
        self.scale = scale    # 特征缩放因子
        self.eps = 1e-7       # 防止数值不稳定

    def forward(self, x, labels):
        # 特征和权重归一化（关键步骤！）x_norm = F.normalize(x, p=2, dim=1)
        W_norm = F.normalize(self.W, p=2, dim=0)

        # 计算 cosθ（矩阵乘法）cos_theta = torch.mm(x_norm, W_norm)
        cos_theta = torch.clamp(cos_theta, -1+self.eps, 1-self.eps)

        # 计算 θ（反余弦）theta = torch.acos(cos_theta)

        # 仅对正确类别添加 margin
        one_hot = F.one_hot(labels, num_classes=self.W.shape[1]).float()
        cos_theta_m = torch.cos(theta + self.margin * one_hot)

        # 缩放后的 logits
        logits = self.scale * (one_hot * cos_theta_m + (1-one_hot) * cos_theta)
        loss = F.cross_entropy(logits, labels)

        return loss

关键实现细节：

双归一化：对特征 $x$ 和权重 $W$ 都进行 L2 归一化，确保计算的是纯角度关系
数值稳定：通过 clamp 避免 acos 计算出 NaN
margin 应用：只对 ground truth 类别的角度添加 margin
特征缩放：scale 参数控制特征在超球面上的 ” 半径 ”

ResNet 系列：建议最后一层特征维度≥512
MobileNet 等轻量网络：可适当降低到 256-384 维
过大维度（如 1024+）可能导致计算量激增而收益有限

margin 值	推荐学习率	适用场景
0.1-0.2	1e-3	小规模数据集
0.3-0.4	5e-4	中等规模（1M+）
0.5-0.6	1e-4	困难样本（遮挡等）

梯度裁剪 ：设置torch.nn.utils.clip_grad_norm_ 防止梯度爆炸
混合精度训练：使用 AMP 自动管理精度
类别采样：当类别数 >1M 时，可采用随机负采样

现象：loss 震荡不收敛
解决：检查 forward 中是否执行了F.normalize

现象：训练初期出现 NaN
解决：从 0.1 开始逐步增加，监控梯度幅度

现象：测试准确率异常低

解决：测试时使用纯余弦相似度：

# 测试代码示例
def test_forward(self, x):
    x_norm = F.normalize(x, p=2, dim=1)
    W_norm = F.normalize(self.W, p=2, dim=0)
    return torch.mm(x_norm, W_norm)  # 直接返回 cosθ

我们对比了不同 margin 值在验证集上的表现：