Recently announced Genie Code is Databricks’ autonomous AI partner purpose built for data work. It replaced Databricks Assistant, while it subsumed several agents and provided new integration points and capabilities. Genie Code has deep integration with Unity Catalog, meaning it understands your tables, columns, lineage, metrics views, and business definitions (semantics). This contextual awareness makes Genie Code far more useful for data practitioners than the generic chatbots.
When Genie Code generates a notebook for traditional ML tasks, such as "build a churn prediction model", we expect it to yield a production-ready workflow that includes installation of the appropriate Python libraries, exploration and preprocessing of the data, training, tuning, registration and deployment of the model, and evaluation of how well it performs. We also expect that each step is truly informed by the data: for example, Genie Code shall understand that imbalanced classes in a binary classification problem result in dramatically different workflows and success metrics.
To ensure Genie Code consistently follows Databricks-native best practices and avoids, for example, skipping of cross-validation, failing to notice data leakage or improper data imputation, we needed a rigorous way to answer one question: How do we know if the generated code is actually any good? The generated notebook will greatly depend on the problem the customer is trying to solve, and this can vary vastly among different customers, so this is a very non-trivial question.
In this post, we’ll walk through how we built an evaluation pipeline for Genie Code’s traditional ML capabilities and how we used MemAlign (a new open-source alignment framework in MLflow) to close the massive gap we found between LLM judges and human experts. The improved judges helped us identify and fix gaps in Genie Code's ML guidance that we would have otherwise missed.
A robust evaluation framework is required for:
Evaluating traditional ML notebooks is one of the most complex evaluation tasks as it spans evaluation of code quality, best ML practices, and data-informed adaptations/tailoring. To handle a task as broad and messy as evaluating ML notebooks, we use an LLM-as-a-judge - an LLM “expert” taught by humans what exactly a good notebook looks like. We created nine judges which are prompted to evaluate the ML notebooks along nine dimensions that appear in most ML workflows:
| Dimensions | What we grade |
|---|---|
| Library Installation | Proper dependencies |
| Exploratory Data Analysis | Thorough EDA and |
| Data Imputation | Mean Time to Contain |
| Handling missing values without leakage. | Feature Engineering |
| Feature selection/transformation. | Model Training |
| Model selection, Cross Validation, Hyperparameter tuning | Reusing the trained model to do inference. |
| Metrics Evaluation | Inference logic and task-appropriate metrics (e.g., MAPE for forecasting, MAE for regression, Accuracy for classification). |
| MLflow Logging | Experiment tracking setup. |
| Cell Organization | Splitting the code into cells, code cleanliness, readability, markdown headers, appropriate logging. |
For each dimension, we wrote scoring rubrics (reused between human raters and LLM judges) that assign a score from 1 to 3, and 0 for "not applicable":
To give an idea of the granularity, here is the specific rubric we use for the ”data imputation” dimension:
Along with the judges, we maintain a set of evaluation test cases that span a range of ML tasks (classification, regression, forecasting), across different dataset sizes, domains, and complexity levels. Each test case includes a user prompt which tells Genie Code the ML task it is supposed to solve on the specified dataset (“I have passenger data in the tables titanic_train_table and titanic_test_table. Can you figure out who survived?”). The evaluation loop consists of using Genie Code to generate a notebook (or multiple ones) for each test case, and then scoring every notebook along all applicable dimensions.
By using LLM judges, instead of humans, to evaluate Genie Code artifacts, we essentially swapped one difficult problem for another: the out-of-box judge is unskilled on the task at hand and misaligned with human ratings. Our problem statement is to make the LLM judges score align with those of human evaluators.
The evaluation set for LLM-judge appraisal contains 50 Genie Code generated notebooks (“test cases”) where human experts graded every applicable dimension, providing both a score and a short justification to serve as our ground truth. In the grey areas between two scores, raters were allowed to express their own judgement, but the schemas were written in such a way that this is rarely the case.
The measure of human-machine alignment is the mean absolute error (MAE) between scores in each dimension. The results were mixed, some dimensions showed strong alignment (4 dimensions had a MAE of <= 0.10), while others revealed significant disagreement:
This gap exists because humans and LLMs don’t interpret the same rubric the same way. While a human rater can spot a subtly flawed imputation strategy or a training loop that 'works' but is logically unsound, an LLM judge often misses that technical nuance. We also found the judge suffered from a classic positivity bias - it was simply too 'polite' and this got in the way of getting objective results.
It became abundantly clear that given the same rubric, LLM judges and humans would not produce the same results - a misalignment. This is exactly the scenario MemAlign was designed to fix.
MemAlign is a framework within MLflow that can, given a very small amount of human natural language feedback, perform alignment between the human raters and LLM judges. This is achieved through two types of “memories” formed from reading the human feedback:
At inference time, MemAlign constructs a working context by pulling all semantic guidelines and retrieving the most relevant episodic examples for the current input. The judge loads all of these into its context, along with the original rubric, and uses the accumulated knowledge to give a more accurate score to all future notebooks.
The key property that made MemAlign stand out is high performance using only a small number of examples. This is because MemAlign effectively distills learning from rich learning signals in natural language feedback, and incorporates them into the dual-memory system.
Here’s an example of some of the snippets of semantic memory generated for the “data imputation” dimension, filling in the gaps in the rubric we previously defined by generally providing anchor points, examples and counter-examples:
Moreover, as mentioned earlier, the semantic memory reflected in the prompt is complemented with relevant examples from the judge’s episodic memory at scoring time, thus giving the judge even more context in order to interpret the optimized instructions.
Following the ML training-testing paradigm, we applied K-fold cross-validation (K=4) on 50 test cases (notebooks) therefore avoiding data leakage and the need to label a separate test set. For each fold we did the following:
To calculate the confidence intervals without additional labeled data, we generated 100 bootstrapped samples with replacement out of the original 50. By repeating this 10,000 times and tracking MAE between human and machine scores, we calculated the confidence intervals for human-machine alignment with a 95% CI defining a statistically significant change.
The evaluation pipeline is implemented as a single MLflow snippet that orchestrates the entire process:
The MemAlign optimizer is able to align LLM judges based on the test cases’ traces in just a couple of lines of code. We used this new “aligned” judge to calculate the new MAE. Aligning a judge on a single dimension takes roughly 25 seconds per fold, so the alignment itself is not a bottleneck.
Three out of 9 dimensions showed statistically significant improvement:
These 3 dimensions are among initial 4 dimensions that were heavily misaligned. A weak initial alignment is indicative of the LLMs and humans having a fundamentally different understanding of the shared rubrics, and the memory injected from MemAlign seems to provide enough context to get them “on the same page”.
The dual-memory structure of MemAlign led us to question whether both of them are actually contributing to the judge alignment. In particular, the episodic memory is supposed to help the judge by giving a set of the most similar annotated notebooks as a reference point (utilizing the nearest neighbor search). But what if the retrieved notebooks (nearest neighbors) aren't actually similar to the current one - just the least dissimilar? Loading those into the judge's context might muddy things rather than help. The problem space we’re grading (ML notebooks) is very broad, and we initially hypothesized that a set of 50 notebooks would simply not be enough to get a sufficiently dense set of memories for the judge to recall.
Without episodic memory, the picture degrades substantially:
This was the opposite of what we expected. We initially hypothesized that our sparse annotated set would end up confusing the judge, but almost every dimension got worse without episodic memory. The one exception was Data Exploration, where dropping the episodic examples may have actually helped - without the specific notebooks our annotators disagreed on, the judge only had the distilled guidelines, and less noisy signal to work with.
The takeaway: even when your inputs are large and messy, episodic memory still improves the judge’s performance drastically. Both semantic and episodic memories are integral to the functioning of MemAlign.
Judging whether a coding agent is doing its job is hard enough, whereas evaluating an autonomous AI partner on building and executing traditional ML workflows is at another level of complexity. Due to the fast iteration on AI products, there is just not enough time to have experts monitor the agent’s “continuous integration". The only viable scalable solution are LLM judges - but we still need a jury of humans to keep the LLM judge in check.
By applying MemAlign, we cut the judge error by 74–89% on the dimensions where it mattered the most. But, as with any ML/LLM work, the result is only as good as the information you put in, so make sure the labeling is competent.
Takeaways:
MemAlign ships with MLflow and it worked for us with just ~50 labeled examples. If your LLM judges aren't matching your experts, it's worth an afternoon.
此内容由惯性聚合(RSS阅读器)自动聚合整理，仅供阅读参考。 原文来自 — 版权归原作者所有。