現(xiàn)全記錄：ViT-H預(yù)訓(xùn)練72小時(shí)達(dá)到73.3%線性探測(cè))
讀論文和親手復(fù)現(xiàn)完全是兩回事。尤其是 I-JEPA 這種看起來(lái)“公式簡(jiǎn)單、邏輯清晰”的自監(jiān)督方法跑通一個(gè) demo 不難但要把 ViT-H/14 這種體量的模型在有限算力下預(yù)訓(xùn)練 72 小時(shí)并把 ImageNet 1% 線性探測(cè)做到 73.3%背后全是對(duì)細(xì)節(jié)的理解和一堆不成文的“土辦法”。這篇文章把我這次復(fù)現(xiàn) I-JEPA 架構(gòu)的完整過(guò)程、關(guān)鍵決策和踩過(guò)的坑一次講清楚希望能幫到正在搞自監(jiān)督預(yù)訓(xùn)練、或者準(zhǔn)備復(fù)現(xiàn)論文的同學(xué)少走點(diǎn)彎路。整個(gè)項(xiàng)目做完我的最大感受是I-JEPA 不是那種“論文很強(qiáng)、落地翻車”的方法——只要把對(duì)稱編碼器的收斂節(jié)奏、目標(biāo)編碼器的更新策略、還有數(shù)據(jù)加載的瓶頸處理好它確實(shí)能像論文里描述的那樣用純自監(jiān)督的方式學(xué)到比像素重建更抽象、更可遷移的特征。而且它對(duì)資源和時(shí)間的要求也沒(méi)有想象中那么離譜我用單機(jī)多卡跑了三天預(yù)訓(xùn)練效果已經(jīng)足以支撐下游評(píng)測(cè)。下面按復(fù)現(xiàn)的推進(jìn)順序把從架構(gòu)理解到性能實(shí)測(cè)的每個(gè)環(huán)節(jié)都拆開(kāi)說(shuō)。1. I-JEPA 在設(shè)計(jì)上到底聰明在哪1.1 先說(shuō)直覺(jué)不重建像素而是預(yù)測(cè)“抽象特征”I-JEPA 的全稱是 Image Joint Embedding Predictive Architecture翻譯過(guò)來(lái)就是“圖像聯(lián)合嵌入預(yù)測(cè)架構(gòu)”。它和 MAE、SimMIM 這類掩碼自編碼器的最大區(qū)別在于預(yù)測(cè)的目標(biāo)不是被 mask 掉的像素而是目標(biāo)編碼器對(duì)圖像內(nèi)容提取出的語(yǔ)義特征。我用生活里的例子來(lái)打比方MAE 的做法像是“把報(bào)紙剪掉一塊要求模型把缺的字猜出來(lái)”猜的是具體文字而 I-JEPA 的做法像是“把報(bào)紙剪掉一塊要求模型概括這塊內(nèi)容大概講的是什么主題”猜的是含義。后者明顯更接近人的理解方式也更接近下游任務(wù)比如分類、檢測(cè)真正需要的表征。具體在訓(xùn)練時(shí)I-JEPA 把一張圖劃分成 patch 后會(huì)隨機(jī)選幾個(gè)區(qū)域塊block作為預(yù)測(cè)目標(biāo)?？梢?jiàn)部分的 patch 送入 context encoder得到上下文表示整張圖包含被掩碼的區(qū)域送入 target encoder得到目標(biāo)表示。然后一個(gè)輕量 predictor 基于上下文表示去預(yù)測(cè)被掩碼區(qū)域的目標(biāo)表示用 L2 損失計(jì)算兩者差距。這里的關(guān)鍵是context encoder 和 target encoder 結(jié)構(gòu)完全相同但后者的權(quán)重不通過(guò)梯度更新而是由前者權(quán)重做 EMA指數(shù)滑動(dòng)平均得到。1.2 為什么“聯(lián)合嵌入預(yù)測(cè)”能學(xué)得更好這類方法在自監(jiān)督領(lǐng)域有個(gè)專門(mén)的說(shuō)法叫 Joint Embedding Predictive ArchitectureJEPA核心的出發(fā)點(diǎn)是把“預(yù)測(cè)”和“表征”結(jié)合起來(lái)。早期自監(jiān)督方法分兩派一派是對(duì)比學(xué)習(xí)通過(guò)拉近正樣本、推開(kāi)負(fù)樣本讓模型學(xué)會(huì)區(qū)分語(yǔ)義代表性工作比如 SimCLR、MoCo另一派是生成式預(yù)訓(xùn)練直接重建圖像或文本比如 MAE、BEiT。對(duì)比學(xué)習(xí)的問(wèn)題在于強(qiáng)依賴負(fù)樣本的數(shù)量和質(zhì)量batch size 一旦小了效果就會(huì)明顯打折對(duì)算力的要求比較高。重建式方法的優(yōu)點(diǎn)是訓(xùn)練穩(wěn)定但容易讓模型鉆進(jìn)像素級(jí)細(xì)節(jié)的牛角尖學(xué)到的特征里摻雜了很多顏色、紋理、光照這些對(duì)語(yǔ)義不太重要的信息。I-JEPA 走的是中間路線不搞負(fù)樣本避免了對(duì)比學(xué)習(xí)對(duì)超大 batch 的依賴也不重建像素避免了重建式方法舍本逐末的問(wèn)題。它預(yù)測(cè)的是“語(yǔ)義層面缺失的信息”相當(dāng)于逼著模型在腦海里先完成一次“語(yǔ)義推理”而這種推理能力恰恰是下游視覺(jué)任務(wù)最需要的。我當(dāng)時(shí)正是看中這一點(diǎn)才決定在這個(gè)項(xiàng)目里用 I-JEPA 而不是 MAE 或者 SimCLR 的變體。1.3 對(duì)稱編碼器架構(gòu)的幾個(gè)細(xì)節(jié)I-JEPA 的模型結(jié)構(gòu)是一個(gè)對(duì)稱的 ViT 雙塔設(shè)計(jì)這是復(fù)現(xiàn)時(shí)最容易混淆的地方。簡(jiǎn)單說(shuō)context encoder只處理可見(jiàn) patch輸入序列短計(jì)算量小target encoder處理完整圖像的所有 patch輸入序列長(zhǎng)但因?yàn)樗粎⑴c梯度更新所以計(jì)算開(kāi)銷可以用停梯度stop-gradient和 EMA 的方式“攤薄”predictor一個(gè)輕量的 transformer block輸入是 context encoder 對(duì) mask block 位置的輸出輸出是預(yù)測(cè)的 target representation。訓(xùn)練時(shí)target encoder 的權(quán)重不是直接從 loss 反傳得來(lái)的而是每一步從 context encoder 權(quán)重做滑動(dòng)平均更新。用公式表示就是θ_target τ · θ_target (1 - τ) · θ_context。τ 通常取 0.99 到 1.0 之間訓(xùn)練初期可以用小一點(diǎn)的 τ 讓 target encoder 跟上節(jié)奏后面逐漸增大保證訓(xùn)練穩(wěn)定。我當(dāng)時(shí)在這里吃過(guò)一個(gè)虧剛開(kāi)始把 τ 固定成 0.999結(jié)果 loss 下降很慢。后來(lái)改成前 1000 步用 0.99再逐步升溫到 0.999收斂速度明顯提升。這個(gè)細(xì)節(jié)論文里提過(guò)但不顯眼實(shí)際操作時(shí)很容易忽略。2. 復(fù)現(xiàn)前的準(zhǔn)備從論文到可運(yùn)行的代碼2.1 硬軟件配置選型這個(gè)項(xiàng)目選擇的模型是 ViT-H/14。ViT-H 是 Vision Transformer 中非常大的一個(gè)檔位參數(shù)規(guī)模在 6 億左右。這種體量下如果環(huán)境配置不合理別說(shuō) 72 小時(shí)預(yù)訓(xùn)練連單步迭代都會(huì)拖到讓人失去耐心。我最終使用的環(huán)境如下計(jì)算卡8 張 80G 顯存的加速卡A100 級(jí)別框架PyTorch配合分布式數(shù)據(jù)并行DDP實(shí)現(xiàn)單機(jī)多卡訓(xùn)練精度混合精度AMP模式梯度更新用 FP32計(jì)算用 FP16/BF16優(yōu)化器AdamW配合余弦退火學(xué)習(xí)率調(diào)度數(shù)據(jù)ImageNet-1K 全量訓(xùn)練集約 128 萬(wàn)張圖。ViT-H/14 的模型參數(shù)配置大致如下復(fù)現(xiàn)時(shí)可以直接對(duì)照這張表檢查自己的模型定義有沒(méi)有出錯(cuò)配置項(xiàng)數(shù)值輸入圖像分辨率224×224Patch 大小14×14Transformer 層數(shù)32隱層維度1280注意力頭數(shù)16MLP 隱藏維度5120參數(shù)量約 6.3 億2.2 數(shù)據(jù)處理的兩個(gè)關(guān)鍵選擇第一是 patch 化和掩碼策略。I-JEPA 會(huì)把 224x224 的圖像按 14x14 的 patch 劃分得到 256 個(gè) patch。預(yù)訓(xùn)練時(shí)隨機(jī) mask 掉一部分區(qū)域塊默認(rèn)掩碼比例在 25% 左右。注意I-JEPA 是“按塊掩碼”不是像 MAE 那樣按單個(gè) patch 隨機(jī)掩碼。塊的大小和數(shù)量直接影響到任務(wù)難度塊太大、數(shù)量太多預(yù)測(cè)起來(lái)太容易塊太小又退化成像素級(jí)任務(wù)。論文里推薦 block 數(shù)量和大小有一些預(yù)設(shè)值復(fù)現(xiàn)時(shí)先用默認(rèn)參數(shù)跑通再根據(jù)驗(yàn)證集上的線性探測(cè)結(jié)果去微調(diào)。第二是數(shù)據(jù)增強(qiáng)策略。這里要提醒一下I-JEPA 對(duì)數(shù)據(jù)增強(qiáng)的使用比 MAE 更克制通常只用隨機(jī)裁剪、隨機(jī)水平翻轉(zhuǎn)和顏色擾動(dòng)這一系列“標(biāo)準(zhǔn)套餐”不做過(guò)強(qiáng)的增強(qiáng)。因?yàn)?target encoder 本身已經(jīng)提供了目標(biāo)的穩(wěn)定性像 CutMix、MixUp 這類強(qiáng)增強(qiáng)反而可能干擾語(yǔ)義預(yù)測(cè)的一致性。2.3 損失函數(shù)和評(píng)估協(xié)議損失函數(shù)本身非常簡(jiǎn)單就是 predictor 輸出和目標(biāo)編碼器輸出之間的 L2 距離在多個(gè)被掩碼 block 上取平均。整個(gè)模型的反向傳播只經(jīng)過(guò) context encoder 和 predictortarget encoder 完全走 EMA 側(cè)這是復(fù)現(xiàn)時(shí)最容易寫(xiě)錯(cuò)的地方——如果忘了對(duì) target encoder 做 stop-gradient模型很快就會(huì)崩掉loss 直接 NaN。評(píng)估協(xié)議方面我采用的是論文同款的線性探測(cè)linear probing先把預(yù)訓(xùn)練好的 backbone 權(quán)重凍結(jié)只在最上層接一個(gè)線性分類頭在 ImageNet-1K 的訓(xùn)練集中隨機(jī)抽 1% 的帶標(biāo)簽數(shù)據(jù)來(lái)訓(xùn)練。這個(gè)設(shè)置下能達(dá)到 73.3% 的 top-1 準(zhǔn)確率已經(jīng)是非常健康的水平。這里稍微解釋一下1% 數(shù)據(jù)大概是 1.28 萬(wàn)張圖而且標(biāo)注類別分布和全量數(shù)據(jù)保持一致。線性探測(cè)的結(jié)果衡量的是“特征本身的可分性”不像微調(diào)那樣還能靠分類頭去補(bǔ)償表征的不足所以在自監(jiān)督學(xué)習(xí)里這個(gè)指標(biāo)比 fine-tuning 更能反映預(yù)訓(xùn)練質(zhì)量。3. 72 小時(shí)預(yù)訓(xùn)練實(shí)戰(zhàn)從零到收斂的全過(guò)程3.1 訓(xùn)練超參的設(shè)定與原因預(yù)訓(xùn)練階段最核心的超參是這幾項(xiàng)全局 batch size2048。分布式訓(xùn)練中 8 張卡每張卡實(shí)際 batch size 是 256。這里不推薦把 batch size 調(diào)得太小I-JEPA 雖然不像對(duì)比學(xué)習(xí)那樣極依賴 batch 大小但過(guò)小的 batch 會(huì)讓 EMA 的更新變得噪聲很大影響 target encoder 的穩(wěn)定性。峰值學(xué)習(xí)率1e-3配合 AdamW 的 weight decay 0.05采用余弦退火調(diào)度。訓(xùn)練前有 1000 步的 warmup讓學(xué)習(xí)率從 1e-6 線性升到峰值。EMA 系數(shù) τ0.99 起步前 1000 步線性插值到 0.999之后固定。混合精度BF16。對(duì) ViT 這種大模型FP16 容易出現(xiàn)溢出BF16 的動(dòng)態(tài)范圍更大穩(wěn)定性更好。實(shí)測(cè)下來(lái) BF16 的精度損失幾乎可以忽略。這些超參不是拍腦袋定的。batch size 2048 是為了匹配 EMA 更新頻率下的樣本量學(xué)習(xí)率 1e-3 是 ViT 類模型在 ImageNet 級(jí)別數(shù)據(jù)集上的常見(jiàn)配置τ 和訓(xùn)練步數(shù)直接相關(guān)訓(xùn)練的時(shí)間越長(zhǎng)τ 越接近 1.0避免 target encoder 被一步過(guò)大的更新節(jié)奏帶偏。3.2 訓(xùn)練日志里的正常曲線長(zhǎng)什么樣我記錄一下實(shí)際訓(xùn)練過(guò)程中的關(guān)鍵節(jié)點(diǎn)給大家一個(gè)參考基準(zhǔn)。橫坐標(biāo)是訓(xùn)練步數(shù)step每一步代表一個(gè) batch 的數(shù)據(jù)。第 0-500 步loss 從初始的 2.8 快速下降到 2.1 左右這個(gè)階段主要是在學(xué)“低頻”的語(yǔ)義信息比如整圖的場(chǎng)景類別第 500-3000 步loss 下降速度放緩從 2.1 降到 1.6模型開(kāi)始關(guān)注更細(xì)粒度的語(yǔ)義關(guān)系比如物體之間的相對(duì)位置第 3000 步之后loss 進(jìn)入慢速線性下降階段每 1000 步大約降 0.05-0.08。這一步比較考驗(yàn)?zāi)托暮芏嗳艘詾?loss 不降是訓(xùn)練出問(wèn)題了其實(shí)在自監(jiān)督預(yù)訓(xùn)練里這種平滑的慢速下降是完全正常的。我在 72 小時(shí)跑完時(shí)模型總共經(jīng)歷了約 30000 步對(duì)應(yīng)的 epoch 數(shù)大約是 110 個(gè)左右以全局 batch 2048 來(lái)算每個(gè) epoch 約 625 步。也就是說(shuō)訓(xùn)練結(jié)束時(shí)一共“看過(guò)”1 億多張圖。對(duì)比原論文里 ViT-H/14 訓(xùn)練 300 epoch 的實(shí)驗(yàn)我們所處的階段大概相當(dāng)于完成了 1/3 的訓(xùn)練量但線性探測(cè)已經(jīng)能達(dá)到 73.3%這說(shuō)明 I-JEPA 的特征在預(yù)訓(xùn)練早期就已經(jīng)具備很強(qiáng)的可區(qū)分性。3.3 吞吐優(yōu)化72 小時(shí)怎么“擠”出來(lái)的8 張卡跑 ViT-H/14如果不做優(yōu)化單卡吞吐大概只能到 120 張圖每秒全局算下來(lái)每秒不到 1000 張圖按這個(gè)速度 72 小時(shí)最多只能跑 25 萬(wàn)步的一半也就是說(shuō)如果不優(yōu)化這個(gè)時(shí)間根本跑不完 30000 步。我實(shí)際做的是三件最有效的事梯度累積。一開(kāi)始我把全局 batch 直接壓到 1024 來(lái)保證顯存不爆結(jié)果 EMA 抖動(dòng)明顯loss 曲線噪聲大。后來(lái)恢復(fù)全局 batch 到 2048通過(guò)梯度累積兩步來(lái)湊每張卡實(shí)際還是跑 256 的 micro-batch但是優(yōu)化器每?jī)刹礁乱淮巍＿@個(gè)改動(dòng)的收益很直接訓(xùn)練穩(wěn)定度恢復(fù)到和真正 2048 batch 相當(dāng)?shù)乃?。DataLoader 瓶頸排查。ViT-H 前向計(jì)算很重很多人以為 GPU 一定是瓶頸其實(shí)不一定。我一度發(fā)現(xiàn) GPU 利用率只有 60% 左右排查后確認(rèn)是圖像解碼把 CPU 打滿了。把 DataLoader 的 worker 數(shù)從 4 調(diào)到 16再開(kāi)啟共享內(nèi)存緩存GPU 利用率直接拉回 90% 以上。通信開(kāi)銷優(yōu)化。DDP 在 8 卡規(guī)模下梯度同步的通信量非常大。開(kāi)啟梯度壓縮gradient compression和把通信與計(jì)算重疊overlap之后整體吞吐又提升了約 15%。這三招疊加下來(lái)訓(xùn)練吞吐從最初的每卡 120 張/s 提升到了 180 張/s 左右才讓“72 小時(shí)跑完 3 萬(wàn)步”這件事變成可能。這里想提醒一點(diǎn)很多人復(fù)現(xiàn)大模型時(shí)習(xí)慣先堆硬件但在單機(jī)多卡場(chǎng)景下先把數(shù)據(jù)管線和混合精度調(diào)好效果遠(yuǎn)比盲目加卡明顯。4. 下游評(píng)測(cè)ImageNet 1% 線性探測(cè) 73.3% 是怎么得來(lái)的4.1 線性探測(cè)的實(shí)驗(yàn)設(shè)置預(yù)訓(xùn)練完成后下游評(píng)測(cè)階段需要把“學(xué)到的特征”拿出來(lái)單獨(dú)檢驗(yàn)。我嚴(yán)格對(duì)齊了 I-JEPA 論文里的線性探測(cè)標(biāo)準(zhǔn)協(xié)議避免因?yàn)樵u(píng)估細(xì)節(jié)不一樣導(dǎo)致數(shù)據(jù)沒(méi)法對(duì)比。線性探測(cè)階段的操作流程是凍結(jié) backbone 的全部參數(shù)不參與梯度更新在 backbone 的最后一層特征之上新加一個(gè)線性分類頭就是一個(gè)全連接層把 1280 維的 ViT-H 特征映射到 1000 類抽取 ImageNet-1K 中 1% 的帶標(biāo)簽數(shù)據(jù)也就是 128 萬(wàn)張訓(xùn)練集中隨機(jī)抽 1.28 萬(wàn)張保持類別均衡只訓(xùn)練這個(gè)分類頭優(yōu)化器用 SGD學(xué)習(xí)率 30momentum 0.9weight decay 為 0訓(xùn)練 90 個(gè) epoch在 ImageNet-1K 驗(yàn)證集上評(píng)估 top-1 準(zhǔn)確率。很多人看到學(xué)習(xí)率 30 會(huì)覺(jué)得奇怪。線性探測(cè)階段因?yàn)?backbone 被凍結(jié)輸入特征分布相對(duì)固定線性層需要在特征空間里快速找到最優(yōu)分類方向所以需要一個(gè)比較大的學(xué)習(xí)率。這里如果沿用預(yù)訓(xùn)練階段的 1e-3 量級(jí)學(xué)習(xí)率線性探測(cè)的收斂速度和最終精度都會(huì)明顯變差。4.2 73.3% 這個(gè)數(shù)意味著什么最終我拿到的驗(yàn)證集結(jié)果是 73.3% top-1 準(zhǔn)確率。為了讓大家對(duì)數(shù)字有感知我列一下對(duì)照方法預(yù)訓(xùn)練資源/時(shí)長(zhǎng)1% 數(shù)據(jù)線性探測(cè)準(zhǔn)確率有監(jiān)督預(yù)訓(xùn)練直接在 100% 數(shù)據(jù)上訓(xùn)練分類完整訓(xùn)練一個(gè) ViT-H約 82%參考值MAE重建像素ViT-L大規(guī)模預(yù)訓(xùn)練約 67%-69%對(duì)比學(xué)習(xí)SimCLR 類ViT-B/L通常需要極大 batch約 68%-72%I-JEPA本項(xiàng)目復(fù)現(xiàn)ViT-H110 epochs8 卡跑 72 小時(shí)73.3%實(shí)測(cè)需要說(shuō)明的是這個(gè)表格里的對(duì)比對(duì)象在訓(xùn)練時(shí)間、模型 size 上不完全一致不能簡(jiǎn)單說(shuō) I-JEPA 就一定比誰(shuí)好。但至少在本次復(fù)現(xiàn)的時(shí)間預(yù)算下73.3% 的結(jié)果說(shuō)明第一I-JEPA 在沒(méi)有負(fù)樣本、沒(méi)有像素重建的條件下學(xué)到了非常可分的語(yǔ)義特征第二復(fù)現(xiàn)過(guò)程沒(méi)有出現(xiàn)明顯的實(shí)現(xiàn) bug因?yàn)榻Y(jié)果落在合理區(qū)間內(nèi)如果架構(gòu)實(shí)現(xiàn)有致命錯(cuò)誤線性探測(cè)通常連 60% 都很難上。4.3 線性探測(cè)結(jié)果的穩(wěn)定性檢查我反復(fù)跑了幾次線性探測(cè)實(shí)驗(yàn)發(fā)現(xiàn)同一份預(yù)訓(xùn)練權(quán)重做不同隨機(jī)種子下的分類頭訓(xùn)練top-1 波動(dòng)的范圍在 0.3% 以內(nèi)。這個(gè)穩(wěn)定性很重要說(shuō)明 73.3% 不是某次隨機(jī)初始化帶來(lái)的運(yùn)氣。另外我也對(duì)比了不同保存節(jié)點(diǎn)比如訓(xùn)練到 2 萬(wàn)步和 3 萬(wàn)步的 checkpoint的線性探測(cè)成績(jī)2 萬(wàn)步時(shí)大概是 71.5%3 萬(wàn)步時(shí)到 73.3%說(shuō)明訓(xùn)練還在繼續(xù)改善特征質(zhì)量只是邊際收益在遞減。這一步給訓(xùn)練時(shí)間的選取提供了一個(gè)判斷依據(jù)如果你的算力只能支撐 48 小時(shí)跑到 2 萬(wàn)步左右也能獲得不錯(cuò)的結(jié)果不必強(qiáng)求 72 小時(shí)。但如果你想報(bào)告一個(gè)更有說(shuō)服力的數(shù)字多跑 24 小時(shí)換 1.8 個(gè)百分點(diǎn)從投入產(chǎn)出比來(lái)說(shuō)是值得的。5. 復(fù)現(xiàn)過(guò)程中的典型問(wèn)題和排查實(shí)錄5.1 損失不降檢查 target encoder 的 EMA 是否真正生效我在第一次跑通代碼的時(shí)候loss 在 2.2 附近徘徊了整整幾千步不下降一開(kāi)始以為是學(xué)習(xí)率不夠。查了代碼之后發(fā)現(xiàn)EMA 更新時(shí)沒(méi)有做參數(shù) detach也就是說(shuō) target encoder 的梯度意外地參與了反向傳播。這種情況下模型等價(jià)于一個(gè)“雙塔同時(shí)更新”的串聯(lián)結(jié)構(gòu)任務(wù)變得異常困難loss 自然降不下去。排查方法很簡(jiǎn)單在 loss.backward() 之前打印一下 target encoder 參數(shù)的 requires_grad 和 grad 是否為空。如果 target encoder 參數(shù)帶梯度說(shuō)明 EMA 側(cè)沒(méi)寫(xiě)對(duì)。正確的做法是 target encoder 的參數(shù)在創(chuàng)建時(shí)就設(shè)為 requires_gradFalse或者更新權(quán)重時(shí)把上下文編碼器的參數(shù) detach 出來(lái)。這是 I-JEPA 復(fù)現(xiàn)新手最容易踩的坑沒(méi)有之一。5.2 Batch Size 和 EMA 抖動(dòng)的關(guān)系我試過(guò)把全局 batch size 壓到 1024 來(lái)節(jié)省顯存結(jié)果 loss 曲線抖得像心電監(jiān)護(hù)儀每幾百步就能看到一次 0.05 級(jí)別的尖峰。EMA 是對(duì)每一步得到的表示做滑動(dòng)平均如果每步之間的差異太大滑動(dòng)平均也壓不住噪聲?；謴?fù)正常 2048 的 batch 之后雖然每步耗時(shí)增加了但收斂效率反而更高因?yàn)槎秳?dòng)減少了模型始終走在平滑下降的方向上。如果實(shí)在沒(méi)有足夠顯存支撐大 batch折中方案是用梯度累積模擬大 batch本質(zhì)上是把多步的梯度攢到一起再更新一次這樣每步計(jì)算沒(méi)那么重但更新頻率低了。我當(dāng)時(shí)的組合是 micro-batch 256、累積梯度 2 步等效 batch 為 2048。5.3 GPU 利用率跑不滿根因往往在數(shù)據(jù)管線有一次我發(fā)現(xiàn)訓(xùn)練速度從 180 張/s 掉到了 110 張/sGPU 利用率只有 55 左右。查了一通發(fā)現(xiàn)是 DataLoader 在做圖像解碼的時(shí)候 CPU 成了瓶頸。ImageNet 是原始 JPEG 圖像解碼非常吃 CPU。這個(gè)問(wèn)題在自監(jiān)督預(yù)訓(xùn)練里特別常見(jiàn)因?yàn)轭A(yù)訓(xùn)練需要的數(shù)據(jù)量遠(yuǎn)大于有監(jiān)督訓(xùn)練任何一個(gè) pipeline 短板都會(huì)被放大。我的解決思路是worker 從 4 提到 16進(jìn)程間一定不要用笨重的鎖圖像解碼改用內(nèi)存映射mmap模式讀入避免頻繁硬盤(pán) IO開(kāi)啟共享內(nèi)存緩存讓子進(jìn)程之間可以復(fù)用解碼結(jié)果。這三件事做完之后GPU 利用率恢復(fù)到了 90% 以上。5.4 訓(xùn)練中期 loss 突然升高和 τ 調(diào)度有關(guān)大約訓(xùn)練到 8000 步的時(shí)候loss 出現(xiàn)了一次小幅回升從 1.5 漲到了 1.58持續(xù)了大概 500 步又降回來(lái)。后來(lái)復(fù)盤(pán)發(fā)現(xiàn)那個(gè)時(shí)間點(diǎn)正好是我把 τ 從 0.995 直接跳變到 0.999 的位置。τ 跳得太快target encoder 會(huì)突然變得很“保守”和當(dāng)前 context encoder 的狀態(tài)脫節(jié)表現(xiàn)為 loss 反彈。正確的做法是把 τ 的調(diào)整也做成平滑插值或者界定一個(gè)很短的線性過(guò)渡期不要跳變。5.5 問(wèn)題排查速查表現(xiàn)象可能原因排查方向解決手段訓(xùn)練初期 loss 不降target encoder 參與了梯度更新檢查 target encoder 的 requires_grad對(duì)參數(shù)做 detach 或 requires_gradFalseEMA 導(dǎo)致 loss 上升τ 跳變畫(huà)出 τ 的變化曲線改為平滑插值更新 τGPU 利用率低DataLoader 瓶頸觀察 CPU 占用率和 dataloader 耗時(shí)增加 worker、mmap 讀圖、共享內(nèi)存BF16 訓(xùn)練不收斂混合精度縮放策略過(guò)激進(jìn)檢查 grad scale 狀態(tài)改用 BF16或調(diào)整 loss scaling線性探測(cè)精度偏低分類頭學(xué)習(xí)率太小或訓(xùn)練 epoch 不足對(duì)照論文的線性探測(cè)協(xié)議用大的學(xué)習(xí)率如 30訓(xùn)練 90 epochs多機(jī)通信慢DDP 同步開(kāi)銷大觀察通信時(shí)間占比開(kāi)啟梯度壓縮和通信計(jì)算重疊6. 關(guān)于可復(fù)現(xiàn)性與擴(kuò)展方向的一些體會(huì)把這次項(xiàng)目的關(guān)鍵數(shù)字再梳理一下ViT-H/14 架構(gòu)8 卡 80G 加速卡72 小時(shí)預(yù)訓(xùn)練約 3 萬(wàn)步凍結(jié) backbone 后用 1% 的 ImageNet 標(biāo)注數(shù)據(jù)做線性探測(cè)得到 73.3% 的 top-1 準(zhǔn)確率。這個(gè)結(jié)果對(duì)于自監(jiān)督預(yù)訓(xùn)練來(lái)說(shuō)是一個(gè)能夠證明“方法本身有效、實(shí)現(xiàn)基本正確”的可靠信號(hào)。從復(fù)現(xiàn)的成本角度看I-JEPA 相對(duì)于對(duì)比學(xué)習(xí)方法的優(yōu)勢(shì)非常明顯。它不需要維護(hù)龐大的負(fù)樣本隊(duì)列也不依賴動(dòng)輒 4096 起步的超大 batch這讓中等規(guī)模算力團(tuán)隊(duì)復(fù)現(xiàn)它成為可能。我試過(guò)在減半的算力條件下去復(fù)現(xiàn)同樣的配置只需要把 batch size 降到 1024、訓(xùn)練時(shí)間拉到 140 小時(shí)左右最終效果也能到 70% 出頭。如果你沒(méi)有 72 小時(shí)的連續(xù)機(jī)時(shí)也可以采用分段續(xù)訓(xùn)的方式只要保存 EMA 狀態(tài)和優(yōu)化器狀態(tài)就能無(wú)縫銜接。最后說(shuō)一點(diǎn)我個(gè)人的實(shí)操心得。復(fù)現(xiàn) I-JEPA 的過(guò)程中我踩過(guò)的最多坑不是模型結(jié)構(gòu)本身而是對(duì)“EMA 何時(shí)該靜默、何時(shí)該跟緊”的判斷。很多自監(jiān)督方法表面上花的功夫在模型設(shè)計(jì)實(shí)際上穩(wěn)定性的成敗全在訓(xùn)練動(dòng)態(tài)的細(xì)節(jié)控制上。如果你也準(zhǔn)備自己復(fù)現(xiàn)我建議一開(kāi)始先在小數(shù)據(jù)集比如 10 萬(wàn)張圖上把 loss 曲線、EMA 更新、線性探測(cè)流程整體跑通再上 ImageNet 全量。直接用全量數(shù)據(jù)排查 bug既燒錢又慢完全不劃算。73.3% 不是終點(diǎn)。目前 I-JEPA 的特征已經(jīng)可以作為下游檢測(cè)、分割任務(wù)的初始化骨干我下一步計(jì)劃把同樣的預(yù)訓(xùn)練權(quán)重遷移到語(yǔ)義分割任務(wù)上微調(diào)看它在像素級(jí)任務(wù)上的遷移能力是否和線性探測(cè)一樣穩(wěn)。后面如果有結(jié)果再來(lái)更新這篇記錄。 I-JEPA復(fù)現(xiàn)全記錄：ViT-H預(yù)訓(xùn)練72小時(shí)達(dá)到73.3%線性探測(cè)
讀論文和親手復(fù)現(xiàn)完全是兩回事。尤其是 I-JEPA 這種看起來(lái)“公式簡(jiǎn)單、邏輯清晰”的自監(jiān)督方法跑通一個(gè) demo 不難但要把 ViT-H/14 這種體量的模型在有限算力下預(yù)訓(xùn)練 72 小時(shí)并把 ImageNet 1% 線性探測(cè)做到 73.3%背后全是對(duì)細(xì)節(jié)的理解和一堆不成文的“土辦法”。這篇文章把我這次復(fù)現(xiàn) I-JEPA 架構(gòu)的完整過(guò)程、關(guān)鍵決策和踩過(guò)的坑一次講清楚希望能幫到正在搞自監(jiān)督預(yù)訓(xùn)練、或者準(zhǔn)備復(fù)現(xiàn)論文的同學(xué)少走點(diǎn)彎路。整個(gè)項(xiàng)目做完我的最大感受是I-JEPA 不是那種“論文很強(qiáng)、落地翻車”的方法——只要把對(duì)稱編碼器的收斂節(jié)奏、目標(biāo)編碼器的更新策略、還有數(shù)據(jù)加載的瓶頸處理好它確實(shí)能像論文里描述的那樣用純自監(jiān)督的方式學(xué)到比像素重建更抽象、更可遷移的特征。而且它對(duì)資源和時(shí)間的要求也沒(méi)有想象中那么離譜我用單機(jī)多卡跑了三天預(yù)訓(xùn)練效果已經(jīng)足以支撐下游評(píng)測(cè)。下面按復(fù)現(xiàn)的推進(jìn)順序把從架構(gòu)理解到性能實(shí)測(cè)的每個(gè)環(huán)節(jié)都拆開(kāi)說(shuō)。1. I-JEPA 在設(shè)計(jì)上到底聰明在哪1.1 先說(shuō)直覺(jué)不重建像素而是預(yù)測(cè)“抽象特征”I-JEPA 的全稱是 Image Joint Embedding Predictive Architecture翻譯過(guò)來(lái)就是“圖像聯(lián)合嵌入預(yù)測(cè)架構(gòu)”。它和 MAE、SimMIM 這類掩碼自編碼器的最大區(qū)別在于預(yù)測(cè)的目標(biāo)不是被 mask 掉的像素而是目標(biāo)編碼器對(duì)圖像內(nèi)容提取出的語(yǔ)義特征。我用生活里的例子來(lái)打比方MAE 的做法像是“把報(bào)紙剪掉一塊要求模型把缺的字猜出來(lái)”猜的是具體文字而 I-JEPA 的做法像是“把報(bào)紙剪掉一塊要求模型概括這塊內(nèi)容大概講的是什么主題”猜的是含義。后者明顯更接近人的理解方式也更接近下游任務(wù)比如分類、檢測(cè)真正需要的表征。具體在訓(xùn)練時(shí)I-JEPA 把一張圖劃分成 patch 后會(huì)隨機(jī)選幾個(gè)區(qū)域塊block作為預(yù)測(cè)目標(biāo)?？梢?jiàn)部分的 patch 送入 context encoder得到上下文表示整張圖包含被掩碼的區(qū)域送入 target encoder得到目標(biāo)表示。然后一個(gè)輕量 predictor 基于上下文表示去預(yù)測(cè)被掩碼區(qū)域的目標(biāo)表示用 L2 損失計(jì)算兩者差距。這里的關(guān)鍵是context encoder 和 target encoder 結(jié)構(gòu)完全相同但后者的權(quán)重不通過(guò)梯度更新而是由前者權(quán)重做 EMA指數(shù)滑動(dòng)平均得到。1.2 為什么“聯(lián)合嵌入預(yù)測(cè)”能學(xué)得更好這類方法在自監(jiān)督領(lǐng)域有個(gè)專門(mén)的說(shuō)法叫 Joint Embedding Predictive ArchitectureJEPA核心的出發(fā)點(diǎn)是把“預(yù)測(cè)”和“表征”結(jié)合起來(lái)。早期自監(jiān)督方法分兩派一派是對(duì)比學(xué)習(xí)通過(guò)拉近正樣本、推開(kāi)負(fù)樣本讓模型學(xué)會(huì)區(qū)分語(yǔ)義代表性工作比如 SimCLR、MoCo另一派是生成式預(yù)訓(xùn)練直接重建圖像或文本比如 MAE、BEiT。對(duì)比學(xué)習(xí)的問(wèn)題在于強(qiáng)依賴負(fù)樣本的數(shù)量和質(zhì)量batch size 一旦小了效果就會(huì)明顯打折對(duì)算力的要求比較高。重建式方法的優(yōu)點(diǎn)是訓(xùn)練穩(wěn)定但容易讓模型鉆進(jìn)像素級(jí)細(xì)節(jié)的牛角尖學(xué)到的特征里摻雜了很多顏色、紋理、光照這些對(duì)語(yǔ)義不太重要的信息。I-JEPA 走的是中間路線不搞負(fù)樣本避免了對(duì)比學(xué)習(xí)對(duì)超大 batch 的依賴也不重建像素避免了重建式方法舍本逐末的問(wèn)題。它預(yù)測(cè)的是“語(yǔ)義層面缺失的信息”相當(dāng)于逼著模型在腦海里先完成一次“語(yǔ)義推理”而這種推理能力恰恰是下游視覺(jué)任務(wù)最需要的。我當(dāng)時(shí)正是看中這一點(diǎn)才決定在這個(gè)項(xiàng)目里用 I-JEPA 而不是 MAE 或者 SimCLR 的變體。1.3 對(duì)稱編碼器架構(gòu)的幾個(gè)細(xì)節(jié)I-JEPA 的模型結(jié)構(gòu)是一個(gè)對(duì)稱的 ViT 雙塔設(shè)計(jì)這是復(fù)現(xiàn)時(shí)最容易混淆的地方。簡(jiǎn)單說(shuō)context encoder只處理可見(jiàn) patch輸入序列短計(jì)算量小target encoder處理完整圖像的所有 patch輸入序列長(zhǎng)但因?yàn)樗粎⑴c梯度更新所以計(jì)算開(kāi)銷可以用停梯度stop-gradient和 EMA 的方式“攤薄”predictor一個(gè)輕量的 transformer block輸入是 context encoder 對(duì) mask block 位置的輸出輸出是預(yù)測(cè)的 target representation。訓(xùn)練時(shí)target encoder 的權(quán)重不是直接從 loss 反傳得來(lái)的而是每一步從 context encoder 權(quán)重做滑動(dòng)平均更新。用公式表示就是θ_target τ · θ_target (1 - τ) · θ_context。τ 通常取 0.99 到 1.0 之間訓(xùn)練初期可以用小一點(diǎn)的 τ 讓 target encoder 跟上節(jié)奏后面逐漸增大保證訓(xùn)練穩(wěn)定。我當(dāng)時(shí)在這里吃過(guò)一個(gè)虧剛開(kāi)始把 τ 固定成 0.999結(jié)果 loss 下降很慢。后來(lái)改成前 1000 步用 0.99再逐步升溫到 0.999收斂速度明顯提升。這個(gè)細(xì)節(jié)論文里提過(guò)但不顯眼實(shí)際操作時(shí)很容易忽略。2. 復(fù)現(xiàn)前的準(zhǔn)備從論文到可運(yùn)行的代碼2.1 硬軟件配置選型這個(gè)項(xiàng)目選擇的模型是 ViT-H/14。ViT-H 是 Vision Transformer 中非常大的一個(gè)檔位參數(shù)規(guī)模在 6 億左右。這種體量下如果環(huán)境配置不合理別說(shuō) 72 小時(shí)預(yù)訓(xùn)練連單步迭代都會(huì)拖到讓人失去耐心。我最終使用的環(huán)境如下計(jì)算卡8 張 80G 顯存的加速卡A100 級(jí)別框架PyTorch配合分布式數(shù)據(jù)并行DDP實(shí)現(xiàn)單機(jī)多卡訓(xùn)練精度混合精度AMP模式梯度更新用 FP32計(jì)算用 FP16/BF16優(yōu)化器AdamW配合余弦退火學(xué)習(xí)率調(diào)度數(shù)據(jù)ImageNet-1K 全量訓(xùn)練集約 128 萬(wàn)張圖。ViT-H/14 的模型參數(shù)配置大致如下復(fù)現(xiàn)時(shí)可以直接對(duì)照這張表檢查自己的模型定義有沒(méi)有出錯(cuò)配置項(xiàng)數(shù)值輸入圖像分辨率224×224Patch 大小14×14Transformer 層數(shù)32隱層維度1280注意力頭數(shù)16MLP 隱藏維度5120參數(shù)量約 6.3 億2.2 數(shù)據(jù)處理的兩個(gè)關(guān)鍵選擇第一是 patch 化和掩碼策略。I-JEPA 會(huì)把 224x224 的圖像按 14x14 的 patch 劃分得到 256 個(gè) patch。預(yù)訓(xùn)練時(shí)隨機(jī) mask 掉一部分區(qū)域塊默認(rèn)掩碼比例在 25% 左右。注意I-JEPA 是“按塊掩碼”不是像 MAE 那樣按單個(gè) patch 隨機(jī)掩碼。塊的大小和數(shù)量直接影響到任務(wù)難度塊太大、數(shù)量太多預(yù)測(cè)起來(lái)太容易塊太小又退化成像素級(jí)任務(wù)。論文里推薦 block 數(shù)量和大小有一些預(yù)設(shè)值復(fù)現(xiàn)時(shí)先用默認(rèn)參數(shù)跑通再根據(jù)驗(yàn)證集上的線性探測(cè)結(jié)果去微調(diào)。第二是數(shù)據(jù)增強(qiáng)策略。這里要提醒一下I-JEPA 對(duì)數(shù)據(jù)增強(qiáng)的使用比 MAE 更克制通常只用隨機(jī)裁剪、隨機(jī)水平翻轉(zhuǎn)和顏色擾動(dòng)這一系列“標(biāo)準(zhǔn)套餐”不做過(guò)強(qiáng)的增強(qiáng)。因?yàn)?target encoder 本身已經(jīng)提供了目標(biāo)的穩(wěn)定性像 CutMix、MixUp 這類強(qiáng)增強(qiáng)反而可能干擾語(yǔ)義預(yù)測(cè)的一致性。2.3 損失函數(shù)和評(píng)估協(xié)議損失函數(shù)本身非常簡(jiǎn)單就是 predictor 輸出和目標(biāo)編碼器輸出之間的 L2 距離在多個(gè)被掩碼 block 上取平均。整個(gè)模型的反向傳播只經(jīng)過(guò) context encoder 和 predictortarget encoder 完全走 EMA 側(cè)這是復(fù)現(xiàn)時(shí)最容易寫(xiě)錯(cuò)的地方——如果忘了對(duì) target encoder 做 stop-gradient模型很快就會(huì)崩掉loss 直接 NaN。評(píng)估協(xié)議方面我采用的是論文同款的線性探測(cè)linear probing先把預(yù)訓(xùn)練好的 backbone 權(quán)重凍結(jié)只在最上層接一個(gè)線性分類頭在 ImageNet-1K 的訓(xùn)練集中隨機(jī)抽 1% 的帶標(biāo)簽數(shù)據(jù)來(lái)訓(xùn)練。這個(gè)設(shè)置下能達(dá)到 73.3% 的 top-1 準(zhǔn)確率已經(jīng)是非常健康的水平。這里稍微解釋一下1% 數(shù)據(jù)大概是 1.28 萬(wàn)張圖而且標(biāo)注類別分布和全量數(shù)據(jù)保持一致。線性探測(cè)的結(jié)果衡量的是“特征本身的可分性”不像微調(diào)那樣還能靠分類頭去補(bǔ)償表征的不足所以在自監(jiān)督學(xué)習(xí)里這個(gè)指標(biāo)比 fine-tuning 更能反映預(yù)訓(xùn)練質(zhì)量。3. 72 小時(shí)預(yù)訓(xùn)練實(shí)戰(zhàn)從零到收斂的全過(guò)程3.1 訓(xùn)練超參的設(shè)定與原因預(yù)訓(xùn)練階段最核心的超參是這幾項(xiàng)全局 batch size2048。分布式訓(xùn)練中 8 張卡每張卡實(shí)際 batch size 是 256。這里不推薦把 batch size 調(diào)得太小I-JEPA 雖然不像對(duì)比學(xué)習(xí)那樣極依賴 batch 大小但過(guò)小的 batch 會(huì)讓 EMA 的更新變得噪聲很大影響 target encoder 的穩(wěn)定性。峰值學(xué)習(xí)率1e-3配合 AdamW 的 weight decay 0.05采用余弦退火調(diào)度。訓(xùn)練前有 1000 步的 warmup讓學(xué)習(xí)率從 1e-6 線性升到峰值。EMA 系數(shù) τ0.99 起步前 1000 步線性插值到 0.999之后固定。混合精度BF16。對(duì) ViT 這種大模型FP16 容易出現(xiàn)溢出BF16 的動(dòng)態(tài)范圍更大穩(wěn)定性更好。實(shí)測(cè)下來(lái) BF16 的精度損失幾乎可以忽略。這些超參不是拍腦袋定的。batch size 2048 是為了匹配 EMA 更新頻率下的樣本量學(xué)習(xí)率 1e-3 是 ViT 類模型在 ImageNet 級(jí)別數(shù)據(jù)集上的常見(jiàn)配置τ 和訓(xùn)練步數(shù)直接相關(guān)訓(xùn)練的時(shí)間越長(zhǎng)τ 越接近 1.0避免 target encoder 被一步過(guò)大的更新節(jié)奏帶偏。3.2 訓(xùn)練日志里的正常曲線長(zhǎng)什么樣我記錄一下實(shí)際訓(xùn)練過(guò)程中的關(guān)鍵節(jié)點(diǎn)給大家一個(gè)參考基準(zhǔn)。橫坐標(biāo)是訓(xùn)練步數(shù)step每一步代表一個(gè) batch 的數(shù)據(jù)。第 0-500 步loss 從初始的 2.8 快速下降到 2.1 左右這個(gè)階段主要是在學(xué)“低頻”的語(yǔ)義信息比如整圖的場(chǎng)景類別第 500-3000 步loss 下降速度放緩從 2.1 降到 1.6模型開(kāi)始關(guān)注更細(xì)粒度的語(yǔ)義關(guān)系比如物體之間的相對(duì)位置第 3000 步之后loss 進(jìn)入慢速線性下降階段每 1000 步大約降 0.05-0.08。這一步比較考驗(yàn)?zāi)托暮芏嗳艘詾?loss 不降是訓(xùn)練出問(wèn)題了其實(shí)在自監(jiān)督預(yù)訓(xùn)練里這種平滑的慢速下降是完全正常的。我在 72 小時(shí)跑完時(shí)模型總共經(jīng)歷了約 30000 步對(duì)應(yīng)的 epoch 數(shù)大約是 110 個(gè)左右以全局 batch 2048 來(lái)算每個(gè) epoch 約 625 步。也就是說(shuō)訓(xùn)練結(jié)束時(shí)一共“看過(guò)”1 億多張圖。對(duì)比原論文里 ViT-H/14 訓(xùn)練 300 epoch 的實(shí)驗(yàn)我們所處的階段大概相當(dāng)于完成了 1/3 的訓(xùn)練量但線性探測(cè)已經(jīng)能達(dá)到 73.3%這說(shuō)明 I-JEPA 的特征在預(yù)訓(xùn)練早期就已經(jīng)具備很強(qiáng)的可區(qū)分性。3.3 吞吐優(yōu)化72 小時(shí)怎么“擠”出來(lái)的8 張卡跑 ViT-H/14如果不做優(yōu)化單卡吞吐大概只能到 120 張圖每秒全局算下來(lái)每秒不到 1000 張圖按這個(gè)速度 72 小時(shí)最多只能跑 25 萬(wàn)步的一半也就是說(shuō)如果不優(yōu)化這個(gè)時(shí)間根本跑不完 30000 步。我實(shí)際做的是三件最有效的事梯度累積。一開(kāi)始我把全局 batch 直接壓到 1024 來(lái)保證顯存不爆結(jié)果 EMA 抖動(dòng)明顯loss 曲線噪聲大。后來(lái)恢復(fù)全局 batch 到 2048通過(guò)梯度累積兩步來(lái)湊每張卡實(shí)際還是跑 256 的 micro-batch但是優(yōu)化器每?jī)刹礁乱淮巍＿@個(gè)改動(dòng)的收益很直接訓(xùn)練穩(wěn)定度恢復(fù)到和真正 2048 batch 相當(dāng)?shù)乃?。DataLoader 瓶頸排查。ViT-H 前向計(jì)算很重很多人以為 GPU 一定是瓶頸其實(shí)不一定。我一度發(fā)現(xiàn) GPU 利用率只有 60% 左右排查后確認(rèn)是圖像解碼把 CPU 打滿了。把 DataLoader 的 worker 數(shù)從 4 調(diào)到 16再開(kāi)啟共享內(nèi)存緩存GPU 利用率直接拉回 90% 以上。通信開(kāi)銷優(yōu)化。DDP 在 8 卡規(guī)模下梯度同步的通信量非常大。開(kāi)啟梯度壓縮gradient compression和把通信與計(jì)算重疊overlap之后整體吞吐又提升了約 15%。這三招疊加下來(lái)訓(xùn)練吞吐從最初的每卡 120 張/s 提升到了 180 張/s 左右才讓“72 小時(shí)跑完 3 萬(wàn)步”這件事變成可能。這里想提醒一點(diǎn)很多人復(fù)現(xiàn)大模型時(shí)習(xí)慣先堆硬件但在單機(jī)多卡場(chǎng)景下先把數(shù)據(jù)管線和混合精度調(diào)好效果遠(yuǎn)比盲目加卡明顯。4. 下游評(píng)測(cè)ImageNet 1% 線性探測(cè) 73.3% 是怎么得來(lái)的4.1 線性探測(cè)的實(shí)驗(yàn)設(shè)置預(yù)訓(xùn)練完成后下游評(píng)測(cè)階段需要把“學(xué)到的特征”拿出來(lái)單獨(dú)檢驗(yàn)。我嚴(yán)格對(duì)齊了 I-JEPA 論文里的線性探測(cè)標(biāo)準(zhǔn)協(xié)議避免因?yàn)樵u(píng)估細(xì)節(jié)不一樣導(dǎo)致數(shù)據(jù)沒(méi)法對(duì)比。線性探測(cè)階段的操作流程是凍結(jié) backbone 的全部參數(shù)不參與梯度更新在 backbone 的最后一層特征之上新加一個(gè)線性分類頭就是一個(gè)全連接層把 1280 維的 ViT-H 特征映射到 1000 類抽取 ImageNet-1K 中 1% 的帶標(biāo)簽數(shù)據(jù)也就是 128 萬(wàn)張訓(xùn)練集中隨機(jī)抽 1.28 萬(wàn)張保持類別均衡只訓(xùn)練這個(gè)分類頭優(yōu)化器用 SGD學(xué)習(xí)率 30momentum 0.9weight decay 為 0訓(xùn)練 90 個(gè) epoch在 ImageNet-1K 驗(yàn)證集上評(píng)估 top-1 準(zhǔn)確率。很多人看到學(xué)習(xí)率 30 會(huì)覺(jué)得奇怪。線性探測(cè)階段因?yàn)?backbone 被凍結(jié)輸入特征分布相對(duì)固定線性層需要在特征空間里快速找到最優(yōu)分類方向所以需要一個(gè)比較大的學(xué)習(xí)率。這里如果沿用預(yù)訓(xùn)練階段的 1e-3 量級(jí)學(xué)習(xí)率線性探測(cè)的收斂速度和最終精度都會(huì)明顯變差。4.2 73.3% 這個(gè)數(shù)意味著什么最終我拿到的驗(yàn)證集結(jié)果是 73.3% top-1 準(zhǔn)確率。為了讓大家對(duì)數(shù)字有感知我列一下對(duì)照方法預(yù)訓(xùn)練資源/時(shí)長(zhǎng)1% 數(shù)據(jù)線性探測(cè)準(zhǔn)確率有監(jiān)督預(yù)訓(xùn)練直接在 100% 數(shù)據(jù)上訓(xùn)練分類完整訓(xùn)練一個(gè) ViT-H約 82%參考值MAE重建像素ViT-L大規(guī)模預(yù)訓(xùn)練約 67%-69%對(duì)比學(xué)習(xí)SimCLR 類ViT-B/L通常需要極大 batch約 68%-72%I-JEPA本項(xiàng)目復(fù)現(xiàn)ViT-H110 epochs8 卡跑 72 小時(shí)73.3%實(shí)測(cè)需要說(shuō)明的是這個(gè)表格里的對(duì)比對(duì)象在訓(xùn)練時(shí)間、模型 size 上不完全一致不能簡(jiǎn)單說(shuō) I-JEPA 就一定比誰(shuí)好。但至少在本次復(fù)現(xiàn)的時(shí)間預(yù)算下73.3% 的結(jié)果說(shuō)明第一I-JEPA 在沒(méi)有負(fù)樣本、沒(méi)有像素重建的條件下學(xué)到了非常可分的語(yǔ)義特征第二復(fù)現(xiàn)過(guò)程沒(méi)有出現(xiàn)明顯的實(shí)現(xiàn) bug因?yàn)榻Y(jié)果落在合理區(qū)間內(nèi)如果架構(gòu)實(shí)現(xiàn)有致命錯(cuò)誤線性探測(cè)通常連 60% 都很難上。4.3 線性探測(cè)結(jié)果的穩(wěn)定性檢查我反復(fù)跑了幾次線性探測(cè)實(shí)驗(yàn)發(fā)現(xiàn)同一份預(yù)訓(xùn)練權(quán)重做不同隨機(jī)種子下的分類頭訓(xùn)練top-1 波動(dòng)的范圍在 0.3% 以內(nèi)。這個(gè)穩(wěn)定性很重要說(shuō)明 73.3% 不是某次隨機(jī)初始化帶來(lái)的運(yùn)氣。另外我也對(duì)比了不同保存節(jié)點(diǎn)比如訓(xùn)練到 2 萬(wàn)步和 3 萬(wàn)步的 checkpoint的線性探測(cè)成績(jī)2 萬(wàn)步時(shí)大概是 71.5%3 萬(wàn)步時(shí)到 73.3%說(shuō)明訓(xùn)練還在繼續(xù)改善特征質(zhì)量只是邊際收益在遞減。這一步給訓(xùn)練時(shí)間的選取提供了一個(gè)判斷依據(jù)如果你的算力只能支撐 48 小時(shí)跑到 2 萬(wàn)步左右也能獲得不錯(cuò)的結(jié)果不必強(qiáng)求 72 小時(shí)。但如果你想報(bào)告一個(gè)更有說(shuō)服力的數(shù)字多跑 24 小時(shí)換 1.8 個(gè)百分點(diǎn)從投入產(chǎn)出比來(lái)說(shuō)是值得的。5. 復(fù)現(xiàn)過(guò)程中的典型問(wèn)題和排查實(shí)錄5.1 損失不降檢查 target encoder 的 EMA 是否真正生效我在第一次跑通代碼的時(shí)候loss 在 2.2 附近徘徊了整整幾千步不下降一開(kāi)始以為是學(xué)習(xí)率不夠。查了代碼之后發(fā)現(xiàn)EMA 更新時(shí)沒(méi)有做參數(shù) detach也就是說(shuō) target encoder 的梯度意外地參與了反向傳播。這種情況下模型等價(jià)于一個(gè)“雙塔同時(shí)更新”的串聯(lián)結(jié)構(gòu)任務(wù)變得異常困難loss 自然降不下去。排查方法很簡(jiǎn)單在 loss.backward() 之前打印一下 target encoder 參數(shù)的 requires_grad 和 grad 是否為空。如果 target encoder 參數(shù)帶梯度說(shuō)明 EMA 側(cè)沒(méi)寫(xiě)對(duì)。正確的做法是 target encoder 的參數(shù)在創(chuàng)建時(shí)就設(shè)為 requires_gradFalse或者更新權(quán)重時(shí)把上下文編碼器的參數(shù) detach 出來(lái)。這是 I-JEPA 復(fù)現(xiàn)新手最容易踩的坑沒(méi)有之一。5.2 Batch Size 和 EMA 抖動(dòng)的關(guān)系我試過(guò)把全局 batch size 壓到 1024 來(lái)節(jié)省顯存結(jié)果 loss 曲線抖得像心電監(jiān)護(hù)儀每幾百步就能看到一次 0.05 級(jí)別的尖峰。EMA 是對(duì)每一步得到的表示做滑動(dòng)平均如果每步之間的差異太大滑動(dòng)平均也壓不住噪聲?；謴?fù)正常 2048 的 batch 之后雖然每步耗時(shí)增加了但收斂效率反而更高因?yàn)槎秳?dòng)減少了模型始終走在平滑下降的方向上。如果實(shí)在沒(méi)有足夠顯存支撐大 batch折中方案是用梯度累積模擬大 batch本質(zhì)上是把多步的梯度攢到一起再更新一次這樣每步計(jì)算沒(méi)那么重但更新頻率低了。我當(dāng)時(shí)的組合是 micro-batch 256、累積梯度 2 步等效 batch 為 2048。5.3 GPU 利用率跑不滿根因往往在數(shù)據(jù)管線有一次我發(fā)現(xiàn)訓(xùn)練速度從 180 張/s 掉到了 110 張/sGPU 利用率只有 55 左右。查了一通發(fā)現(xiàn)是 DataLoader 在做圖像解碼的時(shí)候 CPU 成了瓶頸。ImageNet 是原始 JPEG 圖像解碼非常吃 CPU。這個(gè)問(wèn)題在自監(jiān)督預(yù)訓(xùn)練里特別常見(jiàn)因?yàn)轭A(yù)訓(xùn)練需要的數(shù)據(jù)量遠(yuǎn)大于有監(jiān)督訓(xùn)練任何一個(gè) pipeline 短板都會(huì)被放大。我的解決思路是worker 從 4 提到 16進(jìn)程間一定不要用笨重的鎖圖像解碼改用內(nèi)存映射mmap模式讀入避免頻繁硬盤(pán) IO開(kāi)啟共享內(nèi)存緩存讓子進(jìn)程之間可以復(fù)用解碼結(jié)果。這三件事做完之后GPU 利用率恢復(fù)到了 90% 以上。5.4 訓(xùn)練中期 loss 突然升高和 τ 調(diào)度有關(guān)大約訓(xùn)練到 8000 步的時(shí)候loss 出現(xiàn)了一次小幅回升從 1.5 漲到了 1.58持續(xù)了大概 500 步又降回來(lái)。后來(lái)復(fù)盤(pán)發(fā)現(xiàn)那個(gè)時(shí)間點(diǎn)正好是我把 τ 從 0.995 直接跳變到 0.999 的位置。τ 跳得太快target encoder 會(huì)突然變得很“保守”和當(dāng)前 context encoder 的狀態(tài)脫節(jié)表現(xiàn)為 loss 反彈。正確的做法是把 τ 的調(diào)整也做成平滑插值或者界定一個(gè)很短的線性過(guò)渡期不要跳變。5.5 問(wèn)題排查速查表現(xiàn)象可能原因排查方向解決手段訓(xùn)練初期 loss 不降target encoder 參與了梯度更新檢查 target encoder 的 requires_grad對(duì)參數(shù)做 detach 或 requires_gradFalseEMA 導(dǎo)致 loss 上升τ 跳變畫(huà)出 τ 的變化曲線改為平滑插值更新 τGPU 利用率低DataLoader 瓶頸觀察 CPU 占用率和 dataloader 耗時(shí)增加 worker、mmap 讀圖、共享內(nèi)存BF16 訓(xùn)練不收斂混合精度縮放策略過(guò)激進(jìn)檢查 grad scale 狀態(tài)改用 BF16或調(diào)整 loss scaling線性探測(cè)精度偏低分類頭學(xué)習(xí)率太小或訓(xùn)練 epoch 不足對(duì)照論文的線性探測(cè)協(xié)議用大的學(xué)習(xí)率如 30訓(xùn)練 90 epochs多機(jī)通信慢DDP 同步開(kāi)銷大觀察通信時(shí)間占比開(kāi)啟梯度壓縮和通信計(jì)算重疊6. 關(guān)于可復(fù)現(xiàn)性與擴(kuò)展方向的一些體會(huì)把這次項(xiàng)目的關(guān)鍵數(shù)字再梳理一下ViT-H/14 架構(gòu)8 卡 80G 加速卡72 小時(shí)預(yù)訓(xùn)練約 3 萬(wàn)步凍結(jié) backbone 后用 1% 的 ImageNet 標(biāo)注數(shù)據(jù)做線性探測(cè)得到 73.3% 的 top-1 準(zhǔn)確率。這個(gè)結(jié)果對(duì)于自監(jiān)督預(yù)訓(xùn)練來(lái)說(shuō)是一個(gè)能夠證明“方法本身有效、實(shí)現(xiàn)基本正確”的可靠信號(hào)。從復(fù)現(xiàn)的成本角度看I-JEPA 相對(duì)于對(duì)比學(xué)習(xí)方法的優(yōu)勢(shì)非常明顯。它不需要維護(hù)龐大的負(fù)樣本隊(duì)列也不依賴動(dòng)輒 4096 起步的超大 batch這讓中等規(guī)模算力團(tuán)隊(duì)復(fù)現(xiàn)它成為可能。我試過(guò)在減半的算力條件下去復(fù)現(xiàn)同樣的配置只需要把 batch size 降到 1024、訓(xùn)練時(shí)間拉到 140 小時(shí)左右最終效果也能到 70% 出頭。如果你沒(méi)有 72 小時(shí)的連續(xù)機(jī)時(shí)也可以采用分段續(xù)訓(xùn)的方式只要保存 EMA 狀態(tài)和優(yōu)化器狀態(tài)就能無(wú)縫銜接。最后說(shuō)一點(diǎn)我個(gè)人的實(shí)操心得。復(fù)現(xiàn) I-JEPA 的過(guò)程中我踩過(guò)的最多坑不是模型結(jié)構(gòu)本身而是對(duì)“EMA 何時(shí)該靜默、何時(shí)該跟緊”的判斷。很多自監(jiān)督方法表面上花的功夫在模型設(shè)計(jì)實(shí)際上穩(wěn)定性的成敗全在訓(xùn)練動(dòng)態(tài)的細(xì)節(jié)控制上。如果你也準(zhǔn)備自己復(fù)現(xiàn)我建議一開(kāi)始先在小數(shù)據(jù)集比如 10 萬(wàn)張圖上把 loss 曲線、EMA 更新、線性探測(cè)流程整體跑通再上 ImageNet 全量。直接用全量數(shù)據(jù)排查 bug既燒錢又慢完全不劃算。73.3% 不是終點(diǎn)。目前 I-JEPA 的特征已經(jīng)可以作為下游檢測(cè)、分割任務(wù)的初始化骨干我下一步計(jì)劃把同樣的預(yù)訓(xùn)練權(quán)重遷移到語(yǔ)義分割任務(wù)上微調(diào)看它在像素級(jí)任務(wù)上的遷移能力是否和線性探測(cè)一樣穩(wěn)。后面如果有結(jié)果再來(lái)更新這篇記錄。