
我最早开始重新审视softmax分类头是因为在模拟项目X上遇到一个非常现实的问题干净数据上评测指标还行只要数据里混入一点噪声或者图像上出现之前没见过的局部遮挡模型立刻变得不稳定。一开始我以为是骨干网络不够强换了更大的模型能缓解但有限。后来把注意力放到最后一层的全连接分类头才意识到问题不在特征图上而在“最后的判决方式”。今天想聊的CPLConvolutional Prototype Learning卷积原型学习就是一套把分类头从“全连接层 softmax”换成卷积原型学习的方案。核心思路一句话就能讲明白每个类别不再用一个向量点当原型而是用一组可学习的卷积模板在骨干输出的特征图上逐位置滑动匹配最后根据匹配证据的累积结果做类别判决。最终概率化还是要softmax所以更准确的说法是它替代的是“全连接线性层 softmax”这套方案里线性判决的那部分。这篇文章我会把这套思路的原理、最小实现、损失设计、以及我在复现过程中踩过的坑一次说清楚。适合那些觉得当前分类模型泛化不够稳、想换个有可解释性的分类头、或者单纯对原型学习感兴趣的开发者和研究者。1. 先看softmax作为分类头时潜在的问题1.1 线性分类头等于把类别原型压成一个方向向量绝大多数卷积网络做分类时最后一层都是一个全连接层输出各个类别的logitss_y W_y^T h(x) b_y这里的W_y如果拆开来看其实就是第y类的“原型向量”。每个类别只有一个这样的方向向量网络对一张图的判决依据就是特征h(x)和这组方向向量之间谁内积最大。问题在于真实世界里的视觉类别几乎都是多模态的。一个类别内部存在大量亚模式不同姿态、不同环境、不同材质形态。你让一个全连接层用一根向量去表达所有这些亚模式必然造成特征空间挤压。模型只能选择一个对所有子模式都还算说得通的方向这个方向往往落在几个子模式之间导致判决边界被拉平对不常见的子模式覆盖不足。类比一下就很好理解好比让保安只看一张“证件照”来认人来人换个角度、换个着装他就要犹豫可如果保安手里有一组不同方位的照片识别就会稳健很多。CPL要做的就是把每个类别的判决依据从“一张证件照”变成“一组模板照片”。1.2 空间结构信息在全局池化后已经丢了CNN骨干输出的特征图最后通常要接全局平均池化变成一个向量h(x)再接全连接层。这一步操作把空间位置信息全部丢掉了某个局部纹理出现在图像左上角还是右下角在全局平均之后只剩一个均值。对纯粹的图像语义分类来说这可能不算坏事因为“物体出现在哪”往往不重要重要的是“是什么”。但对鲁棒分类来说这是实打实的损失。一个很直接的场景是遮挡图像左上角被遮掉一块全连接头看到的特征向量只是若干维度数值发生了变化它不知道这种变化是来自某个局部区域被污染还是整个图像分布发生了偏移。模型只能被动地把这个变化当作全局信息去调整判决。更麻烦的是细小的判别性局部模式在均值池化中很容易被淹没。比如工业质检场景里一个产品表面只有一小片缺陷但这个缺陷恰恰是分类的关键证据全局平均池化会把这个局部响应稀释到几乎看不见。1.3 为什么这种状况在分布偏移下更敏感softmax训练的目标是最小化交叉熵梯度几乎完全由误分类样本的“方向偏差”驱动。当测试分布发生偏移——亮度变化、对比度变化、噪声模式变化——特征空间整体会跟着偏转但全连接判决面不会因此自适应。局部模式如果发生了变化全局特征向量可能直接从判决区域的这一边滑到另一边。这种情况下模型表现出的鲁棒性本质上完全依赖数据增强是否覆盖了这些偏移。你加多少增强它就防多少种变化这是一种被动防御。增强覆盖不到的风格变化、遮挡形式、局部扰动在分类头层面没有任何结构性的缓冲机制。这个认知让我换了一个判断标准如果一个分类头在判决时本身就能保留并利用局部证据可能就不需要依赖无穷无尽的数据增强去兜底。顺着这个思路去找“在特征图上做原型匹配”的方案就落到了CPL上。2. CPL的核心思想与整体设计2.1 卷积原型的定义CPL把类别原型定义成卷积核。对第y个类别有G个原型卷积核P_{y,g}每个原型形状是C×h×w其中C是骨干输出的特征图通道数h×w是原型尺寸通常取3×3、5×3、3×5这样小的局部模板。模型把特征图F和这G个原型分别做卷积E_{y,g} F * P_{y,g}得到G张证据图。每一张证据图上的每一个空间位置都表示“以这个位置为中心的局部区域在多大程度上匹配这个原型模板”。如果某个位置恰好出现原型对应的局部模式响应就会高否则响应低。这里的核心差别在于全连接分类头是在全局特征向量上做内积卷积原型是在特征图上做滑窗内积。后者显式利用了空间位置信息而且天然具备平移等变性——一个局部模式出现在图片的任意位置都能被同一个模板匹配到。模型要学习的不是“某个特征整体长什么样”而是“哪些局部模式组合在一起最有可能代表这个类别”。2.2 证据图到类别分数的转换有了证据图之后怎么把它变成每个类别的标量logits这是整个设计方案里最讲究的地方。如果直接拿证据图做全局平均池化去掉了空间信息和普通分类头的区别就小了一大半如果直接取全局最大值又只依赖响应最强的那个位置容易被背景伪影带偏。我实际用的折中方案是Top-k平均。对每个类别先在G个原型维度上做排序取每张特征图位置上匹配强度最高的k个原型响应做平均然后在空间维度上取均值。这样得到的类别得分含义就是“这张图里最自信的那几个局部原型匹配的平均强度”。这里面有一个很直观的鲁棒性逻辑如果某张图的某个局部区域被遮挡或污染那么对应位置的证据响应会被压低但其它正常位置仍然可以提供强的原型匹配最后在空间维度上取平均时被污染位置的影响被稀释。这个机制和全连接头“把所有特征维度混合后判决”完全不同。2.3 设计参数怎么选参数常见取值选择理由原型尺寸3×3、5×5、3×5太小抓不到足够上下文太大容易过拟合局部共现每类原型数G3~8再往上收益边际递减显存和计算量线性增长Top-k2~5兼顾整体证据和局部最强响应温度t0.05~0.15控制logits尺度影响早期训练梯度关于原型尺寸还有个小细节。奇数尺寸的原型在卷积padding时是对称的3×3配padding15×5配padding2可以有效保持特征图尺寸不变。偶尺寸核要么让特征图尺寸漂移要么padding不对称边缘效应会比较烦。我自己一般用3×3起步先把效果跑通再尝试非对称尺寸。每类原型数的选择取决于类别内多样性程度。工业缺陷类别、细粒度物种这些类内形态差异很大的可以多给几组原型相对规整的类别3组就已经够用。多出来的原型如果一直得不到有效激活本质上就是浪费。3. 最小可复现实现与损失函数拆解3.1 前向计算的最小实现用PyTorch这类主流框架实现CPL头非常直接。核心就是把原来接在全局池化后的全连接层换成一组对特征图做卷积的原型参数。import torch import torch.nn as nn import torch.nn.functional as F class ConvPrototypeHead(nn.Module): def __init__(self, in_channels, num_classes, num_protos5, proto_size(3, 3), init_temperature0.1): super().__init__() self.num_classes num_classes self.num_protos num_protos ph, pw proto_size self.ph, self.pw ph, pw # 每个类别一组原型每组num_protos个卷积核 self.prototypes nn.Parameter( torch.randn(num_classes * num_protos, in_channels, ph, pw) * 0.05 ) self.temperature nn.Parameter(torch.tensor(init_temperature)) def forward(self, feature_map): pad_h, pad_w self.ph // 2, self.pw // 2 evidence F.conv2d( feature_map, self.prototypes, stride1, padding(pad_h, pad_w) ) B, _, H, W feature_map.shape evidence evidence.view(B, self.num_classes, self.num_protos, H, W) # 先取top-k原型响应再在空间维度上平均 k min(3, self.num_protos) sorted_ev, _ evidence.sort(dim2, descendingTrue) topk_avg sorted_ev[:, :, :k].mean(dim2) # [B, C, H, W] logits topk_avg.mean(dim(2, 3)) / self.temperature return logits, evidence这里有个需要解释的点topk操作的索引本身不可导但模型对选出来的values继续求均值时梯度是可以正常回传到这些被选中的证据图上的。所以整个前向过程对训练完全友好不需要什么特殊的“软化”技巧。PyTorch的sort默认在最后一维排所以我把G个原型放到维度2排序后切片即可。3.2 概率化与温度参数分类头的最后logits还要除以一个可学习的温度t再进softmax。这个操作参考了对比学习里的温度设计本质是控制概率分布的锐利程度。初始温度我建议设在0.05到0.15之间。原因很简单特征图和原型卷积出来的响应数值可能很大如果直接除以1再进softmax概率会立刻饱和到接近one-hot。梯度在饱和区几乎消失原型层更新缓慢。反过来温度太小会让早期概率分布过于尖锐模型在犹豫的时候仍然受到很大惩罚训练不稳定。一点实操心得温度参数一定要设成可学习的而不是固定死的。训练到中后期特征图尺度和响应分布会变化固定温度可能成为瓶颈。让温度自己适应通常后期会收敛到0.05到0.3之间的某个值这个值也反映了模型整体对证据的自信程度。3.3 损失函数交叉熵加两个正则项CPL头可以直接用交叉熵训练但直接训会有一个问题组内多个原型可能退化成同一个模板。如果所有原型变得几乎一样那跟只用一个原型没什么区别投票机制就失效了。所以我加了两个正则项。第一个是类内原型多样性正则。做法是把每类原型展平并做l2归一化然后计算类内原型两两之间的相似度让相似度尽量接近0。def diversity_loss(prototypes, num_classes, num_protos): proto_view prototypes.view(num_classes, num_protos, -1) proto_view F.normalize(proto_view, dim-1) sim torch.matmul(proto_view, proto_view.transpose(-1, -2)) # 只看类内非对角元素 mask 1 - torch.eye(num_protos, deviceprototypes.device) loss (sim * mask.unsqueeze(0)).sum() / (num_classes * num_protos * (num_protos - 1)) return loss第二个是证据图熵正则。把证据图在类别维度上做softmax然后约束每个空间位置的分布尽可能尖。含义是一个局部区域不应该对多个类别都产生模糊的高响应而应该明确地属于某个类别的局部模式。def evidence_entropy_loss(evidence, num_classes): B, _, H, W evidence.shape ev evidence.view(B, num_classes, -1) p ev.softmax(dim1) entropy -(p * torch.log(p 1e-8)).sum(dim1).mean() return entropy总损失是loss ce_loss alpha * diversity_loss beta * entropy_lossalpha我建议0.01到0.1beta建议0.001到0.01。不要贪心。我试过把alpha调到0.3原型多样性确实保住了但骨干特征被挤得厉害干净样本的logits反而退化。正则的目的是“唤醒所有原型”不是“强制每个原型都独立到极致”。3.4 训练配置的推荐训练CPL和直接换softmax头的常规训练不太一样有几个配置值得注意优化器SGD和AdamW都能跑我在模拟项目X上的对比实验里SGD配momentum0.9更稳尤其是后期收敛阶段。学习率建议比常规分类任务略低批次256时从0.01到0.03起步比较合适。原型卷积层对学习率偏敏感太大会让证据图尺度剧烈波动。骨干学习率如果是预训练骨干替换最后几层建议特征提取器的lr乘0.1CPL头用正常lr避免预训练特征被快速破坏。数据增强CPL对强增强的兼容性很好因为证据图对局部变化宽容。Cutout、随机遮挡这类增强可以直接上不会像某些分类头那样因为结构信息被切断而剧烈波动。训练轮数原型层的收敛比单纯全连接头慢一点点我一般会多训10%的epoch。4. 和相近思路的对比以及影响范围4.1 和原型网络、度量学习的区别不少文章把CPL和原型网络Prototypical Networks搞混。原型网络的做法是用support set里的样本特征均值作为每个类的原型然后按距离度量最近类。它本质是metric-based适合少样本任务需要设计episode来采样support集和query集。CPL不是样本均值原型而是可学习的卷积模板原型。它不依赖support set端到端直接训练更接近一个“分类头”的替换方案。少样本场景里可以把两者结合但核心机制并不相同一个比的是整体embedding距离一个比的是局部模板匹配证据。4.2 和ArcFace、CosFace、Softmax margin系列的区别角度损失ArcFace、CosFace这类本质是在softmax上加margin让同一类特征更聚拢类别中心向量仍然是每类一个原型向量。它提高的是特征的判别性但对“判决时如何利用局部结构证据”没有帮助边际收益在分布偏移下会衰减。CPL和这些margin损失不冲突。如果你想要更紧凑的特征也可以在CPL的logits上加margin只是我实测下来CPL带来的证据累积改变已经让特征在局部模板层面更规整margin的必要性没有那么高。4.3 适用与不适用场景CPL的适用场景核心特征是“存在二维空间结构的特征图”。图像分类、细粒度分类、工业表面缺陷检测、医学影像分类这类任务CPL换头的侵入很小效果通常比较明显。不太适用的场景也要说清楚文本分类token序列本质上是一维没有二维空间结构卷积原型在这上面优势不大。表格数据特征本身就是分散的列不存在空间邻域卷积模板无从谈起。骨干输出已经退化成1×1或接近1×1的特征图比如一些轻量网络最后一层通道数高但分辨率只有1这时卷积原型等于在1×1上做1×1卷积退化成全连接没有意义。类别数量极大的场景万类级别时每组原型再乘类别数参数量会爆炸性价比很低。影响范围方面CPL对训练流程的侵入很小相当于只替换了分类头不依赖特殊算子普通CNN骨干都能接。这也是我比较愿意推它的原因可作为可插拔选项先用常规softmax头跑基线再换CPL头对比很快就能看出结构差异带来的变化。5. 鲁棒性提升的微观机制5.1 多个原型投票让标签噪声的影响被稀释softmax头用一根向量方向去拟合所有被标为同一类的样本少量错误标签会直接拉动这一根方向向量污染全局。CPL组内的每个原型对应不同局部模式一个错误标签样本只会激活它最像的那少数几个模板对类别分数的拉动会被其它原型的正常响应稀释。这有点像随机森林里多棵弱树的投票效应单个原型可以犯错但多数原型持有正确证据时最终判决还是会回到正确的一侧。我在模拟数据上把10%的标签随机打乱测试softmax头的验证损失波动明显加大而CPL头相对更平稳损失曲线没有出现那种尖锐的峰值。5.2 局部遮挡对证据图的影响是局部的全连接头对遮挡没有空间感知图像左上角被遮掉一块和中间被遮掉一块在特征向量里造成的数值变化可能完全不同但判决逻辑不会对这种空间位置做出区分。CPL的证据图天然保持空间位置信息遮挡会压低对应位置的响应但其它正常位置的原型匹配仍然存在Top-k平均会在空间维度上自动寻找最自信的响应位置。这里有个不太好听但很重要的提醒CPL真正带来的提升不是某种“魔法”而是把“局部证据累积”这个归纳偏置显式放进了判决过程。骨干网络本身如果学得一塌糊涂特征图里根本没有有意义的局部模式换什么头都没用。5.3 哪些提升是真实的哪些是故事聊到这里我得诚实说一句。CPL换头在不同任务上的收益差异非常大。如果骨干已经极其判别比如用超大模型在大量干净数据上训练那么换CPL对干净测试集准确率几乎不会有提升它改善的更多是鲁棒性和分布偏移下的下限。反过来如果骨干偏弱、数据量少CPL的结构约束反而表现出明显的正则化作用验证损失更平滑过拟合更晚出现。不要把这个东西神话成“解决鲁棒分类的银弹”。它是一个归纳偏置合理的分类头能帮你在不换骨干的前提下获得更稳、更有解释性的证据累积过程。至于最终能改善多少取决于现有分类头的短板有多大。6. 复现与调优中容易踩的坑6.1 证据图数值尺度失控特征图的数值范围并不固定。如果骨干输出没有经过BN或者数值本身偏大原型卷积出来的evidence可能上千除以温度后还是非常大softmax直接饱和梯度消失。我的习惯是在原型卷积之后不看原始证据图而是观察logits的范围。训练早期如果logits绝对值经常大于10就说明温度初始化太小或者原型初始化偏大。把原型初始化scale从0.1降到0.05通常能解决。也可以对特征图先过一个LayerNorm再进原型卷积效果稳定只是多一层开销。6.2 原型死亡和局部模式坍缩训练过程中经常出现一类原型始终得不到高响应最后梯度越来越小变成“死亡原型”。这个和ReLU死亡类似但原因不太一样特征图如果过于稀疏某些随机初始化的原型可能一开始就没有和任何局部模式对上网络也没机会修正它。我的对策是在骨干输出特征图之后加一个ReLU或温和的激活保证特征图非负且不稀疏。如果死亡原型仍然存在尝试把原型的初始化改成从骨干预训练特征上随机切片初始化这样一开始就有合法的局部模式信号。6.3 位置泄露我踩过最隐蔽的坑是原型学到了“类别和位置”的组合而不是“类别和局部形态”的组合。比如某个类别在数据集中总出现在中心位置原型可能会偷偷把中心位置的高响应当成类别证据而不是真正学到局部模式。排查方法很简单把证据图按空间位置做均值打印出来。如果某个类别的证据图激活位置一直集中在图像中心或固定角落说明发生了位置泄露。缓解办法是加强随机平移、随机裁剪这类增强让类别和位置解耦。6.4 Top-k操作带来不稳定的梯度Top-k本身不会断掉梯度但会带来阶梯状的响应切换某张图的第k1个原型响应突然超过第k个选中的原型集合发生跳变logits会有轻微抖动。这个问题在训练后期尤其明显。如果不希望看到这种抖动可以考虑先用soft top-k近似或者直接把k设得大一点减少选择集合的频繁变动。k3时抖动通常可忽略k1时偶尔会有明显跳变这点在调参时可以留意。6.5 可视化是调试的核心手段CPL比全连接头好调试因为原型就是卷积核可以直接把每类原型可视化出来。做法是把原型展平成图像大小或者统计原型在训练集证据图上的高响应区域。我常用的小技巧是每个epoch结束后打印每个类别的原型平均激活率和证据图热区。如果某个类别的原型激活率长期偏低基本就是该类原型初始化不好或者类别本身证据不足应该考虑增加该类原型数量或调整正则在类间的权重。复现CPL没有什么玄学数据流清晰、参数可解释、故障能用可视化定位这是它比普通分类头更友好的地方。只要把温度、原型数量、Top-k这三个核心参数调顺剩下的交给骨干网络就行。最后再说一点个人体会。我在实际测试中发现CPL这类分类头最值得推广的用法不是直接替换softmax成为唯一方案而是作为基线之外的对比项。每次做一个带空间结构的分类任务先跑一发全连接softmax再跑一发CPL头两个结果摆在一起往往能暴露出原有分类头到底在依赖整体特征还是局部证据。这种诊断价值比单纯涨点更让人觉得踏实。后面如果你要往增量学习、小样本方向走这个“局部模板匹配多原型投票”的结构也天然有扩展空间。