Message passing all the way up

paper Link 1 Link 2

A message passing GNN over this graph

\[\mathbf{h}_u = \phi \left( \mathbf{x}_u, \bigoplus_{v∈\mathcal{N}_u} \psi(\mathbf{x}_u, \mathbf{x}_v) \right) \tag{1}\]
  • Graph: \(\mathcal{G}\)

  • Nodes: \(\mathcal{V}\)

  • Edges: \(\mathcal{E}\)

  • one-hop neighbourhoods: \(\mathcal{N}_u \{ v ∈ \mathcal{V} | (v, u) ∈ \mathcal{E} \}\)

  • node feature matrix: \(\mathbf{X} ∈ \mathbb{R}^{|\mathcal{V}|×k}\)

  • features of node u: \(\mathbf{x}_u\)

  • number of node features: \(k\)

  • message function: \(\psi : \mathbb{R}^k ×\mathbb{R}^k \to \mathbb{R}^l\) can be MLPs

  • new size of node features: \(l\)

  • readout function: \(\phi : \mathbb{R}^k ×\mathbb{R}^l \to \mathbb{R}^m\) can be MLPs

  • new size of message features: \(m\)

  • permutation-invariant aggregation function: \(\bigoplus\)

Simple example: Master nodes

master node, \(µ\), connecting it with all other nodes, and then performing message passing as usual. Mathematically, \(\mathcal{V}' = \mathcal{V} ∪ \{µ\}, \mathcal{N}'_u = \mathcal{N}_u ∪ \{µ\}, \text{ and }\mathcal{N}'_µ = V\).

Broadly, augmented message passing techniques can be categorised into six distinct types.

Feature augmentation

once the features are computed, the GNN computations proceed as before—hence, trivially expressible using message passing.

examples:

  • a one-hot encoding (Murphy et al., 2019) for detecting patterns

  • a random feature (Sato et al., 2021) for detecting patterns

  • count subgraphs of interest (Bouritsas et al., 2020)

  • provide the graph’s Laplacian eigenvectors (Dwivedi & Bresson, 2020)

Message passing modulation

non-binding: the model is not forced to use these features

while the message function, \(ψ\), is modified, the blueprint of Equation 1 remains, and this case is also trivially an instance of message passing.

examples

  • directional graph networks (Beaini et al., 2021)

    • define“flows”on a graph

    • guide how various incoming messages are scaled

  • LSPE (Dwivedi et al., 2021) computed positional features

Graph rewiring

many methods modify the edges of the input graph to compensate. Such graph rewiring methods leave V unchanged, but make direct changes to \(\mathcal{N}_u\). The underlying message function setup of Equation 1 is unchanged, hence these methods are still expressible using message passing.

examples

  • fully connected graph, i.e. setting \(\mathcal{N}_u = \mathcal{V}\).

    • graph Transformers (Ying et al., 2021; Kreuzer et al., 2021; Mialon et al., 2021)

    • graph Fourier transform (Bruna et al., 2013) spectrally defined graph convolutions

  • Nontrivial changes to \(\mathcal{N}_u\)

    • multi-hop layers (Defferrard et al.,2016)

    • rewiring based on diffusion (Klicpera et al., 2019)

    • curvature (Topping et al., 2021)

    • subsampling (Hamilton et al., 2017)

  • alter the adjacency in a learnable fashion

    • (Kipf et al., 2018; Wang et al., 2019; Kazi et al., 2020; Velickovi ˇ c et al., 2020)

Subgraph aggregation

An extension of graph rewiring methods.

replicating the nodes for every subgraph of interest, and connecting the copies of every node together.

Mathematically, if we are learning over K subgraphs, let \(\mathcal{V}' = \mathcal{V}' × \{1, 2, . . . , K\}\) and, assuming that subgraph \(i\)’s edges define neighbourhoods \(\mathcal{N}^{(i)}_u\) , we can redefine neigbourhoods as follows: \(\mathcal{N}'_{u,i} = \{(v, i) | v ∈ \mathcal{N}^{(i)}_u \} ∪ \{(u, j) | j ∈ \{1, 2, . . . , k\}\}\).

examples

  • Papp et al. (2021); Cotta et al. (2021); Zhao et al. (2021); Bevilacqua et al. (2021)

Substructure based methods

naturally-occurring phenomena are sometimes best described by interactions of groups of entities. an equivalent architecture using only the pairwise message passing primitive.

  • the functional groups in a molecule can often strongly influence its properties (Duvenaud et al., 2015)

  • methods have been developed to support computing representations

    • junction trees (Fey et al., 2020)

    • spectral signals (Stachenfeld et al., 2020)

    • simplicial complexes (Bodnar et al., 2021b)

    • cellular complexes (Bodnar et al., 2021a; Hajij et al., 2020)

    • general k-tuples of nodes (Morris et al., 2019; 2020)

    • hypergraphs (Huang & Yang, 2021; Chien et al., 2021; Georgiev et al., 2022)

junction trees: also known as Tree decomposition, a mapping of a graph into a tree that can be used to define the tree width of the graph and speed up solving certain computational problems on the graph.

if the groups of interest are greater than two nodes, over the original graph, it can be provably impossible to use pairwise messaging to simulate such interactions (Neuhauser et al., 2021). But what about modifing the graph structure?

The trick is to create new nodes for every substructure that we want to model, and appropriately connecting them to their constituent nodes. How easy this is to do depends on whether the functions of interest are permutation invariant.

establish bidirectional edges between the constituent nodes and the corresponding “substructure node”. Mathematically, we assume that we have \(\mathit{K}\) substructures, \(\mathcal{S}_1, \mathcal{S}_2, . . . , \mathcal{S}_\mathit{K}\) each operating in a permutation invariant way over its constituent nodes \((\mathcal{S}_i ⊆ \mathcal{V})\). Then, we augment the graph by creating new substructure nodes: \(\mathcal{V}' = \mathcal{V} ∪ \{µ_1, µ_2, . . . , µ_\mathit{K}\},\) and modifying the neighbourhoods to connect every constituent to its substructure(s): \(\mathcal{N}'_u = \mathcal{N}_u ∪ \{µ_i| u ∈ \mathcal{S}_i\}\), and \(\mathcal{N}'_{µi} = \mathcal{S}_i\). Note that this is a more general case of the master node approach of Section 2, which can be seen as having just one substructure, \(\mathcal{S}_1 = \mathcal{V}\).

When interactions within a substructure are permutation sensitive, a more intricate gadget is required.

examples

  • create \(\mathcal{O}(\mathcal{S}_i)\) new nodes to process nodes’ features one at a time acccording to the permutation (not unlike a long short-term memory (Hochreiter & Schmidhuber, 1997))

  • using carefully constructed message functions to materialise a concatenation of the inputs which respects the permutation.

General equivariant GNNs

characterise all possible linear permutation-equivariant layers over an input graph, and use them as a basis for building equivariant GNNs (Maron et al., 2018)

Maron et al. (2018) discover

  • 2 linear invariant layers

  • 15 linear equivariant layers

represented as matrices in \(\mathbb{R}^{\mathcal{V}^2×\mathcal{V}^2}\), which get multiplied with edge features.

Mathematically, assume we want to multiply with a basis matrix \(\mathbf{B} ∈ \mathbb{R}^{\mathcal{V}^2×\mathcal{V}^2}\). Inventing new edge-based nodes corresponds to \(\mathcal{V}' = \mathcal{V} ∪ \{e_{uv} | u, v ∈ \mathcal{V}\}\). Then, we need to connect these edges to the nodes incident to them, and also to other edges, whenever the entries of \(\mathbf{B}\) mandate it. Hence the neighbourhoods update as follows: \(\mathcal{N}'_u = \mathcal{N}_u ∪ \{e_{ab} | a = u ∨ b = u\}\) for nodes, and \(\mathcal{N}'_{e_{uv}} = \{u, v\} ∪ \{e_{ab} | \mathbf{B}_{(a,b),(u,v)} \neq 0\}\) for edges. Similar gadgets (potentially “tensored up” to k-tuple nodes) would also apply for follow-up work, including but not limited to Maron et al. (2019); Keriven & Peyre (2019); Albooyeh et al. (2019); Azizian & Lelarge (2020).

General equivariant GNNs

이전 섹션에서는 모두 GNN 표현력을 향상시킬 수 있는 특정 계산이나 모티프를 식별하려고 했지만, 그 반대 접근 방식은 입력 그래프에서 가능한 모든 선형 순열 등변수 레이어를 특성화하고 이를 등변수 GNN을 구축하기 위한 기초로 사용하는 것입니다(Maron et al., 2018). 이미지 데이터에 대한 유사한 분석에 따르면 이미지에는 정확히 한 가지 유형의 선형 변환 등변수 레이어, 즉 컨볼루션이 존재한다는 사실이 밝혀졌습니다(Bronstein et al., 2021).

이 프레임워크를 사용하여 Maron 등(2018)은 에지 값 입력이 있는 설정에 대해 2개의 선형 불변층과 15개의 선형 등변층으로 구성된 기초를 발견했습니다(노드의 k-튜플에 대해 정의된 데이터의 경우 차원은 k번째 및 2k번째 벨 번호로 정의됨). 이 15개의 레이어는 그래프 구조의 대칭성을 유지하면서 에지 쌍에 걸쳐 정보를 효과적으로 재조합할 수 있게 해줍니다. 따라서 이들은 에지 피처를 곱한 \(\mathbb{R}^{\mathcal{V}^2×\mathcal{V}^2}\)의 행렬로 표현됩니다.

이러한 접근 방식은 명백한 텐서럴 의미론에도 불구하고 쌍방향 메시지 전달 언어로 표현할 수 있습니다. 모든 정사각형 행렬 곱셈 연산은 해당 행렬의 0이 아닌 항목이 암시하는 그래프에 대한 컨볼루션 GNN 인스턴스를 나타냅니다(Bronstein et al., 2021). 핵심적인 차이점은 이 경우 메시지가 노드가 아닌 에지를 통해 전달되므로 에지에 해당하는 새로운 노드를 생성하고 그에 따라 연결해야 한다는 것입니다!

수학적으로, 기본(basis) 행렬 \(\mathbf{B} ∈ \mathbb{R}^{\mathcal{V}^2×\mathcal{V}^2}\)를 곱하고 싶다고 가정해봅시다. 새로운 에지 기반 노드를 발명하면 \(\mathcal{V}' = \mathcal{V} ∪ \{e_{uv} | u, v ∈ \mathcal{V}\}\)에 해당합니다. 그런 다음 이러한 에지를 해당 에지와 접하는 노드에 연결해야 하며, \(\mathbf{B}\)의 항목이 요구할 때마다 다른 에지에도 연결해야 합니다. 따라서 이웃은 다음과 같이 업데이트됩니다: 노드의 경우 \(\mathcal{N}'_u = \mathcal{N}_u ∪ \{e_{ab} | a = u ∨ b = u\}\), 에지의 경우 \(\mathcal{N}'_{e_{uv}} = \{u, v\} ∪ \{e_{ab} | \mathbf{B}_{(a,b),(u,v)} \neq 0\}\) 입니다. 후속 작업에도 유사한 가젯(잠재적으로 K-튜플 노드로 “텐서드 업tensored up”)이 적용될 수 있으며, 여기에는 Maron et al. (2019), Keriven & Peyre (2019), Albooyeh et al. (2019), Azizian & Lelarge (2020) 등이 포함되나 이에 국한되지 않습니다.

이 특별한 경우, 메시지 전달의 의미론은 텐서럴(tensorial) 접근 방식과 잘 맞지 않습니다. 하지만 에지를 노드로 표현하면 메시지 전달을 직접적으로 보강하는 여러 제안, 특히 선 그래프(line graph)와 관련된 개념과의 연결이 드러납니다(Monti et al., 2018; Cai et al., 2021). 따라서 메시지 전달로 전환하는 것이 가장 실용적이지 않더라도 서로 관련이 없어 보이는 제안들 사이에서 놀라운 연관성을 발견하고 향후 연구를 자극하는 유용한 연습이 될 수 있습니다.

구성(constituent) 노드와 대응하는 하위구조(substructure) 노드 사이에 양방향 에지를 설정하면 됩니다. 수학적으로는, k개 서브구조 \(\mathcal{S}_1, \mathcal{S}_2, . . . , \mathcal{S}_\mathit{K}\) 가 각각 구성노드 \((\mathcal{S}_i ⊆ \mathcal{V})\) 에 대해 순열 불변(permutation invariant) 방식으로 작동한다고 가정합니다.

그런 다음 새로운 하위 구조 노드를 생성하여 그래프를 보강합니다.

\(\mathcal{V}' = \mathcal{V} ∪ \{µ_1, µ_2, . . . , µ_\mathit{K}\},\)

그리고, 모든 구성(constituent) 요소를 해당 구성 요소의 하위 구조에 연결하도록 이웃을 수정합니다.

\(\mathcal{N}'_u = \mathcal{N}_u ∪ \{µ_i| u ∈ \mathcal{S}_i\}\), and \(\mathcal{N}'_{µi} = \mathcal{S}_i\)

이는 섹션 2의 마스터 노드 접근법의 보다 일반적인 경우로, \(\mathcal{S}_1 = \mathcal{V}\) 라는 하나의 하위 구조만 있는 것으로 볼 수 있습니다.

Appendicts A

노드의 k-튜플에 대한 함수 f(x1, x2, … , xk)를 모델링하고, 이 함수가 처음 n개의 인자에서는 순열에 민감하고 나머지 m개에서는 순열 불변이라고 가정해 봅시다(단, n + m = k). 이전과 마찬가지로 메인 논문에서와 같은 방식으로 m 개의 불변 파라미터를 맞출 수 있습니다(해당 m 개의 노드와 메시지를 교환하는 새로운 “마스터 노드”인 µ를 발명). 이렇게 하면 방정식 1을 다시 사용해 수정된 함수를 얻을 수 있습니다:

2

여기에서 마스터 노드의 피처를 설정할 수 있으며, 사용할 수 있는 피처가 없는 경우 초기에는 xµ를 0 벡터로 설정할 수 있습니다.

이제 순서 의존적인 방식으로 이 모든 n+1 입력을 g로 처리해야 합니다. 이를 수행하는 가장 간단한 방법은 “concat-aggregation”을 사용하는 새 노드를 만드는 것이지만, 이것이 방정식 1에서 명시적으로 지원되는지 여부는 불분명합니다. 각 xi 벡터를 결과 벡터의 별도 “슬롯”에 복사하는 메시지 함수를 준비한 다음 합계 집계를 적용하여 이 연결을 구체화하는 방법이 존재합니다. 그런 다음 연결된 노드는 결과를 저장하는 노드에 자신의 특징을 메시지로 전송할 수 있습니다.

λi = 0으로 초기화(또는 학습 가능으로 설정)한 다음 다음과 같이 각각을 업데이트합니다:

3

방정식 1의 메시지 전달 프레임워크와 일치합니다 (Nλi = {λi-1, i}, 특히 Nλn+1 = {λn, µ}로 설정). 또한 이러한 방정식은 순환 신경망과도 일치합니다(예: φ2와 ψ2의 특수한 선택에 대해 LSTM을 복구할 수 있음). 이러한 모델의 n + 1 단계가 수행되면 모든 특징이 전파되기에 충분한 시간이 경과한 것이며, λn+1을 목표 함수 f의 최종 표현으로 사용하여 해당 출력을 관련 노드에 다시 공급할 수 있습니다.