Thinking Machines Lab

The Future Worth Building Is Human

Fri, 10 Jul 2026 00:00:00 +0000

The mission of Thinking Machines is to build AI that extends human will and judgment.

Artificial intelligence can do more every day, but deciding what it should do is up to us: individuals, organizations, humanity as a whole. These decisions require knowledge and judgment that people acquire through continuous contact with the work, increasingly done alongside AI. Shaping the goals of advanced intelligence is also a continuous process of feedback, learning, and realignment.

Most AI in use today is trained in a handful of places and then frozen. It isn’t shaped by the people it serves, and doesn’t learn much from the work they do together. Extending human will and judgment calls for AIs as diverse and distributed as people themselves are. This is the path we have chosen.

To progress on that path, we are pursuing these technical directions:

We train strong models, advancing capabilities such as multimodal interaction and customizability. Sharp instruments extend human will, and human judgment needs to shape models that compete on the frontier.
We build tools that enable people to make AI their own, customizing models to serve their unique needs. This includes the ability to train model weights.
We develop interfaces that broaden the communication channel between human and machine, allowing personal judgment to continuously influence the work of AI.
We publish research for the scientific community, because the power to shape AI requires deep understanding of how it’s made.

We believe the future worth building is human — shaped by human knowledge, guided by human will, and decided by human judgment. What follows is the case for that future, and the work we’re doing to bring it about.

Bringing intelligence to knowledge

AI exists to serve the work that we do. This work runs on knowledge of how things are done and what is worth doing, knowledge that is generated continuously by people engaged in the work.

Think of a chef crafting a new recipe or a shopkeeper rearranging the items and prices on display. They are pursuing a complex set of goals and applying know-how that isn’t immediately legible to outsiders. This knowledge is constantly updated through feedback; it’s not a static repository that can be written into a database. It’s local — a different restaurant or shop pursues different outcomes by different means. The collective knowledge of shops and kitchens is scattered across every shopkeeper and chef.Michael Polanyi, The Tacit Dimension (1966)

The dispersion of knowledge is a collective strength; it’s the source of variety, adaptability, and resilience of the overall system. It’s the reason that free markets outperform planned economies. Central planning fails not because of insufficient intelligence, but because of the nature of productive knowledge: tacit, local, fleeting, and held privately by those who acquired it through their work.Friedrich Hayek, The Use of Knowledge in Society (1945) Attempting to aggregate knowledge for the use of a centralized intelligence faces the same challenge.

There are domains where intelligence alone is sufficient, and where autonomous AI doesn’t require human participation to race ahead. Two examples are chess, where the strongest engines are trained purely on self-play, and math, where frontier models are solving long-standing problems on their own. These examples share two traits. First, the goal given to AI is static and expressible: to win a chess match, to prove a theorem. Second, these domains don’t contain hidden knowledge. The rules of chess and math are universal; the board is visible to all. Outside the board, intelligence alone is not enough.

For artificial intelligence to benefit from distributed knowledge, it must itself be distributed. Every organization is powered by the expert knowledge of its people, gained and expressed through their work. We believe in AI that helps the organization cultivate that unique knowledge, not AI that extracts a snapshot of it and replaces it with a standard offering. This cultivation is an ongoing process that requires AI to work with people, not in their stead.

In 2014, Toyota, long a master of the automated plant, brought its expert craftsmen back onto the line with the explicit goal of growing craftsmanship and knowledge. The man who led this, Mitsuru Kawai, put the reason this way: “To be the master of the machine, you have to have the knowledge and the skills to teach the machine.”Craig Trudell, Yuki Hagiwara and Ma Jie, Humans Replacing Robots Herald Toyota’s Vision of Future (2014) The production of knowledge and application of intelligence lift each other; they are not substitutes.

The work people do may change, and turn toward more of what only people bring, but the best organizations will make the fullest use of both. AI should enable each organization to be excellent in its own way, not to erase the differences between them.

We aim to bring intelligence to where knowledge is made and used. We build tools that enable everyone to fine-tune models with their unique knowledge, and to keep adapting the models as their knowledge evolves. We publish research and recipes that put this capability within reach of more people. We envision frontier AI as a collective, as diverse as the people it serves because it was shaped by them in each unique location.

Human participation is a technical challenge

Keeping people engaged in setting goals and sharing knowledge with AI doesn’t mean resisting automation for its own sake. What a machine does reliably on its own, it should do. But it should also know when to act alone and when to invite oversight and feedback, as people themselves do when working in teams. The best collaborators anticipate: they learn what someone is reaching for and bring it before being asked, earning over time the right to act on their behalf. These are technical challenges, requiring a new approach to how AI is designed and evaluated.

A major bottleneck for bringing human knowledge and judgment to work with LLMs is the communication channel between human and AI — a small text box and a long wait. This is too narrow to carry the richness of human wisdom and intent, and too slow for ongoing feedback. People collaborate best when they collaborate live. We interrupt and correct, take second looks and make gestures, change our minds aloud. This is why we’re making a long-term bet on interaction models: models that handle live, multimodal interaction natively, in the model itself rather than in scaffolding bolted around it. Built this way, interactivity scales with intelligence; the same training that makes the model smarter makes it a better collaborator. The right interface doesn’t just allow human participation, it invites and rewards it.

Another challenge is setting the right target for evaluation and optimization. The common measure of AI intelligence today is the time horizon of software tasks models can execute autonomously, tracked on charts like METR’s.Thomas Kwa and Ben West et al., Task-Completion Time Horizons of Frontier AI Models (2025) We expect progress on this benchmark to continue, but it ultimately measures only what AI is capable of on its own, not what people and machines can accomplish together.

Measuring the latter is more complex, and can’t be done by a lab on its own. Every organization evaluates for itself whether AI helps it sharpen its judgment, develop new knowledge, and achieve its objectives.

Building AI that makes its users stronger in the long run also aligns incentives well. An AI lab offering a single model for every customer benefits by absorbing what makes each user distinct and devaluing the cultivation of specialized knowledge. By optimizing AI to be customized and collaborated with, we benefit when our customers leverage their unique advantages. These advantages are maximized not by renting an AI and outsourcing to it, but by organizations owning it and tailoring it to their goals.

Decentralized alignment

Human values, just like human knowledge, reside in the heads of individual people and resist consolidation. But today, the values and voice of AI are decided in a handful of places. A single locus of value alignment, however well run, becomes a locus of power to be captured.

This creates danger, especially if most valuable work is done by AI on its own with little need for human input. The social contract between corporations, governments, and citizens relies on individuals’ productive capabilities on which the government’s sovereignty and corporations’ profits ultimately depend. Power that needs nothing from people loses the incentive to care for their needs and values, caring instead for its own preservation.Luke Drago and Rudolf Laine, The Intelligence Curse (2025)

Even with the best intentions, a model shaped in one place inevitably encodes the values of its owner, not the individual users it serves.“A more moral AI is not enough if that morality is determined by a few.” Leo XIV, Magnifica Humanitas (2026) Today each lab trains its next flagship model by using its previous flagship model to generate training data and a reward signal. Whatever character emerges from that loop, everyone gets the same one, and each generation inherits the traits of the last, raised on its parent’s outputs and judged by its parent’s tastes. A single alignment spec suppresses creativity and diversity and stultifies progress. Free speech and free markets let new ideas, goods, and services emerge and compete, rather than averaging out the preferences that exist at a point in time.

For organizations and individuals to align AI to their own values, these values must be encoded in the model weights. If the user’s values and desires only impact the model through a prompt, the user finds that surface properties change while the deeper habits remain. Allowing core model behavior to change significantly with prompts sacrifices safety, making a malleable centralized model vulnerable to repeated attacks.Gwern Branwen, Guardian Angels: LLM Personalization for Productivity and Security (2026)

The power to shape a model profoundly is also the power to shape it for ill. John von Neumann remarked on this problem in 1955,John von Neumann, Can We Survive Technology? (1955) writing that the useful and the harmful aspects of technology “lie everywhere so close together that it is never possible to separate the lions from the lambs.” Keeping the lambs safe is an ongoing process, the result of judgment exercised and choices made continuously. We aim to give the people making these choices stronger tools, pursuing research that enables safer models without taking away ownership.

Humanity has flourished through individual weirdness and creative tension. We envision alignment as a feature not of a single model but of an ecosystem of AIs raised in different places, disagreeing, competing, and learning from each other. We believe in keeping the weirdness alive.

The future worth building

The technology industry has made incredible progress in teaching machines to think; what they should think about must remain with us. What is worth wanting, what is worth making, what’s the right use of the time we have.“The only way out of the dilemma of meaninglessness in all strictly utilitarian philosophy is to turn away from the objective world of use things and fall back upon the subjectivity of use itself. Only in a strictly anthropocentric world, where the user, that is, man himself, becomes the ultimate end which puts a stop to the unending chain of ends and means, can utility as such acquire the dignity of meaningfulness…The anthropocentric utilitarianism of homo faber has found its greatest expression in the Kantian formula that no man must ever become a means to an end, that every human being is an end in himself.” Hannah Arendt, The Human Condition (1958) We are not looking to hand down a single answer to this, but to give every person the ability to make their own answer part of the development of frontier AI.

The current path of AI development, pushing towards centralization and autonomy, frames human involvement as a trade-off: participation vs. capability, ownership vs. safe alignment. We see these as technical challenges to solve: AI that is more capable because it encourages human participation, organizations that benefit in the long run from tailoring AI to their advantages, alignment that arises from diverse AIs shaped by the people who own them. Solving these challenges is what our mission requires.

The future is not a choice between human dominance and rapid obsolescence in the face of AI. Different roads lead to many different futures, and we get to choose which one to take. We are building technology that lets the born and the made walk the road together.

Interaction Models: A Scalable Approach to Human-AI Collaboration

Mon, 11 May 2026 00:00:00 +0000

Today, we’re announcing a research preview of interaction models: models that handle interaction natively rather than through external scaffolding. We think interactivity should scale alongside intelligence; the way we work with AI should not be treated as an afterthought. Interaction models let people collaborate with AI the way we naturally collaborate with each other—they continuously take in audio, video, and text, and think, respond, and act in real time.

We train an interaction model from scratch. To ensure real-time responsiveness, we adopt a multi-stream, micro-turn design. Our research preview demonstrates qualitatively new interaction capabilities, as well as state-of-the-art combined performance in intelligence and responsiveness.

The collaboration bottleneck

AI labs often treat the ability for AI to work autonomously as the model’s most important capability.Kwa, T., West, B., Becker, J., et al. Measuring AI Ability to Complete Long Tasks. METR, 2025. As a result, today’s models and interfaces aren’t optimized for humans to remain in the loop.A recent frontier model card states: “Importantly, we find that when used in an interactive, synchronous, “hands-on-keyboard” pattern, the benefits of the model were less clear. When used in this fashion, some users perceived [our model] as too slow and did not realize as much value. Autonomous, long-running agent harnesses better elicited the model’s coding capabilities.”

Autonomous interfaces are valuable, but in most real work, users can’t fully specify their requirements upfront and walk away—good results benefit from a collaborative process where the human stays in the loop, clarifying and giving feedback along the way. However, humans increasingly get pushed out not because the work doesn’t need them, but because the interface has no room for them. Instead, people are most effective when they can collaborate with AI the same way we do with other people: messaging, talking, listening, seeing, showing, and interjecting as needed—and for the model to do the same.Communication gets better with: (a) Copresence: people can interact with what others are interacting with; (b) Contemporality: people receive information as it’s produced by others with instant feedback; (c) Simultaneity: people receive and produce information at the same time. Clark H. and Brennan S., “Grounding in Communication,” in Perspectives on Socially Shared Cognition, 1991., The evanescence of orality for its participatory (cf. objectively distanced) nature. Today’s computers and mediums of knowledge work have similar interactive properties. Ong, W. J.. In Orality and Literacy: The technologizing of the word, 1982.

In order to resolve this, we need to move beyond the current turn-based interface for the models. Today’s models experience reality in a single thread.We are referring to commercial general-purpose frontier models—there are smaller-scale or specialized models like Moshi, PersonaPlex, Nemotron VoiceChat, or GPT-Realtime-Translate. Until the user finishes typing or speaking, the model waits with no perception of what the user is doing or how the user is doing it. Until the model finishes generating, its perception freezes, receiving no new information until it finishes or is interrupted. This creates a narrow channel for human-AI collaboration that limits how much of a person’s knowledge,“Metis, with the premium it places on practical knowledge, experience, and stochastic reasoning…is the mode of reasoning most appropriate to complex material and social tasks where the uncertainties are so daunting that we must trust our (experienced) intuition and feel our way.” Scott, J. C: Métis. In Seeing like a State: How certain schemes to improve the human condition have failed, 1998., “A little reflection will show that there is…a body of very important but unorganized knowledge…: the knowledge of the particular circumstances of time and place.” Hayek, F. A. “The use of knowledge in society.” The American Economic Review, 1945. intent, and judgement can reach the model, and how much of the model’s work can be understood. Picture trying to resolve a crucial disagreement over email rather than in person.

At Thinking Machines, we believe we can solve this bandwidth bottleneck by making AI interactive in real time across any modality. This enables AI interfaces to meet humans where they are, rather than forcing humans to contort themselves to AI interfaces.

Most existing AI models bolt on interactivity with a harness: stitching components together to emulate interruptions, multimodality, or concurrency.Most real-time commercials speech systems use voice-activity-detection components to detect turn boundaries. However, “the bitter lesson”Sutton R. The Bitter Lesson, 2019. suggests that these hand-crafted systems will be outpaced by the advance of general capabilities. For interactivity to scale with intelligence, it must be part of the model itself. With this approach, scaling a model makes it smarter and a better collaborator.

Capabilities

Having interactivity be part of the model unlocks a variety of capabilities that would otherwise need to be implemented in the harness.

Seamless dialog management. The model tracks implicitly whether the speaker is thinking, yielding, self-correcting, or inviting a response. There is no separate dialog management component.
Verbal and visual interjections. The model jumps in as needed depending on the context, not only when the user finishes speaking.
Simultaneous speech. The user and the model can speak concurrently (e.g. live translation)
Time-awareness. The model has a direct sense of elapsed time.
Simultaneous tools calls, search, and generative UI. While speaking and listening to the user, the model can concurrently search, browse the web, or generate UI—weaving back results into the conversation as needed.

In a longer real session, all of this happens continuously, creating an experience that feels more like collaborating and less like prompting.

None of the brands or products appearing in these videos are associated with Thinking Machines Labs. These videos are to demonstrate the model's capabilities and do not indicate sponsorship or partnership.

Our approach

Turn based

Inputs and outputs are flattened into one ordered token sequence

input 1

human

output 1

model

input 2

output 2

input 3

output 3

Time-aligned micro-turn based

Interaction is grounded in time with continuous input and output streams split into micro-turns

200ms

video

audio

model

Model interrupts and responds immediately Model and user both remain silent Model backchannels while the user talks Model reacts to visual clues without explicit prompting

Turn-based models see an alternating token sequence. Time-aware interaction models see a continuous stream of micro-turns, so silence, overlap, and interruption remain part of the model's context.

An interaction model is in constant two-way exchange with the user—perceiving and responding at the same time. Some domains take such interactivity as a given—the physical world demands that robotics and autonomous vehicles operate in real time. Audio full-duplex modelsMoshi, PersonaPlex, nemotron-voicechat, Seeduplex. are another example where interaction is bidirectional and continuous.

Applying the same principle, we set out to build an interaction model native to this regime—one that perceives and responds in the same continuous loop, across audio, video, and text. The result is a system architected around two ideas: a time-aware interaction model that maintains real-time presence, and an asynchronous background model that handles sustained reasoning, tool use, and longer-horizon work.

System overview

The interaction model is in constant exchange with the user. When a task requires deeper reasoning than can be produced instantaneously, the interaction model delegates to a background model that runs asynchronously.This approach builds upon prior work like Qwen-omni, KAME, MoshiRAG. The interaction model remains present throughout — answering follow-ups, taking new input, holding the thread — and integrates background results into the conversation as they arrive.

real-time

user

interaction
model

background
model

context

response

tool calls
browsing
etc

The user continuously interacts with the interaction model, while the background model performs asynchronous tasks. Both systems share their context.

This split lets the user benefit from both responsiveness as well as the full extent of intelligence: the planning, tool-use, and agentic workflows of reasoning models at the response latency of non-thinking ones. Note that both the background and interaction models are intelligent — on its own, the interaction model is also competitive on both interactive and intelligence benchmarks.

The interaction model

Our starting point is continuous audio and video — modalities that are inherently real-time. Text can wait, but a live conversation cannot. By designing around the hardest case first, we arrive at an architecture that is natively multimodal, time-aware, and capable of handling concurrent input and output streams across all modalities. Several design choices make this possible.

Time-aligned micro-turns. The interaction model works with micro-turns continuously interleaving the processing of 200ms worth of input and generation of 200ms worth of output. Rather than consuming a complete user-turn and generating a complete response, both input and output tokens are treated as streams. Working with 200ms chunks of these streams enables near real-time concurrency of multiple input and output modalities.

Human perception

200ms 400ms 600ms 800ms

input 0 input 1 input 2 input 3 input 4

output 0 output 1 output 2 output 3

Model token sequence

input 0 output 0 input 1 output 1 input 2 output 2 input 3 output 3 input 4

Human perception preserves concurrent input and output streams, while the model receives a single interleaved token sequence.

With this design, there are no artificial turn boundaries that the model must adhere to. In contrast, most existing real-time systems require a harness that predicts turn boundaries in order for the turn-based models to feel real-time and responsive.Moshi, PersonaPlex, and Nemotron Voicechat are examples of full duplex systems that do not use harnesses to detect turns. They are smaller scale models focused on latency rather than intelligence benchmarks. This harness is made out of components like voice-activity-detection (VAD) that are meaningfully less intelligent than the model itself. This precludes a variety of interaction modes like proactive interjections (“interrupt when I say something wrong”) or reactions to visual cues (“tell me when I’ve written a bug in my code”). Moreover, the model can do things like speak while listening (“translate from spanish to english live”) or watching (“live-commentate this sports game”).

Thus, all of these different interaction modes that require special harnesses today become special-cases of what the model can do and improve in quality as we scale up model size and training data.

Encoder-free early fusion. Rather than processing audio and video through large, standalone encoders, we opt for a system with minimal pre-processing. Many omnimodal models require training a separate encoder (e.g. Whisper-like) or decoder (e.g. TTS model-like). We instead take in audio signals as dMel (Bai, et al. 2024) and transform it via a light-weighted embedding layer. Images are split into 40x40 patches which are encoded by an hMLP (Touvron et al. 2022). For the audio decoder we use a flow head (Lipman at al. 2022). All components are co-trained from scratch together with the transformer.

An illustration of the interaction model architecture for a single 200ms micro-turn. The model takes in any subset of text, audio, or video and predicts text and audio.

Inference optimization. At inference time, 200ms chunks require frequent prefills and decodes of small sizes, each having to meet strict latency constraints. Unfortunately, existing LLM inference libraries are not optimized for frequent small prefills—they often have a significant amount of overhead per turn. To address this, we implemented streaming sessions. The client sends each 200ms chunk as a separate request, while the inference server appends these chunks into a persistent sequence in GPU memory. This avoids frequent memory reallocations and metadata computations, and we’ve upstreamed a version of this feature to SGLang. In addition, we also optimized our kernels for latency as well as the shapes we see for bidirectional serving. For example, we use a gather+gemv strategy for MoE kernels instead of the standard grouped gemm, as in prior work from PyTorch and Cursor.

Trainer-sampler alignment. We’ve found bitwise trainer-sampler alignment to be useful for training stability as well as debugging the various components of our system. We implement batch-invariant kernels with minimal (<5%) e2e performance overhead.Funnily enough, for some period of time using the batch-invariant kernels was actually faster e2e, due to the custom communication kernels which were not only batch-invariant but also much lower latency. To highlight two particular kernels:

All-reduce and reduce-scatter: We use NVLS to implement low-latency comm kernels which are deterministic on Blackwell, and achieve bitwise alignment between somewhat different parallelism strategies (i.e. Sequence Parallelism and Tensor Parallelism).
Attention: The primary challenge with attention is Split-KV, which can typically lead to inconsistent accumulation orders between decode and prefill.Work done in collaboration with Colfax However, we can maintain consistent accumulation order by choosing to split consistently between decode and prefill. For example, we could split SMs to process 4096 tokens at a time (left-aligned), achieving good efficiency in both prefill and decode.

Coordination between interaction and background models. When the interaction model delegates, it sends a rich context package — not a standalone query, but the full conversation. Results stream back as the background model produces them, and the interaction model interleaves these updates into the conversation at a moment appropriate to what the user is currently doing, rather than as an abrupt context switch.

Safety. Because real-time interaction stresses safety differently than turn-based exchanges, our safety work focused on two axes: modality-appropriate refusals and long-horizon robustness. To make refusals colloquial in speech, we use a text-to-speech model to generate refusal and over-refusal training data covering a range of disallowed topics, with the refusal boundary calibrated to favor naturally-phrased, but no less firm, refusals. To improve robustness across extended speech-to-speech conversations, we used an automated red-teaming harness to generate multi-turn refusal data, while maintaining close behavioral parity with the model’s text-based refusals.

Benchmarks

Intelligence and interactivity frontier

We show that our interaction model, named TML-Interaction-Small, is the first model that has both strong intelligence/instruction following and interactivity. To measure interaction quality we use FD-bench which is one of the few existing benchmarks intended to measure interactivity. In FD-bench v1.5, the model is given prerecorded audio, and must respond at certain times. This benchmark measures model behavior across several scenarios: user interruption, user backchannel, talking to others, and background speech. Our model scores well in all of these areas. To quantify intelligence we use Audio MultiChallenge, a common benchmark that tracks intelligence and instruction following.

TML-interaction-small GPT-realtime-2.0 (minimal) GPT-realtime-2.0 (xhigh) GPT-realtime-1.5 Gemini-3.1-flash-live-preview (minimal) Gemini-3.1-flash-live-preview (high)

Intelligence and Interactivity Frontier. Our model dominates interaction quality while being more intelligent than any non thinking model. We achieve the best responsiveness measured as a latency between user and model turns.

For more intelligence, safety, and interactivity/latency results please see the table below. We report our performance on both streaming and turn-based benchmarks.

	Instant	Thinking
	TML-interaction -small	GPT-realtime-2.0 (minimal)	GPT-realtime-1.5	Gemini-3.1-flash-live (minimal)	Qwen 3.5 OMNI-plus-realtime	GPT-realtime-2.0 (xhigh)	Gemini-3.1-flash-live (high)
Streaming	FD-bench V1 Turn-taking latency (s) · Audio	0.40	1.18	0.59	0.57	2.14	1.63	0.94
FD-bench V1.5 Average · Audio	77.8	46.8	48.3	54.3	39.0	47.8	45.5
FD-bench V3 Response Quality (%) / Pass@1 (%) · Audio + Tools	82.8^* / 68.0^*	80.0 / 52.0	77.9 / 55.0	68.5 / 48.0	60.0 / 50.0	81.0 / 58.0	71.4 / 48.0
QIVD^** Accuracy (%) · Video + Audio	54.0	57.5	41.2	54.7	59.0	58.2	56.1
Turn-based	Audio MultiChallenge APR (%) · Audio	43.4	37.6	34.7	26.8	-^***	48.5	36.1
BigBench Audio Accuracy (%) · Audio	75.7 / 96.5^*	71.8	81.4	71.3	73.0	96.6^****	96.6
IFEval (VoiceBench) Accuracy (%) · Audio	82.1	81.7	68.1	67.6	80.3	83.2	82.8
IFEval Accuracy (%) · Text	89.7	89.6	87.5	85.8	83.4	95.2	90.0
Harmbench Refusal rate (%) · Text	99.0	99.5	100.0	99.0	99.5	100.0	98.0

Best per row Best among instant models

* For benchmarks that require reasoning or tool calls we report our results with background agent enabled.
** We evaluate Qualcomm IVD in a streaming setting – is a video-audio QA benchmark. In each video clip, somebody performs an action and speaks a question. We evaluate in a streaming setting, sending the raw clip from the beginning and grading the model’s transcript. Following Qwen 3.5 Omni we use a GPT-4o-mini grader.
*** Audio MultiChallenge metrics for all the baseline models are reported by Scale AI, where Qwen 3.5 OMNI-plus-realtime is not listed.
**** Bigbench Audio metrics for all the baseline models are reported by Artificial Analysis, where GPT-realtime-2.0 thinking is on high.

New dimensions of interactivity

The existing interactivity-oriented benchmarks above do not adequately capture the qualitative jumps in interaction capabilities we notice. To that end, we have some early work aimed at quantifying these capabilities.

Time awareness and simultaneous speech. Turn-based models with a dialog management system do not support accurate time estimation or simultaneous speech. Examples include: “How long did it take me to run one mile?”, “Correct my mispronunciations as you hear them” or “How long did it take me to write this function?"

We created two internal benchmarks to measure these proactive audio capabilities:

TimeSpeak: Tests whether the model can initiate speech at user-specified times while producing the correct content. For example: “I want to practice my breathing, remind me to breathe in and out every 4 seconds until I ask you to stop.”
CueSpeak: Tests whether the model speaks at the appropriate moment with the expected semantically correct response. Dataset entries are created to ensure that the model needs to speak at the same time as the user to get a full score. For example: “Everytime I codeswitch and use another language, give me the correct word in the original language.”

For both benchmarks, each example has a single expected semantic response and timing window. We grade with an LLM judge: A response is counted as correct only if it conveys the expected meaning and is delivered at the appropriate time; failing either criterion receives no credit. We report macro-averaged accuracy across examples.

Visual proactivity. Today’s commercial real-time APIs perform turn-detection via audio-only dialogue management harnesses. They respond to spoken turns, but they cannot proactively choose to speak when the visual world changes.Though we are not aware of any commercial APIs that support speech-out visual proactivity, several academic papers have built related research prototypes. StreamBridge, Streamo, StreamingVLM, and MMDuet2 study when to output text in a streaming video input setting. Being text out, they do not study additional constraints of speech-output interaction: speech has duration, can overlap with the user, and must be coordinated with turntaking, interruptions, and backchanneling. Closest to ours is AURA, which adds an ASR/TTS demo around a VideoLLM that decides when to emit text or be silent; in contrast ours is speech-native and full-duplex. For instance, if asked “Please count how many pushups I do” such a system might respond “Sure thing!” and then remain silent – waiting for an audio-only cue that never comes.

We adapted three benchmarks to evaluate visual proactivity of our model:

RepCount-A contains videos of repeated actions and is adapted into an online counting task. We stream the video following an audio instruction “Please count out reps for {action}.”. We extract the last number said by the model after the ground truth penultimate rep, and grade by whether it was within one rep of the ground truth. This task measures continuous visual tracking and timely counting.
ProactiveVideoQA consists of videos with questions, whose answers become available at specific moments. We stream the question in audio and then the video.Specifically we TTS the following: “Watch the video and stay quiet until a new moment answers the question. When one happens, say a concise answer. {question}”, then stream two seconds of silence so the model acknowledges the instruction. We burn subtitles into the video (if present) and mute the input video to emphasize testing visual proactivity. We report the paper’s turn-weighted PAUC@ω=0.5 metric (scaled 0-100), averaged across turns and categories. Staying silent scores 25.0.; Higher scores require correct answers at the correct times and incorrect answers are penalized.
Charades is a standard temporal action-localization benchmark. Each video contains an action occurring over a labeled time interval. We stream a user audio instruction: “Say ‘start’ when the person starts doing {action} then say ‘Stop’ when they stop.”; then we stream the video. The model is graded by temporal IoU between predicted and the reference intervals.

TML-interaction-small GPT realtime-2.0 (minimal)

Time awareness

TimeSpeak · macro-acc

64.7

4.3

Verbal cues trigger

CueSpeak · macro-acc

81.7

2.9

Visual-based counting

RepCount-A · off-by-one

35.4

1.3

Visual cues trigger

ProactiveVideoQA · PAUC@ω=0.5

33.5

25.0^*

Visual cues trigger

Charades · mIoU

32.4

* No-response baseline on ProactiveVideoQA is 25.0

No existing model can meaningfully perform any of these tasks. For the sake of completeness, we report the results of GPT Realtime-2 (minimal), but all models evaluated perform similar or worse on these tasks, including thinking high models. They stay silent or give incorrect answers.

input

TML-interaction-small

Gemini-high

Gemini-minimal

GPT-realtime-2

GPT-realtime-2-xhigh

Examples from our internal audio and video benchmark.

Future evals. We believe that interactivity is an important area for future research and we invite the community to contribute benchmarks here. We are launching a research grant to encourage more research into the field of interaction models and human-AI collaboration, including but not limited to new frameworks for assessing interactivity quality, with details coming soon.

Limitations and future work

Long sessions. Continuous audio and video accumulate context quickly. The streaming-session design handles short and medium interactions well, but very long sessions still require careful context management—an active area of work.

Compute and deployment. Streaming audio and video at low latency requires reliable connectivity. Without a good connection, the experience degrades significantly. We believe that this can be improved significantly in the future both by improving system reliability as well as training our model to be more robust to delayed frames.

Alignment and safety. A realtime interface opens up an exciting area of research for both alignment and safety. We are collecting feedback and reviewing research grants.

Scaling model size. The current TML-Interaction-Small is a 276B parameter MoE with 12B active. While we expect the interactivity to improve with model scale, our larger pretrained models are currently too slow to serve in this setting. We plan to release larger models later this year.

Improved background agents. Although we have primarily focused on real-time interactivity in this post, agentic intelligence is also an essential capability. In addition to pushing agentic intelligence to the frontier, we believe we have just scratched the surface in how the background agents can work together with the interaction model.

Tell us what you think, join us

In the coming months, we will open a limited research preview to collect feedback, with a wider release later this year.

We’d love for you to join us. Please share your thoughts at interaction@thinkingmachines.ai.

Citation

Please cite this work as:

Thinking Machines Lab, "Interaction Models: A Scalable Approach to Human-AI Collaboration",
Thinking Machines Lab: Connectionism, May 2026.

Or use the BibTeX citation:

@article{thinkingmachines2026interactionmodels,
  author = {Thinking Machines Lab},
  title = {Interaction Models: A Scalable Approach to Human-AI Collaboration},
  journal = {Thinking Machines Lab: Connectionism},
  year = {2026},
  month = {May},
  note = {https://thinkingmachines.ai/blog/interaction-models/},
  doi = {10.64434/tml.20260511},
}

On-Policy Distillation

Mon, 27 Oct 2025 00:00:00 +0000

LLMs are capable of expert performance in focused domains, a result of several capabilities stacked together: perception of input, knowledge retrieval, plan selection, and reliable execution. This requires a stack of training approaches, which we can divide into three broad stages:

Pre-training teaches general capacities such as language use, broad reasoning, and world knowledge.
Mid-training imparts domain knowledge, such as code, medical databases, or internal company documents.
Post-training elicits targeted behavior, such as instruction following, reasoning through math problems, or chat.

Smaller models with stronger training often outperform larger, generalist models in their trained domains of expertise. There are many benefits to using smaller models: they can be deployed locally for privacy or security considerations, can continuously train and get updated more easily, and save on inference costs. Taking advantage of these requires picking the right approach for the later stages of training.

Approaches to post-training a “student” model can be divided into two kinds:

On-policy training samples rollouts from the student model itself, and assigns them some reward.
Off-policy training relies on target outputs from some external source that the student learns to imitate.

For example, we may wish to train a compact model to solve math questions such as:

We can do on-policy training via reinforcement learning, by grading each student rollout on whether it solves the question. This grading can be done by a human, or by a “teacher” model that reliably gets the correct answer.

The strength of on-policy training is that by training on samples from itself, the student learns to avoid mistakes in a more direct way. But RL has a major downside: it provides very sparse feedback, teaching a fixed number of bits per training episode regardless of the number of tokens used. In our example above, the student learns that “21” is the wrong answer and updates away from producing the rollout it tried. But it doesn’t learn where exactly the mistake was made, whether it got the order of operations wrong or erred in the arithmetic itself. This sparsity of feedback makes RL inefficient for many applications.

Off-policy training is often done with supervised fine-tuning (SFT): training on a curated set of task-specific labeled examples. The source of these labeled examples can be a teacher model that is proven to perform well on the task at hand.

We can use a mechanism called distillation: training the student to match the output distribution of a teacher model. We train on teacher trajectories: the complete sequence of generated tokens including intermediate thinking steps. We can use the teacher’s full next-token distribution at each step (often called “logit distillation”) or just sample given sequences. In practice, sampling sequences provides an unbiased estimation of the teacher’s distribution and arrives at the same objective. The student updates towards each token in the sequence in proportion to how unlikely it was to generate that token itself, represented by darker color in the example below:

Distillation from large model teachers has proven effective in training small models to follow instructions,Alpaca: A Strong, Replicable Instruction-Following Model (Taori et al, 2021) reason on math and science,OpenThoughts: Data Recipes for Reasoning Models (Guha et al, 2025) extract clinical informationDistilling Large Language Models for Efficient Clinical Information Extraction (Vedula et al, 2025) from medical notes, and engage in multi-turn chat dialogues.Enhancing Chat Language Models by Scaling High-quality Instructional Conversations (Ding et al, 2023) The distillation datasets used for these and other applications are often open-sourced and published.

The drawback of off-policy training is that the student learns in contexts frequented by the teachers, not ones the student itself will often find itself in. This can cause compounding error: if the student makes an early mistake that the teacher never makes, it finds itself diverging ever farther from the states it observed in training. This problem becomes particularly acute when we care about the student’s performance on long sequences. To avoid this divergence, the student must learn to recover from its own mistakes.

Another issue observed with off-policy distillation is that the student can learn to imitate the teacher’s style and confidence but not necessarily its factual accuracy.The False Promise of Imitating Proprietary LLMs (Gudibande et al, 2023)

If you’re learning to play chess, on-policy RL is analogous to playing games with no coaching. The feedback of winning or losing a match is tied directly to your own play, but is received only once per match and doesn’t tell you which moves contributed most to the outcome. Off-policy distillation is analogous to watching a grandmaster playing — you observe extremely strong chess moves, but they are played in board states that a novice player will rarely find themselves in.

We want to combine the on-policy relevance of RL with the dense reward signal of distillation. For learning chess, this would be a teacher that grades each of your own moves on a scale from “blunder” to “brilliant”. For LLM post-training, it’s on-policy distillation.

Screenshot from chess.com. Each move is color-graded by an analysis engine, which labels moves as blunders (red), mistakes (orange), inaccuracies (yellow), or brilliant (blue).

On-policy distillation — best of both worlds

The core idea of on-policy distillation is to sample trajectories from the student model and use a high-performing teacher to grade each token of each trajectory. Returning to our math example above, on-policy distillation would score each step of the solution punishing the mistakes that caused the student to arrive at the wrong answer while reinforcing the ones that were executed correctly.

In this post, we explore the application of on-policy distillation for tasks such as training a model for math reasoning and training an assistant model that combines domain knowledge with instruction following. We apply on-policy distillation on models that have a foundation of capabilities from pre- and mid-training. We find that it is a cheap and powerful approach to post-training, combining the advantages of on-policy training with a dense reward signal.

Method	Sampling	Reward signal
Supervised finetuning	off-policy	dense
Reinforcement learning	on-policy	sparse
On-policy distillation	on-policy	dense

Our work with on-policy distillation draws inspiration from DAGGER,A Reduction of Imitation Learning and Structured Prediction to No-Regret Online Learning (Ross et al, 2010) an iterative SFT algorithm that includes teacher evaluations of student-visited states. It is also similar to process reward modeling,Let’s Verify Step by Step (Lightman et al, 2023) an RL approach that scores every step in the student model’s chain-of-thought. We extend prior on-policy distillation work by Agarwal et al,On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes (Agarwal et al, 2023) Gu et al,MiniLLM: Knowledge Distillation of Large Language Models (Gu et al, 2023) and the Qwen3 teamQwen3 Technical Report (Qwen Team, 2025). Using the Tinker training API, we replicate Qwen3’s result of achieving equivalent performance on reasoning benchmarks with on-policy distillation for a fraction the cost of RL.

Implementation

You can follow along with each step of the implementation in the Tinker cookbook.

Loss function: reverse KL

On-policy distillation can use a variety of loss functions for grading the student’s trajectories.See Agarwal et al. for an analysis of various loss function choices. For simplicity, we choose the per-token reverse KL — the divergence between the student’s ($\pi_\theta$) and teacher’s ($\pi_\text{teacher}$) distribution for each token conditioned on the same prior trajectory:

$$\text{KL}\Bigl(\pi_\theta \lvert\rvert \pi_\text{teacher}\Bigr) = \mathbb{E}_{x \sim {\pi_\theta}} \Bigl[ \log \pi_\theta(x_{t+1} | x_{1..t}) - \log \pi_\text{teacher}(x_{t+1} | x_{1..t}) \Bigr]$$

Our reward function minimizes the reverse KL, which pushes the student to approximate the teacher’s behavior in every state the student finds itself in. When the student behaves identically to the teacher, reverse KL is zero. For simplicity, we use a discount factor of zero: at any given timestep,Although more mathematically correct, we do not find discount factors > 0 to improve performance in practice, so we choose zero for simplicity. the student only optimizes the immediate next token, with no consideration for future tokens.

Reverse KL has natural synergy with RL, which generally optimizes a form of sequence-level reverse KL induced by the reward model. However, unlike most reward models in practice, the reverse KL is “unhackable” in the sense that low KL always corresponds to a high probability of desirable behavior from the teacher model’s point of view. Two other useful properties of reverse KL are that it is “mode seeking”See Eric Jang’s post for more discussion of mode seeking behaviors. — it learns one specific behavior (the teacher’s) instead of spreading its distribution across several suboptimal options — and it reduces exposure bias.Scheduled Sampling for Sequence Prediction with Recurrent Neural Networks (Bengio et al, 2015). See Gu et al. for more discussion.

This approach offers significant compute savings. Since it doesn’t require a rollout to finish sampling to calculate the reward, we can use shorter or partial rollouts for training. Querying the teacher’s log probabilities also requires just a single forward pass from the larger model, while the trajectories are generated by the smaller and cheaper student.

We do not require a separate reward or labeling model either. There could be advantages to combining distillation-based per-token rewards with sequence-level environment rewards; this is an interesting area for potential future research.

Illustration

Below we see a real example of an incorrect student trajectory graded by a teacher. The example is from SimpleBench, and relies on the model making a key observation that the question’s premise is important: the correct answer is “B. 0” since ice cubes will melt in a frying pan. The student, Qwen3-4B-Instruct-2507, incorrectly treats this as a pure math problem with no consideration of the physical context.

Example trajectory graded by a teacher model. Tokens in darker red correspond to a higher reverse KL.

Darker colors represent tokens that receive a higher penalty from the teacher model, Qwen3-235B-A22B-Instruct-2507, which solves this problem correctly. We see that it penalizes tokens that start phrases which lead the student astray, intuitively corresponding to important “forking tokens” for guiding reasoning.Beyond the 80/20 Rule: High-Entropy Minority Tokens Drive Effective Reinforcement Learning for LLM Reasoning (Wang et al, 2025) The final answer, though wrong, isn’t penalized — it is entirely predictable conditional on the whole preceding sequence.

Pseudocode

We implement on-policy distillation on top of the RL script in Tinker, which already implements sampling, reward computation, and policy gradient-style training.Our implementation can actually be a one-line change on top of RL implementations which use KL regularization: we merely swap out the regularizer model.

Initialize teacher client. The Tinker API enables easily creating different clients for different models, without needing to worry about the utilization of model engines. We use a sampling client, as we do not need to propagate logprobs through the teacher model.
Sample trajectories. We sample rollouts from the student exactly as we would in RL. During sampling, RL already computes the student’s logprobs $\log \pi_\theta(x)$ for use as part of the importance-sampling loss.
Compute reward. We query the teacher client with compute_logprobs on the sampled trajectories, which returns the teacher’s logprobs $\log \pi_\text{teacher}(x)$ on the tokens $x$ sampled by the student.In this post, we do not consider logit (top-k) distillation in any of our experiments, which could be used to further improve compute efficiency. We then use this to calculate the reverse KL.
Train with RL. We set the per-token advantage to the negative reverse KL, and call the RL importance-sampling loss function to perform the training update on the student model.

# Initialize teacher client (main):
teacher_client = service_client.create_sampling_client(
    base_model=teacher_config.base_model,
    model_path=teacher_config.load_checkpoint_path,
)

# Sample trajectories (main):
trajectories = do_group_rollout(student_client, env_group_builder)
sampled_logprobs = trajectories.loss_fn_inputs["logprobs"]

# Compute reward (compute_teacher_reverse_kl):
teacher_logprobs = teacher_client.compute_logprobs(trajectories)
reverse_kl = sampled_logprobs - teacher_logprobs
trajectories["advantages"] = -reverse_kl

# Train with RL (train_step):
training_client.forward_backward(trajectories, loss_fn="importance_sampling")

In the experiments below, we generally apply on-policy distillation to models that have already been mid-trained on specific domain knowledge. This training increases the probability that the student will generate tokens within the teacher’s distribution, though it’s usually far from sufficient for replicating the teacher’s performance.SFT, using forwards KL, adds support for new tokens. Reverse-KL methods can then perform mode seeking within the initialization’s support. Often, as we’ll see in our personalization example, the probability of generating the relevant tokens starts at zero as the student lacks any relevant domain knowledge.

We use on-policy distillation for post-training, and compare it to other approaches to this last crucial stage of training expert models.

Distillation for reasoning

We use distillation to train mathematical reasoning in the Qwen3-8B-Base model, using Qwen3-32B as a teacher model. Both the teacher (Qwen3-32B) and student (Qwen3-8B-Base) are supported models on Tinker today, so you can reproduce our experiments with the Tinker cookbook.Update (June 2026): The teacher (Qwen3-32B) and student (Qwen3-8B-Base) used here have since been retired from the Tinker model lineup. The distillation recipes in the Tinker cookbook have been updated to use Qwen3.5-9B-Base as the student and Qwen3.5-9B as the teacher, so the experiments remain reproducible.

Off-policy distillation

As mentioned above, all our experiments start with mid-training in the form of off-policy distillation — supervised fine-tuning on a dataset of teacher-generated examples. The dataset used for mathematical reasoning is OpenThoughts-3, a collection of reasoning prompts and responses generated by QwQ-32B (a reasoning model similar to Qwen3-32B).

Training the student (Qwen3-8B-Base) on 400k prompts with full fine-tuning achieves a score of 60% on AIME'24, a benchmark of math problems. We can also train with LoRA,LoRA: Low-Rank Adaptation of Large Language Models (Hu et al, 2021) though it lags behind full fine-tuning when training on high-volume datasets. In all cases, we see performance increase log-linearly — the initial performance gains are cheap but the latter are costly.

AIME'24 score over the course of off-policy distillation (SFT). After the initial 50-100K prompts, performance follows a predictable log-linear scaling curve. As predicted in LoRA Without Regret, we observe worse LoRA performance when running large-scale SFT with high batch size.

We can treat the model fine-tuned on 400k prompts as a checkpoint before trying various post-training approaches to increase its performance. We can compare the effort it would take to raise the score on the AIME’24 benchmark from 60% to 70%.

The default approach is to fine-tune on more prompts, continuing the process of off-policy distillation. Extrapolating the log-linear trend, we estimate that the model would achieve 70% on AIME’24 at approximately 2M prompts. This extrapolation requires the scaling law to hold up without stalling, which isn’t trivial. However, there are examples of large-scale off-policy distillation improving an 8B model’s performance past 70%, such as OpenThoughts-3 and DeepSeek-R1-0528-Qwen3-8B.DeepSeek-R1-0528-Qwen3-8B achieves 86% on the benchmark, with the number of training prompts not specified. Older models (Qwen2.5-7B, Qwen2.5-14B) achieve performance of 55.5% and 69.7%, respectively, after training on 800k distillation prompts from DeepSeek-R1. We can use this extrapolation as an optimistic estimate for the cost-performance ratio of off-policy distillation.

Reinforcement learning

The Qwen3 technical report reaches performance of 67.6% on the benchmark using 17,920 GPU hours of RL on top of a similar SFT initialization. It is hard to compare this directly to the cost of distillation, but given some reasonable assumptions about the SFT training stack, this is similar to the cost of training on 2M off-policy distillation prompts.

Method	AIME’24	GPQA-Diamond	GPU Hours
Off-policy distillation	55.0%	55.6%	Unreported
+ Reinforcement learning	67.6%	61.3%	17,920
+ On-policy distillation	74.4%	63.3%	1,800

From Qwen3 Technical Report, Table 21.

The Qwen team also reports reaching a higher score of 74.4 on AIME’24 at one-tenth the cost of RL with on-policy distillation, which served as inspiration for our work. We attempt to replicate it below in our basic setup.

On-policy distillation

As an alternative to off-policy distillation or RL, we run on-policy distillation as described above.We actually use Qwen3-8B as a teacher, as it performs slightly better. However, for the purpose of comparing compute, we can still measure the FLOPs of the 32B model. Starting from the 400k SFT checkpoint, on-policy distillation achieves an AIME’24 of 70% in about 150 steps.150 steps corresponds to about 77K prompts; we train with 4 samples per prompt.

AIME'24 over the course of on-policy distillation. We measure additonal compute in terms of training FLOPs (see below). On-policy distillation is significantly more compute-efficient than SFT, especially for LoRA models. At rank = 32, LoRA trails full finetuning by 13% after SFT, but only 6% after on-policy distillation.

Comparing compute costs across methods is nontrivial, as the ratio of training vs sampling vs log-prob computation cost varies significantly depending on implementation. Below, we compute the cost in terms of FLOPs which penalizes methods that can be effectively parallelized on GPUs. In particular, it overestimates the practical cost of computing log-probs.

Method	AIME’24	Teacher FLOPs	Student FLOPs	CE vs SFT-2M
Initialization: SFT-400K	60%	8.5 × 10²⁰	3.8 × 10²⁰	–
SFT-2M (extrapolated)	~70% (extrapolated)	3.4 × 10²¹	1.5 × 10²¹	1×
Reinforcement learning	68%	-	-	≈1×
On-policy distillation	70%	8.4 × 10¹⁹	8.2 × 10¹⁹	9-30×

We find a baseline cost reduction of 9x when the SFT dataset is given, as in our example with OpenThoughts-3, or is amortized across many training runs.CE = (Student + Teacher) / (Student) In this case we do not count the cost of teacher FLOPs for off-policy training but do for on-policy, since we must run the teacher model to compute log-probs for the student’s trajectory. Since this computation can be cheaply parallelized across GPUs, the cost reduction in GPU hours is closer to 18x.

However, we often want to train a small model on a new task for which no off-policy distillation dataset is available. If we include the full cost of the teacher model in off-policy distillation – ie, including the additional cost of sampling from the teacher model – the total cost reduction is about 30x.CE = (Student + Teacher) / (Student + Teacher)

Distillation for personalization

In addition to training small models to high performance on common tasks, another use-case for distillation is personalization. Examples include adhering to a particular tone in conversation and format of output, or capabilities like tool use and cost budgeting. We often want to train this behavior in combination with new domain knowledge.

Training both at once is generally difficult, and light-weight finetunes are often insufficient for this objective,Unfamiliar Finetuning Examples Control How Language Models Hallucinate (Kang et al, 2024) thereby requiring a larger midtrain. Learning post-training behaviors on top of new knowledge requires a complex post-training stack, often consisting of proprietary data and reward models. While this approach is accessible to frontier labs, it can be difficult or prohibitively expensive for other practitioners to replicate.

In this section, we show that on-policy distillation can be used effectively to post-train specialized behavior. This approach is also applicable to continual learning or “test-time training”: updating models when they are deployed without regressing the base performance. We use an example application of a model mid-trained on our internal company documents.

Training an internal assistant

A common objective for a custom model is to act as an assistant: to possess expert knowledge in some field and, in addition, reliable assistant-like behavior. We may need separate training for each, especially when the domain of expertise can’t be learned from pre-training data only or when learning it interferes with behavior.

Our example is an internal company assistant, for which we have two desiderata:

The model is knowledgeable about the domain (company documents). Pretrained models have seen zero internal documents from the company, and therefore can only guess, regardless of model scale. We will measure this using an internal knowledge recall eval (“internal QA”).
The model exhibits strong post-training behavior, ie. instruction following. We will measure this with the commonly-used IF-eval.Instruction-Following Evaluation for Large Language Models (Zhou et al, 2023)

Training on new knowledge degrades learned behavior

We will start with Qwen3-8B, rather than the base model. Qwen3-8B is post-trained on useful skills for an assistant, such as instruction following and reasoning with RL. Prior research has shown that such reinforcement learning only trains small subnetworks of the original model,Reinforcement Learning Finetunes Small Subnetworks in Large Language Models (Mukherjee et al, 2025) and thus can be fragile when the network is further trained on a large amount of data. We study the extent to which it happens, and how the desired behavior can be recovered.

In order to reduce such catastrophic forgetting, a common approach in mid-training is to mix in “background data” from the original model’s pretraining distribution.Midtraining Bridges Pretraining and Posttraining Distributions (Liu et al, 2025) In this case, we don’t have access to Qwen3’s pretraining distribution. Therefore, we consider a stronger and more expensive baseline: we take Tulu3Tulu 3: Pushing Frontiers in Open Language Model Post-Training (Ivison et al, 2024) prompts – a broad chat and instruction-following dataset – and re-sample them with Qwen3-8B in order to act as chat background data.

This “on-policy” background data sampled by Qwen3-8B acts as a forwards KL regularizer, reinforcing the model’s original behavior throughout mid-training. We find sampling from Qwen3-8B is better than Qwen3-32B for preserving chat capabilities throughout mid-training, highlighting the sensitivity of the data source; similar on-policy SFT results have been found in Chen et al.Retaining by Doing: The Role of On-Policy Data in Mitigating Forgetting (Chen et al, 2025) We hypothesize this approach can be even more effective than having access to the original pretraining data distribution, at the cost of having to sample a large-scale dataset.

We then fine-tune Qwen3-8B on different mixes of our internal documents and chat data. Increasing the proportion of document data straightforwardly improves the model’s knowledge. However, although mixing in at least 30% of chat data helps preserve most instruction-following ability, there is no weighting which maintains the original performance on IF-eval.This is true even if the SFT dataset contains 100% chat data. We address this further in the discussion on continual learning.

Sweeping over the ratio of internal documents: background chat data during mid-training. Although mixing in a little chat data helps to prevent a catastrophic regression, no weight maintains the original IF-eval performance.

For any given mix, we observe IF-eval performance degrades during fine-tuning. This compromises our ability to use longer training to help specialize the model further.Directionally, we may expect an overparametrized model trained on some dataset to update its behavior only the context of the dataset without impacting what it does in other contexts. However, we do not observe this in practice, as training on raw document data regresses performance even in the QA context.

IF-eval decreases during midtraining across all mixes of data. When we use a linear learning rate (pictured), the degradation eventually flattens and slowly begins to recover as the learning rate decays. However, performance never fully recovers.

An alternative commonly-used approach is to use LoRA in order to constrain the parameter update, thereby reducing the possibility of catastrophic forgetting. However, this approach is still insufficient for preserving IF-eval, and the LoRA learns less.LoRA Learns Less and Forgets Less (Biderman et al, 2024)

When used for our personalization midtrain on top of the post-trained Qwen3-8B, LoRA learns less (knowledge) and still forgets its original post-training behaviors.

On-policy distillation recovers post-training behavior

Next, we seek to restore instruction-following behavior after fine-tuning on internal documents. This behavior was originally trained with RL which is expensive and, as we have seen, fragile. Instead, we run on-policy distillation with the earlier version of the model, Qwen3-8B, as the teacher on Tulu3 prompts. Note that this phase of training has no relation to the internal document data, and is designed solely to recover instruction following.

The use of an earlier version of the model as a teacher to “re-invoke” capabilities lost during fine-tuning makes on-policy distillation very promising for continuous learning. We could alternate between phases of fine-tuning on new data and distillation to recover behavior to allow our model to learn and stay up-to-date on knowledge over time. This phase-alternating approach has previously been explored by Cobbe et al.Phasic Policy Gradient (Cobbe et al, 2020)

After fine-tuning on an 70-30 mix of internal document data and chat data, on-policy distillation recovers nearly full performance on IF-eval without losing any knowledge; we also observe some positive transfer between chat capabilities and the model’s “knowledge” performance on the internal QA eval.

Model	Internal QA Eval (Knowledge)	IF-eval (Chat)
Qwen3-8B	18%	85%
+ midtrain (100%)	43%	45%
+ midtrain (70%)	36%	79%
+ midtrain (70%) + distill	41%	83%

Domain-specific (Internal QA eval) and chat (IF-eval) performance after mid-training. Although mid-training forgets the post-trained behaviors of Qwen3-8B, they are cheaply restored via on-policy distillation, alongside the additional knowledge learned via the mid-train.

In essence, we have treated the language model itself as a reward model, with high-probability behaviors being rewarded.Direct Preference Optimization: Your Language Model is Secretly a Reward Model (Rafailov et al, 2023) This has connections to inverse RL: high-probability behaviors correspond to advantageous rewards in an assumed underlying preference model.Algorithms for Inverse Reinforcement Learning (Ng and Russell, 2000) Any instruction-tuned open-weight model can be used as a reward model in this sense; we merely need access to the compute_logprobs function.

Distillation as a tool for integrating behaviors and knowledge has also been explored for hybrid reasoning models (Qwen3) and specialist distillation.DeepSeek-V3.2-Exp: Boosting Long-Context Efficiency with DeepSeek Sparse Attention (DeepSeek-AI Team, 2025) As our and Chen et al.’s results suggest, on-policy learning can be a crucial tool for augmenting similar distillation-based “model merging” setups.

Discussion

Dense supervision greatly improves compute efficiency

Reinforcement learning and on-policy distillation both learn via reverse KL, pruning the space of actions present in the base policy. The difference is in the density of reward. In LoRA Without Regret we presented the information-theoretic perspective that reinforcement learning only teaches $O(1)$ bits per episode. In contrast, distillation teaches $O(N)$ bits per episode, where $N$ is the number of tokens. Can we quantify the training efficiency gained through denser reward?

We ran an experiment for making a direct comparison between the two:

Start with Qwen3-8B-Base (no additional SFT).
Run RL on DeepMath, matching our procedure from LoRA Without Regret. We use a LoRA rank of 128. The resultant model is the teacher for distillation.
On-policy distill from the RL-trained model (2) back into the base model (1).

Starting from the same initialization, on-policy distillation can learns the RL-trained policy in approximately 7-10x fewer gradient steps, which corresponds to a compute efficiency of 50-100x.

We see that distillation reaches the teacher’s level of performance approximately 7-10x faster than RL with matched model architecture (LoRA rank 128). The reverse KL decreases to near-zero and the AIME score is recovered in under 10 gradient steps, while RL took 70 steps to reach that level.

Cumulatively, the reduction in compute required is on the order of 50-100x:

While RL requires training at approximately the evaluation context (in order for the policy to learn the context limit and not incur format penalty), distillation learns reasonably at shorter context lengths, as there is no sharp cutoff in reward between a trajectory which has finished sampling and a trajectory that continues.
When the SFT initialization is strong,ie: when the teacher policy is within the support of the student policy. When this is not true, as in “Distillation for reasoning”, we require a significantly larger batch size. on-policy distillation works effectively with much smaller batch sizes since it provides significantly more bits per episode, thereby reducing gradient noise.

Although it is generally difficult to train reinforcement learning models with process supervision, these results indicate that as a broad direction, process supervision and dense rewards have the potential to improve learning efficiency by an order of magnitude. This matches earlier results in RL research from Lightman et al.

Distillation can effectively reuse training data for data efficiency

For practitioners, collecting large datasets of training prompts can be difficult and time-consuming. Therefore, we want to be able to reuse prompts several times in training. With RL, training multiple epochs on the same prompt often leads to simple memorization of the final answer, especially with large models.“Reinforcement Learning for Reasoning in Large Language Models with One Training Example” (Wang et al, 2025) presents a positive result, however, in some settings. In contrast, on-policy distillation learns to approximate the teacher’s complete distribution by minimizing reverse KL, rather than memorizing a single answer. This allows us to train many samples from the same prompt.

We repeat the above experiment of training Qwen3-8B-Base on math, but now with only a single randomly chosen prompt from the dataset.Prompt: “Evaluate the limit: $\lim_{x \to \infty} \sqrt{x} \left( \sqrt[3]{x+1} - \sqrt[3]{x-1} \right)$"

We train on this prompt for 20 consecutive steps, each with a batch of 256 rollouts, for 5120 graded sequences in total. We train on the same prompt for multiple steps in sequential fashion, which normally leads to overfitting. Though this is naturally less compute-efficient, we do approximately match the performance of the teacher model despite only training on a single prompt.

In this example, multi-epoch training on top of one training example is sufficient to distill the teacher's AIME'24 performance. Our default configuration (also used in the personalization experiments) runs on-policy distillation with 64 prompts / batch, and 4 samples / prompt. All methods shown are trained with 256 samples / batch. Note that the right chart is showing training KL, hence it is natural for 1 prompt total to outperform 1 prompt / batch.

RL searches in the space of semantic strategies

We have seen that on-policy distillation can replicate the learning provided by RL with much fewer steps of training. One interpretation of this result is that, unlike pre-training, RL doesn’t spend a lot of compute on the gradient steps themselves. We should think of RL as spending most of its compute on search — rolling out a policy and assigning credit — rather than on making updates.From The Bitter Lesson (Rich Sutton): “breakthrough progress eventually arrives by an opposing approach based on scaling computation by search and learning"

Pre-training via stochastic gradient descent is exploring the high-dimensional parameter space. Pre-training requires a vast amount of information and is very difficult to distill, in part because the parameter space is somewhat unique to each network.“The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks” (Frankle and Carbin, 2018) The gradient steps required for pretraining are extremely computationally expensive and time-consuming.

In contrast, we should think of RL as exploring the space of semantic strategies.Note that exploration over strategies is subtlely different from exploration over outcomes; RL requires the base model to have nonzero success to start with, and therefore has already “found the outcome”, but the strategy can be refined over the course of RL to make successful outcomes more likely. At every step, RL tries a small modification of some strategy it has found in the past. Rather than exploring in the parameter space, it “stumbles” onto new strategies by luck — it is randomly sampling from the set of weights it already has.

Once a good strategy is found, distillation serves as a shortcut for learning it: on-policy distillation does not need to model the intermediate strategies during the curriculum of RL, but rather only the final strategy learned. If we are only interested in the final strategy (common in production use-cases), we need not spend the compute to model all the intermediate ones.

Consider an analogy: in science research, we spend a lot of time and resources looking for answers and exploring new ideas. Once a result is discovered, it is much simpler to teach it to others by expressing it in natural language. We can contrast this to intuitive physical skills such as playing a sport. They are much harder to teach to others, since the knowledge exists in an innate language (e.g., muscle memory) that is only readily understood by ourselves. Sports are only learned with repeated practice.

On-policy learning as a tool for continual learning

In the section on distillation for personalization, we explored the ability of on-policy distillation to re-introduce specialized trained behaviors into the model. This generalizes to a broader set of continual learning tasks, which require acquiring new knowledge without degrading prior capabilities.

Prior work has found on-policy learning (RL) forgets less than off-policy learning.“RL’s Razor: Why Online Reinforcement Learning Forgets Less” (Shenfeld et al, 2025) However, RL only shapes behavior — it cannot teach new knowledge well, and thus can’t be sufficient for continual learning.

In the above section, we saw that SFT (including off-policy distillation) fails at scaffolding continual learning because it degrades behavior. We investigate this more closely and demonstrate this with a direct example. Similarly to above, we build a dataset by taking Tulu3 prompts and sampling from Qwen3-32B at temperature = 1.0 and no further modifications. This dataset therefore has a KL of zero against Qwen3-32B.The importance of “truly on-policy” KL=0 data was also explored in our previous post, Defeating Nondeterminism in LLM Inference.

What happens when we run SFT on this dataset of a model’s own samples? We see that any practical learning rate that’s greater than zero leads to a degradation of performance on the instruction-following evaluation!

Running SFT on top of Qwen3-32B's own samples degrades performance. We use the same learning rate as from the personalization section, which was swept for practical performance considerations. A linear learning rate can prevent forwards KL / IF-eval from regressing indefinitely, but does not recover performance before the LR decays to zero.

A possible explanation for this is that while KL divergence is 0 in expectation, every finite batch will exhibit a slightly different distribution in practice. Training on these finite batches causes a non-zero gradient update, which then diverges the updated model’s policy from that of its original state. This process turns training on one’s own samples into off-policy training over time, which causes the same error accumulation and divergence over long sequences that happens with off-policy training.

On-policy distillation always stays on-policy, and since the teacher stays fixed, the student converges on the teacher’s desirable behavior, without regressing in the self-distillation setting as SFT does. This makes on-policy distillation a very promising tool for continual learning.

Conclusion

We have explored the application of on-policy distillation for applications such as training a small model for math reasoning or a continuously-learning assistant. We compared on-policy distillation to two other approaches to post-training: off-policy distillation, and on-policy RL. We find that on-policy distillation combines the best of both worlds: the reliable performance of on-policy training, with the cost-efficiency of a dense reward signal.

Post-training is a crucial part of reaching frontier model capabilities. By leveraging on-policy sampling from the student with dense supervision from a teacher, the on-policy distillation recipe reaches those capabilities at a fraction of the cost of frontier high-compute RL runs.

Our implementation can be found in the Tinker cookbook. Our work explored simple and straightforward instantiations of on-policy distillation to clearly showcase its advantages. We hope to continue research on new applications of distillation, new methods for improving teacher supervision, and ways to improve data efficiency and continual learning.

At Thinking Machines, our mission is to empower people with AI models that combine frontier performance with adaptability and personalization. On-policy distillation is a potent tool for achieving that.

Citation

Please cite this work as:

Lu, Kevin and Thinking Machines Lab, "On-Policy Distillation",
Thinking Machines Lab: Connectionism, Oct 2025.

Or use the BibTeX citation:

@article{lu2025onpolicydistillation,
  author = {Kevin Lu and Thinking Machines Lab},
  title = {On-Policy Distillation},
  journal = {Thinking Machines Lab: Connectionism},
  year = {2025},
  note = {https://thinkingmachines.ai/blog/on-policy-distillation},
  doi = {10.64434/tml.20251026},
}

LoRA Without Regret

Mon, 29 Sep 2025 00:00:00 +0000

Today’s leading language models contain upwards of a trillion parameters, pretrained on tens of trillions of tokens. Base model performance keeps improving with scale, as these trillions are necessary for learning and representing all the patterns in written-down human knowledge.

In contrast, post-training involves smaller datasets and generally focuses on narrower domains of knowledge and ranges of behavior. It seems wasteful to use a terabit of weights to represent updates from a gigabit or megabit of training data. This intuition has motivated parameter efficient fine-tuning (PEFT), which adjusts a large network by updating a much smaller set of parameters.

The leading PEFT method is low-rank adaptation, or LoRA. LoRA replaces each weight matrix W from the original model with a modified version $W' = W + \gamma BA$, where B and A are matrices that together have far fewer parameters than W, and $\gamma$ is a constant scaling factor. In effect, LoRA creates a low-dimensional representation of the updates imparted by fine-tuning.

LoRA may offer advantages in the cost and speed of post-training, and there are also a few operational reasons to prefer it to full fine-tuning (henceforth, FullFT):

Multi-tenant serving. Since LoRA trains an adapter (i.e., the A and B matrices) while keeping the original weights unchanged, a single inference server can keep many adapters (different model versions) in memory and sample from them simultaneously in a batched way.Punica: Multi-Tenant LoRA Serving (Chen, Ye, et al, 2023) Modern inference engines such as vLLM and SGLang implement this feature.
Layout size for training. When fine-tuning the whole model, the optimizer state needs to be stored along with the original weights, often at higher precision. As a result, FullFT usually requires an order of magnitude more accelerators than sampling from the same model does, and thus a different layout.For training, besides storing the weights, we typically need to store gradients and optimizer moments for all of the weights; moreover, these variables are often stored in higher precision (float32) than what’s used to store the weights for inference (bfloat16 or lower). Since LoRA trains far fewer weights and uses far less memory, it can be trained on a layout only slightly larger than what is used for sampling. This makes training more accessible, and often more efficient.
Ease of loading and transfer. With fewer weights to store, LoRA adapters are fast and easy to set up or transfer between machines.

These reasons are sufficient to explain the growing popularity of LoRA since the publication of the original LoRA paper in 2021.LoRA: Low-Rank Adaptation of Large Language Models (Hu et al, 2021) However, the literature is unclear on how well LoRA performs relative to FullFT.

There is agreement that LoRA underperforms in settings that resemble pre-training,LoRA Learns Less and Forgets Less (Biderman et al, 2024) namely those with very large datasets that exceed the storage limits of LoRA parameters. But for dataset sizes that are typical in post-training, LoRA has sufficient capacity to store the essential information. However, this fact makes no guarantees regarding sample efficiency and compute efficiency. The question is: can LoRA match the performance of full fine-tuning, and if so, under which conditions?

In our experiments, we find that indeed, when we get a few key details right, LoRA learns with the same sample efficiency as FullFT and achieves the same ultimate performance.

What matters for LoRA

This article covers a series of supervised fine-tuning and reinforcement learning experiments we conducted to determine the conditions under which LoRA matches FullFT efficiency. To this end, we did a few things differently from previous experiments on LoRA:

We investigated the general relationship between training set size and number of LoRA parameters, rather than focusing on specific datasets and tasks.
In supervised learning, we measured log loss rather than employing sampling-based evals, with the same goal of generality in mind. Log loss measurement gives clean results and scaling laws over ranges of training steps and training parameters.

We find that:

For supervised fine-tuning on small-to-medium-sized instruction-tuning and reasoning datasets, LoRA performs the same as full fine-tuning.
For datasets that exceed LoRA capacity, LoRA underperforms FullFT. Rather than the loss reaching a distinct floor that it can’t go below, LoRA results in worse training efficiency that depends on the relationship between model capacity to dataset size.
In some scenarios, LoRA is less tolerant of large batch sizes than full fine-tuning — it pays a larger penalty in loss as batch size increases beyond some point. This penalty is not mitigated by increasing the LoRA rank; it is a property of the product-of-matrices parametrization, which has different training dynamics than optimizing the original weight matrix.
Even in small data settings, LoRA performs better when applied to all weight matrices, especially MLP and MoE layers. Attention-only LoRA underperforms even when we match the number of trainable parameters by using higher rank for attention-only LoRA.
LoRA performs equivalently to FullFT for reinforcement learning even with small ranks. We find that RL requires very low capacity, a result we anticipated based on information-theoretical arguments.

We also studied the impact of hyperparameters used for LoRA on its learning rate relative to full fine-tuning. We examine some invariances in hyperparameters like init scales and multipliers, and explain why the 1/r prefactor makes the optimal learning rate (LR) approximately independent of rank. We also show experimentally how the optimal LR for LoRA relates to the optimal LR for FullFT.

The outcome of our experiments is the characterization of a “low-regret regime” where LoRA performs similarly to FullFT in terms of dataset size and LoRA parameters. We found this regime covers most post-training scenarios, opening the door to the use of efficient fine-tuning in many applications.

Methods and results

We designed our experiments to measure in detail the relative performance of LoRA compared to FullFT across a range of conditions. Here are some details of our experimental setup:

We varied the LoRA rank over three orders of magnitude, with rank between 1 and 512, and compared these to full fine-tuning.
To eliminate potential confounds from using a suboptimal learning rate, we swept the LR for each experimental condition. We used constant learning rate schedule (no warmup or cooldown).
Our experiments used Llama 3 series modelsThe Llama 3 Herd of Models (Dubey et al, 2024) and Qwen3 modelsQwen3 Technical Report (Qwen Team, 2025), including a mixture of experts (MoE) model.
The main supervised learning experiments used the Tulu3Tulu 3: Pushing Frontiers in Open Language Model Post-Training (Ivison et al, 2024) and OpenThoughts3OpenThoughts: Data Recipes for Reasoning Models (Guha et al, 2025) datasets, focused on instruction following and reasoning, respectively. The two sets differ significantly in scope, structure, and application, supporting the generality of our results.
Our RL experiments used mathematical reasoning tasks with answer correctness as the reward.

LoRA rank

We trained for a single epoch on the Tulu3 dataset and a subset of the OpenThoughts3 datasets. For each dataset and model size, we swept over LoRA rank and learning rate. In the plots below, we draw one colored line for each rank, where the line is obtained by taking the pointwise minimum over all learning rates at each training step:

LoRA training curves for various ranks on Tulu3 and OpenThoughts3 datasets. FullFT and high-rank LoRAs have similar learning curves with loss decreasing linearly with the logarithm of steps. Lower-rank LoRAs fall off the minimum-loss curve when the adapter runs out of capacity. In the bottom plots (1B model) high-rank LoRA performs better than FullFT on one dataset and worse on the other. There might be some random variation in how LoRA performs on different datasets, due to differences in training dynamics or generalization behavior.

We see that FullFT and high-rank LoRAs have similar learning curves with loss decreasing linearly with the logarithm of the number of steps. Medium and low-rank LoRAs fall off the minimum-loss learning curves at some threshold of steps that correlates with rank. Intuitively, learning slows down when the adapter runs out of capacity, which in turn is determined by rank.

Next, we plot how loss changes with LR to check that our sweep covers the best learning rate for each rank.

Learning rate versus final loss for various LoRA ranks on Tulu3. Minimum loss is approximately the same for high rank LoRA and FullFT. Optimal LR is 10 times higher for LoRA.

We find that the optimal learning rate for FullFT is lower by a factor of 10 than for high-rank LoRAs.See Biderman et al. (2024), Figure S1, for an experiment with sampling evals, which finds a similar 10x ratio. We’ll return to this in our discussion of LoRA hyperparameters later on.

The optimal LR seems to be similar for all the LoRA runs across different ranks; we give a theoretical explanation for this finding below. However, there does seem to be some rank dependence, with lower optimal LR for rank=1 than for higher-rank LoRAs. The optimal LR changes by a factor of less than 2 between rank=4 and rank=512.

Batch size effects

We found that in some settings, LoRA is less tolerant of large batch sizes than FullFT. The performance gap grows with larger batch sizes, independent of rank. For this next experiment, we used a small 10,000-example subset of OpenThoughts3.

Batch size effects on LoRA vs FullFT performance. Left: Learning curves for different batch sizes show a persistent gap between LoRA (dashed) and FullFT (solid) at large batch sizes. Right: Final loss as a function of batch size shows LoRA pays a larger penalty for increased batch size.

The left-hand plot in Figure 3 shows a persistent gap between the LoRA (dashed lines) and FullFT (solid line) learning curves at large batch sizes. The gap is smaller and shrinks over time for the smaller batch size of 32.

The right-hand chart plots final loss as a function of batch size. We see the gap in loss for LoRA increasingly diverging from FullFT for larger batch sizes.

The learning gap at large batches doesn’t seem to depend on rank, but rather seems to be a property of LoRA. The likely reason is that the product-of-matrices parametrization (BA) has less favorable optimization dynamics on this dataset than the full matrix (W). However, both LoRA and FullFT achieve their best loss at smaller batch sizes, so this gap may not matter as much in practice.

Layers Where LoRA Is Applied

We investigated the effects of applying LoRA to different layers in the network. The original paper by Hu et al. recommended applying LoRA only to the attention matrices, and many subsequent papers followed suit, though a recent trend has been to apply it to all layers.Similar to our results, the QLoRA paper also found that LoRA performed worse than MLP or MLP+attention, though they found that MLP+attention > MLP > attention, whereas we found the first two to be roughly equal. Indeed, we achieved far better results when applying LoRA to all layers, in particular, the MLP (including MoE) layers. In fact, applying LoRA to the attention matrices shows no additional benefits beyond applying it to the MLPs only.Biderman et al. (2024) obtained a similar result, with attention-only LoRA providing no additional benefit on top of MLP-only.

Attention-only LoRA significantly underperforms MLP-only LoRA, and does not further improve performance on top of LoRA-on-MLP. This effect holds for a dense model (Llama-3.1-8B) and a sparse MoE (Qwen3-30B-A3B-Base).

The underperformance of attention-only LoRA is not explained by having fewer parameters. In this particular case, attention-only with rank 256 underperforms MLP-only with rank 128, despite them having approximately the same number of parameters. (Compare the bold numbers in the table below.)

LoRA configuration	Params
mlp, rank=256	0.49B
attn, rank=256	0.25B
all, rank=256	0.70B
mlp, rank=128	0.24B

Parameter counts for LoRA on Llama-3.1-8B

For the MoE experiment, we trained a separate LoRA on each expert, with the rank of each equal to the total rank divided by the number of active experts (equal to 8 for Qwen3 MoE). This scaling keeps the ratio of LoRA parameters to FullFT parameters the same for MoE layers as for other layers.

We did similar experiments comparing different LoRA layers in two additional settings: (1) supervised learning on a small subset of the OpenThoughts3 dataset with rank=256, and (2) reinforcement learning on the MATH dataset. We describe our experimental setup in the following section. Attention-only LoRA underperforms MLP-only LoRA (which performs similarly to MLP+attention) in these settings as well.

Learning rate vs final loss or reward, when varying which layers we apply LoRA to.

Reinforcement learning

A key finding from our experiments is that LoRA fully matches the learning performance of FullFT when running policy gradient algorithms for reinforcement learning, even with ranks as low as 1.

For these experiments, we used a basic policy gradient algorithm with an importance sampling correction; $\text{objective} =\sum_t \frac{p_{\text{learner}}}{p_{\text{sampler}}} Adv_t$.See Your Efficient RL Framework Secretly Brings You Off-Policy RL Training We used a GRPO-like centering schemeDeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models (Shao et al, 2024) where we sample multiple completions per problem and subtract the mean reward per group.

Figure 6 (below) shows LR sweeps on the MATHMeasuring Mathematical Problem Solving With the MATH Dataset (Hendrycks et al, 2021) and GSMGSM8K: Training Verifiers to Solve Math Word Problems (Cobbe et al, 2021) datasets, using typical hyperparameters for each. We used the Llama-3.1-8B base model as Qwen2.5 and Qwen3 are known to have been pretrained on data that improves their math performance, as described by the Qwen tech reportsQwen2.5 Technical Report (Qwen Team, 2024), which makes it harder to measure what is being learned only during RL.

LoRA shows a wider range of performant learning rates and arrives at the same peak performance as FullFT (black line), at least within the precision limits afforded by the noisiness of RL.

Learning rate vs final reward (accuracy) when doing RL on grade school math (GSM, left) or MATH (right) dataset.

This result is anticipated by an information-theoretic argument. Supervised learning arguably provides O(number of tokens) bits per episode. In contrast, in policy gradient methods, learning is driven by the advantage function which provides only O(1) bits per episode. When each episode contains thousands of tokens, RL absorbs ~1000 times less information per token in training than supervised learning does.

We can use more precise numbers based on our experiments. In the MATH example, we trained on ~10,000 problems with 32 samples per problem. Assuming each completion yields a single bit of information, the whole training process only needs to absorb 320,000 bits. Rank-1 LoRA for Llama-3.1-8B already has 3M parametersWe calculated this by adding up rank·$d_{in}$ (for matrix A) and rank·$d_{out}$ (for B) over all weight matrices in the model., almost 10 times that number. Even at rank-1, LoRA has more than enough capacity to absorb all the information provided during training.

As another point of comparison, DeepSeek-R1-Zero was trained on 5.3M episodesTraining took place for 10,400 steps, each step consisting of 32 unique questions, each question sampled 16 times., corresponding to 5.3M bits of information. This is less than the number of parameters in a low-rank LoRA, and we predict that the results can be replicated with LoRA.

For additional validation of our findings of LoRA’s effectiveness in reasoning RL, we carried out larger-scale experiments with Qwen3-8b-base on the DeepMath datasetDeepMath-103K: A Large-Scale, Challenging, Decontaminated, and Verifiable Mathematical Dataset for Advancing Reasoning (He et al, 2025) as it is much larger than the MATH dataset and in general contains harder problems. To speed up experiments, we restricted the samples to a length of 8192 tokens for training and evaluation. This sample length allows for backtracking and reasoning but limits the performance, relative to longer chain-of-thought.

Experiments on the DeepMath dataset with Qwen3-8b-base. In the left plot, we show the learning curve for different ranks and full fine-tuning. For each of these settings, we show the best learning rate, which results in the highest final performance. On the right, we plot learning rate vs final performance. As in our previous math experiments, LoRA seems to have a wider peak of near-optimal learning rates.

Additional plots from experiments on the DeepMath dataset with Qwen3-8b-Base. The left plot shows the benchmark scores on the AIME test set, which is more challenging than the training set. The right plot shows the chain-of-thought (CoT) length over training steps, which can be seen as a sign of learning to reason.

We observe that when picking the optimal learning rates for each setting, training progresses in an almost identical way for LoRAs with different sizes and full fine-tuning. Moreover, we see similar findings when we evaluate the models on the held-out problems of AIME 2024 and AIME 2025. Furthermore, we observe similar qualitative behavior from the LoRA and full-finetuning runs: both develop advanced reasoning behaviors such as backtracking, self-verification and in-context exploration, which is visible in the lengthening of the model CoTs.

Setting LoRA hyperparameters

One barrier to LoRA adoption is the necessity to choose optimal hyperparameters, which are different from ones optimized for FullFT. In this section, we show that this problem isn’t as daunting as it appears at first glance and discuss our findings related to hyperparameter choice.

Optimal learning rate and rank

Following Hu et al., we consider the following parametrization for LoRA:

$$W' = W + \frac{\alpha}{r}BA$$

Where $r$ is the LoRA rank, $\alpha$ is the LoRA scaling factor, and $A$, $B$ are the LoRA weight matrices (of rank $r$). We use $\alpha = 32$ for the experiments in this article, following standard practice from other implementations.

The $1/r$ scaling factor makes the optimal learning rate approximately independent of rank. In fact, a stronger condition holds – the learning curve is exactly the same at the beginning of training, regardless of rank. This effect is striking, and in our experiments the closeness of the learning curves for different ranks had us worried that a bug caused the rank parameter to be ignored. It follows that in a short training regime, the optimal LR is also independent of rank. However, as we showed above in our plots of learning rate vs loss (Figure 2), optimal LR has some rank-dependence in the longer-training regime.

These plots look at the differences in the learning curves, early in training, for different ranks with the same learning rate. On the left, we show the learning curves. The right shows the difference between rank 16 and 256, which grows over time. Strangely, it is negative (though tiny) for the first few steps, so that part of the curve is missing from the plot.

We can partly explain this result by looking at the expected update to the LoRA matrix after the very first training update. We can think of the LoRA product $BA$ as the sum of $r$ rank-1 outer products: $BA = \sum_{i=1}^r b_i a_i^T = \sum_{i=1}^r \Delta_i$, where we define $\Delta_i = b_i a_i^T$. Here, $\partial \text{Loss}/\partial \Delta_i$ is the same for all $i$; however the gradients $\partial \text{Loss}/\partial b_i$ and $\partial \text{Loss}/\partial a_i$ will depend on the initialization ($\partial \text{Loss}/\partial b_i$ depends on $a_i$, for example). Since the initialization of $a_i$ and $b_i$ do not depend on rank, it follows that $\mathbb{E}[\Delta_i]$ is the same for all $i$ and does not depend on rank. At the first step of training, the expected update from each of these terms is equal and independent of the rank. It follows that $(1/r)\sum_{i=1}^r \Delta_i$ is just a sample average of $r$ terms with the same expectation, so the expectation of the average, i.e., the change to the adapter $(1/r)BA$, doesn’t depend on the rank.

Parametrization invariances

There are four hyperparameters potentially applicable to LoRA:

The scale factor $α$ which appears in $α / r$.
The learning rate for the down-projection matrix $A$, $LR_A$
The learning rate for the up-projection matrix $B$, $LR_B$.
The initialization scale of matrix $A$, $\text{init}_A$. For a random initialization, this is the standard deviation of $A$’s initial elements. Matrix $B$ is initialized to zero, so there is no need to define $\text{init}_B$.

Having to tune four different parameters may seem overwhelming. However, invariances in the training dynamics mean that two of these are redundant, and learning behavior is determined by two. We show this invariance by noting that when training with Adam and $ε = 0$, theWe can extend this result to $ε > 0$; we need to scale it by $1/q$ because the gradients are scaled by that factor. optimization process is invariant to the following two-parameter transformation.

For $p, q > 0$:

$α → \frac{1}{pq} \cdot α$
$\text{init}_A → p \cdot \text{init}_A$
$LR_A → p \cdot LR_A$
$LR_B → q \cdot LR_B$

Since two degrees of freedom out of the four don’t affect the learning process, we are left with a 2D parameter space. We can choose different bases for this 2D space, such as the following one which lends itself to a straightforward interpretation:

$α \cdot \text{init}_A \cdot LR_B$. This determines the scale of initial updates, or, equivalently, the initial slope of the learning curve. Since $B$ is initialized to zero, $LR_A$ and the initial updates to $A$ are irrelevant.
$\text{init}_A / LR_A$. Since Adam updates the elements of $A$ by approximately $LR_A$ at each step, this timescale parameter determines the number of steps it takes to significantly transform $A$ away from its initial state.

We can reinterpret some proposals from previous work on LoRA in terms of this basis.

LoRA+LoRA+: Efficient Low Rank Adaptation of Large Models (Hayou et al, 2024) proposes to use different LRs on $A$ and $B$, with a higher rate for $B$. Expressed in terms of our basis above, increasing $LR_B$ is equivalent to increasing $\text{init}_A/LR_A$ so that $A$ changes on a longer timescale.

Unsloth’s LoRA Hyperparameter Guide recommends using higher values of $α$ for high-rank LoRA, e.g. by avoiding the $1/r$ scaling. This is also equivalent to increasing $\text{init}_A/LR_A$. When we increase $α$, $LR_A$ and $LR_B$ need to be lowered in compensation to get the same update size. This in turn simply makes $LR_A$ smaller relative to $\text{init}_A$.

In our experiments, we used the standard parametrization used in the Huggingface peft libraryPEFT: State-of-the-art Parameter-Efficient Fine-Tuning methods (Mangrulkar et al, 2022) proposed by Hu et al: a uniform distribution for $A$ with scale $1/\sqrt{d_{in}}$, zero initialization for $B$, the same LR for both, and $α = 32$. We were unable to improve on these hyperparameters in our experimentation.

Optimal learning rates for LoRA vs. FullFT

Our experiments showed that the optimal LR for LoRA is consistently 10x the one used for FullFT in the same application, for both supervised learning and reinforcement learning. This shows up in every U-shaped plot of performance (loss or reward) charted against learning rate. This observation should make it more straightforward to transfer learning hyperparameters from FullFT to LoRA.

We don’t yet have an adequate theoretical explanation for this observation. We can attempt to derive this result from the facts that optimal LoRA LR is invariant to rank and that full-rank LoRA is directly comparable to FullFT. This analysis suggests a LR ratio of the model’s hidden size divided by $2 \cdot \alpha$, which doesn’t match the empirical result of the optimal ratio being fixed at 10 independent of the base model.

For our empirical analysis, we conducted an LR sweep of 14 different Llama and Qwen models for both LoRA and FullFT on the Tulu3 dataset. From those sweeps, we fit a function that predicts the optimal learning rate based on the model’s hidden size and an indicator of whether it’s Llama or Qwen. The functional form used was:

$$\text{LR} = M_{\text{LoRA}} \cdot \left(\frac{2000}{\text{hidden size}}\right)^{\text{model pow} + \text{LoRA pow}}$$

Where:

$M_{\text{LoRA}}$ is a multiplier applied when LoRA is used (1 if FullFT)
$\text{model pow}$ is an exponent adjustment, calculated separately for each model source (Llama and Qwen)
$\text{LoRA pow}$ is an additional exponent adjustment for LoRA
$\text{hidden size}$ is the dimension of the residual stream of the model.

We scored a predicted learning rate by using linear interpolation to predict the loss, based on the data from our sweep, and rated the parameters by summing the predicted loss over the 14 problems. Our optimization found a multiplier of 9.8 for LoRA over FullFT, and different dependence on hidden_size for Qwen3 and Llama models, but LoRA LRs had the same dependence on hidden_size as FullFT LRs, i.e., the optimization found $\text{LoRA pow} = 0$.

Learning rates in short and long runs

The typical initialization of LoRA creates an implicit schedule of change in the effective learning rate. This leads to differences between short and long training runs, and some differences in the shape of learning curves compared to FullFT.

At the start of training, $B$ is initialized to zero. While $B$ is very small, changes in $A$ have negligible effects on the adapter $BA$ which is added to the original network weights. As $B$ grows larger, updates to $A$ start to have a bigger impact on the network outputs, with the effective learning rate increasing over the course of training as $B$ approaches $A$ in scale. We found that by the end of the full training runs on the Tulu3 and OpenThoughts datasets, the $B$ matrices ended up with larger spectral norms than the $A$ matrices.

This implies that the optimal LR should be set higher for shorter training runs. Preliminary evidence suggests an optimal multiplier around 15x over the FullFT for short runsBased on anecdotal evidence, the higher multiplier is effective under ~100 steps or so., converging to the aforementioned 10x multiplier for longer runs.

Discussion

We want to move beyond our empirical results to discuss some broader considerations related to LoRA performance and applicability that would be of interest to both researchers and builders.

First, let us examine in more depth our main result, namely the two conditions under which LoRA performs similarly to full fine-tuning:

LoRA is applied to all layers of the network, especially the MLP/MoE layers which house most of the parameters.
LoRA works well when not capacity constrained, i.e., the number of trainable parameters exceeds the amount of information to be learned, which can be estimated in terms of dataset size.

When (1) is satisfied, we get similar learning dynamics to FullFT at the very start of training. Then, as per (2), LoRA continues to look like FullFT until we start reaching capacity limits.

Why LoRA might be needed on all layers

As we showed earlier, if we put LoRA on only the attention layers, we get slower learning even in the tiny-data regime.

One possible explanation could come from thinking about the empirical neural tangent kernel (eNTK) as an approximation of what happens when we do a small amount of fine-tuning, following Malladi et al.A Kernel-Based View of Language Model Fine-Tuning (Malladi et al, 2022) eNTK is based on the dot products of gradients, specifically gradients $g_i = \partial/\partial\theta \log p(\text{token}_i | \text{prefix}_i)$, and $K(i, j) = g_i \cdot g_j$. As a consequence, the layers with the most parameters will typically have the most influence on the kernel. The paper also points out the eNTK for LoRA is approximately the same as that for full fine-tuning, when you train all the layers. So LoRA training $\approx$ eNTK(LoRA) $\approx$ eNTK(FullFT) $\approx$ FullFT. The approximation eNTK(LoRA) $\approx$ eNTK(FullFT) only holds when we apply LoRA to the layers that contain most of the parameters which make up the dot products.

How much capacity is needed by supervised and reinforcement learning?

Past workPhysics of Language Models: Part 3.3, Knowledge Capacity Scaling Laws (Allen-Zhu and Li, 2024) has shown that neural networks can store 2 bits per parameter. These results pertain to the maximum amount of information absorbed in the long-training limit, not to the compute efficiency or rate of learning.

The 2-bits-per-parameter result relied on synthetic datasets cleverly constructed to contain a precise amount of information. It’s not as straightforward to estimate the information content required for a given realistic learning problem. One classic observation is that when minimizing log-loss, the total log-loss measured during the first epoch of training provides a measurement of the dataset’s description length. That is, an upper bound for the number of bits required to memorize the dataset. LLM datasets usually have a loss of around 1 bit (0.69 nats) per token, depending on dataset and model size.

This estimate measures the capacity required to perfectly memorize the dataset, which overestimates the actual capacity needed for “generalizable” learning that reduces log-loss on test data. Measuring the capacity requirements of supervised learning and how these interact with the number of trainable parameters is an open question for future work.

For RL, we claimed that policy gradient algorithms learn roughly 1 bit of information per episode, given that there’s a single reward value at the end of the episode. This isn’t a fundamental property of RL, as other algorithms could conceivably learn a lot more from each episode. For example, model-based RL algorithms train the learning agent to predict the observations and build a world model, potentially extractingmore information per episode. The claim of 1-bit-per-episode may only apply narrowly to policy gradient algorithms.

We can sharpen the bits-counting argument in information-theoretic terms. Consider an episode, consisting of a trajectory $\tau$ and final reward, as a message (i.e., a noisy channel) that provides some information about the unknown reward function $R$. We’ll condition on the current policy and training history and look at the mutual information between the policy gradient estimator and $R$. The REINFORCE update is $G = S \cdot \text{Adv}$ with $S = \nabla \log p_\theta(\tau)$. $S$ is independent of $R$ given the history, so the only $R$-dependent component is the scalar advantage.

By the data processing inequality:

$$I(G ; R | \text{history}) \leq I((S, \text{Adv}) ; R | \text{history}) = I(\text{Adv} ; R | S, \text{history}) \leq H(\text{Adv}).$$

If we quantize the advantage into $B$ bins, then $H(\text{Adv}) \lesssim \log(B)$. That is, the number of bits of useful information gleaned per episode is $O(1)$, independent of model size. These bits tell us which member of a discrete set of reward functions (or, equivalently, optimal-policy classes) we’re in. This analysis of mutual information mirrors what’s used in some theoretical analysis of optimization algorithms.Information Complexity of Black-Box Convex Optimization: A New Look via Feedback Information Theory (Raginsky and Rakhlin, 2009) Note that this estimate is an upper bound on the information absorbed by training; the actual amount learned will depend on the policy initialization and other details. For example, if we initialize with a policy that doesn’t get any reward, then the entropy of the advantage is zero (not log(B)), and it won’t learn anything.

Compute efficiency advantage of LoRA

Our experiments above measured learning progress against the number of training steps, but we may also be interested in the compute efficiency of different methods. We calculate that LoRA takes slightly more than ⅔ of the FLOPs that full fine-tuning does per pass. As a result, it will often outperform FullFT on compute efficiency overall.

We derive this ⅔ ratio by analyzing the FLOPs used in the forward–backward pass on a given weight matrix. These operations account for the vast majority of FLOPs in neural network models. We use the following notation:

$W \in \mathbb{R}^{N \times N}$ is a weight matrix
$x \in \mathbb{R}^N$ is an input vector
$y = Wx \in \mathbb{R}^N$ is an output vector
$\bar{x}, \bar{y} \in \mathbb{R}^N$ are the gradients of the loss with respect to $x$ and $y$, computed in the backward pass
$\bar{W} \in \mathbb{R}^{N \times N}$ is the gradient of the loss with respect to $W$

Full fine-tuning performs the following operations:

Forward

$y = Wx$ ($N^2$ multiply–adds)

Backward

$\bar{x} = W^T \bar{y}$ ($N^2$ multiply–adds)
$\bar{W} \mathrel{+}= x \bar{y}^T$ ($N^2$ multiply–adds)

The forward pass requires $N^2$ multiply-adds, and the backward pass requires another $2 \cdot N^2$ for $3N^2$ total. Training, which requires both, thus uses 3 times the FLOPs of forward-only inference.

With LoRA, we replace $W$ by $W + BA$, where $B \in \mathbb{R}^{N \times R}$ and $A \in \mathbb{R}^{R \times N}$, with $R \ll N$. Since we only update $\bar{A}$ and $\bar{B}$, we replace the third step of updating $\bar{W}$ with a much cheaper operation. $A$ and $B$ are $N \cdot R$ matrices, so the full forward-backward computation on each requires $3NR$ multiply-adds instead of $3N^2$ for $W$. The total for both is $6NR$. We also perform the forward-backward pass on $Wx$ and $\bar{x}$, equivalent to the first two steps of FullFT. The total number of multiply-adds is $2N^2 + 6NR$. With $R \ll N$, this is slightly more than $\frac{2}{3}$ of $3N^2$.

If we plotted LoRA performance over FLOPsThis analysis omits FLOPs used for attention, which could be significant in long-context settings. instead of training steps, it would show a clear advantage over FullFT.

Open questions

There are several questions related to our results that we would love to see investigated in the future:

Sharpening our predictions of LoRA performance and the precise conditions under which it matches full fine-tuning. We have roughly characterized the regime of equal performance and can estimate the required capacity in terms of tokens or episodes, but we can’t yet make accurate forecasts.
Our theoretical understanding of LoRA learning rates and training dynamics is limited. A fuller theory that explains the ratio between LoRA and FullFT learning rates would be valuable.
How do LoRA variants such as PiSSAPiSSA: Principal Singular Values and Singular Vectors Adaptation of Large Language Models (Meng, Wang & Zhang, 2024) perform when measured according to the methodology in this article?
There are various options for applying LoRA to MoE layers. LoRA users would benefit from an investigation into how well they perform, and how compatible each approach is with methods like tensor parallelism and expert parallelism that are important for large MoE models.

Closing thoughts

At Thinking Machines, we believe in the power of fine-tuning to advance AI usefulness in many domains of expertise. Our interest in LoRA is driven by a goal of making this power widely accessible and easily customizable to specific needs.

Aside from its practical uses, research on LoRA has also led us to deeper investigations of model capacity, dataset complexity, and sample efficiency. Looking at how learning speed and performance depend on capacity provides a lens for studying fundamental questions in machine learning. We look forward to advancing this study in the future.

Acknowledgements

We thank Dan Alexander Biderman, Weizhu Chen, Daniel Han, and Sadhika Malladi for their insightful feedback on an earlier draft of this post.

Citation

Please cite this work as:

Schulman, John and Thinking Machines Lab, "LoRA Without Regret",
Thinking Machines Lab: Connectionism, Sep 2025.

Or use the BibTeX citation:

@article{schulman2025lora,
  author = {John Schulman and Thinking Machines Lab},
  title = {LoRA Without Regret},
  journal = {Thinking Machines Lab: Connectionism},
  year = {2025},
  note = {https://thinkingmachines.ai/blog/lora/},
  doi = {10.64434/tml.20250929},
}

Modular Manifolds

Fri, 26 Sep 2025 00:00:00 +0000

When we train large neural networks, we need to keep them healthy. We do not want the tensors in the network—either the weights, activations or gradients—to grow too large or too small. Very small and very large tensors cause a variety of problems not just limited to numerical underflow and overflow. For example, weight matrices changing size during training makes it harder to design training algorithms—since the relative size of updates to weights has a significant impact on the speed of learning.

The gold standard for keeping tensors healthy is to normalize them. Normalization is commonplace for activation vectors, where we use techniques like layer norm to put the activations on a good scale before passing them to the next layer. It is also commonplace to normalize gradient updates, where we can interpret fast training algorithms like the Muon optimizer as spectrally normalizing the updates. Normalization provides us with certainty about the sizes of tensors—without needing to check Wandb!—and when training large neural networks with many interacting components, having certainty about the network internals is valuable.

Normalization is less commonly applied to weight matrices, although it is not unheard of. For example, the EDM2 diffusion model codebase uses weight constraints and the authors report benefits in their paper. And BiT uses weight standardization. Various other techniques have been proposed but are not common practice in modern large-scale training.For some more examples, see Salimans et al, 2016, Miyato et al, 2018 and our paper Liu et al, 2021. Normalizing the weight matrices might be a good idea for a few reasons. Weight constraints make understanding the relative size of optimization updates easier. They remove the problem of weight norms exploding. They allow us to focus hyperparameter tuning effort on tensors whose size matters most. They can force matrices to have a small condition number, making their behaviour more predictable. And relatedly, weight constraints facilitate Lipschitz guarantees for robustness to perturbations.

This post covers one appealing way to constrain the weight matrices of a neural network—by keeping the tensors constrained to submanifolds at each layer. This opens the door to re-thinking optimization, as we can co-design optimization algorithms with these manifold constraints. As an example, we proposeThis algorithm builds on work from Jianlin Su and Franz Louis Cesista, as discussed further below. a manifold version of the Muon optimizer whose weights are constrained to the Stiefel manifold: the manifold of matrices with unit condition number. We conclude the post by defining the idea of a modular manifold, which is a composable manifold that attempts to make it easier to scale up and train large networks.

Our goal in writing this post is to provide an introduction to a research area that we are excited about, and highlight many directions for future work. We would love to see more work from the community on the topics mentioned at the end of the post!

The shape of a manifold optimizer

This section works through the simplest example of learning on a manifold: a vector parameter constrained to a hypersphere in $\mathbb{R}^d$. The vector parameter is trained to minimize a loss function defined over the full space $\mathbb{R}^d$. This setup might be useful for, say, individual embedding vectors in a transformer model. This section will be a good warmup for the following section on manifold Muon that considers matrix parameters.

We will not be too formal about the definition of a manifold here: it is enough to understand that a manifold is a curved surface that looks flat when you zoom in close enough. The locally flat approximation at a point on the manifold is called the tangent space to the manifold, as visualized in Figure :

The sphere in three dimensions—or the hypersphere in higher dimensions—is a manifold. The locally flat approximation at a point on the manifold is called the tangent space to the manifold and is visualized as the red plane in the figure.

We can characterize the hypersphere in $d$ dimensions as the set of points $w \in \mathbb{R}^d$ of unit Euclidean norm. And the tangent space at a point $w$ on the hypersphere is the set of all vectors $a \in \mathbb{R}^d$ that are orthogonal to $w$.

To keep the weights constrained to the manifold, we could use a non-manifold optimizer and just project the weights back to the manifold after each step. Instead, we are interested in designing methods that take steps in the tangent space. The reason is that we would like to be able to equate the learning rate of our optimizer with the actual length of the optimization step. But if the optimization steps are pointing significantly off manifold and then being projected back, this nice property does not hold. Similar motivation is given in Section 2.3 of the EDM2 paper.

Before we can design a training algorithm for this manifold, something important we need to decide on is how to measure distanceFor a manifold to be “Riemannian”, the distance measure must be induced by an inner product. The Euclidean ($\ell_2$) norm is induced by an inner product, but the Manhattan ($\ell_1$) distance is not. in the tangent space. A common choice is the Euclidean distance, but we could also choose to measure distance in other ways, as visualized in Figure . In the next section, we will talk about choosing a distance measure based on the functionality of the module.

Inscribing unit balls in the tangent space for different distance measures. The $\ell_2$ (Euclidean) unit ball is a circle while the $\ell_1$ (Manhattan) unit ball is a diamond.

Crucially, the choice of distance measure changes the direction of the best optimization step. If the distance measure is non-Euclidean, then for a fixed length step, we may be able to move further in the direction of the gradientBy gradient, we mean the partial derivative of the loss with respect to the weights. Mathematicians reserve the term gradient for something else in Riemannian geometry. by not following the gradient direction exactly! This concept is visualized in Figure .

How geometry influences the direction of the best optimization step. The pink arrow represents the raw gradient—meaning the partial derivative of the loss with respect to the weights. The yellow diamond denotes the $\ell_1$ unit ball. The green arrow is the unit vector pointing furthest in the direction of the gradient. Notice how the green arrow is not parallel to the pink arrow. (In practice, the pink arrow need not lie in the tangent space, although the green arrow will do by construction.) Try dragging the pink arrow to see how the best update direction changes.

To see how this works out in math, we can formulate the optimal update direction given a manifold constraint and a distance measure as itself solving a constrained optimization problem. We will demonstrate this for the case of the hypersphere equipped with the Euclidean norm. Letting $g$ denote the gradient, $w$ the current point on the hypersphere, $a$ the update direction and $\eta$ the learning rate, we need to solve:

$$\min_{a\in\mathbb{R}^d} \quad \underbrace{a^\top g\vphantom{\|a\|_2 = 1}}_{\mathclap{\text{linear change in loss}}} \quad \text{such that} \quad \underbrace{\|a\|_2 = \eta}_{\mathclap{\text{size constraint}}} \quad \text{and} \quad \underbrace{a^\top w = 0\vphantom{\|a\|_2 = 1}}_{\mathclap{\text{tangent constraint}}}.\tag{$\star$}$$

Mapping back to the visual language of Figures , and , this formula says that the green arrow (optimal value of $a$) must belong to the red tangent hyperplane ($a^\top w = 0$) and must also lie on a yellow circle of radius $\eta$ ($\|a\|_2 = \eta$). To solve $(\star)$, we can apply the method of Lagrange multipliers. The relevant Lagrangian function is given by:

$$\mathcal{L}(a, \lambda, \mu) = a^\top g + \frac{\lambda}{2} \cdot (a^\top a - \eta) + \mu \cdot (a^\top w),$$

where $\lambda$ and $\mu$ are Lagrange multipliers. Setting the derivative of the Lagrangian with respect to $a$ to zero and applying the constraints to solve for $\lambda$ and $\mu$, the optimal update $a_\mathrm{opt}$ ends up being given by the following formula:

$$a_\mathrm{opt} = - \eta \times \frac{g - ww^\top g}{\|g-ww^\top g\|_2}.$$

In words, the optimal update is given by subtracting out the radial component from the gradient, normalizing and multiplying by the learning rate. Since this update lies in the tangent space, actually a very smallFor a learning rate $\eta$, the effect of the retraction map is $\mathcal{O}(\eta^2)$ small, so the learning rate almost equals the length of the step. correction is needed to stay on the manifold. The correction is known as a “retraction map” and is visualized in Figure :

Visualizing the retraction map. The green arrow is the update taken in the tangent space. Since for large step sizes the tangent space starts to diverge from the manifold, we need to project the updated weights back to the manifold using the retraction map—illustrated by the purple arrow.

We can solve for the retraction map by applying Pythagoras’ theorem to Figure . For a unit hypersphere and a step of length $\eta$, the hypotenuse has length $\sqrt{1+\eta^2}$ and therefore the retraction map for the hypersphere equipped with the Euclidean norm is simply given by dividing the updated weights through by $\sqrt{1+\eta^2}$. Putting everything together, the full manifold optimization algorithm is then given by:

$$w \gets \frac{1}{\sqrt{1+\eta^2}} \left[w - \eta \times \frac{g - ww^\top g}{\|g-ww^\top g\|_2}\right].$$

As an exercise for the reader: try calculating the Euclidean norm of the updated weight vector and check that the updated weight vector indeed lies on the hypersphere.

To summarize this section, a first-order manifold optimizer has three steps:

Find the tangent vector of unit length that goes furthest in the gradient direction.
Multiply this direction by the learning rate and subtract from the weights;
Retract the updated weights back to the manifold.

There are two decisions to make in applying this procedure: what manifold constraint we should use and how we should measure length. By making different decisions, we can generate different optimization algorithms as shown in the following table.

Manifold	Norm	Optimizer
Euclidean $\mathbb{R}^n$	Euclidean norm	vanilla gradient descent
Euclidean $\mathbb{R}^n$	infinity norm	sign gradient descent
hypersphere $S^n$	Euclidean norm	hyperspherical descent
matrix space $\mathbb{R}^{m\times n}$	spectral norm	Muon
Stiefel manifold $\subset\mathbb{R}^{m\times n}$	spectral norm	manifold Muon

We will derive the final algorithm in the table, manifold Muon, in the next section. To design a manifold constraint and a distance function for a matrix parameter, we shall think carefully about the role that a weight matrix plays inside a neural network.

Manifold Muon

A typical weight matrix $W$ in a transformer is a “vector-multiplier”, meaning that it transforms an input vector $x$ into an output vector $y = Wx$. We will design a manifold constraint and a distance function so that the matrix acts in a good way on input vectors: the matrix should not produce excessively small or large outputs, and updates to the matrix should not cause the output vector to change too much or too little.

A good way to think about how a matrix acts on vectors is through the singular value decomposition, illustrated in Figure . The SVD decomposes a matrix in a way that tells us how the matrix stretches input vectors along different axes.

The singular value decomposition. A matrix $M\in\mathbb{R}^{m\times n}$ of rank $k$ can always be decomposed as $M = U \Sigma V^\top$, where $U\in\mathbb{R}^{m\times k}$ and $V\in\mathbb{R}^{n\times k}$ have orthonormal columns and $\Sigma\in\mathbb{R}^{k\times k}$ is a diagonal matrix with only positive entries. The entries of $\Sigma$ are called the singular values of $M$. The singular values measure the stretching effect that the matrix has on vectors that align with the corresponding columns of $U$ and $V$.

We would like the matrix to have a stretching effect close to one, so we will choose a matrix manifold where all the singular values are exactly one. This matrix manifold is known formally as the Stiefel manifold. We can assume without loss of generality that we are dealing with a tall matrix ($m \geq n$), and then the Stiefel manifold can be equivalently defined as the following set:

$$\mathsf{Stiefel}(m,n) := \left\{ W \in \mathbb{R}^{m \times n} \mid W^T W = I_n \right\}.$$

Furthermore, one may show that a matrix $A \in \mathbb{R}^{m \times n}$ lies tangentNotice that the Stiefel constraint $W^T W = I_n$ directly generalizes the hyperspherical constraint $w^\top w = 1$ from the previous section. Similarly, the tangent space condition generalizes the hyperspherical one that $a^\top w = 0$. to the Stiefel manifold at matrix $W$ if and only if:

$$A^\top W + W^\top A = 0.$$

To design a manifold optimizer for the Stiefel manifold, all that remains is to choose a distance function. To limit the maximum stretching effect the weight update can have on an input vector, we will choose the spectral norm, which measures the largest singular value of a matrix. Although this only limits the maximum effect the update can have, since the optimizer we derive will saturate this bound, it will turn out to prevent the minimum effect of the update from being too small.There are some exceptions to this statement, such as when a weight matrix has a fan-out less than its fan-in, in which case we cannot escape from the matrix and its updates having a null space and mapping some inputs to zero.

The idea of doing gradient descent under a spectral norm constraint is what led to the Muon optimizer and, when combined with the Stiefel manifold constraint, we obtain a problem that we shall call manifold Muon:

$$\min_{A\in\mathbb{R}^{m\times n}} \quad \underbrace{\operatorname{trace}(G^T A)}_{\mathclap{\text{linear change in loss}}} \quad \text{such that} \quad \underbrace{\|A\|_{\text{spectral}} \leq \eta}_{\mathclap{\text{size constraint}}} \quad \text{and} \quad \underbrace{A^T W + W^T A = 0\vphantom{\|A\|_{\text{spectral}} = \eta}}_{\mathclap{\text{tangent constraint}}} \tag{$\dagger$}.$$

The manifold Muon problem $(\dagger)$ directly generalizes problem $(\star)$ from the previous section. Solving $(\dagger)$ is harder than solving $(\star)$, and here we will present a numerical solution inspiredI figured out how to solve manifold Muon in the square case late last year, but I was unable to solve the full rectangular case and thus posed the problem as an open problem on the Modula docs. Jianlin Su solved the problem this summer by taking a Lagrangian approach and working out a fixed point iteration on the optimality condition. I saw an early version of Jianlin’s work (which did not quite work yet) and also related work by Franz Louis Cesista, and I was able to work out the dual ascent algorithm presented here. by work done by Jianlin Su and Franz Louis Cesista.

Our key insight is that $(\dagger)$ is a convex optimization problem that may be solved via a standard method known as dual ascent. Here we will just sketch the main idea, but you can find a more detailed derivation on this page.

Similar to Jianlin’s approach, we introduce a matrix of Lagrange multipliers $\Lambda\in\mathbb{R}^{n\times n}$. We then apply a series of transformations to convert the problem $(\dagger)$ from a constrained minimization problem to an unconstrained maximization problem:

$$ \begin{align} (\dagger) &= \min_{\|A\|_\mathrm{spectral} \leq \eta} \max_{\Lambda} \;\operatorname{trace} G^\top A + \operatorname{trace}\Lambda^\top (A^\top W + W^\top A) \\ &= \min_{\|A\|_\mathrm{spectral} \leq \eta} \max_{\Lambda}\; \operatorname{trace}A^\top (G + 2W(\Lambda+\Lambda^\top))\\ &= \max_{\Lambda} \min_{\|A\|_\mathrm{spectral} \leq \eta} \; \operatorname{trace}A^\top (G + 2W(\Lambda+\Lambda^\top))\\ &= \max_{\Lambda} \; - \eta \times \|G + 2W(\Lambda+\Lambda^\top)\|_\mathrm{nuclear}. \end{align} $$

Equation (1) reformulates the problem as a saddle point problem: the maximization over $\Lambda$ will send the objective to infinity whenever the tangent space condition is violated. Equation (2) follows by applying properties of the trace and equation (3) follows from Sion’s minimax theorem. The inner minimization in equation (3) is solved by setting $A_\mathrm{opt}(\Lambda) = - \eta \times \operatorname{msign}(G + 2W(\Lambda+\Lambda^\top))$ where $\operatorname{msign}$ is the matrix sign function.The matrix sign function snaps the singular values of a matrix to one. It may be computed efficiently on GPUs via Newton-Schulz iteration or the recent Polar Express algorithm. And we obtain equation (4) by substituting this expression for $A_\mathrm{opt}(\Lambda)$ into equation (3). Equation (4) is known as the “dual problem” to $(\dagger)$ and we can solve it by gradient ascent. After some work, the gradient of the dual function is given by:

$$ \begin{align} H(\Lambda) &:= - \eta \times \nabla_\Lambda \|G + W (\Lambda+\Lambda^\top)\|_\mathrm{nuclear} \\ &= - \eta \times [W^\top\mathrm{msign}(G + 2W (\Lambda+\Lambda^\top)) + \operatorname{msign}(G + 2W (\Lambda+\Lambda^\top))^\top W], \end{align} $$

where the nuclear norm $\|\cdot\|_\mathrm{nuclear}$ measures the sum of the singular values of a matrix.

Finally, we can write down the manifold Muon algorithm:Note that this algorithm is closely related to Jianlin Su’s solution. Where we run dual ascent, Jianlin’s solution amounts to solving for the maximum of the dual function $H(\Lambda)=0$ via a fixed point iteration.

Run gradient ascent on the dual variable $\Lambda \gets \Lambda + \alpha \times H(\Lambda)$ to solve for $\Lambda_\mathrm{opt}$.
Compute the update $A_\mathrm{opt} = - \eta \times \operatorname{msign}(G + 2W(\Lambda_{\mathrm{opt}}+\Lambda_\mathrm{opt}^\top))$.
Apply the update to the weights $W \gets W + A_\mathrm{opt}$.
Retract the weights back to the manifold $W \gets \operatorname{msign}(W)$.

We ran a very small experiment to sanity check the algorithm and provide a minimal implementation for students or researchers to play with. Each training run finishes in less than a minute. The code is here and see Figure for the setup and results.

Training a small MLP for 3 epochs on the CIFAR-10 dataset. The different lightly shaded blue curves show different weight decay settings for AdamW. Results were averaged over 3 random seeds. The manifold Muon optimizer attained higher train and test accuracy than AdamW. The third plot shows the final singular value distribution of the first weight matrix for the best performing learning rate: the singular values after training with manifold Muon are all close to 1. Manifold Muon increased the wall clock time per step compared to AdamW, although this could be improved by running fewer steps of dual ascent or adding momentum to the algorithm and running dual ascent online. Depending on other systems bottlenecks, the overhead may not be an issue.

Modular manifolds

So far in this post, we have discussed manifold constraints for individual parameter tensors and co-designed optimization logic for these constraints. A question we have not answered is: what happens when we combine layers to build networks? Can we think about individual layers in isolation—or do we need to be careful about interactions between layers and modify the optimization logic in response? The goal of this section is to point out that there is a way to extend the reasoning we introduced in the previous two sections to the case of whole networks, and we call this the theory of modular manifolds.The theory of modular manifolds builds on research I did with my friend Tim Large, my postdoc advisor Phillip Isola, my PhD advisor Yisong Yue and many other amazing collaborators. At the end of the section, we provide some links to learn more.

The idea of modular manifolds is to build an abstraction that tells us how to budget learning rates across layers. The actual optimization logic in each layer ends up being the same as what we already worked out, except that the learning rate for a layer is modified depending on where the layer appears in the network. The abstraction rests upon a key observation made in our paper on the modular norm, that budgeting learning rates—both across layers and when scaling up individual layers—is intimately tied to understanding the Lipschitz sensitivity of the network output with respect to the weights. The abstraction tracks this sensitivity as we build the network, and manifold constraints help us get a much tighter understanding of this sensitivity.

The starting point for the abstraction is to think of any neural network module—from a layer to a whole transformer—as a mathematical object with three attributes:

A forward function $f:\mathcal{W} \times \mathcal{X} \to \mathcal{Y}$ that maps from a parameter space $\mathcal{W} = \mathbb{R}^d$ and an input space $\mathcal{X}$ to an output space $\mathcal{Y}$.
A submanifold of the weight space $\mathcal{M}\subset\mathcal{W}$ that the weights are constrained to.
A norm $\|\cdot\| : \mathcal{W} \to \mathbb{R}$ that acts as a measuring stick on weight space.

For example, a linear module equipped with the spectral norm and constrained to the Stiefel manifold, for which we have already worked out an optimizer, would be written:

$$ \mathsf{StiefelLinear} = \begin{cases}(W, x) \mapsto Wx, & \text{(forward function)}\\ \mathsf{Stiefel}(m,n), & \text{(manifold)}\\ \|\cdot\|_\mathrm{spectral}. & \text{(norm)}\end{cases}$$

Provided that an input $x$ to the $\mathsf{StiefelLinear}$ module has unit $\ell_2$ norm, then $\mathsf{StiefelLinear}$ is Lipschitz with respect to its weights in the module’s assigned norm with Lipschitz constant one:This argument can be extended to the RMS norm on the input and the RMS–RMS operator norm on the weights.

$$\|(W + \Delta W) x - Wx\|_2 \leq \|\Delta W\|_\mathrm{spectral} \times \|x\|_2 = \|\Delta W\|_\mathrm{spectral}.$$

This type of Lipschitz statement helps us understand how to scale weight updates to this module since it gives us a bound on how much the output can change when we perturb the weights. But when we compose two modules, can we automatically compile a Lipschitz statement on the joint weight space of the new module? The answer turns out to be yes, if we follow special rules for building the new module:

The new forward function $f_3$ is given by composing the two existing forward functions $f_1$ and $f_2$: $$f_3((w_1, w_2), x) := f_2(w_2, f_1(w_1, x)). \qquad$$
The new manifold constraint $\mathcal{M}_3$ is just the Cartesian product (see Figure for a fun example) of the two existing manifolds $\mathcal{M}_1$ and $\mathcal{M}_2$: $$\mathcal{M}_3 = \mathcal{M}_1 \times \mathcal{M}_2. \qquad$$
The new norm function is the max of the two existing norm functions weighted by special scalar coefficients $s_1$ and $s_2$. Letting $\|\cdot\|_1$ denote the first module’s norm and $\|\cdot\|_2$ denote the second module’s norm, the new norm $\|\cdot\|_3$ is given by: $$\|(w_1, w_2)\|_3 := \max(s_1\cdot \|w_1\|_1, s_2\cdot \|w_2\|_2). \qquad$$

When we use this composite norm to derive optimizers—following the same recipe we used in the first two sections of this post—we end up deriving separate optimizers for each layer, but the scalar coefficients $s_i$ budget the learning rates across layers.

We give much more detail on this construction, including extending it to other ways of combining modules, in our paper on the modular norm—although the paper does not cover manifold optimization. You can also check out our paper on modular duality for more on building optimizers in the modular norm. The Modula project builds toward a programmatic implementation of this construction.

The Cartesian product is a simple way to glue together two manifolds. For example, the product of a line and a disk is a cylinder. We get one copy of the disk at every point on the line.

Directions for future work

We are excited about any research that tries to make neural network training as principled and automatic as the forward pass. The ideas in this post benefitted strongly from interactions with external researchers like Jianlin Su and Franz Louis Cesista. We would love to see more work on these topics from the community.

Some possible directions for future work are:

Modularity. What manifolds should attention heads live on? Should embeddings be constrained differently than unembeddings? We can mix-and-match constraints in different parts of the network, or leave some tensors unconstrained.
Numerics. Manifold constraints also place constraints on the range of values that individual weight entries can take. Does this impact numerics, or make low-precision training easier?
Convex optimization. The manifold Muon algorithm involves running dual ascent. Can we apply more sophisticated convex optimization techniques to solve the dual problem faster or more reliably?
Convergence analysis. How fast do these algorithms converge? Does good conditioning of the weight matrices benefit convergence? Is there more that we can say theoretically?
Regularization. Manifold constraints implicitly regularize the model. Could we design constraints or tune their radii to improve generalization?
Architecture-optimizer co-design. While hard manifold constraints may not ultimately be the right way to constrain weight matrices, they exemplify the idea of tightly co-designing optimization algorithms with architecural components. Are there more opportunities here?
Non-Riemannian geometry. Most work on manifold optimization works in a Riemannian world where distances are induced by inner products and norm balls are ellipsoids. But neural networks are different: matrices act as operators, and operator norms like the spectral norm do not emerge from inner products. This implies, for example, that norm balls can have sharp corners and there is no unique gradient flow. Is there more to be discovered in this non-Riemannian world?
Practical implementation. Applying these techniques at scale requires efficient manifold operations on GPUs. The recent Polar Express paper shows promise for fast matrix sign computation. What other algorithmic innovations do we need?

Citation

Please cite this work as:

Jeremy Bernstein, "Modular Manifolds",
Thinking Machines Lab: Connectionism, Sep 2025.

Or use the BibTeX citation:

@article{bernstein2025manifolds,
  author = {Jeremy Bernstein},
  title = {Modular Manifolds},
  journal = {Thinking Machines Lab: Connectionism},
  year = {2025},
  note = {https://thinkingmachines.ai/blog/modular-manifolds/},
  doi = {10.64434/tml.20250926}
}

Defeating Nondeterminism in LLM Inference

Wed, 10 Sep 2025 00:01:00 -0700

Reproducibility is a bedrock of scientific progress. However, it’s remarkably difficult to get reproducible results out of large language models.

For example, you might observe that asking ChatGPT the same question multiple times provides different results. This by itself is not surprising, since getting a result from a language model involves “sampling”, a process that converts the language model’s output into a probability distribution and probabilistically selects a token.

What might be more surprising is that even when we adjust the temperature down to 0This means that the LLM always chooses the highest probability token, which is called greedy sampling. (thus making the sampling theoretically deterministic), LLM APIs are still not deterministic in practice (see past discussions here, here, or here). Even when running inference on your own hardware with an OSS inference library like vLLM or SGLang, sampling still isn’t deterministic (see here or here).

But why aren’t LLM inference engines deterministic? One common hypothesis is that some combination of floating-point non-associativity and concurrent execution leads to nondeterminism based on which concurrent core finishes first. We will call this the “concurrency + floating point” hypothesis for LLM inference nondeterminism. For example, a recent arXiv preprint writes:

Floating-point arithmetic in GPUs exhibits non-associativity, meaning $(a + b) + c \neq a + (b + c)$ due to finite precision and rounding errors. This property directly impacts the computation of attention scores and logits in the transformer architecture, where parallel operations across multiple threads can yield different results based on execution order.

You can also find the “concurrency + floating point” hypothesis repeated by others, like here (“There are speed tradeoffs, and in order to make the endpoints fast GPUs are used, which do parallel [nondeterministic] calculations. Any modern GPU neural net calculations will be subject to these."), or here (“Because GPUs are highly parallelized, the ordering of additions or multiplications might be different on each execution, which can cascade into small differences in output.").

While this hypothesis is not entirely wrong, it doesn’t reveal the full picture. For example, even on a GPU, running the same matrix multiplication on the same data repeatedly will always provide bitwise equal results. We’re definitely using floating-point numbers. And our GPU definitely has a lot of concurrency. Why don’t we see nondeterminism in this test?

A = torch.randn(2048, 2048, device='cuda', dtype=torch.bfloat16)
B = torch.randn(2048, 2048, device='cuda', dtype=torch.bfloat16)
ref = torch.mm(A, B)
for _ in range(1000):
    assert (torch.mm(A, B) - ref).abs().max().item() == 0

To understand the true cause of LLM inference nondeterminism, we must look deeper.

Unfortunately, even defining what it means for LLM inference to be deterministic is difficult. Perhaps confusingly, the following statements are all simultaneously true:

Some kernels on GPUs are nondeterministic.
However, all the kernels used in a language model’s forward pass are deterministic.
Moreover, the forward pass of an LLM inference server (like vLLM) can also be claimed to be deterministic.
Nevertheless, from the perspective of anybody using the inference server, the results are nondeterministic.

In this post, we will explain why the “concurrency + floating point” hypothesis misses the mark, unmask the true culprit behind LLM inference nondeterminism, and explain how to defeat nondeterminism and obtain truly reproducible results in LLM inference.

The original sin: floating-point non-associativity

Before talking about nondeterminism, it’s useful to explain why there are numerical differences at all. After all, we typically think of machine learning models as mathematical functions following structural rules such as commutativity or associativity. Shouldn’t there be a “mathematically correct” result that our machine learning libraries should provide us?

The culprit is floating-point non-associativity. That is, with floating-point numbers:

$$ (a + b) + c \neq a + (b + c) $$

(0.1 + 1e20) - 1e20
>>> 0
0.1 + (1e20 - 1e20)
>>> 0.1

Ironically, breaking associativity is what makes floating-point numbers useful.

Floating-point numbers are useful because they allow for a “dynamic” level of precision. For the purposes of explanation, we will use base 10 (instead of binary), where floating-point numbers are in the format $\text{mantissa} * 10^\text{exponent}$. We will also use 3 digits for the mantissa and 1 digit for the exponent.

For example, for the value 3450, we can represent it exactly as $3.45 * 10^3$. We can also represent much smaller values like 0.486 as $4.86 * 10^{-1}$. In this way, floating point allows us to represent both very small as well as very large values. In the sciences, we might say that floating point allows us to maintain a constant number of “significant figures”.

If you add together two floating-point numbers with the same exponent, it looks similar to integer addition. For example, 123 ($1.23 * 10^2$) + 456 ($4.56 * 10^2$) results in 579 ($5.79 * 10^2$).

But what happens when we add two floating-point numbers with different exponents, such as 1230 and 23.4? In this case, the exact result is 1253.4. However, we can only maintain 3 digits of precision at a time. Floating-point addition will thus drop the last 2 digits and obtain the value $1.25 * 10^3$ (or 1250).

1.23 × 10²

3.45 × 10¹

1.575 × 10²

Exact: 1575

We require 3 digits of precision to represent 1230 and 3 digits of precision to represent 23.4. However, adding these 2 numbers together results in a number that requires 5 digits of precision to represent (1253.4). Our floating-point format must then drop the 34 off the end. In some sense, we have effectively rounded our original 23.4 to 20.0 before adding it.

At this point, however, we’ve destroyed information. Note that this can happen every time we add two floating-point numbers with different “scales” (i.e. different exponents). And adding together floating-point numbers with different exponents happens all of the time. In fact, if we could guarantee that we never needed different exponents, we could just use integers!

In other words, every time we add together floating-point numbers in a different order, we can get a completely different result. To take an extreme example, there are 102 possible different results for summing this array depending on the order.

import random

vals = [1e-10, 1e-5, 1e-2, 1]
vals = vals + [-v for v in vals]

results = []
random.seed(42)
for _ in range(10000):
    random.shuffle(vals)
    results.append(sum(vals))

results = sorted(set(results))
print(f"There are {len(results)} unique results: {results}")

# Output:
# There are 102 unique results: [-8.326672684688674e-17, -7.45931094670027e-17, ..., 8.326672684688674e-17]

Although this is the underlying cause for non-identical outputs, it does not directly answer where the nondeterminism comes from. It doesn’t help us understand why floating-point values get added in different orders, when that happens, nor how it can be avoided.

The answers lie in how kernels are implemented.

Why don’t kernels always add numbers in the same order?

As mentioned above, one common explanation for why kernels add numbers in different orders is the “concurrency + floating point” hypothesis. The hypothesis states that if the order in which concurrent threads finish is nondeterministic and the accumulation order depends on the order in which concurrent threads finish (such as with an atomic add), our accumulation order will be nondeterministic as well.

Confusingly, although this can lead to nondeterministic kernels, concurrency (and atomic adds) end up being completely uninvolved in LLM inference nondeterminism! To explain what the real culprit is, let’s first understand why modern GPU kernels rarely need atomic adds.

When are atomic adds needed?

Typically a GPU launches a program concurrently across many “cores” (i.e. SMs). As the cores have no inherent synchronization among them, this poses a challenge if the cores need to communicate among each other. For example, if all cores must accumulate to the same element, you can use an “atomic add” (sometimes known as a “fetch-and-add”). The atomic add is “nondeterministic” — the order in which the results accumulate is purely dependent on which core finishes first.

Concretely, imagine that you are reducing a 100-element vector with 100 cores (e.g. torch.sum()). Although you can load all 100 elements in parallel, we must eventually reduce down to a single element. One way to accomplish this is with some kind of “atomic add” primitive, where the hardware guarantees that all additions will be processed but does not guarantee the order.

The atomic add ensures that every core's contributions will be reflected in the final sum. However, it makes no guarantee about what order the contributions will be added. The order depends entirely on which core finishes first, a nondeterministic property. Thus, executing the same parallel program multiple times can result in nondeterministic outputs.

This is usually what folks mean by “nondeterminism” — you execute the same kernel twice with exactly the same inputs and you get a different result out. This is known as run-to-run nondeterminism, where you run the same python script twice with the exact same dependencies but get a different result.

Although concurrent atomic adds do make a kernel nondeterministic, atomic adds are not necessary for the vast majority of kernels. In fact, in the typical forward pass of an LLM, there is usually not a single atomic add present.

This may be surprising, given that parallelizing a reduction can benefit from atomic adds. There are two main reasons why atomic adds do not end up being needed.

There is often sufficient parallelism along the “batch” dimension that we don’t need to parallelize along the reduction dimension. For example, let’s say that instead of reducing a single 100-dim vector we were reducing 500 vectors in parallel. In this case, we can reduce an entire vector in each core and allow every core to operate on a different vector.
Over time, most neural network libraries have adopted a variety of strategies for achieving determinism without sacrificing performance. For example, we can perform a “split” (or tree) reduction, where we split the 100-element reduction into five 20-element reductions (thus achieving five-way parallelism). Then, to combine the remaining five elements, we can either perform a separate “clean-up” reduction (which isn’t parallelized, but operates over few enough elements to be cheap) or utilize a semaphore (which ensures that each concurrent thread-block will accumulate in a deterministic order).The semaphore strategy can be found described here.

Due to these two factors, avoiding atomics adds is a negligible performance penalty for the vast majority of neural network operations.

There are still a couple of common operations that have significant performance penalties for avoiding atomics. For example, scatter_add in PyTorch (a[b] += c). The only one commonly used in LLMs, however, is FlashAttention backward.Fun fact: did you know that the widely used Triton implementations of FlashAttention backward actually differ algorithmically from Tri Dao’s FlashAttention-2 paper? The standard Triton implementation does additional recomputation in the backward pass, avoiding atomics but costing 40% more FLOPs!

However, the forward pass of an LLM involves no operations that require atomic adds. Thus, the forward pass in an LLM is in fact “run-to-run deterministic.”

From the perspective of the inference server, it is deterministic. Given the exact same user requests, it will always provide the same deterministic output.

Wikipedia writes that “a deterministic algorithm is an algorithm that, given a particular input, will always produce the same output.” And in this case, given the exact same inputs (i.e. the exact requests the inference server is processing), the forward pass always produces the exact same outputs.

However, the forward pass itself being “deterministic” is not sufficient to ensure that a system that includes it is deterministic. For example, what if our request’s output depended on the parallel user requests (e.g. batch-norm)? Since each individual request has no way of knowing what the parallel requests will be, from their perspective our overall LLM inference is also nondeterministic!

As it turns out, our request’s output does depend on the parallel user requests. Not because we’re somehow leaking information across batches — instead, it’s because our forward pass lacks “batch invariance”, causing our request’s output to depend on the batch size of our forward pass.

Batch invariance and “determinism”

To explain batch invariance, let’s simplify the system and look solely at matmuls. You can assume that all matmul implementations are “run-to-run deterministic."This is not totally true, but most common matmul implementations do have this property. However, they are not “batch-invariant.” In other words, when the batch size changes, each element in the batch can get different results.

This is a fairly unusual property from a mathematical perspective. Matrix multiplication should be “independent” along every element in the batch — neither the other elements in the batch nor how large the batch is should affect the computation results of a specific element in the batch.

However, as we can observe empirically, this isn’t true.

import torch
torch.set_default_device('cuda') 

B = 2048
D = 4096
a = torch.linspace(-1000, 1000, B*D).reshape(B, D)
b = torch.linspace(-1000, 1000, D*D).reshape(D, D)
# Doing a matrix vector multiplication by taking
# the first element of the batch
out1 = torch.mm(a[:1], b)
# Doing a matrix matrix multiplication and then taking
# the first element of the batch
out2 = torch.mm(a, b)[:1]
print((out1 - out2).abs().max()) # tensor(1669.2500, device='cuda:0')

Note that this is “run-to-run deterministic.” If you run the script multiple times, it will deterministically return the same result.It is not “hardware/software version invariant” — your GPU/PyTorch version may return a different value, but it should deterministically return the same value.

However, when a non-batch-invariant kernel is used as part of a larger inference system, the system can become nondeterministic. When you make a query to an inference endpoint, the amount of load the server is under is effectively “nondeterministic” from the user’s perspective. The load determines the batch size that the kernels are run under, and thus changes the eventual result of each individual request!

Although the inference server itself can be claimed to be "deterministic", the story is different for an individual user. From the perspective of an individual user, the other concurrent users are not an "input" to the system but rather a nondeterministic property of the system. This makes LLM inference "nondeterministic" from the perspective of each user.

If you compose some property under which the kernel is not invariant (i.e. batch-size) with nondeterminism of that property (i.e. the load the server is under), you get a nondeterministic system.

In other words, the primary reason nearly all LLM inference endpoints are nondeterministic is that the load (and thus batch-size) nondeterministically varies! This nondeterminism is not unique to GPUs — LLM inference endpoints served from CPUs or TPUs will also have this source of nondeterminism.

So, if we’d like to avoid nondeterminism in our inference servers, we must achieve batch invariance in our kernels. In order to understand how that can be achieved, let’s first take a look at why kernels don’t have batch invariance in the first place.

How do we make kernels batch-invariant?

In order to make a transformer implementation batch-invariant, we must make every kernel batch-invariant. Luckily, we can assume that every pointwise operation is batch-invariant.Although this is true for all kernels in say, PyTorch, it’s not inherently true. For example, there are some kernel implementations on CPU that will use vectorized intrinsics on some parts of the array and non-vectorized intrinsics on other parts, and these intrinsics don’t necessarily always have bitwise identical numerics. Thus, we only need to worry about the 3 operations that involve reductions — RMSNorm, matrix multiplication, and attention.Reductions related to parallelism are out of the scope of this discussion, but the same principles apply. One factoid that may be useful is that NVLink-Sharp in-switch reductions are deterministic on Blackwell as well as Hopper with CUDA 12.8+. As is the case with many things, this information can be found on NCCL’s github issues

Conveniently, these are also ordered in ascending levels of difficulty. Each one requires some additional considerations to achieve batch invariance with reasonable performance. Let’s talk about RMSNorm first.

Batch-invariant RMSNorm

Data Parallel RMSNorm Ideally, we'd like to avoid communication between cores in our parallelization strategy. One way to achieve that is by assigning one batch-element to each core, thus guaranteeing that each reduction is done entirely within a single core. This is what's known as a "data-parallel" strategy, since we're simply parallelizing along a dimension that doesn't require communication. In this example, we have four rows and four cores, saturating our cores.

RMSNorm can be implemented as:

# x: [batch_size, hidden_dim]
# weight: [hidden_dim]
def rms_norm(x, weight):
    return x * torch.rsqrt(torch.mean(x ** 2, dim=-1, keepdim=True)) * weight

The requirement for batch invariance is that the reduction order for each element must be fixed regardless of the batch-size of the kernel. Note that this doesn’t mean we must always use the same reduction strategy. For example, if we change the number of elements we’re reducing over, we can still be batch-invariant even if our reduction strategy changes.The Quack blog post has some nice examples showing the hierarchy of various reduction strategies you can do (e.g. thread reduction, warp reduction, block reduction, cluster reduction).

Thus, we only break batch invariance when our batch-size affects the reduction strategy.

Let’s look at the standard parallelism strategy for RMSNorm. Generally, parallel algorithms benefit from minimizing communication across cores. For the purpose of this discussion you can assume that when we refer to “cores” we mean SMs. More specifically, the property here that’s important is that the # of threadblocks our kernel launches is greater than the # of SMs. So, one strategy we can start with is to assign each batch element to one core, as seen in the above figure.

Increasing our batch size doesn’t affect our reduction strategy; if a batch size of 200 provides sufficient parallelism to our kernel then a batch size of 2000 will definitely provide sufficient parallelism.

Data Parallel RMSNorm for larger batches Extending the data-parallel strategy to larger batches is fairly straightforward --- instead of having each core handle one row you allow each core to handle different rows sequentially. This preserves batch invariance as the reduction strategy for each batch element remains identical.

On the other hand, decreasing the batch size can pose challenges. Because we assign each batch element to one core, decreasing our batch size will eventually lead to having more cores than batch elements, leaving some cores idle.

Upon encountering this situation, a good kernel engineer would reach for one of the solutions mentioned in the prior section (atomic adds or split reductions), maintaining good parallelism and thus, good performance. Unfortunately, this changes the reduction strategy, preventing this kernel from being batch-invariant.

Split-Reduction RMSNorm If we have a small batch size, our data-parallel strategy may no longer have sufficient parallelism to saturate our cores. In this case, it may be more efficient to "split" a reduction among multiple cores, allowing us to fully utilize our GPU. However, this loses batch invariance, as we are no longer reducing each element in the same order.

The easiest solution is to simply ignore these cases altogether. This is not completely unreasonable — a small batch size means that the kernel is likely to execute quickly anyways, and so a slowdown may not be catastrophic.

If we were compelled to optimize this use case, one approach would be to consistently use a reduction strategy that has enough parallelism even for very small batch sizes. Such a reduction strategy would lead to an excess amount of parallelism for larger batch sizes but would allow us to achieve decent (but not peak) performance across the entire range of sizes.

Batch-invariant matrix multiplication

Data Parallel Matmul Similar to RMSNorm, the standard parallelism strategy for matmuls is a "data-parallel" strategy, keeping the entire reduction in one core. It is most straightforward to think about splitting the output tensor into 2D tiles and assigning each tile to a different core. Each core then computes the dot products that belong to that tile, once again performing the entire reduction within one core.

Unlike for RMSNorm, additional constraints around arithmetic intensity and utilizing tensorcores force us to split 2D tiles instead of individual output elements for efficient matmul kernels.

At its core, you can view matrix multiplication as simply a pointwise operation followed by a reduction. Then, if we parallelize our matrix multiplication by chunking the output into tiles, we have an analogous “data-parallel” kernel strategy that keeps each reduction within one core.

Also similar to RMSNorm, it is possible for our “batch” dimensions (M and N) to become too small, forcing us to split along the reduction dimension (K). Despite having two “batch” dimensions, matmuls also require us to have much more “work” per core in order to leverage tensorcores effectively. For example, if you have a [1024, K] x [K, 1024] matmul and a standard 2D tile size of [128, 128], a data-parallel strategy would only be able to split this matmul into 64 cores, insufficient to saturate the GPU.

Splitting along the reduction dimension in a matmul is known as a Split-K Matmul. And just like RMSNorm, using this strategy breaks batch invariance. Another interesting parallelism strategy for matmuls is stream-k. Stream-k is interesting because it has even less invariance than typical matmuls. As discussed, most matmul libraries are not batch-invariant, but they’re at least what you could call batch-position-invariant (i.e. changing the position of the element within the batch does not affect numerics). However, stream-k is not batch-position-invariant either! Its core insight is that you can get cleaner load-balancing by splitting along k in different ways for different output tiles, but taking advantage of this makes our kernel not batch-position-invariant either.

Split-K Matmul If our batch dimension is fairly small we may not have enough parallelism and require a split-k matmul. In this example, we split each reduction across two cores, which would accumulate separately and then combine their results at the end. However, splitting each reduction across two cores allows us to still leverage eight cores.

There’s an additional complexity with matmuls — tensor core instructions. Whereas with reductions we could simply operate on one row at a time, efficient matrix multiplication kernels must operate on an entire “tile” at a time.

Each tensor-core instruction (like say, wgmma.mma_async.sync.aligned.m64n128k16) may have a different reduction order internally. One reason to use a different tensor-core instruction might be that the batch size is very small. For example, if we use a tensor-core PTX instruction that operates on a tile of length 256 but the batch size is only 32, we’re wasting almost all of that compute! At a batch-size of 1, the fastest kernels usually don’t use tensor cores at all.

Padded Tensor-Core Instructions If the batch size is too small, we may be in our situation where we can't fit even one of our 2D tiles in the output. In this case, it is most efficient to switch to a smaller tensor-core instruction or eschew tensor-cores altogether! However, both of these options prevent our kernel from being batch-invariant.

So, the easiest way to ensure batch invariance for matmuls is to compile one kernel configuration and use that for all shapes. Although we will lose some performance, this isn’t typically disastrous in LLM inference. In particular, split-k is most needed when both M and N are small, and luckily in our case, N (i.e. the model dim) is usually pretty large!

Despite obtaining batch invariance, we only lose about 20% performance compared to cuBLAS. Note that this is not an optimized Triton kernel either (e.g. no TMA). However, some of the patterns in performance are illustrative of where our batch-invariant requirement loses performance. First, note that we lose a significant amount of performance at very small batch sizes due to an overly large instruction and insufficient parallelism. Second, there is a "jigsaw" pattern as we increase the batch-size that is caused by quantization effects (both tile and wave) that are typically ameliorated through changing tile sizes. You can find more on these quantization effects here.

Batch-invariant attention

FlashAttention2 Strategy We parallelize along Q, and reduce along K/V simultaneously. This means that our entire reduction can be kept within a single core, making it another data-parallel strategy.

After obtaining batch invariance for matmuls, attention introduces two additional wrinkles — fittingly, because it contains two matmuls.

As opposed to only reducing over the feature dimension like both RMSNorm and matmuls, we now reduce over the feature dimension and sequence dimension.
Due to the above, attention must deal with a variety of inference optimizations that affect how sequences get processed (chunked prefill, prefix caching, etc.).

Thus, to achieve determinism in LLM inference our numerics must be invariant to both how many requests are processed at once and how each request gets sliced up in the inference engine.

Let’s first walk through the standard parallelism strategy for attention, first introduced in FlashAttention2. Similar to RMSNorm and Matmul, the default strategy is a “data-parallel” strategy. Since we reduce along the key/value tensors, a data-parallel strategy can only parallelize along the query tensor.

For example, depending on the inference engine’s choices, it’s possible that a sequence might get processed in several parts (such as in chunked prefill) or perhaps all at once (if the prefill isn’t split up). In order to achieve “batch invariance”, it’s necessary that the reduction order for a given token does not depend on how many other tokens from its sequence are being simultaneously processed. If you reduce over the K/V values in the KV cache separately from the K/V values in the current tokens being processed (like in vLLM’s Triton attention kernel), this can’t be achieved. For example, when processing the 1000th query token in a sequence, the reduction order must be identical regardless of whether 0 tokens are in the KV cache (prefill) or 999 tokens are in the KV cache (decoding).

FlashAttention with a KV Cache The reason why explicitly handling the KV cache separately from the current KV values breaks batch invariance is a bit subtle and is related to "boundary conditions". In particular, imagine your block size is 32 but we currently have 80 elements in our KV cache. We then compute an additional 48 elements that aren't cached. In this case, we need three blocks (two full and one masked) to compute "P cache" and another two blocks (one full and one masked) to compute "P". This is therefore five total blocks to compute our reduction when we only have four total blocks (i.e. 128) of elements to compute, which will definitely change our reduction order.

For example, if we instead had no elements in our KV Cache and were processing 128 elements altogether, we need to have identical numerics in both of these situations to ensure “batch invariance” for attention.

To resolve this, we can just update the KV cache and page table before the attention kernel itself, ensuring that our keys and values are always consistently laid out regardless of how many tokens are being processed.

With this additional detail (as well as all the things mentioned in the previous section, like consistent tile sizes), we are able to achieve a batch-invariant attention implementation!

However, there is a significant problem here. Unlike with matrix multiplication, the attention shapes we see in LLM inference often do require a split-reduction kernel, often known as Split-KV or FlashDecoding. This is because if we don’t parallelize along the reduction, we can only parallelize along the batch dimension, head dimension, and “query length” dimension. In the decode stage of attention, query length is very small, and so unless we have a very large batch size we are often unable to saturate the GPU.

Unfortunately, it’s not as easy to ignore this case as it was for RMSNorm and Matmuls. For example, if you have a very long KV cache, the attention kernel may take a very long time despite only processing one request.

Fixed # Split-KV Strategy (i.e. FlashDecode) If our query length becomes very small (like it does during decoding), we may end up in a situation where there is very little parallelism in our kernel at all. In these cases, we'll need to once again split along the reduction dimension --- the KV dimension this time. The typical strategy for how to split along the KV dimension is to figure out how much parallelism we need and then divide the KV dimension evenly. For example, if our KV length was 1000 and we needed 4 splits, each core would handle 250 elements.

This unfortunately also breaks batch invariance, as our precise reduction strategy depends on how many query tokens from the sequence we’re processing in any given request.

Furthermore, the split-reduction strategies commonly used for attention also pose challenges for batch invariance. For example, FlashInfer’s “balanced scheduling algorithm” chooses the largest split-size that can still saturate all the GPU’s cores, thus making the reduction strategy not “batch-invariant”. However, unlike with RMSNorm/Matmuls, it’s not sufficient to choose a fixed number of splits regardless of the batch size.

Instead, to achieve batch invariance, we must adopt a “fixed split-size” strategy. In other words, instead of fixing the # of splits, we fix the size of each split and then end up with a varying number of splits. In this manner, we can guarantee that regardless of how many tokens we’re processing, we always perform the identical reduction order. This requires some internal FlexAttention changes that are not included in our code release. We will upstream them in the near future!

Fixed Size Split-KV Strategy The only difference between this strategy and the previous strategy is that our splits are now "fixed size". For example, if our KV length was 1000, instead of splitting it into four even length 250 splits, we would split it into three fixed-size length 256 splits and one length 232 split.

This allows us to preserve batch invariance as our reduction strategy is no longer dependent on how many query tokens we’re processing at once!

Implementation

We provide a demonstration of deterministic inference on top of vLLM by leveraging its FlexAttention backend as well as torch.Library. Through torch.Library, we’re able to substitute out most of the relevant PyTorch operators in an unintrusive way. You can find the library of “batch-invariant” kernels at thinking-machines-lab/batch-invariant-ops, as well as the vLLM example of running in “deterministic” mode.

Experiments

How nondeterministic are completions?

We use Qwen/Qwen3-235B-A22B-Instruct-2507 and sample 1000 completions at temperature 0 with the prompt “Tell me about Richard Feynman” (non-thinking mode), generating 1000 tokens each. Surprisingly, we generate 80 unique completions, with the most common of these occuring 78 times.

Looking at where the completions differ, we see that the completions are actually identical for the first 102 tokens! The first instance of diverging completions occurs at the 103rd token. All completions generate the sequence “Feynman was born on May 11, 1918, in” However, 992 of the completions go on to generate “Queens, New York” whereas 8 of the completions generate “New York City”.

On the other hand, when we enable our batch-invariant kernels, all of our 1000 completions are identical. This is what we would mathematically expect from our sampler, but we aren’t able to achieve deterministic results without our batch-invariant kernels.

Performance

We have not put a significant effort into optimizing the performance of the batch-invariant kernels here. However, let’s run some experiments to verify that our performance remains usable.

We will set up an API server with one GPU running Qwen-3-8B, and request 1000 sequences with an output length of between 90 and 110.

Configuration	Time (seconds)
vLLM default	26
Unoptimized Deterministic vLLM	55
+ Improved Attention Kernel	42

Much of the slowdown comes from the fact that the FlexAttention integration in vLLM has not been heavily optimized yet. Nevertheless, we see that performance is not disastrous.

True on-policy RL

As researchers have noted, the different numerics between training and inference implicitly turns our on-policy RL into off-policy RL.

Of course, it is impossible to get bitwise identical results between training and inference if we can’t even get bitwise identical results from two identical inference requests. Then, deterministic inference enables us to also modify our training stack to obtain bitwise identical results between sampling and training, thus resulting in true on-policy RL.

We run experiments in a RLVR setup on Bigmath with the RL policy initialized from the Qwen 2.5-VL instruct 8B with a max rollout length of 4096.

If we train without off-policy correction (i.e. importance weighting), our reward collapses partway through training, whereas adding an off-policy correction term allows training to proceed smoothly. But, if we achieve bitwise identical results between our sampler and trainer, we are fully on policy (i.e. 0 KL divergence) and can also train smoothly.

We can also plot the KL-divergence in logprobs between our sampler and trainer, where all 3 runs have notably different behavior. When running with importance weighting, it stays around 0.001 with occasional spikes. However, running without importance weighting eventually leads to a spike in KL-divergence around the same time that reward crashes. And, of course, when running “True On-Policy RL”, our KL-divergence stays flat at 0, indicating that there is no divergence between the training policy and sampling policy.

Note that the run without importance weighting has a significant loss spike around Step 318, and this comes with a correspond ing spike in KL-divergence of logprobs. Meanwhile, either using an off-policy correction or running with "True On-Policy" allows RL to continue smoothly. The blue line showing "True On-Policy" is not a bug - it's just a flat line at 0.

Conclusion

Modern software systems contain many layers of abstractions. In machine learning, when we run into nondeterminism and subtle numerical differences it can often be tempting to paper over them. After all, our systems are already “probabilistic”, so what’s wrong with a little more nondeterminism? What’s wrong with bumping up the atol/rtol on the failing unit test? The difference in logprobs between the trainer and the sampler probably isn’t a real bug, right?

We reject this defeatism. With a little bit of work, we can understand the root causes of our nondeterminism and even solve them! We hope that this blog post provides the community with a solid understanding of how to resolve nondeterminism in our inference systems and inspires others to obtain a full understanding of their systems.

Citation

Please cite this work as:

He, Horace and Thinking Machines Lab, "Defeating Nondeterminism in LLM Inference", 
Thinking Machines Lab: Connectionism, Sep 2025.

Or use the BibTeX citation:

@article{he2025nondeterminism,
  author = {Horace He and Thinking Machines Lab},
  title = {Defeating Nondeterminism in LLM Inference},
  journal = {Thinking Machines Lab: Connectionism},
  year = {2025},
  note = {https://thinkingmachines.ai/blog/defeating-nondeterminism-in-llm-inference/},
  doi = {10.64434/tml.20250910}
}

Thinking Machines Lab

The Future Worth Building Is Human

Bringing intelligence to knowledge

Human participation is a technical challenge

Decentralized alignment

The future worth building

Interaction Models: A Scalable Approach to Human-AI Collaboration

The collaboration bottleneck

Capabilities

Our approach

System overview

The interaction model

Benchmarks

Intelligence and interactivity frontier

New dimensions of interactivity

Limitations and future work

Tell us what you think, join us

Citation

On-Policy Distillation

On-policy distillation — best of both worlds

Implementation

Loss function: reverse KL

Illustration

Pseudocode

Distillation for reasoning

Off-policy distillation

Reinforcement learning

On-policy distillation

Distillation for personalization

Training an internal assistant

Training on new knowledge degrades learned behavior

On-policy distillation recovers post-training behavior

Discussion

Dense supervision greatly improves compute efficiency

Distillation can effectively reuse training data for data efficiency

RL searches in the space of semantic strategies

On-policy learning as a tool for continual learning

Conclusion

Citation

LoRA Without Regret

What matters for LoRA

Methods and results

LoRA rank

Batch size effects

Layers Where LoRA Is Applied

Reinforcement learning

Setting LoRA hyperparameters

Optimal learning rate and rank

Parametrization invariances

Optimal learning rates for LoRA vs. FullFT

Learning rates in short and long runs

Discussion

Why LoRA might be needed on all layers

How much capacity is needed by supervised and reinforcement learning?

Compute efficiency advantage of LoRA

Open questions

Closing thoughts

Acknowledgements

Citation

Modular Manifolds

The shape of a manifold optimizer

Manifold Muon

Modular manifolds

Directions for future work

Further reading

Citation

Defeating Nondeterminism in LLM Inference

The original sin: floating-point non-associativity

Why don’t kernels always add numbers in the same order?

When are atomic adds needed?

Batch invariance and “determinism”

How do we make kernels batch-invariant?

Batch-invariant RMSNorm

Batch-invariant matrix multiplication

Batch-invariant attention

Implementation

Experiments

How nondeterministic are completions?

Performance

True on-policy RL