[Paper Review] Zero-th Order Algorithm for Softmax Attention Optimization

Yichuan Deng, Zhihang Li+2 more|arXiv (Cornell University)|Jul 17, 2023

Topic ModelingComputer Science3 citations

TL;DR

This paper proposes a zeroth-order optimization algorithm tailored for softmax attention mechanisms in large language models (LLMs), enabling gradient estimation via forward passes only. It establishes global convergence with a sample complexity of $ O(M \cdot (1 + d^{1.5} \kappa^2(A)/k) \cdot \mu^{-2} B^{-1} \log((L(x_0)-L^*)/epsilon)) $, where $ M = \exp(O(R^2 + \log n)) $, demonstrating efficient training without backpropagation in high-dimensional settings.

ABSTRACT

Large language models (LLMs) have brought about significant transformations in human society. Among the crucial computations in LLMs, the softmax unit holds great importance. Its helps the model generating a probability distribution on potential subsequent words or phrases, considering a series of input words. By utilizing this distribution, the model selects the most probable next word or phrase, based on the assigned probabilities. The softmax unit assumes a vital function in LLM training as it facilitates learning from data through the adjustment of neural network weights and biases. With the development of the size of LLMs, computing the gradient becomes expensive. However, Zero-th Order method can approximately compute the gradient with only forward passes. In this paper, we present a Zero-th Order algorithm specifically tailored for Softmax optimization. We demonstrate the convergence of our algorithm, highlighting its effectiveness in efficiently computing gradients for large-scale LLMs. By leveraging the Zeroth-Order method, our work contributes to the advancement of optimization techniques in the context of complex language models.

Motivation & Objective

• Address the high computational cost of backpropagation in large-scale LLMs due to massive parameter counts.
• Enable efficient optimization of softmax attention mechanisms without gradient computation via backpropagation.
• Develop a zeroth-order method that approximates gradients using only forward pass evaluations, reducing memory and computational overhead.
• Establish theoretical convergence guarantees for the proposed algorithm under realistic assumptions on model parameters and loss landscape.
• Provide a scalable optimization framework suitable for large-scale LLMs where traditional backpropagation is infeasible.

Proposed method

• Proposes a zeroth-order gradient estimator based on simultaneous perturbation stochastic approximation (SPSA), using only forward loss evaluations.
• Derives an unbiased gradient estimate via loss difference between perturbed and unperturbed parameters: $ \widehat{g}(x_0)_i = \frac{1}{2\epsilon p_i}(L(x_0 + \epsilon p) - L(x_0 - \epsilon p)) $.
• Applies the gradient estimate within a projected gradient descent (GD) framework to optimize the softmax attention objective.
• Introduces a loss function $ L(x) $ defined as the squared error between normalized exponential outputs and target vectors, modeling attention as a softmax regression problem.
• Uses concentration and matrix norm bounds to control variance in gradient estimation, ensuring convergence under smoothness and PL conditions.
• Establishes convergence via a recursive expectation bound involving the loss decrease per iteration, leveraging $ \mu $-PL condition and trace of Hessian covariance.

Experimental results

Research questions

• RQ1 Can a zeroth-order optimization method achieve global convergence for softmax attention in large language models?
• RQ2 What is the sample complexity of zeroth-order optimization for attention mechanisms under realistic assumptions on model parameters?
• RQ3 How does the proposed method compare to traditional backpropagation in terms of computational efficiency and convergence rate?
• RQ4 What theoretical guarantees can be provided for gradient estimation in high-dimensional, non-convex attention optimization problems?
• RQ5 How do model parameters such as condition number $ \kappa(A) $, dimension $ d $, and batch size $ B $ affect convergence?

Key findings

• The proposed zeroth-order algorithm achieves global convergence to an $ \epsilon $-optimal solution with $ O(M \cdot (1 + d^{1.5} \kappa^2(A)/k) \cdot \mu^{-2} B^{-1} \log((L(x_0)-L^*)/\epsilon)) $ iterations, where $ M = \exp(O(R^2 + \log n)) $.
• The method ensures convergence under $ \mu $-PL, $ l $-smoothness, and bounded Hessian effective rank, with variance controlled via SPSA perturbations.
• The convergence rate depends on the inverse of the PL constant $ \mu $, batch size $ B $, and condition number $ \kappa(A) $, with improved sample efficiency for larger $ B $.
• The algorithm avoids backpropagation entirely, relying only on forward pass evaluations, making it suitable for black-box or gradient-obscured LLMs.
• Theoretical analysis shows that the gradient estimator's variance is bounded by $ \operatorname{tr}[\Sigma(x)] \leq \alpha (L(x) - L^*) $, enabling stable descent.
• The final convergence bound scales with $ \exp(O(R^2)) $, where $ R $ is the parameter norm bound, indicating exponential dependence on model size in worst-case scenarios.