张量网络的自动微分(文献翻译)

本文介绍了自动微分技术在张量网络中的应用,详细阐述了如何利用自动微分来优化张量网络收缩算法,包括张量重整化群(TRG)和角转移矩阵重整化群(CTMRG)方法。通过反向自动微分计算高阶导数,实现了张量网络状态的高效优化,并在自由能高阶微分和iPEPS基于梯度的优化中展示了其能力。同时,文章讨论了稳定后向传播的策略,如通过洛伦兹展宽处理数值稳定性问题,以及使用checkpointing和不动点迭代来节省内存。

摘要生成于 C知道 ,由 DeepSeek-R1 满血版支持, 前往体验 >

本文为文献 Differentiable Programming Tensor Networks 的原创译文,转载请注明出处。
Author:小岚岚

I.引言

张量网络是研究经典统计物理和量子多体物理问题的重要方法[1-3]。近年来,它的应用迅速扩展到不同的领域,包括量子电路的仿真与设计[4-7]、量子误差校正[8,9]、机器学习[10-14]、语言建模[15,16]、量子场论[17-20]、nolographic二象性[21,22]。

张量网络在一般情况下的优化是众多研究方向的核心问题之一。尽管一维矩阵积状态的优化方案非常成功[23-28],但二维或高维张量网络的优化一直是一个具有挑战性的课题。这个困难部分原因是张量收缩的计算成本高,还要部分原因是在高维情况下缺乏有效的优化方案。对于无限平动不变量子系统,优化张量网络状态的挑战尤其紧迫。在这些情况下,同一个张量以多种方式影响变分能量,从而导致高度非线性的优化问题。基于近似虚时算符的优化方案[29-33]努力解决目标函数中的非局部依赖性。文献[34,35]采用基于梯度的优化,结果明显改善。
然而,即使对于一个简单的物理哈密顿量,从解析的角度推导张量网络状态梯度也很麻烦,而且容易出错。这一事实限制了基于梯度的张量网络状态优化在更复杂系统中的广泛应用。替代方法如使用数值导数计算梯度,精度和效率有限;因此,它们只适用于变化参数较少的情况[36,37]。使用链式法则手动推导梯度只适用于精心设计的简单张量网络结构[38]。

可微编程通过以完全可微的方式组合整个张量网络程序并使用自动微分来优化它们,为这些问题提供了优雅而有效的解决方案。在本文中,我们提出基本的自动微分技术,计算张量网络程序的(高阶)导数有效地数值精度。这一进展为一般情况下基于梯度(甚至基于hessian)的张量网络状态优化打开了大门。此外,计算张量网络算法输出的(高阶)梯度提供了一个简单的方法来计算物理量,如比热和磁化率。可微规划方法不知道详细的格点几何、哈密顿量和张量网络收缩方案。因此,该方法的通用性足以支持广泛的张量网络应用。我们的重点是涉及二维无限张量网络的应用,与传统方法相比,可微规划技术提供了显著的好处。我们证明,在通过奇异值分解(SVD)解决数值稳定微分和不动点迭代的内存高效实现等重大技术挑战后,可以获得张量网络状态变分优化的最先进结果。
本文的组织结构如下。在第二节中,我们在张量网络算法的背景下引入自动微分,并将张量网络收缩表示为计算图。在第三节中,我们给出了张量网络算法的稳定和可扩展可微规划的关键技术。在第四节中,我们证明了该方法的能力,并将其应用于无限正方晶格上的经典Ising模型和量子海森堡模型。最后,我们在第五节对未来的研究方向进行了展望。我们的代码实现是公开的(参见参考[39])。

II.一般理论

A.自动微分

通过计算图自动微分是一个统一的框架,涵盖了机器学习训练神经网络,优化量子物理张量网络等等。我们首先回顾自动微分的核心思想,然后解释它在以计算图形式表述的各种张量网络收缩算法中的应用。
自动微分机械地计算用计算机程序[40]表示的计算过程的导数。与数值微分不同,自动微分计算的导数值精确到机器精度。其性能有、普适的理论保证,不超过原程序的算法复杂度[41,42]。
在这里插入图片描述

图1所示。计算图的反向自动微分。黑色箭头表示从输入到输出的正向评价函数。红色箭头表示伴随反向传播的后退步骤。(a)链图。(b)一个更一般的计算图。在向后传递过程中,给定节点的伴随点根据公式(2)计算。

自动微分是现代深度学习应用的计算引擎[43,44]。在量子最优控制[45]和量子化学计算中也有应用,如计算力[46],优化基参数[47]等。自动微分的中心是计算图的概念。计算图是由基本计算步骤组成的有向无环图。图的节点表示数据,数据可以是标量、向量,矩阵,或者张量[48]。图连通性表示计算过程中数据流的依赖性。最简单的计算图是图1(a)所示的链。例如,从输入参数向量theta开始,可以计算一系列中间结果,直到最终的输出L,我们假设它是一个标量。所谓的正向求值就是按顺序遍历链图,即 θ → T 1 → . . . → T n → L \theta \rightarrow T^1\rightarrow...\rightarrow T^n \rightarrow \mathcal{L} θT1...TnL
为了计算目标函数相对于输入参数的梯度,我们可以利用链式法则
∂ L ∂ θ = ∂ L ∂ T n ∂ T n ∂ T n − 1 . . . ∂ T 2 ∂ T 1 ∂ T 1 ∂ θ . ( 1 ) \frac{\partial \mathcal{L} }{\partial \theta}= \frac{\partial \mathcal{L} }{\partial T^n} \frac{\partial T^n }{\partial T^{n-1}}... \frac{\partial T^2 }{\partial T^1} \frac{\partial T^1}{\partial \theta}. (1) θL=TnLTn1Tn...T1T2θT1.(1)
由于我们考虑的情况是,输入维数大于输出维数,在Eq.(1)中通过使用一系列雅可比向量乘积从左到右相乘来计算梯度是更有效的。根据图1(a)所示的计算图,将图向后遍历,并将梯度信号从输出端传播回输入端。用这种方法计算导数叫做反向自动微分。这种方法通常称为反向传播算法[43],被认为是训练深度神经网络最成功的方法。引入伴随变量 T ˉ = ∂ L / ∂ T \bar T=\partial \mathcal{L} /\partial T Tˉ=L/T 来表示最终输出L相对于变量T的梯度。我们可以看到,反向自动微分将伴随变量从 T n ‾ = L ˉ ∂ L / ∂ T n \overline {T^n} =\bar \mathcal{L} \partial \mathcal{L}/\partial T^n Tn=LˉL/Tn L ˉ = 1 \bar \mathcal{L}=1 Lˉ=1
一直传播到

<think>我们讨论的是在PyTorch中进行掩码操作后梯度是否会停止的问题。首先需要明确:在PyTorch中,大多数标准的掩码操作(例如使用布尔掩码或数值掩码进行元素选择)不会自动停止梯度,除非显式地使用`detach()`或设置`requires_grad=False`。 但是,有一些特殊情况需要注意: 1. **掩码操作本身是可微的**:如果掩码操作是通过可微运算生成的(例如使用sigmoid生成软掩码),那么梯度会通过掩码操作传播到生成掩码的参数上。 2. **硬掩码(不可微操作)**:如果掩码是二进制的(0和1)并且是通过不可微操作生成的(例如通过`>0.5`的阈值将概率掩码变为二进制),那么梯度在掩码本身会停止(因为二进制掩码的生成过程不可微),但是掩码作用到的张量如果需要梯度,则梯度仍然会传播到该张量上。 具体来说,考虑以下两种情况: **情况1:使用掩码进行元素选择(如`masked_select`)** 这种操作会返回一个新的张量,该张量是原始张量的一个子集。由于这个操作只是选择元素,所以梯度会从结果张量传播到原始张量的被选中元素上。但是,未被选中的元素不会接收到梯度(因为它们在计算图中没有参与运算)。注意,这个操作本身不会停止梯度,但会改变梯度的传播路径。 **情况2:使用掩码进行元素置零(如`mask * tensor`)** 这是最常见的掩码操作。在这种情况下,梯度会正常传播到原始张量上,因为乘法操作是可微的。具体来说: - 对于掩码中为1的位置,梯度等于原始梯度。 - 对于掩码中为0的位置,梯度也为0(因为乘以0),所以这些位置的参数不会更新。 因此,掩码操作并不会自动停止整个计算图的梯度,它只是改变了梯度值(将一部分梯度置零)。如果希望完全停止梯度,必须使用`detach()`。 下面通过代码示例说明: 示例1:掩码乘法(梯度部分置零) ```python import torch # 创建需要梯度的张量 x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True) # 创建掩码(不需要梯度,因为它是布尔类型或者整数类型,但即使它是浮点数且需要梯度,乘法操作也会将梯度传播到掩码上) mask = torch.tensor([1, 0, 1], dtype=torch.float32) y = x * mask # 结果为 [1, 0, 3] y.sum().backward() # 计算梯度 print(x.grad) # 输出: tensor([1., 0., 1.]) ``` 示例2:使用掩码选择元素(梯度只传播到被选中的元素) ```python x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True) mask = torch.tensor([True, False, True]) y = x[mask] # 等价于 x[mask].sum() 得到 [1,3] y.sum().backward() print(x.grad) # 输出: tensor([1., 0., 1.]) ``` 示例3:如果掩码本身需要梯度(例如,掩码是通过学习得到的概率值,然后被二值化) ```python # 生成掩码的参数 mask_param = torch.tensor([0.6, 0.4, 0.9], requires_grad=True) # 通过阈值二值化(不可微操作) hard_mask = (mask_param > 0.5).float() # 得到 [1,0,1],注意:这步操作不可微,所以梯度在hard_mask处停止 x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True) y = x * hard_mask # 结果为 [1,0,3] # 计算梯度 y.sum().backward() # 输出梯度 print(x.grad) # tensor([1.,0.,1.]),x的梯度正常 print(mask_param.grad) # 输出None或者全0?实际上,由于二值化操作不可微,所以mask_param不会得到梯度 ``` 但是,如果我们使用Gumbel-Softmax等技巧来生成可微的硬掩码,则梯度可以传播到mask_param上。不过,标准的二值化操作会断开梯度。 **总结**: - 在PyTorch中,掩码操作(如乘法或索引选择)不会自动停止梯度,但会改变梯度的值(例如,将一部分梯度置零)。 - 如果掩码是通过不可微操作生成的(如二值化),那么生成掩码的过程会阻止梯度传播到生成掩码的参数上,但不会阻止梯度传播到被掩码作用的其他张量上。 - 如果希望完全停止梯度,应该使用`detach()`方法。 因此,回到用户的问题:使用掩码操作后梯度是否会自动停止?答案是否定的,除非掩码操作中包含了显式的梯度停止操作(如`detach()`)或者掩码生成过程是不可微的(导致生成掩码的参数无法得到梯度),但被掩码作用的数据张量仍然可能得到梯度(部分位置)。 参考文献: [^1]: 在TPU上运行PyTorch时,将张量移至CPU会降低性能,但有时为了使用numpy操作而不得不这样做。 [^2]: 在TPU上运行PyTorch可能会遇到一些限制和性能问题。 [^3]: DataLoader用于封装数据,但用户问题不直接相关。 相关问题: 1. 如何在PyTorch中正确停止梯度? 2. 掩码操作在反向传播中是如何影响梯度计算的? 3. 在TPU上使用掩码操作是否有特殊注意事项?</think>在 PyTorch 中使用掩码操作后,**梯度不会自动停止**,但掩码会影响梯度计算的方式。以下是详细分析: ### 1. 掩码操作对梯度的影响原理 PyTorch 的自动微分系统会追踪所有张量操作,包括掩码操作。当使用掩码时: - **梯度仍会传播**:掩码操作(如乘法、索引选择)本身是可微操作,梯度会继续向后传播 - **梯度值被修改**:掩码会将部分位置的梯度置零,但其他位置的梯度仍正常计算 - **计算图保持完整**:掩码操作作为计算图的一部分被记录 ```python import torch # 创建需要梯度的张量 x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True) mask = torch.tensor([1, 0, 1]) # 掩码张量 # 应用掩码 y = x * mask # 输出: tensor([1., 0., 3.], grad_fn=<MulBackward0>) loss = y.sum() loss.backward() print(x.grad) # 输出: tensor([1., 0., 1.]) ``` *输出说明:索引1位置的梯度被置零,但其他位置梯度正常计算* ### 2. 需要显式停止梯度的场景 在以下情况下需要手动停止梯度: - **数据预处理**:输入数据预处理管道 - **特征提取**:使用预训练模型提取特征时 - **梯度屏蔽**:实现自定义梯度修改逻辑 - **混合精度训练**:某些操作需在FP32下进行 ### 3. 显式停止梯度的方法 使用 `.detach()` 或 `torch.no_grad()` 上下文管理器: ```python # 方法1: .detach() masked_data = input_data.detach() * mask # 方法2: torch.no_grad() with torch.no_grad(): processed = complex_operation(input_data) masked_output = processed * mask ``` ### 4. 特殊场景:TPU上的掩码操作 在TPU环境中使用掩码时需特别注意: ```python # TPU优化建议 (参考引用[1][2]) def tpu_safe_mask(x, mask): # 将计算移至TPU设备 x = x.to(xla_device) mask = mask.to(xla_device) # 对于不需要梯度的部分使用detach with torch.no_grad(): static_mask = mask.clone().detach() return x * static_mask ``` ### 5. 性能优化建议 1. **避免不必要的CPU转移**: ```python # 不推荐 (引用[1]) # my_tensor.cpu().detach().numpy() # 推荐保持张量在设备上 masked = x * mask.to(x.device) ``` 2. **使用原地操作**: ```python mask = mask.bool() # 转换为布尔掩码节省内存 x[mask] = 0 # 原地修改 ``` 3. **DataLoader优化**: ```python # 在Dataset中预计算掩码 (引用[3]) class MaskedDataset(Dataset): def __init__(self, data): self.data = data self.mask = self._generate_mask() def __getitem__(self, idx): return self.data[idx] * self.mask[idx] ``` ### 总结 | 操作类型 | 梯度行为 | 是否需要显式停止 | |---------|---------|----------------| | 标准掩码乘法 | 梯度部分传播 | 否 | | 索引掩码选择 | 梯度部分传播 | 否 | | 数据预处理 | 不应传播梯度 | 是 | | 特征提取 | 通常需要停止 | 是 | | 梯度屏蔽 | 自定义梯度流 | 部分需要 | 在PyTorch中,掩码操作**不会自动停止梯度**,但会修改梯度值。需要根据具体场景决定是否使用`.detach()`或`torch.no_grad()`显式停止梯度,特别是在TPU等特殊硬件环境或性能敏感场景中[^1][^2]。
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值