Source-linked AI summary

Diversity-Rewarded CFG Distillation

Geoffrey Cideron, Andrea Agostinelli, Johan Ferret, Sertan Girgin, Romuald Elie, Olivier Bachem, Sarah Perrin, Alexandre Ramé

arXiv:2410.06084v1cs.LG

TL;DR

Creative generation needs outputs that fulfill prompts while remaining surprising, but quality and diversity can be difficult to balance. This paper distills CFG into model weights, adds an RL diversity reward, and interpolates between model weights to control the trade-off. In text-to-music experiments, human evaluations favor the finetuned-then-merged model, and interpolation at λ=0.5 outperforms intermediate-diversity finetunes.

  • Problem

    Directly optimizing diversity across one model’s generations is under-explored, despite the competing goals of quality and diversity in creative generation.

    Source: §1 · pp. 1–2
  • Method

    The approach distills CFG-augmented predictions into model weights, adds an RL reward for diverse generations, and interpolates between quality- and diversity-focused weights.

    Source: §1 · pp. 2–3
  • Results

    Human evaluations find more diverse music while maintaining high quality than the CFG-augmented previous state-of-the-art; λ=0.5 interpolation beats intermediate-diversity finetunes.

    Source: §1 · p. 3
  • Takeaways & Limitations

    Weight interpolation provides deployment-time control of the quality-diversity trade-off without inference overhead.

    Source: §1 · p. 3

    Embedding-based diversity can be hacked: diversity-maximizing models sometimes generate excessive background drums, likely due to embedding-model bias.

    Source: §5 · p. 12

Abstract

from arXiv · show

Generative models are transforming creative domains such as music generation, with inference-time strategies like Classifier-Free Guidance (CFG) playing a crucial role. However, CFG doubles inference cost while limiting originality and diversity across generated contents. In this paper, we introduce diversity-rewarded CFG distillation, a novel finetuning procedure that distills the strengths of CFG while addressing its limitations. Our approach optimises two training objectives: (1) a distillation objective, encouraging the model alone (without CFG) to imitate the CFG-augmented predictions, and (2) an RL objective with a diversity reward, promoting the generation of diverse outputs for a given prompt. By finetuning, we learn model weights with the ability to generate high-quality and diverse outputs, without any inference overhead. This also unlocks the potential of weight-based model merging strategies: by interpolating between the weights of two models (the first focusing on quality, the second on diversity), we can control the quality-diversity trade-off at deployment time, and even further boost performance. We conduct extensive experiments on the MusicLM (Agostinelli et al., 2023) text-to-music generative model, where our approach surpasses CFG in terms of quality-diversity Pareto optimality. According to human evaluators, our finetuned-then-merged model generates samples with higher quality-diversity than the base model augmented with CFG. Explore our generations at https://google-research.github.io/seanet/musiclm/diverse_music/.

1. Introduction

The introduction presents quality and diversity as competing goals in creative generation and investigates increasing diversity in a single model. It combines CFG distillation and a diversity reward, then uses weight interpolation to control their trade-off at deployment.

1. Introduction

  • The introduction identifies direct optimization of diversity across one model’s generations as under-explored amid the trade-off between quality and diversity.It frames diverse, unexpected outputs as important to creativity.Source: §1 · pp. 1–2
  • The method distills CFG-augmented predictions into model weights and adds an RL objective that rewards diversity across generations for a prompt.The diversity reward compares generation embeddings using negative cosine similarity.Source: §1 · p. 2
  • Training separate quality-focused and diversity-focused models makes their quality-diversity trade-off adjustable at deployment through linear weight interpolation.The models use low and high diversity-reward coefficients, respectively.Source: §1 · p. 3
  • Interpolating quality- and diversity-focused weights yields a steerable quality-diversity front, with λ=0.5 outperforming models finetuned at intermediate diversity settings.Figure 1 compares LERP, CFG guidance values, and training trajectories across diversity coefficients.Source: Figure 1 · p. 2, §1 · p. 3
  • Human evaluations find the MusicLM models generate more diverse music while maintaining high quality than the CFG-augmented previous state-of-the-art.The finetuned-then-merged model performs best according to human evaluations.Source: §1 · p. 3

2. Diversity-rewarded CFG distillation

The section distills CFG’s quality benefits into model weights and adds reinforcement learning with a diversity reward. It then uses weight interpolation to steer the quality-diversity trade-off at deployment.

2.1. CFG distillation for quality

  • CFG distillation trains the student to match CFG-augmented predictions from the same initialized model, rather than a larger teacher.Completions are sampled online from the student, a strategy described as reducing train-test mismatch.Source: §2.1 · p. 4

2.2. RL for diversity

  • CFG guidance and alignment can reduce output diversity, motivating a strategy that encourages diverse samples for a given prompt.The section reports that diversity from the CFG-distilled student steadily decreases as training progresses.Source: §2.2 · p. 4
  • The diversity reward is one minus the cosine similarity between embeddings of two generations, and RL increases the likelihood of generations with higher diversity rewards.The approach uses multiple generations for the diversity objective; for text-to-music, its embedding model is trained contrastively on audio segments from the same or different songs.Source: §2.2 · pp. 4–5

2.3. Diversity-rewarded CFG distillation

  • The combined objective adds a scaled RL diversity term to CFG distillation, with larger β prioritizing diversity at a slight quality cost.β scales the diversity term relative to the KL quality-distillation term.Source: §2.3 · p. 5

2.4. Model merging for Pareto-optimal quality-diversity

  • Linear interpolation between quality-focused and diversity-focused models finetuned from a shared initialization provides a steerable quality-diversity trade-off at deployment.Interpolation coefficients near 0 favor the quality-focused model, while values near 1 favor the diversity-focused model; the approach requires only two finetunings.Source: §2.4 · p. 5

3. Experiments

Experiments on text-to-music test CFG quality distillation, diversity-rewarded RL, and weight merging as ways to improve quality-diversity trade-offs. Human evaluation finds that the merged model offers higher diversity than CFG while maintaining comparable quality.

3. Experiments

  • The experiments ask whether CFG quality can be distilled, diversity traded off through RL, and a steerable quality-diversity front formed by weight interpolation.Source: §3 · p. 6

3.1. Text-to-music generation: task, setup and metrics

  • MusicRL-R text-to-music experiments use mixed-source prompts and online student sampling with REINFORCE; quality combines MuLan and User Preference scores, unused during training.The prompt sources include synthetic prompts, user interactions, and MusicCaps.Source: §3.1 · p. 6

3.2. CFG distillation for quality

  • With β=0, CFG distillation lowers KL and improves quality to match CFG, but reduces diversity to CFG's level after 1,000 training steps.The reported diversity at that point is around 0.37.Source: §3.2 · pp. 6–7, Figure 2 · p. 7

3.3. RL for diversity

  • Diversity rewards with β=10 or 15 raise diversity toward the empirical upper bound, but larger β increases KL and lowers quality.The reward increases diversity across generations during training.Source: §3.3 · p. 7, Figure 2 · p. 7

3.4. Model merging for Pareto-optimal quality-diversity

  • Two RL finetunings plus negligible-cost merging form a strong front; with one finetuning, interpolation toward the base beats CFG, while four-model uniform averaging performs better.The four-model average is reported as above the other fronts.Source: §3.4 · p. 8
  • Interpolating weights yields steerable quality-diversity fronts that surpass individual intermediate-β models; for β=0/15, quality exceeds the linear expectation while diversity stays close to it.The β=0/15 comparison also reports that CFG performs worse than merged models.Source: §3.4 · p. 7, Figure 4 · p. 8, Figure 3 · p. 8

3.5. Human evaluation

  • LERP(0, 15) delivers higher diversity than CFG while maintaining comparable quality; human raters preferred it for diversity in 57% and quality in 51% of comparisons.The evaluations used side-by-side comparisons of model outputs.Source: Figure 5 · p. 9, §3.5 · pp. 9–10

4. Related work

Related work covers CFG and distillation, quality-diversity methods, model merging, and music-generation systems. The paper distinguishes its approach through on-policy CFG distillation, diversity-rewarded RL, and weight merging for music generation.

Distillation

  • Unlike Meng et al. (2023)'s two-stage offline diffusion distillation, this work uses single-stage on-policy CFG distillation for an LLM and adds diversity-promoting RL and model merging.Source: §4 · p. 10

Quality-diversity in LLMs

  • Prior RLHF work reports diversity reductions, while earlier single-agent methods used distinct n-grams and mutual-information objectives; this work introduces RL to optimize diversity.The cited quality-diversity algorithms instead seek populations of agents with diverse behaviors.Source: §4 · pp. 10–11

Model merging for Pareto-optimality

  • Model merging has been used to combine strengths and interpolate policies finetuned with different rewards, but its empirical success remains incompletely understood.Source: §4 · p. 11

Music generation

  • For music generation, the work uses the Transformer-based approach of Agostinelli et al. (2023) and applies CFG distillation, RL diversity rewards, and model merging.Source: §4 · p. 11

5. Discussions and limitations

The discussion connects CFG distillation to amplification-and-distillation approaches, considers extensions beyond text-to-music, and identifies limitations in measuring diversity.

5. Discussions and limitations

  • The approach echoes amplification and distillation: a model learns from predictions generated by an augmented version of itself, removing inference overhead.The discussion relates this pattern to AlphaGo distilling MCTS and language models distilling Chain-of-Thought predictions.Source: §5 · p. 11
  • The authors propose adapting distillation, diversity-rewarded RL, and model merging to other architectures and text, image, or video generation.They note that diversity in those domains could support fairer representations, pluralistic opinions, diverse reasoning paths, or multi-sample decoding.Source: §5 · p. 11
  • The diversity reward can be hacked, with excessive background drums offered as an example likely tied to embedding-model bias.The discussion also notes that one diversity metric may miss distinct forms of variation, such as changes in voice, key, or tempo.Source: §5 · p. 12

6. Conclusion

The conclusion presents diversity-rewarded CFG distillation as a finetuning approach for improving the quality-diversity trade-off, combining online CFG distillation, diversity-rewarded RL, and model merging. In text-to-music experiments, the finetuned-then-merged model performed best in human evaluations.

6. Conclusion

  • The finetuned-then-merged model performed best in human evaluations in the text-to-music experiments.The conclusion describes the approach as a promising direction for tasks where alignment and creativity are key.Source: §6 · p. 12
  • Online distillation of CFG removes its computational overhead at deployment.The conclusion says the model distills CFG before inference.Source: §6 · p. 12
  • An RL procedure optimises a diversity reward based on similarity embeddings to encourage diversity across generations.Source: §6 · p. 12
  • Model merging enables dynamic control over the quality-diversity trade-off at deployment time.Source: §6 · p. 12

A. Policy gradient for the diversity reward

The appendix formulates a symmetric reward over multiple generations and derives a policy-gradient estimator for optimizing it. The resulting estimator can be used with standard policy-gradient algorithms, including REINFORCE.

A. Policy gradient for the diversity reward

B. CFG for text-to-music generation

On MusicRL-R, CFG’s adherence setting traces a quality-diversity front, with higher quality accompanied by lower diversity. Negative prompting produces a higher front than unconditional prompting.

B. CFG for text-to-music generation

C.1. Embedding model for diversity

The embedding model operationalizes musical diversity as whether two excerpts come from the same recording. It learns an audio embedding where same-recording segments are close, and uses one minus cosine similarity as diversity.

C.1. Embedding model for diversity

  • Diversity is framed as whether two audio excerpts derive from the same recording, encompassing musical differences such as instrumentation, melody, and tempo.The paper uses this framing to estimate dissimilarity between music segments.Source: §C.1 · p. 20
  • A self-supervised contrastive model maps non-overlapping chunks from the same audio segment nearby and chunks from different recordings farther apart.The diversity measure is one minus cosine similarity in this embedding space.Source: §C.1 · p. 20
  • The embedding model processes four-second, 16 kHz clips into mel-spectrograms and uses a 12-layer ViT with 25 million parameters.Its time-averaged embeddings are projected to 192 dimensions; training uses semi-hard triplet loss and audio augmentations.Source: §C.1 · pp. 20–21
  • RL diversity reward at β=15 increases embedding-model and human-rated diversity, but not entropy; inference-time CFG reduces all three measures.Figure 7 compares the embedding measure, human Elo rankings, and entropy per token.Source: Figure 7 · p. 21

C.2. Comparison of our diversity metric with human evaluation and entropy

The paper compares its diversity metric with human diversity Elo scores and entropy per generated token. The embedding metric correlates more with human ratings than entropy, although it remains imperfect and can be hacked.

C.2. Comparison of our diversity metric with human evaluation and entropy

  • The embedding diversity measure correlates more with human raters’ notion of diversity than entropy per generated token.At β=15, Elo and the embedding measure score high, while entropy is as low as for CFG.Source: §C.2 · p. 21
  • The embedding measure remains imperfect and can be hacked.The passage refers to further discussion in Section 5.Source: §C.2 · p. 21
Loading 2410.06084v1…