How AI collaboration during training can improve reasoning

Researchers have built frameworks that allow large language models (LLMs) to collaborate when they are prompted with a question. The idea is that two models working together are better than one, and it can improve the ability of models to reason about difficult questions. But there are drawbacks to this method. Running more than one model at inference time results in added costs. And while multiple models may be able to answer questions together that they wouldn’t be able to answer alone, the models themselves don’t learn anything new in the process.

A new study from MBZUAI researchers proposes a different approach. Instead of collaborating during inference, they have two models work together during training.

The framework, called collaborative reasoning via cross teaching, or CoRe, suggests that training two small language models in a collaborative setting improves both their collective and individual reasoning abilities.

“It’s like having two people sitting in a room solving problems together,” says Kshitij Mishra, a postdoctoral associate in Machine Learning at MBZUAI and lead author of the study about CoRe.

“One person might be able to answer some of the questions, while the other person might be able to answer different questions. If they transfer their knowledge during training and are taken out of the room, they should perform better than they would have if they never worked together.”

A study about CoRe was recently presented at the 43rd International Conference on Machine Learning (ICML) in Seoul. Mirat Aubakirov, Martin Takáč, Nils Lukas, and Salem Lahlou are co-authors of the study.

When collaboration succeeds – and fails

One limitation of inference-time collaboration is due to a phenomenon known as correlated failure. If a model misses a key insight on a problem, sampling more answers from the model, or even adding more instances of the model, will likely produce similar mistakes. It’s like a committee of people who all share the same blind spots agreeing that they are making the best decision when in fact they are all wrong in the same way.

The researchers present a theorem in the study that describes correlated failure. They explain that in a two-model setup of a stronger and weaker model, the improvement that the two-model system will have over the stronger model working alone is the probability that the weaker model is right where the stronger model is wrong. Collaboration only pays off if the models’ errors don’t overlap.

For this reason, the researchers pursue a framework that emphasizes “team success,” meaning that they try to make the models accurate in different ways, encouraging a kind of diversity in reasoning.

How CoRe works

CoRe’s training process resembles a classroom exercise that follows two rounds. In the first, both models attempt to solve a problem independently and receive no help or hints. If one of the models succeeds, the solution becomes training material for the other model in the second round. The successful reasoning is extracted and inserted into a prompt for the model that got it wrong and it tries to answer the question again.

The researchers implement two conditions that keep the learning process honest. First, the final answer isn’t included in the hint; the model that got it wrong only sees how the other model arrived at the correct answer. It’s like seeing a classmate’s work but not the solution. Second, the hint is withheld some of the time, a technique the researchers call “hint dropout.” This prevents the models from becoming overly dependent on support from their peers, since during inference, each model answers on its own and can’t benefit from receiving hints.

The researchers use a reinforcement learning framework to reward learning. Models earn credit for correct answers, a “rescue bonus” when a model that fails in the first round answers correctly in the second round, and an exploration reward for producing reasoning that differs from what’s already been tried, the goal being to avoid correlated failure.

“If some tokens lead to a correct outcome, those tokens are maximized,” Mishra says. “If one of the reasoning strategies is correct, we focus on that reasoning strategy instead of the other one.”

Results and the future of collaboration

The researchers used two small language model pairs that combined had no more than approximately 7 billion parameters. They tested two pairs of models: a Qwen instruct pair (Qwen2.5-3B-Instruct and Qwen3-4B-Instruct) and a compact reasoning pair (Phi-4-mini-reasoning and Ministral-3-3B-reasoning). They tested the pairs on four reasoning benchmarks that cover math, problem solving, and scientific question-answering. They split the benchmarks into training and test sets and used no more than 1,000 training examples for any one benchmark.

The Qwen instruct models went from solving roughly 60% of math questions from the GSM8K dataset as an untrained team to solving 99.54% after training with CoRe. The pair of reasoning models reached 92% accuracy on the MATH benchmark, nearly 80% accuracy on the AIME problem-solving benchmark, and more than 77% accuracy on GPQA, a dataset of graduate-level science questions. These performances beat single-model training and a baseline in which two independently trained models were combined at test time.

Perhaps even more interesting is that when the models were tested alone after training, they performed better than they did before training, illustrating that the models learned through the process. For example, Phi-4-mini-reasoning’s solo accuracy on MATH increased from 57% to 85% after training with CoRe.

Beyond the results, Salem Lahlou, assistant professor of Machine Learning at MBZUAI and co-author of the study about CoRe, says that this work has been informative in a couple ways. First, he says, “CoRe taught the benefits of baking diversity into the training objective.” The second insight is more conceptual but may have broader implications for the future of collaborative frameworks. Researchers have previously explored test-time collaboration, and CoRe identified benefits of training-time collaboration. “Can we start to develop a theory about when it’s better to use collaboration at training time versus doing it at test time?” he asks. This remains, however, an open question.