本文第一作者 Haiyu Wu 博士畢業(yè)于美國頂級(jí)私立學(xué)府圣母大學(xué)，現(xiàn)在就職于 Altos Labs —— 一家擁有 30 億美元啟動(dòng)資金，專注研究細(xì)胞重激活的初創(chuàng)公司。他的研究方向?yàn)椋澜缒Ｐ秃投嗄B(tài)模型的穩(wěn)定性以及泛化能力。其余兩位作者分別是布朗大學(xué)教授，LeJEPA/SIGReg 的第一作者 Randall Balestriero 和前耶魯大學(xué)教授，現(xiàn) Altos Labs 的副總裁，Morgan Levine。
近日，自監(jiān)督學(xué)習(xí)新工作VISReg（Variance-Invariance-Sketching Regularization）獲圖靈獎(jiǎng)得主 Yann LeCun 連續(xù)轉(zhuǎn)發(fā)并給予高度認(rèn)可 —— 他在轉(zhuǎn)發(fā)時(shí)評(píng)價(jià)道「VICReg begat SIGReg which begat VISReg」（VICReg 孕育了 SIGReg，SIGReg 又孕育了 VISReg），一句話點(diǎn)明了這條正則化路線的技術(shù)傳承。能獲得 LeCun 如此認(rèn)可，VISReg 究竟強(qiáng)在哪里？答案在于：它精準(zhǔn)命中了 LeCun 長期押注的 JEPA 世界模型的核心難題 ——表征坍塌（representation collapse）。VISReg 將防止坍塌的正則項(xiàng)解耦為「尺度」與「形狀」兩個(gè)獨(dú)立目標(biāo)，在不依賴任何啟發(fā)式訓(xùn)練技巧、也不依賴海量數(shù)據(jù)的前提下，于 15 個(gè)數(shù)據(jù)集上綜合表現(xiàn)超過 7 種主流自監(jiān)督學(xué)習(xí)方法；其中僅用約 1/10 的訓(xùn)練數(shù)據(jù)，即在分布外（OOD）基準(zhǔn)上追平 DINOv2。
【資料圖】
- 論文標(biāo)題：VISReg: Variance-Invariance-Sketching Regularization for JEPA training
- 論文地址：https://arxiv.org/abs/2606.02572
- 代碼 / 預(yù)訓(xùn)練權(quán)重：https://github.com/HaiyuWu/visreg
- 項(xiàng)目主頁：https://haiyuwu.github.io/visreg/
一、研究背景：表征坍塌與正則化的兩難
JEPA 世界模型的底層是 Yann LeCun 自 2017 年起持續(xù)倡導(dǎo)的自監(jiān)督學(xué)習(xí)（Self-Supervised Learning, SSL）。SSL 無需人工標(biāo)注即可從海量數(shù)據(jù)中學(xué)習(xí)通用表征，但普遍面臨一個(gè)核心難題 ——表征坍塌（representation collapse）：模型傾向于把不同輸入映射到相同或極少數(shù)幾個(gè)向量上，看似完成了訓(xùn)練，實(shí)則未學(xué)到有判別力的表征。
為抑制坍塌，主流方法大多依賴一系列啟發(fā)式技巧（EMA、教師 - 學(xué)生網(wǎng)絡(luò)、停止梯度、凍結(jié)層等）。這些技巧使訓(xùn)練變得脆弱、難以調(diào)參，也削弱了方法的可解釋性與可擴(kuò)展性。
另一條路線是通過正則項(xiàng)直接約束表征分布。LeCun 團(tuán)隊(duì)提出的VICReg將學(xué)習(xí)目標(biāo)拆為方差、不變性、協(xié)方差三項(xiàng)，用協(xié)方差約束各維度之間的相關(guān)性；但協(xié)方差僅刻畫二階統(tǒng)計(jì)量，無法區(qū)分「均值、方差相同，而分布形狀迥異」的兩種表征。其后提出的 SIGReg 基于 Cramér–Wold 定理，用 sketching 技術(shù)將整個(gè)嵌入分布對(duì)齊到標(biāo)準(zhǔn)高斯，從而約束完整的分布形狀。然而SIGReg仍存在兩個(gè)關(guān)鍵缺陷：
- 坍塌時(shí)梯度消失：當(dāng)表征開始坍塌時(shí)，SIGReg 的梯度隨之衰減 —— 坍塌越嚴(yán)重、修正信號(hào)越弱，模型難以自行恢復(fù)；
- 尺度與形狀耦合：未將「幅度大小（尺度）」與「分布形態(tài)（形狀）」兩個(gè)獨(dú)立屬性分離，二者在優(yōu)化中相互干擾，導(dǎo)致在長尾、低質(zhì)量、低秩數(shù)據(jù)上適配性較差。
也就是說，在模型最需要梯度信號(hào)來逃離坍塌狀態(tài)時(shí)，SIGReg 的梯度恰恰趨近于消失。
這正是 VISReg 要解決的核心問題。
二、核心方法：將正則項(xiàng)解耦為尺度與形狀
VISReg 對(duì) VICReg 與 SIGReg 取長補(bǔ)短：保留 VICReg 的方差項(xiàng)來控制尺度，同時(shí)用基于切片 Wasserstein 距離（Sliced Wasserstein Distance, SWD）的 sketching 目標(biāo)替代協(xié)方差項(xiàng)來控制形狀，并通過停止梯度將二者徹底解耦。整個(gè)正則目標(biāo)由三部分組成。
1. 尺度正則（Scale Regularization）
第一部分約束每一維的方差，防止幅值坍縮：
其關(guān)鍵性質(zhì)在于：當(dāng)模型坍縮時(shí)，該項(xiàng)的梯度趨近于一個(gè)常數(shù)，從而保證模型能夠穩(wěn)定地恢復(fù) —— 這恰好彌補(bǔ)了 SIGReg 梯度消失的缺陷。
2. 形狀正則（Shape Regularization）
第二部分先歸一化以消除尺度影響，再單獨(dú)約束形狀。關(guān)鍵的一步是帶停止梯度（stop-gradient, sg）的歸一化：
這里對(duì)標(biāo)準(zhǔn)差 σ 施加停止梯度，使得形狀損失的優(yōu)化不會(huì)反過來改變尺度—— 這正是「尺度」與「形狀」兩個(gè)目標(biāo)真正解耦、互不干擾的機(jī)理所在。
歸一化之后，再用切片 Wasserstein 距離將分布的幾何形狀對(duì)齊到各向同性高斯：
其中 q_N 為標(biāo)準(zhǔn)高斯分位數(shù)，w_k 為隨機(jī)投影方向（即「切片 /sketching」）。其理論依據(jù)是Cramér–Wold 定理（論文 Lemma 3.1）：* 兩個(gè)分布相等，當(dāng)且僅當(dāng)它們沿單位球面上所有方向的一維投影都相等 *。因此，只要把高維表征沿足夠多的隨機(jī)一維方向切片后逐一對(duì)齊到高斯，就等價(jià)于在高維空間對(duì)齊了整個(gè)分布 —— 這使得可以用廉價(jià)的一維排序操作刻畫完整的分布形狀，而非僅僅二階統(tǒng)計(jì)量。
3. 合并目標(biāo)
第三部分是一個(gè)將 batch 均值 μ 拉向原點(diǎn)的中心化損失 L_center = ‖μ‖_2^2。三個(gè)正則項(xiàng)按權(quán)重組合：
預(yù)測損失沿用 JEPA / LeJEPA 的不變性目標(biāo) —— 讓各視角（global + local，共 V 個(gè)）的嵌入 z_i 都向全局視角的均值 μ_g 對(duì)齊：
最后用單一超參 λ 在預(yù)測與正則之間平衡，得到完整目標(biāo)：
與 VICReg 的對(duì)比：VICReg 同樣將正則解耦為方差 + 協(xié)方差，但協(xié)方差只刻畫二階統(tǒng)計(jì)量；VISReg 用基于切片 Wasserstein 的 sketching 目標(biāo)完整刻畫了分布形狀，同時(shí)保留方差項(xiàng)做尺度控制 —— 既保留了 VICReg 的靈活性，又獲得了分布層面的嚴(yán)格性。
僅需約 15 行 PyTorch 代碼
該正則目標(biāo)在實(shí)現(xiàn)上非常輕量，核心邏輯只需約 15 行：
計(jì)算復(fù)雜度與擴(kuò)展性
在計(jì)算與擴(kuò)展性上，VISReg 同樣具備優(yōu)勢。其正則部分的計(jì)算復(fù)雜度為 O (NDK)（N 為 batch、D 為維度、K 為切片數(shù)），對(duì)所有擴(kuò)展因子都是線性的；相比之下，VICReg 的協(xié)方差項(xiàng)為 O (ND^2)，隨維度平方增長。在同等 batch 規(guī)模下，VISReg 在單塊 H100 GPU 上的運(yùn)行速度與顯存占用均優(yōu)于 SIGReg。
更重要的是，K 個(gè)隨機(jī)切片可以分?jǐn)偟蕉鄩K GPU 上：在 M 塊 GPU 上每塊各生成 K/M 個(gè)切片，效果等價(jià)于單卡生成全部 K 個(gè)。實(shí)驗(yàn)中，當(dāng)單卡切片數(shù)不足時(shí)，改用8 卡、每卡 128 個(gè)切片（合計(jì) 1024），即可把與「單卡 1024 切片」之間的精度差距從約2.4% 縮小到 0.22%。這意味著擴(kuò)大訓(xùn)練規(guī)模時(shí) K 可保持常數(shù)，幾乎不增加單卡負(fù)擔(dān)。
圖：在固定 K 與 D 時(shí)，增加 GPU 數(shù)量帶來的線性探測精度變化。當(dāng) K 不足（K = ?D）時(shí)，用 8 倍的 GPU 數(shù)量即可把精度補(bǔ)齊到 K = 2D 的水平 —— 這使得在大規(guī)模訓(xùn)練中保持常數(shù) K 成為可能。
三、實(shí)驗(yàn)結(jié)果
回到標(biāo)題的問題 ——VISReg 到底強(qiáng)在哪里？研究團(tuán)隊(duì)在15 個(gè)數(shù)據(jù)集（8 個(gè)域內(nèi) + 6 個(gè)分布外 + ADE20K 稠密預(yù)測）上，將 VISReg 與 MoCoV3、DINO、iBOT、I-JEPA、MAE、data2vec 等 7 種主流自監(jiān)督方法進(jìn)行了對(duì)比，場景涵蓋天文、醫(yī)療、遙感、紋理、花卉等。答案體現(xiàn)在從識(shí)別到分割、生成的多個(gè)維度上。
1. 域內(nèi)（In-Domain）線性探測
為保證比較公平，實(shí)驗(yàn)按是否使用啟發(fā)式技巧分為兩組。在不使用任何啟發(fā)式技巧的一組中，VISReg 領(lǐng)先：ViT-B/16 的域內(nèi)線性探測精度達(dá)75.7%，高于 MAE（75.1%）；ViT-L/14 進(jìn)一步提升至77.0%，高于 LeJEPA（75.6%）。與使用啟發(fā)式技巧的 iBOT、DINO 相比，VISReg 在常規(guī)數(shù)據(jù)集上僅略低，但在紋理數(shù)據(jù)集DTD上反超全部方法 —— 這表明其跨域泛化能力源于方法本身，而非人工技巧的堆疊。
2. 分布外（OOD）泛化：全面最優(yōu)
分布外泛化是比域內(nèi)精度更嚴(yán)格的檢驗(yàn)：依賴啟發(fā)式的方法常在 ImageNet 域內(nèi)被充分調(diào)優(yōu)，卻未必能遷移到差異較大的新分布。研究團(tuán)隊(duì)在覆蓋醫(yī)療（ChestXRay、RetinaMNIST、OrganAMNIST）、天文（Galaxy10）、遙感（AID）、紋理（DTD） 的 6 個(gè) OOD 數(shù)據(jù)集上評(píng)測，這些數(shù)據(jù)集與 ImageNet 訓(xùn)練域完全無關(guān)。結(jié)果顯示，VISReg 在所有方法、所有骨干規(guī)模上都取得了最佳的平均 OOD 精度，甚至超過部分使用啟發(fā)式技巧、且骨干更大的方法。
圖 4：平均 OOD 線性探測精度。VISReg 全面優(yōu)于 iBOT、DINO、MoCoV3、I-JEPA、MAE、data2vec 等方法。
如圖 4 所示，ViT-B/16 的 VISReg 平均 OOD 精度為 70.19%，ViT-L/14 為 70.63%，明顯高于 MAE（67.85%），并優(yōu)于 MoCoV3（69.46%）、DINO（69.56%）、I-JEPA（68.55%）等方法。
3. 數(shù)據(jù)效率：以 1/10 數(shù)據(jù)比肩 DINOv2
將 VISReg（ViT-L/14）在 ImageNet-22K（約 1400 萬張圖像）上預(yù)訓(xùn)練后，其 6 個(gè) OOD 數(shù)據(jù)集的平均精度達(dá)72.94%，與在10 倍規(guī)模的 LVD-142M（1.42 億張圖像） 上訓(xùn)練的 DINOv2（72.93%） 基本持平。也就是說，VISReg 以約1/10的數(shù)據(jù)達(dá)到了同等水平。（作為對(duì)照，同為 ViT-L/14、但僅用 ImageNet-1K 預(yù)訓(xùn)練的 VISReg 平均精度為 70.63%。）這說明其學(xué)到的表征具有很強(qiáng)的通用性。
圖 5：在 ImageNet-22K 上預(yù)訓(xùn)練的 VISReg，在 OOD 基準(zhǔn)上比肩用 10 倍數(shù)據(jù)（LVD-142M）訓(xùn)練的 DINOv2。
4. 遷移微調(diào)：全面超過 DINO
盡管 VISReg 在部分域內(nèi)數(shù)據(jù)集上的線性探測精度略低于 DINO，但經(jīng)過微調(diào)后，它在 CIFAR-10、CIFAR-100、Flowers、ImageNet-1K、Galaxy10 全部五個(gè)數(shù)據(jù)集上均超過 DINO 與有監(jiān)督預(yù)訓(xùn)練—— 這表明其表征分布更均勻、冗余更低、可遷移性更強(qiáng)。
圖：遷移學(xué)習(xí)對(duì)比。微調(diào)后，VISReg 在所有測試數(shù)據(jù)集（CIFAR-10、CIFAR-100、Flowers、ImageNet-1K、Galaxy10）上都優(yōu)于 DINO 與有監(jiān)督預(yù)訓(xùn)練。
5. 稠密預(yù)測與生成引導(dǎo)
VISReg 的優(yōu)勢不局限于分類。在ADE20K 線性語義分割上（ViT-B/16），其 mIoU 為，高于 DINO（29.40）與 MAE（23.60），僅次于 MoCoV3（31.69）；在不使用任何啟發(fā)式技巧的前提下，這一結(jié)果具有競爭力。論文亦坦言，稠密預(yù)測與最佳方法仍有差距，是后續(xù)優(yōu)化的重點(diǎn)。
圖 7：ADE20K 上的線性語義分割。在不使用任何啟發(fā)式技巧的情況下，VISReg 取得了具有競爭力的 mIoU，僅次于 MoCoV3。
生成引導(dǎo)上（SiT-B/2，iREPA 框架，10 萬步訓(xùn)練），由 VISReg 特征引導(dǎo)的生成在四項(xiàng)指標(biāo)中的三項(xiàng)優(yōu)于 DINO：gFID（DINO 41.15）、Precision（DINO 50.51）、Recall（DINO 60.70），IS 基本持平（33.48 vs 33.47）。這說明 VISReg 學(xué)到的表征作為生成引導(dǎo)信號(hào)同樣更優(yōu)。
圖 8：使用 SiT-B/2、分別由 VISReg 與 DINO 特征引導(dǎo)的圖像生成。VISReg 在多數(shù)指標(biāo)上都提供了更好的引導(dǎo)（更低的 gFID、更高的 Precision 與 Recall）。
6. 低質(zhì)量數(shù)據(jù)上的魯棒性
在長尾分布（ImageNet-LT）與低秩（Galaxy10） 等低質(zhì)量數(shù)據(jù)集上，VISReg 能穩(wěn)定地防止坍塌并學(xué)到有意義的表征，而 DINO 在缺乏精細(xì)調(diào)參時(shí)直接失敗。
表 1：ImageNet-LT 上的線性探測精度（ViT-S/8，從頭訓(xùn)練 400 epoch；* 表示增大形狀損失的權(quán)重）
DINO 在長尾數(shù)據(jù)上幾近完全失敗（Overall 僅 5.13%），而 VISReg* 取得了全面最優(yōu)。
表 2：Galaxy10 上的域內(nèi)線性探測精度（從頭訓(xùn)練，測試低秩任務(wù)；* 表示增大形狀損失的權(quán)重）
SIGReg、SWD、VISReg 都能成功避免訓(xùn)練坍縮并取得良好精度，而 DINO 難以學(xué)到有意義的特征。
四、結(jié)論
VISReg 表明：將表征正則解耦為「尺度」與「形狀」兩個(gè)獨(dú)立組件，可以得到一種比現(xiàn)有方法更穩(wěn)定、更高效、泛化性更強(qiáng)的自監(jiān)督學(xué)習(xí)方法。在不使用任何訓(xùn)練啟發(fā)式技巧的前提下，它在圖像識(shí)別、分割與生成引導(dǎo)等多個(gè)維度上取得了領(lǐng)先或接近領(lǐng)先的結(jié)果，并以約 1/10 的數(shù)據(jù)達(dá)到了 DINOv2 的 OOD 水平。這為 JEPA 世界模型長期存在的表征坍塌問題提供了一種新的正則化解法。
