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 評委與專家不匹配，值得花一個下午試試。