Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Target Concrete Score Matching: A Holistic Framework for Discrete Diffusion

Ruixiang Zhang, Shuangfei Zhai, Yizhe Zhang, James Thornton, Zijing Ou, Joshua Susskind, Navdeep Jaitly

TL;DR

Target Concrete Score Matching (TCSM) introduces a unified, data-space-centered objective for discrete diffusion models, enabling both pre-training from data or parametric targets and versatile post-training via rewards, preferences, and distillation. By leveraging the concrete score in the clean data space and flexible neighborhood definitions, TCSM unifies prior discrete diffusion methods under a common framework and provides practical training objectives with Monte Carlo estimators. The authors demonstrate competitive language modeling results and improved sample efficiency across pre- and post-training scenarios, including reward-guided tuning, preference-based fine-tuning, and AR-to-diffusion distillation. Overall, TCSM offers a versatile, theory-grounded approach to improving discrete diffusion models for real-world tasks.

Abstract

Discrete diffusion is a promising framework for modeling and generating discrete data. In this work, we present Target Concrete Score Matching (TCSM), a novel and versatile objective for training and fine-tuning discrete diffusion models. TCSM provides a general framework with broad applicability. It supports pre-training discrete diffusion models directly from data samples, and many existing discrete diffusion approaches naturally emerge as special cases of our more general TCSM framework. Furthermore, the same TCSM objective extends to post-training of discrete diffusion models, including fine-tuning using reward functions or preference data, and distillation of knowledge from pre-trained autoregressive models. These new capabilities stem from the core idea of TCSM, estimating the concrete score of the target distribution, which resides in the original (clean) data space. This allows seamless integration with reward functions and pre-trained models, which inherently only operate in the clean data space rather than the noisy intermediate spaces of diffusion processes. Our experiments on language modeling tasks demonstrate that TCSM matches or surpasses current methods. Additionally, TCSM is versatile, applicable to both pre-training and post-training scenarios, offering greater flexibility and sample efficiency.

Target Concrete Score Matching: A Holistic Framework for Discrete Diffusion

TL;DR

Target Concrete Score Matching (TCSM) introduces a unified, data-space-centered objective for discrete diffusion models, enabling both pre-training from data or parametric targets and versatile post-training via rewards, preferences, and distillation. By leveraging the concrete score in the clean data space and flexible neighborhood definitions, TCSM unifies prior discrete diffusion methods under a common framework and provides practical training objectives with Monte Carlo estimators. The authors demonstrate competitive language modeling results and improved sample efficiency across pre- and post-training scenarios, including reward-guided tuning, preference-based fine-tuning, and AR-to-diffusion distillation. Overall, TCSM offers a versatile, theory-grounded approach to improving discrete diffusion models for real-world tasks.

Abstract

Discrete diffusion is a promising framework for modeling and generating discrete data. In this work, we present Target Concrete Score Matching (TCSM), a novel and versatile objective for training and fine-tuning discrete diffusion models. TCSM provides a general framework with broad applicability. It supports pre-training discrete diffusion models directly from data samples, and many existing discrete diffusion approaches naturally emerge as special cases of our more general TCSM framework. Furthermore, the same TCSM objective extends to post-training of discrete diffusion models, including fine-tuning using reward functions or preference data, and distillation of knowledge from pre-trained autoregressive models. These new capabilities stem from the core idea of TCSM, estimating the concrete score of the target distribution, which resides in the original (clean) data space. This allows seamless integration with reward functions and pre-trained models, which inherently only operate in the clean data space rather than the noisy intermediate spaces of diffusion processes. Our experiments on language modeling tasks demonstrate that TCSM matches or surpasses current methods. Additionally, TCSM is versatile, applicable to both pre-training and post-training scenarios, offering greater flexibility and sample efficiency.

Paper Structure

This paper contains 43 sections, 10 theorems, 68 equations, 7 figures, 9 tables, 7 algorithms.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. Target Concrete Score Matching
  4. Pre-training with TCSM
  5. TCSM with Data Samples ${\mathbf{x}}_1 \sim p_1$
  6. TCSM with Parametric Model $p_1$
  7. Post-training with TCSM
  8. TCSM Fine-tuning with a Parametric Model $p_{1|t}$
  9. TCSM Fine-tuning with Reward Optimization
  10. Direct Preference Fine-tuning
  11. AR $\to$ Diffusion distillation
  12. Conclusion
  13. Appendix
  14. Extended Preliminaries
  15. Proofs
  16. ...and 28 more sections

Key Result

Proposition 1

Let $\mathcal{N}$ define a neighborhood structure that induces a weakly connected graph $G$ over the support of $p_{1|t}(\cdot|{\mathbf{x}}_t)$. Assuming mild regularity conditions on the divergence measure $\mathcal{D}$, the global minimum of the TCSM objective $\mathcal{L}_{\textsf{TCSM}}$ in eq:t

Figures (7)

  • Figure 1: Comparison of perplexity on the OpenWebText validation set after training for 26B tokens: TCSM vs. baseline models.
  • Figure 2: TCSM Reward vs. Entropy in IMDB sentiment fine-tuning (\ref{['sec:post_training_dpo']}).
  • Figure 3: TCSM toxicity vs. generative perplexity in \ref{['sec:post_training_reward']}.
  • Figure 4: Bits Per Character (BPC) on Text8 test set. CD=Continuous Diffusion, DD=Discrete Diffusion, AR=Autoregressive, AO=Any-Order.
  • Figure 4: TCSM Reward vs. Entropy in IMDB sentiment fine-tuning (\ref{['sec:post_training_dpo']}).
  • ...and 2 more figures

Theorems & Definitions (22)

  • Definition 3.1: Concrete Score meng2022concrete
  • Proposition 1
  • proof
  • Definition 3.2: $k$-Hamming Neighborhood
  • Proposition 2
  • proof
  • Proposition 3
  • proof
  • Proposition 4
  • proof
  • ...and 12 more