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

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

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

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

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

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

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

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

TF-GNN 1.0首面世

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

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

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

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

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

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

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

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

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

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

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

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

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

GNN：對上下文中的對象進行預(yù)測

接下來，進一步解釋下TF-GNN，可以看下其中一個典型的應(yīng)用：

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

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

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

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

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

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

而TF-GNN提供了，通過動態(tài)和交互采樣來改進這一點的工具。

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

谷歌發(fā)現(xiàn)這些選擇，為我們的許多內(nèi)部問題提供了強有力的基線。模板實現(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ù)的原語，從頭開始編寫GNN模型，比如將數(shù)據(jù)從節(jié)點廣播到其所有傳出邊，或?qū)?shù)據(jù)從其所有傳入邊匯集到節(jié)點中。

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

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

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

訓(xùn)練編排

雖然高級用戶可以自由地進行定制模型訓(xùn)練，但TF-GNN Runner還提供了一種簡潔的方法，在常見情況下協(xié)調(diào)Kera模型的訓(xù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)練（如上所示）外，它還支持多個（兩個或更多）任務(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還包括用于模型歸因的集成梯度實現(xiàn)。

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

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