Using MemAlign to Improve Evaluation of Traditional Machine Learning in Genie Code

Databricks built LLM judges to evaluate Genie Code's generated ML notebooks across nine dimensions (library installation, EDA, data imputation, feature engineering, model training, etc.). Initial judges showed significant misalignment with human expert scores. Using MemAlign, an open-source alignment framework in MLflow, they injected dual-memory (semantic guidelines + episodic examples) into the judges using only ~50 labeled notebooks. This reduced judge error by 74–89% on the most misaligned dimensions. Key findings: episodic memory proved critical even with sparse training data, and K-fold cross-validation was used to avoid data leakage in the alignment evaluation itself.