LLM 修复错误的能力如何？使用 Keras 和 TPUs 进行的聊天机器人竞技场实验

TL;DR

Hugging Face 进行了一项聊天机器人竞技场实验，要求 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）。这种内存允许作者通过将模型分片到所有核心来同时加载多个模型。使用 Keras（现在原生运行在 JAX 上），他们以 bfloat16 格式加载了多达五个约 8B 参数的模型和三个约 2B 参数的模型，总共同时有七个 LLM 驻留在内存中。模型并行依赖于 Keras 内置的布局映射（例如，keras_hub.models.Llama3Backbone.get_layout_map(device_mesh)）将每个模型分片到八个 TPU 核心上。作者指出，为了实现稳定加载，需要进行调试和偶尔的布局映射调整。

评估的模型

实验重点关注参数量在 10B 以下的指令微调模型，选择这些模型是因为它们可以共同放入 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”提示词才能赢得它的西兰花。
  • 在线 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 模型由于格式问题表现不佳，这表明指令忠实度和 Token 级别的控制比原始规模更重要。较小的模型 (≤2B) 和较旧的模型经常根本无法产生有效的 API 调用，即使产生了，它们也很难在没有大量提示的情况下整合纠正性反馈。该实验表明，对于用户可以用 plain English 纠正错误的简单工具使用场景，当前的中间规模指令微调 LLM 是有用的，但可靠性仍取决于模型规模、训练数据以及抑制无关文本的能力。

读者可以使用公开可用的 Keras Chatbot Arena Space 复制或扩展该研究，并鼓励通过微调来改进错误修复行为。

Sources