SiamJEPA:论 JEPA 中孪生学生编码器的作用 SiamJEPA: On the Role of Siamese Student Encoders in JEPA
提出SiamJEPA,首次系统揭示JEPA中孪生学生编码器作为有效正则化器,提升表示质量并加速收敛。
前置知识
JEPA(联合嵌入预测架构)
Joint Embedding Predictive Architecture 是一类自监督表示学习方法,不重建像素,而是在潜空间用一个预测器网络去预测被掩码区域的潜表示。代表方法 I-JEPA(图像)和 V-JEPA(视频)通过 stop-gradient 和 EMA 教师网络防止表示坍缩。
本文的研究对象就是 JEPA 框架,理解其「用潜空间预测代替像素重建」的核心思想是读懂全文的前提。
EMA 教师网络
Exponential Moving Average 教师网络是指教师编码器的参数不通过梯度更新,而是按动量公式 $\theta_{ema} \leftarrow m \cdot \theta_{ema} + (1-m)\cdot\theta$ 对学生参数做指数滑动平均。配合 stop-gradient 可防止表示坍缩,是 BYOL、DINO、I-JEPA 的标配。
SiamJEPA 的教师编码器 $f_{ema}$ 就是 EMA 网络,提供预测目标;其动量调度(0.99→0.999→0.9999)是关键训练细节。
孪生编码器(Siamese Encoders)
孪生编码器指两个共享权重、但接收不同输入视图的编码器,输出在同一潜空间中对齐。受大脑海马体-新皮层回路启发的 PhiNet 采用孪生学生编码器建模「时间预测假设」,SiamMAE、CropMAE、RSP 也在视频任务中使用孪生结构。
本文的核心创新正是把「是否使用孪生学生编码器」作为研究变量,读者必须先明白孪生结构与单编码器的区别。
表示坍缩与正则化
表示坍缩指编码器输出退化为常量、丢失信息的失败模式。联合嵌入方法常用对比学习、stop-gradient、EMA、方差-协方差正则(VICReg)、KL 散度加 free-bit 等手段防止坍缩。free-bit 通过对 KL 设置下界阈值,允许两分布保留最低限度的差异。
SiamJEPA 用 KL 散度对齐两个孪生分支并用 free-bit 防止完全坍缩,理解这些防坍缩机制才能理解 λKL 的作用。
ImageNet 线性探针
线性探针是评估预训练表示质量的协议:冻结预训练骨干网络,在其输出上只训练一个线性分类器,报告 Top-1 准确率。它衡量表示的「线性可分性」,是 SSL 论文的标准评测。本文还比较了 CLS token、第 10/12 层 mean pooling 等不同表示来源。
全文几乎所有实验指标都是 ImageNet 线性探针 Top-1 准确率,读者需理解该协议才能判断 63.4%/70.2% 等数字的含义。
研究动机
JEPA(联合嵌入预测架构)作为自监督表示学习的新兴范式,通过在潜空间预测被掩码区域的表示而非重建像素来学习表示,代表方法 I-JEPA 和 V-JEPA 已展现出优于重建式方法(如 MAE)的训练效率。然而,现有 JEPA 模型沿袭 BYOL 和 DINO 的学生-教师范式,仅在学生网络中使用单一编码器。与此同时,受大脑海马体和新皮层生物回路启发的 PhiNet 框架采用了孪生学生编码器(两个共享权重但输入不同的编码器)配合 EMA 教师网络,在抗坍缩和持续学习上表现出优势,Siamese MAE、CropMAE、RSP 等方法也在视频任务中广泛使用孪生编码器。但关键问题是:尽管孪生编码器在自监督学习中被越来越多地采用,它们在 JEPA 这类纯潜空间预测架构中究竟扮演什么角色,至今缺乏系统研究。具体而言,孪生编码器是否能提升 JEPA 的表示质量、是否会加速收敛、其作用机理如何,都是未解之谜。
本文的目标是本文的目标是系统揭示孪生学生编码器在 JEPA 框架中的作用机理。为此作者提出 SiamJEPA——一种基于孪生学生 Transformer 并配备 EMA 教师网络的掩码潜空间预测架构,可视为 PhiNet 架构的 JEPA 化版本(即 PhiNetv2 的掩码输入变体)。通过引入一个正则化参数 $\lambda_{KL}$ 来连续控制孪生编码器的影响,使传统单编码器 JEPA 成为特例(当 $\lambda_{KL}=0$ 时退化为单编码器 JEPA),从而在统一框架内系统比较两者。作者在 ImageNet 线性探针基准上用 ViT-Base 进行全面消融研究,从 KL 正则化权重、free-bit 阈值、掩码策略、权重衰减、预测目标和预测器深度等多维度剖析孪生编码器的贡献,定性与定量地回答其究竟是有效的归纳偏置,还是无关紧要的架构选择。
与已有工作不同的是,本文的独特切入角度在于:它没有追逐 SOTA 性能,而是把孪生编码器这个常被当作默认架构选择的组件当作科学问题来审视。与之前在像素/特征重建层面使用孪生编码器(SiamMAE、CropMAE)或在 PhiNet 视频预测中使用不同,本文首次将其置于纯潜空间预测(JEPA)的图像表示学习场景下,并通过可调的 $\lambda_{KL}$ 实现了单编码器 JEPA 到孪生 JEPA 的连续插值,使两者对比不受实现差异干扰。此外,作者从训练效率角度切入,强调在有限训练预算下潜空间预测相对重建式方法的样本效率优势,并从表示可分性、早期收敛速度等机理层面给出解释,填补了架构相似但角色未明这一知识空白。
核心方法
SiamJEPA 的整体直觉是:让两个共享权重的孪生编码器各自处理同一图像的两个不重叠掩码视图,再用一个 EMA 教师编码器提供完整图像的表示作为预测目标,从而在对齐两个视图全局表示和预测被掩码区域潜表示双重目标下学习。技术路线上,输入图像 $X\in\mathbb{R}^{m\times d_{in}}$ 被分块后,学生编码器输出 $H^{(1)}=f(\mathrm{Mask}(X,M_1))$ 和 $H^{(2)}=f(\mathrm{Mask}(X,M_2))$,其中两个掩码集合 $M_1\cap M_2=\emptyset$ 互不相交以防止捷径学习;教师编码器输出 $Y=f_{ema}(X)$。每个编码器输出被分解为 CLS token 与 patch token。随后用一个线性预测器 $h$ 对齐两视图的全局 CLS 表示,用一个 Transformer 预测器 $g$ 配合概率化的先验/后验潜变量 $Z$ 来预测被掩码 patch 的潜表示。最终用 Sim-1(KL 散度)和 Sim-2(归一化均方误差)两个损失联合优化。
核心创新点在于用 KL 正则化把两个孪生学生编码器的输出拉到一致,并使该约束的强度连续可调。具体地,预测器 $g$ 采用概率公式:后验 $q(Z\mid[h(H_{cls}^{(1)}),H_{cls}^{(2)}])$ 同时看两视图,先验 $p(\hat{Z}\mid h(H_{cls}^{(1)}))$ 只看单视图,最小化两者 KL 散度使单视图推断逼近双视图推断,并通过对先验分支施加 stop-gradient 提升稳定性。这与普通 JEPA(单学生编码器直接预测教师目标)的本质区别是:孪生编码器加 KL 约束引入了额外的归纳偏置,迫使编码器丢弃视图特异信息、保留跨视图共享信息。通过设置 $\lambda_{KL}=0$ 可退化为单编码器 JEPA 变体(作者称之为 JEPA-like method,取 $\lambda_{KL}=10^{-4}$),从而在同一实现框架内干净地隔离出孪生这一因素的作用。
方法步骤详情
完整步骤:(1) 图像分块为 $X$;(2) 用 block 或 random 掩码,按 0.7/0.75/0.8 比从可见区域不相交采样 $M_1$、$M_2$($M_1\cap M_2=\emptyset$);(3) 共享权重的孪生编码器分别处理 $\mathrm{Mask}(X,M_1)$、$\mathrm{Mask}(X,M_2)$ 得 $H^{(1)}$、$H^{(2)}$;(4) EMA 教师对完整 $X$ 输出目标 $Y$;(5) 线性预测器 $h$ 处理 CLS,构造后验 $q$ 与先验 $p$,算带 stop-gradient 的 KL 散度(Sim-1);(6) Transformer 预测器 $g$ 结合潜变量 $Z$ 预测掩码 patch 表示,算归一化 MSE(Sim-2);(7) 总损失 $\mathcal{L}{=}\tfrac12(MSE^{(1)}{+}MSE^{(2)}){+}\tfrac{\lambda_{KL}}2(KL_{sg}^{(1)}{+}KL_{sg}^{(2)})$,用 free-bit 防坍缩;(8) EMA 动量按 0.99→0.999→0.9999 分阶段提升。
技术新颖性
技术新颖性体现在三点。第一,首次将 PhiNet 的孪生学生加 EMA 教师范式从视频帧预测迁移到图像 JEPA,并通过正则化参数 $\lambda_{KL}$ 把单/双编码器 JEPA 统一到同一损失族中,使孪生成为可连续调节而非 0/1 开关的设计维度。第二,预测器 $g$ 借鉴 Jang 等人的概率先验/后验设计,使潜变量 $Z$ 建模给定上下文时表示的不确定性,再配以 free-bit 机制防止两分布完全坍缩为同一,这是对标准 JEPA 确定性预测的灵活化。第三,作者发现 SiamJEPA 仅需 1 层浅 Transformer 预测器即可工作良好,远浅于常用 JEPA 预测器,提示了参数效率上的潜在优势。整体而言,新颖性不在于刷分,而在于把一个被忽视的架构选择提炼为可解释、可调控的正则化机制。
实验结果
主对比(Table 2):ViT-Base、ImageNet 线性探针上 SiamJEPA 训练 100/200/300/400 epoch 达 63.4%/67.8%/69.6%/70.2% Top-1,而 MAE 400 epoch 仅 61.9%、1600 epoch 才 68.0%,CAE 1600 epoch 达 70.4%——SiamJEPA 用不到四分之一 MAE 的 epoch 超越 MAE、追平 CAE。KL 正则化(Table 3):$\lambda_{KL}=0.0001$ 时 epoch300 仅 66.78%,$\lambda_{KL}=0.01$ 达 69.30% 且 100 epoch 就达 63.72%,证明 KL 既提升质量又加速收敛。掩码策略(Table 5):block 掩码 0.75 比达 69.30%,优于 random 的 65.10%。权重衰减(Table 6):0.1 在 400 epoch 达 70.15%,优于 0.05 的 69.33%。预测目标(Table 7):第 10 层 mean pooling 达 69.30%,优于 CLS 的 67.39%。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| ImageNet 线性探针(ViT-Base,400 epoch) | Top-1 准确率 | SiamJEPA 70.2% | MAE 400 epoch 61.9% | +8.3 个百分点,且用不到 MAE 一半算力 |
| ImageNet 线性探针(与 CAE 对比) | Top-1 准确率 | SiamJEPA 70.2%(400 epoch) | CAE 70.4%(1600 epoch) | 用四分之一 epoch 追平 CAE |
| KL 正则化消融(epoch 300) | Top-1 准确率 | $\lambda_{KL}=0.01$ 时 69.30% | $\lambda_{KL}=0.0001$ 时 66.78% | +2.52 个百分点,证明孪生正则有效 |
| 掩码策略消融(epoch 300,0.75 比) | Top-1 准确率 | Block 掩码 69.30% | Random 掩码 65.10% | +4.20 个百分点 |
局限与改进
作者明确承认:实验规模有限,仅用 ViT-Base,未扩展到 ViT-Large/Huge,扩展性未知;最终性能对实现细节和超参高度敏感,仍有大幅提升空间;仅做了图像预训练,未涉及视频;与 I-JEPA 对比不直接(I-JEPA 训练 600 epoch、本文 400 epoch),SiamJEPA 最终 70.2% 低于 I-JEPA 的 72.9%;预测器深度的研究因算力不足而未完成。我的额外观察:论文未提供与不带 EMA、不带概率预测器的更朴素孪生 baseline 的对比,难以完全剥离各组件贡献;所有消融在单一数据集 ImageNet 上完成,跨数据集泛化能力(如检测、分割、视频)缺乏验证;论文未开源代码,且自述基于 MAE codebase 而非 I-JEPA 官方实现,复现门槛和可比性存疑;free-bit 和 $\lambda_{KL}$ 的最优值在不同 epoch 不一致,实际调参成本可能较高。
独立分析的弱点
弱点一:实验仅在 ImageNet 线性探针单一指标上验证,未报告下游检测/分割/视频任务,孪生编码器的通用表示优势证据不足——改进方向是补充 COCO 检测、ADE20K 分割等密集预测评测。弱点二:组件解耦不彻底,KL 正则、概率预测器、stop-gradient、EMA、不相交掩码五者同时引入,缺乏逐一去除的消融,难以归因——改进方向是设计更细粒度的因子实验。弱点三:与 I-JEPA 非公平对比(不同 epoch、不同 codebase),说服力受限——改进方向是在同一 codebase、相同 epoch 下复现 I-JEPA 单编码器版本直接对决。弱点四:论文声称浅预测器(1-2 层)足够,却未给出参数量/FLOPs 对比,工程价值存疑——改进方向是补充效率分析。弱点五:缺少表示坍缩的定量监控(如特征秩、有效维度),无法直接佐证正则化器的机理解释——改进方向是加入表示几何分析。
未来方向
作者提出的方向:扩展到 ViT-Large/Huge 验证规模可扩展性;系统化超参与架构精修以缩小与 I-JEPA 差距;迁移到视频表示学习(SiamMAE、PhiNetv2 已显示孪生编码器在视频上有效,近期 Daithankar 2026、Rao 2026 已开始将孪生设计引入视频预测);探索更多神经科学启发机制。基于成果可延伸的方向:将 $\lambda_{KL}$ 的连续调控思想迁移到 V-JEPA、LeJEPA 等 JEPA 变体,统一研究编码器数量这一维度;把概率先验/后验加 free-bit 的坍缩防护与 LeJEPA 的 SIGReg(朝各向同性高斯正则化)结合,比较两种坍缩防护机理;在第 10 层表示最优这一发现启发下,研究分层表示蒸馏或动态选择最优层;鉴于浅预测器可行,探索极简预测器能否进一步降低 JEPA 训练成本。
复现评估
复现评估中等偏难。有利因素:作者基于开源的官方 MAE codebase 实现,骨干网络为标准 ViT-Base,数据集为公开 ImageNet-1K,关键超参(batch size 8192、学习率 $1.5\times10^{-4}$、EMA 动量 0.99→0.999→0.9999、掩码比、$\lambda_{KL}$、free-bit)在 Table 1 及正文给出较完整记录,消融表格齐全。不利因素:论文截至发表未见明确代码/权重开源链接(正文仅引用 MAE codebase),概率预测器 $h/g$ 的精确结构、projector head 配置需从 PhiNet/PhiNetv2 文献间接推断;训练需 NVIDIA V100/A100/H100,effective batch 8192 对算力要求高,单卡难以复现;作者自述性能对实现细节高度敏感,意味着小的工程差异可能导致结果波动;EMA 动量分阶段调度、不相交掩码采样等细节若实现不当易影响收敛。综合判断:有经验团队可在中等算力上大致复现趋势,但精确复刻数值有难度。
论文图表
对比 MAE 与 SiamJEPA 的预训练与线性探针配置。两者骨干均为 ViT-Base,数据集 ImageNet-1K,输入 224²,优化器 AdamW,基础学习率 $1.5\times10^{-4}$,预训练 400 epoch。关键差异:解码器深度 MAE 为 8、SiamJEPA 仅 1;有效 batch size MAE 4096、SiamJEPA 8192;掩码比 MAE 固定 0.75、SiamJEPA 为 {0.7,0.75,0.8};EMA 仅 SiamJEPA 使用({0.99,0.999,0.9999});权重衰减 SiamJEPA 为 {0.05,0.1}。
明确实验设置,是解读后续所有结果表的基准,也揭示了 SiamJEPA 用更浅预测器和更大 batch 的工程选择。