2005年，劃時代之作「The Graph Neural Network Model」的問世，將圖神經網絡帶到每個人面前。

在此之前，科學家處理圖數據的方式是，在數據預處理階段，將圖轉換為一組「向量表示」。

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

今天，谷歌正式官宣發布TensorFlow GNN 1.0（TF-GNN）——用于大規模構建GNN的經過生產測試的庫。

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

TF-GNN是專為異構圖從頭開始構建的，其中對象和關系的類型由不同的節點和邊集合來表示。

現實世界中的對象及其關系以不同的類型出現，而TF-GNN的異構焦點，使得表示它們變得非常自然。

谷歌科學家Anton Tsitsulin表示，復雜的異構建模又回來了！

TF-GNN 1.0首面世

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

而關系對于理解一個對象的重要性，不亞于孤立地看待對象本身的屬性，比如交通網絡、生產網絡、知識圖譜或社交網絡。

離散數學和計算機科學長期以來一直將這類網絡形式化為圖，由「節點」以各種不規則方式通過邊任意連接而成。

然而，大多數機器學習算法只允許輸入對象之間存在規則統一的關系，如像素網格、單詞序列，或完全沒有關系。

圖形神經網絡，簡稱GNN，是一種強大的技術，既能利用圖的連通性（如早期算法DeepWalk和Node2Vec），又能利用不同節點和邊輸入特征。

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

除了對圖形進行預測之外，GNN還是一個強大的工具——用于彌合與更典型的神經網絡用例之間的鴻溝。

它們以連續的方式對圖的離散關系信息進行編碼，從而可以將其自然地納入另一個深度學習系統。

谷歌在今天正式宣布用于大規模構建GNN的經過生產測試的庫——TensorFlow GNN 1.0（TF-GNN）。

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

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

它既能存儲圖結構，也能存儲節點、邊和整個圖的特征。

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

GNN：對上下文中的對象進行預測

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

預測一個龐大數據庫中，由交叉引用表定義的圖中某類節點的屬性

舉個例子，計算機科學（CS）的引文數據庫arxiv論文中，有一對多的引用和多對一的引用關系，可以預測每篇論文的所在的主題領域。

與大多數神經網絡一樣，GNN也是在許多標記樣本（約數百萬個）的數據集上進行訓練的，但每個訓練步驟只包含一批小得多的訓練樣本（比如數百個）。

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

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

現有的大多數工具都是以批方式完成采樣，生成用于訓練的靜態子圖。

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

子圖抽樣過程，即從一個較大的圖中抽取小的、可操作的子圖，為GNN訓練創建輸入示例

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

具體來說，對存儲在單個訓練主機主內存中的小型數據集進行「高效采樣」，或通過Apache Beam對存儲在網絡文件系統中的龐大數據集（多達數億節點和數十億條邊）進行分布式采樣。

在這些相同的采樣子圖上，GNN的任務是，計算根節點的隱藏（或潛在）狀態；隱藏狀態聚集和編碼根節點鄰域的相關信息。

一種常見的方法是「消息傳遞神經網絡」。

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

在n輪之后，根節點的隱藏狀態反映了，n條邊內所有節點的聚合信息（如下圖所示，n=2)。消息和新的隱藏狀態由神經網絡的隱層計算。

在異構圖中，對不同類型的節點和邊使用單獨訓練的隱藏層通常是有意義的。

圖為一個簡單的「消息傳遞神經網」，在該網絡中，每一步節點狀態都會從外部節點傳播到內部節點，并在內部節點匯集計算出新的節點狀態。一旦到達根節點，就可以進行最終預測

訓練設置是，通過將輸出層放置在已標記節點的GNN的隱藏狀態之上、計算損失（以測量預測誤差）并通過反向傳播更新模型權重來完成的，這在任何神經網絡訓練中都是常見的。

除了監督訓練之外，GNN也可以以無監督的方式訓練，可以讓我們計算節點及其特征的離散圖結構的連續表示（或嵌入）。

然后，這些表示通常在其他ML系統中使用。

通過這種方式，由圖編碼的離散關系信息，就能被納入更典型的神經網絡用例中。TF-GNN支持對異構圖的無監督目標進行細粒度規范。

構建GNN架構

TF-GNN庫支持構建和訓練，不同抽象層次的GNN。

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

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

谷歌發現這些選擇，為我們的許多內部問題提供了強有力的基線。模板實現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)

在最低層，用戶可以根據用于在圖中傳遞數據的原語，從頭開始編寫GNN模型，比如將數據從節點廣播到其所有傳出邊，或將數據從其所有傳入邊匯集到節點中。

當涉及到特征或隱藏狀態時，TF-GNN 的圖數據模型對節點、邊和整個輸入圖一視同仁。

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

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

訓練編排

雖然高級用戶可以自由地進行定制模型訓練，但TF-GNN Runner還提供了一種簡潔的方法，在常見情況下協調Kera模型的訓練。

一個簡單的調用可能如下所示：

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提供了現成的解決方案，如分布式訓練和云TPU上固定形狀的 tfgnn.GraphTensor 填充。

除了單一任務的訓練（如上所示）外，它還支持多個（兩個或更多）任務的聯合訓練。

例如，非監督任務可以與監督任務混合，以形成具有特定于應用的歸納偏差的最終連續表示（或嵌入）。調用方只需將任務參數替換為任務映射：

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還包括用于模型歸因的集成梯度實現。

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

總之，谷歌希望TF-GNN將有助于推動GNN在TensorFlow中的大規模應用，并推動該領域的進一步創新。