將激活、權重和梯度量化為4位，有望加速神經網絡訓練。

然而，現有的4位訓練方法需要自定義數字格式，而現代硬件不支持這種格式。

最近，清華朱軍團隊提出了一種使用INT4算法實現所有矩陣乘法的Transformer訓練方法。

使用超低INT4精度進行訓練，是非常具有挑戰性的。為了實現這一目標，研究者仔細分析了Transformer中激活和梯度的具體結構，為它們提出專用的量化器。

對于前向傳播，研究者確定了異常值的挑戰，并提出了Hadamard量化器來抑制異常值。

對于后向傳播，他們通過提出位分割，來利用梯度的結構稀疏性，并利用分數采樣技術來準確量化梯度。

這種新的算法，在自然語言理解、機器翻譯和圖像分類等廣泛任務上，都實現了具有競爭力的準確性。

原型線性算子運算速度比FP16同類算子快2.2倍，訓練速度提高了35.1%。

圖片

論文地址：https://arxiv.org/abs/2306.11987

代碼地址：https://github.com/xijiu9/Train_Transformers_with_INT4

全新的INT 4訓練算法

訓練神經網絡對計算的要求很高。使用低精度算術進行訓練（完全量化訓練/FQT）有望提高計算和內存效率。

FQT方法在原來的全精度計算圖中添加了一些量化器和反量化器，并用消耗更小的低精度浮點運算，代替了消耗更高的浮點運算。

FQT的研究旨在降低訓練數值精度，而不犧牲太多的收斂速度或精度。

所需的數值精度已從FP16降低到FP8、INT32+INT8和INT8+INT5。

FP8訓練是在帶有Transformer引擎的Nvidia H100 GPU中實現的，加速了大規模Transformer的訓練。最近的訓練數值精度，已經降到了4位。

然而，這些4位訓練方法不能直接用于加速，因為它們需要自定義數字格式，而現代硬件不支持這些格式。

首先，前向傳播中的不可微量化器，會使損失情況變得崎嶇不平，基于梯度的優化器很容易陷入局部最優。

其次，梯度僅僅以低精度近似計算。這種不精確的梯度會減慢訓練過程，甚至導致訓練不穩定或發散。

而在這項工作中，研究者為Transformer提出了一種新穎的INT4訓練算法。

圖片

訓練Transformer的所有高消耗的線性運算，都可以寫在矩陣乘法（MM）的形式中。

這種MM形式，可以讓我們設計更靈活的量化器，通過利用Transformer中激活、權重和梯度的特定結構，就可以更好地近似于FP32矩陣乘法。

隨機數值線性代數 (RandNLA) 領域的進步，被這種量化器充分利用。

對于前向傳播，研究者發現，激活中的異常值是精度下降的主要原因。

為了抑制異常值，他們提出了Hadamard量化器，它會對激活矩陣的變換版本進行量化。這種變換是塊對角Hadamard矩陣，它將離群值中攜帶的信息傳播到矩陣的鄰近條目，從而縮小了離群值的數值范圍。

對于后向傳播，他們利用了激活梯度的結構稀疏性。研究者發現，一些token的梯度非常大。同時，其余大多數token的梯度非常均勻，甚至比較大梯度的量化殘差更均勻。

圖片

因此，與其計算所有梯度，不如節省計算較大梯度殘差的計算資源。

為了利用這種稀疏性，研究者提出了位分割，將每個token的梯度分割為高4位和低4位。

然后，通過杠桿分數采樣（leverage score sampling）來選擇信息最豐富的梯度，這是RandNLA的一種重要采樣技術。

圖片

結合前向和后向傳播的量化技術，研究者提出了一種使用INT4MM進行Transformer中所有線性運算的算法， 并且評估了在各種任務上訓練Transformer的算法，包括自然語言理解、問答、機器翻譯和圖像分類。

與現有的4位訓練算法相比，他們的算法實現了有競爭力的或更高的精度。

此外，這種算法與GPU等當代硬件兼容，因為它不需要FP4或對數格式等自定義的數字格式。

這種原型量化+INT4 MM算子實現，速度比FP16MM基線快2.2倍，并且將訓練速度提高了35.1%。

相關工作

完全量化訓練

完全量化訓練 (FQT) 方法通過將激活、權重和梯度量化為低精度來加速訓練，因此訓練期間的線性和非線性算子可以用低精度算術來實現。

FQT的研究設計了新穎的數值格式和量化算法，可以更好地逼近全精度張量。

目前的研究前沿是4位FQT。由于梯度的數值范圍很大以及從頭開始訓練量化網絡的優化問題，FQT具有挑戰性。

由于這些挑戰，現有的4位FQT 算法在某些任務上的精度仍然下降了1-2.5%，并且無法支持當代硬件。

圖片

其他有效的訓練方法

混合專家在不增加訓練預算的情況下提高了模型容量。

結構性dropout利用計算有效的方法來正則化模型。高效的注意力降低了計算注意力的二次時間復雜度。

分布式訓練系統通過利用更多的計算資源，減少了訓練時間。

研究者降低數值精度的工作與這些方向具有正交性。

圖片

前向傳播

神經網絡訓練是一個迭代優化過程，通過前向和后向傳播計算隨機梯度。

研究團隊使用4位整數(INT4)算法加速前向和后向傳播。

正向傳播能以線性和非線性(GeLU, normalization, softmax等)算子的組合來實現。

在我們的訓練過程中，我們用INT4算術加速所有線性運算符，并將所有計算量較小的非線性運算符保留在16位浮 點(FP16)格式中。

Transformer中的所有線性運算都可以寫成矩陣乘法(MM)的形式。

為了便于表述，本文考慮以下簡單矩陣乘法的加速：

圖片

這種MM的最主要用例是全連接層。

考慮一個輸入形狀為(批量大小S，序列長度T，維度D)的Transformer。

全連接層可以表述成上邊的公式，其中X是N = STtoken的激活，W是權重矩陣。

對于注意力層，可能需要批量矩陣乘法(BMMS)。

我們提出的技術可以應用于BMMS。

學習步長量化（Learned Step Quantization）

為了加速訓練，必須使用整數運算來計算前向傳播。

研究人員為此目的，利用學習步長量化器（LSQ）。

LSQ是靜態量化，他的量化尺度不依賴于輸入的方法，因此比動態方法消耗更小，量化方法，需要在每次迭代時動態計算量化尺度。

激活異常值

簡單地將LSQ應用到具有4位激活/權重的FQT會導致精度下降，因為會激活異常值。

圖片

如上圖所示，激活有一些離群值條目，它們是其規模比其他條目大得多。

不幸的是，Transformers傾向于將信息存儲在這些異常值中，而且這樣的截斷會嚴重損害準確性。

當訓練任務是在一些新的下游任務上微調預訓練模型時，異常值問題尤為明顯。

因為預訓練模型比隨機初始化包含更多的異常值 。

Hadamard量化

我們提出了Hadamard量化（HQ）來解決異常值問題。

其主要思想是將另一個具有較少異常值的線性空間中的矩陣進行量化。

激活矩陣中的異常值形成了一個特征結構（feature-wise structure）。

他們通常集中在幾個維度上，也就是說X中只有幾列顯著大于其他列。

哈達瑪變換（Hardamand transform）是一個線性變換，它可以將異常值分攤到其他條目中。

后向傳播

現在我們考慮使用INT4操作來加速線性層的后向傳播。

我們將在本節中討論激活梯度/權重梯度的計算。

梯度的結構稀疏性

我們注意到，在訓練過程中梯度矩陣往往非常稀疏。

而且稀疏性具有這樣的結構：

的幾行（比如tokens）具有較大的條目，而大多數其他行卻接近全零向量。

圖片

這種結構稀疏性源于現代神經網絡的嚴重過度參數化。

幾乎在整個訓練過程中，網絡都以超參數化方案運行，除了一些困難的例子之外，它可以很好地適應大多數訓練數據。

因此，對于擬合良好的數據點，（激活）梯度將接近于零。

研究人員發現對于預訓練任務，例如，經過幾個訓練周期后，結構稀疏性很快就會出現。

對于微調任務，梯度整個訓練過程中始終是稀疏的。

位分割（Bit Splitting）和杠桿分數采樣(Leverage Score Sampling)

如何設計梯度量化器，以利用結構稀疏性在反向傳播期間準確計算MM呢？

高級的思路是：梯度的許多行都是如此小，對參數梯度影響很小，但浪費了大量的計算量。

另一方面，大行無法用INT4精確表示。

我們放棄掉一些小行并使用節省下來的計算能力來更準確地表示大行。

實驗

研究人員在包括語言模型在內的各種任務上評估我們的INT4訓練算法微調、機器翻譯和圖像分類。

研究人員用CUDA和cutlass執行了他們提出的HQ-MM和LSS-MM算法。

研究人員用INT4實現替換所有浮點線性運算符，但沒有簡單地使用LSQ來嵌入層，并保持最后一個分類器層的精度。

最后研究人員對所有評估的模型采用了默認架構、優化器、調度器和超參數。

收斂模型精度

研究人員在下表中比較了收斂模型在各種任務上的準確性。

圖片

作為對照的方法包括全精度訓練（FP）、INT8訓練（INT8）、FP4訓練（「超低」），使用LSQ進行激活和權重(LSQ+LUQ)的4 位對數量化，以及我們這種利用HQ進行前向傳播，利用LSS進行反向傳播（HQ+LSS）的算法。

「超低」沒有公開的實現，因此我們僅列出了它在機器上的原始論文中的性能翻譯任務。

除了大型機器翻譯任務和大型視覺Transformer任務之外，我們將每次運行重復三次，并將標準差報告為表中的下標。

研究人員沒有進行任何類型的知識蒸餾或數據增強。

消融實驗

研究人員進行的消融實驗目的是展示前向和后向方法的有效性。

研究不同量化器的前向傳播的有效性，我們將后向傳播留在FP16中。

結果如下圖所示。

圖片

計算和內存效率

最后，研究人員通過評估他們的原型實現，展示了他們的方法加速神經網絡訓練的潛力。

而且他們的實施還沒有完全優化。

研究人員也沒有將線性算子與非線性和歸一化進行融合。

因此，結果不能完全反映INT4訓練算法的潛力。

完全優化的實施需要大量工程，超出了我們論文的討論范圍。

結論

研究人員提出了一種對硬件很友好的Transformer INT4的訓練方法。

通過分析Transformer中MM的屬性，研究人員提出了HQ和LSS方法來量化激活和梯度，同時保持準確性。

在幾個重要任務上，我們的方法與現有的INT4方法表現相當，甚至更好。

研究人員的這些工作可能會擴展到除了Transformers之外的其他MM架構中，例如 MLP-Mixer、圖神經網絡和循環神經網絡網絡。

這是他們未來的研究方向。

更廣泛的影響：研究人員的算法可以提高效率并減少訓練神經網絡的能源消耗，這有助于減少深度學習造成的碳排放。

但是，高效的訓練算法還可能促進那些，對于人來安全存在隱患的大語言模型和惡意人工智能應用程序的開發。

比如，會被用于虛假內容生成的相關模型和應用。

限制：這項工作的主要限制是它只能加速具有較大規模的矩陣乘法（線性層）的大模型，但不能加速卷積層。

而且，所提出的方法還不能很好地適用于OPT-175B等超大模型。

據我們所知，即使是INT8訓練對于這些超大型模型來說仍然是尚待解決的問題。