iADD: Improving Alignment and Diversity in Diffusion Policy Optimization

iADD. Improving Alignment and Diversity in Diffusion Policy Optimization. By Ashok Prasad Neupane and four others, from Tribhuvan University, NAAMII, and one other institution. Published on the first of October. Its code repository is public and real: about forty source files.

Diffusion models can be post-trained with reinforcement learning. Denoising Diffusion Policy Optimization, or DDPO, is one example: it optimises the reverse diffusion process under a reward function. The difficulty is what that optimisation costs. Put simply, the reward goes up and something else comes down. Current approaches raise the reward at the expense of diversity and of quality, and the authors set out to find better trade-offs through careful theory and method design.

The work has two parts, and the first is analytical. Examining the theoretical framework, the authors demonstrate mathematically that updating a diffusion model only at the later timesteps may be harmful for diversity. They state that this runs contrary to the conclusions of a previous work. The point is a narrow one but it has consequences: if where in the denoising process the updates land affects diversity, then that choice belongs in the design of any fix.

The second part is a method: an incremental Feynman-Kac training procedure, which the authors rest on strong theoretical foundations and propose as a way to reach the best alignment and diversity trade-off yet. The abstract does not detail its steps, and neither will this summary.

Experiments compare the method against related diffusion policy optimisation approaches on three different tasks, and ablations are provided for each component. The authors report strong performance gains in alignment and in diversity together, rather than one bought with the other.

The order matters. A mathematical argument first questions an earlier conclusion about which timesteps are safe to update, and only then does a method follow, built to claim alignment and diversity at once instead of trading one away.