How good are LLMs at fixing their mistakes? A chatbot arena experiment with Keras and TPUs

TL;DR

Hugging Face ran a chatbot‑arena experiment where LLMs had to generate API calls for a calendar tool and then correct mistakes based on short English feedback. Gemma 2 9B performed best, fixing errors reliably; Llama 3.1 8B was close but needed extra prompts; smaller models (1‑2B) and older models like Vicuna often failed or produced garbled output.

Experiment setup

The test used a two‑function Python API: action.add_calendar_entry(title, date="YYYY-MM-DD", time="HH:MM", duration=m) and action.remove_calendar_entry(title, date, time). The system prompt instructed the model to act as a helpful vocal assistant and to reply with a single line of executable code. The dialog started with an imprecise request (“Add a meeting with Fred on Nov 11 at 5PM”) that omitted the year, forcing the model to make a mistake that could later be corrected with the follow‑up statement “The current year is 2024.” Subsequent turns added events, changed durations, and cancelled entries, providing multiple opportunities for mistake fixing.

Technical implementation: TPUs, JAX, Keras, and model sharding

The arena was built on Hugging Face Spaces using Gradio. It ran on a TPU v5e 2x4, which provides eight cores and 16 GB of RAM per core (128 GB total). This memory allowed the authors to load several models simultaneously by sharding them across all cores. Using Keras (which now runs natively on JAX) they loaded up to five ∼8B‑parameter models and three ∼2B‑parameter models in bfloat16 format, for a total of seven LLMs resident in memory at once. Model parallelism relied on Keras’ built‑in layout maps (e.g., keras_hub.models.Llama3Backbone.get_layout_map(device_mesh)) to shard each model across the eight TPU cores. The authors note that debugging and occasional layout‑map adjustments were required to achieve stable loading.

Models evaluated

The experiment focused on instruction‑tuned models under 10 B parameters, chosen because they could fit together in TPU memory and the task was simple enough for them to handle. The families tested were Gemma, Llama 3, Mistral, Vicuna, and CodeGemma. Specific checkpoints included:

  • 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 All models were accessed via their Hugging Face Hub URLs listed in the post.

Reliability on the first question

Each model was asked the initial prompt five times. Results (✓ = correct API call, 🍄 = mostly correct but with a mistake, 🔥 = garbage/no recognizable call):

  • 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: ✓ ✓ 🍄 ✓ 🍄 Only the 1‑2B models and the older Vicuna consistently failed; the larger models answered correctly every time.

Full dialogue – fixing mistakes

When the complete six‑turn dialog was run, the authors gave feedback after each mistake and recorded whether the model could fix it (🥦 = successful fix). Key outcomes:

  • Gemma 2 9B‑instr and Llama 3.1 8B‑instr completed the dialog with only one or two minor slips; Llama needed an extra “fix it” prompt to earn its broccoli.
  • An online Gemini run (much larger model) required special prompting to output API calls and still made several mistakes, showing that size alone does not guarantee better performance on this task.
  • Among the small models, only Gemma 2B‑instr managed to finish the dialog, albeit with a recurrent tendency to add chatty text (“Sure, here’s the updated code…”) and to mix up dates/times; it could fix mistakes when asked.
  • Vicuna 1.5 7B‑instr often devolved into repeated or junk output (🔥🔥) before recovering with remaining errors.
  • Mistral 7B‑instr made many mistakes but was able to fix them, earning six broccoli symbols after numerous corrective turns.
  • CodeGemma 7B‑instr acquired a sticky recurrent mistake: it inserted a space in the year (“20 24”) and could not eliminate it despite correction.

Additional mistake‑fixing tests

To isolate correction ability, the authors fed each model pre‑generated erroneous outputs from other models and checked whether the model could repair them. Results (🥦 = correct fix, 🍄 = error persisted, 🔥 = multiple errors):

  • Gemma 2 9B‑instr: fixed wrong time and “API only” prompts perfectly; struggled with “Wrong API” prompts, outputting an apology alongside the correct call (a recurrent mistake that could be removed by asking for “API call only please”).
  • Llama 3.1 8B‑instr: similar to Gemma, fixed most errors but had trouble reliably correcting a wrong API call.
  • CodeGemma 7B‑instr: behaved like Gemma, persisting with the apology pattern.
  • Mistral 7B‑instr: fixed time errors well, but showed inconsistency with API‑only and wrong‑API prompts.
  • All smaller models (Llama 3.2 3B, Llama 3.2 1B, Gemma 2B) and Vicuna produced many 🔥 or 🍄 outcomes, indicating they could not reliably repair the supplied mistakes.

Recap and implications

The authors expected the trivial two‑API setting to be easy for all models, but the results showed a clear split. Gemma 2 9B‑instr was the only model that navigated the full dialog with near‑perfect correctness and minimal correction prompts. Llama 3.1 8B‑instr was a close second, needing slightly more guidance. Larger size helped, but the much larger Gemini model underperformed due to formatting issues, indicating that instruction fidelity and token‑level control matter more than raw scale. Smaller models (≤2 B) and older models frequently failed to produce a valid API call at all, and when they did, they struggled to incorporate corrective feedback without extensive prompting. The experiment demonstrates that, for a simple tool‑use scenario where users can correct mistakes in plain English, current mid‑size instruction‑tuned LLMs can be useful, but reliability still depends on model scale, training data, and the ability to suppress extraneous text.

Readers can replicate or extend the study using the publicly available Keras Chatbot Arena Space, and they are encouraged to experiment with fine‑tuning to improve mistake‑fixing behavior.

Sources