Why Are AI Answers So Predictable? Uncovering the Math Behind Diversity Loss
When humans are asked to pick a random number between one and one hundred, they offer a wide spread of choices, but when a large language model is asked the exact same question, it will almost always spit out the identical number (27!). CDS PhD student Anthony GX-Chen and his colleagues discovered the mathematical reason behind phenomena like this. Their work shows that “diversity loss” is a direct consequence of the model’s objective function.
Generative models rely on reinforcement learning, a trial-and-error training process that maximizes rewards for correct outputs. To prevent the model from forgetting its baseline knowledge, such as how to speak English while learning to solve math problems, developers use a technique called Kullback-Leibler (KL) regularization. This keeps the new behavior close to the original base model.
In their paper, “KL-Regularized Reinforcement Learning Is Designed to Mode Collapse,” the research team investigated the optimal solution to these competing objectives. Co-authors include CDS Associate Professor of Computer Science and Data Science Rajesh Ranganath, Courant PhD student Jatin Prakash, CDS Associated Professor Rob Fergus, and Jeff Guo, then a PhD student at École Polytechnique Fédérale de Lausanne.
GX-Chen explained that researchers usually treat optimization based on intuition. They tinker with exploration tactics or gradients to fix repetitive outputs. By viewing reinforcement learning through the lens of variational inference, the team analyzed the theoretical target distribution directly.
“Assuming you do all of that perfectly, and you get to the best possible solution you can have under this objective that we’ve all written down and use, what is the actual outcome?” GX-Chen asked. “It turns out the outcome is that we essentially lose all of the diversity, and we concentrate all of the probability mass of the generations onto just a few very, very good answers.”
The findings apply to more than text generators. The study also tested chemical language models, which design molecular structures for drug discovery. Researchers need a diverse set of candidate molecules for clinical experimentation. Guo’s chemistry background helped the team realize that diversity loss in text generation and drug discovery were two sides of the same coin.
To address the issue, the team proposed a method called Mode Anchored Reward Augmentation. Instead of relying on optimization tricks, the method redefines the optimal distribution so the model learns to assign equal, high probability to all good answers.
“Diversity is also an expression of uncertainty now that these models are interfacing with this extremely complex world and the many diverse kinds of people that live in it,” GX-Chen said.
By Stephen Thomas
