Converge Faster, Talk Less: Hessian-Informed Federated Zeroth-Order Optimization

Area of Science:

• Machine Learning
• Optimization Algorithms
• Distributed Systems

Background:

• Federated learning (FL) enables collaborative model training without sharing raw data.
• Zeroth-order (ZO) optimization offers dimension-free communication in FL, ideal for large language models (LLMs).
• Existing ZO-FL methods often ignore curvature information, limiting convergence speed.

Purpose of the Study:

• To develop a Hessian-informed ZO federated optimization method (HiSo) for accelerated LLM fine-tuning.
• To leverage global diagonal Hessian approximations for faster convergence while maintaining scalar-only communication.
• To theoretically and empirically validate the effectiveness of Hessian information in ZO-FL.

Main Methods:

• Proposed HiSo, a novel Hessian-informed ZO federated optimization algorithm.
• Utilized global diagonal Hessian approximations to inform optimization steps.
• Preserved strict scalar-only communication, avoiding transmission of second-order information.
• Conducted theoretical analysis for non-convex functions under Hessian approximation assumptions.
• Performed empirical evaluations on diverse LLM fine-tuning benchmarks.

Main Results:

• HiSo achieves an accelerated convergence rate independent of Lipschitz constant and model dimension under specific Hessian approximation assumptions.
• Theoretical analysis provides insights into why ZO convergence can exceed its worst-case bound.
• Empirically, HiSo demonstrated a 1-5x speedup in communication rounds compared to state-of-the-art ZO-FL baselines.
• Demonstrated significant communication cost reduction in LLM fine-tuning.

Conclusions:

• Hessian information acts as an effective accelerator in federated ZO optimization.
• HiSo offers a practical and efficient approach for fine-tuning LLMs in FL settings.
• The method successfully balances communication efficiency with accelerated convergence.