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

Table of contents 1. Introduction 2. The experiment 3. Keras chatbot arena tech: Spaces, Gradio, TPUs, JAX and Keras 3.1 Why TPUs? 3.2 Why JAX and Keras? 3.3 Sharding Models? 3.4 Which models? 4. Results 4.1 Reliability 4.2 The complete chat - fixing mistakes 4.3 More mistake fixing 5. Recap

I'm not interested in having LLMs solve big problems, quite the opposite. I want them to dispatch drudgery, and if they don't get it right on the first try, a short English sentence should be enough to fix it. In short, I want an assistant, like the computers in old sci-fi movies, minus the " I'm sorry Dave, I'm afraid I can't do that " bit 😅.

This paper explores such a tool for coding. Setting aside the creative title claim (No, AI is not beating Kaggle gandmasters yet), what the paper authors did was to manually break various Kaggle problems into micro-tasks, have an LLM generate code for them and iterate until unit tests pass. An example micro-task could be, for an image classification problem, to "figure out the format of the input data and reformat it into a CSV with columns 'id', 'image_filename' and 'class'".

For their fast inference and large memory. A TPU v5e 2x4 has 8 cores and 16GB of RAM per core for an aggregate 128GB of memory. With this much memory, we can load multiple LLMs at once, provided we shard them across all cores, and switch between them at will in the UI. In this experiment, I have been able to load five ∼8B params models (one more would OOM) and three ∼2B models for a total of 7 LLMs in memory at once, in bfloat16 format.

I like that approach because that's how I would like to work on my projects with AI in the future. Have AI generate the boring pieces of code, like data reformatting, so I can focus on the interesting bits: correctly framing the problem and devising the steps that will lead to a solution.

But this interactive coding assistant must be able to listen to feedback in plain English and fix mistakes in its code. With LLM's ability to infer information from knowledge and context, this could be a very efficient computer interface. But if LLM's quirks like hallucinations or lack of formal logic get in the way, we could end up with a case of "artificial stupidity" rather than AI.

So I decided to run a little test with today's LLMs. A super-simplified one, to see how effectively LLMs fix their mistakes when you point them out to them.