Ternary sparse method reduces memory for fine-tuning language models
TerMeZO: Ternary Sparse Zeroth-Order Optimization for Fine-tuning BitNet Models at the Edge
Machine LearningArtificial Intelligence
Summary
Fine-tuning large language models usually needs a lot more memory than just using them for basic tasks. The paper focuses on a special kind of small-memory fine-tuning method that estimates changes without extra memory-heavy calculations but hits trouble with certain models called BitNet. The authors introduce TerMeZO, which smartly picks which weight parts to update based on how those weights are structured, reducing memory needs without hurting performance. Their method can actually speed up fine-tuning and works well across different tasks using BitNet models.
What this means in practice
- •For edge ai developers: Fine-tune BitNet language models on devices with limited memory while maintaining performance.
- •For mobile app engineers: Deploy smarter language model updates that reduce memory use and latency for apps on mobile devices.
Authors
Houssem Sifaou, Prabodh Katti, Bipin Rajendran, Osvaldo Simeone
Abstract
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.