Mạng Nơ-ron Đồ Thị trong TensorFlow

Trong thế giới chúng ta sống, các đối tượng và mối quan hệ giữa chúng tồn tại khắp nơi. Hiểu về các đối tượng không chỉ liên quan đến đặc điểm của chính chúng mà còn liên quan đến cách chúng tương tác với nhau. Các ví dụ điển hình bao gồm mạng lưới hậu cần, mạng sản xuất, kiến thức đồ thị hoặc mạng xã hội. Từ lâu, toán học rời rạc và khoa học máy tính đã biến các loại mạng này thành đồ thị, tức là kết nối các nút bằng các cạnh theo nhiều cách không đều.

Nhưng phần lớn các thuật toán học máy (ML) chỉ cho phép các đối tượng đầu vào có mối quan hệ đồng nhất và quy tắc, như mạng lưới pixel, chuỗi từ, hoặc không có mối quan hệ nào. Mạng nơ-ron đồ thị (GNN) là một kỹ thuật mạnh mẽ có thể sử dụng sự kết nối của đồ thị (ví dụ như các thuật toán DeepWalk và Node2Vec ban đầu) cũng như các tính năng trên các nút và cạnh. GNN có thể dự đoán toàn bộ đồ thị (một phân tử cụ thể có phản ứng dạng nào không?), một nút riêng lẻ (theo trích dẫn, chủ đề của tài liệu này là gì?), hoặc một cạnh tiềm năng (khách hàng có thể mua cùng lúc một sản phẩm và sản phẩm khác không?). Ngoài việc dự đoán đồ thị, GNN cũng là một công cụ mạnh mẽ giúp thu hẹp khoảng cách giữa các trường hợp sử dụng mạng nơ-ron điển hình.

Gần đây, Google đã chính thức ra mắt thư viện TensorFlow GNN 1.0 – một thư viện được thử nghiệm trong sản xuất để xây dựng GNN quy mô lớn. Theo Google, TensorFlow GNN 1.0 hỗ trợ mô hình hóa và đào tạo trong TensorFlow, cũng như trích xuất đồ thị đầu vào từ kho dữ liệu lớn. TF-GNN là một đồ thị đa dạng, sử dụng các tập hợp nút và cạnh khác nhau để biểu diễn các loại đối tượng và mối quan hệ của chúng. Trong thế giới thực, các đối tượng và mối quan hệ thường thuộc nhiều loại khác nhau, có thể được thể hiện một cách tự nhiên thông qua điểm nhấn đa dạng của TF-GNN.

Trong TensorFlow, các đồ thị này được biểu diễn dưới dạng đối tượng tfgnn.GraphTensor. Đây là một loại tensor phức tạp (tức là một tập hợp tensor trong một lớp Python), và cũng là loại đối tượng chính trong tf.data.Dataset, tf.function, v.v. Nó có thể lưu trữ cấu trúc đồ thị, cũng như các tính năng của nút, cạnh và toàn bộ đồ thị. Các biến đổi có thể đào tạo của GraphTensors có thể được xác định dưới dạng lớp Layers trong API Keras, hoặc trực tiếp sử dụng nguyên thủy tfgnn.GraphTensor.

Ví dụ điển hình về việc sử dụng TF-GNN, Google đã đề cập đến việc dự đoán thuộc tính của một loại nút trong đồ thị được định nghĩa bởi bảng chéo tham chiếu trong cơ sở dữ liệu lớn. Ví dụ, cơ sở dữ liệu trích dẫn về bài báo khoa học CS arXiv liên quan đến các mối quan hệ trích dẫn một đối một và nhiều đối một, mục đích là để dự đoán lĩnh vực chủ đề của mỗi bài báo.

Tương tự như hầu hết các mạng nơ-ron, GNN được đào tạo trên các tập dữ liệu lớn chứa hàng triệu ví dụ được đánh dấu, nhưng mỗi bước đào tạo chỉ liên quan đến một số nhỏ ví dụ đào tạo (ví dụ như vài trăm). Để mở rộng quy mô, GNN đào tạo trên các tiểu đồ từ đồ thị lớn hơn. Mỗi tiểu đồ đều chứa đủ dữ liệu gốc để tính toán kết quả GNN của nút được đánh dấu trung tâm và đào tạo mô hình. Quá trình này (thường được gọi là lấy mẫu tiểu đồ) rất quan trọng đối với việc đào tạo GNN. Hầu hết các công cụ hiện tại hoàn thành việc lấy mẫu theo cách tĩnh và tạo ra tiểu đồ tĩnh để đào tạo. TF-GNN cải tiến quá trình này bằng cách lấy mẫu động và tương tác.

TF-GNN 1.0 lần đầu tiên mang lại API Python linh hoạt, có thể cấu hình lấy mẫu tiểu đồ động hoặc tĩnh ở mọi cấp độ tỷ lệ, bao gồm: lấy mẫu hiệu quả cho dữ liệu nhỏ lưu trữ trong bộ nhớ chính của máy chủ đào tạo trong ghi chú Colab, hoặc lấy mẫu phân tán cho dữ liệu lớn lưu trữ trong hệ thống tệp mạng (có thể chứa hàng trăm triệu nút và hàng tỷ cạnh).

Trên các tiểu đồ lấy mẫu này, chức năng của GNN là tính toán trạng thái ẩn (hoặc trạng thái tiềm năng) tại các nút gốc; trạng thái ẩn này tập hợp và mã hóa thông tin liên quan từ vùng lân cận của nút gốc. Một phương pháp ban đầu là sử dụng mạng nơ-ron truyền tin. Trong mỗi vòng truyền tin, các nút nhận tin từ các nút liền kề thông qua các cạnh truyền vào, và cập nhật trạng thái ẩn của chúng. Sau n vòng truyền tin, trạng thái ẩn của nút gốc đã phản ánh thông tin tổng hợp từ tất cả các nút trong phạm vi n cạnh (trong hình dưới đây, n=2). Tin nhắn và trạng thái ẩn mới được tính toán bởi các lớp ẩn của mạng nơ-ron. Trong đồ thị đa dạng, thường nên sử dụng các lớp ẩn được đào tạo riêng biệt cho các loại đỉnh và cạnh khác nhau.

Trong việc thiết lập đào tạo cụ thể, lớp đầu ra được đặt trên trạng thái ẩn GNN của nút đánh dấu, tính toán mất mát (để đo lường sai số dự đoán), và cập nhật trọng số mô hình thông qua backpropagation. Quá trình này giống như đào tạo mạng nơ-ron thông thường.

Ngoài việc đào tạo có giám sát (tức là tối thiểu hóa mất mát được định nghĩa bởi nhãn), GNN còn có thể đào tạo không giám sát (tức là không phụ thuộc vào nhãn). Chúng ta có thể tính toán đại diện liên tục (còn gọi là embedding) cho cấu trúc đồ thị rời rạc và các tính năng của nút. Những đại diện này thường được sử dụng trong các hệ thống học máy khác. Bằng cách này, thông tin về mối quan hệ rời rạc được mã hóa vào đồ thị có thể được bao gồm trong các trường hợp sử dụng mạng nơ-ron điển hình hơn. TF-GNN hỗ trợ việc mục tiêu không giám sát được chuyên biệt hóa cho đồ thị đa dạng.

Thư viện TF-GNN hỗ trợ việc xây dựng và đào tạo GNN ở các mức độ trừu tượng khác nhau.

Ở cấp độ cao nhất, người dùng có thể sử dụng bất kỳ mô hình tiền định nào được biểu diễn dưới dạng lớp Keras được gắn kèm với thư viện. Ngoài các mô hình được đề cập trong tài liệu nghiên cứu, TF-GNN còn cung cấp một bộ các mẫu mô hình có thể cấu hình cao, cung cấp các tùy chọn mô hình được chọn lọc. Chúng tôi thấy rằng các tùy chọn mô hình này cung cấp các nền tảng mạnh mẽ cho nhiều vấn đề nội bộ của Google. Với việc triển khai các lớp GNN này, người dùng chỉ cần khởi tạo từ lớp Keras.

Ở cấp độ thấp nhất, người dùng có thể viết mô hình GNN từ đầu dựa trên các nguyên thủy được sử dụng để truyền dữ liệu trong đồ thị, như gửi dữ liệu từ nút đến tất cả các cạnh truyền ra, hoặc thu thập dữ liệu từ tất cả các cạnh truyền vào đến nút. Khi liên quan đến các tính năng hoặc trạng thái ẩn, TF-GNN xử lý tất cả các nút, cạnh và toàn bộ đồ thị đầu vào một cách bình đẳng, điều này không chỉ cho phép biểu diễn trực tiếp các mô hình tập trung vào nút (như MPNN đã đề cập trước đó), mà còn biểu diễn các dạng đồ thị tổng quát hơn. Quá trình này có thể (nhưng không nhất thiết) được thực hiện thông qua việc sử dụng Keras làm khung mô hình chính trên TensorFlow.

Người dùng nâng cao có thể tự do đào tạo mô hình tùy chỉnh, ngoài ra TF-GNN Tunner còn cung cấp một phương pháp gọn gàng để điều phối đào tạo mô hình Keras trong các trường hợp phổ biến. Dưới đây là một ví dụ đơn giản về lệnh gọi:

Runner cung cấp giải pháp sẵn có cho các vấn đề ML như đào tạo phân tán và điền vào tfgnn.GraphTensor có hình dạng cố định trên TPU đám mây. Ngoài việc đào tạo nhiệm vụ đơn lẻ đã đề cập trước đó, Runner còn hỗ trợ đào tạo song song nhiều (hai hoặc hơn) nhiệm vụ. Ví dụ, có thể trộn nhiệm vụ không giám sát với nhiệm vụ có giám sát để tạo ra một biểu diễn liên tục cuối cùng (tức là embedding) với thiên vị suy luận cụ thể cho ứng dụng. Người gọi chỉ cần thay thế tham số nhiệm vụ bằng bản đồ nhiệm vụ:

Bên cạnh đó, TF-GNN Runner bao gồm một thực hiện tích phân gradient cho việc quyến lỗi mô hình. Đầu ra tích phân gradient là một GraphTensor, kết nối giống như GraphTensor được quan sát, nhưng thay vì tính năng, nó sử dụng giá trị gradient. Trong dự đoán GNN, các giá trị gradient lớn thường đóng góp nhiều hơn so với giá trị gradient nhỏ. Người dùng có thể kiểm tra giá trị gradient để hiểu rõ hơn về những tính năng nào đóng vai trò lớn hơn trong GNN.