PaperlayerView exact revision on arXiv

Source-linked AI summary

AdaZeta: Adaptive Zeroth-Order Tensor-Train Adaption for Memory-Efficient Large Language Models Fine-Tuning

Yifan Yang, Kai Zhen, Ershad Banijamal, Athanasios Mouchtaris, Zheng Zhang

arXiv:2406.18060v3cs.CLcs.AIcs.LG

TL;DR

LLM fine-tuning remains memory-intensive, while zeroth-order methods still face performance gaps and divergence risks. AdaZeta combines fast-forward tensorized adapters with an adaptive query schedule, and the paper reports improved fine-tuning results and convergence across tasks and models.

  • Problem

    Fine-tuning LLMs consumes increasing GPU memory as models grow, while zeroth-order methods still face performance gaps versus first-order methods and divergence risks.

    Source: §1 · p. 1

  • Method

    AdaZeta combines low-parameter tensorized adapters with parallel contraction and an adaptive query schedule for zeroth-order fine-tuning.

    Source: §1 · p. 2, §2.1 · p. 3, §3.2 · p. 4

  • Results

    AdaZeta achieves significantly better fine-tuning results across tasks and models, with faster convergence than other zeroth-order methods.

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

  • Takeaways & Limitations

    Theoretical analysis and experiments support AdaZeta's improved convergence across the evaluated RoBERTa-Large and Llama-2 tasks.

    Source: §5 · p. 9

    Multiple queries are executed sequentially in a for-loop at each training step, constraining further speed improvements; parallel or distributed GPU execution is proposed as a possible optimization.

    Source: Limitations · pp. 9–10

Abstract

from arXiv · show

Fine-tuning large language models (LLMs) has achieved remarkable performance across various natural language processing tasks, yet it demands more and more memory as model sizes keep growing. To address this issue, the recently proposed Memory-efficient Zeroth-order (MeZO) methods attempt to fine-tune LLMs using only forward passes, thereby avoiding the need for a backpropagation graph. However, significant performance drops and a high risk of divergence have limited their widespread adoption. In this paper, we propose the Adaptive Zeroth-order Tensor-Train Adaption (AdaZeta) framework, specifically designed to improve the performance and convergence of the ZO methods. To enhance dimension-dependent ZO estimation accuracy, we introduce a fast-forward, low-parameter tensorized adapter. To tackle the frequently observed divergence issue in large-scale ZO fine-tuning tasks, we propose an adaptive query number schedule that guarantees convergence. Detailed theoretical analysis and extensive experimental results on Roberta-Large and Llama-2-7B models substantiate the efficacy of our AdaZeta framework in terms of accuracy, memory efficiency, and convergence speed.

1 Introduction

AdaZeta combines fast-forward tensorized adapters with an adaptive query schedule to improve zeroth-order fine-tuning accuracy and convergence. It targets the memory burden of backpropagation-based fine-tuning as well as performance gaps and divergence risks in zeroth-order methods.

1 Introduction

  • MeZO avoids a backpropagation graph and reduces memory use, but zeroth-order fine-tuning still faces a performance gap versus first-order methods and divergence risks.Memory use remains high in conventional fine-tuning because it relies on a backpropagation graph.Source: §1 · p. 1
  • Its tensorized adapters offer high representational ability with fewer trainable parameters than LoRA adapters.The paper introduces them as part of a task-specific parameter-efficient fine-tuning approach for zeroth-order methods.Source: §1 · p. 2
  • An adaptive query schedule sub-linearly increases queries to address divergence while supporting better accuracy and faster convergence in steps and time.AdaZeta combines the schedule with fast-forward tensorized adapters.Source: §1 · p. 2
  • AdaZeta outperforms MeZO, MeZO-LoRA, and Sparse-MeZO across tasks, with faster convergence.The paper also reports better fine-tuning results across tasks and models, supported by theoretical analysis and experiments on RoBERTa-Large and Llama-2.Source: §1 · p. 2

2 Background

The background motivates task-specific parameter-efficient methods for zeroth-order fine-tuning and introduces tensorized adapters as a low-parameter option. AdaZeta uses parallel contraction to address the slow forward pass associated with tensor-train factors.

2.1 Parameter-Efficient Fine-tuning

  • Prior zeroth-order PEFT work considered LoRA and prefix tuning, but reported limited improvement and did not detail task-specific adapter tuning.Tensorized adapters compress adapter-layer weight matrices with Tensor-Train decomposition, offering an ultra-low-parameter alternative.Source: §2.1 · p. 3
  • Sequentially contracting small tensor factors slows tensorized-adapter forward passes, so the paper proposes parallel contraction to improve inference speed.This matters for zeroth-order methods, which require two forward passes per step.Source: §2.1 · p. 3

2.2 Tensorized Adapters

  • Tensorized linear layers represent reshaped weights as a sequence of TT factors, which are contracted and reshaped into a weight matrix during the forward pass.The weights are initialized, stored, and updated in TT format; the first and last ranks are set to 1.Source: §2.2 · p. 3
  • Each adapter has two tensorized layers around a nonlinearity and is attached after attention and feed-forward layers; layer norm is frozen during zeroth-order tuning.The adapters reduce trainable parameters by over 80×.Source: §2.2 · p. 3

3 Methods

AdaZeta combines randomized zeroth-order estimation with low-parameter tensorized adapters and an adaptive query schedule. Its analysis relates convergence to the query schedule and the number of trainable parameters.

3.1 Zeroth-order Estimation

  • Randomized zeroth-order estimation uses perturbed loss evaluations to estimate gradients, requiring only forward passes and eliminating the backpropagation graph.The method defines the loss over a mini-batch and applies vector-wise random perturbations to trainable adapter parameters.Source: §3.1 · p. 4

3.2 The AdaZeta Framework

  • Tensorized adapters reduce trainable dimensionality to improve zeroth-order estimation, while parallel contraction groups tensor factors to speed forward passes and avoid high-dimensional tensors.For larger models, the tensor factors can be arranged in tripartite or quadripartite structures.Source: §3.2 · pp. 4–5
  • AdaZeta increases query count sublinearly at each epoch, averages multiple query results, and caps queries; the authors report it resolves observed divergence and is faster than increasing batch size.The query count stays fixed within each epoch; medium-size models use one query.Source: §3.2 · p. 5

3.3 Theoretical Analysis

  • Under A1–A2, AdaZeta’s convergence bound depends on the query schedule; a sublinear schedule can drive the true gradient toward zero, while fixed Q = 1 incurs O(C(d, ϵ)/Kϵ).The analysis also relates convergence to model dimension d, motivating fewer trainable parameters.Source: §3.3 · pp. 5–6

4 Experiments

Experiments span RoBERTa-Large and Llama-2-7B across language-understanding and generation tasks, comparing AdaZeta with inference-only, zeroth-order, and first-order baselines. Results assess accuracy, convergence, memory use, and training time.

4.2 Large-scale Llama-2 Models

  • AdaZeta outperforms zeroth-order baselines in 8 of 10 Llama-2-7B tasks, surpasses ICL and zero-shot accuracy, and exceeds FO-AdamW on RTE, CB, and COPA.The FO-AdamW comparisons require 8× more GPU memory.Source: Table 2 · p. 7, §4.2 · pp. 7–8

4.1 Medium-size Roberta-Large Models

  • AdaZeta beats MeZO-LoRA in 5 of 7 RoBERTa-Large tests at both 16 and 64 batch sizes, with the authors attributing the gain to fewer trainable parameters.The authors connect the result to dimension-related ZO convergence rates.Source: §4.1 · p. 7
  • AdaZeta produces more consistent results across 16- and 64-sized batches than MeZO-LoRA, whose performance declines under the 16-shot setup.The authors attribute the difference to reduced ZO gradient noise and report superior convergence at batch size 16.Source: §4.1 · p. 7

4.3 Memory Training Time Efficiency

  • AdaZeta converges on-par or faster than other ZO methods in GPU hours across SST2, WIC, CB, and MultiRC, including better results than MeZO-LoRA and Sparse-MeZO at large batch size.The 64-batch experiments did not use gradient accumulation, which may increase training time.Source: Table 3 · p. 8, §4.3 · p. 8

4.2 Large-scale Llama-2 Models

  • Nearly 6K steps versus less than 1K lets MeZO-LoRA and AdaZeta, respectively, reach loss 0.4 at the same 1e-4 learning rate; adaptive queries mitigate divergence without increasing training memory.The schedule improves training outcomes in large-scale fine-tuning.Source: §4.2 · p. 8

4.3 Memory Training Time Efficiency

  • 14GB suffices for Llama-2-7B SST2 fine-tuning, yielding over 8× memory reduction versus FT; AdaZeta uses similar or less memory than other ZO methods.The accuracy–memory comparison identifies AdaZeta as the most accurate memory-efficient method.Source: Figure 3 · p. 8, §4.3 · p. 8

4.4 Further Comparison with LoRA

  • In first-order LoRA, lowering rank barely reduces memory because the backpropagation graph remains; even against LoRA/r=1/BS=1, AdaZeta/r=8/BS=16 uses nearly 50% less.Reducing batch size lowers memory for both first-order and zeroth-order methods.Source: Table 4 · p. 9, §4.4 · p. 9

5 Conclusion

AdaZeta combines tensor-train decomposition with an adaptive query schedule for zeroth-order fine-tuning. The paper reports better results across tasks and models, while identifying speed improvements as future work.

5 Conclusion

  • AdaZeta combines tensor-train decomposition with an adaptive query schedule for zeroth-order fine-tuning.The framework also uses fast-forward tensorized adapters.Source: §5 · p. 9
  • The method achieves significantly better fine-tuning results across tasks and models, with theoretical convergence results consistent with experiments on Roberta-Large and Llama-2.Source: §5 · p. 9
  • Future work could accelerate AdaZeta by running independent queries concurrently across multiple GPUs.The paper also proposes exploring the adaptive query schedule with other PEFT methods.Source: §5 · p. 9

Limitations

AdaZeta's queries run sequentially at each training step, which restricts further speed enhancements.

Limitations

  • Queries run sequentially in a for-loop at each training step, restricting further speed enhancements.The paper suggests parallel or distributed GPU execution because the queries use different random seeds and are independent.Source: Limitations · pp. 9–10

Potential Risks

The paper notes that prolonged training times, especially with multiple GPUs, can pose environmental challenges.

Potential Risks

  • Prolonged training times, especially when using multiple GPUs, can pose environmental challenges.The authors say they are pursuing more efficient training methods with ecological considerations.Source: Potential Risks · p. 10

A.1 Dataset Setup

The experiments evaluate Roberta-Large on sentiment-analysis and natural-language-inference tasks, and Llama-2-7B on SuperGLUE and generation tasks under a low-resource condition. The section also describes the baselines and hyperparameter choices used in the experiments.

A.1 Dataset Setup

  • Roberta-Large evaluation covers sentiment-analysis and natural-language-inference tasks, while Llama-2-7B tests include SuperGLUE and generation tasks.The Llama-2-7B condition uses 1,000 training, 500 validation, and 1,000 testing samples.Source: §A.1 · p. 13

A.2 Baselines

  • The experiments compare against full-model first-order fine-tuning, linear probing, MeZO, MeZO-LoRA, and Sparse-MeZO.Sparse-MeZO was reproduced using a fixed sparsity ratio for each layer, selected based on its paper's reported overall outcome.Source: Full-model First-Order Fine-Tuning (FT) · pp. 13–14

Full-model First-Order Fine-Tuning (FT)

A.3 Hyperparameters

  • Roberta-Large tasks ran for 80K steps except full-model fine-tuning, and checkpoints were selected by validation loss every 200 steps.Learning rates were generally kept consistent across methods for the same task, with lower rates for full-model fine-tuning to ensure convergence; tensor-factor shapes were also selected experimentally.Source: Table 7 · p. 14, §A.3 · p. 14

B.1 Additional Momeory Comparison results

Additional profiling reports 5–8× memory reduction for AdaZeta across tasks, with memory similar to or below other MeZO variants. The reported comparison also pairs this reduction with variance reduction that improves ZO fine-tuning accuracy.

B.1 Additional Momeory Comparison results

  • The profiling extends the SST2 measurements to WIC, CB, and MultiRC tasks.Source: §B.1 · p. 14
  • 5–8× memory reduction is reported for AdaZeta across different tasks.Table 10 presents quantitative memory profiling results.Source: §B.1 · p. 14, Table 10 · p. 15
  • AdaZeta uses similar or less memory than MeZO, MeZO-LoRA, and Sparse-MeZO, while adding variance reduction that improves ZO fine-tuning accuracy.Source: §B.1 · p. 14
  • The Llama-2-7B hyperparameter table specifies 5K AdaZeta steps, 10K steps for other ZO methods, and 5 epochs for full-model fine-tuning.The best checkpoint is recorded from validation loss every 200 training steps.Source: Table 8 · p. 15

C Proof of Theorem 1

The proof of Theorem 1 bounds gradient noise from mini-batch sampling and zeroth-order estimation, then uses the AdaZeta update and smoothness to derive a convergence bound.

C Proof of Theorem 1

  • The proof separates mini-batch sampling noise h_k from zeroth-order estimation noise e_k, then bounds h_k under Assumption A2.The lemma considers uniform sampling without replacement and an i.i.d. mini-batch.Source: C Proof of Theorem 1 · p. 16
  • The proof bounds the zeroth-order contribution through the variance of the gradient estimate under i.i.d. perturbation vectors.Source: C Proof of Theorem 1 · pp. 16–17
  • Starting from the AdaZeta update, the proof applies Taylor’s theorem and Lipschitz smoothness, takes expectations, and sums across optimization steps.Source: C Proof of Theorem 1 · pp. 17–18
  • The final bound follows by decomposing the estimated gradient into the true gradient and its noise, then applying the preceding lemma and averaging over steps.Source: C Proof of Theorem 1 · pp. 18–19

Related papers

Loading 2406.18060v3…