Using MemAlign to Improve Evaluation of Traditional Machine Learning in Genie Code | Databricks Blog
Skip to main content
Genie Code generates full ML notebooks from natural language prompts — we built nine LLM judges to evaluate their quality across dimensions like model training, data imputation, and feature engineering.
Human annotation revealed that judges disagreed with experts by up to 0.68 MAE on a 3-point scale. MemAlign, an open-source alignment framework in MLflow, closed this gap using just ~50 labeled examples.
On the three worst-aligned dimensions, MemAlign reduced judge error by 74-89%, and a followup study showed that both semantic and episodic memory are essential to the result.
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.
Building the Evaluation Framework
A robust evaluation framework is required for:
Hillclimbing: quantify how prompts, tools, skills and architecture changes affect output.
Guarding against regressions: Ensure that improving "Model Training" doesn't accidentally degrade "Data Exploration."
Benchmarking: Measure how different foundation models (LLM backends) impact notebook quality.
CI: Monitor how changes in the underlying agentic loop ripple through to the final ML tasks.
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:
DimensionsWhat we grade
Library InstallationProper dependencies
Exploratory Data AnalysisThorough EDA and
Data ImputationMean Time to Contain
Handling missing values without leakage.Feature Engineering
Feature selection/transformation.Model Training
Model selection, Cross Validation, Hyperparameter tuningReusing the trained model to do inference.
Metrics EvaluationInference logic and task-appropriate metrics (e.g., MAPE for forecasting, MAE for regression, Accuracy for classification).
MLflow LoggingExperiment tracking setup.
Cell OrganizationSplitting 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":
3 (Good): The notebook meets a high bar for a dimension. It demonstrates best practices, covers the expected scope, and handles edge cases appropriately.
2 (Average): Acceptable but with gaps. The basics are present, but the notebook misses refinements that an experienced practitioner would expect.
1 (Bad): Fundamental problems. Key steps are missing, incorrect, or applied in a way that would lead to wrong conclusions.
N/A (Not Applicable): This dimension isn’t applicable for this prompt (e.g. the dimension data imputation can’t be applied if the data set isn’t missing any values).
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.
Evaluating the evaluation system
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 quantity: When human annotators disagree with each other (as we saw in our Data Exploration regression), alignment has no coherent signal to learn from.
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.
Get the latest posts in your inbox
Subscribe to our blog and get the latest posts delivered to your inbox.
Sign up
View all blogs