AMSoftmax（Additive Margin Softmax）是在人脸识别领域被广泛使用的度量学习损失函数，它通过在余弦相似度上叠加一个固定间隔，强迫模型把同类样本拉得更近、异类样本推得更远。不过不少人在实际训练中发现一个现象：训练集上的损失一路下降，验证集上的准确率却在若干个 epoch 之后开始回落，闭集测试还好，开集验证指标明显变差。这基本可以判定为过拟合。AMSoftmax 因为引入了间隔约束，对特征的判别性要求更高，一旦训练数据量不足或者类别分布不均衡，模型就会把训练样本的噪声细节也学进去，泛化能力随之下降。本文围绕正则化和数据增强两条主线，系统梳理应对 AMSoftmax 过拟合的实用方法。
为什么 AMSoftmax 更容易过拟合
先从原理层面理解问题，才能对症下药。普通 Softmax 交叉熵只需要类别可分即可，而 AMSoftmax 的损失形式为：
import torch
import torch.nn as nn
import torch.nn.functional as F
class AMSoftmaxLoss(nn.Module):
def __init__(self, embedding_dim, num_classes, scale=30.0, margin=0.35):
super().__init__()
self.scale = scale
self.margin = margin
self.weight = nn.Parameter(torch.randn(num_classes, embedding_dim))
def forward(self, embeddings, labels):
# 特征与权重都做 L2 归一化，计算余弦相似度
cos = F.linear(F.normalize(embeddings), F.normalize(self.weight))
target_cos = cos.gather(1, labels.view(-1, 1)).squeeze(1)
# 对目标类别的余弦值减去 margin，拉开类间距离
target_cos = target_cos - self.margin
logits = cos.clone()
logits.scatter_(1, labels.view(-1, 1), target_cos.view(-1, 1))
logits = logits * self.scale
return F.cross_entropy(logits, labels)
从代码可以看出，AMSoftmax 会持续挤压特征空间：为了满足间隔约束，模型必须为每个类别寻找更加紧凑的特征表示。当训练样本数量有限时，这种挤压会把样本中的个体差异（光照、姿态、背景噪声）也编码进特征，形成所谓的记忆化。此时模型在训练集上的余弦间隔可以轻松拉满，但在新样本上特征漂移严重。
另外两个加剧因素也值得注意：一是 scale 参数设置过大（例如 64 以上）会让 logits 变得尖锐，梯度集中在少量困难样本上，容易放大噪声；二是分类权重矩阵的维度等于类别数乘以特征维度，类别很多时参数量本身就不小，进一步增加了过拟合风险。
正则化手段：从权重衰减到特征约束
权重衰减与分层设置
权重衰减是最直接的正则化方式，但用在 AMSoftmax 上有个细节：分类头中的 self.weight 通常会做 L2 归一化，对它再加权重衰减意义不大，甚至可能干扰训练。推荐的做法是把骨干网络和分类头分开设置优化器参数组：
optimizer = torch.optim.AdamW([
{"params": model.backbone.parameters(), "weight_decay": 5e-4},
{"params": loss_fn.weight, "weight_decay": 0.0}, # 归一化权重不做衰减
], lr=1e-3)
骨干网络的衰减系数可以从 5e-4 起步尝试，如果过拟合依然严重，可以逐步加大到 1e-3。观察训练曲线时，重点看验证损失而不是训练损失，前者回升才是衰减不足的信号。
Label Smoothing 缓解过度自信
AMSoftmax 训练到后期往往过度自信，目标类别的余弦值趋近 1，其余类别趋近 0，这种极端分布本身就是过拟合的表现。把损失中的硬标签替换为平滑标签能有效缓解：
def am_softmax_with_smoothing(logits, labels, scale, margin, smooth=0.1):
num_classes = logits.size(1)
# 构造平滑后的软标签分布
soft_labels = torch.full_like(logits, smooth / (num_classes - 1))
soft_labels.scatter_(1, labels.view(-1, 1), 1.0 - smooth)
log_probs = F.log_softmax(logits * scale, dim=1)
return -(soft_labels * log_probs).sum(dim=1).mean()
平滑系数一般取 0.05 到 0.1 之间，太大反而会削弱间隔的判别力。需要注意的是，间隔机制与标签平滑存在一定张力：平滑鼓励概率分布更均匀，间隔要求目标类占据主导，因此 margin 不宜同时设置过大，建议在 margin=0.3、smooth=0.1 的组合附近微调。
Dropout 与特征归一化的配合
在特征提取层之后、进入 AMSoftmax 之前加入 Dropout，可以在批次内制造特征扰动，迫使模型学习更鲁棒的表示。由于 AMSoftmax 前通常有 L2 归一化操作，Dropout 率不宜过高，0.1 到 0.3 即可，过高的丢弃率会让归一化后的特征方向剧烈抖动，训练不稳定。此外可以尝试在 embedding 层使用较小的输出维度（如 128 或 256），维度过高在数据量不足时本身就是过拟合的温床。
数据增强：比正则化更根本的解法
图像层面的增强组合
正则化本质上是在惩罚模型复杂度，而数据增强是在扩充有效样本，后者对度量学习往往效果更明显。针对人脸、行人等常见场景，推荐以下增强组合：
- 随机水平翻转：几乎无副作用的基础增强，必开。
- 随机裁剪与轻微缩放：模拟目标位置变化，裁剪比例控制在 0.8 到 1.0 之间。
- 色彩抖动：亮度、对比度、饱和度的扰动幅度建议不超过 0.2，幅度过大会改变身份相关纹理。
- 随机擦除：以较小概率（0.25 左右）遮挡局部区域，迫使模型不依赖单一局部特征。
- Mixup 需谨慎：直接对图像做线性混合会破坏身份标签的准确性，如需使用建议只在同类别样本间混合，或改用流形上的特征级增强。
使用 PyTorch 可以这样组织增强流水线：
from torchvision import transforms
train_transform = transforms.Compose([
transforms.RandomResizedCrop(112, scale=(0.8, 1.0)),
transforms.RandomHorizontalFlip(),
transforms.ColorJitter(0.2, 0.2, 0.2),
transforms.RandomGrayscale(p=0.05),
transforms.ToTensor(),
transforms.RandomErasing(p=0.25),
])
类别均衡与困难样本采样
数据层面的另一个隐蔽问题是类别不均衡。AMSoftmax 对尾部类别的间隔学习往往不充分，而头部类别又因样本过多被反复拟合。采用 P 类每类 K 张的 PK 采样策略，可以保证每个批次内类别分布相对均衡。同时配合困难样本挖掘时要注意设置挖掘上限，比如每个批次内只取最难的 25% 样本参与额外加权，否则极端困难样本（多为标注噪声）会主导梯度方向，反而把模型带偏。
组合方案与调参建议
把上述手段整合起来，一个经过验证的基线配置大致是：骨干网络权重衰减 5e-4，分类头不衰减；Dropout 率 0.2，特征维度 128；scale 取 30，margin 取 0.3，标签平滑 0.1；数据增强采用翻转加随机裁剪加轻度色彩抖动加随机擦除；采样策略为 PK 采样，P=16、K=4。在这个基础上，优先调整的变量是 margin 和增强强度，两者对泛化能力的影响最直接。
调参过程中建议固定随机种子做对照实验，每次只改一个变量，同时记录闭集准确率与开集验证指标（如 TPR@FPR）两条曲线。如果发现训练后期开集指标震荡下降而闭集准确率仍在上行，说明间隔约束已经开始记忆训练样本，此时应该回退增强强度或加大正则化力度，而不是继续增大 margin 去追求训练集上的更小损失。
最后补充一点：如果数据量实在太小（比如每类只有几张图），再强的正则化和增强也难以完全弥补，这时可以考虑先用大规模公开数据集预训练骨干网络，再在自己的小数据集上用较小的学习率微调，同时把 margin 从 0 逐步退火到目标值，让模型有一个从宽松到严格的过渡过程，通常能显著缓解过拟合并提升最终指标。