Diamond Maps: Efficient Reward Alignment for Generative Models

1. 0:00 – 0:07

   Intro

   1. AGAnkit Gupta0:00

      [upbeat music]

2. 0:07 – 0:38

   Diamond Maps overview: steering a base generative model toward a target task/style

   1. AGAnkit Gupta0:07

      Cool. Doug, tell us what you're working on.

   2. DCDouglas Chen0:09

      Yeah. So this is Diamond Maps. Uh, Diamond Maps aims to solve this class of problems called reward alignment, and what reward alignment does is it essentially takes, uh, you know, a, a model, a base model, and it tries to steer or align it to a specific task that you might have. So for example, you might have a great, you know, image generation model, which generates great images, but you're not really able to steer it or prompt it to, you know, generate images that are interesting to you. Um-

   3. AGAnkit Gupta0:35

      You might have, like, a specific style you want, let's say.

   4. DCDouglas Chen0:36

      Exactly, a specific style.

   5. AGAnkit Gupta0:38

      Yeah.

3. 0:38 – 1:49

   How reward guidance depends on a value function

   1. DCDouglas Chen0:38

      Yeah. So that's, uh, the problem that Diamond Maps is trying to solve, and it's built on top this previous work in flow matching and also Flow Maps, and the actual, uh, main advantage of Diamond Maps is that it's extremely efficient. So, uh, going into more of the details here, reward alignment relies on something called reward guidance, and to, to, to run this kind of reward guidance process, you need to estimate something called a value function. So a value function would be something that you define that kind of, uh, you know, a high va- a high value or high reward would maybe align more with a sample that you care about more. So for example, with prompt alignment, your value function could be, you know, uh, some type... a difference between, you know, the image and this prompt that you give it. So estimating this value function is kind of a core, uh, problem that's in reward alignment. Now, we have, you know, a lot of methods currently that help us estimate this value function, but the two, uh, main drawbacks are that they're either slow or they're also, uh, not great estimators with the value function because of... because they're deterministic.

   2. AGAnkit Gupta1:48

      Mm.

4. 1:49 – 2:00

   Why value estimation is hard: intermediate diffusion/flow states don’t look like real samples

   1. DCDouglas Chen1:49

      And exactly why is, uh, this value function is not defined over this generative process. So generally, when you have this generative process, you start from some noise, and you actually get to your final

5. 2:00 – 2:11

   Limitations of existing approaches: slow or overly deterministic value estimates

   1. DCDouglas Chen2:00

      distribution. Now, it's not actually defined on any of these intermediate steps because they don't actually look like a sample.

   2. AGAnkit Gupta2:06

      Yeah, when you have some, like, noisy image or whatever-

   3. DCDouglas Chen2:09

      Exactly

   4. AGAnkit Gupta2:09

      ... it doesn't really look like any of the samples you can find in the value function.

6. 2:11 – 2:56

   Core idea of Diamond Maps: stochastic sampling of multiple finals from the same intermediate state

   1. DCDouglas Chen2:11

      Exactly, yeah. So you kind of need to get to your final sample first. So that's what Flow Maps and flow matching do, but they do it deterministically. So that's a problem because your value function estimate should really be multiple samples because that point that you're kind of starting at can, can end up or, or can look like many different types of images. So to get a great estimate of this value function, you want many samples of this. So deterministic, uh, deterministic Flow Maps, let's say, at your intermediate step, you can, you know, sample this final sample that we were talking about before, and you can estimate this value function only with one sample. But Diamond Maps allows you to sample multiple final samples from this same, you know, intermediate-

   2. AGAnkit Gupta2:56

      At the intermediate step.

7. 2:56 – 3:36

   Empirical results: improved alignment metrics over vanilla flow matching

   1. DCDouglas Chen2:56

      Exactly, yeah, or given the intermediate step, it allows you to sample multiple final images, and this allows you to calculate reward over multiple different samples and get a much better estimate of this value function. Now, what does this actually look like? Well, you know, we have, we have results here that show basically this better estimate does actually, you know, tangibly improve your actual alignment performance, right? So for example, in this graph, you can see that Diamond Maps, uh, achieves better alignment. So this is, um, a metric called LPIPS, which measures alignment basically, and achieves better performance than if you were to use, you know, just determine-

   2. AGAnkit Gupta3:35

      Vanilla flow matching.

8. 3:36 – 3:55

   Qualitative demos: prompt-aligning an unconditioned ImageNet model (out-of-distribution prompts)

   1. DCDouglas Chen3:36

      Exactly, and you can also see some cool samples here of an ImageNet model, which, uh, you know, is not conditionally, uh, steered or, you know, aligned at all, and it actually is prompt align now. So you can see we generate a lighthouse in the middle of a desert, and these prompts are chosen so that... you know, these are not images that will occur naturally.

9. 3:55 – 4:14

   Two usage modes: fine-tuning vs training-free inference-time alignment

   1. DCDouglas Chen3:55

      You know, you won't have a lighthouse in the middle of a desert, and you can see that it does extremely well on these, uh, examples. And we actually propose two different approaches here. One is a training, you know, a training or a fine-tuning method where you fine-tune off a pre-trained Flow Map, and one is a training-free method, uh, that's only at inference time.

10. 4:14 – 4:33

    Plug-and-play integration: attach Diamond Maps to an existing Flow Map model

    1. AGAnkit Gupta4:14

       With an existing Flow Map-

    2. DCDouglas Chen4:16

       Exactly

    3. AGAnkit Gupta4:16

       ... trained model.

    4. DCDouglas Chen4:17

       So this basically plugs onto any existing Flow Map, and you can, you know, basically sample stochastically with, uh, with this method and not need to do any training at all. So this is kind of a trade-off between, you know, inference time compute versus training time compute. And-

11. 4:33 – 4:40

    Can you combine training + inference? Practical compute considerations

    1. AGAnkit Gupta4:33

       And presumably you could do both also, right? You could have it as part of your training process and then still use it at inference time. And does that, does that make a difference?

12. 4:40 – 5:30

    Scaling motivation: inference-time method enables experiments on larger models/datasets

    1. DCDouglas Chen4:40

       Um, you could do that. I'd assume that basically the whole point of, uh, the inference time or, or they... the results are the same. So if you did that, that would probably just, you know, you'll be spending extra compute, and you wouldn't really be gaining a lot from that. But, uh, yeah, if you already trained a model, it's probably better to use that. But for example, training these models are not that easy. So if you see here, like, our text-to-image experiments here, they're actually on much higher scale, uh, data sets, you know, we're using flux, and compared to here, which is just ImageNet, and we're able to do this because we didn't have enough compute to, you know, fine-tune these models. Takes a lot of compute. This inference time work only takes that inference time compute, and we're able to, you know, then scale up easier and sh- show, you know, better results here as well. Yeah.

13. 5:30 – 5:33

    Wrap-up and acknowledgments

    1. AGAnkit Gupta5:30

       Cool. Doug, this is really awesome work. Thanks so much for showing us.

    2. DCDouglas Chen5:32

       Thank you.

    3. AGAnkit Gupta5:32

       This was awesome. Thanks.

Episode duration: 5:35