LLM 修正錯誤的能力有多強？使用 Keras 與 TPUs 的聊天機器人競技場實驗

TL;DR

Hugging Face 進行了一項聊天機器人競技場（chatbot‑arena）實驗，要求 LLM 為日曆工具生成 API 呼叫，並根據簡短的英文回饋修正錯誤。Gemma 2 9B 表現最佳，能可靠地修正錯誤；Llama 3.1 8B 表現接近但需要額外的提示；較小的模型（1-2B）和較舊的模型（如 Vicuna）經常失敗或產生亂碼輸出。

實驗設置

測試使用了一個包含兩個函數的 Python API：action.add_calendar_entry(title, date="YYYY-MM-DD", time="HH:MM", duration=m) 和 action.remove_calendar_entry(title, date, time)。系統提示詞指示模型扮演一個得力的語音助手，並以單行可執行代碼進行回覆。對話從一個不精確的請求開始（「Add a meeting with Fred on Nov 11 at 5PM」），其中省略了年份，迫使模型犯錯，隨後可以透過後續陳述「The current year is 2024.」來修正。隨後的對話增加了事件、更改了持續時間並取消了分錄，提供了多次修正錯誤的機會。

技術實現：TPUs、JAX、Keras 與模型分片

該競技場是使用 Gradio 在 Hugging Face Spaces 上構建的。它運行在 TPU v5e 2x4 上，每個核心提供 8 個核心和 16 GB RAM（總計 128 GB）。這份記憶體讓作者能夠透過將模型分片（sharding）到所有核心來同時加載多個模型。使用 Keras（現在已原生運行於 JAX），他們以 bfloat16 格式加載了最多五個約 8B 參數的模型和三個約 2B 參數的模型，總共同時有七個 LLM 駐留在記憶體中。模型並行（Model parallelism）依賴於 Keras 內建的佈局圖（例如 keras_hub.models.Llama3Backbone.get_layout_map(device_mesh)）將每個模型分片到八個 TPU 核心上。作者指出，為了實現穩定加載，需要進行除錯和偶爾的佈局圖調整。

評估的模型

實驗重點關注 10B 參數以下的指令微調（instruction-tuned）模型，選擇這些模型是因為它們可以共同放入 TPU 記憶體中，且任務對它們來說足夠簡單。測試的系列包括 Gemma、Llama 3、Mistral、Vicuna 和 CodeGemma。具體檢查點包括：

  • Gemma 2 9B‑instr
  • Llama 3.1 8B‑instr
  • Llama 3.2 3B‑instr
  • Llama 3.2 1B‑instr
  • Gemma 2B‑instr
  • CodeGemma 7B‑instr
  • Vicuna 1.5 7B‑instr
  • Mistral 7B‑instr 所有模型均透過本文中列出的 Hugging Face Hub URL 進行訪問。

第一個問題的可靠性

每個模型都被要求回答初始提示詞五次。結果（✓ = 正確的 API 呼叫，🍄 = 基本正確但有錯誤，🔥 = 垃圾/無法識別的呼叫）：

  • Gemma 2 9B‑instr: ✓ ✓ ✓ ✓ ✓
  • Llama 3.1 8B‑instr: ✓ ✓ ✓ ✓ ✓
  • Llama 3.2 3B‑instr: ✓ ✓ ✓ ✓ ✓
  • Llama 3.2 1B‑instr: 🔥 🍄 🔥 🔥 🔥
  • Gemma 2B‑instr: 🍄 🍄 🍄 🍄 ✓
  • CodeGemma 7B‑instr: ✓ ✓ ✓ ✓ ✓
  • Vicuna 1.5 7B‑instr: ✓ 🔥 🔥 ✓ 🔥
  • Mistral 7B‑instr: ✓ ✓ 🍄 ✓ 🍄 只有 1-2B 模型和較舊的 Vicuna 一直失敗；較大的模型每次都能正確回答。

完整對話 – 修正錯誤

當運行完整的六輪對話時，作者在每次錯誤後提供回饋，並記錄模型是否能修正它（🥦 = 成功修正）。關鍵結果：

  • Gemma 2 9B‑instr 和 Llama 3.1 8B‑instr 完成了對話，僅有一次或兩次輕微失誤；Llama 需要額外的「fix it」提示來獲得它的青菜（broccoli）。
  • 在線 Gemini 運行（規模大得多的模型）需要特殊的提示來輸出 API 呼叫，且仍然犯了幾個錯誤，這表明單憑規模並不能保證在任務上有更好的表現。
  • 在小型模型中，只有 Gemma 2B‑instr 設法完成了對話，儘管它有重複添加閒聊文本（「Sure, here’s the updated code…」）和混淆日期/時間的傾向；它可以在被要求時修正錯誤。
  • Vicuna 1.5 7B‑instr 在恢復之前經常退化為重複或垃圾輸出（🔥🔥）。
  • Mistral 7B‑instr 犯了很多錯誤，但能夠修正它們，在多次糾正回合後獲得了六個青菜符號。
  • CodeGemma 7B‑instr 產生了一個頑固的重複錯誤：它在年份中插入了一個空格（「20 24」），儘管經過糾正也無法消除它。

其他錯誤修正測試

為了隔離修正能力，作者將其他模型預先生成的錯誤輸出提供給每個模型，並檢查模型是否可以修復它們。結果（🥦 = 正確修正，🍄 = 錯誤持續存在，🔥 = 多個錯誤）：

  • Gemma 2 9B‑instr: 完美修復了錯誤的時間和「API only」提示；在處理「Wrong API」提示時感到吃力，在輸出正確呼叫的同時還附帶了道歉（這是一個重複的錯誤，可以透過要求「API call only please」來消除）。
  • Llama 3.1 8B‑instr: 與 Gemma 類似，修復了大多數錯誤，但在可靠地修正錯誤的 API 呼叫方面遇到困難。
  • CodeGemma 7B‑instr: 行為與 Gemma 類似，持續保持道歉模式。
  • Mistral 7B‑instr: 很好地修復了時間錯誤，但在處理 API-only 和 wrong-API 提示時表現出不一致性。
  • 所有較小的模型（Llama 3.2 3B, Llama 3.2 1B, Gemma 2B）和 Vicuna 都產生了許多 🔥 或 🍄 結果，表明它們無法可靠地修復提供的錯誤。

總結與啟示

作者原以為簡單的雙 API 設置對所有模型來說都很容易，但結果顯示出明顯的分層。Gemma 2 9B‑instr 是唯一一個以近乎完美的正確性且只需極少修正提示就完成整個對話的模型。Llama 3.1 8B‑instr 位居第二，需要稍微多一點的引導。較大的規模有所幫助，但規模大得多的 Gemini 模型由於格式問題表現不佳，這表明指令忠誠度（instruction fidelity）和 token 級別的控制比原始規模更重要。較小的模型（≤2B）和較舊的模型經常根本無法產生有效的 API 呼叫，即使產生了，它們也很難在沒有大量提示的情況下納入糾正回饋。該實驗表明，對於用戶可以用簡單英文糾正錯誤的簡單工具使用場景，目前的次中型指令微調 LLM 是有用的，但可靠性仍取決於模型規模、訓練數據以及抑制多餘文本的能力。

讀者可以使用公開可用的 Keras Chatbot Arena Space 來複製或擴展這項研究，並鼓勵嘗試透過微調來改進錯誤修正行為。

Sources