Machine Learning - Weight Ensembling Improves Reasoning in Language Models

Hey PaperLedge crew, Ernis here, ready to dive into some seriously fascinating research! Today, we're tackling a paper that shines a light on a tricky problem that pops up when we're training AI to think and reason like us. Think of it as teaching a kid to solve a puzzle – sometimes they get stuck in a rut, and we need to shake things up!

This paper looks at what happens when we're training these big language models to, say, write code or solve math problems. The researchers noticed something weird: As they kept training the model, it got better at getting the first answer right (they call this "Pass@1," like getting the first shot in basketball), but it got worse at coming up with a whole bunch of different, potentially correct answers (that's "Pass@k"). Imagine the kid only learning one way to solve the puzzle, even if other ways exist!

So, what's going on? Well, the researchers figured out that the model's "brain" – its internal settings – starts to become too specialized. It loses the ability to explore different possibilities. They call this a "collapse of diversity." Think of it like a musician who only knows one song – they might play it perfectly, but they can't improvise or adapt!

---

Now, here's the cool part: They found a surprisingly simple fix! It's like having the kid show their work on the puzzle, and then comparing their work with earlier attempts. The researchers took the model's current "brain" and mixed it with an earlier version of its "brain" from earlier in the training process. It's like blending the experience of a seasoned player with the fresh perspective of a rookie! They call this mixing technique "WiSE-FT."

And guess what? It worked like a charm! Mixing the "brains" almost completely fixed the problem of the model getting worse at generating diverse solutions. In fact, it even improved the model's ability to get the first answer right! It's like the musician suddenly being able to improvise and play their signature song even better!

"WiSE-FT almost completely recovers Pass@k while also improving Pass@1."

---

The researchers then went a step further. They showed that using this "brain-mixing" trick made the model better at learning from even less data when they used reinforcement learning to fine-tune it. And even better, it gave them performance gains that couldn't be achieved by simply tweaking how the model generates its answers, using things like "temperature scaling."

To understand why this works, they used some fancy math to explain that "Pass@k" involves a tradeoff between what the model expects to get right ("bias") and how much its performance varies ("variance"). They found that WiSE-FT can reduce both bias and variance simultaneously. Temperature scaling, on the other hand, is inherently a tradeoff between bias and variance.

---

Why does this matter?

• For AI researchers: This paper provides a valuable insight into a common failure mode in training reasoning models and offers a simple, effective solution.
• For developers building AI applications: This technique can help improve the reliability and robustness of AI systems, especially in tasks that require creative problem-solving.
• For anyone interested in AI: It highlights the challenges of training AI to think like humans and the importance of finding ways to encourage diversity and exploration.

Think about it this way: Imagine training a self-driving car. You want it to reliably get you from point A to point B ("Pass@1"), but you also want it to be able to handle unexpected situations and find alternative routes ("Pass@k"). This research suggests a way to train the car to do both!

So, here are a couple of things I'm pondering after reading this paper:

• Is this "collapse of diversity" a fundamental problem with how we train AI, or is it specific to certain types of models or tasks?
• Could this "brain-mixing" technique be applied to other areas of AI, like image recognition or natural language processing?

That's it for this week's deep dive! I hope you found this paper as thought-provoking as I did. Until next time, keep learning, keep exploring, and keep pushing the boundaries of what's possible!

Credit to Paper authors: Xingyu Dang, Christina Baek, Kaiyue Wen, Zico Kolter, Aditi Raghunathan