Ben Kunkle on Building Zed's Zeta2 Prediction Model

Ben Kunkle, Lead at Zed Industries, detailed the process of building Zeta2, an AI model designed to predict a user's next edit as they type. In his presentation, Kunkle explained the technical pipeline and data considerations involved in training such a model, emphasizing the challenges and solutions encountered in production environments.

Understanding Edit Prediction

Kunkle began by defining edit prediction as the task of providing the model with context around a user's cursor and recent edits, along with type or variable definitions and any diagnostics or errors, to predict the subsequent edit. This process must be extremely fast, operating on every keystroke with a latency budget under 300 milliseconds, necessitating a small, specialized model.

The Training Pipeline

The core of the training process involves a pipeline that ingests both 'production data' (snapshots of user activity) and 'synthetic data' (git commits). This data is fed into a 'teacher frontier' model, which generates predictions. These predictions are then evaluated, and any that fail are sent to a 'repair' stage, where a teacher model attempts to correct them. The corrected data is then fed back into the distillation process to train the student model. Kunkle highlighted that each stage in this pipeline enriches the data, converting JSONL inputs into enriched 'examples' and outputting JSONL, which is crucial for managing large datasets efficiently across experiments.

Data Considerations and 'Settled Data'

A significant challenge in training edit prediction models is the inherent noisiness of the data. Kunkle explained that to address this, they use a concept called 'settled data'. This involves waiting for the prediction region to stabilize and then capturing the final state of the code as the 'answer'. By comparing the model's predictions against this 'settled state', they can filter out noisy examples and identify high-quality training data. This method allows for training on ideal examples where matches between predictions and the final code are clear and unambiguous.

Offline Evaluation and Production Monitoring

For offline evaluation, Kunkle mentioned metrics such as 'deltaChrF' (character F-score), exact lines matched, reversal ratio, and kept rate. These metrics are used to assess the model's performance on a held-out test set. He also touched upon the importance of tracking model performance in production after deployment. This includes using structured logs for latency, kept rate, and token counts, as well as dashboards that monitor acceptance rates and A/B test results across different model versions. The goal is to continuously monitor and improve the model's effectiveness in real-world usage.