[e3nn] 等变神经网络 | 线性层o3.Linear

第4章：等变神经网络模块

欢迎回来～

在我们探索e3nn的旅程中，我们已经揭示了一些基本概念：

在第1章：不可约表示（Irreps）中，我们学习了Irreps作为等变数据的标签，告诉我们数据在旋转和反射下如何变换。
在第2章：球谐函数中，我们看到了这些抽象的Irreps如何通过方向性的"波模式"具体表示。
在第3章：张量积中，我们发现了e3nn如何结合不同的Irreps来创建新的、更复杂的等变特征。

现在，我们有了这些惊人的构建块。但我们如何实际构建一个神经网络呢？在传统神经网络中，我们有"层"，如Linear变换和激活函数（如ReLU或Sigmoid）。我们如何创建这些层的等变版本？

这就是等变神经网络模块的用武之地。

这些是您用来构建一个固有尊重3D对称性的神经网络的实际构建块。

它们在底层利用Irreps和TensorProduct来确保当输入旋转或反射时，网络的输出正确变换，就像物理定律一样。

本章的目标是理解如何使用e3nn来创建：

一个等变线性变换，类似于torch.nn.Linear层。
一个等变非线性（激活函数），类似于torch.nn.ReLU。

让我们开始吧～

等变线性层：`e3nn.o3.Linear`

就像标准线性层将其输入乘以权重矩阵一样，e3nn.o3.Linear层对其输入特征执行线性变换。

关键区别在于它以等变的方式执行。

`e3nn.o3.Linear`的功能

e3nn.o3.Linear层接受由irreps_in描述的输入特征，并将其转换为由irreps_out描述的输出特征。

关键约束是Linear层只能组合相同Irrep类型的组件。

例如，输入中的所有0e组件可以贡献给输出中的所有0e组件，输入中的所有1o组件可以贡献给输出中的所有1o组件。但0e输入不能通过简单的线性层变为1o输出。

这确保了数据的基本对称属性得以保留。

可以把它想象成整理袜子：

可以将白袜子与其他白袜子组合，彩色袜子与其他彩色袜子组合。

但不能通过简单的重新排列将白袜子"变成"彩色袜子。e3nn.o3.Linear保持了这些"类型"。

如何使用`e3nn.o3.Linear`

通过使用不可约表示（Irreducible Representations）对象指定其irreps_in和irreps_out来定义e3nn.o3.Linear层。

import torch
from e3nn import o3# 定义输入数据的Irreps：
# 一个标量（0e）和一个向量（1o）的混合
irreps_in = o3.Irreps("0e + 1o")
print(f"输入Irreps: {irreps_in}")# 创建符合这些Irreps的随机输入数据
x = irreps_in.randn(10, -1) # 10个样本，其中-1被irreps_in.dim（4）替换
print(f"输入数据形状: {x.shape}")# 定义输出数据的Irreps：
# 两个标量（0e）和两个向量（1o）的混合
irreps_out = o3.Irreps("2x0e + 2x1o")
print(f"输出Irreps: {irreps_out}")# 创建一个将'irreps_in'数据转换为'irreps_out'数据的线性层
# e3nn确保此层尊重旋转/反射
linear_layer = o3.Linear(irreps_in=irreps_in, irreps_out=irreps_out)
print(f"\n线性层创建: {linear_layer}")# 将线性层应用于输入数据
y = linear_layer(x)
print(f"输出数据形状: {y.shape}")

输出：

输入Irreps: 1x0e+1x1o
输入数据形状: torch.Size([10, 4])
输出Irreps: 2x0e+2x1o线性层创建: Linear(1x0e+1x1o -> 2x0e+2x1o | 10 weights)
输出数据形状: torch.Size([10, 8])

如所见，Linear层成功地将输入（维度4）转换为输出（维度8），就像常规的torch.nn.Linear一样。

e3nn的魔力在于这种转换保证是等变的。

底层实现：`e3nn.o3.Linear`

当创建e3nn.o3.Linear层时，e3nn智能地确定输入Irreps的哪些部分可以连接到输出Irreps的哪些部分。

具体来说，它知道只有相同l和p的Irrep可以连接。

在这里插入图片描述

连接逻辑：e3nn识别irreps_in和irreps_out之间所有可能的连接，其中Irrep类型（如0e到0e，1o到1o）匹配。
权重分配：对于每个有效连接，它分配一个可学习的权重矩阵。此矩阵的大小取决于输入和输出Irreps的multiplicity。例如，如果输入中有2x0e，输出中有3x0e，则需要一个2x3权重矩阵用于该连接。
高效计算：e3nn的内部代码（在e3nn/o3/_linear.py中）生成高度优化的torch.einsum操作，而不是手动拆分输入并为每个Irrep类型应用矩阵乘法。einsum是一个强大的函数，允许简洁高效的张量收缩，有效地执行给定输入张量所需的所有乘法和求和，确保等变性和良好的性能。

Linear模块确保每个输入Irrep组件（mul_in x ir_in）仅当ir_in == ir_out时转换为每个输出Irrep组件（mul_out x ir_out）。然后使用大小为mul_in x mul_out的权重矩阵进行此转换。

等变非线性：`e3nn.nn.Gate`

非线性对于神经网络学习复杂模式至关重要。

然而，直接将标准激活函数（如ReLU或Sigmoid）应用于向量组件会破坏等变性。

如果旋转一个向量(x, y, z)，然后对每个组件应用ReLU，这与先应用ReLU再旋转不同。

e3nn.nn.Gate模块提供了一种巧妙的方式引入非线性，同时严格保持等变性。

`e3nn.nn.Gate`的功能

Gate通过将高阶Irreps（如向量，l=1，或张量，l=2）与标量（0e Irreps的组件）相乘来实现非线性，这些标量首先通过标准激活函数传递。

可以这样理解：

**标量（0e Irreps）**可以直接应用正常激活函数（如tanh、sigmoid、relu），因为标量在旋转下不会改变。
**高阶Irreps（l > 0）**不能直接应用标准激活。相反，e3nn.nn.Gate使用另一个标量（我们称之为"门"）来乘以这些高阶Irreps。这种乘法缩放高阶特征，但不改变其方向或旋转方式。就像灯的调光器开关——它改变亮度（大小），但不改变光的颜色或方向。

这种标量乘法是张量积的特殊情况：0e x Irrep -> Irrep。当您将标量乘以任何其他Irrep时，结果Irrep具有与非标量Irrep相同的l和p。这就是如何保持等变性。

如何使用`e3nn.nn.Gate`

code: https://github.com/lvy010/AI-exploration/tree/main/neural_network

Gate模块接受五个主要参数：

irreps_scalars：仅作为标量并接收直接激活的特征的Irreps。
act_scalars：irreps_scalars的激活函数列表。
irreps_gates：用作"门"的特征的Irreps。这些必须是标量（0e）并接收直接激活。
act_gates：irreps_gates的激活函数列表。
irreps_gated：将被激活门乘以的特征的Irreps。irreps_gates中的Irrep组数必须匹配irreps_gated。

import torch
from e3nn import o3
from e3nn.nn import Gate# 示例：假设输入特征是两个标量和一个向量
# 输入将是：1x0e（标量特征）+ 1x0e（门标量）+ 1x1o（向量特征）
irreps_in = o3.Irreps("1x0e + 1x0e + 1x1o")
x = irreps_in.randn(1, -1)
print(f"输入Irreps: {irreps_in}")
print(f"输入数据形状: {x.shape}")# 定义Gate模块：
# - 第一个0e将直接激活（无门控）。
# - 第二个0e将被激活并用作门。
# - 1o将由激活的第二个0e门控。
equivariant_gate = Gate(irreps_scalars="1x0e",       # 一个0e标量act_scalars=[torch.tanh],    # 对其应用tanhirreps_gates="1x0e",         # 一个0e标量作为门act_gates=[torch.sigmoid],   # 对门标量应用sigmoidirreps_gated="1x1o"          # 一个1o向量被门控
)
print(f"\n等变Gate创建: {equivariant_gate}")
print(f"输出Irreps: {equivariant_gate.irreps_out}")# 将Gate应用于输入数据
y = equivariant_gate(x)
print(f"输出数据形状: {y.shape}")# 让我们手动拆分输入以检查输出（简化说明）
# scalars_input = x[:, irreps_in.slices_by_irreps[0]]
# gates_input = x[:, irreps_in.slices_by_irreps[1]]
# gated_input = x[:, irreps_in.slices_by_irreps[2]]# print(f"\n(手动说明):")
# print(f"激活标量: {torch.tanh(scalars_input)}")
# print(f"激活门: {torch.sigmoid(gates_input)}")
# print(f"门控向量: {torch.sigmoid(gates_input) * gated_input}")

输出：

输入Irreps: 1x0e+1x0e+1x1o
输入数据形状: torch.Size([1, 5])等变Gate创建: Gate(1x0e+1x0e+1x1o -> 1x0e+1x1o)
输出Irreps: 1x0e+1x1o
输出数据形状: torch.Size([1, 4])

注意输出Irreps 1x0e+1x1o如何组合激活的irreps_scalars和被irreps_gates乘以后的irreps_gated。

irreps_gates本身不出现在输出中，因为它们在乘法中被消耗。总维度从5缩小到4（1x0e（激活标量）+ 1x1o（门控向量））。

底层实现：`e3nn.nn.Gate`

e3nn.nn.Gate模块（位于e3nn/nn/_gate.py）执行以下步骤：

在这里插入图片描述

特征拆分：首先将输入张量（组合了所有irreps_scalars、irreps_gates和irreps_gated组件）拆分为其相应部分。这是由内部_Sortcut实用程序处理的，以高效提取张量的正确部分。
标量激活：irreps_scalars和irreps_gates（均为0e类型）通过标准的e3nn.nn.Activation模块传递，应用其各自的非线性函数（如tanh、sigmoid）。
门控（张量积）：然后，激活的标量门与irreps_gated特征逐元素相乘。这是由e3nn.o3.ElementwiseTensorProduct（张量积的一种专门形式）执行的，确保标量-Irrep乘法保持等变性。
连接：最后，激活的标量和新的"门控"特征被连接在一起，形成输出张量。