There have been significant advancements made by large language models (LLMs)
in various aspects of our daily lives. LLMs serve as a transformative force in
natural language processing, finding applications in text generation,
translation, sentiment analysis, and question-answering. The accomplishments of
LLMs have led to a substantial increase in research efforts in this domain. One
specific two-layer regression problem has been well-studied in prior works,
where the first layer is activated by a ReLU unit, and the second layer is
activated by a softmax unit. While previous works provide a solid analysis of
building a two-layer regression, there is still a gap in the analysis of
constructing regression problems with more than two layers.
In this paper, we take a crucial step toward addressing this problem: we
provide an analysis of a two-layer regression problem. In contrast to previous
works, our first layer is activated by a softmax unit. This sets the stage for
future analyses of creating more activation functions based on the softmax
function. Rearranging the softmax function leads to significantly different
analyses. Our main results involve analyzing the convergence properties of an
approximate Newton method used to minimize the regularized training loss. We
prove that the loss function for the Hessian matrix is positive definite and
Lipschitz continuous under certain assumptions. This enables us to establish
local convergence guarantees for the proposed training algorithm. Specifically,
with an appropriate initialization and after $O(\log(1/\epsilon))$ iterations,
our algorithm can find an $\epsilon$-approximate minimizer of the training loss
with high probability. Each iteration requires approximately $O(\mathrm{nnz}(C)
+ d^\omega)$ time, where $d$ is the model size, $C$ is the input matrix, and
$\omega < 2.374$ is the matrix multiplication exponent.

我们对两层回归问题进行了分析，使用了 softmax 激活单元作为第一层，并分析了近似牛顿法用于最小化正则化训练损失的收敛性质，证明了 Hessian 矩阵的损失函数是正定和 Lipschitz 连续的，在适当的初始化和迭代次数后，我们的算法可以高概率地找到训练损失的 ε- 近似最小化器。