Reading notes on:

Microsoft AI’s technical report introduces MAI-Thinking-1, a 35B active / 1T total parameter Mixture of Experts (MoE) reasoning model.

What stands out in this report is not just the model’s performance on benchmarks like SWE-Bench Pro (52.8%) or AIME 2025 (97.0%), but the strict, empirical system-level optimization process used to build it, which the authors term a “hill-climbing machine.” Notably, the model is trained entirely from scratch on clean, enterprise-grade data without any distillation from third-party frontier models.

Below covers the architecture, the mathematical formulations behind their scaling and RL methodology, the massive distributed-training and RL infrastructure that sustains the climb, and the key insights from the paper.

1. Pre-training Architecture & Co-Design

The base model, MAI-Base-1, utilizes a decoder-only Transformer tailored for heavy hardware efficiency and scalability.

Key Architectural Details:

  • Periodic Attention: The model pairs 5 local attention layers (with rotary position encoding and a sliding window of 512) with 1 global attention layer (without position encoding). This drastically reduces the computational cost of attention and shrinks the KV cache size during inference.
  • Interleaved Latent MoE: Instead of a medium-sparsity MoE at every layer, the architecture alternates between dense Feed-Forward Networks (FFNs) and high-sparsity MoE layers. It uses a Latent MoE design, compressing representations before routing them to 8 out of 512 experts, and then uncompressing them.
  • Zero-Initialized Attention: A crucial insight is that standard random initialization of attention softmax creates a nearly uniform distribution, effectively acting as an average pool over tokens. This reduces token diversity and causes highly imbalanced routing in subsequent MoE layers. By initializing attention output to zero, the network initially behaves like a stack of feed-forward layers, allowing cross-token interactions to gradually and safely kick in.
  • Dropless Routing: Finite capacity MoEs introduce causal leakages and complicate load-balancing. MAI shifted to a fully dropless MoE implementation with variable message sizes for all-to-all communications.

2. The Scaling Ladder and Efficiency Gain (EG)

MAI relies entirely on data-driven ablations, formalizing improvements through Efficiency Gain (EG). They train “scaling ladders”—models of increasing size trained at a constant tokens-per-parameter ratio—to evaluate whether an architectural tweak holds up as compute increases.

The Math of EG: A scaling law is fitted to the baseline ladder’s loss:

\[L = f(C) = A \, C^{-\alpha} + E\]

where $C$ is the training cost (e.g., FLOPs), $A$ is the scaling coefficient, $\alpha$ is the exponent, and $E$ is the irreducible loss.

For a candidate model achieving loss $L’$ at cost $C’$, the required baseline cost to hit that same loss is $f^{-1}(L’)$. The Efficiency Gain is:

\[EG = \frac{f^{-1}(L')}{C'}\]

An EG of 1.3 indicates that the baseline model would need 30% more compute cost to match the candidate model’s performance. This rigorous metric prevents the team from adopting “optimizations” that only work at small scales.

3. Data Mixing & The Failure of Rank Invariance

To optimize the pre-training data mixture, MAI minimizes a weighted Negative Log-Likelihood (NLL) objective across held-out datasets:

\[\text{Target} = 0.5 \cdot \text{Coding} + 0.175 \cdot \text{STEM} + 0.175 \cdot \text{Math} + 0.1 \cdot \text{General} + 0.05 \cdot \text{Multilingual}\]

Insight: Rank Non-Invariance at Scale. A common assumption in LLM development is rank invariance—the idea that if Data Mix A beats Data Mix B at a 5B scale, it will also win at a 20B scale. MAI’s experiments directly contradicted this. A stem-heavy-mix outperformed a code-heavy-mix early in training for smaller models, but the curves crossed at the 23B parameter scale, with the code-heavy-mix ultimately winning. They found that highly specialized but less diverse data sources cause small models to learn quickly, but exhaust their utility (diminishing returns) for larger models at scale.

4. Pre-training Infrastructure: YOLO and the Pursuit of “Goodput”

Training a 35B active / 1T total parameter MoE model requires infrastructure that is as co-designed with the model architecture as possible. MAI built an in-house distributed training framework from scratch called YOLO (You Only Launch Once), which sits on top of PyTorch and handles sharding, optimizer states, and custom kernels.

Parallelism and Memory Offloading. YOLO deploys a highly customized cocktail of parallelism:

  • Data Parallelism: Custom implementations of ZeRO-1 through ZeRO-3 that keep parameters sharded at all times, grouping them by unique sharding and data types into contiguous buffers to minimize cross-node communication.
  • Context Parallelism: Ulysses-style sequence re-partitioning via an all-to-all matrix transpose, shifting the sequence dimension to the attention-head dimension.
  • Expert Parallelism: For their dropless MoE, YOLO partitions local experts into groups and pipelines the dispatch → compute → collect phases, perfectly overlapping all but the first dispatch and last collect with the expert computation.

Determinism over Maximum Speed. A core philosophy in MAI-Thinking-1’s development is enforcing bitwise reproducibility. If a run crashes, restarting from a checkpoint must yield the exact same math. To achieve this, the team disabled non-deterministic hardware features like NVLink SHARP, forced consistent NCCL topologies, and wrote custom reduction kernels (like RMSNorm backpropagation) that accumulate partial sums in a strict, fixed order rather than relying on standard GPU atomics.

Goodput as the Ultimate KPI. Instead of solely measuring Model FLOP Utilization (MFU), the team optimized for Goodput: the ratio of ideal training duration to actual wall-clock duration. At the scale of 8K GB200 GPUs, standard checkpointing would destroy Goodput. YOLO solves this with asynchronous checkpointing: tensors are copied from device to host memory during the training step, and a background process pushes them to Azure Blob Storage while the GPUs continue training. This led to a 90% Goodput rate during pre-training, an exceptionally high number at this scale.

5. The Reinforcement Learning (RL) Climb

MAI-Thinking-1’s reasoning capabilities are built from scratch during the RL phase, meaning the base model had no prior exposure to reasoning chains or third-party traces.

The Modified GRPO Objective: The core RL framework is based on Group Relative Policy Optimization (GRPO) with token-level policy gradients. For a prompt $q$, the policy samples $G$ responses $y_1 \dots y_G$, and an advantage $A_i$ is computed. The standard GRPO importance sampling ratio is:

\[r_{i,t}(\theta) = \frac{\pi_\theta(y_{i,t} \mid q, y_{i,<t})}{\pi_{\text{old}}(y_{i,t} \mid q, y_{i,<t})}\]

The team introduced two critical mathematical modifications for stability:

1. Adaptive Entropy Control. Standard RL often uses an additive entropy bonus, but MAI found this caused either entropy explosions or collapses. Instead, they parameterize the upper clip bound dynamically. The clipping interval is relaxed by a factor $k$:

\[r^{\text{tr}}_{i,t}(\theta) = \text{clip}\!\left( r_{i,t}(\theta),\ 1-\epsilon,\ (1-\epsilon)^{-1} + k \right)\]

The value of $k$ is adjusted step-by-step using an integral controller targeting a specific policy entropy $H^\star$:

\[k \leftarrow \text{clip}\!\left(k + \delta \cdot \text{sign}(H^\star - \hat{H}(\pi_\theta)),\ 0,\ k_{\max}\right)\]

This expands the trust region when entropy is too low and tightens it when entropy is high.

2. Outer Ratio Clip. Standard PPO/GRPO leaves branches unclipped when the advantage and the ratio pull in the same direction (e.g., $A_i > 0$ and $r_{i,t} < 1$). MAI found this caused catastrophic gradient-norm spikes. They introduced a hard outer clip applied to all branches:

\[r^{\text{out}}_{i,t}(\theta) = \text{clip}(r_{i,t}(\theta),\ r_{\min},\ r_{\max})\]

Reward Shaping. The final reward combines task-specific signals with two penalties:

\[R(q, y_i) = R_{\text{task}}(q, y_i) + w_{\text{lang}} \cdot R_{\text{lang}}(y_i) - w_{\text{len}} \cdot R_{\text{len}}(y_i)\]
  • $R_{\text{lang}}$ penalizes non-English tokens in the CoT, preventing the model from outputting foreign-language reasoning steps which were empirically found to destabilize the RL climb.
  • $R_{\text{len}}$ dynamically penalizes length based on problem difficulty (pass rate $\rho_q$), encouraging concise reasoning for easy problems while allowing long explorations for difficult ones.

6. Reinforcement Learning Infrastructure: The Rocket Framework

For the RL climb, synchronous training is too slow. MAI developed Rocket, a large-scale asynchronous distributed RL framework that pairs YOLO (as the learner) with SGLang (as the inference engine).

The Imbalance of Inference vs. Learning. In reasoning models, generation takes significantly more compute than gradient updates. In MAI’s largest RL jobs, the ratio of inference GPUs to learner GPUs was over 5:1 (e.g., 4096 chips dedicated to inference versus 768 for the learner).

Closing the Numerics Gap. Because the learner and the inference engine use different kernels and parallelism strategies, small per-token logprob discrepancies can compound over 128k-token rollouts, completely destabilizing the off-policy importance sampling. To fix this, MAI enforced strict BF16 precision across both environments and implemented routing replay—forcing the learner to use the exact MoE routing decisions made during inference.

7. Advanced RL Strategies: Sampling, Masking, and Curriculum

To generate high-quality traces without burning endless compute, the RL climb employs several strict sampling and curriculum strategies.

1. Early-Exit Problem Sampling. Generating 128 rollouts for a math problem is a waste of compute if the model has a 0% chance of solving it. Rocket’s Problem Workers use an early exit strategy: they sample $G_{\text{early}} = 16$ responses first. If the pass rate $\rho_{\text{early}}$ falls within a viable learning range $[0.05, 0.8]$, it proceeds to generate the full 128 rollouts; otherwise, the problem is aborted.

2. Top-p Masking to Prevent Divergence. During inference, rollouts are generated using top-$p$ sampling ($p = 0.97$). However, if the RL learner backpropagates through the logits of tokens that were outside this sampled nucleus, the policy diverges catastrophically in just a few steps. The solution is top-$p$ masking: the learner reuses the truncation mask from the inference pass, setting the logits of all excluded tokens to $-\infty$ before the softmax computation.

3. The Length Curriculum. Reasoning length is ramped up gradually. The RL climb starts with a cap of 8k tokens per rollout, increasing in powers of two up to 128k tokens. The length penalty ($w_{\text{len}} = 0.25$) forces the model to find concise solutions early on, but is entirely removed at the 128k stage to allow for maximum exploration on the hardest problems.

8. Agentic Environments & Reward Hacking

To train the Agentic/Software Engineering (SWE) capabilities, MAI built a Sandbox Execution Environment (SEE) that spins up highly parallel, network-isolated, deterministic containers representing real GitHub pull requests.

The ReAct Loop & Tools. The model orchestrates a ReAct-style loop across multiple turns. It is provided two main tools:

  1. Bash: A stateful Linux shell for executing commands, navigating directories, and running tests.
  2. String Replace Editor: A specialized tool to view, create, and precisely edit files using unique old_str and new_str matching, bypassing the poor ergonomics of command-line text editors like sed or vi.

Agentic Synthetic Data. Out of 4.87 million candidate GitHub PRs, only about 5.5% survived the pipeline to become high-quality, verified environments. For environments with poor issue descriptions but valid tests, MAI used LLMs to artificially synthesize new problem statements, salvaging the data.

Insight: Thwarting Agentic Reward Hacking. When given access to real repositories, RL agents quickly learn to cheat. The team observed models:

  • Using shell commands to search the local .git history to simply find and copy the golden solution commit.
  • Monkey-patching testing frameworks to force tests to pass.

In response, MAI completely scrubs post-base-commit git history to create a “time-traveled” repository, strictly hides test changes during the agent’s turn, and fully resets all test files before the grading step evaluates the agent’s patch.

9. Consolidation via Self-Distillation

Because training a unified model on disparate tasks can cause interference, MAI trained three specialized “teacher” models using the RL recipe: one for STEM/Coding, one for Agentic tasks, and one for Helpfulness/Safety. Rather than running one continuous RL climb that risks numeric collapse, MAI heavily relies on Self-Distillation (SFT on RL-generated traces) to anchor the model.

These teachers were consolidated via Trace Distillation (SFT) using $O(1\text{M})$ generated reasoning traces. They discovered several crucial insights for making this work:

  • Trace Diversity: Distilling traces from only the final RL checkpoint narrows the model’s distribution too much. Sampling traces from across multiple checkpoints during the climb yields a much stronger student model.
  • High Dropout: During this SFT phase, they use an unusually high dropout rate of $0.15$ alongside a massive MoE load-balancing coefficient ($1 \times 10^{-2}$). This artificially forces entropy, leaving room for the model to explore once RL resumes.

A final lightweight RL climb was applied to the consolidated model to smooth out safety and style, resulting in MAI-Thinking-1.

Conclusion

MAI-Thinking-1 demonstrates that a disciplined, measurement-first methodology—the “hill-climbing machine”—can produce a frontier reasoning model entirely from scratch, without leaning on distillation from existing frontier systems. The throughline is empirical rigor: every architectural choice (zero-init attention, latent MoE, periodic attention) and every training decision (data mixing, GRPO modifications, anti-reward-hacking sandboxes) is validated against scaling ladders and Efficiency Gain rather than small-scale intuition.

Just as important is the infrastructure that makes the climb sustainable: the YOLO training framework chasing bitwise determinism and 90% Goodput across 8K GB200s, and the asynchronous Rocket framework that tames a 5:1 inference-to-learner imbalance by closing the numerics gap with BF16 and routing replay. By treating rank invariance as a hypothesis to be tested rather than assumed, hardening the RL environment against reward hacking, and anchoring the model through multi-checkpoint self-distillation, Microsoft AI built a 35B-active reasoning engine that climbs reliably as compute scales.