Visualising Policy-Reward Interplay to Inform Zeroth-Order Preference Optimisation of Large Language Models
This work addresses the challenge of reducing memory usage in LLM fine-tuning for researchers and practitioners, though it is incremental as it adapts existing methods to a new task.
The paper tackled the problem of computationally intensive fine-tuning of Large Language Models (LLMs) by developing ZOPrO, a zeroth-order algorithm for preference optimisation, which achieved convergence times comparable to first-order methods in tasks like summarisation and machine translation.
Fine-tuning Large Language Models (LLMs) with first-order methods like back-propagation is computationally intensive. Zeroth-Order (ZO) optimisation uses function evaluations instead of gradients, reducing memory usage, but suffers from slow convergence in high-dimensional models. As a result, ZO research in LLMs has mostly focused on classification, overlooking more complex generative tasks. In this paper, we introduce ZOPrO, a novel ZO algorithm designed for Preference Optimisation in LLMs. We begin by analysing the interplay between policy and reward models during traditional (first-order) Preference Optimisation, uncovering patterns in their relative updates. Guided by these insights, we adapt Simultaneous Perturbation Stochastic Approximation (SPSA) with a targeted sampling strategy to accelerate convergence. Through experiments on summarisation, machine translation, and conversational assistants, we demonstrate that our method consistently enhances reward signals while achieving convergence times comparable to first-order methods. While it falls short of some state-of-the-art methods, our work is the first to apply Zeroth-Order methods to Preference Optimisation in LLMs, going beyond classification tasks and paving the way for a largely unexplored research direction. Code and visualisations are available at https://github.com/alessioGalatolo/VisZOPrO