Machine Learning System Basic II: Basic nn.Module Layers
作者Author Yikai Zhu
发布于Published October 21, 2023
标签Tags machine learning systemclass note

序

这篇文章是CMU 10-414/714的学习笔记第二篇。主要概括课程Lec6-Lec10的内容。本文重点如下，欢迎大家评论补充。

  1. Initialization: 神经网络的初始化实现
  2. Optimization: 机器学习优化器实现

Initialization

教授强调了一点：初始化方式会对最终的结果产生非常大的影响，因为模型最终收敛时的参数与初始参数的距离是很小的。我们通常会认为机器学习是一个天资聪颖的神童，只要给他学习资料他就能无师自通；但是事实上模型更像是一个平庸的普通人，终其一生奋斗缺依然逃不了太远。另外，错误的初始化方式还可能会导致下面两个灾难性的结果:

img

事实上，有很多工作是关于如何初始化参数的。在 hw2中，我们会分别实现下面四种初始化方式。在给定了数学公式之后，用 python 实现起来是非常简单的。

但是，我们其实并不希望初始化方式有这么重要的影响。这里的一个思想就是引入一个新的层来解决这个问题(“fix” the normalization of the activations to be whatever we want.)。下面介绍的两个思想就是解决这个问题：

Layer Normalization LayerNorm 的思想是，既然 activation 会爆炸/消失，那我们就引入一个层来 activation 的均值和方差，让其大小保持在一个合理的范围内。具体公式如下，效果如下图。 \(y = w \circ \frac{x_i - \textbf{E}[x]}{((\textbf{Var}[x]+\epsilon)^{1/2})} + b\)

img

Batch Normalization 还有一个很奇怪的想法是 BatchNormalization 。其采用同 Layer Normalization 的公式，只是均值和方差是 Batch 这个维度的。计算公式如下: $y = \frac{(x - \hat{m}u)}{((\hat{\sigma}^2_{i+1})_j+\epsilon)^{1/2}}$。这里的 $\hat{m}, \hat{\sigma}$ 不是需要计算梯度的参数，仅仅是代表一个需要更新的状态 $\hat{m} = \beta \hat{m} + (1-\beta)\mathbf{E}\left[z\right], \hat{\sigma} = \beta \hat{\sigma} + (1-\beta)\mathbf{Var}\left[z\right], \text{where} \beta is \text{momentum}$.

Other normalization

Regularization Regularization的基本思想是限制权重的大小(“limiting the complexity of the function class”)，也就是在 loss 里加上 $\frac{\lambda}{2} \sum_{i=1}^D\left|W_i\right|_2^2$。这个在数学上等效于在 optimizer step 的时候加一个 weight decay。证明过程如下： \(W_i:=W_i-\alpha \nabla_{W_i} \ell(h(X), y)-\alpha \lambda W_i=(1-\alpha \lambda) W_i-\alpha \nabla_{W_i} \ell(h(X), y)\)

Optimization

作为一个类，Optimizer 只需要实现一个 step 方法，其会更新模型的参数。下面介绍几个常常见的优化器，最值得关注的是 SGD 和 Adam。

Gradient Descent: $\theta_{t+1}=\theta_t-\alpha \nabla_\theta f\left(\theta_t\right)$，其中 $\alpha$是 learning rate。

SGD: SGD 是 Gradient Descent 的一种变形。基本思想是: Instead of taking a few expensive, noise-free, steps, we take manycheap, noisy steps, which ends having much strong performance per compute. \(\theta_{t+1}=\theta_t-\frac{\alpha}{B} \sum_{i \in B} \nabla_\theta \ell\left(h\left(x^{(i)}\right), y^{(i)}\right)\)

Momentum: 下面公式中的$u_{t+1}$就代表动量(momentum)，是一个需要存储的变量: \(\begin{array}{l} u_{t+1}=\beta u_t+(1-\beta) \nabla_\theta f\left(\theta_t\right) \\ \theta_{t+1}=\theta_t-\alpha u_{t+1} \end{array}\)

Adam: 相较于 Momentum, Adam 额外引入了 $v_{t+1}$ 来估计 scale of gradients，其值为梯度平方的级数求和。个人理解这个值代表了梯度的 norm: \(\begin{aligned} u_{t+1} & =\beta_1 u_t+\left(1-\beta_1\right) \nabla_\theta f\left(\theta_t\right) \\ v_{t+1} & =\beta_2 v_t+\left(1-\beta_2\right)\left(\nabla_\theta f\left(\theta_t\right)\right)^2 \\ \theta_{t+1} & =\theta_t-\alpha u_{t+1} /\left(v_{t+1}^{1 / 2}+\epsilon\right) \end{aligned}\)

附录