Insider Brief
- Salesforce AI Research has developed LaTRO, a framework that enables large language models (LLMs) to enhance their reasoning abilities by independently evaluating reasoning paths.
- LaTRO’s “self-rewarding” method allows LLMs to reinforce effective strategies without external feedback, making the models more efficient and less dependent on manually labeled data.
- Tests showed significant accuracy improvements, suggesting LaTRO’s potential for broader applications in complex problem-solving areas like math and logic.
A Salesforce AI Research team has introduced a framework that could transform how large language models (LLMs) handle complex reasoning tasks. Calling it LaTent Reasoning Optimization, or LaTRO, the team says the approach allows models to “self-reward”– a process enabling them to independently refine their reasoning steps without needing human feedback.
The study, published on ArXiv, demonstrates how LaTRO significantly boosts reasoning accuracy, with potential applications in diverse fields such as advanced mathematics, logical reasoning, and scientific problem-solving.
Enhanced Reasoning Through Self-Rewarding Mechanism
LaTRO tackles a core challenge in AI development: improving a model’s reasoning skills during the training phase. For background, despite remarkable progress in AI, LLMs still struggle with tasks requiring multi-step problem-solving. Traditionally, prompt-based methods like Chain-of-Thought (CoT) prompting have provided models with step-by-step guidance during inference to improve reasoning. However, LaTRO offers a breakthrough by optimizing the reasoning process within the model’s training itself, enhancing both its ability to generate logical steps and evaluate their quality in real-time.
The research team’s approach modeled reasoning as a sampling process from a “latent distribution.” By treating the generation of reasoning paths as a variable within the model, LaTRO encourages LLMs to select high-quality reasoning steps autonomously. This self-rewarding framework allows the models to reinforce effective reasoning paths without relying on external reward models or task-specific examples, reducing dependency on manual training data.
It may not be exactly correct, but the process is a little like a student who learns to solve problems without a teacher constantly guiding or grading them. Instead of relying on external feedback, the student learns by trying multiple methods to answer each question. Each time they try, they keep track of which methods worked best and gradually become better at selecting the most effective strategies on their own. In other words, the student doesn’t need someone else to point out mistakes; they can learn independently by reinforcing the paths that lead to correct answers.
Substantial Performance Gains on Benchmark Tests
The researchers validated LaTRO using popular reasoning benchmarks, including the GSM8K math problem dataset and the ARC-Challenge dataset for logical reasoning. In head-to-head comparisons with baseline models and standard fine-tuning, LaTRO showed substantial improvements. On GSM8K, LaTRO’s models achieved an average 12.5% increase in accuracy over base models, with the highest improvement reaching 19.5% for the Mistral-7B model. Similarly, LaTRO demonstrated a 9.6% advantage over traditional supervised fine-tuning approaches. These gains suggest that pre-trained LLMs possess latent reasoning capabilities that can be unlocked and enhanced through LaTRO’s optimization.
LaTRO’s ability to improve reasoning across different LLM architectures — such as Salesforce’s Phi-3.5-mini, Mistral-7B, and Llama-3.1-8B — further reinforces its versatility. By enhancing reasoning accuracy without altering base model configurations, LaTRO shows promise as a flexible tool adaptable to various model types and task domains.
Implications for Resource Efficiency in AI Models
A significant advantage of LaTRO is its potential to streamline resource usage. Inference-time techniques — which are are methods applied when an AI model is actively generating answers or solving problems, rather than during its initial training — can improve reasoning, but often demand high processing power, as they require multiple reasoning paths during each inference. By optimizing reasoning during training, LaTRO shifts this computational burden to the model’s development phase, potentially reducing energy and cost requirements for practical applications. This efficiency could make advanced AI solutions more accessible and affordable, encouraging their use in areas where AI can solve intricate problems—such as drug discovery, data analysis, and automated decision-making.
Salesforce’s researchers emphasize that LaTRO’s self-rewarding mechanism creates a new baseline for efficient LLM reasoning.
The researchers write: “Our findings suggest that pre-trained LLMs are not only capable reasoners but also possess the potential to act as explicit reward models for evaluating reasoning paths.”
A Technical Overview of LaTRO
We’ll go a little bit into the weeds here. LaTRO’s framework is built on variational principles, with the model sampling reasoning paths as hidden “latent variables.” During training, the model generates multiple reasoning paths for a single query, evaluates their likelihood of yielding the correct answer, and updates its parameters to favor high-quality reasoning paths. This iterative process enables the model to improve not only its reasoning skills but also its ability to assess the quality of its reasoning paths. The model learns to prioritize reasoning strategies that maximize accuracy while minimizing unnecessary computational steps.
The Salesforce team’s experiments confirmed that LaTRO’s approach is particularly effective for complex tasks. For instance, in one GSM8K test, the model accurately calculated the amount of fabric required to create a garment by following step-by-step logic, demonstrating reasoning previously achievable only with supervised examples or human prompts.
Challenges and Future Directions
The researchers explain some of LaTRO limitations and detail ways researcher might approach improvements. For example, training with multiple reasoning paths is computationally intensive, especially for large-scale models. To mitigate this, the researchers suggest future work on more efficient sampling techniques or adaptive rationale generation.
Expanding LaTRO’s capabilities to broader reasoning domains also remains a priority, according to the researchers. The team plans to explore how LaTRO can be applied beyond math and science tasks, potentially using it in areas like legal reasoning, ethical decision-making, and natural language understanding.
Another direction for future work is to refine the latent variable sampling method to enable multi-step learning. While LaTRO significantly improves single-step reasoning, multi-step processes require more sophisticated methods. Researchers are optimistic that LaTRO’s core principles can be extended to accommodate such challenges, potentially transforming how AI models handle deeply layered reasoning tasks.
As AI technologies progress, the team suggests that methods like LaTRO will likely play an essential role in building models that are not only powerful but also efficient, versatile, and capable of operating with minimal external oversight. This advance brings AI closer to solving complex, real-world problems, from optimizing logistics to advancing scientific research, where dependable reasoning skills are crucial.
The research team includes: Haolin Chen, Yihao Feng, Zuxin Liu, Weiran Yao, Akshara Prabhakar, Shelby Heinecke, Ricky Ho, Phil Mui, Silvio Savarese, Caiming Xiong, Huan Wang – Salesforce AI Research.