揭秘大模型推理非确定性:Thinking Machines Lab的批次不变性解决方案 Best Partners TV 2025-09-12

开篇:大模型推理的“老大难”问题

大家好,这里是最佳拍档,我是大飞。今天将和大家探讨一个大模型开发与应用领域普遍面临的难题:为何在输入完全相同内容时,大模型会产生不同的输出?即使已固定随机种子(random seed: 用于初始化伪随机数生成器,确保每次运行生成相同的随机序列),结果依然可能“失控”。

近日,一家公司公布的研究成果直接揭示了这一问题的核心真相,这家公司便是由OpenAI前CTO米拉·穆拉蒂(Mira Murati)创立的Thinking Machines Lab。首先,我们来了解一下这家公司背景,它自成立之初便自带“顶流”光环。今年7月,Thinking Machines Lab完成了一轮高达20亿美元的种子轮融资,公司估值直接冲到了120亿美元。值得注意的是,在获得这笔巨额融资时,该公司尚未发布任何产品。这种“未出产品先获重投”的情况,在AI行业中极为罕见。其投资方名单更是星光熠熠,领投方为知名风投A16z,跟投方包括英伟达、AMD、思科(Cisco)等科技巨头。为何资本市场愿意给予一家“零产品”公司如此高的估值?答案可能就隐藏在其团队背景及最新发布的研究中。

9月11日,Thinking Machines Lab正式推出了其研究博客,命名为Connectionism,即联结主义(Connectionism: 人工智能领域的一个子领域,研究神经网络与生物大脑的相似性)。博客首篇文章直接瞄准了大模型推理中最令人困扰的“非确定性”问题,标题为《击败大语言模型推理中的非确定性》。这篇文章不仅揭示了许多人对大模型非确定性的误解,还提供了可行的解决方案,并附上了完整的实验数据。本文将从现象、原理到解决方案,对这篇论文进行深入解读。

大模型非确定性之谜:常见误解与核心真相

首先,我们需要明确一个前提:大模型推理的非确定性究竟是如何产生的?许多从事算法开发和工程落地的人员可能都有过类似经历:向大模型输入完全相同的提示词,例如“请介绍一下理查德·费曼(Richard Feynman)”,并在代码中将随机种子设置为同一值。按常理,此时模型的输出应完全一致,但实际情况是,有时运行10次会得到5到6种不同的结果。更令人困惑的是,即使使用vLLM(vLLM: 一个高性能的开源大语言模型推理库)、SGLang(SGLang: 一个用于大模型推理优化的开源库)等开源推理库在自有硬件上运行,这种情况依然会发生。

面对这种现象,行业内最常见的解释是什么呢?许多人会将原因归咎于“GPU并发执行(GPU Concurrent Execution: 图形处理器并行执行多个计算任务的能力)”或“浮点数运算”。例如,有人认为GPU是并行计算的,不同核心的计算完成顺序不同,自然会导致结果差异;还有人指出,浮点数运算本身存在误差,误差累积后便导致输出不一致。这些说法并非完全错误,但它们并未触及问题的核心。至少在Thinking Machines Lab的研究中,这些并非导致大模型推理非确定性的“真凶”。

那么,真正的原因是什么呢?研究团队通过大量实验发现,核心问题在于大模型推理过程中的“批次不变性缺失(Batch Invariance Missing: 指模型在不同批次大小下,相同输入产生不同输出的现象)”。

批次处理、浮点数非结合性与归约化顺序

在此,需要先解释一下“批次处理”的概念。当你向大模型推理服务器发送请求时,服务器不会每次只处理你一个人的请求,它会根据当前负载情况,将你的请求与其他用户的请求打包成一个“批次”(batch)来处理。例如,在低峰期发送请求,服务器可能将你的请求和另外2个请求打包成一个“批次大小3”的任务;而在高峰期,可能会与另外15个请求打包成一个“批次大小16”的任务。问题便在于此:相同的输入,在不同批次大小下,会产生不同的输出。这好比你在餐厅点一道番茄炒蛋,厨房却告知你,今天同时炒3道菜,你的番茄炒蛋会偏咸;如果同时炒10道菜,你的番茄炒蛋则会偏甜。这听起来荒谬,但这正是目前绝大多数大模型推理系统的真实状况。

导致这种“批次依赖”的根本原因,还要从浮点数的一个特性说起,那就是浮点数非结合性(Floating-Point Non-Associativity: 浮点数运算中,改变运算顺序可能导致结果不同的特性)。什么是浮点数的非结合性?简单来说,对于浮点数运算,(a+b)+c的结果并不等于a+(b+c)。举一个具体例子:计算(0.1 + 1e20) - 1e20,结果会是0;但如果计算0.1 + (1e20 - 1e20),结果会是0.1。明明是同样的三个数,只是改变了加法的顺序,结果却完全不同。

为何会如此?因为浮点数的设计初衷是“动态精度”,通过“尾数×10^指数”的形式,既能表示非常大的数,也能表示非常小的数。然而,其代价是当两个数值尺度差异过大时,小数的精度会被大数“吞噬”。例如,1e20是一个极其大的数,0.1与它相加时,0.1的精度会被完全忽略,结果仍是1e20;而1e20减1e20等于0,再加上0.1才会保留0.1的精度。

这种特性在大模型的核心操作中会被无限放大。在模型推理过程中,矩阵乘法(Matrix Multiplication: 两个矩阵相乘的运算)、RMSNorm(Root Mean Square Normalization: 均方根归一化,一种常用的神经网络层归一化方法)和注意力机制(Attention Mechanism: 神经网络中一种允许模型关注输入序列不同部分的技术)等关键步骤,本质上都涉及大量的浮点数加法和乘法。当服务器处理不同批次大小的请求时,这些操作的归约化(Reduction: 将一组数值通过某种运算(如求和、求均值)聚合成单个数值的过程)顺序,即浮点数相加的顺序,会发生变化。例如,当批次较大时,内核会采用“分割归约”策略,将大任务拆分成小任务并行计算;当批次较小时,可能会采用“单核心完整归约”策略。归约顺序的变化,自然导致最终计算结果的不同。

在此需要澄清一个误区:许多人以为“GPU并发执行”是关键,但实际上,即使在GPU上,只要归约顺序固定,相同的计算任务重复执行1000次,结果也会完全一致,甚至每一位都相等。Thinking Machines Lab团队曾进行一项实验:使用PyTorch定义两个2048×2048的随机矩阵A和B,计算它们的乘积,并重复1000次。每次的结果都完全相同。这说明GPU本身的并发能力并不会导致非确定性,真正的问题在于内核的归约策略会随着批次大小而变化,即“批次不变性缺失”。

解决方案:实现批次不变内核

既然找到了核心原因,解决方案便明确了:实现“批次不变内核(Batch-Invariant Kernels: 无论批次大小如何,都采用相同归约顺序的计算内核)”,让大模型的核心操作,如RMSNorm、矩阵乘法和注意力机制,在任何批次大小下都采用相同的归约顺序。接下来,我们将具体探讨这三类操作如何实现批次不变性。

RMSNorm的批次不变实现

首先看RMSNorm。RMSNorm的计算公式为:输出 = 输入 × (1/根号(输入平方的均值)) × 权重。这里的关键在于输入平方的均值,因为计算均值需要对输入的每个元素做平方后求和,即归约操作。常规的RMSNorm内核在批次较大时,会采用数据并行的策略,每个核心处理一个批次元素,并在核心内部完成完整的归约,从而使归约顺序固定;然而,当批次变小时,核心数量多于批次元素,为避免算力浪费,内核会切换到“分割归约”策略,将一个批次元素的归约任务拆分给多个核心,这便改变了归约顺序。

Thinking Machines Lab的解决方案十分直接:无论批次大小如何,都强制采用数据并行的策略。即使批次很小导致部分核心闲置,也要保证每个批次元素的归约在单个核心内完成。这样一来,无论批次是3还是16,归约顺序都不会改变,RMSNorm的结果自然也就实现了批次不变。

矩阵乘法的批次不变实现

再来看矩阵乘法,它是大模型中计算量最大的操作。常规内核为追求性能,会根据批次大小调整策略。例如,当矩阵维度(M、N)较小时,内核会采用“Split-K(Split-K: 一种矩阵乘法优化策略,将K维度拆分并行计算)”策略,将矩阵乘法的归约维度(K)拆分成多个部分,并行计算后再合并结果。然而,Split-K会改变归约顺序,导致批次不变性缺失。此外,GPU中张量核心(Tensor Core: 英伟达GPU中专门用于加速矩阵乘法和累加运算的硬件单元)的指令选择也会影响归约顺序。例如,批次较小时,内核可能使用小尺寸的张量核心指令(如m64n128k16);批次较大时,可能使用大尺寸指令。而不同指令的内部归约顺序是不同的。

针对矩阵乘法,Thinking Machines Lab的方案是固定内核配置。首先,禁用Split-K策略,无论矩阵维度如何,都不拆分归约维度。其次,固定张量核心指令的尺寸,例如统一使用128×128×32的块大小。即使批次较小时会有算力浪费,也要保证归约顺序一致。实验数据显示,这种方案虽然会损失一定的性能,但在批次较大的场景下,性能损耗仅为20%左右,完全在可接受范围内。

注意力机制的批次不变实现

最后是注意力机制,这也是最复杂的部分。注意力机制涉及两次矩阵乘法(Q×K^T和注意力权重×V)以及一次softmax,其中Q×K^T的归约过程最容易受批次影响。常规的注意力内核,例如FlashAttention2(FlashAttention2: 一种高效的注意力机制算法,优化了内存访问和计算效率),在批次较大时,会沿Q维度并行,每个核心处理一部分Q,并在核心内部完成Q与K^T的归约;但在解码阶段,Q维度的并行性会消失,内核会切换到Split-KV(Split-KV: 一种注意力机制优化策略,将KV缓存拆分并行处理)策略,即将KV缓存拆分成多个部分,并行计算后合并,这便改变了归约顺序。更甚者,许多推理引擎,如vLLM,会对KV缓存进行“分块预填充”,将一个序列的KV拆分成多个块处理,这同样会导致不同块的归约顺序出现差异。

Thinking Machines Lab的解决方案是固定分割大小,即不按分割数量,而是按固定分割大小来拆分KV缓存。例如,无论KV缓存总长度是1000还是2000,都按照每个分割256个元素的大小来拆分。这样一来,1000个元素会拆成3个256的块和1个232的块;2000个元素则拆成7个256的块和1个208的块。如此,无论序列长度、批次大小如何,KV的归约顺序都由固定的分割大小决定,自然就实现了批次不变。同时,他们还在注意力内核之前统一更新KV缓存的布局,确保预填充和解码阶段的KV数据格式一致,从而避免因数据布局差异导致的归约顺序变化。

实验验证与性能考量

那么,这些解决方案的效果究竟如何呢?Thinking Machines Lab进行了一组非常有说服力的实验。他们使用的模型是Qwen/Qwen3-235B-A22B-Instruct-2507,提示词为“介绍一下理查德·费曼(Tell me about Richard Feynman)”,温度设置为0,即贪婪采样(Greedy Sampling: 一种文本生成策略,每次选择概率最高的下一个词,不引入随机性)。理论上应输出唯一结果,总共采样1000次。

在使用常规内核的情况下,结果令人意外:1000次采样竟然产生了80个不同的完成结果。更有意思的是,所有结果的前102个token(token: 文本经过分词器处理后得到的最小语义单元)完全相同,均为“费曼出生在1918年五月11日(Feynman was born on May 11, 1918, in)”,但从第103个token开始出现分歧。例如,跑完992步输出的是纽约皇后区,而跑完8步输出的却是纽约市。这种局部分歧正是归约顺序差异累积所导致的。

然而,当启用“批次不变内核”之后,奇迹发生了:1000次采样的结果完全一致,没有任何差异。这表明,只要解决了批次不变性问题,大模型推理的非确定性便可被彻底消除。当然,性能方面确实存在一定损耗。未经优化的批次不变内核会导致推理速度下降约2倍,但经过优化(例如调整核心调度、优化内存访问)后,性能损耗已控制在可接受范围内。对于需要结果一致性的场景,如金融AI、医疗诊断和强化学习训练等,这种性能上的权衡是完全值得的。

对在线策略强化学习的关键影响

谈及强化学习,这项研究的一个重要价值在于它使大模型的“在线策略强化学习(On-Policy Reinforcement Learning: 强化学习的一种,学习和执行策略是同一个策略)”,即On-Policy RL成为可能。此前在线策略RL难以落地的原因在于:训练时的模型输出与推理时的模型输出不一致。训练时使用小批次甚至单批次,而推理时使用大批次,两者之间的差异会导致策略偏移,最终造成训练崩溃。有了批次不变内核后,训练和推理的输出完全一致,KL散度(Kullback-Leibler Divergence: 衡量两个概率分布之间差异的指标),即衡量两个概率分布差异的指标,可以保持在0,真正实现了在线策略。

Thinking Machines Lab团队在Bigmath上进行了实验:在没有批次不变内核的情况下,不使用离线策略校正(例如重要性加权的RL训练),会在约318步时出现奖励崩溃。而启用批次不变内核后,即使不使用校正,训练也能顺利进行,奖励的稳定性大幅上升。

研究的深远意义与未来展望

讨论完技术方案,我们来总结一下这项研究的意义。它不仅解决了大模型推理非确定性的技术难题,更重要的是,它为大模型的“可重复性”和“可靠性”提供了科学的解决方案。在AI行业快速发展的今天,我们常常被性能、参数规模等指标吸引,却忽略了“可重复性”这一科学研究的基石。如果一个模型的输出无法复现,那么即便其性能再强,也难以在需要严谨性的领域落地应用。Thinking Machines Lab的研究恰恰填补了这一空白。

正如托马斯在社交媒体上的总结:“大模型推理非确定性,不只是浮点数非结合性或者GPU并发执行的问题,核心的罪魁祸首是批次方差,服务器负载不可预测地改变了数值计算。批次不变内核解锁了真正的可重复性,终于让强化学习的在线策略变得可行。”这句话也点出了这项研究的核心价值:它并非对现有技术的小修小补,而是从底层逻辑上解决了大模型落地的一个关键障碍。

对于普通用户而言,这项研究也意味着未来我们使用的AI产品将变得更加可靠。例如,AI生成的报告不会再因服务器负载变化而出现不同版本,AI辅助的决策也不会因计算顺序差异而产生偏差。对于AI从业者来说,这篇文章也提供了一个全新的思考视角:在追求模型性能的同时,也要关注底层计算的确定性。这或许会成为未来大模型技术竞争的一个新焦点。

最后,还有一个值得关注的小细节,即博客的名字“联结主义”。这个词并非凭空创造,它实际上是1980年代AI领域的一个重要子领域,核心研究方向是“神经网络与生物大脑的相似性”,简单来说,就是试图让AI的神经网络像人脑的神经元一样工作。米拉·穆拉蒂在社交媒体上解释说,Thinking Machines Lab的使命之一是提高人们对AI的科学理解,并与更广泛的研究社区合作,推出这个博客正是为了分享他们的科学见解。更有意思的是,根据联合创始人Lilian Weng的补充,Thinking Machines Lab的第一代旗舰产品就叫“Connection Machine”。

至此,关于这篇文章的解读告一段落。相关链接将放置在视频简介中。如果您有任何疑问或想讨论更多技术细节,欢迎在评论区留言。感谢收看本期视频,我们下期再见。

📌 文中提及的人物和组织

媒体/书籍: best-partners-tv

关键字: llm non-determinism technology