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.

Large language models have become ubiquitous in modern life, finding applications in various domains such as natural language processing, language translation, and speech recognition. Recently, a breakthrough work [Zhao, Panigrahi, Ge, and Arora Arxiv 2023] explains the attention model from probabilistic context-free grammar (PCFG). One of the central computation task for computing probability in PCFG is formulating a particular tensor low rank approximation problem, we can call it tensor cycle rank. Given an $n\times n\times n$ third order tensor A, we say that A has cycle rank-k if there exists three $n\times k^{2}$ size matrices $U, V$, and W such that for each entry in each \begin{equation*}A_{a,b,c}=\sum_{i=1J^{=1}}^{k}\sum_{\prime}^{k}\sum_{l=1}^{k}U_{a,i+k\left(J-1\right)}\prime\otimes V_{b_{J}'+k\left(l-1\right)}\otimes W_{c,l+k\left(i-1\right)}\end{equation*}for all $a\in[n], b\in[n], c\in[n]$. For the tensor classical rank, tucker rank and train rank, it has been well studied in [Song, Woodruff, Zhong SODA 2019]. In this paper, we generalize the previous “rotation and sketch” technique in [Song, Woodruff, Zhong SODA 2019] and show an input sparsity time algorithm for cycle rank.