DEV Community: MemoraX AI

Are We Optimizing the Training Process—or the Model’s Ability to Reason?

MemoraX AI — Tue, 27 Oct 2026 10:03:59 +0000

In large language model reinforcement learning, better training metrics do not always translate into better reasoning performance.

A lower training loss, higher reward, or continuously updated policy may look like progress. But when the model is evaluated in a real inference environment, the improvement can be much smaller—or may not appear at all.

Why does this happen?

One possible explanation is the mismatch between training and inference. Differences in sampling strategies, numerical precision, inference engines, and system configurations can cause the model to behave differently during training and deployment.

This raises a fundamental question:

Should we optimize the training policy itself, or should we directly optimize how the model performs during inference?

MemoraX AI, together with Tianjin University and Alibaba, explores this question in a research project for NeurIPS 2026:

The Mirage of Optimizing Training Policies: Monotonic Inference Policies as the Real Objective for LLM Reinforcement Learning

Paper:

https://arxiv.org/abs/2606.29526

Project page:

https://anitaleungxx.github.io/MIPU

The problem: training improvements do not always transfer to inference

Most reinforcement learning pipelines focus on improving the training process through reward design, policy updates, sampling strategies, and optimization objectives.

However, the training environment and the inference environment are not always identical.

Training may involve systems such as vLLM, SGLang, Megatron, or FSDP, while inference may use a different runtime or sampling configuration. Even small differences in implementation can affect the model’s output distribution.

As a result, a policy that looks better during training may not necessarily produce better answers during inference.

This training–inference mismatch can lead to several problems:

Training rewards may not fully reflect real-world reasoning quality;
Some policy updates may improve local metrics without improving final answers;
Model performance may fluctuate during training;
Improvements may not transfer consistently across inference systems.

The key question is therefore not simply whether the training policy has improved.

The more important question is:

Did the model actually become better at solving the task?

MIPI: making inference quality the real objective

MIPI starts from a simple idea:

Training updates should be evaluated by how they affect the model’s actual inference behavior.

Let μₖ denote the inference policy after the k-th update, πₖ denote the training policy at step k, and J denote the model’s performance on the target inference task.

MIPI aims to enforce:

J(μₖ₊₁) − J(μₖ) ≥ 0

In other words, each effective update should preserve or improve the model’s inference performance.

The change in inference performance can be decomposed as:

J(μₖ₊₁) − J(μₖ)

= [J(μₖ₊₁) − J(πₖ₊₁)]

+ [J(πₖ₊₁) − J(πₖ)]

+ [J(πₖ) − J(μₖ)]

This decomposition highlights three different factors:

The gap between the updated inference policy and the updated training policy;
The improvement introduced by the training update itself;
The gap between the previous training policy and the previous inference policy.

This perspective suggests that improving the training policy alone is not enough. We also need to understand how training updates affect the model’s actual inference behavior.

MIPU: turning the principle into a practical method

To make MIPI actionable, the paper introduces MIPU.

MIPU consists of two main stages:

Generate candidate updates;
Evaluate candidate inference behavior and decide whether to accept the update.

Step 1: Generate candidate updates

MIPU compares the behavior of the policy before and after an update.

It considers the probability assigned to candidate responses by both the current training policy and the updated policy:

πᵏθ / μₖ = (πₖ / μₖ) × (πθ / πₖ)

This allows MIPU to separate changes caused by the training update from differences between training-time and inference-time behavior.

The first stage produces a set of candidate models or candidate updates for further evaluation.

Step 2: Evaluate candidate responses

MIPU then evaluates the actual responses generated by each candidate.

An update is more valuable when it:

Increases the probability of high-quality answers;
Reduces the probability of low-quality answers;
Improves performance across multiple tasks;
Preserves training stability.

The goal is not to reject as many updates as possible.

Rejecting too many candidates can prevent the model from exploring new reasoning paths. Accepting too many low-quality updates can introduce instability.

MIPU therefore aims to balance reliable improvement with continued exploration.

Experiments

The method was evaluated on Qwen3-1.7B and Qwen3-4B models under FP8 training settings.

The experiments covered several mathematical reasoning benchmarks, including:

MATH-500;
AIME24;
OlympiadBench;
Minerva;
AMC23.

MIPU was compared with standard GRPO, MIS, and learning-rate decay.

Model	GRPO	MIS	Learning-rate decay	MIPU	Improvement over GRPO
Qwen3-4B	64.42%	63.42%	65.66%	66.71%	+2.29 points
Qwen3-1.7B	50.86%	51.73%	52.23%	53.97%	+3.11 points

MIPU achieved higher average accuracy on both model sizes.

On Qwen3-4B, MIPU reached 66.71% average accuracy, improving over standard GRPO by 2.29 percentage points.

On Qwen3-1.7B, MIPU reached 53.97%, improving over GRPO by 3.11 percentage points.

The method also showed more stable training trajectories, suggesting that the improvement was not limited to a single evaluation point.

Why candidate selection matters

The paper also evaluates the contribution of each stage:

Method	Average accuracy
Standard GRPO	64.42%
Step 1: candidate generation only	65.36%
Step 2: candidate evaluation only	62.81%
Full MIPU	66.71%

The results suggest that candidate generation and candidate evaluation play complementary roles.

Step 1 helps identify potentially useful updates. Step 2 helps determine whether those updates actually improve inference behavior.

Together, they allow MIPU to filter ineffective updates while preserving the model’s ability to explore new solutions.

Looking ahead

This work focuses on 1.7B and 4B models and mathematical reasoning tasks. Several directions remain open:

Scaling MIPU to larger models;
Applying the method to code generation and tool use;
Evaluating it across different inference engines;
Reducing the computational cost of candidate evaluation;
Designing adaptive acceptance and rejection strategies;
Extending the framework to more complex reward functions.

The broader message is simple:

Training optimization should ultimately be measured by the quality of the model’s real-world behavior.

MIPI and MIPU offer one possible direction: making the connection between training updates and inference performance more explicit, measurable, and stable.

Paper:

https://arxiv.org/abs/2606.29526

Project page:

https://anitaleungxx.github.io/MIPU

MemoraX AI at NeurIPS 2026: Ten Papers on Reliable Learning, Efficient Reasoning, and Long-Term Agent Memory

MemoraX AI — Thu, 08 Oct 2026 07:19:22 +0000

We are excited to share that 10 papers from MemoraX AI have been accepted to NeurIPS 2026.

These papers cover a range of topics, including robust self-training, reinforcement learning, difficulty-aware reasoning, diffusion models, long-term Agent Memory, and the continuous evolution of AI systems.

Although the papers address different technical problems, they share a common research question:

How can AI systems learn from experience, use computation more effectively, and continue improving over time?

Beyond Bigger Models

As AI systems become capable of handling longer and more complex tasks, simply increasing model size is no longer enough.

A capable AI system should also be able to:

recognize when a problem requires deeper reasoning;
learn from reliable experience rather than noisy feedback;
avoid repeating failed strategies;
align training improvements with real inference performance;
retain and reuse useful information over long-running tasks.

This perspective guides much of MemoraX AI’s research.

Building More Robust Learning Systems

One line of work focuses on improving the reliability of self-training for large language models.

In “Beyond Correctness: Robustness-Driven Evolutionary Self-Training for Large Language Models,” the authors study how correct solutions can be selected and evolved more reliably during training.

The central idea is that a useful solution should not only be correct once. It should also be robust, stable, and transferable to related problems.

The work introduces a robustness-driven evolutionary self-training framework that uses correctness feedback to guide the generation and selection of new solution trajectories. This helps the model learn from high-quality reasoning paths while reducing the risk of overfitting to fragile solutions.

Aligning Training with Real Inference

Another research direction studies the gap between training-time improvements and inference-time behavior.

In “The Mirage of Optimizing Training Policies: Monotonic Inference Policies as the Real Objective for LLM Reinforcement Learning,” the authors examine a common problem in reinforcement learning for language models: improving the training policy does not always lead to better behavior when the model is actually deployed.

The paper introduces the idea of optimizing for monotonic inference policy improvement. Its proposed framework separates training-side updates from inference-side validation, allowing potentially harmful updates to be detected and rejected.

This line of research asks an important question:

Are we optimizing the metric used during training, or are we truly improving the behavior users experience at inference time?

Making Reasoning More Compute-Efficient

AI systems often spend the same amount of computation on problems with very different levels of difficulty.

“Does Your Large Language Model Have an Intuitive Sense of the Difficulty of a Question?” investigates whether a model can estimate problem difficulty before committing to a reasoning strategy.

The work proposes a difficulty-perception mechanism that helps models dynamically select different reasoning intensities. Easier problems can be handled with a faster path, while harder problems receive additional computation.

This direction is especially important for practical AI systems, where accuracy and efficiency must be balanced. Better difficulty awareness can improve token efficiency without simply reducing the quality of responses.

Exploring New Policy Models for Reinforcement Learning

MemoraX AI also studies how diffusion models can be used in online reinforcement learning.

“What Kind of Diffusion Models Do We Need in Online Reinforcement Learning?” explores the challenges of using diffusion policies in high-dimensional action spaces, including computational efficiency, expressiveness, and distribution coverage.

The paper introduces Coupled Flow, a structured approach designed to preserve expressive policy modeling while improving optimization efficiency.

This work contributes to a broader effort to understand how generative models can support more capable decision-making systems in complex environments.

Long-Term Agent Memory

One of MemoraX AI’s central research areas is long-term Agent Memory.

As AI Agents move from short interactions to longer-running tasks, memory becomes more than a convenience. It becomes part of the system’s ability to learn from experience.

An Agent working on a real software project may need to remember:

why a particular architecture was chosen;
which implementation paths have already failed;
how a previous bug was fixed;
which project rules are stable;
how the team prefers to structure and review code.

This is the motivation behind MemoraX Code, our long-term memory plugin for Coding Agents.

MemoraX Code helps Coding Agents retain and reuse project context, historical decisions, debugging experience, failed approaches, verified fixes, and reusable development workflows.

It is designed to work with Coding Agents such as Codex, Claude Code, OpenCode, Trae, and WorkBuddy. The goal is not to replace a developer’s preferred Agent, but to help the Agent maintain continuity across long-running development work.

From Research to Evaluation

In addition to building memory systems, MemoraX AI is also working on how Agent Memory should be evaluated.

Through the Agent Memory Leaderboard, we are exploring an open and reproducible evaluation framework for comparing memory systems across different scenarios.

These scenarios include:

long conversations and cross-session history;
coding memory and historical engineering experience;
multimodal memory;
factual recall and multi-hop reasoning;
temporal understanding;
personalization and rule following;
memory governance and safe usage.

A shared evaluation framework can help researchers and developers move beyond vague claims and better understand what different memory systems can actually do.

Toward Continuously Evolving AI

The ten accepted papers represent different parts of a larger research agenda.

Some focus on how models learn. Some study how they allocate computation. Others examine how policies improve during reinforcement learning or how Agents can retain useful experience over time.

Together, they point toward a broader direction for AI development:

Future AI systems should not only generate answers. They should learn from experience, understand when to reason more deeply, reuse reliable knowledge, and improve continuously through interaction.

At MemoraX AI, we are continuing to explore this direction through both fundamental research and practical systems such as MemoraX Code and Agent Memory Leaderboard.

We look forward to sharing more details about the accepted papers, released benchmarks, and open research projects in the coming weeks.

Learn More

MemoraX AI: Official Website
MemoraX Code: Product Website
GitHub: MemoraX AI on GitHub

ai #machinelearning #agents #reinforcementlearning #llm #opensource