研究线性模型训练中损失变化的规律和最优学习率的影响

探究一维线性模型训练中，测试损失随训练步数变化的缩放定律及其最优学习率影响，并研究多维线性模型训练的缩放定律，确定参数以符合特定损失衰减模式。

研究大模型的缩放定律对减少其训练开销至关重要，即最终的测试损失如何随着训练步数和模型大小的变化而变化？本题中，我们研究了训练线性模型时的缩放定律。

1. 在本小问中，考虑使用梯度下降学习一个一维线性模型的情况。

• 定义数据分布$\mathcal{D}$为一个${\mathbb{R}}^2$上的分布，每个数据是一个数对$(x,y)$，分别代表输入和输出，并服从分布$x\sim N(0,1),y\sim N(3x,1)$。

• 用梯度下降算法学习线性模型$f_w(x)=w\cdot x$，其中$w,x\in \mathbb{R}$。初始化${\omega }_0=0$并进行多步迭代。每次迭代时，从$\mathcal{D}$中采样$(x_t,y_t)$，然后更新$w_t$为$w_{t+1}\leftarrow w_t-\eta \nabla l_t(w_t)$，其中$l_t(w)=\frac{1}{2}(f_w(x_t)-y_t{)}^2$是平方损失函数，$\eta >0$是学习率。

设学习率$\eta \in (0,\frac{1}{3}]$，那么$T\ge 0$步迭代之后的测试损失的期望

${\overline{\mathcal{L}}}_{\eta ,T}={\mathbb{E}}_{w_T}{\mathbb{E}}_{(x,y)\sim D}[\frac{1}{2}(f_{w_T}(x)-y{)}^2]$

是多少？

2. 现在我们在第一小问的设定下，考虑学习率$\eta$被调到最优的情况，求函数$g(T)$，使得当$T\to +\infty$时，以下条件成立：

$\left|\underset{\eta \in (0,\frac{1}{3}]}{\inf }{\mathcal{I}}_{n,T}-g(T)\right|=O(\frac{(\log T{)}^2}{T^2})$

3. 一个常常被观测到的实验现象是大语言模型的预训练过程大致遵循Chinchilla缩放定律：

${\overline{\mathcal{L}}}_{N,T}\approx \frac{A}{N^{\alpha }}+\frac{B}{T^{\beta }}+C$，

其中${\overline{\mathcal{L}}}_{N,T}$是在经过$T$步训练后具有$N$个参数的模型的测试损失的期望，$A$，$B$，$a$，$\beta$，$C$是常数。现在我们举一个训练多维线性模型的例子，使其也遵循类似的缩放定律。

• 固定$a>0,b\ge 1$，每个数据$(x_{\cdot },y)$由一个输入和输出组成，其中输入$x_{\cdot }$是一个无限维向量（可看作一个序列），输出$y$满足$y\in \mathbb{R}$。定义数据分布$\mathcal{D}$如下。首先，从Zipf分布中采样$k$，$\mathrm{Pr}[k=i]\propto i^{-(a+1)}(i\ge 1)$。令$j:=[k^b]$，然后，从$mathcalN(0,1)$中采样得到$x_{\cdot }$的第$j$个坐标$x_j$，并令其余坐标为0。最后，$y\sim N(3x_j,1)$。这样得到的$(x_{\cdot },y)$的分即数据分布$\mathcal{D}$。

• 我们研究一个仅关注前$N$个输入坐标的线性模型。定义函数${\varphi }_N(x{x}_{\cdot })=(x_1,...,x_N)$。我们研究的线性模型具有参数$\mathbf{w}\in {\mathbb{R}}^N$，输出为$f_{\mathbf{w}}(x)=(\mathbf{w},{\varphi }_N(x_{\cdot }))$。

• 我们使用梯度下降算法学习该线性模型。初始化${\mathbf{w}}_0=0$并进行多步迭代。每次迭代时，从$\mathcal{D}$中采样$(x_{t,\cdot },y_t)$，然后更新${\mathbf{w}}_t$为${\mathbf{w}}_{t+1}\leftarrow {\mathbf{w}}_t-\eta \nabla l_t({\mathbf{w}}_t)$，其中$l_t(\mathbf{w})=\frac{1}{2}(f_{\mathbf{w}}(x_{t,\cdot })-y_t{)}^2$。

令${\overline{\mathcal{L}}}_{\eta ,T}={\mathbb{E}}_{{\mathbf{w}}_T}{\mathbb{E}}_{(x,y)\sim D}[\frac{1}{2}(f_{{\mathbf{w}}_T}(x)-y{)}^2]$为以学习率$\eta \in (0,\frac{1}{3}]$对其有N个参数的线性模型进行$T\ge 0$步训练后的测试损失的期望。

请求出$\alpha$，$\beta$，$C$，使得$\forall \gamma >0,\forall c>0$，当$T=N^{c+o(1)}$且$N$足够大时，以下条件成立：

$ϵ(N,T):=\frac{{\inf }_{\eta \in (0,\frac{1}{3}]}{\overline{\mathcal{L}}}_{N,T}-C}{\frac{A}{N^{\alpha }}+\frac{B}{T^{\beta }}}$，

$(\log N+\log T{)}^{-\gamma }\le ϵ(N,T)\le (\log N+\log T{)}^{\gamma }$。即${\inf }_{\eta \in (0,\frac{1}{3}]}{\overline{\mathcal{L}}}_{N,T}=\tilde{\Theta }(N^{-\alpha }+T^{-\beta })+C$，其中$\tilde{\Theta }$表示忽略任何关于$\log N$和$\log T$的多项式。

解：

1. 首先，我们来计算测试损失的期望${\overline{\mathcal{L}}}_{\eta ,T}$。

由于$x$和$y$是独立的随机变量，且$y$的条件分布是$N(3x,1)$，我们可以写出测试损失的期望为：

$${\overline{\mathcal{L}}}_{\eta ,T}={\mathbb{E}}_{(x,y)\sim D}[\frac{1}{2}(w_{T}x-y{)}^2]$$

由于$y=3x+ϵ$，其中$ϵ\sim N(0,1)$且独立于$x$，我们可以将$y$替换为$3x+ϵ$：

$${\overline{\mathcal{L}}}_{\eta ,T}={\mathbb{E}}_{x,ϵ}[\frac{1}{2}(w_{T}x-(3x+ϵ){)}^2]$$

展开并利用$\mathbb{E}[{ϵ}^2]=1$和$\mathbb{E}[x^2]=1$（因为$x\sim N(0,1)$）：

$${\overline{\mathcal{L}}}_{\eta ,T}={\mathbb{E}}_x[\frac{1}{2}(w_T^2{x}^2-6w_T{x}^2+9x^2+{ϵ}^2-6w_{T}xϵ+3w_T^2{x}^2)]$$

由于$ϵ$和$x$是独立的，我们可以分别计算期望：

$${\overline{\mathcal{L}}}_{\eta ,T}=\frac{1}{2}(w_T^2-6w_T+9)\mathbb{E}[x^2]+\frac{1}{2}\mathbb{E}[{ϵ}^2]$$

$${\overline{\mathcal{L}}}_{\eta ,T}=\frac{1}{2}(w_T^2-6w_T+9)+\frac{1}{2}$$

现在我们需要计算$w_T$的期望值。由于$w_t$的更新规则是$w_{t+1}=w_t-\eta \nabla l_t(w_t)$，我们有：

$$\nabla l_t(w_t)=w_t{x}_t-y_t=w_t{x}_t-(3x_t+ϵ)$$

因此，更新规则变为：

$$w_{t+1}=w_t-\eta (w_t{x}_t-3x_t-ϵ)$$

取期望并利用$\mathbb{E}[x_t]=0$和$\mathbb{E}[ϵ]=0$：

$$\mathbb{E}[w_{t+1}]=\mathbb{E}[w_t]-\eta (3\mathbb{E}[x_t^2])$$

由于$x_t^2$的期望是1，我们有：

$$\mathbb{E}[w_{t+1}]=\mathbb{E}[w_t]-3\eta$$

由于$w_0=0$，我们可以递归地计算$w_T$：

$$\mathbb{E}[w_T]=-3\eta T$$

将$\mathbb{E}[w_T]$代入测试损失的期望中：

$${\overline{\mathcal{L}}}_{\eta ,T}=\frac{1}{2}((-3\eta T{)}^2-6(-3\eta T)+9)+\frac{1}{2}$$

$${\overline{\mathcal{L}}}_{\eta ,T}=\frac{1}{2}(9{\eta }^2{T}^2+18\eta T+9)+\frac{1}{2}$$

$${\overline{\mathcal{L}}}_{\eta ,T}=\frac{9{\eta }^2{T}^2+18\eta T+10}{2}$$

2. 接下来，我们需要找到$g(T)$。

首先，我们需要最小化${\overline{\mathcal{L}}}_{\eta ,T}$关于$\eta$。我们可以通过设置$\frac{d{\overline{\mathcal{L}}}_{\eta ,T}}{d\eta }=0$来找到最优的学习率${\eta }^{\ast }$：

$$\frac{d}{d\eta }(\frac{9{\eta }^2{T}^2+18\eta T+10}{2})=9\eta T^2+18T=0$$

解得：

$${\eta }^{\ast }=\frac{2}{3T}$$

将${\eta }^{\ast }$代入${\overline{\mathcal{L}}}_{\eta ,T}$中，我们得到最小化测试损失的表达式：

$${\overline{\mathcal{L}}}_{{\eta }^{\ast },T}=\frac{9(\frac{2}{3T}{)}^2{T}^2+18(\frac{2}{3T})T+10}{2}$$

$${\overline{\mathcal{L}}}_{{\eta }^{\ast },T}=\frac{9(\frac{4}{9T^2})T^2+18(\frac{2}{3T})T+10}{2}$$

$${\overline{\mathcal{L}}}_{{\eta }^{\ast },T}=\frac{4+12+10}{2}$$

$${\overline{\mathcal{L}}}_{{\eta }^{\ast },T}=\frac{26}{2}$$

$${\overline{\mathcal{L}}}_{{\eta }^{\ast },T}=13$$

现在，我们需要找到$g(T)$，使得当$T\to +\infty$时，以下条件成立：

$$\left|\underset{\eta \in (0,\frac{1}{3}]}{\inf }{\mathcal{I}}_{n,T}-g(T)\right|=O\left(\frac{(\log T{)}^2}{T^2}\right)$$

由于我们已经找到了最优的学习率${\eta }^{\ast }$，我们可以将${\overline{\mathcal{L}}}_{{\eta }^{\ast },T}$视为${\mathcal{I}}_{n,T}$的下界。因此，我们需要找到一个函数$g(T)$，使得当$T$趋向于无穷大时，${\overline{\mathcal{L}}}_{{\eta }^{\ast },T}$和$g(T)$之间的差异满足上述条件。

考虑到${\overline{\mathcal{L}}}_{{\eta }^{\ast },T}$是一个常数13，我们可以推断$g(T)$应该也是一个常数，因为测试损失的期望在最优学习率下不随$T$变化。因此，我们可以选择$g(T)=13$。

现在，我们需要验证这个选择是否满足条件：

$$\left|\underset{\eta \in (0,\frac{1}{3}]}{\inf }{\mathcal{I}}_{n,T}-g(T)\right|=O\left(\frac{(\log T{)}^2}{T^2}\right)$$

由于${\mathcal{I}}_{n,T}$的最小值是13，我们有：

$$\left|13-13\right|=0$$

显然，$0=O\left(\frac{(\log T{)}^2}{T^2}\right)$，因为当$T$趋向于无穷大时，$\frac{(\log T{)}^2}{T^2}$趋向于0。因此，我们的选择$g(T)=13$是正确的。

综上所述，$g(T)=13$满足题目中的条件。

3.为了解决这个问题，我们需要推导出多维线性模型在给定数据分布下的缩放定律。根据题目描述，我们有一个线性模型，其参数遵循特定的缩放定律。我们将通过以下步骤来解决这个问题：

步骤 1: 理解数据分布

数据分布 $\mathcal{D}$ 是通过 Zipf 分布来选择输入向量的非零坐标，然后根据该坐标的值来生成输出 $y$。这意味着大部分的数据集中在较少的非零坐标上。

步骤 2: 定义损失函数

损失函数 ${\overline{\mathcal{L}}}_{\eta ,T}$ 是在给定学习率 $\eta$ 和训练步数 $T$ 后，模型参数 $\mathbf{w}$ 的测试损失的期望。

步骤 3: 推导缩放定律

我们需要找到 $\alpha$，$\beta$，和 $C$ 使得损失函数符合 ${\overline{\mathcal{L}}}_{N,T}\approx \frac{A}{N^{\alpha }}+\frac{B}{T^{\beta }}+C$ 的形式。

对于 $\alpha$ 的推导：

• 参数 $N$ 表示模型考虑的输入向量的维度。由于数据分布的特性，大部分的权重不会接收到有效的梯度更新，因为它们对应的输入坐标为零。因此，增加 $N$ 的数量不会显著改善模型的性能，但也不会损害它，因为只有少数权重会被更新。

• Zipf 分布的特性意味着非零坐标的数量随着 $N$ 的增加而减少。因此，我们可以预期 $\alpha$ 大于 0，但小于 1，因为增加维度对于模型性能的提升是有上限的。

对于 $\beta$ 的推导：

• 参数 $T$ 表示训练步数。随着训练步数的增加，模型将获得更多的机会来更新其权重，从而减少损失。因此，我们可以预期 $\beta$ 大于 0。

• 由于数据分布的特性，并不是每一步都会对所有权重进行有效更新。因此，$\beta$ 可能不会是 1，而是小于 1 的某个值。

对于 $C$ 的推导：

• 常数 $C$ 表示当 $N$ 和 $T$ 趋于无穷大时，测试损失的最低值。这是由于数据本身的噪声和模型的能力限制导致的。

步骤 4: 确定 $\alpha$，$\beta$，和 $C$

为了确定 $\alpha$，$\beta$，和 $C$，我们需要进行以下分析：

• 对于 $\alpha$：考虑到只有少数权重会被更新，我们可以假设 $\alpha$ 在 0 和 1 之间。更具体地，由于 Zipf 分布的特性，我们可以假设 $\alpha$ 接近于 1，但小于 1，因为随着 $N$ 的增加，额外维度的边际贡献会减少。一个合理的猜测是 $\alpha =\frac{1}{b}$。

• 对于 $\beta$：考虑到每一步并不是对所有权重都进行有效更新，我们可以假设 $\beta$ 小于 1。一个合理的猜测是 $\beta =\frac{1}{2}$，这是因为通常情况下，梯度下降的收敛速度与步数的平方根成反比。

• 对于 $C$：这是数据噪声和模型表达能力限制的结果。在没有更多信息的情况下，我们无法精确确定 $C$，但可以假设它是一个正数。

步骤 5: 验证条件

我们需要验证 $ϵ(N,T)$ 的条件是否成立。这通常涉及到对 ${\overline{\mathcal{L}}}_{N,T}$ 进行详细的分析，并证明它符合给定的缩放形式。这通常需要数学上的证明和/或实验验证。

综上所述，我们可以假设 $\alpha =\frac{1}{b}$，$\beta =\frac{1}{2}$，$C$ 是一个正数。然而，为了得到精确的值，我们需要更深入的分析和实验数据。在实际应用中，这些参数通常是通过实验来确定的。