Graph Neural Networks for Learning Equivariant Representations of Neural Networks
Chen Kai Architect

本文提出了一种创新的方法,通过引入神经图(Neural Graphs)的概念,将神经网络的参数(如权重和偏置)表示为图结构。这种图结构不仅能够表示神经网络的计算过程,还能够捕捉神经网络架构之间的差异。通过这种表示,模型能够自然地保持对称性,即使在面对多种不同的神经网络架构时,也能保证等变性(equivariance)。

背景介绍

神经图的概念

神经图(Neural Graph)的核心思想是将神经网络的参数和架构表示为图,其中:

  • 节点(Nodes) 代表神经元及其特征(如偏置)。
  • 边(Edges) 代表神经元之间的连接权重。

通过这种图结构的表示方式,模型在处理不同架构的神经网络时,能够保持对这些神经网络排列方式的对称性。这种对称性通过图神经网络(GNN)的特性来实现,GNN能够天然地处理图的拓扑结构,而不受节点排列顺序的影响。

传统神经网络的参数表示

在讨论神经网络的排列对称性之前,了解传统方法是如何表示和操作神经网络的参数是很重要的。传统上,神经网络的表示方法主要包括以下几种:

  1. 平坦表示(Flattened Representations)

平坦表示是指将神经网络的所有参数(例如权重矩阵和偏置向量)展平为一个长向量或高维矩阵。这种方法在一些深度学习框架中被广泛使用,因为它可以简化操作,如计算梯度和优化参数。

  • 优点: 简化了参数的存储和操作,尤其在实现自动微分和参数更新时非常有效。
  • 缺点: 平坦表示忽略了神经网络内部的结构信息,特别是忽略了神经元排列的顺序。因此,如果重新排列网络层中的神经元,这种表示方法无法捕捉到这种变化,也无法保持排列对称性。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
import torch
import torch.nn as nn

# 定义一个简单的多层感知机(MLP)
class SimpleMLP(nn.Module):
def __init__(self):
super(SimpleMLP, self).__init__()
self.fc1 = nn.Linear(784, 128)
self.fc2 = nn.Linear(128, 64)
self.fc3 = nn.Linear(64, 10)

def forward(self, x):
x = torch.flatten(x, start_dim=1)
x = torch.relu(self.fc1(x))
x = torch.relu(self.fc2(x))
x = self.fc3(x)
return x

# 初始化模型
model = SimpleMLP()

# 平坦表示,将所有参数展平成一个向量
flattened_params = torch.cat([p.view(-1) for p in model.parameters()])
print(flattened_params)
  1. 参数共享和卷积(Parameter Sharing and Convolutional Representations)

在卷积神经网络(CNN)中,参数共享是一种常见策略,即使用相同的卷积核(参数)在图像的不同部分应用相同的计算。这种表示方式利用了数据的局部性和空间结构,减少了参数的数量。

  • 优点: 提高了计算效率并减少了模型的参数量,使得卷积操作非常适合处理图像数据。
  • 缺点: 尽管卷积操作可以在一定程度上保持排列对称性(例如,卷积核的移动不改变图像特征的局部结构),但是这种方法通常依赖于数据的空间结构,对于非图像数据或没有明确空间结构的神经网络架构(如全连接层),它并不适用。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
import torch.nn as nn
import torch

# 定义一个简单的卷积神经网络(CNN)
class SimpleCNN(nn.Module):
def __init__(self):
super(SimpleCNN, self).__init__()
self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)
self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
self.fc1 = nn.Linear(64 * 7 * 7, 128)
self.fc2 = nn.Linear(128, 10)

def forward(self, x):
x = torch.relu(self.conv1(x))
x = torch.max_pool2d(x, 2)
x = torch.relu(self.conv2(x))
x = torch.max_pool2d(x, 2)
x = torch.flatten(x, start_dim=1)
x = torch.relu(self.fc1(x))
x = self.fc2(x)
return x

# 初始化模型
cnn_model = SimpleCNN()
print(cnn_model)
  1. 权重共享机制(Weight-Sharing Mechanisms)

权重共享是一种确保不同部分的网络使用相同参数的策略,常用于神经网络的对称性保持。例如,在循环神经网络(RNN)中,不同时间步之间共享相同的权重,以捕捉序列数据的时间相关性。

  • 优点: 能有效捕捉数据的时间序列特征,并减少参数数量。
  • 缺点: 权重共享机制需要精心设计的共享模式,且通常是手动实现的。它的适用性通常限制在特定类型的网络(如RNN),并且不能自然地适应所有类型的网络架构。此外,这种方法仍然没有完全解决排列对称性的问题,因为它不考虑节点排列的不同变化。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
import torch.nn as nn
import torch

# 定义一个简单的循环神经网络(RNN)
class SimpleRNN(nn.Module):
def __init__(self):
super(SimpleRNN, self).__init__()
self.rnn = nn.RNN(input_size=10, hidden_size=20, num_layers=2, batch_first=True)
self.fc = nn.Linear(20, 1)

def forward(self, x):
out, _ = self.rnn(x)
out = self.fc(out[:, -1, :]) # 使用最后一个时间步的输出
return out

# 初始化模型
rnn_model = SimpleRNN()
print(rnn_model)

权重共享的概念:

  • 权重共享是一种技术,其中不同的网络层或网络部分共享相同的权重矩阵。这种方法可以减少模型的参数数量,从而提高训练效率和减少过拟合的风险。
  • 在神经网络的等变性任务中,权重共享被用来确保模型对不同的神经元排列方式保持不变。具体来说,模型会在多个神经元之间共享权重,使得无论这些神经元如何排列,输出都保持不变。

复杂性来源:

  • 这种权重共享机制可能会非常复杂,因为它需要手动指定哪些神经元之间共享权重,以及如何根据网络的架构变化来调整权重共享模式。对于每一个新的神经网络架构,研究者需要重新设计权重共享的规则,以确保等变性。
  • 比如,在卷积神经网络(CNN)中,研究者可能需要确保不同卷积核的通道在前一层和后一层之间按照某种方式共享权重,以保持排列对称性。这种手工设计的过程非常复杂且耗时,特别是在面对多种网络架构时。
  1. 这些方法存在的问题
    1. 结构信息的缺失: 传统方法(如平坦表示)无法表示神经网络中的层次结构和连接信息,这些信息在神经元重新排列时非常关键。
    2. 硬编码的结构假设: 参数共享和卷积表示方法依赖于硬编码的结构假设(如空间局部性),这些假设不适用于所有类型的神经网络架构。它们假定了一种特定的排列和连接方式,因此无法应对神经元或层重新排列的变化。
    3. 手动设计和调整: 复杂的权重共享机制需要手动设计和调整,不具备通用性和灵活性,无法自然地适应不同的网络架构变化。

神经图的数学表示

假设我们有一个神经网络,包含若干层,每层的参数可以表示为权重矩阵 和偏置向量 。神经图 可以定义如下:

  • 节点特征矩阵 :每个节点的特征可以是对应层的偏置
  • 边特征矩阵 :每条边的特征可以是连接相应层的权重

例如,对于一个包含 层的多层感知机(MLP),其权重和偏置可以表示为:

等变性(Equivariance)与排列对称性

等变性指的是模型在对输入进行某种特定变换(如神经元的排列)后,输出也会相应地变化,但整体功能保持不变。在神经网络的上下文中,这意味着即使神经元在隐藏层中的排列顺序改变了,网络的功能(输出)应该保持不变。

等变性的数学定义

如果一个函数 对一个变换 是等变的,那么对于任意输入 ,有:

在神经网络中,等变性确保了当我们对神经元重新排列(通过对权重矩阵进行相应的置换操作)时,神经网络的输出保持不变。

  1. 神经图对称性(Neural Graph Symmetries)

    神经图作为一种图结构,可以运用许多现有的图属性。具体来说,图的对称性由对称群 表示,其中 是节点的总数。设 为将排列 映射到相应排列矩阵的群表示。基于此,图 上的群作用 定义为:

    这意味着,我们希望使用神经网络来处理神经图,同时保持对这些特定对称性的等变性。一个典型的例子是标准的消息传递图神经网络(MPNNs),其操作天然具有等变性(Bronstein et al., 2021)。

  2. 神经元置换群(Neuron Permutation Group)

    神经元置换群(NP组) 被定义为 ,其中每个 代表第 层的神经元的排列。群的元素表示为 。对于特定层 ,动作 被定义为:

    从这两个方程中中可以看出, 的一个子群,因为 是由各对称群的直积定义的,其度数加起来正好等于 。任何 的动作都对应 的一个动作。因此,满足 -等变的模型必然也满足 -等变。

  3. 讨论(Discussion)

    • 为什么要选择 而不是 作为对称群?一个主要原因是, 包含了多个 的选择。使用一个单一的 -等变模型,我们的方法可以自动适应于任意选择的 ,只要这些数加起来等于 ;同样的模型可以处理多种不同类型的架构。相比之下,先前的工作(如 Navon et al., 2023; Zhou et al., 2023a)只针对特定的 进行等变,限制了模型只能处理特定架构。
    • 本文提出的模型在每个节点应用相同的更新函数,不考虑该节点所属的层,这对于单一、固定架构上的性能特别相关。而在需要处理不同架构的情况下,神经图表示更合适。

    实验上,作者的模型不仅在基准测试上表现优异,还具有更好的泛化能力。例如在MNIST INR分类任务中,尽管训练损失差不多,NG-T的测试损失较低,表明使用 等变的参数共享方式不仅没有损害性能,反而可能有助于提升泛化能力。

    通俗理解神经图对称性与神经元置换群的关系

    要理解神经图对称性(Neural Graph Symmetries)和神经元置换群(Neuron Permutation Group, NP Group)之间的关系,我们首先需要了解它们各自的定义和在神经网络中的作用。

    1. 神经图对称性

    神经图对称性涉及的是整个神经网络结构如何可以在不改变其功能的情况下重新排列。具体来说,当我们将神经网络表示为一个图时,图的节点表示网络的神经元,边表示神经元之间的连接或权重。这个图的对称性由对称群 描述,其中,是所有层中神经元总数。

    对称群 是一个非常大的群,因为它考虑了神经网络中所有神经元可能的排列方式。因此,使用 描述的神经图对称性可以表示神经网络中所有神经元任意重排的可能性。这种对称性有助于确保无论神经网络的神经元如何排列,图神经网络(如消息传递神经网络)对这种排列都是等变的。

    1. 神经元置换群

    神经元置换群 是一种特定的对称群,它关注的是神经网络每一层内部的神经元排列,而不是整个网络。它是多个小的对称群的直积,即 ,其中 表示第 层神经元的排列。

    不同, 只关注在同一层中的神经元置换,而不考虑不同层之间神经元的关系。因此, 的一个子群。换句话说, 的一部分,但并不包含所有可能的排列(因为它只限制在每一层内部的排列)。

    3. 二者的关系

    神经图对称性(由 描述)涵盖了整个神经网络中所有可能的排列和置换,而神经元置换群(由 描述)则只关注每一层内的神经元排列。由于 仅仅描述了每层内部的排列,它比 更具体(或者说,更局限)。

    在文中,作者通过对神经图进行对称性分析,指出对于一个神经图的操作,如果它对整个对称群 是等变的,那么它对其任意子群 也是等变的。也就是说,如果一个模型对神经网络中所有神经元的排列()保持等变性,那么它在处理每层内部神经元排列()时,自然也保持等变性。这说明,选择一个 -等变的模型可以自动适应任何特定层次的排列,这使得模型更具通用性,能够处理更多类型的神经网络架构。

    通过这种方式,神经图可以以统一的方式处理不同的神经网络架构,同时保持与传统模型相同的对称性和等变性。这是神经图对称性和神经元置换群之间的核心关系。

等变性在图神经网络中的应用

在图神经网络中,等变性通常与图的对称性(如节点排列的顺序不影响图的性质)相关联。GNN 通过消息传递机制从邻居节点收集信息,学习节点的表示或图的整体表示。在这个过程中,等变性保证了以下几点:

  1. 对称性不变性: 图是无序的结构,因此节点的顺序不应影响网络的输出。GNN 保证等变性,确保无论节点输入顺序如何变化,网络的输出保持不变。例如,如果两个图在结构上相同,GNN 的输出也应该相同,尽管节点的索引可能不同。
  2. 信息的综合: 当聚合来自不同邻居节点的信息时,GNN 需要确保这种聚合操作是等变的。换句话说,改变节点的顺序不会改变节点特征的聚合结果。

具体细节

神经网络表示为神经图

论文详细描述了如何将神经网络(如多层感知机MLP和卷积神经网络CNN)转换为神经图。

MLP → 图

对于多层感知机(MLP),可以将其权重矩阵和偏置向量分别作为边特征和节点特征。具体来说,神经图 的构建如下:

假设一个MLP有 层,每层的权重矩阵为 ,偏置为 ,其中

  • 节点特征矩阵 包含所有层的偏置向量:

  • 边特征矩阵 包含所有层的权重矩阵:

通过这种表示,每个层的神经元及其连接关系可以表示为图中的节点和边。这样,模型可以处理不同的神经网络架构,而不需要对每种架构进行手动调整。

这么做有什么好处呢?

首先是统一表示不同架构的能力:传统的神经网络表示方法通常是为特定的网络架构量身定制的。如果想在不同架构之间迁移模型,往往需要重新设计或调整模型参数。然而,图结构的表示方法能够捕捉到神经网络的基本计算结构,而不依赖于具体的神经元排列或网络层次。这意味着,无论神经网络有多少层、每层有多少神经元,或者连接方式如何变化,只要这些信息可以表示为图结构,模型就可以处理它们。假设有两个不同架构的MLP,一个有3层,另一个有5层,使用传统的神经网络模型,很难直接在这两种不同架构之间迁移学习。但通过图表示,我们可以将这两种架构转换为图结构,其中每一层的权重和偏置分别作为边和节点特征。这样,无论是3层还是5层的MLP,图神经网络(GNN)都可以在相同的框架下处理它们,而不需要重新设计模型。

其次,图结构表示方法能够保持神经网络中的排列对称性。在神经网络中,神经元的排列方式并不会影响最终的输出,只要连接关系保持不变。例如,重新排列同一层内的神经元,只会导致权重矩阵和偏置向量的重新排列,但不影响网络的整体功能。传统的方法在神经元重新排列时,可能会误以为这是一个全新的架构,从而错误地重新训练模型。而通过图结构表示,节点和边的特征(即偏置和权重)只与其拓扑位置相关,而不是排列顺序。GNN天生具有处理图结构的能力,因此对同一结构的不同排列方式具有鲁棒性。

使用图结构表示MLP的另一个优势在于,图神经网络模型可以很自然地扩展和泛化。由于图神经网络处理的是节点和边的特征,而不是具体的参数矩阵,这使得它可以学习到更普遍的特征表示,适用于各种复杂的网络结构。比如,当处理多种类型的神经网络(如卷积神经网络CNN、循环神经网络RNN)时,GNN可以将这些网络的计算流程转换为图结构,并使用相同的机制进行学习和推理。这种通用性使得图神经网络成为更为灵活和强大的工具,适用于更广泛的深度学习任务。

CNN → 图

CNN 的图表示

对于卷积神经网络(CNN),其图表示方式与MLP类似,但需要考虑卷积核的空间维度以及通道之间的关系。论文中提到,对于每一个卷积层,卷积核(即滤波器)可以通过将其空间维度展平后表示为边特征。这样做的原因是为了保证不同尺寸的卷积核在同一个网络中具有一致的表示。

假设我们有一个CNN,其中每个卷积层的卷积核大小为 ,偏置为 ,其中 表示卷积核的数量(即输出通道数), 表示输入通道数, 为卷积核的宽度和高度。为了表示这些卷积核,论文建议如下步骤:

  1. 卷积核展平:将每个卷积核的空间维度展平为一个向量。例如,若卷积核大小为 ,则展平后得到一个长度为9的向量
  2. 零填充:为了在图中表示不同尺寸的卷积核,我们首先将所有卷积核零填充至一个最大尺寸 。这样,我们可以在不失信息的情况下,将不同尺寸的卷积核统一表示为相同长度的向量
  3. 边特征矩阵:展平后的卷积核可以视为多维等效的线性层权重,然后将其作为边特征 ,其中

这种表示方式的一个关键好处是,它能够统一处理不同架构的卷积层,特别是当卷积核大小和通道数发生变化时。通过将卷积层的卷积核表示为固定长度的向量,图神经网络能够灵活地适应不同的卷积神经网络架构。

Flattening层和线性层

论文进一步讨论了如何处理CNN中的Flattening层和线性层:

  • Flattening层:CNN通常会在最后一个卷积层之后进行特征的展平操作,将其转化为一个一维向量。传统的Flattening方法依赖于固定的输入分辨率,无法处理任意尺寸的输入图像。为了解决这个问题,神经图不绑定任何空间分辨率,而是通过自适应池化来集成这种变化。

自适应池化层能够将任意大小的输入特征图缩放到指定的输出大小。这种层非常适合处理不定尺寸的输入,因为它在池化过程中动态调整窗口大小,以确保输出尺寸固定。自适应池化有两种常见类型:自适应平均池化(AdaptiveAvgPool2d)和自适应最大池化(AdaptiveMaxPool2d)。

  • 优点:自适应池化不依赖输入的尺寸,可以处理任意大小的输入图像。它的输出尺寸是固定的,这使得后续的全连接层能够接受不同大小的输入图像而不需要重新训练。
  • 缺点:由于自适应池化层会将所有输入特征图缩放到相同的输出尺寸,这可能导致丢失一些局部细节,尤其是在输入尺寸非常大或者非常小的情况下。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
import torch
import torch.nn as nn

class AdaptivePoolingExample(nn.Module):
def __init__(self):
super(AdaptivePoolingExample, self).__init__()
self.conv1 = nn.Conv2d(1, 16, kernel_size=3, padding=1)
self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
self.adaptive_pool = nn.AdaptiveAvgPool2d((7, 7)) # 自适应平均池化到7x7的大小
self.fc1 = nn.Linear(32 * 7 * 7, 128)
self.fc2 = nn.Linear(128, 10)

def forward(self, x):
x = torch.relu(self.conv1(x))
x = torch.max_pool2d(x, 2)
x = torch.relu(self.conv2(x))
x = torch.max_pool2d(x, 2)
x = self.adaptive_pool(x) # 自适应池化层
x = x.view(x.size(0), -1)
x = torch.relu(self.fc1(x))
x = self.fc2(x)
return x

model = AdaptivePoolingExample()
input_tensor = torch.randn(1, 1, 32, 32) # 输入为1张1通道的32x32图片
output = model(input_tensor)
print("输出尺寸:", output.size())
# 输出尺寸: torch.Size([1, 10])
  • 线性层:线性层通常用于在Flattening或自适应池化之后生成图像的最终特征向量。最直接的方式是将线性层视为一个特殊的1×1卷积,允许线性层和卷积层之间具有统一的表示。这种处理方式使得模型能够灵活地处理各种线性变换和卷积操作。

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    import torch
    import torch.nn as nn

    # 定义一个简单的CNN,将线性层视为1x1卷积
    class CNNWith1x1Conv(nn.Module):
    def __init__(self):
    super(CNNWith1x1Conv, self).__init__()
    self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)
    self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
    self.adaptive_pool = nn.AdaptiveAvgPool2d((1, 1))
    self.conv1x1 = nn.Conv2d(64, 10, kernel_size=1) # 使用1x1卷积代替全连接层

    def forward(self, x):
    x = torch.relu(self.conv1(x))
    x = torch.max_pool2d(x, 2)
    x = torch.relu(self.conv2(x))
    x = torch.max_pool2d(x, 2)
    x = self.adaptive_pool(x)
    x = self.conv1x1(x)
    x = x.view(x.size(0), -1) # 将输出展平成一维向量
    return x

    # 初始化模型并打印输出尺寸
    model = CNNWith1x1Conv()
    print(model)

    # 创建一个随机输入张量,模拟输入图片
    input_tensor = torch.randn(1, 1, 28, 28) # 1个批次,1个通道,28x28的图片
    output = model(input_tensor)
    print("输出尺寸:", output.size())
    # 输出尺寸: torch.Size([1, 10])

图表示异构架构

一个神经图表示的主要优势在于,它可以直观地表示多种不同的网络架构,并且都可以由相同的图神经网络进行处理。这使得我们无需对每种架构进行特定的调整,就能在不同的网络架构之间实现统一的学习和推理。

非线性变换

在神经网络中,非线性变换是通过激活函数(如ReLU、Sigmoid等)实现的。论文建议将这些非线性变换作为节点特征进行编码,具体方法是为一组常见的激活函数创建嵌入,并将它们添加到节点特征中。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
import torch
import torch.nn.functional as F

class ActivationEmbedding(torch.nn.Module):
ACTIVATION_FN = [
"none",
"relu",
"gelu",
"silu",
"tanh",
"sigmoid",
"leaky_relu",
]

def __init__(self, embedding_dim):
super().__init__()
self.activation_idx = {k: i for i, k in enumerate(self.ACTIVATION_FN)}
self.idx_activation = {i: k for i, k in enumerate(self.ACTIVATION_FN)}
self.embedding = torch.nn.Embedding(len(self.ACTIVATION_FN), embedding_dim)

def forward(self, activations, layer_layout, device):
indices = torch.tensor(
[self.activation_idx[act] for act in activations],
device=device,
dtype=torch.long,
)
emb = self.embedding(indices)
emb = emb.repeat_interleave(
torch.tensor(layer_layout[1:-1], device=device), dim=0
)
emb = F.pad(emb, (0, 0, layer_layout[0], layer_layout[-1]))
return emb

# 创建激活嵌入实例
embedding_dim = 4 # 假设嵌入维度为4
activation_embedding = ActivationEmbedding(embedding_dim)

# 示例输入
activations = ["relu", "gelu", "sigmoid"]
layer_layout = [2, 3, 2] # 表示第一层有2个节点,第二层有3个,第三层有2个
device = torch.device("cpu")

# 前向传播
emb_output = activation_embedding.forward(activations, layer_layout, device)

print("激活函数嵌入向量:")
print(emb_output)

# 激活函数嵌入向量:
# tensor([[ 0.0000, 0.0000, 0.0000, 0.0000],
# [ 0.0000, 0.0000, 0.0000, 0.0000],
# [-0.7482, 1.7909, -0.1332, -0.3995],
# [-0.7482, 1.7909, -0.1332, -0.3995],
# [-0.7482, 1.7909, -0.1332, -0.3995],
# [ 0.7152, -0.6278, -0.2240, -0.1032],
# [ 0.7152, -0.6278, -0.2240, -0.1032],
# [ 0.7152, -0.6278, -0.2240, -0.1032],
# [ 0.6765, -1.4103, 1.7681, 1.2673],
# [ 0.6765, -1.4103, 1.7681, 1.2673],
# [ 0.6765, -1.4103, 1.7681, 1.2673],
# [ 0.0000, 0.0000, 0.0000, 0.0000],
# [ 0.0000, 0.0000, 0.0000, 0.0000]],

残差连接

残差连接(Residual Connections)是现代CNN架构中的一个重要组成部分。残差连接直接将某一层的输入连接到其输出,通过公式 来实现。在神经图中,我们可以通过在神经图架构中包含从 的边来直接表示残差连接。论文中提到,残差连接可以重写为 ,其中 是单位矩阵,边特征对于每个连接的神经元为1。

归一化层

转换方式

归一化层(如BatchNorm和LayerNorm)用于标准化输入特征,以加速训练并提高模型的稳定性。它们通常表示为:

其中:

  • 是输入特征向量,维度为
  • 是可学习的缩放参数(scale),维度也是 ,用于控制输入特征的幅度。
  • 是可学习的偏移参数(bias),维度为 ,用于控制特征的偏移。

这里的 表示逐元素乘法。我们可以将上面的表达式重写为类似线性层的形式:

在这个表示中:

  • 是一个对角矩阵,它的对角线元素是 中的各个元素。因为 维的,所以这个对角矩阵的维度是
  • 对于每一个输入特征 ,输出特征 ,表示为从输入 到输出 的线性变换。

这个重写形式将归一化层看作是具有对角权重矩阵和偏移量的线性层。具体来说,对角矩阵 表示对输入特征进行逐元素的缩放,而 则表示逐元素的偏移。

神经图表示

在神经图中,我们可以这样表示归一化层:

  • 节点表示:
    • 输入节点: 的每一维作为一个节点,总共有 个输入节点。
    • 输出节点: 的每一维作为一个节点,总共有 个输出节点。
  • 边的特征:
    • 从每个输入节点 到对应输出节点 的边特征表示 ,即一个缩放因子。
    • 边特征的维度为 ,每条边的特征为单一标量值
  • 节点特征:
    • 额外的输出节点特征用于表示偏移项 ,每个输出节点都有一个偏移量。

通过这种表示方法,我们可以在神经图中直观地捕捉到归一化层的运算,保持图的结构完整性,并且支持进一步的图操作和学习。

Transformers的图表示

Transformers依赖于多头自注意力机制(Multi-Head Self-Attention),包括线性投影、自注意力计算、头的连接和最终的线性投影。

多头自注意力机制的步骤
  1. 线性投影: 输入 通过 个独立的线性层进行投影,产生查询()、键()和值():

    • 输入维度: ,其中 是输入序列的长度, 是输入的特征维度。
    • 输出维度(每个头): ,其中
  2. 点积注意力: 每个头的输出通过点积计算注意力得分,并进行 Softmax 激活处理:

    • 输出维度(每个头):
  3. 头的连接与线性投影: 将所有头的输出连接后,再通过一个线性层进行投影,得到最终输出:

    • 连接后维度:
    • 最终输出维度:
神经图表示
  • 节点表示:
    • 输入节点: 对于每个输入维度添加 个节点,用于表示输入特征的每一维。
    • 注意力头输出节点: 对于每个头的每个维度,添加 个节点。
    • 输出节点: 对于每个输出维度,添加 个节点。
  • 边的特征:
    • 线性投影边特征: 三种不同的线性投影()通过多维边特征表示。例如,针对每个注意力头
    • 边特征维度: 每个边特征维度为
  • 注意力计算:
    • 注意力计算是无参数操作,不需要在神经图中显式建模,由神经图网络自行近似。
  • 头的连接与线性投影:
    • 所有头的连接通过连接适当的节点,使用输出权重 映射到相应的输出节点。
    • 输出节点的边特征: 被视为标准的线性层边特征。
总结

在神经图中,Transformers的多头自注意力机制通过节点和边特征表示如下: - 输入节点数: - 注意力头输出节点数: - 输出节点数: - 线性投影边特征数: (每种投影 个)

节点和边的表示

神经图表示法为我们提供了灵活性,可以选择哪些数据作为节点和边的特征。虽然我们主要关注权重和偏置,但在某些情况下,也可以使用其他特征,如梯度。

  1. 边的方向:基础编码只考虑神经网络的前向传播,形成一个有向无环图(DAG)我们可以在神经图中添加反向边。例如,包含 作为额外的边特征。同样,可以将 作为无向特征。

    添加反向边(reverse edges)主要有以下几个原因:

    1. 增强信息流动性

    在传统的神经网络前向传播过程中,信息从输入层流向输出层,构成一个有向无环图(Directed Acyclic Graph, DAG)。这种结构限制了信息只能沿着一个方向流动,通常是从低层到高层。通过添加反向边,神经图允许信息在图中的不同层之间进行双向流动。这种信息的双向流动可以增强图神经网络捕捉复杂关系和相互依赖性的能力。例如,反向边可以帮助模型理解上下文中的逆向关系,或者使信息从输出节点“反馈”到输入节点。

    1. 改进特征传播和学习效率

    在没有反向边的图神经网络中,节点只能接收到来自前序节点的消息。反向边允许每个节点接收到来自后续节点的消息,从而提供更多的上下文信息。这种双向信息流动有助于更高效地传播特征,使得每个节点能够综合考虑来自其前后节点的所有信息,改进模型的学习和推理能力。

    1. 提高模型的表达能力

    反向边可以提升模型对图的表示能力(representation power)。通过添加反向边,图神经网络可以捕捉更复杂的关系,这些关系在仅有单向边的情况下可能会被忽略。例如,在社交网络分析中,用户A与用户B的互相关注可以被更准确地建模;在知识图谱中,双向关系(如“父母”和“子女”)也能够被更完整地表示。

    1. 应对反向传播中的梯度流失

    在神经网络的训练过程中,反向传播是更新模型参数的关键步骤。通过添加反向边,神经图可以帮助在反向传播过程中更好地维护梯度信息。这对于深度网络尤为重要,因为它们容易出现梯度消失或梯度爆炸的问题。反向边的引入可以促进梯度的稳定传播,提高训练的稳定性和效率。

    1. 实现无向图的表示

    在许多应用场景中,图可能是无向的(如社交网络、化学分子结构等)。为了在神经图中准确表示这些无向图,可以通过添加反向边来实现无向边的效果。这种方法使得每条边都可以被视为双向的,确保了模型能够对图的无向性质进行正确的建模和处理。

  2. 探测特征(Probe Features):探测特征的灵感来自于人类如何理解复杂的函数。在人类分析复杂的数学函数时,通常会选择一些输入样本,并观察这些样本通过函数后的输出,以此来推测和理解函数的行为。类似地,在神经图中,我们通过添加探测特征,使图神经网络能够“探测”神经网络的行为。要生成探测特征:

    1. 首先我们选择一组特定的输入样本,这些输入样本可以是训练集中具有代表性的样本,或者是通过某种方式生成的样本
    2. 接下来,我们将这些样本输入到待分析的神经网络中,记录每一层的中间激活值以及最终的输出
    3. 然后,我们将这些激活值作为额外的节点特征,添加到神经图的相应节点中

    在神经图中,探测特征可以理解为特定输入样本经过神经网络的处理后,在各层之间传播的中间激活值。这些激活值作为额外的节点特征被添加到神经图中。通过这些额外的特征,图神经网络可以在训练过程中学习到更多关于神经网络各层之间关系的有用信息。具体而言,我们学习一组样本输入值,通过输入神经网络并保留所有中间激活值和输出值。例如,对于简单的输入神经网络:

    我们为每个 获取一个额外的节点特征:

    这些特征作为附加的节点特征被包括进来。探测特征为每个节点提供了更多的上下文信息,使得图神经网络在学习过程中可以结合更多的信息,从而更好地学习到神经网络结构和参数之间的关系。

    值得注意的是,探测特征对于所有神经网络参数的增强是保持不变的,只要它们保持相同的输出和隐藏层功能。即使神经网络的参数(如权重和偏置)发生变化,只要这些变化不会影响神经网络在相同输入下的输出和中间层的激活值(即网络的功能保持不变),探测特征的值也不会变化。换句话说,这种不变性是因为探测特征关注的是网络的功能表现,而不是具体的参数值。只要网络对给定的输入仍然产生相同的输出和中间激活,探测特征就会保持一致。

    代码大致思路为:

    1. 模型实例化和参数处理:使用给定的模型配置实例化一个神经网络,并将其转换为函数式模型(Functional Model),以便将模型的参数作为输入。
    2. 探测输入初始化:根据需要初始化探测输入,要么使用提供的初始值,要么随机生成。
    3. 特征生成:通过对模型的参数和探测输入进行处理,生成特征。这个过程包括将参数展平并使用向量化后的模型计算输出特征。
    4. 特征投影和归一化(可选):如果指定了投影维度,对生成的特征进行线性投影和层归一化处理。
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    32
    33
    34
    35
    36
    37
    38
    39
    40
    41
    42
    43
    44
    45
    46
    47
    48
    49
    50
    51
    52
    53
    54
    55
    56
    57
    58
    59
    60
    61
    62
    63
    64
    65
    66
    67
    68
    69
    70
    import hydra
    import torch
    import torch.nn as nn
    from einops.layers.torch import Rearrange
    from nn.inr import make_functional, params_to_tensor, wrap_func

    class GraphProbeFeatures(nn.Module):
    def __init__(self, d_in, num_inputs, inr_model, input_init=None, proj_dim=None):
    super().__init__()

    # 使用 Hydra 实例化给定的 INR(隐式神经表示)模型
    inr = hydra.utils.instantiate(inr_model)

    # 将模型转换为函数式模型,便于参数处理
    fmodel, params = make_functional(inr)

    # 将模型参数转换为单个张量和形状列表
    vparams, vshapes = params_to_tensor(params)

    # 使用 vmap 向量化函数模型以便进行批量操作
    self.sirens = torch.vmap(wrap_func(fmodel, vshapes))

    # 如果未提供输入初始化,随机生成探测输入
    inputs = (
    input_init
    if input_init is not None
    else 2 * torch.rand(1, num_inputs, d_in) - 1 # 范围在 [-1, 1] 之间
    )
    # 将输入注册为参数,并设置是否可训练
    self.inputs = nn.Parameter(inputs, requires_grad=input_init is None)

    # 定义用于重排权重和偏置的 einops 操作
    self.reshape_weights = Rearrange("b i o 1 -> b (o i)")
    self.reshape_biases = Rearrange("b o 1 -> b o")

    # 如果指定了投影维度,初始化投影层
    self.proj_dim = proj_dim
    if proj_dim is not None:
    self.proj = nn.ModuleList(
    [
    nn.Sequential(
    nn.Linear(num_inputs, proj_dim), # 线性变换
    nn.LayerNorm(proj_dim), # 层归一化
    )
    for _ in range(inr.num_layers + 1) # 对每层执行投影
    ]
    )

    def forward(self, weights, biases):
    # 将权重和偏置张量展平
    weights = [self.reshape_weights(w) for w in weights]
    biases = [self.reshape_biases(b) for b in biases]

    # 将展平的权重和偏置连接成一个大张量
    params_flat = torch.cat(
    [w_or_b for p in zip(weights, biases) for w_or_b in p], dim=-1
    )

    # 使用探测输入和向量化模型计算输出特征
    out = self.sirens(params_flat, self.inputs.expand(params_flat.shape[0], -1, -1))

    # 如果指定了投影维度,对特征进行投影和归一化
    if self.proj_dim is not None:
    out = [proj(out[i].permute(0, 2, 1)) for i, proj in enumerate(self.proj)]
    out = torch.cat(out, dim=1)
    return out
    else:
    # 否则直接返回展平的特征
    out = torch.cat(out, dim=-1)
    return out.permute(0, 2, 1)
  3. 归一化:现有的关于参数空间网络的研究中,对特征归一化的处理是分别计算训练集中每个神经元的均值和标准差。然而,这种操作会破坏神经元对称性,因为神经元可以进行置换。因此,标准化神经网络的神经元是没有意义的。论文提出了一种简单的替代方案:对每一层的所有神经元的权重和偏置计算一个共同的均值和标准差,而不是对每个神经元单独计算。这意味着每一层有一个单一的均值 和标准差 用于权重,以及一个单一的均值 和标准差 用于偏置。 这种方法的效果在于:

    • 保证神经元对称性:同一层中的神经元是对称的,标准化操作不会打破这种对称性。
    • 简化归一化计算:减少了归一化计算的复杂性,不需要为每个神经元单独计算均值和标准差,只需为每一层计算一次。
  4. 位置嵌入(Positional Embeddings):在处理神经图之前,我们为每个节点增加了学习到的位置嵌入。为了保持隐藏层中的置换对称性,对应于同一中间层的节点共享相同的位置嵌入。虽然这些层信息在邻接矩阵中已隐含存在,但使用位置嵌入可以立即识别层信息,而无需多个局部消息传递步骤。然而,我们区分了输入和输出节点的对称性与隐藏节点的对称性。在神经网络中,重新排列输入或输出节点通常会改变网络的基础功能。为了解决这种差异,我们为每个输入和输出节点引入唯一的位置嵌入,从而打破它们之间的对称性,使得GNN和Transformer能够区分它们。具体来说,位置嵌入的作用是:

    为每个节点增加位置信息:在神经图中,每个节点代表神经网络中的一个特征或一个操作。为了让模型更好地理解这些节点在网络中的位置,我们为每个节点增加了一个位置嵌入(Positional Embedding)。位置嵌入可以看作是节点的“位置标签”,它告诉模型这个节点在哪一层,并帮助模型理解不同层次的节点如何相互作用。

    保持隐藏层的置换对称性:对于同一层中的节点(比如隐藏层中的神经元),它们共享相同的位置嵌入。这是为了保持它们之间的置换对称性,也就是说,不管这些节点的顺序如何,它们的功能是等价的。这样做的好处是,即使节点顺序发生变化,模型依然能够正确地理解和处理它们,因为它知道这些节点来自同一层。

    立即识别层信息:虽然层信息可以从邻接矩阵中隐含地得到,但使用位置嵌入可以使模型立即知道节点所在的层次,而不需要经过多个步骤的消息传递来推断。这样可以加快模型的理解和计算速度。

    区分输入、输出节点与隐藏节点的对称性:在神经网络中,输入节点和输出节点的排列顺序通常会影响网络的功能,因为改变输入或输出节点的顺序相当于改变网络的输入输出映射。为了解决这一问题,位置嵌入为每个输入和输出节点提供了独特的标识。这种做法打破了输入和输出节点之间的对称性,使得模型能够明确区分它们,并正确处理它们的相对位置。

学习神经图

图神经网络(GNN)和 Transformer 在图的排列对称性方面是等变的,作者提出了每种方法的一种变体,并对其进行调整以处理神经图,具体如下:

  1. GNN(图神经网络):GNN是一种以消息传递神经网络的形式应用的图神经网络,它在每个节点上应用相同的局部消息传递函数。在各种GNN变体中,很少有设计用于支持边特征的,更新隐藏层中的边特征的则更少。在本文的设定中,边特征是主要的特征,需要在每层中更新。因此作者选择了PNA(Principal Neighborhood Aggregation)作为基础,这是一个支持边特征的最先进的图网络。但是,PNA并不更新其边特征。为了解决这一问题,作者对其进行扩展,通过轻量级神经网络 在每层中更新边特征:

    其中, 是网络中的层索引。

    此外,为了增强消息传递步骤,我们将FiLM(Feature-wise Linear Modulation)应用于消息传递步骤,使其包含节点和边特征之间的乘性交互:

  2. Transformer:Transformer编码器可以看作是一个在全连接图上操作的图神经网络。与GNN类似,原始的Transformer编码器和常见的变体不支持边特征。作者使用了具有关系注意力的Transformer变体,它在自注意力计算中添加了边特征。与GNN一样,我们进一步增强Transformer,通过调制使节点和边特征之间的乘性交互成为可能。特别地,我们修改了自注意力模块中的值矩阵的更新:

通过这些方法,神经图表示可以灵活地编码和处理复杂的神经网络结构,并且在神经图网络和Transformer中有效地执行消息传递和特征更新。

若干问题

  1. 如何确保神经图的表示和操作在大规模网络中的计算效率和内存利用率

可能可以采取的策略:

  • 稀疏化处理: 在大规模网络中,许多神经图可能是稀疏的,因此可以考虑使用稀疏矩阵来表示神经图,从而减少内存使用。
  • 图压缩技术: 通过图压缩技术,如节点聚合、边压缩等,可以减少图的规模。例如,可以使用聚类方法将多个节点聚合成一个超节点,简化图结构,同时保留重要的拓扑信息。
  • 局部计算与特征重用: 在进行图神经网络计算时,可以考虑使用局部计算,即只对需要更新的节点和边进行计算,而不是对整个图进行全局计算。此外,重用之前计算的特征来减少重复计算,也可以提高计算效率。
  • 分布式计算: 对于非常大的图,可以采用分布式计算的方法,将图划分为多个子图,分别在多个计算节点上进行计算,然后合并计算结果。这种方法可以充分利用现代分布式计算框架的优势,处理大规模神经图。
  1. 神经图如何处理动态网络架构的变化,例如在神经架构搜索(NAS)中的应用?

神经架构搜索(Neural Architecture Search, NAS)是寻找最优神经网络架构的过程,它可以生成具有不同层次结构和连接的动态网络。神经图在固定架构下表现良好,但如何有效地适应动态变化的网络架构?研究如何将神经图表示与NAS相结合,探索在架构变化时自动调整图的表示和学习过程,是否需要引入动态图调整机制或自适应的图神经网络模型来应对这种挑战?可能可以采用的策略:

  • 动态图调整机制: 为了适应动态变化的网络架构,可以开发一种能够动态调整图结构的机制。例如,使用可微的图生成模型,这些模型能够在架构搜索过程中根据不同的结构生成对应的神经图。
  • 自适应图神经网络(Adaptive GNN): 可以设计一种自适应的图神经网络,它能够在训练过程中根据输入的架构自动调整自身的参数和结构。这样可以确保在处理不同架构时,GNN能够有效捕捉每个架构的特征和信息。
  • 层次化表示: 在NAS中,网络架构的变化通常表现为不同层次结构的组合和替换。使用层次化的图表示方法,可以更好地捕捉这些变化。例如,将神经网络的每一层或模块视为一个超节点,在图的高层次表示中进行操作。
  • 元学习(Meta-Learning)策略: 通过元学习策略,模型可以学习如何根据不同的架构调整其自身参数。这种策略可以帮助神经图在面对不断变化的架构时快速适应,并有效学习。
  1. 在多任务学习中,神经图表示如何应对不同任务的多样性和特异性?
  • 任务特定的子图: 在神经图中,可以为每个任务创建任务特定的子图。这些子图包含了与该任务相关的节点和边特征,并通过共享图结构的公共部分来捕捉任务间的共性信息。
  • 多层次表示: 可以使用多层次图表示来处理任务间的多样性。例如,底层子图用于捕捉任务的特定特征,而高层子图用于捕捉任务的共性特征。通过这种分层表示,可以更好地学习和泛化到多种任务。
  • 多任务学习的联合损失函数: 通过设计联合损失函数,可以在优化过程中权衡不同任务的需求。例如,可以通过加权策略来平衡任务间的损失,从而确保每个任务都能得到充分的训练。
  • 任务特定的边特征: 可以在神经图中添加任务特定的边特征,这些特征用于表示不同任务之间的关联关系。例如,在图的边特征中添加任务特定的信息,模型可以更好地学习不同任务之间的共享特征和独特特征。
  1. 神经图在处理具有噪声数据或不确定性数据的神经网络时,是否有独特的优势或挑战?

优势:

  • 鲁棒性增强: 通过引入边特征的不确定性建模,神经图可以更好地应对噪声数据。例如,可以使用概率图模型(如贝叶斯神经网络)来对边特征进行建模,从而自然地处理噪声和不确定性。
  • 多视角特征表示: 神经图能够捕捉数据的多视角特征,这些特征可能包含噪声和不确定性。通过聚合多视角信息,模型可以更好地过滤掉噪声数据,提取有用特征。
  • 自适应特征学习: 神经图具有自适应特征学习的能力,可以动态调整其特征表示,以应对不同的数据分布和不确定性。这使得神经图能够更好地处理噪声数据。

挑战:

  • 高维度不确定性建模: 在处理高维度数据时,对每个特征和边进行不确定性建模可能会增加计算复杂度和内存需求。
  • 噪声传播问题: 在图神经网络中,信息通过消息传递机制在节点和边之间传播。如果数据中存在噪声,噪声可能会在图中传播,导致误差累积。因此,设计一个能有效抑制噪声传播的机制是一个重要挑战。
  • 缺乏标注的监督: 在处理带有噪声的数据时,缺乏标注的监督信息可能会使得神经图难以区分噪声和有用信号。需要设计有效的无监督学习策略,以便在缺乏监督信息的情况下,仍能从噪声数据中学习有用特征。
  1. 神经图的等变性在处理跨领域迁移学习任务时,有何独特贡献和局限性?

贡献:

  • 保持特征一致性: 等变性确保了模型对输入的排列和变换保持不变,这在迁移学习中非常有用。例如,当从一个领域迁移到另一个领域时,等变性可以帮助模型保持对特征的理解和处理方式的一致性,即使输入数据的排列方式发生了变化。
  • 提高泛化能力: 通过等变性,模型可以学习到更通用的特征表示,这些表示不仅适用于特定领域的数据,还能在其他领域中表现良好。这有助于模型在跨领域迁移学习中快速适应新领域的数据分布。

局限性:

  • 对领域差异的适应性不足: 等变性可能会忽略不同领域之间的显著差异。在某些情况下,两个领域的数据特征和关系结构可能完全不同,此时单纯的等变性可能无法捕捉这些领域特异的特征。这可能导致模型在迁移到新领域时,无法充分利用新领域的特定信息。
  • 潜在的过拟合风险: 如果等变性过于严格,模型可能会对训练数据中的排列和变换过于敏感,导致过拟合。在跨领域迁移时,这种过拟合可能会限制模型的泛化能力,因为它过于依赖原始领域的数据特征,而无法充分适应新领域的数据。
  • Post title:Graph Neural Networks for Learning Equivariant Representations of Neural Networks
  • Post author:Chen Kai
  • Create time:2024-08-30 09:00:00
  • Post link:https://www.chenk.top/Graph Neural Networks for Learning Equivariant Representations of Neural Networks/
  • Copyright Notice:All articles in this blog are licensed under BY-NC-SA unless stating additionally.
 Comments