CMU DLSys Course Homework 0: Implementation and Reflection
Jan. 01, 2025 - Jan. 04, 2025 · Qiyao Wang · Topic #MLSys
新年新气象,今年的新目标是不只做专利相关的任务,也不只做 Prompt 相关的任务。目前对 MLSys 和 Reasoning 比较感兴趣,从本文开始开一个新坑,从 CMU 的 DLSys 课程开始,学习一些底层的内容,希望未来也能做一些和 LLM 底层相关的工作。欢迎各位大佬指导我的不足!
1月1日看完了 Lecture 1 和 Lecture 2,本文是对 Homework 0 的实现和思考。由于没有课程账号,所有的实验内容均在我的个人电脑上进行。
Question 1: 编写代码、测试代码和提交代码的基本方式
由于不提交,我这里就只进行到测试阶段。
首先应该在官方的仓库中下载 Homework 0 的相关源码,虽然不使用 mugrade 提交到平台上,但是还是建议去官方仓库中下载 mugrade, 避免其中发生不必要的依赖项报错。下载好后在命令行里运行下面的代码,并且下载一些相关的包,其中需要用 pytest 来测试代码。
python setup.py install
pip install pybind11 numpy numdifftools pytest
测试环境配置,需要进入 src/simple_ml.py,将其中的待填写的代码块换为 return x + y
,之后在命令行中运行 python -m pytest -k "add"
。这其中的基本逻辑是 pytest 会去 tests 文件夹中寻找包含 “add” 的测试函数,在本作业中即为 tests/test_simple_ml.py 中的 test_add 函数,它通过调用 simple_ml.py 中的 add 函数,并通过 assert 来进行正确性的验证。
Question 2: 加载 MNIST 数据
首先需要前往 MNIST 的官方网站下载相应的数据(本作业的仓库中已经给出,在 data 文件夹中)。官方网站同时披露了 MNIST 数据集的相关情况,根据 hw0.ipynb 中的提示,需要使用 gzip 包对数据进行解构和读取,因此需要明晰数据集中 Label 和 Image 的格式。
如图1所示,为 MNIST 训练数据集中标签的格式,绿框中的内容为数据的描述性内容,共 8 个字节,在读取时需要先读取描述字段再读取后续的内容。由 src/simple_ml.py 中的代码可以知道,所需的标签格式为 numpy.ndarray[dtype=np.uint8]
。
如图2所示,为 MNIST 训练数据集中图片的格式,绿框中的内容是图片的描述字段,共16个字节,该数据集中的每个图片的大小为 28x28,因此在后续的形状调整中,将图片的像素拉成一行时,其大小为 28x28 = 784
维,其对应的格式为 numpy.ndarray[np.float32]
,并且需要进行全局归一化,即将像素的 0-255 取值放缩到 0-1.0 区间内,最终形成形状为 num_examples x input_dim
的 numpy 数组。
下面展示上述逻辑对应的代码,即读取 MNIST 数据集中的标签和图片,并对图片进行全局归一化和形状修改,最终返回相应的图片和标签。
# Loading Label Dataset
with gzip.open(label_filename, "rb") as label:
# Each label’s descriptor field is 8 bytes.
magic, n = struct.unpack(">II", label.read(8))
# Read to the end.
y = np.frombuffer(label.read(), dtype=np.uint8)
# Loading Image Dataset
with gzip.open(image_filename, "rb") as image:
# Each image's description field is 16 bytes.
magic, num, rows, cols = struct.unpack(">IIII", image.read(16))
# Read to the end.
X = np.frombuffer(image.read(), dtype=np.uint8)
# Reshape the image matrix set to align with the labels.
X = X.reshape(len(y), 784)
以读标签数据为例,其描述字段为 8 字节,通过 struct.unpack()
方法能够解构该二进制文件,其中参数 ">II" 中的 ">" 表示大端序(高位字节存储在低地址,低位字节存储在高地址),"I" 表示数据类型 int,每个 I 占用 4 个字节,因此占据 8 个字节的描述字段使用 II。read(8)
函数表示读取该二进制文件的前 8 位(由于无前序读取,因此从头开始),下次使用 read 函数时并非从头开始,此次读取会将指针移至第 9 位。因此下一行的 read 函数为从该二进制文件的第 9 位读至文件末尾。同时需要注意的是 np.fromstring()
函数在我当前的 numpy 版本 2.2.1 下已被弃用,应使用 np.frombuffer()
代替,其数据类型为 8 位无符号整数 uint8。
X = X.astype(np.float32) / 255.0
return X, y
根据函数描述,图片数据的格式应为 np.float32
,并且进行全局归一化,需要与单个图片的归一化进行区分,由于该数据集较为简单,可以使用 min-max 归一化,由于像素值区间为 0-255,因此这里就直接进行 /255.0 的操作。
Question 3: Softmax Loss
Softmax Loss 原理概述
以 ImageNet 的 1000 种图片分类为例,假设使用 AlexNet,在最后一层的卷积层后接入的是两层全连接层,最后一层全连接层的输出维度应与图片的种类相同,即为 1000 维的列向量(此处忽略 batchsize),该列向量可以称为 logits;在大语言模型中,预测下一个 token 时的最后一层也是一个与字典大小相匹配的向量,该向量也称为 logits。但是这样的 logits 向量并非每一类或每一个 token 的概率,其元素之和也并不满足概率的性质,即非负性、总和为1。在真实使用过程中,需要使用激活函数(二分类时通常使用 sigmoid,多分类时通常使用 softmax)来将其映射到概率空间。
Softmax 函数的定义如下:在一个具有 $k$ 类的分类任务中,第 $i$ 类的概率值为 $z_i$,h 为假设函数(能够预测该 input 对应的 logits)
$$z_i = p(label=i) = \frac{\exp{(h_{i}(x))}}{\sum_{j=1}^k \exp{(h_{j}(x))}} = normalize(\exp{(h(x))}) = softmax(h(x))$$其中指数函数能够确保将 logit 映射到非负区间内,通过求和进行归一化,从而满足概率在 0-1 区间内的要求。在分类任务中,真实分布通常是 one-hot(独热编码) 的形式,即假设该样本属于第 $i$ 类,则 $y = [..., 0, 1_{i-th}, 0, ...]$。最小化损失,则是希望能够最小化预测的分布与真实分布之间的差值,而对于分类任务而言,则通常使用 softmax 损失,在此上下文中,softmax loss 与 cross-entropy loss(交叉熵损失)是等价的。此处不证明交叉熵相关的内容,需要从熵的定义写起,互联网上已经有很多资料。
Softmax 损失(交叉熵损失)的具体定义如下:在一个具有在一个具有 $k$ 类的分类任务中,真实分布为 $y$,预测分布为 $h(x)$,与 hw0.ipynb 中的形式相对应。
$$\ell_{ce}(h(x), y) = -\log{p(\text{label} = y)} = -h_y(x) + \log{\sum_{j=1}^k\exp{(h_{j}(x))}} = -z_y + \log{\sum_{i=1}^k\exp{(z_i)}}$$代码实现
下面实现 simple_ml.py 中的 softmax_loss()
函数。其中输入 Z 为二维数组,即 batchsize 个样本,和每个样本对应的 logits。真值标签 y 是一个一维数组,有 batchsize 个元素,每个元素的值为标签对应的索引。在本次实现中并未对 softmax 进行缩放,但在实际计算时,如果 $z_j$ 数值过大,指数运算后会过大,可以采取减去 $z$ 的最大值如 $e^{z_j-\max(z)}$ 的方式来进行缩放,在本次实现中并不考虑这个问题。
def softmax_loss(Z, y):
# Compute e^x for each element in Z .
exp_Z = np.exp(Z)
# Create an index array of size equal to the batchsize.
index_ = range(exp_Z.shape[0]) # list: [0, 1, 2, ..., batchsize-1]
result = np.log(np.sum(exp_Z, axis=1)) # axis=1 for each sample (row)
result = result - Z[index_, list(y)] # 1D array
return np.average(result)
其中 Z[index_, list(y)]
为多行索引,能够提取对应的内容,下面是一个示例:
其中 result 为从 Z 中提取的对应的预测分布中的 logit,需要注意的是这是一个 batch 的数据,在随机梯度下降(SDG)中,通常针对一个 minibatch/batch 进行优化,因此在返回时需要计算 batch 中的平均损失。
Question 4: Stochastic Gradient Descent for Softmax Regression
对 Softmax loss function 进行随机梯度下降,最重要的是能够求出该目标函数的梯度。随机梯度下降一般使用一个 minibatch/batch 的样本的平均 loss 进行优化,在讨论多个样本下的梯度前,先推导单个样本时的梯度。对于假设函数 $h$ 和标签真值为 $y$ 的第 $i$ 个样本:
$$ \frac{\partial{\ell_{ce}(h,y)}}{\partial{h_i}}=\frac{\partial}{\partial h_i}(-h_y(x) + \log{\sum_{j=1}^k\exp{(h_{j}(x))}}) = -\mathbb{I}\{i=y\}+\frac{\exp{(h_i)}}{\sum_{j=1}^k\exp{(h_j)}} $$其中 $\mathbb{I}\{i=y\}$ 为判别函数,当 $i=y$ 即条件满足时为 1,否则为 0。在该公式中,$h_i$ 在 $i \neq y$ 的情况下,$h_y$ 对 $h_i$ 而言为常数,其导数为 0,当 $i=y$ 时,则导数为 1。对括号中第二项则使用链式法则进行求导即可。因此对 $h$ 求导为:
$$ \nabla_h\ell_{ce}(h,y)=z-e_y, \quad\text{where }z=normalize(\exp(h)) = softmax(h) $$其中 $e_y$ 为独热编码,即该向量仅在 $y$ 对应的索引处元素的值为 1,其余位置元素的值均为 0。作为假设函数,$h$ 是与模型参数相对应的,在上面的式子中均忽略了决定 $h$ 的重要因素,即参数 $\theta$,下式为本文所讨论的线性假设函数
$$ h_\theta(x)=\theta^Tx $$梯度下降真正的优化目标是模型的参数,通过假设函数 $h_\theta$ 与参数 $\theta$ 建立联系,因此需要推导在单一样本 $x$ 下,目标函数对于参数的导数,即
$$ \frac{\partial}{\partial\theta}\ell_{ce}(\theta^Tx,y)=\frac{\partial\ell_{ce}(\theta^Tx,y)}{\partial\theta^Tx}\cdot\frac{\partial\theta^Tx}{\partial\theta}=(z-e_y)x $$其中 $(z-e_y)\in\mathbb{R}^{k\times1}$,$x\in\mathbb{R}^{n\times1}$。由于 假设函数 $h$ 为将 $n$ 维的列向量 $x$ 映射为 $k$ 维的向量(本任务中假设分类类别有 $k$ 种),因此 $\theta^T\in\mathbb{R}^{k\times n}$,即 $\theta\in\mathbb{R}^{n\times k}$,损失函数对于 $\theta$ 导数的维度也是如此。而上式中两个部分相乘并不满足矩阵乘法法则,因此需要进行调整才可能形成 $n\times k$ 维的矩阵。
$$ \nabla_{\theta}\ell_{ce}(\theta^Tx,y)\in\mathbb{R}^{n\times k}=x(z-e_y)^T $$随机梯度下降一般采用 minibatch/batch 进行优化,即当 minibatch 的 batchsize 为 $B$ 时,输入为 $X\in\mathbb{R}^{B\times n}$,标签为 $y\in\{1,\dots,k\}^B\in\mathbb{R}^{B\times k}$,更新参数时需要使用批次的平均损失来更新,即
$$ \theta := \theta - \frac{\alpha}{B}\sum_{i=1}^B\nabla_{\theta}\ell(h_\theta(x^{(i)}),y^{(i)}) $$其中 $\alpha$ 为学习率(步长),因此与单个样本 $x$ 类似,需要推导在批次样本 $X\in\mathbb{R}^{B\times n}$ 下的损失函数的导数
$$ \frac{\partial}{\partial\theta}\ell_{ce}(X\theta,y) = \frac{\partial\ell_{ce}(X\theta,y)}{\partial X\theta}\cdot\frac{\partial X\theta}{\partial\theta}=(Z-I_y)\cdot X $$其中 $X\theta\in\mathbb{R}^{B\times k}, (Z-I_y)\in\mathbb{R}^{B\times k}$,而 $X\in\mathbb{R}^{B\times n}$,对 $\theta$ 的导数的维度为 $n\times k$,因此需要进行调整
$$ \nabla_\theta\ell_{ce}(X\theta,y) \in \mathbb{R}^{n\times k} = X^T\cdot(Z-I_y),\quad Z = normalize(\exp{(X\theta)}) = softmax(X\theta) $$其中 $I_y$ 为标签的独热矩阵。在 SGD 时需要求得 minibatch/batch 中 $B$ 个样本的平均损失来进行优化。
def softmax_regression_epoch(X, y, theta, lr = 0.1, batch=100):
# number of all samples: m
m = X.shape[0]
# the size of minibatch/batch: B
B = batch
# compute the number of iterations
iterations = (m + (B - 1)) // B
for i in range(iterations):
# compute the end index of this iteration step
num = min(m, (i + 1) * B)
# get the batch of data of this iteration step
batch_x = X[i*batch:num, :]
batch_y = y[i*batch:num]
# compute the differentiation based on the equation
h_y = batch_x @ theta
Z = np.exp(h_y) / np.sum(np.exp(h_y), axis=1, keepdims=True)
# Z - I_y with multi-index
Z[range(batch_y.shape[0]), batch_y] -= 1
gradient = batch_x.T @ Z
# compute the average of loss
gradient = gradient / batch_x.shape[0]
theta = theta - lr * gradient
其中迭代次数需要考虑两种情况,能够整除 batch 和不能整除 batch,通过 $\lfloor\frac{m+(B-1)}{B}\rfloor$ 来统一计算,如当 $m=10, B=4$时,迭代次数应为 $3 = \lfloor\frac{10+4-1}{4}\rfloor$次,或引入 math 库 import math
,使用 math.ceil()
对 $\frac{m}{B}$ 进行向上取整即可。在每一次迭代过程中,需要计算该迭代步骤中批次的最后一个元素的索引,由于不一定能够整除,因此需要与样本数 $m$ 进行比较。在 numpy 中矩阵乘法可以简单使用 @
符号完成,但需要注意的是维度需要对齐。在 np.sum
函数中,axis=1
是在行上的操作,keepdims=True
能够保证结果数组与输入数组的形状一致,如果不添加该参数,在行聚合后,会形成向量的形状如 [1,2]
,而非多维数组如 [[1],[2]]
。在减去独热矩阵时,使用多维的索引,与 Question 3 中的例子相同,通过二维索引来操作 Z 中的元素的值。注意在计算平均损失时,不要除 $B$,而应除真实的 batch 数据大小,因为有可能最后一个批次中元素的个数小于 $B$。
Question 5: SGD for a two-layer NN
在上面的内容中均未涉及假设函数的具体类型,与此同时,上面的 softmax 相当于作为输出层的激活函数,在下面使用两层的神经网络来进行实践,中间层中也应使用激活函数来提升相应的效果,假设输入样本 $x\in\mathbb{R}^n$。
$$ z = W_2^T\text{ReLU}(W_1^Tx), \quad\text{where }W_1\in\mathbb{R}^{n\times d}, W_2\in\mathbb{R}^{d\times k} $$由上面的维度,可以知道参数的维度为 $n\times k$,而输出的维度为 $k\times1$。站在数据集(含 $m$ 个样本)的全局角度
$$ Z\in\mathbb{R}^{m\times k} = (W_2^T\text{ReLU}(W_1^TX^T))^T=\text{ReLU}(XW_1)W_2, \quad\text{where }X\in\mathbb{R}^{m\times n} $$优化目标对应为
$$ \text{minimize}_{W_1,W_2}\text{ }\frac{1}{m}\sum_{i=1}^m\ell_{softmax}(\text{ReLU}(XW_1)W_2,y),\quad y\in\mathbb{R}^{m\times k} $$根据上面对 softmax loss 的导数推导,$h = \text{ReLU}(XW_1)W_2$, 记 $Z_1\in\mathbb{R}^{m\times d}=\text{ReLU}(XW_1)$,此处的 $Z$ 其实是假设函数 $h$ 的角色, 则 $Z=Z_1W_2$,分别对 $W_1$ 和 $W_2$ 求偏导,此处不管 loss 导数的平均值,在梯度下降时乘 $\frac{1}{m}$。
$$ \nabla_{W_2}\ell_{softmax}(Z,y) = \frac{\partial\ell}{\partial Z}\cdot\frac{\partial Z}{\partial W_2} = Z_1^T(Softmax(\exp{(Z)}) - I_y) = Z_1^TG_2, \quad\text{where }G_2 \in\mathbb{R}^{m\times k}=Softmax(\exp{(Z)}) - I_y $$对 $W_1$ 求导也需要使用链式法则
$$ \nabla_{W_1}\ell_{softmax}(Z,y) = \frac{\partial\ell}{\partial Z}\cdot\frac{\partial Z}{\partial Z_1}\cdot\frac{\partial Z_1}{\partial W_1}=G_2\cdot W_2^T\cdot \frac{\partial Z_1}{\partial W_1} $$由于 $\frac{\partial Z_1}{\partial W_1}$ 仍然是个复合函数,将其单独计算
$$ \frac{\partial Z_1}{\partial W_1} = \frac{\partial \text{ReLU}}{\partial XW_1}\cdot\frac{\partial XW_1}{\partial W_1} $$ $$ \text{ReLU}(x) = \begin{cases} x & \text{if } x > 0 \\ 0 & \text{if } x \leq 0 \end{cases} \quad\text{ReLU}'(x) = \begin{cases} 1 & \text{if } x > 0 \\ 0 & \text{if } x \leq 0 \end{cases} $$$\text{ReLU}$ 的导数可以看到为一个掩码矩阵,即在 $XW_1 > 0$的元素位置处的值为1,否则为0,此处记为 $\mathbb{I}\{Z_1>0\}$,将上面的链式法则进行整理,记 $G_1\in\mathbb{R}^{m\times d}=\mathbb{I}\{Z_1>0\}\circ(G_2W_2^T)$,其中 $\circ$ 为逐元素的运算,在本式中相当于抽取出 $G_2W_2^T$ 对应掩码矩阵中值为 1 的位置的元素则
$$ \nabla_{W_1}\ell_{softmax}(Z,y)=\frac{\partial Z_1}{\partial W_1}G_1=X^TG_1 $$
代码实现如下,其中 *
号在矩阵之间的运算表示逐元素乘法,在计算 $Z$ 及之后的 $Z_1$ 都已变为掩码后的矩阵。
def nn_epoch(X, y, W1, W2, lr = 0.1, batch=100):
m = X.shape[0]
B = batch
iterations = (m+ (B - 1)) // B
for i in range(iterations):
num = min(m, (i + 1) * B)
batch_x = X[i * B:num, :]
batch_y = y[i * B:num]
Z_1 = batch_x @ W1
relu_mask = Z_1 > 0
Z_1 = Z_1 * relu_mask
Z = Z_1 @ W2
G2 = np.exp(Z)/np.sum(np.exp(Z), axis=1, keepdims=True)
G2[range(batch_y.shape[0]), batch_y] -= 1
G1 = G2 @ W2.T
G1 = G1 * relu_mask
W2_grad = Z_1.T @ G2
W1_grad = batch_x.T @ G1
Summary
基于 CMU 的课程又回顾了一下 Softmax 的相关内容和梯度下降的相关内容,发现自己的基础还是比较薄弱,希望通过每周对于 CMU DLSys 课程的学习,在未来能够做一些不止 Prompt 的工作,继续加油!By the way,能坚持下来才是最重要的!不要三分钟热度!
Reference
[1] xx要努力. (2022, 10). Deep Learning System-Homework0. 知乎专栏. https://zhuanlan.zhihu.com/p/576378167.
[2] Some parts of the code or math were consulted from ChatGPT and KimiChat.
Contact
There may be some errors present. If you find any, please feel free to contact me at wangqiyao@mail.dlut.edu.cn
. I would appreciate it!