Official Review · ICLR 2025

This paper presents a new fine-tuning method for discrete diffusion models. The objective is to generate high-reward yet natural sequences of DNA or RNA, where naturalness is measured by the KL divergence with the prior distribution (i.e., the pretrained model). This problem is mainly formulated with KL-constrained reinforcement learning, so they use Gumbel-softmax tricks to enable backpropagation over rewards. The performance is compared with existing RL and classifier guidance methods in DNA and RNA optimization. 1. Using Gumbel-softmax in non-differentiable reinforcement learning in training discrete diffusion model is a novel approach. 1. This work should establish stronger connections with probabilistic inference methods that aim to sample from unnormalized densities using objectives with KL-constrained reinforcement learning. This includes hierarchical variational inference and GFlowNets, as mentioned in recent work [1] on fine-tuning diffusion models, including the discrete case. 2. Is the KL constraint for diffusion models tractable to compute? Based on theoretical results in [2], the KL divergence is intractable due to the compositionality of the diffusion process, so we can only obtain an upper bound for the KL. Please make correction that if my understanding is wrong. 3. Why should this method be better than existing RL methods? Variational inference, REINFORCE, and Gumbel-softmax are rival approaches in many applications that enable training models in non-differentiable settings. Why does the Gumbel trick have to be the best among them? [1] Venkatraman et al. "Amortizing Intractable Inference in Diffusion Models for Vision, Language, and Control," NeurIPS 2024 [2] Fan et al. "Reinforcement Learning for Fine-Tuning Text-to-Image Diffusion Models," NeurIPS 2023. See weaknesses above.