采样是一种运算

引子¶

在 2025 年后谈生成模型，大语言模型(llm) 一定是绕不开的一种生成模型。 它的使用体验和它工作的基本逻辑的一致性很高。 在使用大语言模型类的产品的时候，我们只需要以字符串的形式输入我们想做的事，它就会给我们几个字几个字的慢慢蹦出来，最后变成一个完整的回答。 在相对更细节的模型层面上，模型是这么运作的:

1. 把用户输入的问题（通常称为 token）和模型已经回复的那些词，输入模型。假设这里输入的总长度为 $t$ ，则这些词对应记为 $x_1, x_2, ... x_t$ ，每个 $x_i$ 是可能取值为 $1, 2, ..., V$ 的一个整数，表示对应的词在词表里的编号（$V$即词表总大小）。
2. 把这些输入给一个神经网络 $f_\theta$ ，得到的输出是一个 $V$ 维向量，这里记为 $P_{t + 1}$，表示下一个位置的词的取值概率。 通常来说 $P_{t + 1}$ 成为一个合法的概率的处理方式，和常规的分类模型类似，先让模型前面的层输出一个无约束的 $V$ 维向量，然后做 softmax。 这些步骤，抽象成公式表达即 $P_{t + 1} = f(x_1, x_2, \dots, x_t), P_{t + 1} \in \mathbb{R}_+^V$ 。
3. 根据概率表 $P_{t + 1}$， 采样 出下个词 $x_{t + 1}$
4. 循环重复以上过程，直到采样到词表里的特殊结束词<EOS>，便结束循环，完成作答。

上面的循环里涉及到很多的对象和操作。 深度学习入门课程中，可能会关注怎么处理“词”/token，这种加减乘除等常规无意义的数据元素，给你介绍 embedding 这种概念； 可能会关注模型用什么，也即 $f_\theta$ 的具体构造，从 RNN 说到 Transformer。

但很少有强调根据概率表中采样下一个词，这个过程的。 因为它可能是以上循环涉及的过程和对象里，最直观，不需要太多解释就能理解到过程。

但在这篇文章中，我们就需要详细拆解一下抽样这个过程，了解其中的细节，更重要的是，给之后要串讲的其他的生成模型的设计动机讲解，埋下一个种子。

流程拆解¶

直接给出问题：现有一个 $V$ 维度的概率表 $P$， $P$ 第 $i$ 个通道的值 $P_i$ 表示抽出的随机样本 $x = i$ 的概率；还有一个服从 $[0,1]$ 上均匀分布也即 $U(0, 1)$ 的随机数源，你可以从中任意的采样随机数 （比如 torch.rand 这种函数接口）；设计出返回一个随机样本的程序，这个随机样本服从概率表 $P$ 所代表的类别分布(Categorical distribution)

解：从一种对 $P$ 的可视化方法出发，我们先给出程序的过程，再给出如此过程产生的随机数是服从 $P$ 所代表的分布的说明。 我们可以把 $p_1, p_2, ..., p_{V}$ 转为等长度的线段，然后头尾相接地排开，即下图的上半部分：

$P$ 是一个概率表，这个约束就代表着 $\sum_{i=1}^V p_i = 1$，即上面的拼接线段长度为1，可建立一个数轴，左端点是0，右端点是1。再从$U(0, 1)$中抽出一个随机样本 ε，看 ε 在数轴上的位置，坐落在哪个线段里，假设坐落在 $p_i$ 所对应的线段里，那么就令返回值 $X = i$ ，然后返回 $X$。

最后再验证一下，这样的采样过程，满足我们想要的要求。对任意 $i$, $P(X = i) = P(\epsilon \in [\sum_{n=1}^{i-1} p_n, \sum_{n=1}^{i-1} p_n + p_i]) = p_i$。

1. 第一个等号的根据是： $(X = i)$ 和 $\epsilon \in [\sum_{n=1}^{i-1} p_n, \sum_{n=1}^{i-1} p_n + p_i]$ 是等价事件；
2. 第二个等号的根据是 ε 服从 $U(0, 1)$ ，

从 $i$ 的任意性，以上过程返回的 $X$ 确实服从 $P$ 所代表的类别分布。

把上面的过程，用更紧凑的数学公式总结就是：

服从其他分布的随机数的采样¶

把上面的过程再做一般化的推广，当我们从一个复杂的分布采样的时候，可以这样考虑。

1. 从一个简单的分布(如上一节的 $U(0,1)$ )采样一个随机样本 ε 。这个的简单分布也被成为熵源。
2. 构造函数 $g$ ，取 $X = g(\epsilon)$ ，返回 $X$ 即作为采样样本。这里的 $g$ 的构造需求是使得 $X = g(\epsilon)$ 服从你想要的目标分布。

不难看到，以上过程难点是构造出满足“$X = g(\epsilon)$ 服从目标分布”的函数 $g$ 。 接下来我们给出一些场景下构造函数 $g$ 的思考过程，也给后续的文章埋些伏笔。 没有特殊说明的情况下，在下文几个小节中，熵源总取为 $U(0, 1)$

行文中整体考察分布的思路是从简单的分布到复杂的一般化的分布，从一维的分布到多维的分布。

范围一般化的均匀分布采样¶

若我们要达成在 $U(a, b)$ 上的采样，（左右端点 $a, b$ 为已知参数），类比上面类别分布采样的构造思路，直接能想到的一个构造起点是是一个带放缩和端点平移的变换，在线段 $[a, b]$ 和线段 $[0, 1]$ 上做变换。 先抽样 $\epsilon \sim U(0, 1)$ ， ε 坐落在线段$[0, 1]$内，通过上的平移放缩变换找到在 $[a, b]$ 上的对应点，就作为采样值返回即可。 以上过程所对应的函数，严格用公式写下就是：

最后简单验证 $g(\epsilon) \sim U(a, b)$： 从 $g$ 的单调性和可逆性，我们可以得到以下的事件等价关系 -- 对任意的 $l \geq a, r \leq b$ ，事件 $g(\epsilon) \in [l, r]$ 和事件 $g^{-1}(g(\epsilon)) \in [g^{-1}(l), g^{-1}(r)]$ 等价，代入具体形式，该事件就是 $\epsilon \in [\frac{l - a}{b - a}, \frac{r - a}{b - a}]$ 。

从 $\epsilon \sim U(0, 1)$ ，事件 $\epsilon \in [\frac{l - a}{b - a}, \frac{r - a}{b - a}]$ 发生的概率为区间长 $\frac{r - a}{b - a} - \frac{l - a}{b - a} = \frac{r - l}{b - a}$，这也是 $X = g(\epsilon) \in [l, r]$ 发生的概率。 从 $l, r$ 的任意性，和以上求出的 $X$ 在 $[l,r]$ 上的概率的表达式，可知 $X \sim U(a,b)$

这是比较直接能想到的构造方法。 并且，对于这个简单的问题，我们还可以发现，$g$ 的构造方法并不唯一。 可以验证，以下的构造动机和最终构造出的 $g$ 也是符合 $g(\epsilon) \sim U(a, b)$ 的要求的：

1. 原来的 $g$ 的平移放缩是保证线段 $[0, 1]$ 和 $[a, b]$ 的左端点对左端点，右端点对右端点。 但直觉告诉我们，左端点对右端点，右端点对左端点，放缩倍率不变，也应该是可行的。 即构造 $g(\epsilon) = b - (b - a) \epsilon$
2. 把 $[0, 1]$ 线段等分成若干份等长首尾相连的小线段，每个小线段类似原来的 $g$ 构造平移放缩，左端点映射到 $a$ 右端点映射到 $b$。 即当 $\epsilon \in [0, 0.5)$ 时， $g(\epsilon) = a + (b - a) \frac{\epsilon}{0.5}$， 当 $\epsilon \in [0.5, 1]$ 时， $g(\epsilon) =a + (b - a) \frac{\epsilon - 0.5}{0.5}$

标准正态分布¶

通过上面两个例子，我们已经慢慢有种感觉是，如果把 $g$ 构造范围限制在单调递增函数，那么对任意的 $[0, 1]$ 区间上的小线段 $l$， $g$ 应当把它映射到一个新的区间段，使得这个区间段的取值概率（在目标分布的意义下）就是 $l$。 在 $U(0, 1)$ 到 $U(a, b)$ 的例子中，这就是函数里乘的倍率项 $(b - a)\epsilon$。

在 $U(0, 1)$ 转换为 $\mathcal{N}(0, 1)$ 这个问题上，我们用同样的思路尝试求解。 以下是草稿本上的不严密的尝试过程

首先对于分布 $\mathcal{N}(0, 1)$ 它取值为非零的区间（如果你不排斥更多数学概念，它就是分布的支撑集）是实数集$\mathbb{R}$。 所以需要考虑的映射后的区间段的左右端点可取成 $\forall l, r \in \mathbb{R}, l < r$ ，而不像是 转换为 $U(a, b)$ 的问题中，需要在 $[a, b]$ 中取区间端点。

因为我们已经限制了 $g$ 是单调递增函数，则 $g$ 有反函数 $g^{-1}: \mathbb{R} \to [0, 1]$，把最上方的构造思路翻译成数学语言即为

而 $P(X \in [l, r])$ 的计算就是计算标准正态的概率密度函数 (probability density function, pdf)在 $[l, r]$ 上的积分，或者更紧凑地用标准正态的累计分布函数 (cumulative distribution function, cdf) $\varPhi$ 在 $r, l$ 两点的取值作差来求。 这里我们为了形式上的对仗，使用后者，则化简构造的 $g$ 所需满足的约束是

我们看到上式中的两个要点，1. 式子对 $l, r$ 的任意性， 2. 标准正态的 cdf 也是一个单调递增函数，从这两点出发， $g$ 的一种构造就可以是： $g$ 满足 $g^{-1}(x) = \varPhi(x), \, \forall x \in \mathbb{R}$。再利用第二点，把反函数取回去，就得到 $g(x) = \varPhi^{-1}(x)$

做完了以上不严密的尝试，就要写严密的证明了。 本文中补充这些思考过程，是因为当初学者丢却思考步骤直接看证明，可能会觉得证明里的构造就像天外来物，机械降神，摸不着头脑，但是一步步跟下来就是对的，从而产生挫败感。

以下是严密证明：

取 $g = \varPhi^{-1}$ ，其中 $\varPhi$ 为标准正态分布的累积分布函数。 从 $\varPhi$ 的定义，它是单调递增的，$\mathbb{R} \to [0, 1]$ 的函数，从而 $\varPhi$ 的逆函数 $\varPhi^{-1}$ 的定义是合法的，且 $\varPhi^{-1}$ 是单调递增，$[0, 1] \to \mathbb{R}$ 的函数。

为验证 $g(\epsilon) \sim \mathcal{N}(0, 1)$，取 $\forall l < r \in \mathbb{R}$，构成试探区间 $[l, r]$，下求事件 $g(\epsilon) \in [l, r]$ 的概率，从而能得到关于 $g(\epsilon)$ 的分布信息。

其中第一个等号使用了 $g$ 是单调递增函数，对事件做了等价转换，第二个等号将 $g, g^{-1}$ 做了函数复合相消，第三个等号依据的是 $\epsilon \sim U(0, 1)$ 的假设，第四个等号是代入了 $g = \varPhi^{-1}$ 的构造定义，并在定义式两边取逆函数。

则从 $\varPhi$ 是标准高斯分布的 cdf，以及 $l, r$ 的任意性，证得 $g(\epsilon) \sim \mathcal{N}(0, 1)$

一般的一维分布¶

本节中，更一般化地，考虑如何构造 $g$ 使得 $g(\epsilon)$ 服从任意给定的一维分布。

首先，我们要考虑描述一个目标一维分布的方法有很多中，比如给出它的 pdf, 给出它的 cdf 等等。 但这些描述之间都是可以相互转换的。 最终我们需要的是目标一维分布的 cdf，把它记为 $F$

从此出发，上一节标准正态分布的转换函数 $g$ 的构造思路的思考基本是可以复用的，把思考过程中的标准正态分布，换成目标一维分布，思路是完全通的，所以不再重复思考过程。

则取 $g = F^{-1}$，以下验证 $g(\epsilon)$ 服从目标的一维分布。（你会发现证明过程也是类似的）

从 $F$ 的定义，它是单调递增的，$\mathbb{R} \to [0, 1]$ 的函数，从而 $F$ 的逆函数 $F^{-1}$ 的定义是合法的，且 $F^{-1}$ 也即 $g$ 是单调递增，$[0, 1] \to \mathbb{R}$ 的函数。

为验证 $g(\epsilon)$ 服从目标的一维分布，取 $\forall l < r \in \mathbb{R}$，构成试探区间 $[l, r]$，下求事件 $g(\epsilon) \in [l, r]$ 的概率，从而能得到关于 $g(\epsilon)$ 的分布信息。

则从 $F$ 是目标分布 cdf，以及 $l, r$ 的任意性，证得 $g(\epsilon)$ 服从目标分布。

边长为 1 的正方形上的均匀分布¶

从这个例子开始，我们慢慢过渡到高维空间的分布，为简单期间先考虑一个二维的简单分布。 不是一般性地，我们可以把正方形放入坐标系，认为这是在 $[0, 1] \times [0, 1]$ 这个区间范围内的正方形。

然后我们会在第一步就开始犯难。 观察目标分布，我们需要构造的 $g$ 首先应该满足定义域和值域应该是 $[0, 1] \to [0, 1] \times [0, 1]$，从一维打到二维，且满射。 对于很多读者来说，找到这样的函数都很难，但确实有这样的函数，对于有经验的读者，可以基于尝试 peano 曲线、Hilbert 曲线等函数来构造 $g$，这里不展开这个思路。

但跳出构造从一维到二维的符合要求的 $g$ 这样的约束，假设我们能在 $U(0, 1)$ 做两次独立同分布的采样得到两个样本 $\epsilon_1, \epsilon_2 \stackrel{\text{i.i.d.}}{\sim} U(0, 1)$，那直接用数对 $(\epsilon_1, \epsilon_2)$ 作为最终的返回结果就好。

举出这个例子是想说明，除了直接构造函数 $g$ ，我们还可以考虑用更多次的独立同分布采样并排成有序数组，来得到更高维的结果，这是和一维情形下显著不同的点。 并且我们也可以看到，如果函数 $g$ 是一个多值输入多值输出的函数 $g: \mathbb{R}^m \to \mathbb{R}^n$，当 $g$ 的输入是一个 $m$ 维在 $[0,1]^m$ 上的随机分布采样出的随机变量 ε，则输出 $g(\epsilon)$ 就是一个 $n$ 维的随机变量，这里和一维的时候还是相似的。

单位圆上的均匀分布¶

这个要求看上去只是正方形均匀分布的一个简单拓展，更换了一个形状，但做下去我们就会发现难度陡增加。

首先看我们很难像一维分布的情形下，从 cdf 做切入。 二维随机变量的 cdf 的定义是 $F(x, y) = P(X \leq x, Y \leq y)$ （其中 $(X,Y)$ 从目标分布中抽样而来）。 在二维函数中，就没有单调性这个概念了，也不存在 $F$ 可逆这个性质了。 而且我们也确实能很容易在具体的例子中，寻找到两个不同的点 $(x_1, y_1), (x_2, y_2)$ 使得 $F(x_1, y_1) = F(x_2, y_2)$ ，比如说在单位圆上均匀分布这个案例中，关于直线 $y = x$ 对称的任意两个点对，都是符合这个要求的 $(x_1, y_1), (x_2, y_2)$。

对于本问题，这里给出三种转换思路。

1. 分轴向处理，转换为多个一维分布的采样。将 $p(x, y)$ 用条件分布的定义拆分成 $p(x, y) = p_{Y|X}(y|x)p_X(x)$，右边两项边缘分布和条件分布都能解析求解，从而变成了两个一维随机变量的序贯采样，先采样出 $x$ 再条件于 $x$ 采样出 $y$。只要问题变到了一维随机变量的范畴内，都可以用 cdf 反函数来通用的构造采样变换，从而使问题可解。当然，这里也可以切换拆分方式成 $p(x, y) = p_{X|Y}(x|y)p_Y(y)$，对应地先采样 $y$ 再采样 $x$
2. 构造二维输入输出的变换 $g$，使得从二维简单分布（如上一节中的正方形上均匀分布）采样得到 ε 后，$g(\epsilon)$ 服从目标分布。但这样的 $g$ 具体怎么构造，就要根据问题具体分析了，很难找到一维分布中那种通式通法。
3. 采样放弃法：这同样也是一种利用采样次数的方法，并且可以在更高维做扩展。我们先在 $[-1, 1] \times [-1, 1]$ 这个正方形上的均匀分布做采样，得到采样点 $x, y$ 后计算它到原点的距离 ρ，如果 $\rho \leq 1$ 则作为结果返回，若 $\rho > 1$ 则丢弃，再重新在 $[-1, 1] \times [-1, 1]$ 上的均匀分布采样，直到满足 $\rho \leq 1$。

伏笔：任意高维随机分布¶

在单位圆上的随机分布采样的案例中，我们已经遇到了一些困难，这些困难其实是在高维随机变量的采样中同样存在。 在做高维随机分布的采样时，同样也是三种转换思路，即 分轴向处理, 构造高维输入输出的函数 $g$ 多次采样和采样拒绝。 这里略微展开各种转换思路对应的设问思考，为之后讲各种深度生成模型埋下伏笔。

1. 对于 分轴向处理 类方法：维度变高以后，采样轴向的先后顺序为 $n!$ 种，这就是很多种可行顺序了。这些采样顺序有优劣势之分吗？在不同类型的分布采样里，答案会不一样吗？
2. 对于 构造高维输入输出的函数 类方法：所构造的函数的输入维度一定要和输出维度等维度吗，高一些或者低一些会怎么样？
3. 对于 构造高维输入输出的函数 类方法：在高维空间里，很难像之前对一维和二维简单分布的分析那样，给出构造的函数的解析表达，但在高维情形下这种解析求解大概率很繁复，那如果考虑直接用神经网络来充当分布采样转换函数，这是可行的吗？这时 怎么处理网络的学习过程呢？
4. 对于多次采样和采样拒绝类方法：采样次数、拒绝条件目前看来是一个个针对具体问题的巧思构造，有没有数学上统一的对这类方法的描述语言，帮助我们更好的理解和选取采样次数、拒绝条件?

总结¶

本文介绍了如何根据不同的分布形式，执行具体的样本采样的过程。 在这里，核心想让读者了解到的观点是：

1. 随机样本的采样，实际上是通过对从熵源中得到的简单随机变量进行确定性的运算加工得来的。
2. 这样的确定性的运算是需要根据目标分布的不同，而做不同的构造的。
3. 构造运算的方式，随随机变量/随机向量的维度的升高而难度逐渐增加。
4. 确定性运算的形式很多，可以是各种数学函数，可以是多次的采样操作，当然在本栏目中，我们更会关注的是把神经网络作为这些运算的中间或者全部过程，会擦出什么神奇的火花。

前三点是关于文章主要内容的总结，而关于最后一点，我们将会在接下来的文章直接揭晓。

最后是本文的彩蛋小节，这里会写一些插入文章主线会导致叙事节奏不流畅的内容，但这些内容确实会对理解本文的内容有帮助，可以选择性地阅读。

拓展彩蛋¶

linux 系统视角下的随机数发生机制¶

在正文中，我们默认选择的熵源是均匀分布 $U(0, 1)$。 这是因为在数学上它的性质较好，且能方便地衔接到任意一维随机分布的采样方法。 熵源分布会给人一种所有随机性的根本来源的感觉，地位比其他的随机分布高一档，所有的其他分布的发生与采样，都利用它的采样结果配合合适的确定性变换来完成。 像 numpy, torch 等知名库，也确实把诸如 random, rand 等派生/简写于“随机”含义，并无特指的单词 random 的函数接口，用来表达在均匀分布 $U(0, 1)$ 中抽象，而为高斯分布，类型分布，加上特指的后缀(randn, randint)。 看上去均匀分布确实处于这个根源地位。

但实际上，$U(0, 1)$ 并非处于这个枢轴位置。 $[0, 1]$ 区间上的实数有无穷多个，在计算机中是无法被全部精确表达的，我们只能用浮点数来表达其中某些值。 从这个角度看，$U(0, 1)$ 的抽样结果和计算机的“本土表达”形式格格不入，并不太可能像那个最原初的随机性来源。 而实际上也确实是这样，在计算机中，是先获得了均匀的离散 $0, 1$ 二值分布的 bit 序列，再将其按需转换为浮点数，来实现均匀分布的发生的，这里放一个 numpy 源码中的例子。

如果 bit 的每一位服从 $0, 1$ 二值分布，概率均等，且位与位之间独立，那很直觉地我们就能看到，如果我们每 32 位地来将其视为 uint32，那得数自然服从 $0, 1,\dots, 2^{32} - 1$ 这么多种可取情况，每种情况等概率的类别分布。 反之如果将其均匀随机的 uint32 一位一位地发出去，那我们会得到一个独立同分布等概率的 $0, 1$ 二值序列的 bit 流。

如果只是需要“生成均匀或近似均匀的 32 位的无符号整数序列”（也就是后文加以区分后的伪随机序列），我们可以考虑以下的简单迭代

其中 $m = 2^32$ 使得结果总在 $0, 1,\dots, 2^{32} - 1$ 内，为保证均匀性对 $a, c$ 有一定的要求，可以是 $c$ 与 $m$ 互质（使得 $0, 1,\dots, 2^{32} - 1$ 都能取到）， $a - 1$ 是 $m$ 的倍数（使得数列 ${X_n}$ 的周期足够长），起点 $X_0$ 任取

这其实就是线性同余生成器 这种最经典易懂的伪随机数发生器。 实践中的伪随机数发生器一般会采用不同地更复杂的算法，但都可以看成是类似这样的数列递推公式的形式。 使得随机数发生器达成，需要多少，计算机就给你往下算多少，然后将结果做成字节流返回的效果。

你可以在 linux 中真的访问到这样的东西，伪随机数发生器文件是 /dev/urandom，我们通过如下示例命令，可以从中读取 16 个字节的随机数，并格式化成16进制数

dd if=/dev/urandom bs=16 count=1 status=none | xxd -p

伪随机数发生有个好处，当我们给定起点 $X_0$ 的时候，整个数列 ${X_n}$ 的形态就完全确定下来了。 从它抽出的随机序列，也就是可重复的，用这些随机序列再参与别的确定性运算，结果也是可重复的，而 $X_0$ 也就是通常所说的随机数种子。

但当我们不给随机数种子的时候，我们会发现含有随机运算的程序一样能运行，而且不同次运行的结果并不一样，那我们不免要问，这个作为数列初始值的 $X_0$ 是从哪来的呢？ 答案是，来自真实世界的不可预测的行为，通过将实时的硬件事件（如键鼠输入的坐标，磁盘操作的磁道坐标，传感器探测值等）序列化，并读这个序列，我们就能得到另一组随机序列。 这种情况下，我们会发现，随机性的根源是现实的物理世界，确实很难说，能通过计算机程序的控制，来确定性地改变这些观测值，所以这种序列也被称为真随机序列。

同样你也可以在 linux 中真的访问到这样的东西，真随机数发生器文件是 /dev/random，我们通过如下示例命令，可以从中读取 16 个字节的随机数，并格式化成16进制数

dd if=/dev/random bs=16 count=1 status=none | xxd -p

当你不指定种子/伪随机发生器的初始值$X_0$ 的时候，一般的程序往往会从真随机发生器中取一个值，作为它的种子，这样就能开始迭代了。 而从真随机发生器的原理，我们还可以想到的是，它是基于对客观物理世界的实时观测值得到的，当你试图一次性取用过多值，就只能阻塞地等待新的观测值。 而伪随机数发生器是没有这样的限制的，想取多少，就往下算多少，如果串行地计算伪随机数序列成为了时间瓶颈，也可以开发一些可并行化的伪随机发生算法来做优化。 我们很难去控制用户程序对于随机数使用量的多少，也并不希望随机数抽样成为性能瓶颈，所以各种库大都会使用真随机发生器取种子之后，用伪随机发生器跑完程序生命周期的配置。

网络初始化的耗时¶

在我们(从头)训练网络的代码中，通常会将网络作为一个对象，如果你了解得稍微细一点在这个对象初始化的时候， pytorch 会为网络中的参数申请相应的内存/显存空间，并执行参数初始化。 举个例子，以下是 nn.Linear 的 __init__ 函数（源码）

def __init__(
    self,
    in_features: int,
    out_features: int,
    bias: bool = True,
    device=None,
    dtype=None,
) -> None:
    factory_kwargs = {"device": device, "dtype": dtype}
    super().__init__()
    self.in_features = in_features
    self.out_features = out_features
    self.weight = Parameter(
        torch.empty((out_features, in_features), **factory_kwargs)
    )  ## torch.empty 申请 weight 所需的存储空间
    if bias:
        ## torch.empty 申请 bias 所需的存储空间
        self.bias = Parameter(torch.empty(out_features, **factory_kwargs))
    else:
        self.register_parameter("bias", None)
    self.reset_parameters()  ## reset parameter 里执行参数的初始化

def reset_parameters(self) -> None:
    # Setting a=sqrt(5) in kaiming_uniform is the same as initializing with
    # uniform(-1/sqrt(in_features), 1/sqrt(in_features)). For details, see
    # https://github.com/pytorch/pytorch/issues/57109
    init.kaiming_uniform_(self.weight, a=math.sqrt(5))  ## 用分布中采样的结果，填充 weight
    if self.bias is not None:
        fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
        bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
        init.uniform_(self.bias, -bound, bound) ## 用分布中采样的结果，填充 bias

从代码描述中不难看出，整个过程确实是分成申请内存，和参数初始化两个分离的步骤的。 回忆本篇文章的标题和核心观点是“采样是一种运算”。 而运算本身对于 CPU/GPU 来说，是一个需要消耗时间执行的任务。 如果增大参数的大小，需要执行的采样的次数就变多了，折算成的运算就更多，在这个大小足够大的情况下，我们应该是能在外面观察到人类可感知的程序运行的运算卡顿的，这也从侧面说明了“采样是一种运算”。 简单写一个程序验证一下：

import time

import torch.nn as nn

def init_cbyc_linear(C):
    start = time.time()
    m = nn.Linear(C, C, bias=False)
    end = time.time()
    print(f'init {C} by {C} linear in {end - start} s')

for i in range(15):
    init_cbyc_linear(2 << i)

'''
在我本机运行上的一次示例结果
init 2 by 2 linear in 0.00039649009704589844 s
init 4 by 4 linear in 4.839897155761719e-05 s
init 8 by 8 linear in 4.100799560546875e-05 s
init 16 by 16 linear in 3.528594970703125e-05 s
init 32 by 32 linear in 3.314018249511719e-05 s
init 64 by 64 linear in 4.220008850097656e-05 s
init 128 by 128 linear in 9.322166442871094e-05 s
init 256 by 256 linear in 0.0002980232238769531 s
init 512 by 512 linear in 0.0010356903076171875 s
init 1024 by 1024 linear in 0.0041387081146240234 s
init 2048 by 2048 linear in 0.017512083053588867 s
init 4096 by 4096 linear in 0.06756138801574707 s
init 8192 by 8192 linear in 0.2654736042022705 s
init 16384 by 16384 linear in 1.0543007850646973 s
init 32768 by 32768 linear in 4.210833549499512 s
'''

可以看到在最后两个大参数量的线性层的初始化的过程中，初始化的用时确实是人类比较可感知的时长了。 当然，这里的实验并不够细致，nn.Linear 的 __init__ 有很多步骤，从源码判断主要步骤是 1. 申请内存/显存 2. 随机抽样，并将抽样结果填入相应的内存/显存，完成初始化。 说不定是申请大的内存/显存块，所使用的时间变多了呢？ 这里我们以 $2^{14} \times 2^{14}$ 的线性层为例，去做拆解。

import math
import time

import torch
import torch.nn as nn

C = 2 << 13
start = time.time()
weight = nn.Parameter(torch.empty((C, C)))
end = time.time()
print(f"memory allocation time: {end - start}")

start = time.time()
nn.init.kaiming_uniform_(weight, a=math.sqrt(5))
end = time.time()
print(f"random sampling init time: {end - start}")

"""
在我本机上的一次运行结果
memory allocation time: 0.000118255615234375
random sampling init time: 1.0812816619873047
"""

如此拆解以后，我们确实可以看到，申请存储空间并不需要很多时间，大部分的时间都花在了随机抽样做参数初始化的阶段。 这充分说明了，采样是一种运算，是一种耗时不可小视的运算的观点。 同时我们也可以断定，之前所做实验观察到的“随着参数量增加，初始化线性层的时间增加”的现象，的主要原因就是，做随机抽样初始化参数的时间增加了。

同时，我们也可以看看本例中的具体的数值，初始化一个 $2^{14} \times 2^{14}$ 维的无 bias 的线性层，需要初始化的参数量是2²⁸，大约就是 1B。 1B 参数量在大模型时代之后看来并不大，是一个能在实际实践中常遇到的参数量大小。 在很多 pytorch 项目的训练代码里，对模型往往是这样处理的，1. 初始化 2. 如果传入了预训练权重，就加载预训练权重， 3. 将模型放到对应的 GPU 上。 也即：

model = Model()
if args.pretrained_ckpt:
    model.load_state_dict(torch.load(args.pretrained_ckpt))
model = model.cuda()

但以读完文章后的认识，我们会发现这段代码有些可改进的地方

1. 在 Model() 这里实例化 Model 对象的时候，同时在做的是内存申请 + 模型参数初始化。 那如果后续走进了 pretrained_ckpt 分支，其实我们发现，之前的初始化中的计算完全是多余的，在下一步就会被加载进来的参数覆盖。 当有加载的预训练参数时，模型的实例化应该被优化成不做初始化，直接把权重加载并放置到相应的内存/显存中。
2. 若是打算从头训练模型，不走 pretrained_ckpt 分支从头训练，我们也可以发现，参数初始化的计算是在 cpu 上执行，最后再移动到 gpu 上，而因为参数初始化大部分时候都是逐元素独立同分布的抽样，直接就是逐元素可并行的操作，直觉告诉我们，在 gpu 上做采样计算，是能得到很好的加速的。

这两点如果不改进，按照之前的分析，主要带来的问题是在 cpu 上做初始化采样计算的延迟，而这会在模型参数量较大时，影响更大。

先给出改进点2的简单示例，以之前 $2^{14} \times 2^{14}$ 的线性层的初始化拆分代码为例，我们在 import 后所有代码执行之前，加入 torch.set_default_device("cuda")，在我本机所用的笔记本 3060 上，运行结果就变成了

memory allocation time: 0.08356356620788574
random sampling init time: 0.0012929439544677734

可以对比看到，初始化耗时的开销有明显地降低。

再来看看第一点改进。 按照之前对模型初始化的拆分分析，我们只需要把“随机采样做参数初始化”这一步骤给跳过，其实就能得到很好的加速。 也就是说可以把所有的 nn.init 里的函数的调用都给省略掉，比如取消调用，或者把调用函数动态更换成空函数。 后者就是 huggingface transformers 常用的 from_pretrained 里的实现，具体可见 调用点 和 相关函数实现。

当然 pytorch 还有更进一步的优化方式，会把初始申请内存/显存也优化掉。参考 pytorch tutorials module load state dict tips

if args.pretrained_ckpt:
    with torch.device('meta'):
        model = Model()
    model.load_state_dict(torch.load(args.pretrained_ckpt, map_location=device), assign=True)

以上的代码中，with torch.device('meta') 达成了只做 shape 等元信息写入的初始化，torch.load 的 map_location 参数完成了直接加载模型预训练权重到对应的设备上，model.load_state_dict 的 assign=True 参数完成了将模型的参数权重重新指向这些预训练权重中。 它们配合起来完成了预训练模型的无冗余初始化和内存申请的加载。

随机变量的数学定义¶

回忆一下在数学课上学的，关于概率空间和随机变量的数学定义。 本文中，我们就先摘抄一下这些定义，然后给出一些自己的感性理解，并切入到与本文主题关联的话题。

概率空间：

概率空间是一个三元组 $(\Omega, \mathcal{A}, P)$，各项目的定义如下：

• Ω 被称为样本空间，其中的任意元素 $\omega \in \Omega$ 被称为样本。

• $\mathcal{A}$ 被称为事件集合，它需要满足的要求有：

  1. $\mathcal{A}$ 中的任意元素 $A$ 需要是样本空间 Ω 的子集
  2. $\mathcal{A}$ 是一个 σ 代数，也即除了 $\mathcal{A}$ 中的元素需要关于集合间的运算有如下性质：
     1. $\Omega \in \mathcal{A}$
     2. 如果 $A \in \mathcal{A}$ ,那么 $\bar{A} \in \mathcal{A}$
     3. 对 $A_i \in \mathcal{A}, i = 1,2, \dots$，那么 $\cup_{i=1}^\infty A_i \in \mathcal{A}$

• $P$ 称为概率或概率测度，它是一个 $\mathcal{A} \to {\mathbb{R}}$ 的函数，且需要满足如下性质：

  1. $P(\Omega) = 1$
  2. 若事件列 $A_i \in \mathcal{A} \quad i=1,2\dots\infty$ 两两交集为空，那么 $P(\cup_{i=1}^\infty A_i) = \sum_{i = 1}^\infty P(A_i)$

随机变量:

随机变量是一个 $\Omega \to \mathbb{R}$ 的函数，并且这个函数还需要满足一个约束：

$\forall r \in \mathbb{R}$， $\left\{ \omega \in \Omega | X(\omega) \in (-\infty, r] \right\} \in \mathcal{A}$

即任意形如 $(-\infty, r]$ 的集合关于函数 $X$ 的逆像需要是一个属于概率空间 $(\Omega, \mathcal{A}, P)$ 中的 sigma 代数 $\mathcal{A}$ 可以计算概率 $P$ 的集合。

对于一些当初学完概率论后淡忘掉了这部分内容的同学来说，这段摘抄可能就像天外来物，有点找不到它想说什么。 我会在这里放一个，我觉得做得还行的概率空间和随机变量的讲解举例视频。 那同时呢，也给出我认为对这两个概念有帮助的感性理解。

样本空间 Ω 中的元素 ω，学名称为样本。 但某种意义上，也可以把它看成是“世界线”，比如如果你关心的是马上丢出去的硬币是正是反，那么可能有：

• 一个世界线 $\omega_1$：这条世界线上硬币攥在手心里时是正面，丢出去时转速中等，中间状态如何如何，最后落下来是正面；
• 一个世界线 $\omega_2$：这条世界线上硬币攥在手心里时是正面，丢出去时转速中等，奇怪的蝴蝶在这段时间内扇了扇翅膀，最后落下来是反面；
• 一个世界线 $\omega_3$：这条世界线上硬币攥在手心里时是反，丢出去的时候丢呲了压根没转，垂直起落，最后得到反面；
• ...

总而言之，元素 ω 无奇不可包，你在意的不在意的现象都可以抽象成一个“世界线元素符号” ω 。 但另一方面，因为我们还是希望最后研究的是一个数学对象，可量化，所以我们就会设置各式各样的观察仪器，通过观察仪器观察世界，得到某些数值。 那数学一点来说，就是建立从 $X: \omega \to \mathbb{R}$ 的映射关系。 还是以丢硬币实验为例子，遵从我们最后只在乎观测硬币落地的情况，我们就可以把 $\omega \to \mathbb{R}$ 的映射关系也即函数 $X$ 取成把最后落地为正面的 ω 映射为 1 落地为反面的情况预测为 0 这样的映射。 那因为世界线 ω 的无奇不包和无法预测，我们在观测端看到的数字 $0, 1$ 也就看起来无法预测了。

但这样的描述只涉及到概率空间和随机变量定义里的第一条。 其他的定义条目是来干嘛的呢？ 是用来限制什么东西是可以被合法地定义概率值以及它的值是多少的。这里不多展开，感兴趣的读者可以查阅一下诸如“概率的公理化定义的动机”，“如何理解概率空间三要素”等问题，或者Bertrand 悖论 之类的东西。

这里想强调的事是，数学定义上的随机变量是一个函数，从抽象的样本空间 Ω 映射到实数空间 $\mathbb{R}$。 从这个角度看，上两节提到的计算机随机数的发生方式，和之前提到的各种随机分布的抽样方式，就是在一比一地翻译“随机变量一个是 $\Omega \to \mathbb{R}$ 的映射”，是“从世界线到数的映射”。 我们还是按照例子一个个来说。

linux 的真随机数发生机制是，把磁盘访问，鼠标键盘设备等实时硬件观测相应做序列化，得到比特流。 这里的 Ω 就是可以是所有可能的硬件响应， ω 就可以是要做真随机发生时的硬件观测（当然，中二一点也可以是用户在键盘敲下了某些值的世界线）。 那把一个样本做序列化，就是在把 ω 转成数/数列。

正文里讨论的各种随机数采样方法也是在一比一地翻译“随机变量一个是 $\Omega \to \mathbb{R}$ 的映射”，是“从世界线到数的映射”。 但这里有两种理解，我们一个个说

1. 我们直接建立的是实数到实数的映射 $g: \mathbb{R} \to \mathbb{R}$，但因为 $g$ 的输入是来自于 $U(0, 1)$的，计算机在抽出 $U(0, 1)$ 的时候会做硬件观测到 $[0,1]$ 上的浮点数的一系列转换，即做了 $\omega \to \mathbb{R}$，所以两步映射复合起来，$\omega \to \mathbb{R} \to \mathbb{R}$，最终达成定义了一个随机变量，如果 $g$ 取取好，这个随机变量就能服从我们想要的分布了。
2. 你也可以将 $U(0, 1)$ 上的观测值，作为 ω 看待。某个世界线上，torch.rand 给我返回的是 0.8，另外某个世界线上它返回的是 0.6，通过构造不同的 $g$，做 $\omega \to \mathbb{R}$ 的映射，就能得到不同的随机变量。同样如果 $g$ 取取好，这个随机变量就能服从我们想要的分布了。

而所谓的计算机模拟随机变量的抽样，很自然地就是实现这个函数，备好相应的输入，然后调用它，收到返回值。 没有什么弯弯绕绕，只是数学定义的一比一翻译。