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

Databricks built an evaluation framework using LLM judges aligned with human experts through MemAlign to assess the quality of Genie Code-generated machine learning notebooks.

• •Nine LLM judges evaluate ML notebooks across nine dimensions including data exploration, feature engineering, model training, and metrics evaluation using a 1-3 scoring scale
• •Initial testing showed significant misalignment between LLM judges and human evaluators, with Model Training (MAE 0.680), Model Use (MAE 0.562), and Data Imputation (MAE 0.474) having the largest gaps
• •MemAlign, an open-source alignment framework in MLflow, uses semantic memory (generalized guidelines) and episodic memory (specific examples) to improve judge accuracy
• •After MemAlign training, three dimensions showed statistically significant improvement: Model Training (74% reduction to MAE 0.180), Model Use (78% reduction to MAE 0.125), and Data Imputation (89% reduction to MAE 0.053)