ConvNeXtV2全卷积掩码预训练：GRN与稀疏卷积解析

1. 这个全卷积掩码自编码器到底想解决什么问题

第一次翻到 ConvNeXtV2 的论文时，我盯了半天没反应过来：一个纯卷积的骨干网络，为什么要去做掩码自编码？掩码建模这事儿不是 Transformer 的主场吗——把一块块 patch 挖掉，喂进编码器，让解码器把挖掉的部分补回来，这套流程天生适合注意力机制，因为注意力可以很自然地把"这里有个洞"当作一个 token 处理。卷积怎么搞？卷积核是滑窗，你把某块像素填成 0，它和真实的黑色像素、和补零 padding 三者混在一起，网络根本分不清。

ConvNeXtV2 给出的答案是把这件事从头设计一遍。它由三个咬合在一起的部分组成：一套针对全卷积主干重新设计的掩码预训练框架（全卷积掩码自编码器）、一个叫 GRN 的归一化层（全局响应归一化），以及配合稀疏卷积的稀疏化前向实现。三者不是并列关系，而是因果链——因为要用卷积做掩码建模，掩码后的零值会污染卷积核，所以需要稀疏卷积；因为稀疏输入会让通道响应趋同、特征退化，所以需要 GRN；因为有了 GRN，掩码预训练才真正训得起来，微调后才能在 1000 类图像分类这类基准上超过同等规模的监督预训练模型。

这篇内容适合三类人看。第一类是正在选型的算法工程师，手上有个检测或者分割的下游任务，想找一个比 ViT 更省显存、部署更友好的骨干，ConvNeXtV2 是很值得纳入候选的。第二类是想复现掩码预训练的研究型同学，这套流程的每一步都有坑，从掩码比例到损失计算范围，错一处 loss 就是平的。第三类是只想抄一个 GRN 模块插到自己网络里的人——说实话，这是整篇论文里性价比最高的部分，一个二十行的模块，几乎零额外开销，能实打实缓解通道响应塌缩。

我个人认为这篇工作最值得琢磨的地方不在于刷了多少点，而在于它示范了一种思路：当你想把一个为 A 架构设计的方法搬到 B 架构上时，不要只搬流程，要搬完之后检查 B 架构的"生理结构"会不会让流程失效，失效了就针对性打补丁。ConvNeXtV2 就是一次完整的"移植—诊断—打补丁"过程。

2. 整体设计思路：三个组件为什么必须一起出现

2.1 从监督预训练到掩码预训练的路线选择

卷积网络的预训练范式在过去几年基本是被监督学习统治的：拿一个大的有标注数据集训一遍，再迁移到下游。这条路的天花板很明显——标注成本。自监督在 Transformer 阵营已经被掩码建模验证过一轮，效果甚至能超过监督预训练，那卷积为什么不行？

早先的尝试主要有两条。一条是把卷积网络和对比学习结合，问题是需要精心构造正负样本对，对数据增强策略极其敏感，训起来调参成本高。另一条是把卷积网络塞进自编码器框架里做像素重建，早期的做法是直接把整张图编码再解码，不做掩码，结果网络学到的多是低频颜色统计，学不到有判别力的中高频结构，迁移效果一般。

ConvNeXtV2 选的是掩码重建这条最难但天花板最高的路。难点在哪？难在卷积的归纳偏置和掩码这个操作是冲突的。卷积假设输入在空间上是连续、稠密、邻域相关的；掩码人为制造了空洞，破坏了这种连续性。Transformer 没有这个假设，它靠注意力全局建模，位置可以被"挖掉"而不影响其它位置的计算。所以要让卷积吃下掩码输入，必须解决"卷积核看不见被挖掉的区域"这件事。

我踩过一个很典型的坑：早期我自己试着在卷积主干上套标准掩码自编码流程，就是把被 mask 的 patch 全部置零然后正常前向。训了三十个 epoch，loss 降到某个值之后完全不动，重建出来的图是一片模糊的色块。当时以为是学习率的问题，调了半天没用。后来把中间层特征拿出来看才明白——通道之间的激活值几乎一模一样，相当于网络把所有通道学成了同一个东西，特征维度白白浪费了。这就是后来论文里明确讨论的特征坍塌，也是 GRN 要解决的问题。

2.2 三件套的分工与因果链

把三个组件拆开看，各自解决的是不同层级的问题，缺一个都跑不通，我用一张表把它们的职责和失效后果对应起来。

因果链是这样的：为了做掩码建模，必须处理稀疏输入，于是引入稀疏卷积；稀疏卷积让特征图在某些位置"不存在"，反向传播时这些位置拿不到梯度，通道统计分布被压缩，进而出现通道间响应高度相似；GRN 在通道维度做一次对比式归一化，把主导通道的响应放大、把萎靡通道压下去，强行把通道间差异拉开，从而恢复特征的表达能力。

这里有个设计细节值得说：GRN 被放在 block 的升维之后、降维之前，也就是在 4 倍宽的高维空间里做归一化。为什么不在原始维度做？因为高维空间里通道数多，统计量更稳，而且归一化的结果马上要被 1×1 卷积压回低维，等于让这次归一化直接影响降维卷积的输入分布，作用更直接。代价几乎没有，因为高维空间的张量只是中间产物，不落地到显存里长期驻留。

2.3 与 Transformer 版掩码自编码的关键差异

很多人第一反应是"这不就是卷积版 MAE 吗"，方向对，但差异比想象中大，这些差异恰好决定了实现时的选型。

第一个差异是没有位置编码。Transformer 版必须给每个 token 加位置编码，否则注意力是置换不变的，模型不知道谁在哪。卷积天然带位置信息，滑窗的局部性本身就是一种强位置先验，所以 ConvNeXtV2 里完全不需要位置编码，也就省掉了掩码后位置编码怎么处理、变分辨率时位置编码怎么插值这一堆麻烦事。

第二个差异是掩码粒度更粗。Transformer 版常用 16×16 的 patch，掩码比例 75%；ConvNeXtV2 用 32×32 的 patch，掩码比例 60%。为什么？32×32 的 patch 意味着卷积核的感受野能覆盖到的邻域更大，块内信息更完整；掩码比例降到 60% 是因为卷积的感受野有限，遮得太狠会剩下太多"孤岛"，稀疏卷积处理起来效率低且边界效应严重。

第三个差异是编码器和解码器的非对称程度。Transformer 版解码器很轻，几个 block 就够；ConvNeXtV2 的解码器也是轻量的，但要处理的是"从稀疏特征图恢复到稠密像素"这个更别扭的操作，因为它需要在被遮掉的位置上也生成预测，而主干在这些位置根本没有输出。实现上通常是在解码器里先把稀疏特征图补零变稠密，再用轻量卷积上采样。

3. FCMAE：把掩码自编码器塞进全卷积主干的具体做法

3.1 掩码策略与 patchify 的设计取舍

预训练阶段的输入处理和我们平时做分类完全不一样。分类时 stem 通常是 4×4、步长 4 的卷积加一个 LayerNorm，把 224×224 的图降到 56×56。预训练时这个 stem 要换掉，改成 32×32、步长 32 的卷积，把图直接切成 7×7 个 patch，每个 patch 对应 32×32×3 的原始像素。

我一开始觉得这个设计很怪——为什么不用大 kernel 的卷积做 patch embedding？后来自己想明白了：掩码自编码的重建目标必须是原始像素级的信息，如果 stem 用了 4×4 步长 4，那后面所有的下采样都已经把细节丢掉了，解码器根本没法恢复出 32×32 的细粒度纹理。所以预训练 stem 必须一次性切到 32×32，让每个 token 保留完整的块内像素，重建才有可能。

掩码的具体操作是在 patch 维度上做的：一共 49 个 patch（224/32 的平方），随机选 60% 遮掉，剩下约 20 个 patch 送进网络。注意这里"遮掉"的含义和 Transformer 版不同——Transformer 是把 token 直接删掉，序列变短；卷积版因为要保持空间结构，不能真的删，只能把这些位置标记为"无效"，让稀疏卷积跳过它们。

注意：掩码比例 60% 是配合 32×32 的块大小定的，如果你把块大小改小到 16×16，比例要往上调（70%~75%），否则剩下的可见区域太多，任务太简单，学不到东西；反过来块大小改大到 64，比例要降到 50% 左右。这是一个联动参数，不要单独调。

3.2 稀疏卷积到底解决了什么

这是整个流程里最容易被含糊带过的部分，但恰恰是能否复现成功的关键。

问题分两层。第一层：假设你把被 mask 的位置填 0，然后正常做 7×7 深度卷积。卷积核在滑过可见区域和遮罩区域的交界处时，会把那些 0 当成真实的像素值参与加权求和。结果是可见区域边缘的特征被"稀释"了——本来一个通道的响应是 1.0，靠近边界的输出可能变成 0.5。这会让网络学到一个错误的先验：靠近边界的特征就是弱的。而推理时根本没有掩码，这个先验就变成了纯粹的噪声。

第二层更微妙：卷积的零填充本身就是用 0 补边界。mask 填的 0 和 padding 补的 0 在数值上完全一样，网络没法区分"这里是图像外部"和"这里是被我遮掉的"，两者的处理逻辑应该完全不同。

稀疏卷积的做法是改变计算的定义域：只对"有效位置"做卷积，输出也只落在有效位置上，无效位置的输出直接不存在（而不是 0）。这样卷积核在计算时，邻域里只统计真实存在的点，被遮掉的位置不参与加权，边界处也不会引入假的零值贡献。数学上等价于在稀疏张量上定义一个局部算子，实现上需要维护坐标索引。

具体实现有两条路，我在不同项目里都用过：

路线一，用子流形稀疏卷积算子。需要额外装一个稀疏卷积库，构造时把特征和坐标一起打包，卷积层会在坐标集合上运算。优点是语义完全正确，边界干净；缺点是依赖重，版本兼容性差，我自己就遇到过算子库和框架版本对不上、编译失败半天的经历，而且调试起来很痛苦，中间张量没法像普通张量那样随手打印。

路线二，归零近似。正常用稠密卷积，但在每次卷积之后把被 mask 的位置重新置零，同时保证所有卷积层的 bias 设为 False。这样做的效果是：无效位置永远保持严格的零，卷积核读到的邻域里只有真实像素和被强制归零的位置；边界处依然会有稀释，但至少不会把错误信号往远处传播。

归零近似的代价要说清楚：7×7 深度卷积的感受野半径是 3，也就是说在有效区域的边缘 3 个像素内，输出特征仍然是"欠统计"的。对于 32×32 的块来说，边缘 3 个像素占比大约 34%，不算小。实测下来归零近似能跑到稀疏卷积大概 95% 的效果，对于个人复现、小规模验证完全够用；如果要追论文数字，还是得上真正的稀疏算子。

另外有个容易忽略的点：LayerNorm 和 GRN 都是沿通道维计算的（对每个空间位置独立地跨通道做归一化），不会把信息从被遮位置扩散到可见位置，所以在归零近似里这两个层是安全的，不需要特殊处理。GELU 也一样，GELU(0)=0，零进零出。真正会跨空间混合的只有卷积层，只需要盯住卷积就够了。这个观察省了我不少事——一开始我甚至想把 LayerNorm 也改成掩码感知的，后来发现完全没必要。

3.3 重建目标与损失范围

重建目标这一块，细节决定成败。核心做法是：在计算损失之前，对每个 patch 内部的像素做一次标准化，减掉这个 patch 的均值、除以标准差。这一步叫逐块归一化。

为什么必须做？因为如果不做，网络会倾向于先学"这块区域的平均颜色"，把大量容量浪费在低频信息上。逐块归一化把每个 patch 的均值方差都抹平了，网络被迫去学块内的相对结构，也就是纹理和边缘，这才是迁移到下游任务时有用的东西。

损失只在被遮掉的 patch 上计算，可见 patch 不参与。这一点很多人会写错，包括我自己第一版。如果全程都算损失，那 40% 的可见区域网络可以直接抄输入，梯度被这些"简单题"主导，被遮区域的梯度被稀释，模型学不到真正的补全能力。

伪代码大致是这样：

# 预训练阶段每个 step 的核心逻辑 images = images.cuda(non_blocking=True) # (B, 3, 224, 224) target = patchify(images, patch_size=32) # (B, 49, 3*32*32) target = per_patch_normalize(target) # 每个 patch 内做标准化 x, mask = random_masking(target, mask_ratio=0.6) # mask: (B, 49) 布尔，True 表示被遮 pred = model(x, mask) # 只在可见位置做稀疏前向，再补全 loss = ((pred - target) ** 2)[mask].mean() # 只统计被遮位置 loss.backward() optimizer.step()

注意[mask]这个索引，它要求 pred 和 target 都被重排成"先所有可见、再所有被遮"的紧凑顺序，否则布尔索引对不上。这个顺序重排是掩码自编码里最容易写错的地方之一，我在两个不同项目里都栽过，症状是 loss 能降但重建出来的图错位。

提示：调试顺序重排最简单的方法是用小数据过拟合——只取 4 张图，训 200 步，看能不能把重建误差压到接近 0 并且重建图肉眼可辨。如果压不下去，八成是顺序或者掩码索引错了，而不是模型容量不够。

4. GRN：一个二十行的模块，解决通道塌缩

4.1 特征坍塌长什么样，怎么量化

先说清楚现象。掩码预训练跑起来之后，如果你把某个 stage 的输出特征拿出来，计算通道之间的余弦相似度矩阵，正常情况应该是一个接近对角占优的矩阵——不同通道学到不同模式，彼此差异明显。发生坍塌时，这个矩阵会几乎全是接近 1 的值，意味着所有通道的输出几乎线性相关。

更直观的一个指标是"死亡通道比例"：统计每个通道在所有空间位置上的标准差，如果某个通道的方差接近 0，说明它在任何输入下都几乎输出同一个值，这个通道就废了。坍塌严重时，死亡通道能占到 40% 以上。我在一次失败的预训练里测过，中间 stage 的死亡通道比例到了 52%，难怪怎么调学习率都上不去。

坍塌的成因值得稍微展开。随机掩码意味着每次前向的输入分布差异极大——这次这张图遮的是左上角，下次遮的是中间。卷积核要在这种极度不确定的输入分布上工作，最"安全"的策略就是把所有通道都学成接近恒等映射，输出一个平滑的、不依赖输入的响应，这样任何掩码模式下都不会产生剧烈误差。这是典型的局部最优陷阱：loss 确实在缓慢下降，但学到的东西没有判别力。

Transformer 那边很少出现这个问题，因为注意力的 softmax 天然会做竞争性归一化——一部分 token 的权重高，另一部分就被压低，通道间差异是被机制本身保证的。卷积没有这个机制，所以需要人为补一个。

4.2 GRN 的计算过程与实现

GRN 的思路非常直接：在通道维度做一次"谁强谁负责"的对比归一化。分四步：

1. 对每个通道，计算它在空间维度上的 L2 范数，得到一个长度为 C 的向量，代表这个通道整体的响应强度；
2. 用每个通道的范数除以所有通道范数的均值，得到归一化系数；
3. 把原特征乘上这个系数（广播到空间维）；
4. 乘一个可学习的缩放 γ 加一个可学习的偏移 β，再残差加回原输入。

第三步是核心：范数大的通道被放大，范数小的通道被压制，通道之间的相对关系被显式强化。这一步实际上是在做软性的通道竞争。

代码不长，但有几个细节必须写对：

class GRN(nn.Module): def __init__(self, dim, eps=1e-6): super().__init__() self.eps = eps self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim)) self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim)) def forward(self, x): # x: (B, H, W, C)，channels-last 布局 gx = torch.norm(x, p=2, dim=(1, 2), keepdim=True) # (B, 1, 1, C) nx = gx / (gx.mean(dim=-1, keepdim=True) + self.eps) return self.gamma * (x * nx) + self.beta + x

细节一：gamma和beta都初始化为 0。这意味着训练开始时 GRN 的输出恒等于输入x，整个模块是一个纯恒等映射。这是非常有必要的——预训练早期特征分布极不稳定，如果 GRN 一开始就大幅改写特征，很容易把训练带偏。随着训练推进，γ 从 0 慢慢长起来，GRN 的作用是渐进注入的。这个"零初始化 + 残差"的模式在 ConvNeXt 系列里反复出现（层缩放也是同一个思路），核心就是让新增模块在训练初期"不添乱"。

细节二：dim=(1, 2)对应的是 H 和 W。如果你用的是 channels-first 布局，维度要改成dim=(2, 3)，gamma的形状也要改成(1, C, 1, 1)。这个我踩过，直接抄代码不改布局，结果归一化在错误的维度上做，等于把 batch 和通道混在一起算了，效果比不加还差。

细节三：分母用的是通道范数的均值而不是求和。论文里写的是求和形式，实现里用均值，两者只差一个常数因子 C，而这个因子会被可学习的 γ 完全吸收，所以效果等价。但如果你从论文公式直接推导代码，会困惑一下。

4.3 为什么放在这个位置，以及它能带来什么

GRN 在 block 里的位置是固定的：深度卷积 → LayerNorm → 1×1 升维（4 倍）→ GELU →GRN→ 1×1 降维 → 层缩放 → 残差。放在 GELU 之后、降维之前。

选这个位置的理由我琢磨过一阵。放在 GELU 之后，是因为 GELU 会把负半轴压到接近零，输出的分布已经比较"干净"，这时候做范数统计更稳定，不会被激活前的正负抵消干扰。放在降维之前，是因为要在高维空间做通道竞争——4 倍宽的通道里有更多"候选者"参与竞争，筛选出来的组合更丰富，如果先降到原始维度再竞争，可选空间小了一半以上。

至于能不能放在 block 开头、LayerNorm 之前，我试过，效果差一些，因为深度卷积的输出还没经过任何归一化，不同通道的数值尺度差异巨大，直接算范数会被极值通道主导。所以位置不是随便放的，它依赖于前面的归一化把尺度拉齐了。

GRN 带来的收益在不同任务上不一样。分类任务上主要是微调后的绝对精度提升，因为通道利用率高了，等效于模型容量被用满。检测和分割任务上感受更明显，因为这些任务依赖细粒度的特征差异，通道塌缩会直接抹掉这些差异。还有一个隐性收益是训练稳定性——加了 GRN 之后，预训练对学习率的敏感度明显下降，我试过在 1e-4 到 2e-4 的区间里变化，loss 曲线都还挺稳，不加 GRN 的时候超过 1.5e-4 就开始抖。

注意：GRN 不替代 LayerNorm，两者是串联关系。LayerNorm 管的是"把每个位置的通道分布拉回标准正态"，GRN 管的是"在通道之间重新分配权重"。有一次我图省事把 LayerNorm 换成了 GRN，结果训练直接崩，因为 GRN 对输入的绝对尺度不做约束，前几层数值会爆。

5. 动手实现：从模块搭建到预训练跑通

5.1 环境准备与工程结构

先把预期说清楚：这套预训练配方在原始工作里是用几十甚至上百张卡、几千的 batch 跑的，个人环境不可能复现同等规模。所以我的建议是分成两个层次——模块级和流程级要严格对齐，规模级要大幅缩减。先把流程跑通、验证每个组件确实起作用，再谈规模。

模块级指的是 GRN、block 结构、损失定义、掩码逻辑，这些必须写对，一处不对就全盘白费。规模级指的是 batch size、图像分辨率、epoch 数、模型宽度，这些可以等比缩。我在一张 24G 显存的卡上跑过完整的预训练流程验证，用的配置是：图像 128×128，patch 32×32（每个样本只有 16 个 patch），模型宽度压到 Base 的四分之一，batch 192，200 个 epoch。loss 能稳定下降、重建图能看出轮廓、微调到一个小规模分类数据集上比随机初始化高 8 个点以上，这就说明流程是对的。

目录结构建议这么分，后面调参的时候好定位：

proj/ models/ grn.py # GRN 模块 block.py # ConvNeXtV2 block backbone.py # 四阶段主干 sparse_ops.py # 稀疏卷积的两种实现封装 data/ patchify.py # 切块 + 逐块归一化 + 随机掩码 pretrain.py # 预训练主循环 finetune.py # 微调流程 configs/ tiny_pretrain.yaml tiny_finetune.yaml

5.2 Block 的完整实现

一个 block 里最容易出错的是通道布局的来回切换。为了性能，卷积希望用 channels-last；但 LayerNorm 和 GRN 的定义都是沿最后一维，所以整个 block 内部统一用 channels-last 最省事，进出 block 的时候不用转。

class ConvNeXtV2Block(nn.Module): def __init__(self, dim, layer_scale_init=1e-6, drop_path=0.0): super().__init__() self.dw = nn.Conv2d(dim, dim, 7, padding=3, groups=dim, bias=False) self.norm = LayerNorm(dim, eps=1e-6, data_format="channels_last") self.pw1 = nn.Conv2d(dim, 4 * dim, 1, bias=True) self.act = nn.GELU() self.grn = GRN(4 * dim) self.pw2 = nn.Conv2d(4 * dim, dim, 1, bias=True) self.gamma = nn.Parameter(layer_scale_init * torch.ones(dim)) \ if layer_scale_init > 0 else None self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() def forward(self, x): inp = x # (B, H, W, C) x = self.dw(x.permute(0, 3, 1, 2)).permute(0, 2, 3, 1) x = self.norm(x) x = self.pw1(x.permute(0, 3, 1, 2)) x = self.act(x) x = self.grn(x.permute(0, 2, 3, 1)).permute(0, 3, 1, 2) x = self.pw2(x) x = x.permute(0, 2, 3, 1) if self.gamma is not None: x = x * self.gamma return inp + self.drop_path(x)

层缩放的初始值是个值得单独说的参数。原始监督预训练里用 1e-6，非常保守，目的是让残差分支在训练初期贡献几乎为零。但掩码预训练的场景不一样——预训练本身提供了很强的正则（掩码就是一种数据增广），而且预训练需要模型快速建立起有效的特征表达，1e-6 会让残差分支很长时间内学不动。实践中我会把预训练阶段的初始值放到 1e-4 量级，微调阶段再改回 1e-6 重新初始化。这个改动在我自己的小规模实验里带来了大约 0.6 到 1.2 个点的提升，不稳定但方向一致。当然这属于超参，不同数据规模下最优值会漂移，建议在小模型上先扫一遍。

还有一个实操细节：pw1和pw2的 bias 建议保留为 True，而深度卷积的 bias 设为 False（因为后面紧跟 LayerNorm，bias 会被消掉，纯属浪费）。在归零近似的稀疏实现里，pw1、pw2的 bias 要额外注意——如果 bias 非零，被 mask 位置的输出就不是严格的 0 了，必须在卷积后重新置零。所以稳妥的做法是：预训练阶段把所有卷积的 bias 全关掉，微调阶段再按常规配置打开。

5.3 预训练循环与关键超参

超参这块我给出一个可参考的配方，同时说明每一项背后的逻辑，方便你按自己的算力缩放。

batch 和学习率的线性缩放这条经验很硬，我试过 batch 从 512 缩到 128 但学习率不变，训练前期直接发散，loss 冲到几百然后一路平着走。反过来把学习率缩了但 batch 不缩，收敛慢得像蜗牛，200 epoch 才到别人 80 epoch 的位置。

warmup 千万别省。掩码预训练的前几百步梯度非常大，因为重建误差初始很高，直接上大学习率会让模型跳进一个很差的局部解，而且自己出不来——你会看到 loss 先降后猛涨然后卡死，调都调不回来。我一般用 10 到 20 个 epoch 做线性 warmup，小规模验证时可以压到 5 个 epoch，但至少要覆盖前 2000 个 step。

EMA 这块很多人会忽略。预训练的主权重在训练末期会有不小的波动，EMA 权重相当于对最近若干步做平滑，通常能比主权重高 0.3 到 0.8 个点的下游精度。实现上注意 EMA 的更新要在梯度同步之后做，别在optimizer.step()之前，否则会把这一步的梯度影响也算进去。另外 EMA 只对浮点参数做，不要把整型 buffer 也卷进去。

预训练循环主体框架如下，重点看损失那几行：

for epoch in range(total_epochs): model.train() for images in loader: images = images.cuda(non_blocking=True) target = patchify(images, 32) # (B, N, 3*32*32) target = per_patch_normalize(target) x, mask = random_masking(target, 0.6) pred = model(x, mask) # 只在可见位置前向 pred = unshuffle(pred, mask) # 重排回原 patch 顺序 loss = ((pred - target) ** 2)[mask].mean() # 只算被遮的 patch optimizer.zero_grad(set_to_none=True) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() ema.update(model) scheduler.step() if epoch % 20 == 0: visualize_reconstruction(model, val_sample) # 定期看重建效果

clip_grad_norm_那行是我自己加的，原始配方里不一定有。理由是掩码预训练早期个别 batch 的梯度范数会突然很大（尤其是随机遮到了全是平坦区域的 patch 时，重建误差会被放大），加个 1.0 的裁剪能明显减少训练中的尖峰。这个操作对最终精度没有负面影响，纯粹是稳定性保险。

另外强烈建议每 20 个 epoch 做一次重建可视化。这是最直观的诊断手段：正常情况应该能看出被遮区域的颜色、大致纹理、物体轮廓；如果重建出来是纯色块或者一片噪点，别继续跑了，先排查问题。我见过太多人盯着 loss 曲线跑了几十个小时，最后发现是掩码索引写反了。

5.4 显存与吞吐的估算方法

这部分给个估算框架，因为每个人的卡不一样，硬给数字没意义。

显存占用主要分四块：模型参数、梯度、优化器状态、激活值。以 Base 规模（约 89M 参数）为例，混合精度训练时，参数 fp16 占 178MB，fp32 主权重 356MB，AdamW 的一阶二阶动量各 356MB，梯度 178MB，加起来大约 1.4GB。这部分是固定的，跟 batch 无关。

大头在激活值。如果开梯度检查点，激活值大概能压到原来的三分之一到四分之一；不开的话，分辨率 224、batch 256 的 Base 模型，激活值能吃到 20GB 以上。所以个人环境上，梯度检查点基本是必开的，代价是训练速度慢 20% 到 30%。

吞吐方面，一个粗略的参考是：24G 显存的卡，128 分辨率、宽度的四分之一版本、batch 192、开梯度检查点，大概能跑到每秒 60 到 90 个样本。这个吞吐下，200 epoch 在 10 万张图上要跑六七个小时。如果你打算在百万级数据上跑，得按周算，这时候建议先在小数据上把流程和超参敲定，再一次性上大规模，别反复试。

6. 微调与下游迁移：权重怎么接、参数怎么调

6.1 stem 替换与权重对接

预训练和微调的模型结构只有一处不同：stem。预训练是 32×32 步长 32 的卷积加 LayerNorm，微调是 4×4 步长 4 的卷积加 LayerNorm，后面四个 stage 的结构完全一致。所以权重对接的做法是：加载预训练权重，忽略掉 stem 的那两个键，其余全部直接拷贝。

ckpt = torch.load("pretrain_ema.pth", map_location="cpu") state = ckpt["model"] msg = model.load_state_dict(state, strict=False) print("missing:", msg.missing_keys) # 应该只有 stem 相关的 2 个键 print("unexpected:", msg.unexpected_keys) # 应该为空

missing_keys里如果出现了 stem 之外的任何键，说明你的 stage 划分或者 block 数量跟预训练时不一致，这时候别急着训，先对齐结构。我遇到过一次是深度卷积的 bias 键，原因是预训练时关了 bias、微调时开了，这种形状不一致会被load_state_dict静默跳过（strict=False下），结果那一层的 bias 一直是随机初始化的值，微调精度莫名低两个点，查了半天。

提示：加载完权重之后，逐层打印每个参数的mean和std，如果某一层全是 0 或者全是一个常数，基本可以断定这层没加载上。这个检查比看missing_keys更可靠，因为形状不匹配的键在strict=False下不会报错。

微调阶段还有两个约定俗成的调整：层缩放的初始值改回 1e-6（因为现在是监督微调，不需要那么激进的残差），以及把所有卷积的 bias 打开。

6.2 分辨率迁移与多尺度特征

卷积主干在分辨率迁移上比 Transformer 舒服太多，这件事值得单独夸一下。因为 LayerNorm 沿通道维、GRN 沿通道维、池化和卷积都在局部空间操作，整个网络对输入的空间尺寸没有任何硬约束。224 训完直接喂 288 或者 384，前向不会报错，精度通常会涨一点（因为测试时可以用更高的分辨率提取更细的特征）。这一点在部署时很实用——同一个权重，端侧用小分辨率、云端用大分辨率，不用改模型。

但有个前提要注意：训练时如果全程固定 224，测试时突然上 384，性能提升可能不如预期，因为归一化的统计量虽然不受影响，但深度卷积 7×7 的感受野在图像中的相对比例变了。我一般的做法是微调的最后 10 个 epoch 里混入一部分 288 分辨率的样本，让模型适应一下尺度变化，测试时再用 288 就能稳定吃到收益。

多尺度特征输出也是现成的：四个 stage 的输出步长分别是 4、8、16、32，直接接到检测或者分割的 neck 上就能用，不需要像 Transformer 那样额外设计特征金字塔的接口。这也是 ConvNeXtV2 在很多下游任务里被选中的实际原因之一——迁移成本低。

6.3 学习率与分层衰减

微调的学习率比预训练低一个数量级，一般 5e-5 起步，batch 1024 的时候可以到 1e-4。AdamW，权重衰减 0.05，100 个 epoch，余弦退火，warmup 5 个 epoch。

分层学习率衰减这个技巧在卷积主干上效果比 Transformer 更明显。做法是给不同深度分配不同的学习率乘子，浅层给小的（乘 0.5 到 0.7），深层给接近 1 的。原因是浅层学到的是通用边缘纹理，这部分知识在预训练里已经很成熟，微调时大改反而破坏；深层学的是任务相关的语义组合，需要充分适配新任务。

我通常按 stage 分组，给四个 stage 分别配 0.5、0.65、0.8、1.0 的乘子，stem 给 0.3。这个配置在中等规模数据集（几万到几十万张）上比较稳。如果是极小的数据集（几千张），更好的策略是直接冻结前两个 stage，只训后两个 stage 加分类头，等 loss 稳定了再解冻全网络用更小的学习率微调。

还有一个反直觉的点：微调阶段的数据增强要比预训练弱。掩码预训练本身就是很强的增广，模型已经见过各种信息缺失的情况，微调时再上重度随机擦除、大比例裁剪，容易让模型"过度适应残缺输入"，在完整图像上的表现反而下降。我一般微调时只用随机裁剪加翻转，轻微的色彩抖动可选。

7. 常见问题与排查实录

7.1 症状到原因的速查表

踩过的坑整理成表，按症状查比较快。

7.2 三个我亲身踩过、文档里不会写的坑

坑一：γ 全零初始化导致的假死。GRN 的 γ 零初始化是对的，但如果优化器把 GRN 的参数放进了 weight decay 组，而 γ 初始是 0，那么梯度更新之后立刻被 weight decay 拉回 0 附近，永远长不起来。表面看训练正常、loss 也在降，但 GRN 实际上完全没起作用，等于没加。排查方法是训几个 epoch 后打印gamma.abs().mean()，正常的曲线应该从 0 稳步爬升到 0.1 往上。解决也简单：把 GRN 的 γ、β 和所有 LayerNorm 的参数从 weight decay 组里排除，单独放一组weight_decay=0。这个坑我怀疑坑过不少人，因为它的症状和"GRN 本来就没用"一模一样。

坑二：稀疏卷积的坐标顺序与稀疏张量的内存布局混淆。用子流形稀疏卷积实现的时候，坐标张量的顺序必须严格是 (batch, z, y, x) 这类固定格式，而特征张量是 (N, C)。我一开始把特征做成了 (C, N) 的通道优先布局，算子库不报错，但输出的值完全不对——因为它在错误的维度上做了运算。这类问题几乎不可能从 loss 曲线上看出来，只能通过对比稀疏前向和稠密前向在"无掩码"情况下的输出是否完全一致来验证。在引入稀疏算子后，第一件事永远是做一次无掩码的一致性对比，误差应该在 1e-5 以内。这个习惯我后来一直保留着。

坑三：EMA 的实现把 buffer 也更新了。有些模型里 LayerNorm 的running_mean之类是 buffer（虽然 LayerNorm 本身没有，但如果你混用了 BN 或者别的统计层就有）。EMA 更新时如果不筛选requires_grad的参数，会把 buffer 也做滑动平均，导致推理时的归一化统计和训练时对不上，精度掉得很突然。正确做法是只对model.parameters()里requires_grad=True的部分做 EMA，或者在保存时只保存参数部分、重建时再补 buffer。

7.3 一个诊断特征坍塌的实用小工具

最后分享一个我自己写的诊断脚本，五分鐘就能判断预训练有没有走上歪路。思路是给指定 stage 的输出挂一个 hook，累积若干 batch 的激活值，然后算两个指标：一是通道间的平均余弦相似度（越接近 1 越危险），二是死亡通道比例（通道方差小于阈值的占比）。

@torch.no_grad() def diagnose(model, loader, stage_idx=2, n_batch=8): feats = [] def hook(m, i, o): feats.append(o.detach().float().mean(dim=(2, 3))) # (B, C) h = model.stages[stage_idx].register_forward_hook(hook) for i, (x, _) in enumerate(loader): model(x.cuda()) if i >= n_batch: break h.remove() f = torch.cat(feats, 0) # (N, C) f = F.normal