Databricks 最近宣布了 Genie Code，这是其专为数据工作构建的自主 AI 合作伙伴。Genie Code 取代了 Databricks Assistant，同时整合了多个代理，并提供了新的集成点和能力。它与 Unity Catalog 深度集成，能够理解用户的表、列、血缘、指标视图和业务定义（语义）。这种上下文感知使 Genie Code 对数据从业者来说比通用聊天机器人有用得多。
当 Genie Code 为传统 ML 任务生成笔记本时，例如“构建一个客户流失预测模型”，我们期望它生成一个生产就绪的工作流程，包括安装适当的 Python 库、数据探索和预处理、训练、调优、注册和部署模型，以及评估其性能。我们还期望每一步都真正由数据驱动：例如，Genie Code 应理解二分类问题中的不平衡类别会导致截然不同的工作流程和成功指标。
为了确保 Genie Code 始终遵循 Databricks 原生最佳实践，并避免例如跳过交叉验证、未能注意数据泄漏或不恰当的数据插补等问题，我们需要一种严格的方法来回答一个问题：我们如何知道生成的代码实际上是好的？生成的笔记本将极大地依赖于客户试图解决的问题，而不同客户之间的问题差异很大，因此这是一个非常不平凡的问题。
在本文中，我们将介绍如何为 Genie Code 的传统 ML 能力构建评估管道，以及如何使用 MemAlign（MLflow 中一个新的开源对齐框架）来缩小我们发现 LLM 评委与人类专家之间的巨大差距。改进后的评委帮助我们识别并修复了 Genie Code 中本可能被遗漏的 ML 指导缺陷。
构建评估框架
一个稳健的评估框架是必需的，用于：
- 爬山优化：量化提示、工具、技能和架构变化对输出的影响。
- 防止回归：确保改进“模型训练”不会意外地降低“数据探索”的质量。
- 基准测试：衡量不同基础模型（LLM 后端）如何影响笔记本质量。
- 持续集成：监控底层代理循环的变化如何波及到最终的 ML 任务。
评估传统 ML 笔记本是最复杂的评估任务之一，因为它涵盖了代码质量、最佳 ML 实践以及数据驱动的适应/定制。为了处理像评估 ML 笔记本这样广泛而杂乱的任务，我们使用 LLM 作为评委——由人类教授何谓好笔记本的 LLM“专家”。我们创建了九个评委，它们被提示沿着大多数 ML 工作流中出现的九个维度评估 ML 笔记本：库安装、探索性数据分析、数据插补、特征工程、模型训练、模型服务、指标评估、MLflow 日志记录和单元格组织。
对于每个维度，我们编写了评分标准（在人类评分员和 LLM 评委之间复用），分配 1 到 3 的分数，0 表示“不适用”：
- 3（好）：笔记本在该维度达到高标准，展示最佳实践，涵盖预期范围，并适当处理边缘情况。
- 2（平均）：可接受但有差距。基础部分存在，但缺少有经验的从业者会期望的改进。
- 1（差）：基本问题。关键步骤缺失、错误或以会导致错误结论的方式应用。
- N/A（不适用）：该维度不适用于此提示（例如，如果数据集没有缺失值，则数据插补维度不适用）。
除了评委，我们还维护一套评估测试用例，涵盖各种 ML 任务（分类、回归、预测），跨越不同的数据集大小、领域和复杂度。每个测试用例包括一个用户提示，告诉 Genie Code 它应该在指定数据集上解决的 ML 任务。评估循环包括为每个测试用例使用 Genie Code 生成一个或多个笔记本，然后沿着所有适用维度对每个笔记本评分。
评估评估系统
通过使用 LLM 评委代替人类来评估 Genie Code 的产物，我们实际上将一个困难问题换成了另一个：开箱即用的评委不擅长当前任务，且与人类评分不一致。我们的问题陈述是使 LLM 评委的评分与人类评估者的评分对齐。
用于 LLM 评委评估的评估集包含 50 个 Genie Code 生成的笔记本（“测试用例”），其中人类专家对每个适用维度进行了评分，并提供了简短的理由作为我们的基础事实。在两种分数之间的灰色区域，评分者被允许表达自己的判断，但模式编写的方式使得这种情况很少发生。
人类与机器对齐的度量是每个维度上分数的平均绝对误差（MAE）。结果参差不齐，一些维度显示出强对齐（4 个维度的 MAE 小于 0.3），但其他维度则不然。最差对齐的维度是“数据插补”（MAE 0.68）、“模型服务”（0.62）和“指标评估”（0.69）。这些维度是 ML 工作流中的关键瓶颈，因此我们无法接受高误差率。
MemAlign 来拯救
MemAlign（记忆对齐）是 MLflow 2.16 及以上版本中提供的一个开源框架，用于将 LLM 评委与人类偏好对齐。它在推理时使用两种类型的记忆来增强 LLM 的上下文：
- 语义记忆：可用配方的长期记忆。这包括显式规则、事实和知识，例如“每个 ML 实验应该包括交叉验证和‘评估’指标”。
- 情景记忆：过去经验的短期记忆。这包括从误判中学习的示例，例如“上次 LLM 因忽略数据泄漏而错误评分——现在它知道要检查数据泄漏。”
要使用 MemAlign，只需将一组带标签的示例（人类评分和解释）添加到 MLflow 中的对标评估集中，并在推理时通过一个简单的 API 调用注入这些记忆。MemAlign 将使用这些记忆自动调整 LLM 评委的输出，使其与人类偏好对齐。它不需要任何微调或提示工程。
我们在 50 个可能最差的维度上应用了 MemAlign：数据插补、模型服务和指标评估。对于每个维度，我们随机选择约 10 个带标签的示例（从 50 个中抽取），并测试对齐效果。结果令人印象深刻：对于数据插补，MAE 从 0.68 下降到 0.18（下降 74%）；对于模型服务，从 0.62 下降到 0.07（下降 89%）；对于指标评估，从 0.69 下降到 0.14（下降 80%）。这个改善非常显著，使 LLM 评委的评估可以安全地用作人类评估的代理，用于检测回归和基准测试。
在随后的研究中，我们测试了每种记忆类型对对齐的影响。我们发现在这个任务上，两种记忆都很重要。语义记忆单独使用时，数据插补的 MAE 为 0.47，而情景记忆为 0.42。两者结合时，MAE 降至 0.18。对于其他维度也观察到类似的模式。
教训
- 一致评分很难：LLM 评委与人类的偏差通常是由于提示过于模糊、评分标准不清晰或上下文不足造成的。确保专家提供准确的评分和清晰的解释，以减少歧义。
- 专家之间的一致性：当人类标注者彼此不一致时（如我们在数据探索回归中看到的），对齐就没有可学习的一致信号。
MemAlign 与 MLflow 一起提供，且仅用约 50 个标注示例就能为我们工作。如果你的 LLM 评委与专家不匹配，值得花一个下午试试。