訓練模型時我們盯著 tensorboard 看著 training loss 一直降低直到收斂, 收斂後每個 checkpoint 的 training loss 都差不多, 那該挑哪一個 checkpoint 呢?
就選 validation loss 最低的那些吧, 由 PAC 我們知道 validation error 約等於 test error (validation set 愈大愈好), 但我們能不能對泛化能力做得更好? 如果 training 時就能讓泛化能力提升, 是否更有效率?

Motivation

很多提升泛化能力的論文和觀點都從 “flat“ optimum 出發. 下圖清楚說明這個想法 ([圖來源]):

直接看 SAM 怎麼 update parameters, 論文的 figure 2:
目前的 weight $w_t$ 的 gradient step 為 $-\eta\nabla L(w_t)$, update 後會跑到 $w_{t+1}$.
SAM 會考慮 $w_t$ locally loss 最大的那個位置 ($w_{adv}$), 用該位置的 gradient descent vector $-\eta\nabla L(w_{adv})$, 當作 weight $w_t$ 的 gradient step, 因此才會跑到 $w_{t+1}^{SAM}$.
先把 SAM 的 objective function 主要目的點出來, SAM 相當於希望找出來的 $w$ 其 locally 最大的 loss 都要很小, 直覺上就是希望 $w$ 附近都很平坦, 有點類似 Support Vector Machine (SVM) 的想法, 最小化最大的 loss.

以下數學推導… 數學多請服用

這是林軒田教授在 Coursera 機器學習基石上 (Machine Learning Foundations)—Mathematical Foundations Week4 的課程筆記.
說明了為什麼我們用 training data 學出來的 model 可以對沒看過的 data 有泛化能力, 因此機器學習才有可能真正應用上.
課程單元的這句話總結得很好 “learning can be probably approximately correct when given enough statistical data and finite number of hypotheses”
以下為筆記內容

總結一下 (up to 2024-02-17) 目前學習的量化技術和流程圖, 同時也記錄在 github 裡.
Post Training Quantization (PTQ) 稱事後量化. Quantization Aware Training (QAT) 表示訓練時考慮量化造成的損失來做訓練
為了得到 fixed point model 可以對事先訓練好的 float model 做 PTQ 或 QAT, 或是直接從 QAT 流程開始
同時 QAT 也可以用 PTQ 來初始化訓練. 如果要從 float model 開始做量化的話, 可以考慮在訓練 float model 時就對之後量化能更加友善的技術 (如 R^2, KURE, PACT)

接著對每個技術點盡量以最簡單的方式解說. 如果對量化還不是那麼熟悉, 建議參考一下文章後半段的”簡單回顧量化”

⚠️ 可能寫的比較瑣碎和雜亂, 主要給自己筆記用

令 $f:\mathbb{R}^{n}\rightarrow\mathbb{R}^{m}$ 的 Jacobian matrix 為 $J_f(x)$ 是 $(m\times n)$ 矩陣, 而 Hessian 為 $H_f(x)$ 是 $(m\times n \times n)$ 高維 tensor
 $\circ$ VJP 稱為 Vector-Jacobian Product, $vJ_f(x)$, 其中 $v$ 是 ($1\times m$) 的 row vector
 $\circ$ JVP 稱為 Jacobian-Vector Product, $J_f(x)v$, 其中 $v$ 是 ($n\times 1$) 的 column vector
 $\circ$ HVP 稱為 Hessian-Vector Product, $H_f(x)v$, 其中 $v$ 是 ($n\times 1$) 的 column vector
計算 $vJ_f(x)$ 不用先把矩陣 $J_f(x)$ 求出來再跟 $v$ 相乘, 而是可以直接得到相乘的結果(這樣做還更快), 聽起來有點矛盾對吧~同樣的 JVP 和 HVP 也是如此
本文會說明怎麼高效率計算 VJP, JVP, Jacobian, Hessian, 以及 HVP

主要參考 PyTorch 文章: JACOBIANS, HESSIANS, HVP, VHP, AND MORE: COMPOSING FUNCTION TRANSFORMS

HVP 可以用來有效率地計算 $tr(H_f(x))$, 而這個 term 有時候會被當作 loss 來用, 舉例來說:
 $\circ$ Sliced Score Matching (SSM) 會用到
 $\circ$ EWGS quantization (Network Quantization with Element-wise Gradient Scaling, arxiv) 會用到
 $\circ$ More and details see: Thoughts on Trace Estimation in Deep Learning, 更多例子且有非常深入的討論
總結可以參考文末 Summary
先把 function $f$ 定義好: (名字為predict)

接續上一篇
現在我們可以真正的來探討以下問題了:
 A. 速度的分析: 加速到什麼程度? 跟小模型的速度和準確度有關聯嗎? (想像如果 draft 一直被拒絕, 則小模型都是多跑的)
 B. 運算量的分析: Operation 數 (計算量) 也會減少嗎? 還是會增加?
 C. Memory bandwidth 的分析: 會減少還是增加?
 D. Performance 能維持住嗎 (PPL, WER, BLEU, … 端看 model task 是什麼): 還是會有 degrade?

這是 Transformer inference 的加速, 有人猜測 GPT-4 也使用這個方法: https://archive.ph/2RQ8X
Speculative decoding 做到了不影響準確率情況下直接加速 (不改 model 架構, 不 fine tune, 不做 PTQ 等)
這麼神奇的操作就是利用了一個小模型來先跑一些 tokens, 再由原來的大模型評估或修正.
論文顯示 LLM 效果無損直接可提速 2~3 倍, 讓我們看下去

Motivation

使用 SongHan 教授的課程 slides.
利用 small model 先提出一些 draft tokens, 然後用 large model 來驗證. 如果大部分都接受, 直覺上可以省去很多 large model 的呼叫次數, 因此加速. 方法十分簡單, 不過其實魔鬼藏在細節裡, 跟原本只使用 large model 的方法比較有幾個問題要回答:
 A. 速度的分析: 加速到什麼程度? 跟小模型的速度和準確度有關聯嗎? (想像如果 draft 一直被拒絕, 則小模型都是多跑的)
 B. 運算量的分析: Operation 數 (計算量) 也會減少嗎? 還是會增加?
 C. Memory bandwidth 的分析: 會減少還是增加?
 D. Performance 能維持住嗎 (PPL, WER, BLEU, … 端看 model task 是什麼): 還是會有 degrade?
Google 這篇論文很精彩的理論分析了以上所有問題, 並有實務驗證
先破題, performance (PPL, WER, BLEU, …) 可以保證維持住! 我們等到本篇筆記最後在討論, 以下會先討論算法流程、加速和運算量的分析.

如同 SmoothQuant 論文裡的圖, 在 memory size 已經跟不上算力和模型大小情況下, memory bandwidth 已經變成 bottleneck.
如何降低 memory 使用量將變的很關鍵, 因此 Activation-aware Weight Quantization (AWQ) 這篇文章就專注在 Weight Only Quantization (WOQ), 顧名思義就是 weight 使用 integer 4/3 bits, activations 仍維持 FP16.
因為 computation is cheap, memory is expensive.

Intel® Neural Compressor 有實作 WOQ 裡面有 AWQ

以下內容直接筆記 MIT SongHan 教授的課程內容[slides], [Video]

這是 MIT SongHan 教授實驗室的論文, 使用 PTQ 對 LLM 做到 W8A8 的量化, 由於 activations 會有比較大的 outliers 導致 quantization 後損失較大, 而一般 weights 的 outliers 很少, 因此透過一些等價的轉換將 activations 的 scale 縮小並放大 weights 的 scale, 使得 activations 變的較容易 quant 而 weights 仍然容易 quant. 如論文的圖顯示:

總歸來說 Data-Free Quantization (DFQ) 的目的是讓 floating model 做 weights 各種調整, 使得不管是 weights or activations 都變得適合 per tensor 量化.這樣理想上就不需用到 per channel 量化, 因為 per channel 雖然效果很好, 但硬體比較不友善, 且花的運算量較高. 另外 DFQ 屬於 Post-Training Quantization (PTQ) 方法. PTQ 對佈署到 edge 端很方便, 但一般來說 PTQ 都不如 Quantization-Aware Training (QAT) 的效果好, 因此 DFQ 嘗試提升效果.

DFQ 共四步, 對照圖看, 需照順序:

1. Cross-Layer Equalization (CLE): 輸入 fused BN 後的 float model $M_f^1$, floating 操作對 weights 做調整使得更均衡方便 per tensor 量化, 為 step 3 的前置作業, 輸出仍為 float model $M_f^2$.
2. Bias Absorption (BA): 輸入 CLE 後的 float model $M_f^2$, floating 操作對 activations 做調整使得更均衡方便 per tensor 量化, 為 step 3 的前置作業, 輸出仍為 float model $M_f^3$.
3. PTQ 量化: 輸入 CLE+BA 後的 float model $M_f^3$, 此時不管 weights or activations 都適合做 per-tensor 量化了, 所以直接 PTQ 輸出 int model $M_i^1$.
4. Bias Correction (BC): 輸入 float model $M_f^1$ 和 step 3 的 $M_i^1$, 並且(option)給一些 unlabeled 的代表 data, BC 會對 $M_i^1$ 的 bias 參數補償因為量化造成的數值 mean 偏移, 輸出為最終 fixed point model $M_i^2$.

Qualcomm AI Lab 的 tool AIMET 說 BC 這一步驟可以用 AdaRound (需要一小部分的 unlabelled training data) 取代

其實認真看完論文, 覺得限制有點多啊! 很多時候不能套 CLE, 有時 BA 也用不了. 把限制條列一下: