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