邁向穩健的小型語言模型強化學習：架構解析與實務判斷

在開發具備自主代理能力（Agentic AI）的系統時，70M 至 500M 參數級別的小型語言模型（SLM）因其極低的推論延遲與邊緣運算（On-Device）適應性，成為業界的熱門選擇。然而，相較於千億參數的巨型模型，要在這個微量級距中利用 PPO（Proximal Policy Optimization） 進行人類偏好對齊（RLHF）一直被視為極度不穩定的玄學。

近期發表的論文《Towards Robust Reinforcement Learning for Small-Scale Language Model Agents》針對此現象進行了系統性的大規模實證，透過 15 組 (模型, 語料庫) 的交叉實驗，成功拆解了 SLM 規模下 PPO 常見的三大崩潰模式，並提出了極具實務價值的「能力空間假說（Capacity-Headroom Hypothesis）」。

實驗設定與模型矩陣

為確保結果具有普適性，研究團隊選擇了兩個不同架構的開源 SLM 家族：

1. Pythia 系列：包含 70M、160M、410M 三種尺寸，基於 GPT-NeoX 架構。
2. SmolLM2 系列：包含 135M、360M 兩種尺寸，基於 Llama 架構（包含 RoPE 與 SwiGLU）。

這些模型分別在三個難度不同的語料庫（TinyStories、CNN/DailyMail、Wikitext-103）上進行了完整的 SFT -> Reward Model -> PPO 訓練循環。

圖一：小型語言模型代理的端到端 RLHF 流程。包含資料處理、SFT 訓練、獎勵模型訓練，以及最終具備三層安全機制的 PPO 穩定微調。

為什麼 SLM 的 PPO 那麼容易崩潰？三大實務地雷

許多開發者在對 500M 以下的 SLM 進行 PPO 訓練時，常會遇到梯度爆炸或模型完全講出亂碼的情況。論文總結了三個在底層架構上最常被忽略的「靜默殺手」：

1. 梯度流動受阻（Gradient Flow Obstruction in PEFT）

在使用 TRL 框架結合 LoRA 進行參數量化微調（PEFT）時，系統可能會因為實作上的漏洞，靜默地將 LoRA 參數凍結（Non-trainable）。這會導致 Policy 模型雖然能不斷生成 Roll-outs 並計算 Loss，但底層權重根本沒有被更新。 解決方案：採用 Merge-and-Reinitialize 技巧。先將 SFT 的 LoRA 權重合併至 Base Model（並將此視為 Reference Model），接著掛載一個全新（Zero-initialized）的 LoRA 進行 PPO 更新，確保梯度流暢。

2. bf16 降精度帶來的數值溢出（Numerical Instability in bfloat16）

在 PPO 演算法中，計算重要性比率（Importance Ratios, $\rho_t$）時需要計算兩個對數機率（log-probabilities）的差值。由於 bfloat16 僅有 7-bit 的尾數精度（Mantissa），在小於 200M 參數的模型中，這極易導致「災難性抵消（Catastrophic Cancellation）」。在訓練的最初幾步，比率值甚至會飆升超過 $10^6$，引發硬體層級的 NaN/Inf 異常。 解決方案：在 PPO 的核心迴圈中（包含 Policy、Reference Model、Value Head 與 Reward Model），強制將所有張量運算切換回 float32 精度。

3. 分布崩潰（Distributional Collapse）

長尾的獎勵分佈若搭配未設限的 KL 散度懲罰（KL Penalty），會促使 Optimizer 將 Policy 推向 Reference Model 認定極低機率的極端區域，最終產生完全無意義的亂碼。 解決方案：實作三層控制系統：

• Reward Whitening 與 $3\sigma$ 截斷：限制 Advantage 估計的極端值。
• Importance-ratio 門檻保護：當 Batch 的平均重要性比率超過 5 時，直接跳過該 Mini-batch 更新。
• 權重回滾（Weight-Rollback）：一旦偵測到 NaN/Inf，立即退回前一步的 Optimizer 狀態。

PPO 與 SFT 獎勵的表現對比

論文在所有 15 組設定中比較了 PPO 訓練後的模型與原本 SFT 模型的獎勵得分。下圖展示了這個對比：

圖二：15 種設定的 SFT 與 PPO 獎勵比較。落在虛線（基準線）上方的標記代表 PPO 成功帶來了效能提升。可以看到 Pythia-410M 與 SmolLM2-360M 在 TinyStories 上有顯著的向右上方偏移。

從結果可以發現，Pythia-410M 與 SmolLM2-360M 模型在 TinyStories 資料集上取得了最大的獎勵增幅（$\Delta = +1.355$ 與 $+0.724$），並且對戰勝率接近 60%。然而，70M 的極小模型則幾乎沒有提升，甚至在某些資料集上出現衰退。

能力空間假說（Capacity-Headroom Hypothesis）：何時該用 PPO？

這篇論文最核心的實務貢獻，是打破了「參數越少，RL 越沒用」的迷思，並提出了清晰的判斷準則——「能力空間假說」。

研究顯示，PPO 能否在小模型上發揮作用，並不取決於模型絕對參數量的多寡，而是取決於兩個先決條件：

1. 流暢的 SFT 先驗（Fluent SFT Prior）
2. 具備鑑別度的獎勵訊號（Discriminative Reward Signal）

圖三：能力空間假說的實證圖。橫軸為 SFT 困惑度（對數尺度），縱軸為 PPO 帶來的獎勵增幅。當 SFT 先驗足夠流暢（PPL < 20）時，PPO 才能有效運作。

PPL < 20 的黃金交叉線

如圖三所示，SFT 模型的困惑度（Perplexity, PPL）與 PPO 能帶來的獎勵增幅呈強烈的負相關，並且在 $\text{PPL} \approx 20$ 處出現明顯的轉折點：

• $\text{PPL} < 20$：模型具備足夠的語言流暢度，能將生成的樣本維持在 Reward Model 可靠的訓練分佈內，此時 PPO 能帶來顯著的成效與獎勵提升。
• $\text{PPL} \in [20, 50]$：預期增幅極其有限，甚至可能出現效能衰退（Regression）。這時的算力資源與其拿去跑 PPO，不如拿去清理 SFT 數據或是提升 LoRA Rank。
• $\text{PPL} > 50$：模型連話都說不清楚，給出的 Gradient 就像雜訊，PPO 極有可能直接崩潰。

此外，論文也透過消融實驗（Ablation Study）證明，若不加上述的三層防護機制，未經保護的 PEFT PPO（即便 PPL 達標）也會在最初的幾個 mini-batch 內發生 NaN 錯誤並宣告訓練失敗。

結論與工程啟示

《Towards Robust Reinforcement Learning for Small-Scale Language Model Agents》為終端 Edge Agent 開發團隊提供了一套極具可操作性的 PPO 實務指南。它告訴我們，在為資源受限環境打造 AI Agent 時，與其盲目擴充參數或放棄 PPO 改用 DPO，不如先審視你的 SFT PPL 是否達標，並老老實實地將浮點運算精度與防護機制（Safety Mechanisms）做好。

只有在穩固的工程腳手架之上，小巧精悍的 SLM 才能真正展現出超越體型的智慧。