How can I improve this toy Graph Neural Network Generative Language model [closed]

Question

Background

I'm an undergraduate student with research interests in a field of physics that has significant overlap with graph theory, and a functioning knowledge of how simple neural nets work and how to build them with TensorFlow and Keras. As many people are, I'm fascinated by the recent advancements in transformer-based language models, and I've spent the last several weeks reading up on them in an attempt to construct my own simple "mini GPT". In doing so I encountered Graph Neural Networks, and decided I'd try instead to construct a generative language model out of these, inspired by the fact that graphs are (in a very hand-wavy sense) perhaps inherently more amenable to encoding relationships like "attention", etc. I'm aware that the task I'm trying to accomplish could probably be much more easily achieved using alternative architectures. This is mostly just a fun "what if" project.

---

Code

I still haven't quite wrapped my head around all the details of how and what a Graph Neural Net learns, nor what the different types of GNN layers do. Nevertheless, I've constructed a simple GNN using the Spektral library, which takes an input string, and predicts the rest of the string word-by-word by predicting the most probable next token. Here's what I have so far:

import tensorflow as tf
from tensorflow.keras.layers import Input, Dense, Embedding
from tensorflow.keras.preprocessing.sequence import pad_sequences

from spektral.layers import GCNConv
from spektral.utils import normalized_adjacency

import numpy as np

import random
import os

def get_training_data(training_data_dir):
        filenames = []

        for filename in os.listdir(training_data_dir):
                filenames.append(os.path.join(training_data_dir, filename))

        random.shuffle(filenames)

        lines = []

        for filename in filenames:
                with open(filename, "r") as file:
                        for line in file:
                                lines.append(line.strip())

        return lines

# Import data
training_corpus = get_training_data("./training_data")
tokens = [line.split() for line in training_corpus]
vocab = set(token for line in tokens for token in line)

vocab_size = len(vocab)

print("Vocabulary size: ", vocab_size, " words")

# Tokenize
word_to_idx = {word: idx for idx, word in enumerate(vocab)}
idx_to_word = {idx: word for word, idx in word_to_idx.items()}

# Pad and truncate sequences to a fixed length
train_data = [[word_to_idx[token] for token in line] for line in tokens]
train_data_padded = pad_sequences(train_data, maxlen=vocab_size, padding='pre', truncating='pre')

# Shift train_data_padded to create train_labels
train_labels = np.roll(train_data_padded, -1, axis=1)
train_labels[train_labels >= vocab_size] = 0

train_data_padded = np.array(train_data_padded)
train_labels = np.array(train_labels)

# Construct token-to-token similarity matrix
similarity_matrix = np.zeros((vocab_size, vocab_size))

for sentence in tokens:
    for i, token1 in enumerate(sentence):
        for j, token2 in enumerate(sentence):
            if i != j:
                num_words_between = pow(abs(j - i), 2)
                similarity_matrix[word_to_idx[token1], word_to_idx[token2]] += num_words_between

adjacency_matrix = normalized_adjacency(similarity_matrix)


# Construct model
input_layer = Input(shape=(None,))
embedding = Embedding(input_dim=vocab_size, output_dim=vocab_size)(input_layer)
gcn_layer = GCNConv(vocab_size)([embedding, adjacency_matrix])

output_layer = Dense(vocab_size, activation='softmax')(gcn_layer)

model = tf.keras.Model(inputs=input_layer, outputs=output_layer)
model.compile(optimizer='adam', loss='categorical_crossentropy')

# Training loop
num_epochs = 100

model.fit(train_data_padded, tf.keras.utils.to_categorical(train_labels, num_classes=vocab_size), epochs=num_epochs)

padding_token = "<PAD>"

def respond_to_text(initial_string):
  initial_tokens = initial_string.split()

  for i in range(len(initial_tokens)):
    if initial_tokens[i] not in word_to_idx:
        initial_tokens[i] = "<PAD>"

  while len(initial_tokens) < vocab_size:
      initial_tokens.insert(0, padding_token)

  generated_tokens = [word_to_idx[token] for token in initial_tokens]
  max_generation_length = 40

  # Number of tokens from the initial string that have been used
  initial_tokens_used = len(initial_tokens)

  for _ in range(max_generation_length):
      current_seq = np.array([generated_tokens[-vocab_size:]])  # Always use the last vocab_size tokens
      next_token_probs = model.predict(current_seq)[0][-1]
      next_token = np.random.choice(np.arange(vocab_size), p=next_token_probs)
      generated_tokens.append(next_token)

      # If there are more initial tokens to use, do that
      if initial_tokens_used < len(initial_tokens):
          generated_tokens[-vocab_size:] = [word_to_idx[token] for token in initial_tokens[initial_tokens_used:]]
          initial_tokens_used = len(initial_tokens)

  # Generate text
  generated_text = [idx_to_word[idx] for idx in generated_tokens]

  # Remove trailing "<PAD>" tokens
  generated_text = [token for token in generated_text if token != padding_token]

  # Join the tokens into text
  generated_text = " ".join(generated_text)
  print("Generated Text:", generated_text)

while True:
  input_string = input()
  respond_to_text(input_string)

It's relatively simple, thus far. I train the network on strings of text from a subset of the WikiQA corpus, which I pre-process by removing all punctuation and capitalization. I define the elements of the adjacency matrix to be the squared distance between tokens (and 0 between a token and itself). I'm using GCNConv(), admittedly without knowing the intimate details of how it is different from other options provided by Spektral. Since the number of tokens in the input string must be identical to the vocabulary size, I pre-pend it with "<PAD>" tokens, and only pass in the last vocab_size tokens each time I generate a new token. I deal with unknown tokens by replacing them with "<PAD>".

If I understand correctly, the GNN learns the "strengths" of connections between nodes (words), i.e. the edge weights and, as I imagine it, the adjacency matrix encodes a sort of very weak form of "attention".

---

Question

1. What and how, precisely, is this learning? I know that in a simple feed-forward neural network, a set of weights "between" perceptrons is learned. Do I understand correctly that what is being learned here are edge weights? How do node features and the adjacency matrix factor into this, and what does the model "do" with some input text? How are the next token probabilities calculated? I only understand this on a very superficial level, sufficient so as to produce this seemingly somewhat functional script. I can see that, when I plot the graph associated with the GCNConv layer using networkx after fitting, words that are similar (e.g. "boat", "ocean", "water") tend to cluster.
2. How can I improve the model?* I've spent a fair amount of time reading GNN research papers, however they're written largely in what appears to be very subfield-specific jargon that is unfamiliar to me as someone very familiar with graph theory and somewhat familiar with machine learning. I'd like to begin by taking smaller steps, hopefully starting with some suggestions provided by the community here. I have no strict, "objective" criteria in mind, outside of producing more realistic, human-like text.

Here's some example input and output:

1. Input: "the traffic" -> Output: "the traffic was invaluable sense day playing urban evening together made past weekends board plot off color cookies calm concert flowers express eye-opening learn garden outside satisfying laughter movie waves how's sunrise of try traffic scratch day captivating hobby live blanket delicious"
2. Input: "i enjoy" -> Output: "i enjoy watching magical journaling awe-inspiring buds tail feeling entertainment resist homemade ones flavors soothing well-being laughter life culture cleanup picnics beauty accomplishment nature mood ocean up satisfying magical contagious joy admiring feeling live marathons beauty views things expression hikes next happiness vacations"

$*$ Outside of the obvious, e.g. more training data, adjusting hyperparameters, etc.

Answer 1

First of all, I would like to encourage you to keep trying new things, it sounds super fun! There are a few things I would like you to clarify about Graph Neural Networks (GNNs) and Graph Convolutional Neural Networks (GCNs), and then I will answer your questions as best I can. Also, as a sidenote, there are too many questions here, so don't expect me to answer all of them. Given that, I will try to answer the main ones.

Adjacency Matrix and Attention Mechanism

I will address the notion that "the adjacency matrix encodes a sort of very weak form of attention". This notion is kind of correct. In fact, it can be shown that the attention mechanism of the transformer is just an instance of the generalization of the message-passing function of the GNN, where the underlying graph is a complete graph, or in other words, its adjacency matrix is a matrix of all ones. If you continue your research on GNNs, you will see another famous architecture called Graph Attention Neural Networks (GAT). This is just an adaptation of the attention mechanism of the transformer for more simple static graph data. If you want to understand how the transformer is a type of GNN, I can explain it to you later.

First Question: What Exactly is Being Learned?

What is learned in a graph neural network depends on the type of neural network that you are implementing. Generally, there are two types of features that you can input into a GNN (for the forward pass). These are the feature data and the edge data of each node. In the case of the GCNConv layer, this only implements the learning of the feature data (which are the embeddings of each token in your case). So, in your code, the GCNConv is learning the weights that are multiplied with the embeddings. In your code, there are no edge weights that are being trained or even initialized because there is no edge data in the first place.

GCN uses the following formula for calculating the message passing: $ \mathbf{h}_i^{(l+1)} = \sigma \left( \sum_{j \in \mathcal{N}(i) \cup \{i\}} \mathbf{A}_{ij} \mathbf{h}_j^{(l)} \mathbf{W}^{(l)} \right) \ $

• $h_i(l)$ is the feature vector of node i at layer $l$.

• $W(l)$ is the weight matrix for layer $l$.

• $N(i)$ is the set of neighbors of node $i$.

• $A_{ij}$ is the element at the $(i,j)$ position of the adjacency matrix $A$. This could be either 1 or 0, indicating whether there is an edge between nodes $i$ and $j$, or it could be a weight if the graph is weighted.

• σ is an activation function, such as ReLU or Sigmoid.

I would recommend watching this video, which in my opinion explains GCNs very well if you have any more doubts.

Second Question: How Can the Model be Improved?

I can understand how you are feeling about GNNs. I did research on them, and they can be very niche with some terminology. I would recommend that you understand the basics of the GNN architecture very well and then follow up with other types of GNNs. Graph theory is definitely useful, but in my opinion, it only helps you to understand the basics of graphs, and the other things are much more specific to the field of GNNs.

As for the question, sadly, deep learning is not a very interpretable field. This means that sometimes we don't know why some architectures seem to work better than others. For example, the transformer is currently the best model for modeling language. We think this is because of the attention mechanism, which is able to take into account all tokens in context at the same time. However, it is also a very scalable model that does not suffer from the gradient vanishing effect, unlike recurrent neural networks.

Recently, there has been a new type of RNN called RWKV that solves the initial gradient vanishing effect of RNNs and does not use the attention mechanism. We have RAVEN, which is an example of how we can build models that are just as performant as transformers with RWKV.

What I'm trying to say is that it is very difficult to say what things you can use to improve the accuracy of your model, besides the basics of training with more data and tuning hyperparameters. If I had to try it, I would probably make my GNN resemble the transformer architecture as much as possible, adding attention, residual layers, and so on. However, this is something that you will have to explore and experiment with.

Final Thoughts

I don't intend to discourage you; quite the opposite. However, I should note that GCNs may not be best suited for this task. You might also explore standard NLP techniques for boosting accuracy, such as improved tokenization and data quality. For more ambitious improvements, you could look into Reinforcement Learning through Human Feedback (RLHF) techniques.