🤖 AI Summary
This study addresses the substantial memory overhead of zeroth-order fine-tuning for BitNet models and the reliance of sparse masks on first-order gradients. To overcome these limitations, this work proposes a gradient-free sparse zeroth-order fine-tuning method. Specifically, it pioneers the exploitation of the geometric structure of ternary quantizers to identify critical weights, enabling the construction of sparse masks without additional data or gradient information, thereby significantly reducing the effective optimization dimensionality. Theoretical convergence guarantees are established for the proposed approach. Empirical evaluations on models ranging from 1B to 3B parameters demonstrate that this method achieves performance comparable to full-parameter fine-tuning while substantially reducing memory consumption, making it highly suitable for deployment on edge devices.
📝 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.