這項由清華大學(xué)電子工程系的邱海權(quán)和姚強(qiáng)明教授團(tuán)隊開展的研究發(fā)表于2025年10月，論文編號為arXiv:2510.04212v1。這個研究團(tuán)隊花費了大量時間來解決一個困擾AI訓(xùn)練領(lǐng)域超過兩年的"神秘故障"，有興趣深入了解技術(shù)細(xì)節(jié)的讀者可以通過這個編號查詢完整論文。

在AI大模型訓(xùn)練的世界里，有一個讓工程師們頭疼不已的"幽靈故障"。你可以把訓(xùn)練大型AI模型想象成培養(yǎng)一個超級聰明的學(xué)生。為了讓這個學(xué)生學(xué)得更快、記得更多，工程師們想出了各種"節(jié)約"的辦法，其中最重要的一招就是使用"簡化數(shù)字"來進(jìn)行計算。這就像是用簡化的數(shù)學(xué)符號來做題，既能節(jié)省時間又能節(jié)省紙張。

然而，這種看似聰明的做法卻帶來了一個詭異的問題。有時候，一個原本訓(xùn)練得好好的AI模型會突然"發(fā)瘋"——損失值猛然飆升，就像一個成績優(yōu)秀的學(xué)生突然在考試中胡寫亂畫一樣。更奇怪的是，這種故障完全無法預(yù)測，可能在訓(xùn)練了幾千步之后突然出現(xiàn)，讓所有的努力付諸東流。

這種現(xiàn)象在使用一種叫做"Flash Attention"（閃電注意力）的技術(shù)時特別常見。Flash Attention就像是AI模型的"超級記憶術(shù)"，能夠讓模型在處理長文本時既快又省內(nèi)存。但是當(dāng)它遇到簡化數(shù)字計算時，就會出現(xiàn)這種神秘的崩潰。

清華團(tuán)隊的這項研究就像是一場精彩的"科學(xué)偵探故事"。他們不滿足于簡單的修補(bǔ)，而是要找出問題的根本原因。經(jīng)過深入調(diào)查，他們發(fā)現(xiàn)這個"幽靈故障"其實有著非常具體的成因，并且首次給出了完整的解釋。

一、訓(xùn)練崩潰的真相：數(shù)字舍入中的"偏見"陷阱

要理解這個問題，我們首先需要知道什么是"簡化數(shù)字計算"。在AI訓(xùn)練中，工程師們通常使用一種叫做BF16的數(shù)字格式，這就像是用簡化版的數(shù)學(xué)符號來做計算。正常情況下，我們可能用32位數(shù)字來表示一個數(shù)值（比如3.14159265），但BF16只用16位（比如3.14），這樣能節(jié)省一半的存儲空間和計算時間。

這種簡化本身并不會造成問題，就像用簡化的小數(shù)進(jìn)行日常計算一樣。但是，研究團(tuán)隊發(fā)現(xiàn)，當(dāng)計算機(jī)進(jìn)行大量連續(xù)的加法運算時，特別是在Flash Attention的某個關(guān)鍵步驟中，會出現(xiàn)一種微妙的"偏見"。

這種偏見可以用一個生活中的例子來理解。假設(shè)你在做大量的購物計算，每次都需要把價格四舍五入到最近的角。如果你買的商品價格大多數(shù)都是類似2.67元、3.78元這樣的數(shù)字，那么四舍五入后，你總是會比實際價格少付一點錢。單獨看每一次購買，這個差異微不足道，但如果你買了幾千件商品，這個差異就會積累成一個顯著的偏差。

在AI訓(xùn)練中也是如此。Flash Attention需要進(jìn)行大量類似的數(shù)字計算，每次計算都涉及舍入操作。研究團(tuán)隊發(fā)現(xiàn)，在特定的條件下，這些舍入誤差不是隨機(jī)分布的，而是有"偏見"的——它們傾向于往同一個方向累積。

二、低秩矩陣的"共謀"：相似模式如何放大誤差

更加復(fù)雜的是，這種數(shù)字偏見還與AI模型內(nèi)部的數(shù)學(xué)結(jié)構(gòu)產(chǎn)生了"共謀"。研究團(tuán)隊發(fā)現(xiàn)，在訓(xùn)練過程中，某些數(shù)學(xué)對象（稱為低秩矩陣）會在不同的訓(xùn)練步驟和不同的數(shù)據(jù)處理中表現(xiàn)出驚人的相似性。

這就像是一個班級里的學(xué)生在做不同的數(shù)學(xué)題時，卻總是犯同樣類型的錯誤。單獨看每個學(xué)生的錯誤，可能只是小問題，但如果全班同學(xué)都在同一個地方犯錯，就會形成系統(tǒng)性的偏差。

在AI模型中，這些相似的"錯誤模式"就是低秩矩陣。當(dāng)數(shù)字舍入的偏見作用到這些相似的模式上時，就像是給每個相似的錯誤都乘以了同一個偏向性的系數(shù)。這樣一來，錯誤不僅不會相互抵消，反而會相互強(qiáng)化，越積越大。

研究團(tuán)隊通過大量的實驗和數(shù)學(xué)分析，證明了這種現(xiàn)象確實會發(fā)生。他們發(fā)現(xiàn)，在訓(xùn)練的某些關(guān)鍵時刻，這種累積的偏差會達(dá)到一個臨界點，導(dǎo)致整個模型的參數(shù)（就像學(xué)生的知識結(jié)構(gòu)）發(fā)生劇烈的扭曲，最終導(dǎo)致訓(xùn)練失敗。

三、深入故障現(xiàn)場：Flash Attention的致命弱點

為了找出確切的故障位置，研究團(tuán)隊進(jìn)行了一系列精密的"解剖"實驗。他們逐步關(guān)閉Flash Attention的各個組件，就像維修工逐個檢查汽車的零部件一樣，直到找出真正的罪魁禍?zhǔn)住?/p>

他們首先發(fā)現(xiàn)，問題并不在Flash Attention的分塊處理策略上。分塊處理就像是把一本厚書分成幾個章節(jié)來閱讀，這種方法本身是安全的。真正的問題出在一個看似無害的計算步驟上。

在Flash Attention的反向傳播過程中（這就像是AI模型的"反思學(xué)習(xí)"階段），需要計算一個叫做δ的數(shù)值。這個δ的計算涉及兩個矩陣的逐元素相乘，然后按行求和。聽起來很簡單，但魔鬼就藏在細(xì)節(jié)中。

當(dāng)這個計算使用BF16格式進(jìn)行時，就會觸發(fā)前面提到的舍入偏見問題。更具體地說，研究團(tuán)隊發(fā)現(xiàn)問題出現(xiàn)在計算過程中的P×V乘法操作上。這里的P是注意力權(quán)重（可以理解為AI模型對不同信息的關(guān)注程度），V是值矩陣（可以理解為具體的信息內(nèi)容）。

四、數(shù)學(xué)"魔法"背后的物理機(jī)制

研究團(tuán)隊進(jìn)一步深入，發(fā)現(xiàn)了一個非常具體的觸發(fā)條件。問題會在P矩陣中出現(xiàn)多個值恰好等于1的情況下發(fā)生。這種情況在使用"安全softmax"技術(shù)時特別容易出現(xiàn)。

安全softmax就像是一種"防止溢出"的計算技巧，它通過調(diào)整基準(zhǔn)值來避免數(shù)值計算中的溢出問題。但正是這種看似安全的操作，在特定條件下會導(dǎo)致多個注意力權(quán)重同時變成1。

當(dāng)P中有多個1，而V矩陣中的對應(yīng)值大多是負(fù)數(shù)時，就會發(fā)生一個微妙的數(shù)學(xué)現(xiàn)象。在進(jìn)行BF16格式的加法時，兩個負(fù)數(shù)相加可能會導(dǎo)致尾數(shù)溢出，需要進(jìn)行右移和重新規(guī)范化。這個過程中的舍入操作就會引入系統(tǒng)性的負(fù)偏差。

研究團(tuán)隊用一個具體的計算例子來說明這個過程。他們展示了兩個負(fù)數(shù)-2.40625和-2.296875在BF16格式下相加的完整過程。由于尾數(shù)溢出，需要右移一位，而被移出的位決定了舍入方向。在這種特定的數(shù)值分布下，舍入總是傾向于讓結(jié)果變得更負(fù)，從而引入系統(tǒng)性偏差。

五、優(yōu)雅的解決方案：動態(tài)調(diào)整的"平衡術(shù)"

找到了問題的根本原因后，研究團(tuán)隊提出了一個既簡單又優(yōu)雅的解決方案。他們的核心思路是：既然問題出在多個注意力權(quán)重同時為1，那就想辦法避免這種情況發(fā)生。

他們設(shè)計了一種"動態(tài)最大值調(diào)整"的方法。這就像是一個聰明的裁判，當(dāng)發(fā)現(xiàn)比賽中出現(xiàn)了可能導(dǎo)致爭議的情況時，會適當(dāng)調(diào)整規(guī)則來維持公平。

具體來說，當(dāng)算法檢測到一行中有多個相同的最大值時，就會動態(tài)調(diào)整歸一化因子。如果最大值是正數(shù)且有重復(fù)，就將歸一化因子調(diào)整為β倍的最大值（其中β>1）；如果最大值是負(fù)數(shù)且有重復(fù)，就將歸一化因子設(shè)置為0。

這種調(diào)整利用了softmax函數(shù)的平移不變性——也就是說，給所有輸入同時加上或減去一個常數(shù)，最終的概率分布是不變的。但是，通過巧妙地選擇這個常數(shù)，可以確保所有的注意力權(quán)重都嚴(yán)格小于1，從而避免觸發(fā)舍入偏見。

研究團(tuán)隊在實際的GPT-2訓(xùn)練中驗證了這個方法的有效性。原本會在幾千個訓(xùn)練步驟后崩潰的模型，在使用了這個簡單的修改后，能夠穩(wěn)定地繼續(xù)訓(xùn)練下去。

六、更廣泛的啟示：AI訓(xùn)練穩(wěn)定性的新理解

這項研究的意義遠(yuǎn)遠(yuǎn)超出了解決一個具體的技術(shù)問題。它揭示了在追求計算效率的過程中可能隱藏的微妙陷阱，并提供了一個系統(tǒng)性分析此類問題的方法論。

研究團(tuán)隊的發(fā)現(xiàn)也解釋了其他一些相關(guān)現(xiàn)象。比如，為什么某些被稱為"注意力沉積"的現(xiàn)象會與訓(xùn)練不穩(wěn)定性相關(guān)。注意力沉積指的是模型傾向于給某些特定位置分配很高的注意力權(quán)重，而這正好符合他們發(fā)現(xiàn)的觸發(fā)條件——容易產(chǎn)生權(quán)重為1的情況。

從更宏觀的角度來看，這項研究展示了現(xiàn)代AI系統(tǒng)的復(fù)雜性。表面上看起來簡單的數(shù)字格式選擇，實際上可能與模型的內(nèi)部數(shù)學(xué)結(jié)構(gòu)產(chǎn)生意想不到的交互效應(yīng)。這提醒我們，在設(shè)計和優(yōu)化AI系統(tǒng)時，需要更加細(xì)致地考慮各個組件之間的相互作用。

研究團(tuán)隊也坦誠地指出了他們工作的局限性。他們的分析主要基于GPT-2這一特定模型，在更大規(guī)模的模型或不同的架構(gòu)中，問題的表現(xiàn)形式可能會有所不同。此外，隨著新的低精度數(shù)字格式（如FP8）的發(fā)展，可能還會出現(xiàn)新的挑戰(zhàn)。

說到底，這項研究就像是為AI訓(xùn)練領(lǐng)域提供了一面"顯微鏡"，讓我們能夠看清那些隱藏在復(fù)雜計算背后的微妙機(jī)制。它不僅解決了一個困擾業(yè)界已久的實際問題，更重要的是，它建立了一套分析此類問題的思路和方法，為未來遇到類似挑戰(zhàn)時提供了寶貴的參考。

對于普通人來說，這項研究的價值在于展示了科學(xué)研究的魅力——即使是看似微不足道的技術(shù)細(xì)節(jié)，背后也可能隱藏著深刻的原理和廣泛的影響。正是這樣一個個看似小小的突破，推動著人工智能技術(shù)不斷向前發(fā)展，最終造福我們每個人的生活。

Q&A

Q1：BF16數(shù)字格式是什么？為什么AI訓(xùn)練要使用它？

A：BF16是一種簡化的數(shù)字表示格式，用16位來存儲原本需要32位的浮點數(shù)。使用它的原因是能節(jié)省一半的存儲空間和計算時間，讓AI模型訓(xùn)練更高效。就像用簡化的數(shù)學(xué)符號做計算一樣，既快又省資源，這對訓(xùn)練大型AI模型來說非常重要。

Q2：Flash Attention技術(shù)為什么會在低精度訓(xùn)練中出現(xiàn)問題？

A：Flash Attention在進(jìn)行矩陣計算時，當(dāng)注意力權(quán)重中出現(xiàn)多個值為1，且對應(yīng)的數(shù)據(jù)多為負(fù)數(shù)時，BF16格式的加法運算會產(chǎn)生系統(tǒng)性的舍入偏差。這種偏差不會相互抵消，反而會累積放大，最終導(dǎo)致訓(xùn)練過程崩潰。

Q3：清華團(tuán)隊提出的解決方案具體是怎么工作的？

A：他們設(shè)計了一種動態(tài)調(diào)整機(jī)制，當(dāng)檢測到可能引發(fā)問題的數(shù)值模式時，會自動調(diào)整計算中的歸一化因子。這樣可以確保注意力權(quán)重始終小于1，避免觸發(fā)舍入偏差的累積。這個方法簡單有效，且不影響模型的最終性能。