AI exercises: First stop - LLM as Judge

I have recently finished reading the book AI engineering by Chip Huyen. One of the readers of the book on Amazon made the analogy that this book is akin to literature review papers. In other words, it’s a book-length review of AI engineering literature. It summarizes at a very high level what generally works in practice. Unlike Chip’s previous book, Design ML Systems, which spells out issues that MLEs need to deal with practical considerations when putting an ML system in place, AI Engineering to me was a lot more theoretical, and not as practical. One of the reasons is my positionality. I engage in putting ML models in deployment on a daily basis, and not necessarily running AI applications every day. When I read AI Engineering, my brain ran wild. Many questions arose. Why eval datasets should be thought of in advance before even designing an AI system? How does an LLM judge work in practice? We know that retrieval augmented generation works, but can it be shown more concretely that with RAG, a system accuracy increases significantly? If so, by how many percentage points? What are the right number of iterations, and experimentation for such a system?

When I was trained as a social scientist, one of the tools that I used to get myself familiarized with a theoretical concept is writing into knowing. In other words, I would not even know whether I understand a theoretical concept or a hypothesis until write something down to delineate the concept, to paint its different contours, and understand its different dimensions.

As a machine learning engineer, my equivalent tool is coding. So over the weekend, I literally coded myself into knowing some of the most important concepts in the book such as LLM as judge, RAG, temperature and top_p sampling, agentic systems.

I thought to myself why not turning what Chip Huyen wrote in the book into bite-size exercises that would be beneficial for AI engineers who want to turn concepts into reusable codes. These exercises serve a few purposes: (1) AI engineers can use them in their production system, (2) interview candidates can practice coding before their job interviews, (3) Hiring managers can use them in the coding rounds of the AI engineering interviews, and (4) of course undergraduate and graduate students can translate concepts that they have learned in class/ read in research papers into something concrete.

There are a few reasons why I thought these exercises would help those four groups. First, as an engineer myself, I find translating theory into practice requires a lot of engineering imagination, which requires time, and thinking that most MLEs cannot afford. A friend of mine said that LLMs has leveled the playing field. Most people have at the most 2 to 3 years of work experience in pre-training, post-training, and building AI applications. So upskilling for MLEs is an absolute must. Yet, I can attest to the difficulty of this upskilling journey as a product MLE working day in day out trying to improve business. Not all ML problems are AI problems, and carving out time to learn new AI methods and applications is challenging. These exercises first and foremost are for my working full time engineer colleagues.

Second, for those who are trying to get a job in AI, and not already developing AI applications, it’s a tough market out there. The main reason is that if you have not developed an AI applications, you might not even be aware of many important engineering decisions. To be able to speak intelligently in an AI engineering interview without real-world experience is like setting yourself up for disappointment. These exercises are for the aspiring AI engineers candidate to go over the fundamentals, as well as the standard practices in industry to prepare for their upcoming interviews.

Third, for hiring managers who want to test candidates real AI knowledge instead of the boring leetcode questions, these exercises are for them to use (freely). The best AI engineering experience I had was at a AI startup, where they asked me to do some data engineering by interacting with their own LLM API. The exercise was quite long, almost 3 hours, but I had so much fun. I was writing prompts to understand the quirks of the LLM, trying to extract passages that I thought were relevant. By the end, I felt like I learned something from the interview. It was not like the typical leetcode interview where the answers were relatively scripted, and I had to practice those a lot. Not coming from a computer science background, I always have anxiety around certain leetcode questions, and praying every time to not have to deal with them during interviews. LLM prompting was different. It felt like that’s what I was doing on a daily basis, and what I would do on the job anyways. Who cares if as an AI engineer I couldn’t reverse a linked list. AI can solve that question much better than I do, but evaluating an AI system sounds much more relevant to the job.

And finally, students and professors in academia are always in my mind. These exercises are for you to simulate what could be done in industry.

In this post, I describe the first exercise in the AI exercise series. The first exercise is LLM as judge. For each exercise, there will be two python files: the problem file, and the answer file that has my example answers. I have posted in this AI exercises repo the first set of files. Over the weekend, I will also convert those files into jupyter notebooks for those who are more familiar with using jupyter notebooks to understand the workflow step by step.

The passage I tried to tease out the meaning in AI Engineering is this one (emphasis is mine):

The challenges of evaluating open-ended responses have led many teams to fall back on human evaluation. As AI has successfully been used to automate many challenging tasks, can AI automate evaluation as well? The approach of using AI to evaluate AI is called AI as a judge or LLM as a judge. An AI model that is used to evaluate other AI models is called an AI judge.¹⁵

Alright the idea that you can prompt one LLM to evaluate its own result or the result of another LLM is simple enough. But how does it really work in practice? How can one quantify the evaluation and the improvement?

I set up a simple exercise, using the GSM8k dataset, which is a grade school math problem that has golden answer to each problem. You can take a look at the description of this dataset on Hugging Face. One might argue that this dataset is probably already used in the pre-training phase of all the frontier models (such as the recent ones from OpenAI, Anthropic). However, the point of this exercise is not to evaluate how much of the benchmark dataset was already remembered by an LLM. The point of this exercise is to create an evaluation pipeline to demonstrate that LLM can evaluate its own answer, and show concretely what the main steps are.

Following is a diagram of the LLM judge evaluator system.

There are essentially four steps in this pipeline. First using prompt engineering, we ask an LLM to generate the result for each math problem. Then you evaluate the accuracy of this first step. Second, send all the questions to an LLM judge to evaluate each answer. You also have to evaluate the LLM judge answers against the previous answers to gauge whether the judge differs or agrees with what the previous generator provided. Third, the self reflection step adds one extra step to generate new answers to any problem that the judge has identified to be wrongly answered previously. Forth, we calculate the accuracy at the final step to evaluate whether the accuracy of the system improves after the self-reflection step. Seeing the accuracy improvement at the end is the entire reward at the end of this exercise.

In the problem file, llm_judge_problem.py, you’ll find in the repo, I have set up those three parts (llm_generator, llm_judge, and self_reflection) as three different functions that you’ll need to implement on your own.

In the llm_judge_answers.py file, you’ll find my proposed answers to the different parts.

One prompting trick here is that this system should be thought of as a continuous chat, where the LLM generator creates the first assistant’s response. Then the assistant’s response is added to the chat, and then the user can ask an LLM judge question, and then the answer keeps being appended until the end of the pipeline. To demonstrate what I described, here’s the piece of code at the end of the pipeline at the self reflection step:

chat = [
    {"role": "system", "content": math_answer_prompt},
    {"role": "user", "content": question},
    {"role": "assistant", "content": model_reasoning},
    {"role": "user", "content": "Is the provided solution correct or not? Check the reasoning and if there is any doubt then call it incorrect. "
                                "Please put your answer inside <answer></answer> tag and only answer 'correct' or 'wrong'"},
    {"role": "assistant", "content": judge_reasoning},
    {"role": "user", "content": "Okay since you think the answer is wrong can you generate a better response?"}
]

This is the messages that will be used as a parameter in an OpenAI chat completion API call at the self reflection step. Essentially, I kept appending the content to the chat history. It took me a while to figure out that LLM as judge pipeline essentially functions as a continuous chat.

Using only 150 examples, after the first step, my LLM generator provided answers with accuracy of 91.33%, LLM judge agrees 92% with the previous generator, and after the self reflection step, the accuracy increases to 93.33%. So with this simple set up, this system accuracy increases by 2 percentage point compared to only prompt engineering.

After the first round of initial success, I increased the number of examples from 150 to the 1,319, which is the full test dataset. The three numbers are 90.30% for the LLM generator, 89.92% for LLM judge, and final accuracy is 90.67%. The improvement is relatively small, only 0.37 percentage point increase. One could reasonably ask should I waste many tokens on this small increase. I would argue it is worth it if the problem requires very high recall such as in the case of medical record extraction, legal document information retrieval. And the token cost might be minimal. In my case I use gpt4.1 nano which didn’t even cost me $1 for this entire exercise.

This exercise demonstrates that one can relatively easily put an LLM as judge evaluator into the AI system that they are building if they have a golden eval dataset. If you have seen some LLM-as-judge system in deployment, open-source code, please send it my way. I would love to see how it is deployed at an industrial scale. What are the considerations in industry you need to take into account when building LLM as judge? Are you using the same LLM in the previous generation and self-reflection steps? Or are you using different LLMs for different steps? How do you reduce latency when the numbers of API calls are in the millions or billions?