PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2610.02199v1cs.LGmath.OC

TL;DR

Full-parameter LLM fine-tuning has substantial optimizer-state memory overhead, and existing approaches trade off first-order information, memory, or update geometry. TACO derives a column-wise sparse update from a dimension-normalized 1→1 operator norm and maintains sparse, low-precision gradient history. It reports reduced persistent and peak memory with competitive performance, and enables fine-tuning models up to 32B parameters on one 80 GB H100.

  • Problem

    Full-parameter LLM fine-tuning incurs substantial optimizer-state memory overhead, while existing approaches expose a trade-off between first-order training memory and optimization efficiency.

    Source: §1 · pp. 1–2

  • Method

    TACO derives a ternary, column-wise one-sparse update from the dimension-normalized 1→1 operator-norm steepest-descent direction and maintains sparse, low-precision gradient history.

    Source: §1 · p. 2, §7 · p. 9

  • Results

    TACO occupies the low-memory end of the accuracy–memory trade-off while retaining competitive performance and practical throughput.

    Source: §6 · p. 8

  • Takeaways & Limitations

    TACO enables full-parameter fine-tuning of models up to 32B parameters on a single 80 GB H100 and transfers across multiple model families and tasks.

    Source: §1 · pp. 2–3

    The continuation-geometry comparison is established only under a fixed-input objective and a common target-output assumption.

    Source: §4 · p. 5

Abstract

from arXiv · show

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.

1 INTRODUCTION

Full-parameter LLM fine-tuning carries substantial optimizer-state memory costs, while existing memory-saving approaches trade off first-order information, dense state, or optimization geometry. TACO uses geometry-induced sparse updates and sparse gradient histories to reduce memory while retaining first-order gradients, with experiments reporting competitive performance and fine-tuning at up to 32B parameters on one 80 GB H100.

1 INTRODUCTION

  • For a 13B model, AdamW's two FP32 moment buffers add about 104 GB, bringing memory to roughly 156 GB before activations, workspaces, or an optional master copy.The model itself occupies about 26 GB in BF16 and gradients contribute about 26 GB; optimizer states alone require roughly four times BF16 parameter memory.Source: §1 · p. 1
  • Muon reduces optimizer memory with a dense momentum matrix, but switching an AdamW-pretrained model to Muon can degrade performance through optimizer mismatch.Zeroth-order methods can approach inference-level memory but require additional forward evaluations, which can increase computation and communication overhead.Source: §1 · p. 2
  • TACO addresses the trade-off between first-order training's memory footprint and standard gradients' optimization efficiency by seeking nearly negligible persistent optimizer-state memory.The motivating question is whether first-order information and competitive wall-clock efficiency can be retained together.Source: §1 · p. 2
  • TACO derives a ternary, column-wise Top-1 update from the exact steepest-descent direction under a dimension-normalized 1→1 operator norm.Its sparsity follows from the optimization geometry rather than post-hoc gradient sparsification.Source: §1 · pp. 2–3
  • Practical TACO keeps low-precision dynamic gradient heavy hitters per column, reducing persistent state from O(mn) to O(n) while retaining first-order gradients.Immediate gradient processing supports this sparse-history design without requiring an accurate dense model of gradient history.Source: §1 · p. 3
  • TACO achieves competitive task performance with substantially reduced training memory, enables OPT-30B and Qwen3-32B fine-tuning on one 80 GB H100, and transfers across five model families.The paper also reports ablations of gradient-history EMA decay, heavy-hitter threshold, update sparsity, and optimizer step budget.Source: §1 · p. 3

2 RELATED WORK

Related work reduces fine-tuning memory through zeroth-order updates or by compressing first-order optimizer state and gradients. Other approaches change the optimization geometry or communicate sign-based updates.

2 RELATED WORK

  • Zeroth-order methods estimate updates from loss changes under parameter perturbations; variants use regenerated perturbations, sparse perturbations, low-dimensional subspaces, or diagonal Hessian information.These approaches include MeZO, Sparse MeZO, ZO-Muon, and HiZOO.Source: §2 · p. 3
  • First-order methods reduce state with factored second moments, quantization, compact master weights, or low-rank gradient and optimizer-state representations.Examples include Adafactor, CAME, 8-bit Adam (Dettmers et al., 2021), FlashAdamW (Gonzalez Ortiz et al., 2026), GaLore, and APOLLO-Mini.Source: §2 · p. 3
  • Operator-norm and linear-minimization-oracle work frames matrix optimization as steepest descent under operator norms and develops stochastic, layer-aware geometries.Bernstein & Newhouse (2024) characterize the 1→p family through maximum column norms; Pethick et al. (2025) instantiate stochastic LMO optimization in Scion.Source: §2 · p. 3
  • Sign-based methods include ternary or coordinate-wise gradient signs for communication compression, and communicating signs of locally constructed Muon directions.Examples are TernGrad, signSGD (Bernstein et al., 2018), GRACE, and SignMuon (Mishra et al., 2026).Source: §2 · p. 3

TACO: TERNARY ABSOLUTE-MAX COLUMN-WISE ONE-SPARSE OPTIMIZER

TACO formulates updates as steepest descent under a dimension-normalized 1→1 operator norm. Because the constraint separates across columns, its exact update selects a signed largest-magnitude gradient coordinate per column, yielding sparse updates and a vanilla convergence guarantee.

TACO: TERNARY ABSOLUTE-MAX COLUMN-WISE ONE-SPARSE OPTIMIZER

  • Under TACO's dimension-normalized 1→1 operator geometry, each column has its own ℓ1 budget rather than a single global element-wise ℓ1 constraint.The maximum-column identity is the p = 1 specialization of the general 1→p characterization and underlies TACO's column-wise one-sparse, low-state design.Source: TACO: TERNARY ABSOLUTE-MAX COLUMN-WISE… · p. 4
  • Since the constraint separates across columns, each column selects a largest-magnitude gradient coordinate, with ties resolved by a fixed rule.This yields the column-wise one-sparse direction used by TACO.Source: TACO: TERNARY ABSOLUTE-MAX COLUMN-WISE… · p. 4
  • The exact geometry-derived update is ternary, column-wise, and one-sparse, with at most n nonzeros in an m × n matrix.The m/n factor follows from dimension normalization rather than adding a layerwise scaling hyperparameter; this differs from conventional Top-1 sparsification.Source: TACO: TERNARY ABSOLUTE-MAX COLUMN-WISE… · p. 4
  • Vanilla TACO reaches an ε-stationary point in O(ε^-2) iterations under lower-boundedness and smoothness assumptions.An ε-stationary point is defined using the dual norm associated with TACO's geometry.Source: TACO: TERNARY ABSOLUTE-MAX COLUMN-WISE… · p. 5

4 CONTINUATION GEOMETRY, IMPLICIT BIAS, AND µP-SCALING OF TACO

Under a simplified fixed-input setting, TACO and Adam follow the same continuation geometry, while Muon follows a different path. TACO’s update also satisfies standard µP spectral scaling under a balanced-winner assumption.

Continuation geometry and implicit bias

  • The analysis compares optimizer paths when the objective depends on a weight matrix through its action on a fixed input and methods share a target output.Under these assumptions, continuation geometry identifies how each optimizer represents the same output correction in parameter space.Source: §4 · p. 5
  • Under the fixed-input assumptions, TACO and Adam follow the same continuation family and reach the same solution, whereas Muon generally follows a different path and solution.The authors note this Adam-like behavior may help avoid optimizer mismatch when fine-tuning AdamW-pretrained models.Source: §4 · p. 5

µP-scaling

  • TACO’s update satisfies standard µP spectral scaling when column winners remain reasonably balanced across rows, providing a theoretical basis for learning-rate transfer across widths.Empirically, the same learning rate remains optimal across model scales.Source: §4 · p. 6
  • Together, the authors characterize TACO as combining Adam-like continuation geometry and µP-related reduced hyperparameter sensitivity with almost no persistent optimizer state.These properties are stated within the paper’s theoretical setting and balanced-winner condition.Source: §4 · p. 6

5 PRACTICAL DESIGN OF TACO

Practical TACO stabilizes its column-wise Top-1 update using a sparse, low-precision history of gradient heavy hitters rather than a dense gradient history. Immediate gradient processing and releasing dense gradients help keep persistent state and peak training memory low.

Gradient history and heavy hitters

  • TACO needs only enough gradient-history information to recover each column’s dominant coordinate, avoiding an accurate dense history model.The practical design addresses the instability of selecting winners directly from a single minibatch gradient without restoring a dense state.Source: §5 · p. 6
  • The practical history retains the k largest-magnitude entries per column after an EMA update, using k = 16 and µ = 0.95 in experiments.Retained values use FP8 E4M3 with int32 row indices; k controls future winner selection, not the column-wise Top-1 update.Source: §5 · p. 6

Persistent state memory

  • For an m × n matrix, TACO stores 5kn bytes with int32 indices, reducing persistent history complexity from O(mn) to O(kn).Each retained entry uses one byte for its FP8 value and four bytes for its int32 row index.Source: §5 · p. 6

Gradient history and heavy hitters

  • The workflow forms a dense autograd gradient, retains column-wise heavy hitters as Mt, applies S to keep each column’s largest-magnitude sign, and persists only Mt.Figure 2 illustrates H2, while experiments use H16.Source: Figure 2 · p. 7

Backward pass memory management

  • Immediate in-place gradient processing, gradient accumulation of 1, and releasing each dense gradient prevent gradients from persisting across microbatches.Gradient checkpointing also reduces activation memory; the authors identify hooks and checkpointing as standard systems techniques, not contributions.Source: §5 · p. 7

6 EXPERIMENTS

Across OPT-13B downstream tasks, TACO uses the least GPU memory while maintaining competitive performance; on SST-2 it reaches 94.2% accuracy at 276 tokens/second with 27.5 GB peak memory. Experiments also report scaling within a single 80 GB H100, architecture transfer, and optimizer ablations.

Memory, trained model quality, and runtime on OPT-13B

  • TACO uses 27.5 GB peak memory versus 52.6 GB for Adafactor, while reaching 94.2% SST-2 accuracy at 276 tokens/second.Across the OPT-13B tasks, TACO consistently requires the least peak GPU memory while maintaining competitive performance.Source: §6 · pp. 7–8, Figure 3 · p. 8

Scaling across model families

  • TACO remains within one 80 GB H100 as OPT scales from 1.3B to 30B parameters and Qwen3 from 8B to 32B, using the same hyperparameters across sizes.OPT SST-2 accuracy is broadly stable, RTE generally improves with scale, and BoolQ is non-monotonic.Source: §6 · p. 8, Figure 5 · p. 9

Transfer across architectures

  • Across four additional architectures, memory efficiency is consistent but performance varies by task; all four OPT scales attain best RTE validation accuracy at 3×10−5.Pythia-12B and Mistral-24B show greater task-dependent variation, particularly on RTE and BoolQ.Source: §6 · p. 8, Figure 6 · p. 9

Gradient history and heavy hitter threshold

  • With k = 16, optimizer state is 38.3 MB, a 34.4× reduction from dense history, and this setting achieves the sweep’s strongest BoolQ and RTE performance.Across eight tasks, µ = 0.95 improves six; the largest gains are MultiRC (+8.9 F1A) and CB (+7.1%), with SQuAD and DROP as exceptions.Source: §6 · pp. 8–9

Optimizer step budget

  • Increasing TACO’s step budget from 2K to 16K improves BoolQ and RTE by +6.5% and +4.3%, respectively, and SQuAD by +2.1 F1.The default 2K-step budget is not TACO’s performance ceiling.Source: §6 · p. 9

Update sparsity

  • Top-1 achieves the strongest final performance across the three evaluated tasks, while denser updates substantially degrade performance, particularly at Top-8 and beyond.The ablation supports retaining the column-wise Top-1 update.Source: §6 · p. 9

7 CONCLUSION

TACO combines a geometry-derived sparse update with a low-precision sparse gradient history for memory-efficient full-parameter fine-tuning. The paper reports competitive performance across scales and families, alongside theoretical results and reproduction materials.

7 CONCLUSION

  • TACO maintains competitive performance with a small memory footprint and enables full-parameter fine-tuning of models up to 32B parameters on one 80 GB H100 GPU.The reported evaluation spans model scales, architectures, and downstream tasks.Source: §7 · p. 9
  • TACO’s dimension-normalized 1→1 operator-norm geometry yields a ternary, column-wise one-sparse update, while practical TACO uses a sparse, low-precision gradient history.The paper also reports a nonconvex convergence guarantee, Adam-proxy equivalence in a simplified fixed-input setting, µP scaling under balanced winner counts, and a sufficient condition for sparse history to preserve the dense-history update.Source: §7 · p. 9
  • The authors report using ChatGPT and Codex as research and writing assistants, while taking responsibility for the final content and results.They state that AI-assisted text, interpretations, and code were reviewed, and code changes were tested against corresponding experiments.Source: AI USE STATEMENT · p. 10
  • Appendix B reports experimental setups and additional optimizer comparisons, scaling results, ablations, and architecture-transfer results; anonymous source code and experiment scripts are available.The paper states that Appendix A contains proofs and additional theoretical details.Source: REPRODUCIBILITY STATEMENT · p. 10

Supplementary Material for TACO: TERNARY ABSOLUTE-MAX COLUMN-WISE ONE-SPARSE OPTIMIZER FOR LLM FINE-TUNING

The supplementary material collects proofs and additional theoretical details alongside experimental setup and results. Its convergence analysis establishes the proof approach for vanilla TACO and describes how the resulting bound balances objective gap and step-size penalty.

Supplementary Material for TACO: TERNARY ABSOLUTE-MAX COLUMN-WISE ONE-SPARSE OPTIMIZER FOR LLM FINE-TUNING

  • Appendix A contains proofs of Theorems 1–4, while Appendix B describes experimental setup and reports additional optimizer comparisons, scaling results, ablations, and architecture-transfer results.Appendix B covers models, datasets, implementation, hyperparameters, and hardware.Source: Supplementary Material for TACO:… · p. 14

A.1 VANILLA TACO CONVERGENCE

  • The vanilla TACO convergence proof uses exact first-order decrease under its direction and a smoothness-based step-size penalty.The proof expresses the first-order decrease through the dual norm of the gradient and the update direction.Source: §A.1 · p. 15
  • The resulting convergence bound balances the initial objective gap against the accumulated step-size penalty.The theorem passage identifies the result as nonconvex convergence under Assumption 1, but the supplied theorem text does not include the full bound.Source: §A.1 · p. 15

A.1.1 MULTILAYER ANALYSIS

The multilayer analysis defines a shared operator-norm constraint across layers and derives a layerwise minimizing direction. Each layer receives a geometry scaling factor, while optimization uses one shared learning rate.

A.1.1 MULTILAYER ANALYSIS

  • Under the multilayer norm constraint, the exact minimizing direction is obtained simultaneously in every layer.The section defines each layer’s norm and combines them using the maximum across layers.Source: §A.1.1 · p. 15
  • Each layer receives a geometry scaling factor mℓ/nℓ, while the optimization learning rate ηt remains a single shared scalar across layers.The paper notes that this distinction suggests the optimization learning rate is invariant to parameter shape.Source: §A.1.1 · p. 16

A.2 CONTINUATION GEOMETRY, IMPLICIT BIAS, AND µP-SCALING ANALYSIS

In the fixed-input setting, TACO and an Adam proxy share continuation geometry and limiting solutions, while a Muon proxy generally follows a different continuation set. The analysis also treats µP scaling and gives a condition under which sparse gradient histories preserve TACO’s dense-history update.

A.2 CONTINUATION GEOMETRY, IMPLICIT BIAS, AND µP-SCALING ANALYSIS

  • For F(W)=ϕ(Wz), the analysis extends prior least-squares implicit-bias work to differentiable objectives and characterizes optimizer continuation geometry and limiting solutions.In the least-squares setting, TACO and the Adam proxy select the same limiting weights as SignGD and inherit its minimum-max-norm correction, while Muon selects a minimum-spectral-norm correction.Source: §A.2 · pp. 16–17
  • TACO and the Adam proxy lie in the same continuation set under the stated assumptions, whereas Muon’s iterates lie in a distinct set.For sequential layerwise fine-tuning of a deep network with other layers fixed, the paper says TACO and the Adam proxy reach the same endpoint at each layer and therefore yield identical final networks.Source: §A.2 · pp. 16–17
  • The µP spectral-scaling result depends on how concentrated TACO’s column-wise winners are across rows, with balanced winner counts providing the stated scaling condition.If one row wins every column, the spectral norm exceeds the µP target by a factor of √m.Source: §A.2 · pp. 18–19

A.3 GRADIENT HISTORY AND HEAVY HITTER ANALYSIS

  • Sparse history preserves the exact dense-history TACO direction when its approximation error stays below half the dense-history winner margin in every column.Exact reconstruction of dense history is therefore unnecessary under this sufficient separation condition.Source: §A.3 · p. 19

B EXPERIMENT SETUP AND ADDITIONAL EXPERIMENT RESULTS

Appendix B describes the shared and experiment-specific protocol, compares optimizers on OPT-13B, and reports scaling results for OPT and Qwen3.

B EXPERIMENT SETUP AND ADDITIONAL EXPERIMENT RESULTS

  • Experiments span five model families and downstream tasks; E1 selects learning rates from a 10^-6–10^-2 sweep on SST-2, RTE, and BoolQ.All runs use the same NVIDIA H100 80 GB GPU model; E1 uses a common seed, batch size 1, and a 500-step budget for the sweep.Source: B EXPERIMENT SETUP AND ADDITIONAL… · pp. 19–20

B.1 OPTIMIZER COMPARISON

  • On OPT-13B, TACO has the lowest peak GPU memory while maintaining competitive downstream performance and practical training throughput.The comparison covers eight downstream tasks and reports SST-2 efficiency measures including persistent optimizer state, peak memory, runtime, and throughput.Source: §B.1 · pp. 20, 22, Table 3 · p. 22, Table 4 · p. 22

B.2 MODEL SCALING

  • Table 5 reports OPT and Qwen3 scaling results on SST-2, RTE, and BoolQ, including downstream performance and peak GPU memory averaged over three seeds.The supplied passages identify the table's scope but do not state its cell values.Source: §B.2 · p. 22, Table 5 · p. 22

B.3 ABLATION STUDIES

The ablations examine training budget, gradient history, update sparsity, batch size, and learning-rate transfer, identifying the settings that preserve or improve performance.

B.3 ABLATION STUDIES

  • 16K steps versus 2K improves BoolQ and RTE by 6.5 and 4.3 percentage points, respectively, and SQuAD by 2.1 F1 points.Most RTE and SQuAD gains arrive by 4K–8K steps, while BoolQ continues improving through 16K; the default 2K budget is not TACO's performance ceiling.Source: §B.3 · p. 23, Table 6 · p. 23
  • µ = 0.95 improves all six classification and multiple-choice tasks, led by MultiRC (+8.9 F1A) and CB (+7.1 percentage points), while SQuAD and DROP favor µ = 0.The generation-task margins favoring µ = 0 are 2.6 and 0.4 F1 points for SQuAD and DROP, respectively.Source: §B.3 · p. 23, Table 7 · p. 23
  • At k = 16, TACO achieves the strongest BoolQ and RTE performance with 38.3 MB of optimizer state, a 34.4× reduction from dense history's 1317.7 MB.Performance is broadly comparable for moderately sparse histories; SQuAD is slightly better with dense or k = 32 history, and very small k begins to degrade performance.Source: §B.3 · p. 23, Table 8 · p. 23
  • Top-1 achieves the strongest final performance across all three tasks; Top-2 remains comparable, while denser updates degrade, especially at Top-8 and beyond.Top-2 and Top-4 can progress faster early on some tasks, but this advantage does not improve final performance; Top-1 is also the 1→1 steepest-descent solution.Source: §B.3 · pp. 23–24, Table 9 · p. 24
  • Larger batches increase peak GPU memory without consistently improving accuracy, so batch size 4 remains the default.The ablation evaluates batch sizes 2–128 on Qwen3-14B and Qwen3-32B using SST-2 and RTE.Source: §B.3 · p. 24, Table 10 · p. 24
  • Every evaluated model-dataset configuration selects 3 × 10−5 as its best learning rate, supporting learning-rate transfer across model scales.The sweep covers Qwen3-8B/14B/32B on SST-2 and OPT-1.3B/6.7B/13B/30B on SST-2 and RTE.Source: Figure 7 · p. 24, §B.3 · p. 24

Related papers

Loading 2610.02199v1…