自动微分：自定义梯度Custom Grad

可视化演示

展示计算图前向反向传播中Custom Grad节点梯度注入、链式法则自定义梯度替换及有限差分梯度检查验证三种核心模式，直观呈现AI框架自定义梯度扩展机制。

需求场景

中央厨房的流水线上，绝大多数算子都带着一份标准菜谱：框架在注册算子时，就把它的求导规则一并写入了自动微分引擎。于是 sin、exp、卷积这些常客，后厨闭着眼都能做出菜来。可总有些“不按常理出牌”的食材——比如从别的厨房借来的酱料（外部库操作）、一遇高温就炸锅的油炸工序（数值极易溢出）、以及按标准流程做要绕一大圈的冗余烹饪（性能瓶颈）。标准菜谱在它们面前哑火了。

PyTorch 这类框架的自动微分，本质上是对计算图中的每个算子预置一份反向函数。它做得再好，也不可能覆盖全世界的所有数学操作。当你在前向中调用了一个 NumPy 函数、一个 C++ 扩展，甚至一段模拟器的物理方程，框架根本看不见它内部的求导结构，默认规则自然无从谈起。更麻烦的是，有些算子虽然能求导，但求导路径上会出现 inf 或 nan，训练直接崩溃。

于是“自定义梯度”登场：它允许你绕开默认规则，亲手给这个特殊算子写一份前向做法和反向反馈。这份能力不是替代自动微分，而是填补它的盲区。

可以说，自定义梯度就是大厨为那些“做不了的菜”单独写的私房菜谱。有了它，厨房的菜单边界一下子被撑开了。

注册机制

私房菜谱光写在纸上没用，得登记到厨房系统里：这道菜叫什么、怎么做、客人不满意时按什么规则反馈，都要一一入库。在 PyTorch 中，这套登记流程就是继承 torch.autograd.Function 并覆写两个静态方法。类定义的一瞬间，菜谱就挂上了号；forward 和 backward 分别是菜谱里的“做法”和“反馈规则”两个栏目。

真正调用时，统一走 .apply() 这个入口。它不是普通的直接函数调用，而是由框架在计算图中登记一个节点，把输入张量、输出张量和上下文 ctx 全部串联起来。框架还会根据输入张量的类型自动分派：CPU 张量走 CPU 实现，CUDA 张量走 GPU 实现，这份“类型分派”让同一道菜在不同灶台上都能做。而 ctx 就是挂在锅边的小纸条，从前向一直传递到反向，贯穿整个生命周期。

至于生命周期，计算图在反向传播结束后会被释放，菜谱本身仍然保留，下次再调用 apply 会重新生成一张新的计算图节点。也就是说，每次做菜都是“现登记、现使用、现销毁”，厨房里不会堆积过期的纸条。

登记完毕的菜谱，才算真正进入了厨房的调度系统。后续所有菜品的传菜、反馈、复盘，都由系统按这份菜谱自动执行。

前向接口

做菜时，大厨会在锅边挂一张小纸条：这道菜用了什么食材、下了多少油、腌制了多久。这纸条不是给客人看的，而是留给自己在反向阶段复盘用的。自定义梯度的前向接口 forward(ctx, *inputs) 正是如此：它负责把输入加工成输出，同时把反向需要的一切关键信息写进 ctx 这张小纸条。

首先，forward 必须是静态方法，第一个参数固定是 ctx，后面跟着任意数量的输入张量。它的返回值就是前向输出，会被自动登记为计算图的下一个节点。其次，如果某个输入在后向阶段还会用到，就必须调用 ctx.save_for_backward(x) 把它存下来；如果某个输入不参与梯度计算（比如固定的超参数），就可以不保存，反向时对应返回 None 即可。

还有一个新手容易踩的坑：不要在 forward 里对输入张量做就地（inplace）修改。因为 save_for_backward 保存的是张量引用，如果前向结束后你把它就地改了，反向时读到的就是“被污染后的菜”，梯度计算随之出错。这就像做完菜后，有人偷偷把锅边的纸条改成了另一份食材清单。

前向做菜只是上半场。真正考验私房菜谱质量的，是反向阶段能否凭这张小纸条，把客人的不满精准翻译成每一份原料的调整建议。

反向接口

客人对一道菜不满意时，前厅会把整桌菜的总反馈单传回厨房。传到这道菜面前时，反馈单上写的是“总损失对这道菜输出的敏感度”——这就是上游梯度。大厨要做的事情很简单：拿着反馈单，对照自己的做法（本地梯度），算出“我该怎么调整当初的原料”，然后把这份新的反馈单继续传给上一道工序。

backward(ctx, grad_output) 就是干这个的。grad_output 是上游梯度 $\partial L / \partial y$；你在函数体内计算本地梯度 $\partial y / \partial x$，两者相乘得到 $\partial L / \partial x$，并以元组形式返回。返回值的数量必须与 forward 的输入数量一一对应；某个输入不需要梯度时，用 None 占位。这套规则，本质上是链式法则的分步实现：每道工序只负责计算自己那一段，剩下的交给整条流水线。

更深一层看，backward 计算的是“雅可比-向量积”。对多维张量来说，完整雅可比矩阵太大，但乘以某个向量（上游梯度）后的结果却很小——自定义梯度正是按这个高效形式写的。这也是为什么它的计算量通常远小于显式构造矩阵。

这样一来，反馈单从后厨一路往前传，每道工序只改动自己那一段的调整建议。自定义梯度，就是这道工序上那位亲手算出本地梯度的大厨。

数值稳定性

有一道菜特别考验火候：指数函数。输入稍微大一点，锅里的油就烧到冒烟——exp(1000) 直接溢出成 inf；反过来输入特别小，又可能下溢成 0。数值不稳定的根源，在于中间结果超出了浮点数的表达范围。大厨的解决办法是：换一根更温和的温度计，把计算搬到对数域去做。

最经典的手法叫 log-sum-exp。它常用于计算 softmax 的分母 $\log \sum e^{x_i}$。直接算，指数会炸；但如果先找出这批输入里的最大值 $m$，把所有输入都减去 $m$ 再取指数，指数部分最多是 $e^0 = 1$，稳稳当当。最后再把减掉的 $m$ 加回来，结果完全等价，数值却安全了。

除了这种数学变换，还有一种暴力但有效的“限压阀”：对输入或梯度做裁剪（clip），防止极端值冲垮管道。比如梯度裁剪把梯度的模长限制在一个阈值内，避免更新步长过大。必要时还会用光滑函数近似不可导点——比如用 softplus 近似 ReLU 在零点的折角，让梯度在拐点处依然平滑可导。

温度计换了个刻度，菜还是那道菜，但再也不会炸锅了。数值稳定性，就是大厨们在一次次“烧糊锅”之后总结出的厨房安全手册。

损失函数定制

一家餐厅卖的不只是一道菜，而是一整套套餐：开胃菜、主菜、甜点，各有各的味道，又要彼此搭配。训练模型时也一样，最终的损失往往是好几份损失拼在一起：主任务损失保证核心目标，辅助损失提供额外指导。自定义梯度在这里派上大用场——它让你能精确控制每一份损失对梯度的贡献。

最常见的做法是加权求和：主损失乘上大权重，辅助损失乘上小权重。权重的意义就是火候，调大调小直接改变模型的学习偏向。更进一步，你可以在自定义 backward 里决定让哪条路径“断流”：比如某些分支不希望被某个损失影响，就在返回梯度时直接返回零。这就是 stop-gradient 技巧的雏形——前向正常传，反向不更新。

对自定义损失函数来说，torch.autograd.Function 给了你一个“控制梯度流动”的水龙头。你可以精确指定哪些变量接收梯度、哪些变量被冻结，这在多任务学习、对比学习里几乎是必需品。没有这层能力，想要精细调控梯度流向，只能靠修改网络结构旁路，笨重又脆弱。

套餐配比得好，客人吃得开心，模型训练也稳定。自定义梯度，就是那位按口味灵活调整菜品配比的总厨。

调试验证

私房菜谱写得再漂亮，也不能直接端上桌——万一反馈规则算错了，整桌菜都会被带偏。所以正式上菜单之前，必须“试菜”。试菜的方法很朴素：用两种不同的方式分别算出梯度，看它们是否吻合。一是你自己的解析梯度，二是用数值近似得到的参考梯度。

数值近似最常用的是中心差分：给输入加一个极小的 $\varepsilon$ 和减一个极小的 $\varepsilon$，分别计算函数值，再用两个函数值的差除以 $2\varepsilon$。当 $\varepsilon$ 足够小时，这个比值就无限逼近真实导数。PyTorch 提供的 gradcheck 工具正是这样做的：它会拿你的 backward 结果与有限差分结果逐项对比，一旦差异超过容差就报警。

调试时常见的问题也都有迹可循：backward 返回的张量数量与 forward 输入数量不一致；ctx 保存的张量在前向之后被就地修改过；backward 里访问了一个从未保存的中间变量。这些错误在 gradcheck 面前几乎无处遁形。养成写自定义梯度后必跑 gradcheck 的习惯，就像大厨正式推出新菜前必先让后厨全员试菜。

试菜合格，私房菜谱才算真正过关。有了这关把守，后厨再也不怕“一道菜毁掉整桌宴席”。

应用实践

私房菜谱一旦成熟，就不再只是个人绝活，而是可以写进集团手册、复制到千家分店的招牌菜。自定义梯度在产业界的应用，正是沿着这条路径展开的。

第一是算子扩展：AI 编译器工程师在实现新算子（比如 FlashAttention 类的高效注意力）时，用自定义梯度把前向和反向一起打包成计算图节点，既保证正确性，又完成算子融合与显存优化。第二是框架桥接：很多科学计算代码跑在 NumPy、SciPy 甚至外部模拟器里，自定义梯度让你可以把这些“没有梯度”的模块包一层壳，把近似梯度或解析梯度手动接进训练循环。第三是性能调优：默认自动微分会保存大量中间张量，自定义梯度可以手动控制保存哪些、丢掉哪些，大幅降低显存占用。

对开源生态而言，每一个优秀的自定义算子都是一份可供社区复用的“私房菜谱”。当它被提交进框架主干，成千上万的开发者都能直接调用。这条路的终极形态，就是你不再被动消费框架能力，而是主动定义框架能力。比如对比学习里常用的 stop-gradient，本质上就是一个前向恒等、反向返回零梯度的自定义算子——简单，却塑造了无数经典算法。

从一项救急的本领，到一套系统的方法论，自定义梯度让每个开发者都能成为计算图世界的“主厨”。而这份私房菜谱，也终将成为整个 AI 生态的共享财富。

AI 应用场景

Python 演示

import torch  class CustomHuber(torch.autograd.Function):    delta = 1.0     @staticmethod    def forward(ctx, x, target):        r = x - target        small = r.abs() <= CustomHuber.delta        loss = torch.where(small, 0.5 * r * r,                           CustomHuber.delta * (r.abs() - 0.5 * CustomHuber.delta))        ctx.save_for_backward(r)        return loss     @staticmethod    def backward(ctx, grad_output):        r, = ctx.saved_tensors        grad_inner = torch.where(r.abs() <= CustomHuber.delta, r,                                 torch.sign(r) * CustomHuber.delta)        return grad_output * grad_inner, None  if __name__ == "__main__":    torch.manual_seed(0)    x = torch.tensor([0.5, 1.2, -0.8], dtype=torch.double, requires_grad=True)    target = torch.zeros_like(x)    loss = CustomHuber.apply(x, target).sum()    loss.backward()    print("x 的梯度:", x.grad.tolist())     from torch.autograd import gradcheck     inputs = (torch.tensor([0.5, 1.2, -0.8], dtype=torch.double, requires_grad=True),              torch.tensor([0.2, -0.4, 0.1], dtype=torch.double))    ok = gradcheck(CustomHuber.apply, inputs, eps=1e-6, atol=1e-4)    print("梯度检查通过:", ok)

1. class CustomHuber(torch.autograd.Function)继承 Function 类完成“菜谱登记”。
2. def forward(ctx, x, target)前向接口：计算 Huber 损失，把残差 r 存入 ctx。
3. ctx.save_for_backward(r)保存中间张量，反向时从 ctx 取回。
4. def backward(ctx, grad_output)反向接口：接收上游梯度，计算本地梯度。
5. return grad_output * grad_inner, None返回与 forward 两个输入对应的梯度，target 不需要梯度返回 None。
6. loss = CustomHuber.apply(x, target).sum()通过 apply 入口构建计算图并反向传播。
7. ok = gradcheck(CustomHuber.apply, inputs, eps=1e-6, atol=1e-4)用有限差分验证自定义梯度是否正确。

常见误区

面试题