AI编程转录稿：20251106编译器前端+优化

幻灯片#

11 Compiler Frontend
12 Optimization (01:34:57)

编译器前端#

AI 编译器的两种路径与前端优化#

我们上次课开始讲 AI 编译器，尤其是 AI 编译器前端的内容。回顾一下，一个 AI 问题主要有两种实现路径。一种是 Runtime 路径，即构建动态图，通过动态图做自动微分，并调动算子库进行计算。算子库可在此基础上做计算优化。这条路径适合研发阶段的算法和用户调试。另一条更高效的路径，则适用于成熟的框架，如大模型。我们希望通过编译器功能，将前端代码转换为灵活的计算图，特别是将常用部分固化为静态子图。静态子图的优势在于可以利用编译器技术做图优化、算子优化，并面向硬件后端做专门的优化。这第二条路很大部分工作要由 AI 编译器完成。我们上节课主要讲了前端方面，即跟计算图优化相关的操作。

图优化概览与常量折叠#

图优化的输入对象是我们构建的计算图。我们基于这个对象，可以一遍遍地去做各种优化工作。我们也推荐了一个优化的顺序，虽然可能没有定论，但推荐的顺序通常是比较好的选择。

我们上次课讲了常量折叠（Constant Folding），它最大的特点是和过去的传统编译器比较接近，是过去的技术在 AI 框架中的延续应用，主要涉及常量传播和常量折叠。这里稍微提一下，在后续的图优化中，我们会经常举 Convolution（卷积）和 Batch Normalization（批归一化）一起的操作为例。当神经网络训练完毕，很多参数，如卷积的权重或 Batch Normalization 的均值和方差，都变成了固定的常数。此时，常量折叠能很高效地把先做 Convolution 再做 Batch Normalization 的两步，直接融合成一个 Convolution 操作，节省大量运算。

算子融合：核心思想与目标#

接下来我们看更多的图优化方法，它们可能适用于神经网络训练等场合。我们很感兴趣的一个部分，也是图优化中研究最多、最起效果的阶段，就是算子融合。

算子融合最核心的思想主要为了解决两类问题。第一类是解决“内存墙”问题。当我们发现计算对于缓存的消耗比较大，需要大量数据传输，且存在生产者和消费者算子间的依赖关系时，我们希望通过算子融合打破内存墙，提升访存的局部性，减少内存读取和数据传输带来的效率问题。

第二类是解决“并行墙”问题。当我们发现朴素的计算设计与当前硬件的并行度不匹配，比如很多算子空闲，只有一两个算子在工作，这就涉及到了并行墙。我们也可以通过算子融合或重编排，使得整体并行度提高。后续的算子融合基本都是为了解决这两个问题。

基于规则的算子融合#

算子融合的手段是多样化的。我们首先介绍一些基于经验总结的、基于规则的算子融合手段，许多工作来源于 XLA。XLA 最初服务于 TensorFlow，主要做静态图，因此在算子融合上有较多积累。

一个简单的例子是竖向融合：如果后一个算子（如 D）完全只依赖于前一个算子（如 C），算力允许的情况下，就可以把这两个算子打包在一起执行。这有助于提高缓存效率，可以认为是一个解决内存墙的用例。类似地，如果遇到一个生产者（A）和两个消费者（B、C）的情况，通常要 A 执行完再执行 B 和 C。一个加速策略是把 A 复制一次，在两个单独的核上各自计算 A，再分别传输给 B 和 C。这种融合方式的前提是，A 是一个计算很快、但生成数据量庞大的操作。与其拷贝庞大的数据，不如重新计算一次 A，这样可能更高效。

此外，两个较小的计算（A、B）也可以融合，再将结果给到 C。好处同样是 A 产生的数据可以减少一次保存，只提供一次就能完成 B 和 C 的计算。这些都属于打破内存墙的高效操作。

我们再举一个更具体的例子。一个网络结构可能有两个分支，一个分支执行两次 Convolution 操作，另一个分支执行另一次 Convolution 操作，最后叠加或归并在一起。通常 Convolution 是计算开销大的工作。有趣的是，这两个分支的 Convolution 可能需要相同的输入数据，且数据维度很高，但它们本身因形状不匹配而很难融合。这里的一个人为设计是，把一个 $1 \times 1$ 的卷积算子扩展成一个 $3 \times 3$ 的卷积算子，因为 $3 \times 3$ 的卷积有能力完成 $1 \times 1$ 的工作。这样做最大的好处是，顶上的两个卷积算子可以被融合成一个卷积计算操作，避免了输入拷贝到两个计算 kernel。为了保证对下游数据有效，融合后补了一个 Split 操作，再执行剩余的 Convolution。对于底下的 Reduction（规约）这类算子，它们很适合与父节点融合。最终，这个图被简洁到只有两次算子调用。

训练中的卷积与 BN 融合策略#

前面讲的是基于规则的融合，我们也关心人为设计的融合思路，特别是之前提过的卷积和 Batch Normalization。在神经网络前向计算时，会产生非常多的中间量，而在反向计算时，很多中间量需要被复用。例如乘法的前向运算，其反向微分一定会复用原来的值。因此，对于训练中的神经网络，如果能共享前向和反向的中间数据，计算效率会有很大提升。

我们之前讲过可以同时做前向图和反向图的联合优化，这里也体现了这一思想。在 Convolution 和 Batch Normalization 这类涉及缓存较多的计算中，尤其是在训练时，常量折叠不再适用。因为此时，无论是输入 $x$ ，还是预测的均值和方差，都不是常量。

接下来介绍的是一个目前 AI 编译器很难自动设计出来的人为优化策略。它不是把 Batch Normalization 吸收到 Convolution 里，而是把 Batch Normalization 操作进行了拆分。Batch Normalization 主要分两部分工作：一是统计（BN_a），关心当前的均值和方差；二是归一化（BN_b），按期望的均值和方差做映射。

拆分后，BN_a 做数据统计，它依赖于前面 Convolution 的输出，存在内存墙。而 BN_b 通常是简单计算，并会生成高维数据，非常适合被吸收到后续的 Convolution 计算中。通过这种人为设计，我们就可以把某一个 Convolution、它下面的 BN_a 以及它前面的 BN_b 融合成一个高效的计算。

这种算子融合很难由 AI 编译器自动完成，因为它打破了算子的界限，进行了跨算子间的融合。如果 AI 编译器的计算图只表达算子节点，没有算子的中间表达，就无法完成这种拆分。这也正是未来 AI 编译器希望打破图优化和算子优化壁垒，实现跨算子融合，做到类似人工设计的高效优化的原因。

此外，前向优化也希望能支撑反向的复用。我们会开辟一个专门的 buffer，存储后续反向计算会用到的中间值，如权重 $W$ 、特征 $x$ 、均值等。对于某些值（如 $Y$ ），原地计算可能比复制拷贝更省事，优化策略可能就是只传递 $Y$ ，反向计算时再原地算一下。

基于图结构的融合范围界定#

前面介绍的都是基于规则的融合策略。人们会担心，这种自动优化是否会改变计算逻辑，计算结果是否安全。TVM 提出了一个方法，用于界定算子融合的最大安全范围。

在一个计算图中，有些节点只有一个父节点，它更容易与父节点等进行融合。但如果一个节点（如 add）的输入来源于两个分支，融合时就需要更小心。为了界定此事，TVM 提出分析计算图中每个节点的“最低公共祖先”（LCA, Lowest Common Ancestor）。通过构建“支配树”（Dominator Tree），我们可以更容易地找到每个节点的 LCA（即其在支配树中的父节点）。

一个节点到它的 LCA 之间的这个范畴，被认为是可以自由做算子融合的范围。例如，底下的 add 节点到它的 LCA（四号 reduction 节点）之间，可以自由融合；但 add 不能轻易和它 LCA 之上的左分支或右分支融合。在这个安全范畴内具体如何融合，则可以参考我们之前讲过的各种融合规则，如解决内存墙或提高并行度的操作。

TVM 也提出了一些基于规则的融合，并做了一些抽象。例如，把一对一映射的操作定义为 Injective（如 ReLU、add），把 conv2d 等计算量大的定义为 Complex-out-fusable，还有 Reduction（规约）和 Opaque（无法融合，如 sort）。基于这些抽象，TVM 能在抽象层面设计更通用的规则，实现算子融合策略。

图优化之布局转换#

除了算子融合，我们也可以提前考虑计算内容的内存布局。在图优化的检查环节（如 Type Check），我们可以确定数据的形状。在算子融合后，我们就可以根据硬件情况做一个内存对齐。这是传统编译器优化技术，例如 32 位系统按 4 字节读取内存，如果数据未对齐，可能导致一次读取跨了两次内存访问，效率低下。

更重要的是布局（Layout）的调整。面向视觉任务时，数据格式有 NCHW（PyTorch 默认）和 NHWC 等。NCHW 将 Channel 放在前面，适合沿着 $H$ 和 $W$ 维度进行规约的操作，在 CPU 上内存连续性好。但对于 Convolution 等需要所有 Channel 参与的计算，尤其是在 GPU 上，NHWC 布局更优，因为它提供了更好的内存连续性。

PyTorch 因为历史原因默认 NCHW，但做一些布局转换，尤其是在 GPU 上，往往能提高性能。PyTorch 也意识到了这点，提供了 channels_last（即 NHWC）的支持。有趣的是，转换后，.shape 属性可能不变，但 .strides（步幅）会改变，这说明它在内存中确实被转换了。图优化时，就可以根据算子需求（是需要 Channel 连续访问还是对 $H$ 、 $W$ 规约）和硬件（CPU/GPU），在计算图中成对地插入布局转换操作。

图优化之内存分配#

计算图优化也可以提前关心内存分配。这里讲的不是真的硬件分配，而是进行共享内存的逻辑设计。这在 AI 框架中非常重要，因为计算图的计算很复杂，除了必须存储的静态内存（如参数），中间计算会占用非常多临时的动态内存。

动态内存的消耗占比非常大，甚至远大于静态内存。以面向移动端的 MobileNet 为例，如果不做任何内存管理优化，按顺序开辟内存，整个网络需要超过 20MB 的存储。但如果做了合理的内存复用，压缩空间非常大。

内存复用策略与并行性#

最简单的内存复用是“原地复用”（In-place）。如果一个节点计算完后，只服务于下一个计算，没有其他依赖，那么下一个运算就可以直接原地复用（覆盖）这块内存。但如果该数据（如 B）后续还被其他计算（如 F）需要，那么 C 在计算时就不能原地复用 B，需要开辟新内存。

对于非 element-wise 的操作，如 Pooling，它依赖邻居关系，不能原地复用上一个算子（B）的内存，需要为 C 开辟新空间。但是，当 C 算完后，B 的内存可能不再需要，后续的节点（如 E）就可以复用 B 的内存空间。

内存复用显然跟图的数据依赖性相关。一个算法实现是，为每个数据节点维护一个标识，即它的出度（被多少个后续节点依赖）。计算开始时，B 需要开辟新空间（如红色）。当 C 开始计算时，B 的出度减 1。如果此时 F 也可以并行计算（只依赖 B），F 也开辟空间，B 的出度再次减 1。当 B 的出度变为 0 时，意味着 B 的内存可以被回收，放回“内存池”。后续的节点（如 E）就可以从内存池中获取内存，而不是开辟新内存。

这个方法比较适合串行设备，它通过贪心策略最大化节省了内存。但代价是，这块内存先后存了 B、E、F、G，导致这些节点间产生了新的依赖关系，即 B 必须被用完，E 和 F 才能用。这可能会影响计算的并行度。

内存分配与并行性#

不合适的内存复用可能导致原本可以并行的左右两个分支，变成必须串行执行。为了最大化利用并行性，一种策略是：先找到计算图中的“最长路径”（Longest Path），在这条路径上做内存复用。因为最长路径上的节点本身就是数据依赖、必须串行的，内存复用不会引入额外的并行瓶颈。优化完最长路径后，再依次优化次长路径，依此类推。这样，内存分配就不会影响整体的并行策略。

总的来说，内存优化在 Inference（预测）模式下应用最广，效果也最好（如减少 4 倍内存使用）。在训练时，由于反向传播需要依赖前向的计算结果，内存复用效率会低一些（如减少 2 倍内存）。但即便如此，内存复用，尤其是对移动端而言，仍是至关重要的优化手段。

AI 编译优化：代数表达式化简#

之前讲过了算子融合、布局排布、内存分配等 AI 编译优化手段。接下来要讲的是代数表达式化简。这类优化很有研究价值，主要是一些专家根据现存算子的情况，提出经典的算子等价化简策略。如果对这个方向感兴趣，推荐阅读 DNN fusion 这篇论文，它总结了许多计算相关的经典化简例子。

在这些研究中，我们关注过去表达式的样子、涉及的运算（如矩阵乘法）及其 flops。接着，我们进行等价的表达式转换，并计算转换带来的维度变化。如果真实使用场景符合这种等价转换，就可以评估开销节省是否可观，并选取适合的经典转换例子。这在实际操作中非常适合 AI 编译器实现，本质上就是查表。根据当前表达式对应的子图，将其置换为一个计算量更小的子图。AI 编译器主要做的是图的匹配工作。

优化策略：规约、广播与代码消除#

举个例子，常见的是规约类操作。我们发现，应尽可能把规约操作提前。例如，如果规约操作内部是一些 element-wise 的操作，就可以把规约提前。这样能提前减少数据维度，再执行 element-wise 操作，结果是等价的。

此外还有更复杂的张量计算，经常涉及不同大小的张量通过广播（broadcast）来进行操作。这就存在如何广播的问题，尤其对可交换的算子。不佳的广播顺序可能导致数据被无效扩大，后续需要更多规约。AI 编译器适合根据静态分析，找到广播次数最少的最佳运算逻辑。我们希望用户只需考虑广播的逻辑，而非计算效率。

除了表达式简化，最后常见的操作是死代码消除。表达式简化和死代码消除这两类优化，在传统编译器中也存在。它们之所以出现，往往是因为前面的优化（如融合）导致了后续出现不优雅的表达式或未被使用的分支。因此，这两个优化策略通常放在优化流程的较后阶段。

优化策略：公共子表达式消除#

接下来讲公共子表达式消除。这与传统编译器的逻辑类似。过去可能是查找公共字符串进行替换，而对于 AI 编译器，则是在计算图中寻找计算量大且一致的部分，进行融合或替换。

子表达式消除的逻辑与传统编译器类似。不同点在于，AI 编译器的输入是计算图，需要先构建子表达式。因此，最开始要做逆后续节点集的创建工作，以构建计算图对应的表达式。

有了表达式后，我们开辟空间建立子表达式候选集。遍历图中节点，将表达式记录到候选集中。判断是否重复出现时，主要采用哈希表结构，通过计算表达式相关信息的哈希值来匹配近似操作。匹配到后，还需详细检查，确保输入输出类型相同，才能复用。复用可能涉及删除某个分支，并将另一分支复制过来。

对于不影响输出的分支，我们可以安全地进行死代码消除。这与传统编译器类似，但额外的操作是，需要先从计算图构建逆后续节点集以获得表达式，然后再用传统编译逻辑删除死节点。

编译器优化顺序与 PyTorch Compile 案例#

回顾一下，计算图优化的常见顺序（Pass）是：首先是算子融合，接着做排布空间优化，然后根据内存消耗做内存布局优化，最后是公共子表达式消除和死代码消除。这个顺序是比较合理的设计。

举个具体例子，torch.compile 是 PyTorch 2.0 的一个主要贡献，包含了很多操作。它能做动态图构建，并通过 AOT (Ahead-of-Time) AutoGrad 实现前向图和反向图的联合构建与优化。PyTorch 在此阶段会选择基于图的 Graph IR 作为中间表达。

完成图构建后，我们希望在算子阶段做工作，这就需要一种图算融合的 Mixed IR。在 PyTorch 中，这种 Mixed IR 就是 PrimIR。我们可以把算子的具体实现套到计算图中，获得图算融合的中间表达。在这个级别上可以做面向硬件后端的优化。在算子层面，我们提过 OpenAI 设计的 Triton，这是一种类似 CUDA 的语言中间表达，用于实现面向 GPU 的后端优化。

PyTorch 后端支持 Triton 后，并没有直接绑定到 CUDA 表达上，这给了 PyTorch 支持非英伟达（如 AMD）显卡的机会。目前 AMD 也在拥抱 Triton，使得 PyTorch 能在支持 CUDA 的同时，应用于 AMD 显卡。此外，TVM 和 XLA 等开源编译器也会积极拥抱 PyTorch，在相应阶段将其 IR 转换为更适合硬件的表达。TVM 的优势是支持更广泛的 GPU（如移动端），而 XLA 更适合 TPU 这类专门的张量计算硬件。整体而言，AI 编译器正呈现出与显卡制造商解耦、互相兼容的开发趋势。

PyTorch Compile 的使用与后端#

我们介绍了 PyTorch 从前到后的编译逻辑。调用 torch.compile 时，这些优化就已打包。如果不使用优化，生成的多个子图会分别调度到 GPU 执行，效率很低。而调度了 PyTorch 默认的 Inductor 优化器，就会执行算子融合，一个简单的程序很可能融合成一个算子。

之前提到 PyTorch 何时决定生成新子图，这主要是一种 “Guard” 机制。它会记录当前生成代码的参数（如 shape）和返回值情况。如果下次调用时输入 shape 发生变化，就会识别为不匹配，从而激发重新构建和优化子图，以匹配新的输入。

torch.compile 默认使用自己的后端加速器 Inductor。PyTorch 在编译方面起步晚于 TensorFlow，且其面向前端的需求灵活，导致后端优化难度大。PyTorch 自身不适合做泛化优化，因此它积极与其他编译器（如 TVM、XLA）合作，这些都可以作为 backend 被调用。

我们的 backend 优化有很多选择，分为两类。第一类是训练时优化，这需要支持 PyTorch 的 AOT AutoGrad（前后向连续构图）。主要 backend 包括通用的 Inductor；如果确定后端是 CUDA，也可以用 CUDA 相关的优化；后端是 CPU 也可以指定。另一类是 inference mode（推理模式），即网络训练完毕。此时优化空间最大（如常量折叠），后端复杂性也高，比如部署到移动端。这时支持的后端有 TVM 和 XLA。PyTorch 目前主要在 inference mode 下支持这些后端优化。

PyTorch Compile 的高级用法#

最终希望大家能在实际中使用 torch.compile。例如，从 Hugging Face 下载的预训练（pretrained）模型，可以直接扔到 torch.compile 里，多数情况下是可行的。唯一的缺点是编译可能需要等很久，但之后就能高效执行。

如前所述，在大模型上调用 torch.compile 会等很久。如果想缩短时间或提前预览，有一些参数可选，比如限制优化时间提前结束。除非你时间充裕且追求最优状态，才选择如 max-autotune 这样的模式。PyTorch 提供这些灵活性，也是因为 torch.compile 运行久的问题常被诟病。

刚才提到 PyTorch 可以接其他编译器做后端，比如 XLA。这些库现在都支持 Python 架构。你可以在 Python 中同时使用 PyTorch 和 PyTorch XLA 版本，并在 torch.compile 时调用 XLA 相关代码。对于 XLA 或 TVM 这类通用编译器，PyTorch 目前普遍支持它们在 inference mode 下使用，此时加速比较理想。对于 training mode，XLA 也在做一些实验性尝试，也有提速，但总的来说 PyTorch 更常在 inference mode 下使用它们。

图的派发与并行#

接下来讲计算图的派发。当我们对后端硬件并行有考量时，就需要做计算图的工作分配。AI 编译器阶段并不执行，但提前安排好派发方案（如子图分割），对部署执行大有裨益。AI 编译器前端需要处理子图分割和构建，真正的部署则在后端硬件执行。

图的部署主要考虑并行性。按顺序执行很耗时。为了更好调度，AI 编译器可以根据用户提供的并发信息（如有多少核可用），依据并行度，来排布一个并发执行队列。它会把没有依赖的节点调入并发队列。然后循环遍历图，将后续节点按并发策略塞入队列。这是一种比较简单、贪心的提高并行性的策略。

除了提高并行性，图调度还涉及分布式。如果在不同设备间并行，就产生子图分割问题。确定分割后，需在子图间补充数据发送和接收节点，使逻辑能在分布式环境中完整执行。最难的是设计图分割策略，以最小化数据传输消耗。这会用到一些基于图的 graph cut 优化算法。

图部署时，我们通常有 CPU 和 GPU。常见做法是把数据相关操作（如读取、裁剪、随机、augmentation）布局到 CPU，把并行算力留给训练相关的计算。朴素的逻辑是顺序执行：CPU 读数据、处理数据，然后 GPU 运行。这会导致 CPU 和 GPU 忙闲不均。考虑到数据处理在 CPU 上，且不同 batch 间依赖性不强，常见的优化策略是异步执行。我们可以采用流水线作业，当上一个 batch 的数据在 GPU 训练时，下一个 batch 可以在 CPU 进行数据读取和处理，实现异步执行。

这些是 CPU 和 GPU 的异步策略，多线程并发也会纳入计算图优化。AI 前端优化完成后，我们得到了一个简便的计算图，并根据硬件情况做了图拆分、内存分配、异步执行调度。到了 AI 编译器后端，主要任务是将算子加入图表达，并做算子级别的优化，尤其是面向特定硬件的专属优化。这便是 AI 编译器前端在整个 AI 框架中要完成的工作。

优化器#

优化器概述：梯度下降#

接下来讲优化（Optimization）。这是 AI 框架很关心的部分。自动微分是为了计算梯度，拿到梯度后如何使用它，就是优化环节关心的事。最基础的是梯度下降，其数学表示为 $\theta_{t+1} = \theta_t - a \nabla_{\theta} f(\theta_t)$ 。优势是梯度计算方便；劣势是步长 $a$ （learning rate）是一个需要人为设定的超参数，其大小的选择非常敏感，由此产生了很多 learning rate 管理策略。

步长选取很敏感。如果步长选得比较小，收敛会很缓慢，但曲线比较均匀光滑，朝着局部最小值（local minimal）逐渐下降，不足是优化时间长。如果不幸选择了较大的步长，在能量 “landscape” 上，可能因走得过长而到达一个与当前梯度不符的点，产生震荡走势。训练曲线上也会看到频繁震荡，这提示我们应减小 learning rate。步长较大时，前期收敛快，但后期震荡，不易收敛到好的点。

牛顿方法#

梯度下降大家比较了解。我们再讲讲传统牛顿方法，这是一个二阶方法，常用在物理仿真等数值计算中。其公式为 $\theta_{t+1} = \theta_t - a \mathrm{H}_{\theta}(f)^{-1} \nabla_{\theta} f(\theta_t)$ ，最大差别在于它对步长的选择。梯度下降只用一阶梯度，牛顿优化用二阶的 Hessian 矩阵（ $\mathrm{H}$ ）来计算最适合的步长，即 $\mathrm{H}^{-1}$ （Hessian 的逆）。 $\mathrm{H}^{-1}$ 乘以梯度就是参数优化的距离。

牛顿优化被称为二阶方法，核心是基于泰勒展开。如果我们把 loss function 做泰勒展开，牛顿优化是假设函数是二次的，然后直接优化到最优点。如果函数不是二次的（含有三阶或更高阶项），牛顿优化也是用二次函数去做拟合，每次迭代优化掉二阶部分。牛顿优化可能有 “overshooting”（越过极值点）的风险，因此我们增加参数 $\alpha$ （damped Newton method）。但对理想的凸优化问题，全牛顿步长（ $\alpha=1$ ）直接使用 Hessian 的逆作为步长选择。

动态图展示了牛顿优化的思路。我们关心的是梯度与 X 轴的焦点（梯度为 0 处为最优解）。牛顿优化假设梯度函数是线性的（即原函数是二次的），然后直接找到这条线性函数与 X 轴的截距，把优化点更新到该截距上，然后循环迭代。牛顿优化最适用于二次函数；如果不是二次但函数是凸的（convex），它通常也比梯度下降效率高。

但在神经网络中却不用牛顿优化。核心问题是，神经网络的参数 $\theta$ 维度 $n$ 非常高，导致 Hessian 矩阵 $\mathrm{H}$ 变得巨大（ $O(n^2)$ ）。计算一阶梯度已经很耗时，计算二阶微分（Hessian）并求逆（ $O(n^3)$ ），这个计算量对于神经网络的参数维度是无法容忍的。尽管牛顿优化不适用，但它的二阶收敛性是我们想要的。在二次函数例子中，若 $\alpha=1$ ，牛顿法一步到位。这种高效收敛性远好于梯度下降的线性收敛性。

拟牛顿方法：BFGS 与 L-BFGS#

如果我们想要牛顿法的高效收敛，又不想要 Hessian 矩阵的复杂计算，就催生了拟牛顿（Quasi-Newton）优化器。核心诉求是克服计算 Hessian 矩阵的复杂度。想法是通过拟合 Hessian 矩阵，比如用一阶微分去拟合二阶的 Hessian。

对所有拟牛顿方法，任务是拟合 Hessian。我们不用 $\mathrm{H}$ ，而是用 $B$ 矩阵来近似。我们希望每次迭代使用一阶梯度信息来更新 $B$ （ $B_{t+1} = B_t + \Delta B_t$ ），使其更接近真实的 Hessian 矩阵。核心原则是利用数学上的割线定理（Secant Condition），即 $B_{t+1}(\mathbf{x}_{t+1} - \mathbf{x}_t) = \nabla f(\mathbf{x}_{t+1}) - \nabla f(\mathbf{x}_t)$ 。算法的核心逻辑是 $\theta_{t+1} = \theta_t - a B_{t+1}^{-1} \nabla f(\theta_t)$ 。它避免了计算 Hessian，但仍涉及大 $B$ 矩阵的求逆。

拟牛顿算法中经典的是 BFGS 算法（以四位数学家命名）。它很经典，因为它解决了几个问题：第一，用 $B$ 矩阵拟合 Hessian；第二，它保证了 $B$ 矩阵是对称正定的，这使得后续求逆时数值计算更稳定高效。BFGS 的核心是 $B_{t+1}$ 的更新公式： $B_{t+1} = B_t + \frac{\mathbf{y}_t \mathbf{y}_t^T}{\mathbf{y}_t^T \mathbf{s}_t} - \frac{B_t \mathbf{s}_t \mathbf{s}_t^T B_t^T}{\mathbf{s}_t^T B_t \mathbf{s}_t}$ 。其中 $\mathbf{s}_t = \mathbf{x}_{t+1} - \mathbf{x}_t$ （自变量变化）， $\mathbf{y}_t = \nabla f(\mathbf{x}_{t+1}) - \nabla f(\mathbf{x}_t)$ （梯度差值）。它还设计了 $B^{-1}$ 的更新式 $B_{t+1}^{-1} = (I - \rho_t \mathbf{s}_t \mathbf{y}_t^T) B_t^{-1} (I - \rho_t \mathbf{s}_t \mathbf{y}_t^T) + \rho_t \mathbf{s}_t \mathbf{s}_t^T$ ，其中 $\rho_t = \frac{1}{\mathbf{y}_t^T \mathbf{s}_t}$ ，这利于直接更新 $B$ 矩阵的逆。

BFGS 策略虽好，但 $B$ 矩阵仍是 $O(n^2)$ 的二阶张量，当神经网络维度 $n$ 很大时，计算和存储依然无法接受。因此，后续提出了更节省内存的 L-BFGS（Limited-memory BFGS）算法。L-BFGS 的核心观察是：我们最终需要的不是 $B^{-1}$ 矩阵本身，而是 $\mathbf{p}_k = -B_k^{-1} \nabla f(\mathbf{x}_k)$ 这个向量。L-BFGS 的想法是不存储 $B_t^{-1}$ 矩阵，而是利用其递推公式，只存储最近 $m$ 步的 $\mathbf{s}_t$ 和 $\mathbf{y}_t$ 向量（内存开销 $O(mn)$ ，通常 $m$ 约等于 10），用它们来直接近似计算 $\mathbf{p}_k$ 。

关键点和注意事项#

常量折叠适用性：Convolution 和 Batch Normalization 的常量折叠融合非常适合推理（Inference）模式，因为此时均值、方差等参数已固定为常量。
训练时融合策略：在训练（Training）模式下，Batch Normalization 的参数是变量，不能使用常量折叠。此时采用的策略是人为地将 Batch Normalization 拆分为统计（BN_a）和归一化（BN_b）两部分，进行跨算子的融合。
跨算子融合的难度：上述训练时的融合策略打破了算子的原子性，属于跨算子融合，目前 AI 编译器很难自动完成，需要人工设计。
算子融合的双重目标：算子融合的核心目标是解决两大问题：一是“内存墙”，通过提升访存局部性减少数据传输；二是“并行墙”，通过提升硬件并行度来加速计算。
融合的安全性：为保证融合后计算结果的正确性，TVM 等编译器使用“最低公共祖先”（LCA）和“支配树”（Dominator Tree）来确定一个节点可以进行自由融合的安全范围（即该节点到其 LCA 的子图）。
布局选择：数据布局（如 NCHW 或 NHWC）的选择至关重要。NCHW 适合 CPU 或规约操作，而 NHWC（channels_last）更适合 GPU 上的卷积等 Channel 密集型计算。
动态内存的挑战：AI 计算中，临时的动态内存开销往往远大于参数等静态内存。因此，内存分配优化（内存复用）是编译器前端的重要工作。
内存复用与并行性的权衡：贪心的内存复用算法（如串行复用）可能会引入虚假的数据依赖，从而破坏计算图原有的并行性。为并行硬件（如 GPU）优化时，应采用“最长路径优先”等策略，避免在可并行路径间共享内存。
拟牛顿（BFGS, L-BFGS）这类近似二阶优化策略，不适合随机（stochastic）梯度下降，即不适合分 batch 的数据训练。