[RSCH] 7 分鐘閱讀OraCore 編輯部

DiT 的異常 token 怎麼救

這篇論文指出,RAE-DiT 管線裡的異常 token 會干擾注意力並造成影像瑕疵,並提出 Dual-Stage Registers 來同時處理編碼器與去噪器的問題。

分享 LinkedIn
DiT 的異常 token 怎麼救

這篇研究指出,Diffusion Transformer 裡的異常 token 會破壞生成品質,作者提出 Dual-Stage Registers 來同時處理編碼器和去噪器的問題。

Diffusion Transformers,簡稱 DiTs,近年是影像生成的熱門架構之一。但這篇論文提醒一個很容易被忽略的痛點:模型裡少數 norm 特別大的 outlier tokens,可能會把注意力拉歪,進而讓生成結果出現瑕疵。換句話說,做 DiT 不只是在比模型大不大、資料多不多,token 在模型內部怎麼流動,也會直接影響成品質感。

這篇論文是 Taming Outlier Tokens in Diffusion Transformers。它聚焦在基於 Representation Autoencoder,簡稱 RAE,的 DiT 管線,發現 outlier 不只出現在預訓練的 ViT encoder,也會出現在 DiT denoiser 本身。這代表問題不是某一個模組單獨出包,而是整條生成流程都可能被影響。

這篇在解什麼痛點

訂閱 AI 趨勢週報

每週精選模型發布、工具應用與深度分析,直送信箱。不定期,不騷擾。

不會寄垃圾信,隨時可取消。

作者要處理的核心問題,是 outlier tokens。這些 token 的 norm 異常大,會吸走很多 attention,但它們本身攜帶的局部資訊卻不一定多。這種現象在 Vision Transformer 裡早就有人注意到,但過去多半是在辨識或表徵學習的脈絡下看待。這篇論文進一步問:如果同樣的現象發生在生成模型裡,會怎樣?

DiT 的異常 token 怎麼救

答案是,影響不小。生成模型不像分類器只要最後判對就好,它需要在 denoising 過程中維持穩定的 token 互動。若少數 token 一直主導 attention,而且主導的理由不是內容本身,而是數值上特別大,那模型就可能把結構帶歪,最後生成出 artifacts,而不是乾淨的影像。

更重要的是,這篇研究不是把 outlier 當成單點故障來看。它指出,在 RAE-DiT 管線裡,pretrained ViT encoder 可能先吐出帶有 outlier 特性的表示;進到 DiT denoiser 後,模型內部也可能在中間層長出自己的 outlier tokens。也就是說,上游一旦不穩,下游很可能跟著放大問題。

這裡還有一個值得注意的負結果:單純把高 norm token mask 掉,並不會讓表現變好。這表示問題不是「抓幾個極端值刪掉」那麼簡單。作者的解讀比較像是,問題根源和局部 patch semantics 被污染有關,這比單純做 outlier clipping 更結構性。

方法到底怎麼運作

為了處理這個問題,論文提出 Dual-Stage Registers,簡稱 DSR。從名字就看得出來,它不是只修一段,而是同時動到兩個階段:encoder 和 denoiser。概念上,register tokens 被拿來當作干預手段,讓模型有一個額外的空間去吸收、轉移或重整那些容易失控的 token 行為。

抽象內容提到 DSR 有三個部分。第一,當有訓練過的 registers 可用時,就直接使用。第二,如果沒有訓練好的 registers,系統會退回 recursive test-time registers。第三,還會加入 diffusion registers,專門處理 denoiser 端的穩定性。白話一點講,作者不是只想在輸入端做修補,而是想讓整個 token 流程在兩個關鍵階段都更穩。

這個設計的工程意義很直接。如果 encoder 已經產生 outlier 表示,denoiser 接收到的輸入本來就不乾淨;如果 denoiser 自己又在中間層產生 outlier tokens,那它的內部 dynamics 也會把問題繼續放大。DSR 的想法,就是把這兩種來源一起管,不要只修一邊。

從實作角度看,這種做法和「把壞 token 直接砍掉」不同。它比較像是在系統裡加入一個緩衝層或調節層,讓模型有機會把原本會失控的訊號導向比較穩定的路徑。這也是為什麼論文把它叫做 dual-stage,而不是單純的 token filter。

論文實際證明了什麼

就 abstract 能確定的部分來看,作者觀察到 outlier token 現象同時存在於現代 RAE-DiT 管線的 encoder 與 denoiser。這一點很重要,因為它把問題從「前處理」或「中間層」的單一故障,提升成整條生成鏈的系統性問題。

DiT 的異常 token 怎麼救

論文也宣稱,這些介入方法可以持續減少 outlier artifacts,並改善 ImageNet 以及大規模 text-to-image 生成的品質。這代表 DSR 不是只在一個小場景裡有效,而是至少在 abstract 提到的幾個設定中都能看到正向效果。

不過,這份來源沒有公開完整 benchmark 細節。沒有看到明確分數、百分比提升,或表格級別的比較結果,所以我們不能從 abstract 推斷改善幅度有多大。能確定的是方向:方法有效;不能確定的是規模:有多強。

即便如此,這個結果還是有實際意義,因為論文不是只說「outlier 存在」。它還指出這些 outlier 會真的影響生成品質,而且可以透過 register-based 的方式被壓下來。對生成模型來說,能減少 artifact,本身就是很有價值的訊號。

  • Outlier tokens 同時出現在 ViT encoder 與 DiT denoiser。
  • 中間層特別容易長出內部 outlier tokens。
  • 單純 mask 高 norm token 並不能解決問題。
  • DSR 用 register tokens 同時干預兩個階段。
  • 摘要宣稱能減少 artifacts,並改善 ImageNet 與 text-to-image 生成品質。

對開發者有什麼影響

如果你在做影像生成,這篇論文提醒了一件事:模型內部的 token 行為,可能跟最終輸出一樣重要。很多時候我們只看 loss、看參數量、看資料規模,但這篇研究把焦點拉回到 token norm、attention 集中程度,以及 local patch semantics 有沒有在 pipeline 裡被破壞。

對工程師來說,這也讓 registers 的角色變得更實際。過去 registers 常被看成 representation learning 的技巧;這篇論文則把它們拿來當生成管線的控制機制。也就是說,當 token semantics 開始漂移、token 互動變得不穩時,加入結構化 token,可能比粗暴地刪 token 更有效。

但這篇摘要也留下不少實務上會在意的問題。它沒有說 DSR 的訓練成本是多少,也沒有說推論時的額外負擔。它沒有交代這方法對模型大小有多敏感,也沒有說離開 RAE-DiT 管線後還能不能用。換句話說,這是一個有方向感的修法,但還不是可以直接當成萬用補丁的結論。

對團隊而言,這類結果比較像是設計提醒:如果你在做 DiT-based 系統,除了看 output 美不美,也要檢查 token 是否出現異常集中、encoder 是否已經把局部語意弄亂、denoiser 是否又把問題放大。這些內部病灶,可能就是 artifacts 的來源。

這篇論文的價值,不只是提出一個新技巧,而是把一個很具體的 failure mode 說清楚。當 diffusion transformers 繼續往更大規模、更高品質的生成模型前進時,這類內部穩定性問題可能會越來越重要。對開發者來說,這是提醒,也是工具箱的擴充方向。

總結

這篇研究把 outlier tokens 從 ViT 的老問題,推進成 DiT 生成管線的實際風險。作者提出的 Dual-Stage Registers,目標是同時穩住 encoder 與 denoiser,讓生成過程少一點異常 token 造成的干擾。就 abstract 來看,方法確實朝著減少 artifacts、提升生成品質的方向前進。

如果你正在看 Diffusion Transformer、RAE、或 register 這類設計,這篇論文值得記住的點很簡單:在生成模型裡,token 不只是資料的載體,它們本身也可能是失敗的來源。把這件事處理好,才有機會讓影像生成更穩、更乾淨。