2005年，劃時(shí)代之作「The Graph Neural Network Model」的問(wèn)世，將圖神經(jīng)網(wǎng)絡(luò)帶到每個(gè)人面前。

在此之前，科學(xué)家處理圖數(shù)據(jù)的方式是，在數(shù)據(jù)預(yù)處理階段，將圖轉(zhuǎn)換為一組「向量表示」。

而CNN的出現(xiàn)徹底改變這種信息丟失的弊端，近20年來(lái)，一代又一代模型不斷演變，推動(dòng)ML領(lǐng)域進(jìn)步。

今天，谷歌正式官宣發(fā)布TensorFlow GNN 1.0（TF-GNN）——用于大規(guī)模構(gòu)建GNN的經(jīng)過(guò)生產(chǎn)測(cè)試的庫(kù)。

它既支持在TensorFlow中的建模和訓(xùn)練，也支持從大型數(shù)據(jù)存儲(chǔ)中提取輸入圖。

TF-GNN是專為異構(gòu)圖從頭開始構(gòu)建的，其中對(duì)象和關(guān)系的類型由不同的節(jié)點(diǎn)和邊集合來(lái)表示。

現(xiàn)實(shí)世界中的對(duì)象及其關(guān)系以不同的類型出現(xiàn)，而TF-GNN的異構(gòu)焦點(diǎn)，使得表示它們變得非常自然。

谷歌科學(xué)家Anton Tsitsulin表示，復(fù)雜的異構(gòu)建模又回來(lái)了！

TF-GNN 1.0首面世

對(duì)象及其相互之間的關(guān)系，在我們的世界中無(wú)處不在。

而關(guān)系對(duì)于理解一個(gè)對(duì)象的重要性，不亞于孤立地看待對(duì)象本身的屬性，比如交通網(wǎng)絡(luò)、生產(chǎn)網(wǎng)絡(luò)、知識(shí)圖譜或社交網(wǎng)絡(luò)。

離散數(shù)學(xué)和計(jì)算機(jī)科學(xué)長(zhǎng)期以來(lái)一直將這類網(wǎng)絡(luò)形式化為圖，由「節(jié)點(diǎn)」以各種不規(guī)則方式通過(guò)邊任意連接而成。

然而，大多數(shù)機(jī)器學(xué)習(xí)算法只允許輸入對(duì)象之間存在規(guī)則統(tǒng)一的關(guān)系，如像素網(wǎng)格、單詞序列，或完全沒(méi)有關(guān)系。

圖形神經(jīng)網(wǎng)絡(luò)，簡(jiǎn)稱GNN，是一種強(qiáng)大的技術(shù)，既能利用圖的連通性（如早期算法DeepWalk和Node2Vec），又能利用不同節(jié)點(diǎn)和邊輸入特征。

GNN可以對(duì)圖的整體（這種分子是否以某種方式做出反應(yīng)？）、單個(gè)節(jié)點(diǎn)（根據(jù)引用，這份文檔的主題是什么？）、潛在的邊（這種產(chǎn)品是否可能與另一種產(chǎn)品一起購(gòu)買？）進(jìn)行預(yù)測(cè)。

除了對(duì)圖形進(jìn)行預(yù)測(cè)之外，GNN還是一個(gè)強(qiáng)大的工具——用于彌合與更典型的神經(jīng)網(wǎng)絡(luò)用例之間的鴻溝。

它們以連續(xù)的方式對(duì)圖的離散關(guān)系信息進(jìn)行編碼，從而可以將其自然地納入另一個(gè)深度學(xué)習(xí)系統(tǒng)。

谷歌在今天正式宣布用于大規(guī)模構(gòu)建GNN的經(jīng)過(guò)生產(chǎn)測(cè)試的庫(kù)——TensorFlow GNN 1.0（TF-GNN）。

在TensorFlow中，這樣的圖形由 tfgnn.GraphTensor 類型的對(duì)象表示。

這是一個(gè)復(fù)合張量類型（一個(gè)Python類中的張量集合），在 tf.data.Dataset 、 tf.function 等中被接受為「頭等對(duì)象」。

它既能存儲(chǔ)圖結(jié)構(gòu)，也能存儲(chǔ)節(jié)點(diǎn)、邊和整個(gè)圖的特征。

GraphTensors的可訓(xùn)練變換可以定義為高級(jí)Kera API中的Layers對(duì)象，或直接使用 tfgnn.GraphTensor 原語(yǔ)。

GNN：對(duì)上下文中的對(duì)象進(jìn)行預(yù)測(cè)

接下來(lái)，進(jìn)一步解釋下TF-GNN，可以看下其中一個(gè)典型的應(yīng)用：

預(yù)測(cè)一個(gè)龐大數(shù)據(jù)庫(kù)中，由交叉引用表定義的圖中某類節(jié)點(diǎn)的屬性

舉個(gè)例子，計(jì)算機(jī)科學(xué)（CS）的引文數(shù)據(jù)庫(kù)arxiv論文中，有一對(duì)多的引用和多對(duì)一的引用關(guān)系，可以預(yù)測(cè)每篇論文的所在的主題領(lǐng)域。

與大多數(shù)神經(jīng)網(wǎng)絡(luò)一樣，GNN也是在許多標(biāo)記樣本（約數(shù)百萬(wàn)個(gè)）的數(shù)據(jù)集上進(jìn)行訓(xùn)練的，但每個(gè)訓(xùn)練步驟只包含一批小得多的訓(xùn)練樣本（比如數(shù)百個(gè)）。

為了擴(kuò)展到數(shù)百萬(wàn)個(gè)樣本，GNN會(huì)在底層圖中合理小的子圖流上進(jìn)行訓(xùn)練。每個(gè)子圖包含足夠多的原始數(shù)據(jù)，用于計(jì)算中心標(biāo)記節(jié)點(diǎn)的GNN結(jié)果并訓(xùn)練模型。

這一過(guò)程，通常被稱為子圖采樣，對(duì)于GNN訓(xùn)練是極其重要的。

現(xiàn)有的大多數(shù)工具都是以批方式完成采樣，生成用于訓(xùn)練的靜態(tài)子圖。

而TF-GNN提供了，通過(guò)動(dòng)態(tài)和交互采樣來(lái)改進(jìn)這一點(diǎn)的工具。

子圖抽樣過(guò)程，即從一個(gè)較大的圖中抽取小的、可操作的子圖，為GNN訓(xùn)練創(chuàng)建輸入示例

TF-GNN 1.0推出了靈活的Python API，用于配置所有相關(guān)比例的動(dòng)態(tài)或批處理子圖采樣：在Colab筆記中交互采樣。

具體來(lái)說(shuō)，對(duì)存儲(chǔ)在單個(gè)訓(xùn)練主機(jī)主內(nèi)存中的小型數(shù)據(jù)集進(jìn)行「高效采樣」，或通過(guò)Apache Beam對(duì)存儲(chǔ)在網(wǎng)絡(luò)文件系統(tǒng)中的龐大數(shù)據(jù)集（多達(dá)數(shù)億節(jié)點(diǎn)和數(shù)十億條邊）進(jìn)行分布式采樣。

在這些相同的采樣子圖上，GNN的任務(wù)是，計(jì)算根節(jié)點(diǎn)的隱藏（或潛在）狀態(tài)；隱藏狀態(tài)聚集和編碼根節(jié)點(diǎn)鄰域的相關(guān)信息。

一種常見的方法是「消息傳遞神經(jīng)網(wǎng)絡(luò)」。

在每一輪消息傳遞中，節(jié)點(diǎn)沿著傳入邊接收來(lái)自鄰節(jié)點(diǎn)的消息，并從這些邊更新自己的隱藏狀態(tài)。

在n輪之后，根節(jié)點(diǎn)的隱藏狀態(tài)反映了，n條邊內(nèi)所有節(jié)點(diǎn)的聚合信息（如下圖所示，n=2)。消息和新的隱藏狀態(tài)由神經(jīng)網(wǎng)絡(luò)的隱層計(jì)算。

在異構(gòu)圖中，對(duì)不同類型的節(jié)點(diǎn)和邊使用單獨(dú)訓(xùn)練的隱藏層通常是有意義的。

圖為一個(gè)簡(jiǎn)單的「消息傳遞神經(jīng)網(wǎng)」，在該網(wǎng)絡(luò)中，每一步節(jié)點(diǎn)狀態(tài)都會(huì)從外部節(jié)點(diǎn)傳播到內(nèi)部節(jié)點(diǎn)，并在內(nèi)部節(jié)點(diǎn)匯集計(jì)算出新的節(jié)點(diǎn)狀態(tài)。一旦到達(dá)根節(jié)點(diǎn)，就可以進(jìn)行最終預(yù)測(cè)

訓(xùn)練設(shè)置是，通過(guò)將輸出層放置在已標(biāo)記節(jié)點(diǎn)的GNN的隱藏狀態(tài)之上、計(jì)算損失（以測(cè)量預(yù)測(cè)誤差）并通過(guò)反向傳播更新模型權(quán)重來(lái)完成的，這在任何神經(jīng)網(wǎng)絡(luò)訓(xùn)練中都是常見的。

除了監(jiān)督訓(xùn)練之外，GNN也可以以無(wú)監(jiān)督的方式訓(xùn)練，可以讓我們計(jì)算節(jié)點(diǎn)及其特征的離散圖結(jié)構(gòu)的連續(xù)表示（或嵌入）。

然后，這些表示通常在其他ML系統(tǒng)中使用。

通過(guò)這種方式，由圖編碼的離散關(guān)系信息，就能被納入更典型的神經(jīng)網(wǎng)絡(luò)用例中。TF-GNN支持對(duì)異構(gòu)圖的無(wú)監(jiān)督目標(biāo)進(jìn)行細(xì)粒度規(guī)范。

構(gòu)建GNN架構(gòu)

TF-GNN庫(kù)支持構(gòu)建和訓(xùn)練，不同抽象層次的GNN。

在最高層，用戶可以使用與庫(kù)綁定在一起的任何預(yù)定義模型，這些模型以Kera層表示。

除了研究文獻(xiàn)中的一小部分模型外，TF-GNN還附帶了一個(gè)高度可配置的模型模板，該模板提供了經(jīng)過(guò)精心挑選的建模選擇。

谷歌發(fā)現(xiàn)這些選擇，為我們的許多內(nèi)部問(wèn)題提供了強(qiáng)有力的基線。模板實(shí)現(xiàn)GNN層；用戶只需從Kera層開始初始化。

import tensorflow_gnn as tfgnn
from tensorflow_gnn.models import mt_albis


def model_fn(graph_tensor_spec: tfgnn.GraphTensorSpec):
  """Builds a GNN as a Keras model."""
  graph = inputs = tf.keras.Input(type_spec=graph_tensor_spec)


  # Encode input features (callback omitted for brevity).
  graph = tfgnn.keras.layers.MapFeatures(
      node_sets_fn=set_initial_node_states)(graph)


  # For each round of message passing...for _ in range(2):
    # ... create and apply a Keras layer.
    graph = mt_albis.MtAlbisGraphUpdate(
        units=128, message_dim=64,
        attention_type="none", simple_conv_reduce_type="mean",
        normalization_type="layer", next_state_type="residual",
        state_dropout_rate=0.2, l2_regularizatinotallow=1e-5,
    )(graph)


  return tf.keras.Model(inputs, graph)

在最低層，用戶可以根據(jù)用于在圖中傳遞數(shù)據(jù)的原語(yǔ)，從頭開始編寫GNN模型，比如將數(shù)據(jù)從節(jié)點(diǎn)廣播到其所有傳出邊，或?qū)?shù)據(jù)從其所有傳入邊匯集到節(jié)點(diǎn)中。

當(dāng)涉及到特征或隱藏狀態(tài)時(shí)，TF-GNN 的圖數(shù)據(jù)模型對(duì)節(jié)點(diǎn)、邊和整個(gè)輸入圖一視同仁。

因此，它不僅可以直接表示像MPNN那樣以節(jié)點(diǎn)為中心的模型，而且還可以表示更一般形式的的圖網(wǎng)絡(luò)。

這可以（但不一定）使用Kera作為核心TensorFlow頂部的建模框架來(lái)完成。

訓(xùn)練編排

雖然高級(jí)用戶可以自由地進(jìn)行定制模型訓(xùn)練，但TF-GNN Runner還提供了一種簡(jiǎn)潔的方法，在常見情況下協(xié)調(diào)Kera模型的訓(xùn)練。

一個(gè)簡(jiǎn)單的調(diào)用可能如下所示：

from tensorflow_gnn import runner


runner.run(
   task=runner.RootNodeBinaryClassification("papers", ...),
   model_fn=model_fn,
   trainer=runner.KerasTrainer(tf.distribute.MirroredStrategy(), model_dir="/tmp/model"),
   optimizer_fn=tf.keras.optimizers.Adam,
   epochs=10,
   global_batch_size=128,
   train_ds_provider=runner.TFRecordDatasetProvider("/tmp/train*"),
   valid_ds_provider=runner.TFRecordDatasetProvider("/tmp/validation*"),
   gtspec=...,
)

Runner為ML Pain提供了現(xiàn)成的解決方案，如分布式訓(xùn)練和云TPU上固定形狀的 tfgnn.GraphTensor 填充。

除了單一任務(wù)的訓(xùn)練（如上所示）外，它還支持多個(gè)（兩個(gè)或更多）任務(wù)的聯(lián)合訓(xùn)練。

例如，非監(jiān)督任務(wù)可以與監(jiān)督任務(wù)混合，以形成具有特定于應(yīng)用的歸納偏差的最終連續(xù)表示（或嵌入）。調(diào)用方只需將任務(wù)參數(shù)替換為任務(wù)映射：

from tensorflow_gnn import runner
from tensorflow_gnn.models import contrastive_losses


runner.run(
     task={
        "classification": runner.RootNodeBinaryClassification("papers", ...),
        "dgi": contrastive_losses.DeepGraphInfomaxTask("papers"),
      },
    ...
)

此外，TF-GNN Runner還包括用于模型歸因的集成梯度實(shí)現(xiàn)。

集成梯度輸出是一個(gè)GraphTensor，其連接性與觀察到的GraphTensor相同，但其特征用梯度值代替，在GNN預(yù)測(cè)中，較大的梯度值比較小的梯度值貢獻(xiàn)更多。

總之，谷歌希望TF-GNN將有助于推動(dòng)GNN在TensorFlow中的大規(guī)模應(yīng)用，并推動(dòng)該領(lǐng)域的進(jìn)一步創(chuàng)新。