TACO: Ternary Absolute-max Column-wise One-sparse Optimizer for LLM Fine-Tuning

📄 arXiv: 2610.02199v1 📥 PDF

作者: Jichao Jiang, Cristian McGee, El Houcine Bergou, Hanqin Cai, Aritra Dutta

分类: cs.LG, math.OC

发布日期: 2026-10-01

备注: 24 pages, 7 figures, 10 tables. Code available at https://github.com/Jichao2357/TACO_optimizer

💡 一句话要点

提出TACO优化器以解决大语言模型微调中的内存瓶颈问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 优化器 大语言模型 微调 内存管理 深度学习 机器学习

📋 核心要点

  1. 现有的优化器在大语言模型微调中面临内存开销过大的问题，限制了模型的规模和训练效率。
  2. TACO优化器通过选择每列最大幅度条目的符号，计算精确的最陡下降方向，从而显著降低优化器状态内存。
  3. 实验表明，TACO在OPT-13B模型上持久优化器状态减少至0.16GB，峰值训练内存减少至27.5GB，同时保持了较高的准确性。

📝 摘要（中文）

大语言模型（LLMs）的全参数微调会产生巨大的优化器状态内存开销，限制了现代GPU上可处理的模型规模。现有方法要么压缩优化器状态，要么放弃一阶梯度，或在保持稠密状态的同时改变更新几何。本文提出的TACO（Ternary Absolute-max Column-wise One-sparse Optimizer）优化器，在保持一阶梯度的同时，通过选择二维权重矩阵每列中最大幅度条目的符号，显著减少了优化器状态内存。实验结果表明，TACO相较于AdamW8bit，持久优化器状态减少了174倍，峰值训练内存减少了2.9倍，同时保持了相似的准确性和运行时间。

🔬 方法详解

问题定义：本文旨在解决大语言模型微调中优化器状态内存开销过大的问题。现有方法如Muon优化器虽然减少了内存，但其几何特性与AdamW不同，可能导致性能下降。

核心思路：TACO优化器通过选择二维权重矩阵每列中最大幅度条目的符号，计算出在维度归一化的$1 o1$算子范数下的最陡下降方向，从而在保持一阶梯度的同时，显著降低优化器状态内存。

技术框架：TACO的整体架构包括权重矩阵的列处理、符号选择和梯度计算等模块。通过这些模块的协同工作，TACO能够在微调过程中有效管理内存。

关键创新：TACO的核心创新在于其在计算最陡下降方向时的符号选择策略，这一策略与传统的优化方法相比，显著降低了内存需求，同时保持了训练的准确性和效率。

关键设计：TACO优化器的设计中，关键参数包括低精度梯度组件的选择，以及优化器状态的存储策略，这些设计确保了在大规模模型微调中的高效性。

🖼️ 关键图片

img_0
img_1
img_2

📊 实验亮点

实验结果显示，TACO优化器在OPT-13B模型上持久优化器状态从27.7GB减少至0.16GB，峰值训练内存从80.6GB减少至27.5GB，性能提升显著，同时保持了与AdamW相当的准确性和运行时间。

🎯 应用场景

TACO优化器在大语言模型的微调中具有广泛的应用潜力，尤其是在需要处理大规模模型时。其显著降低的内存需求使得在单个GPU上进行全参数微调成为可能，推动了自然语言处理、对话系统等领域的研究和应用。

📄 摘要（原文）

Full-parameter fine-tuning of large language models (LLMs) incurs substantial optimizer state memory overhead, limiting the model sizes that fit on modern GPUs. Existing approaches either compress optimizer state, abandon first-order gradients, or change the update geometry while retaining dense state. The recently introduced Muon optimizer reduces optimizer memory through matrix-valued updates. Still, its geometry differs from AdamW and can lead to performance degradation when fine-tuning AdamW-pretrained models. To reduce optimizer memory without sacrificing accuracy or computational efficiency in LLM fine-tuning, we propose Ternary Absolute-max Column-wise One-sparse optimizer, or TACO, which follows Muon's operator-norm steepest-descent view but takes the geometric route further. TACO computes the exact steepest-descent direction under a dimension-normalized $1\to1$ operator norm by selecting the sign of the largest magnitude entry in each column of two-dimensional weight matrices. This retains first-order gradients while making optimizer state memory nearly negligible. Our practical TACO optimizer maintains only a small set of low precision gradient components per column, reducing persistent optimizer state by $174\times$ relative to AdamW8bit (from 27.7 GB to 0.16 GB) and peak training memory by $2.9\times$ (from 80.6 GB to 27.5 GB) on OPT-13B, while achieving comparable accuracy and runtime. TACO further enables full-parameter fine-tuning of 30-32B-parameter models on a single 80 GB H100 GPU across multiple model families and tasks.