Restoring AI Reproducibility: Thinking Machines Lab's Challenge to Overcome Nondeterminism
One of the challenges that AI, particularly Large Language Models (LLMs), often faces is "nondeterminism," where the model returns a different response for the same input every time. This is a significant issue that undermines reliability and reproducibility in scientific research and corporate applications, but TML views this not merely as a constraint, but as a "solvable technical challenge" and is advancing its research accordingly.
1. What is Nondeterminism? Exploring its Causes
1-1. How Nondeterminism Manifests
It is not uncommon for an LLM to provide slightly different answers when asked the same question multiple times. This generally occurs during the "sampling" process, and even if the "temperature parameter" is set to 0 to theoretically select the most probable token, nondeterminism often remains in practice.
1-2. The "Numerical Computation" of GPU Kernels is Key
Many engineers and researchers believe that "the non-associativity (non-commutativity/non-associativity) of parallel processing and floating-point arithmetic is the cause." For example, because the result can differ depending on the order of calculation, such as (a + b) + c ≠ a + (b + c).
2. Root Cause: Lack of Batch Invariance
2-1. Why Kernels Cannot Calculate in the Same Order
Horace He and others point out that the true cause of nondeterminism is that "the calculation order and implementation strategy change depending on conditions such as batch size." Because of this, the results can change depending on server load or batch processing status, even with the same input.
2-2. What is Batch Invariance?
"Batch invariance" refers to the property where each calculation result is obtained consistently regardless of the batch size or the position of its elements. However, in reality, nondeterminism occurs when this property is broken.
2-3. Operations That Specifically Pose Problems
RMSNorm (Normalization Operation)
Matrix Multiplication (MatMul)
Attention Mechanism
Because these include "reduction," the calculation order and results can change depending on the batch size or the method of split processing.
3. Approach to a Solution: Batch-Invariant Kernel Design
3-1. Design That Maintains a Consistent Reduction Order
For example, in RMSNorm, batch invariance can be maintained by processing each batch element with a dedicated GPU core and using a method that reduces them in a consistent order.
3-2. Innovations in Matrix Multiplication and Attention
Matrix multiplication requires a data-parallel strategy that performs reductions in a consistent order while minimizing performance degradation.
Attention is particularly susceptible to the effects of processing splits, such as with the "KV cache," so a "fixed-size split strategy" is necessary to maintain order regardless of batch size.
3-3. Implementation and Results
In experiments incorporating a batch-invariant kernel mode based on vLLM + FlexAttention, it was confirmed that sampling results matched perfectly (i.e., achieving "true reproducibility (complete determinism)").
In addition, while performance is slightly slower than the default mode, it has been shown to remain within a practical range without causing "fatal latency."
4. Practical Impact and Future Outlook
4-1. Benefits for Reinforcement Learning (RL)
Reproducible responses mitigate off-policy issues in reinforcement learning and improve the consistency between training and inference. In fact, it is reported that ensuring this determinism enables "true on-policy RL," leading to a smoother training process.
4-2. Stance on Transparency and Open Research
TML has launched a new blog series called "Connectionism," aiming for a highly open research culture through the regular publication of research findings and the sharing of code.
4-3. Expectations for Commercialization and Social Implementation
Founder Mira Murati has indicated an intention to release "products useful for researchers and startups" in the near future, increasing the likelihood that this deterministic technology will become the foundation for new AI services.
Conclusion
Thinking Machines Lab's "Defeating Nondeterminism in LLM Inference" is a highly concrete and practical approach to the difficult problem of "unpredictable responses" in LLMs. By delving into GPU kernel design to enforce batch invariance, they have restored determinism and improved reproducibility.
This research will serve as a cornerstone for evolving AI into a more reliable tool for research and business applications. We look forward to future developments from TML and their results.