Source-linked AI summary
TerMeZO: Ternary Sparse Zeroth-Order Optimization for Fine-tuning BitNet Models at the Edge
Houssem Sifaou, Prabodh Katti, Bipin Rajendran, Osvaldo Simeone
TL;DR
Fine-tuning BitNet with MeZO still requires perturbing and updating full-precision latent weights, so its memory footprint no longer matches inference; existing sparse ZO methods do not select weights using ternary structure or may require FO gradients. TerMeZO selects boundary-adjacent latent weights, and experiments report that it matches or exceeds full-parameter MeZO with substantially lower fine-tuning memory.
Problem
MeZO avoids storing gradients and activations, but BitNet fine-tuning still perturbs and updates every full-precision latent weight, so its memory footprint no longer matches inference.
Source: §1 · p. 1Method
Results
TerMeZO matches or exceeds full-parameter MeZO across BitNet models from 1B to 3B parameters and downstream tasks while substantially reducing fine-tuning memory.
Source: §1 · p. 2Takeaways & Limitations
TerMeZO's sparse updates make BitNet fine-tuning on edge devices possible with near-inference memory footprint.
Source: §7 · p. 9The convergence analysis assumes a bounded objective, Lipschitz loss and gradients, a masked estimator capturing a c-fraction of gradient energy, and bounded minibatch-gradient variance.
Source: §4 · p. 6
Abstract
from arXiv · showhide
Fine-tuning anguage models (LLMs) with first-order optimizers requires a memory several times larger than that required for inference. Memory-efficient zeroth-order optimization (MeZO) sidesteps this cost by estimating gradients from forward passes only. However, for BitNet architectures, a family of LLMs with ternary {-1,0,1\} weights and 8-bit activations, fine-tuning requires updating full-precision latent weights, and thus the memory footprint of MeZO no longer matches that of inference. A promising solution is to finetune only a subset of the latent weights, but existing sparse zeroth-order (ZO) methods either ignore the ternary structure or require first-order gradient information to build a sparse mask, which is at odds with the purpose of ZO fine-tuning. We propose TerMeZO, a sparse MeZO scheme that exploits the geometry of the ternary quantizer itself to identify the latent weights that are more likely to change values during fine-tuning, at no additional data or memory cost. Our convergence analysis shows that TerMeZO can converge faster than full-parameter MeZO, owing to its optimized reduction of the fine-tuning effective dimension. We run extensive experiments on BitNet models ranging from 1B to 3B parameters, spanning classification, instruction-following, and mathematical reasoning tasks. TerMeZO matches or exceeds the performance of full-parameter MeZO while substantially reducing the fine-tuning memory footprint.
1 Introduction
Fine-tuning BitNet models can erase their inference memory advantage, and existing sparse ZO approaches either ignore ternary structure or rely on first-order gradients. TerMeZO selects boundary-adjacent latent weights without FO gradients; experiments report performance matching or exceeding full-parameter MeZO with reduced memory.
1 Introduction
- Fine-tuning with first-order optimizers uses several times inference memory, while MeZO's inference-level memory parity for standard LLMs does not carry over to BitNet's full-precision latent weights.MeZO perturbs and updates every latent weight at each step.Source: §1 · p. 1
- Prior sparse ZO methods either select parameters using FO gradients, use weight magnitude without accounting for ternary boundaries, or tune only quantization scales.The cited work describes scale-only QZO as too restrictive for harder tasks requiring new knowledge.Source: §1 · p. 2
- TerMeZO updates latent weights nearest ternary quantization boundaries, which are most likely to flip, plus embeddings, normalization layers, and the language-modeling head.Its selection requires a linear scan and no FO gradients.Source: §1 · p. 2
- TerMeZO matches or exceeds full-parameter MeZO across 1B–3B BitNet models and downstream tasks while reducing fine-tuning memory.Experiments include classification, instruction-following, and mathematical reasoning tasks.Source: §1 · p. 2
2 Background and Related Work
The background reviews PEFT, zeroth-order optimization, and extreme quantization as approaches to memory-constrained fine-tuning. TerMeZO is positioned as a sparse ZO method that selects updates from ternary quantization structure rather than FO gradients or scale-only tuning.
Parameter-Efficient Fine-Tuning
- PEFT freezes pretrained weights and adds trainable parameters, but retains much of the activation and gradient overhead that ZO avoids.The review groups PEFT into adapters, soft prompts and prefixes, low-rank reparameterization, and sparse fine-tuning.Source: §2 · pp. 2–3
Zeroth-Order Optimization and MeZO
- MeZO uses SPSA to estimate gradients from function evaluations without backpropagation, while standard LLM fine-tuning can use roughly inference memory.The estimator requires two function evaluations per step.Source: §2 · p. 3
Fine-Tuning Quantized Models
- BitNet trains LLMs with ternary weights and 8-bit activations, while the straight-through estimator treats the non-differentiable quantizer as identity during backpropagation.The review notes that STE is a biased gradient surrogate and can interact poorly with low-precision optimizer states.Source: §2 · p. 3
- TerMeZO selects a sparse full-precision working set from ternary quantization, avoiding FO mask construction and QZO's scale-only update constraint.The related-work comparison describes sparse MeZO variants and TerMeZO as restricting perturbations and updates to a sparse working set.Source: §2 · p. 4
3 TerMeZO
TerMeZO applies sparse MeZO to BitNet by perturbing selected latent weights in ternary linear projections while updating the remaining full-precision parameters. It selects weights near quantization boundaries and uses ZO updates for the active set and full-precision parameters.
3 TerMeZO
- BitNet keeps embeddings, normalization modules, and the language-modeling head full precision, while ternary quantization applies to transformer attention and MLP projections.The ternary latent weights dominate parameter count; for Falcon-E-3B-Instruct, p/d = 0.044.Source: §3 · p. 4
- TerMeZO measures each latent weight's distance to its quantization boundary and selects a target fraction ρ0 of weights using a global distance threshold.The threshold is set from the ρ0-quantile of distances across projection layers.Source: §3 · p. 5
- Weights far from either threshold are effectively frozen in ternary space, whereas weights near thresholds are most sensitive to adaptation.The selected subset and corresponding quantization threshold are updated during fine-tuning.Source: §3 · pp. 4–5
- The initial masks are computed offline in time linear in the number of latent weights, and the active set can only shrink during fine-tuning.A weight that leaves the threshold band is permanently frozen for the remainder of fine-tuning.Source: §3 · p. 5
- At each iteration, TerMeZO performs a standard ZO update on selected latent weights and full-precision parameters using two-sided SPSA estimates averaged over K perturbations.The perturbations use a shared mini-batch and scale ϵ.Source: §3 · p. 5
4 Convergence Analysis
TerMeZO’s convergence analysis studies a Gaussian-smoothed objective because the latent-weight objective is non-differentiable. Under stated boundedness, smoothness, gradient-energy, and minibatch-variance assumptions, its convergence bound depends on active latent weights and boundary crossings.
4 Convergence Analysis
- The analysis uses Gaussian SPSA to study convergence to a stationary point of a smoothed objective, since the latent-weight objective is non-differentiable.The stated analysis assumes one perturbation per iteration; the authors say extension to multiple perturbations follows by standard averaging arguments.Source: §4 · p. 6
- The theorem’s assumptions include bounded loss, Lipschitz behavior, a masked estimator capturing a c-fraction of latent-weight gradient energy, and bounded minibatch-gradient variance.The analysis defines k_t as the active-coordinate count and q_t as the average fraction of active weights that change after perturbation.Source: §4 · p. 6
- TerMeZO’s bound controls the time-averaged gradient norm and includes terms governed by active latent weights and boundary crossings.The theorem assumes constant step sizes and Assumptions 1–4.Source: §4 · p. 6
- The rate recovers sparse MeZO’s O(k/T) rate when ϵ² = 1/√k₀, with the stated perturbation scale matching typical practice for k₀ from 10⁸ to 10¹⁰.The result is discussed in relation to the number of active latent weights and the perturbation scale.Source: §4 · pp. 6–7
5 Memory-Footprint Analysis
TerMeZO stores full-precision copies only for active latent weights, with an additional sparsification mask. GPU measurements on Falcon-E models show substantially lower fine-tuning memory than full-parameter MeZO, approaching inference-level memory.
5 Memory-Footprint Analysis
- TerMeZO retains full-precision copies only for active latent weights, while also storing a sparsification mask.Frozen latent weights are downgraded from full precision to ternary storage; the mask may use a dense bit-mask or an active-index list.Source: §5 · p. 7
- At sparsity ρ₀ = 0.05, TerMeZO reduces memory roughly 3× versus MeZO, and at ρ₀ = 0.01, around 4×, bringing fine-tuning close to inference cost.Measurements cover Falcon-E-1B and Falcon-E-3B on an H100 GPU; the table also reports estimates, assuming 2-bit ternary-weight storage.Source: §5 · p. 7, Table 2 · p. 7
6 Empirical Results
Across classification, instruction-following, and reasoning tasks, TerMeZO generally matches or exceeds full-parameter MeZO while using fewer trainable latent weights. On harder tasks, it closely tracks full-parameter MeZO, while outperforming magnitude-based sparse selection and requiring less memory than first-order fine-tuning.
GLUE/SuperGLUE Classification Tasks
- The evaluation uses six GLUE and SuperGLUE classification benchmarks with validation accuracy, and runs ZO methods for equal optimization steps.Table 3 fixes TerMeZO's trainable latent-weight fraction at ρ0 = 0.05.Source: §6 · p. 8, Table 3 · p. 8
- On most classification tasks TerMeZO matches full-parameter MeZO, and on some it performs better.QZO also performs well on these relatively simple tasks, while LoRA-MeZO lags behind.Source: §6 · p. 8, Table 3 · p. 8
Instruction Following and Reasoning Tasks
- On GSM8K and MagiCoder, TerMeZO closely tracks full-parameter MeZO, while QZO and magnitude-selected S-MeZO fall behind.FIM-S-MeZO matches TerMeZO but requires initial SGD steps on a separate dataset and has a memory footprint comparable to BP.Source: Table 4 · p. 8, §6 · p. 9
- TerMeZO's lower memory footprint allows a 3B model to outperform a 1B BP-fine-tuned model under a tighter memory budget.The paper reports a consistent performance gap between BP and TerMeZO at a given model size.Source: §6 · p. 9
- TerMeZO remains stable at learning rates where FIM-S-MeZO can diverge, for the same fraction of latent weights.The authors attribute this to FIM-S-MeZO selecting weights far from quantization boundaries, where perturbations yield noise.Source: §6 · p. 9
7 Conclusion
TerMeZO uses ternary weight geometry to select a sparse set of latent weights for BitNet fine-tuning, and experiments show performance comparable to or better than full-parameter MeZO with reduced memory. The approach is evaluated on 1B–3B models; practical memory savings require custom hardware-dependent kernels.
7 Conclusion
- Across datasets, TerMeZO recovers and sometimes exceeds full-parameter MeZO performance while substantially reducing fine-tuning memory.The paper evaluates BitNet models ranging from 1B to 3B parameters.Source: §7 · p. 9
- TerMeZO's sparse mask exploits ternary weight geometry and requires only a pass through latent weights to evaluate.The paper also provides a convergence analysis of the method.Source: §7 · p. 9
- Practical memory savings require custom, hardware-dependent kernels, and the approach has only been evaluated on 1B–3B models.Larger-scale validation and real-edge-hardware validation are left for future work.Source: §7 · p. 9
- The paper identifies BitNet linear-attention or hybrid-attention models as a future setting for jointly addressing KV-cache and latent-weight memory.The authors note the relevance of this direction to RL fine-tuning with long generation rollouts.Source: §7 · p. 9
A.1 PEFT methods
The related work surveys parameter-efficient fine-tuning, zeroth-order optimization, and ZO methods for quantized models. It positions TerMeZO against approaches that either ignore ternary quantization geometry or rely on first-order information.
A.1 PEFT methods
- PEFT reduces trainable parameters, with adapter, prompt/prefix, low-rank, and sparse-weight approaches offering different trade-offs.Prompt and prefix methods can be initialization-sensitive and underperform on smaller backbones; LoRA is described as a standard PEFT method.Source: §A.1 · p. 15
- Existing sparse ZO methods either select by latent-weight magnitude, ignoring ternary boundaries, or use first-order gradients to construct masks.The latter can require first-order optimization on the target task or a separate dataset before ZO fine-tuning.Source: §A.1 · p. 15
A.2 Zeroth-Order Optimization and MeZO
- MeZO applies SPSA to fine-tune LLMs at inference-level memory cost, while follow-up methods alter variance, curvature, perturbation structure, or update scheduling.Examples include MeZO-SVRG, HiZOO, LoZO, and block-coordinate variants.Source: §A.2 · p. 15
A.3 ZO Fine-tuning for Quantized Models
- Quantized-model ZO methods include low-precision forward-pass optimization, discrete-space evolution strategies, and tuning quantization scales alone.The paper states that scale-only QZO is insufficient for some fine-tuning tasks.Source: §A.3 · p. 16
Appendix B Empirical Evaluation of the Selection Criterion
The appendix tests whether proximity to ternary decision boundaries identifies latent weights with stronger gradient signals. It reports that this criterion concentrates gradient norm in a small selected subset and outlines the algorithm's distance-based mask initialization.
Appendix B Empirical Evaluation of the Selection Criterion
- Latent weights near ternary decision boundaries have substantially higher gradient magnitudes, supporting boundary proximity as the selection criterion.The analysis compares normalized boundary distance with average per-coordinate gradient magnitude across latent weights and SGD steps.Source: Appendix B · p. 17
- With ρt = 0.01, the selected weights recover a significant fraction of the full gradient norm.This comparison is reported over 100 SGD steps on MultiRC and SST2.Source: Appendix B · p. 17, Figure 1 · p. 17
Appendix C TerMeZO Algorithm: Pseudo-code
- Algorithm 1 initializes its mask using a threshold set at the lower ρ0-quantile of latent-weight distances to quantization boundaries.The pseudocode specifies computing thresholds and distances across layers before forming initial sparsity sets.Source: Appendix C · p. 18
Appendix D TerMeZO for RL-based Fine-tuning
Appendix D extends TerMeZO to RL fine-tuning by applying the evolution-strategy updates only to selected latent weights, and reports faster convergence and better performance than full-parameter MeZO on Countdown. The appendices establish unbiasedness for the Gaussian-smoothed objective and state assumptions used in the convergence proof.
Appendix D TerMeZO for RL-based Fine-tuning
- TerMeZO extends RL-based fine-tuning by restricting the weights to which the evolution-strategy algorithm applies perturbations.The approach uses multiple perturbations, estimates their rewards, and updates with their reward-weighted sum.Source: Appendix D · p. 19
- On Countdown, TerMeZO outperforms full-parameter MeZO and converges noticeably faster, while QZO struggles to solve the task.The RLVR reward reflects correctness and output format, with correctness weighted at 0.9.Source: Appendix D · p. 19, Figure 2 · p. 19
Appendix E Unbiasedness of SPSA for the Gaussian-Smoothed Objective
- The symmetric two-point SPSA estimator with Gaussian perturbations is unbiased for the gradient of the Gaussian-smoothed objective with respect to w.The proof obtains the symmetric form using the fact that z and −z have the same distribution.Source: Appendix E · p. 20
F.1 Assumptions
- The convergence proof assumes bounded loss, Lipschitz and smoothness conditions, mask alignment, and bounded minibatch variance.The conditions include Lipschitz continuity in w and Lipschitz continuity of ∇θF in (w, θ).Source: §F.1 · pp. 20–21
F.2 Proof of Theorem 1
The proof analyzes Gaussian-smoothed objectives to handle the quantized loss and derives bounds for updates restricted to a sparse mask. Its intermediate results establish smoothness and control estimator behavior under stated assumptions.
F.2 Proof of Theorem 1
- For a mask S of size at most k0, the proof tracks the expected squared change of the quantized weights in S through qS(w).The masked change DS(w, z) is the squared difference between quantized weights before and after perturbing coordinates in S.Source: §F.2 · p. 21
- Lemma 1 gives regularity bounds for perturbations supported on S, including bounds on the smoothed objective’s derivatives and its dependence on θ.The bounds use the mask-size limit |S| ≤ k0.Source: §F.2 · p. 21
- The Gaussian-smoothed objective is infinitely differentiable in w for fixed θ, despite the original quantized objective being piecewise constant and discontinuous.It is continuously differentiable jointly in (w, θ).Source: §F.2 · pp. 21–22
- The proof combines the masked and full-precision SPSA updates, with the latent-weight step supported on St and its size bounded by k0.The update analysis invokes the regularity bounds and estimator lemmas to derive an expected objective-change bound.Source: §F.2 · pp. 25–26
G.1 GLUE/SuperGLUE Experiments
The experiments specify shared protocols for GLUE/SuperGLUE and instruction-following and reasoning tasks, then compare validation-loss behavior and QuZO performance. TerMeZO is stable across the tested learning rates in the validation-loss comparison, while QuZO is less stable and performs worse on GSM8K.
G.1 GLUE/SuperGLUE Experiments
- GLUE/SuperGLUE experiments use task-specific prompt templates, the full training split, and test accuracy; the ICL baseline uses 32 sampled demonstrations.The hyperparameter table covers ZO methods and BP on these datasets.Source: §G.1 · p. 28, Table 5 · p. 28
G.2 Instruction Following and Reasoning Experiments
- Instruction-following and reasoning experiments use GSM8K and Magicoder-Evol-Instruct-110K with response-token loss and sequences truncated to 2048 tokens.Training uses a quarter of MagiCoder’s training set; GSM8K validation is carved from its training split.Source: §G.2 · p. 28, Table 6 · p. 28
H.1 Validation loss curves for TerMeZO and FIM-S-MeZO
- TerMeZO decreases validation loss steadily on Falcon-E-3B-Base GSM8K at ρ = 0.01 across η = 10−5, 10−4, and 5 × 10−4; FIM-S-MeZO becomes unstable at 5 × 10−4.FIM-S-MeZO reaches lower loss at η = 10−5 and η = 10−4.Source: §H.1 · p. 29, Figure 3 · p. 29
H.2 Comparison with QuZO
- QuZO barely improves over the Falcon-E-1B-Base GSM8K base model, while TerMeZO clearly outperforms it.QuZO diverges at the higher learning rates used for TerMeZO and the other ZO baselines, so a lower range was searched for QuZO.Source: §H.2 · p. 29, Table 7 · p. 29