MaZO: Masked Zeroth-Order Optimization for Multi-Task Fine-Tuning of Large Language Models

Fuente: arXiv
Gespeichert in:
Bibliographische Detailangaben
Hauptverfasser: Zhang, Zhen, Yang, Yifan, Zhen, Kai, Susanj, Nathan, Mouchtaris, Athanasios, Kunzmann, Siegfried, Zhang, Zheng
Format: Preprint
Veröffentlicht: 2025
Schlagworte:
Online-Zugang:
Tags: Tag hinzufügen
Keine Tags, Fügen Sie den ersten Tag hinzu!
_version_ 1866912233736372224
author Zhang, Zhen
Yang, Yifan
Zhen, Kai
Susanj, Nathan
Mouchtaris, Athanasios
Kunzmann, Siegfried
Zhang, Zheng
author_facet Zhang, Zhen
Yang, Yifan
Zhen, Kai
Susanj, Nathan
Mouchtaris, Athanasios
Kunzmann, Siegfried
Zhang, Zheng
contents Large language models have demonstrated exceptional capabilities across diverse tasks, but their fine-tuning demands significant memory, posing challenges for resource-constrained environments. Zeroth-order (ZO) optimization provides a memory-efficient alternative by eliminating the need for backpropagation. However, ZO optimization suffers from high gradient variance, and prior research has largely focused on single-task learning, leaving its application to multi-task learning unexplored. Multi-task learning is crucial for leveraging shared knowledge across tasks to improve generalization, yet it introduces unique challenges under ZO settings, such as amplified gradient variance and collinearity. In this paper, we present MaZO, the first framework specifically designed for multi-task LLM fine-tuning under ZO optimization. MaZO tackles these challenges at the parameter level through two key innovations: a weight importance metric to identify critical parameters and a multi-task weight update mask to selectively update these parameters, reducing the dimensionality of the parameter space and mitigating task conflicts. Experiments demonstrate that MaZO achieves state-of-the-art performance, surpassing even multi-task learning methods designed for first-order optimization.
format Preprint
id arxiv_https___arxiv_org_abs_2502_11513
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle MaZO: Masked Zeroth-Order Optimization for Multi-Task Fine-Tuning of Large Language Models
Zhang, Zhen
Yang, Yifan
Zhen, Kai
Susanj, Nathan
Mouchtaris, Athanasios
Kunzmann, Siegfried
Zhang, Zheng
Machine Learning
Artificial Intelligence
Large language models have demonstrated exceptional capabilities across diverse tasks, but their fine-tuning demands significant memory, posing challenges for resource-constrained environments. Zeroth-order (ZO) optimization provides a memory-efficient alternative by eliminating the need for backpropagation. However, ZO optimization suffers from high gradient variance, and prior research has largely focused on single-task learning, leaving its application to multi-task learning unexplored. Multi-task learning is crucial for leveraging shared knowledge across tasks to improve generalization, yet it introduces unique challenges under ZO settings, such as amplified gradient variance and collinearity. In this paper, we present MaZO, the first framework specifically designed for multi-task LLM fine-tuning under ZO optimization. MaZO tackles these challenges at the parameter level through two key innovations: a weight importance metric to identify critical parameters and a multi-task weight update mask to selectively update these parameters, reducing the dimensionality of the parameter space and mitigating task conflicts. Experiments demonstrate that MaZO achieves state-of-the-art performance, surpassing even multi-task learning methods designed for first-order optimization.
title MaZO: Masked Zeroth-Order Optimization for Multi-Task Fine-Tuning of Large Language Models
topic Machine Learning
Artificial Intelligence
url https://arxiv.org/abs/2502.11513