arXiv:cs.AI· Tian Liang, Zishan Shao, Yiran Chen·· 5 小时前AI 评分34
如何在压缩扩散语言模型时保留数学推理能力:轨迹感知低秩近似方法
Preserving Mathematical Reasoning in Compressed Diffusion Language Models via Trajectory-Aware Low-Rank Approximation
AI 导读
针对扩散语言模型(dLLM)压缩中校准态与推理态不匹配的问题,研究者提出轨迹感知低秩目标,并设计 Traj-MC 通过蒙特卡洛采样估计轨迹二阶矩,实现采样态最优性与总体一致性。在相同压缩预算下,轨迹感知校准比干净校准更好地重建生成轨迹,并在数学推理基准上保留显著更多推理能力。代码已开源。
正文
Abstract:Diffusion language model (dLLM) compression faces a known challenge because calibration is typically performed on clean, fully visible activations, whereas inference traverses partially masked intermediate states. For low-rank compression, this raises two questions. First, can low-rank optimality still be characterized when approximation quality is measured over trajectory-distributed states, and second, does the choice of calibration states affect mathematical reasoning preservation under compression? We address these questions by formulating a trajectory-aware low-rank objective over corruption levels and masking realizations. To estimate this objective efficiently, we propose Traj-MC, which estimates the trajectory second moment through Monte Carlo sampling and yields exact sampled-state optimality and population consistency. Under matched compression budgets, trajectory-aware calibration improves reconstruction over the generation trajectory and preserves substantially more mathematical reasoning than clean calibration on mathematical reasoning benchmarks. Our results connect trajectory-aware low-rank optimality to the reasoning capability retained after dLLM compression. Our code is available at: this https URL.
| Subjects: | Artificial Intelligence (cs.AI) |
| Cite as: | arXiv:2610.03326 [cs.AI] |
| (or arXiv:2610.03326v1 [cs.AI] for this version) | |
| https://doi.org/10.48550/arXiv.2610.03326 arXiv-issued DOI via DataCite (pending registration) |
Submission history
From: Zishan Shao [view email]
[v1]
Fri, 2 Oct 2026 14:00:55 UTC (51 KB)
来源:arXiv:cs.AI · arxiv.org