通过平滑 MMD 对齐增强大语言模型的数值预测能力

背景

尽管大语言模型（LLMs）在通用语言理解和生成任务中展现了强大的能力，但在需要数值精度的输出场景中，它们的表现往往不可靠。这种“数值幻觉”或计算错误是制约 LLM 在科学计算、金融分析等高精度领域应用的关键瓶颈。

造成这一问题的核心原因在于训练目标函数的设计。标准的交叉熵损失（Cross-Entropy Loss）将数字 token 视为无序的、非结构化的类别（categories）。在这种视角下，数字 "19" 和 "20" 与数字 "1" 和 "100" 在向量空间中被视为同等距离的独立类别，完全忽略了数值本身固有的度量结构（metric structure）和大小关系。因此，模型在预测 "19" 时，即使预测了 "20"，交叉熵也会给予与预测 "1" 相同的惩罚，无法利用数值间的邻近性来提供梯度信号。

为了解决这一不匹配问题，研究人员提出了基于平滑最大均值差异（Smooth Maximum Mean Discrepancy, SMMD）的新方法，旨在将数值的度量结构引入训练过程。

核心内容

本文提出了一种名为 Smooth MMD (SMMD) 的新损失函数，用于增强 LLM 在数值目标任务中的表现。SMMD 建立在经典的 MMD 理论基础上，通过引入基于数值距离的核函数（value-distance kernels）和基于图的平滑机制，重构了预测分布与目标分布之间的对齐方式。

1. 方法论：从经典 MMD 到 SMMD

经典的最大均值差异（MMD）是一种用于衡量两个概率分布之间距离的非参数统计量。SMMD 对其进行了两项关键改进：

• 基于数值子词汇表的核函数定义： SMMD 并非在所有 token 上运行，而是专注于一个“数值子词汇表”（numeric sub-vocabulary）。它定义了特定的核函数，该函数不仅考虑 token 的语义相似性，还显式地融入了数值之间的距离。这意味着，如果模型预测了 "19" 而真实值是 "20"，核函数会捕捉到这两个数值在度量空间中的接近性，从而比预测 "1" 提供更强的梯度信号。

• 基于图的平滑机制（Graph-based Smoothness）： 为了进一步鼓励局部一致性，SMMD 在由核函数诱导的图结构上对“预测-目标残差”进行平滑处理。这种平滑机制确保了模型在数值空间中的预测变化是连续且稳定的，避免了因 token 离散化带来的剧烈波动。

2. 技术实现细节

• 核匹配（Kernel Matching）：SMMD 通过核匹配将预测的数值分布对齐到目标分布。这允许模型在潜在空间中直接优化分布的距离，而非仅仅优化单个 token 的对数似然。
• 残差平滑：通过在诱导的核图上平滑预测与目标之间的残差，SMMD 强制模型在数值相近的 token 之间保持行为的一致性。这有助于缓解 LLM 在处理大数或复杂算术时的不稳定性。

3. 实验评估

研究者在四个具有数值目标的任务上评估了 SMMD 的效果，涵盖了多种开源权重的 LLM 和视觉语言模型（VLM）骨干网络：

1. 数学推理（Mathematical Reasoning）：解决需要多步逻辑和数值计算的数学问题。
2. 算术计算（Arithmetic Calculation）：执行基本的加减乘除运算。
3. 时钟时间识别（Clock-time Recognition）：从图像或文本中识别并输出准确的时间。
4. 图表问答（Chart Question Answering）：从可视化图表中提取精确的数值数据。

实验结果表明，SMMD 在以上所有任务中均一致地提高了准确率，且效果优于标准的交叉熵损失以及近期提出的其他数值目标损失函数。

关键要点

• 痛点明确：标准交叉熵损失将数字视为无序类别，忽略了数值间的度量关系，导致 LLM 在数值任务中表现不佳。
• 核心创新：提出了 SMMD（平滑最大均值差异） 损失函数，结合了基于数值距离的核函数和基于图的平滑技术。
• 机制优势：
  • 度量感知：通过核函数捕捉数值间的距离，使模型在预测接近值时获得更合理的梯度。
  • 局部一致性：通过图平滑机制，确保数值空间中的预测变化具有连续性。
• 广泛适用性：该方法适用于多种架构，包括纯文本 LLM 和多模态 VLM。
• 性能提升：在数学推理、算术、时间识别和图表问答四个基准任务上，SMMD 均显著优于基线方法。
• 互补效应：分析显示，MMD 的全局分布对齐能力与平滑机制的局部一致性约束具有互补效应，共同提升了模型性能。
• 开源支持：相关代码已公开，便于社区复现和进一步研究。

意义与影响

这项工作为解决 LLM 长期存在的“数值不可靠”问题提供了新的视角和有效的工具。

1. 突破训练目标局限：传统上，人们倾向于通过增加数据量或调整架构来改善数值能力。本文证明，通过改进训练目标函数，引入数值本身的度量结构，可以在不增加模型复杂度的情况下显著提升性能。这为后续研究提供了重要的方法论指导。
2. 推动高精度应用落地：数值预测能力的提升直接有助于 LLM 在金融、科学计算、工程仿真等对精度要求极高的领域的应用。SMMD 作为一种通用的损失函数，可以轻松地集成到现有的预训练或微调流程中。
3. 多模态能力的增强：由于实验涵盖了 VLM（视觉语言模型），该方法的成果也表明，通过改进数值对齐，可以提升模型从图表、图像等非文本数据中精确提取数值信息的能力，这对增强多模态理解至关重要。
4. 理论贡献：将 MMD 和图平滑技术引入 NLM 的数值训练，丰富了生成模型损失函数的理论工具箱，展示了核方法在处理结构化数值输出时的潜力。

总之，SMMD 不仅是一个性能提升工具，更是对“如何让 LLM 理解数字”这一根本问题的深刻回应。它强调了在训练阶段显式建模数值关系的重要性，为构建更可靠、更精确的下一代 AI 系统奠定了基础。