How We Built Zeta2: Training an Edit Prediction Model in Production — Ben Kunkle, Zed

Talk transcript

1. 0:00

   [on hold music] I'm Ben Kunkle.

2. 0:15

   I'm the edit predictions lead at Zed. Um, we recently announced our models Zeta 2, and this is how we trained it. I'm gonna go through a lot. This is obviously a pretty short talk, so I'm gonna try and leave enough time for questions at the end.

3. 0:27

   But it's... If you're not familiar with training models, uh, it's gonna be a bit of a whirlwind tour. So if you're not familiar with edit prediction, it's essentially giving the model a region of code around the cursor, asking them to predict the next edit that you're gonna make.

4. 0:41

   We give it various data in, such as your recent edits, your cursor position, the type definitions and variable definitions or, of things around your cursor, as well as diagnostics, errors, et cetera.

5. 0:52

   It also needs to be very fast 'cause it runs on every keystroke, and so it's ideal for a small specialized model, fine-tuned. It can do this task and this task only.

6. 1:02

   Um, so that's what we've, we've done. So the pipeline, in essence, is taking these opt-in production data. Uh, this works really well 'cause it's snapshots, so all of that data that we have collected, related, like related types and definitions and et cetera, all of that gets captured, and then we're able to turn that into training data.

7. 1:23

   In order to do that, we use a process called distillation, where we take a frontier model, we give it all of that input, and we say, "What prediction would you make?"

8. 1:31

   This is a pretty difficult process as it turns out. Even though the frontier models are pretty smart, if you ask them 100,000 times, they're gonna give you 100,001 answers, right?

9. 1:41

   And so there's a bunch of problems there that we've had to like finely tune the prompt that we're giving that frontier model in order to get good things out.

10. 1:51

    One of the things we've done to try and get better, um, predictions to train off of is we run some offline or static evaluations. So, uh, we have some heuristics for, you know, is it just undoing what you just typed?

11. 2:06

    Is it ignoring that editable region boundary that we've given it, et cetera. And if it does, then we send it to another frontier model with a similar prompt like, "Hey, it failed in this way.

12. 2:18

    Can you fix it?" And so that, we call that the repair step, and then once we've repaired the bad predictions, then we can essentially turn what the teacher made into the expected output of the student model, or Zeta 2.

13. 2:36

    Um, up until that point is reusable across experiments. So this is stuff that we can cache. We can train multiple experiments on top of that by turning what the, uh, frontier model predicted into the format that we want the experiment to output.

14. 2:53

    And so that's the next piece, is this prompt formatting. This is experiment specific, i.e., are we including diagnostics this time? Are we not? How much of the edit history are we including?

15. 3:04

    Those are the kinds of experiments we're running. And so we'll turn what the teacher gave us into the prompt to, uh, distill and train our student model, and then we'll f- we'll do our, uh, final set of offline evaluations.

16. 3:21

    Um, the nice part about this whole process, we've designed it in such a way that it's all JSONL or a single line has a giant JSON object. These files get huge.

17. 3:31

    Um, but each stage just adds some more fields to it or moves some fields around, so it's a very like, uh, fluid and dynamic process. Um, we're generally doing 100,000 examples to train a model.

18. 3:46

    Like that's our peak. For the smaller experiments, we'll cut it down lower to 10 to 50K range.

19. 3:53

    Um, one interesting thing that we're trying right now is to use what we call settled data, which is the idea that eventually the user writes the answer, right? When you request a prediction as you're typing, eventually you're going to like write the code in the way that you wanted it.

20. 4:11

    And so we can wait, given that we're the editor, we can just wait until you stop editing that editable region that we gave the model, snapshot it and save it, um, and then use that to inform our training.

21. 4:25

    Uh, this is actually, this is very noisy because by waiting on the edit region to settle, you could change your mind. You could have an agent come in and rewrite it completely.

22. 4:37

    It could be completely different from what it looked like when the prediction was made. So what was maybe a reasonable prediction, it no longer looks reasonable. So we need some way to filter that out.

23. 4:49

    One way that we can do that is by having, generating 10 of the teacher predictions and seeing if any of them are close using like a Levenshtein s- distance type of thing.

24. 5:01

    See if any of those are close to the settled state, and if they are, we know a couple of things. We know it's predictable and that it's not noisy and, you know, completely different than what the input was, right?

25. 5:14

    'Cause we're giving the same input that we gave for the original prediction for this new prediction. Um, that turns out to be quite expensive. For 100,000 examples, you're then doing, you know, a million frontier model requests.

26. 5:29

    That is prohibitively expensive. Fortunately, given that we've now trained models using the original teacher predictions, uh, our student models, or Zeta 2, is approaching the teacher in terms of quality of prediction.

27. 5:43

    So instead of running the teacher, we can run our student checkpoint or something 50 times. That costs us basically nothing, and we can s- do the same process, see if any of them are close to the settled region using Levenshtein or something similar.

28. 5:59

    Um- This gives ideal training examples, right? 'Cause there's a, by looking at the range of distance to the settled state, there's a region that are super far away, we can be confident that that's just noise.

29. 6:14

    There's a con- there's a region that's super close, that's like, it's super obvious what you're gonna be doing, right? You, you typed function add a+. It's obviously b, right?

30. 6:27

    Um, but then there's this interesting section in the middle where it's almost. That's like the ideal, what we want in our training examples. For example, the stuff that's past the training data cutoff of our student model.

31. 6:42

    So new functions, et cetera, that it's never seen before that you actually wanted, and that's gonna show up in these new training examples that we can then, um, train off of.

32. 6:52

    We generally don't train off of the actual settled state, just because it's still noisy, but we can train off of, you know, what was closest to the settled state.

33. 7:03

    So to run those offline evals, we're running on a held-out test set, um, just making sure we're not training the model on the same stuff we're testing it on.

34. 7:13

    Um, delta chrF is our Levenshtein. Essentially it does a, like, N-gram comparison of various sizes of N. Um, and then we're tracking this reversal ratio, reversals being it's undoing exactly what you just typed.

35. 7:28

    Um, then we can also look at kept rate in production. Uh, we're gen- when we're evaling, we're generally running against three teacher predictions, 'cause a lot of these have no one right answer.

36. 7:40

    And so by generating three distinct answers that were all generated by a frontier model, we can be pretty sure that if it's close to one of those, it's a pretty good prediction.

37. 7:54

    So for our experiments, um, this is the training and production part of it. Uh, those evals that we have, um, don't necessarily correlate to what users actually want in their editor.

38. 8:08

    And so we have this page set up of our experiments. These are the two that are live right now. You can see over here, we've got this one being sampled at 15%, and that's gonna get the rest of production traffic.

39. 8:22

    And so we have a dashboard that I can't show you of the acceptance rate, latency, all of that kinda stuff for these experiments. But this is a page that we created so that we can, you know, once we've deployed it, set it to 15% of traffic, set it to 20, make it our, our live running model.

40. 8:40

    Um, so this v0211 Seed-Coder, this is what we released as Zeta 2, uh, last week.

41. 8:48

    Um, and so yeah, like I said, we have these dashboards for, um, the acceptance rate. Uh, we're trying new diagnostics right now, which is kept rate and diagnostic error counts, essentially comparing, for kept rate, comparing, um, what was the original text after the prediction, and then that settled state, and see how many characters between the prediction and

42. 9:14

    the settled state were kept. Um, for diagnostic error counts, it's pretty much exactly what you'd think. We snapshot how many errors there are before the prediction, how many there are after, and then we're then trying to use that to judge the quality of the model.

43. 9:27

    So that's it. I... There was a lot. Happy to answer, uh, questions. I think we have five or eight minutes left, so yeah.

44. 9:37

    Uh, you said it was very noisy to determine, like, the settled state. Are there any signals that you can share that you use, like a, like Git commit, for example, or-

45. 9:43

    Sorry, what was the algorithm?

46. 9:45

    Uh, so like you said that, uh, determining the settled state, like when the user has, is satisfied with that block-

47. 9:51

    Mm-hmm

48. 9:51

    ... of code, for example, is very noisy. Are there any particular signals that you can use already that are useful, like for example, they Git committed something?

49. 9:59

    Sure, yeah.

50. 10:00

    Do you, do you look at any of that stuff?

51. 10:00

    Um, so we don't look at the Git commit. We could, but, uh, right now we just do, like, you stop editing that area for 10 seconds.

52. 10:09

    Gotcha.

53. 10:09

    And that, that serves as a rough enough heuristic that-

54. 10:12

    Okay

55. 10:12

    ... um, so it's only in the cases where you are, like, consistently editing that location for longer, um, without pausing for 10 seconds that we wouldn't snapshot it.

56. 10:23

    Gotcha.

57. 10:24

    Um, yeah. Any other questions? All right. I guess you guys get your time back. [laughs] Thank you for coming. [audience applauding] [outro music]