Zero-th Order Algorithm for Softmax Attention Optimization
This work addresses the optimization bottleneck in training large language models, offering a method to reduce computational expense, though it appears incremental as it applies an existing zero-th order technique to a specific component.
The paper tackles the high computational cost of gradient computation in large language models (LLMs) by proposing a Zero-th Order algorithm for softmax attention optimization, demonstrating its convergence and effectiveness in efficiently computing gradients for large-scale LLMs.
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.