Integrating a language model with ULTRA #9

Activity

1. migalkin commented on Dec 24, 2023

   

   migalkin

   on Dec 24, 2023

   Collaborator

   More actions

   Hi!

   I understand that from that point only the tail representations are used to feed the MLP.

   Those aren't really tail representations anymore because message passing updates all node states and starts with initial node states called boundary (where you can append LLM features):

   ULTRA/ultra/models.py

   Lines 137 to 140 in 33c6e6b

   	# initial (boundary) condition - initialize all node states as zeros
   	boundary = torch.zeros(batch_size, data.num_nodes, self.dims[0], device=h_index.device)
   	# by the scatter operation we put query (relation) embeddings as init features of source (index) nodes
   	boundary.scatter_add_(1, index.unsqueeze(1), query.unsqueeze(1))

   In the GNN layer code, we use those boundary states together with the current message (eg, with the sum aggregation):
   

   ULTRA/ultra/layers.py

   Lines 192 to 194 in 33c6e6b

   	if self.aggregate_func == "sum":
   	update = generalized_rspmm(edge_index, edge_type, edge_weight, relation, input, sum="add", mul=mul)
   	update = update + boundary

   That is, in each GNN layer we actually do have an interaction function of (initial) head and (current) node states.

   Adding other entity/relation features seems quite straightforward, I see two possible ways:

   1. Early interaction - you'd need to re-train the model from scratch because all weight matrices will be of different dimensions

      • You can add LLM features of relations right into the graph of relations (so structural and LLM features will be mixed in the GNN) by appending those extra features to the initialized boundary:
        

        ULTRA/ultra/models.py

        Lines 57 to 65 in 33c6e6b

        	# initialize initial nodes (relations of interest in the batcj) with all ones
        	query = torch.ones(h_index.shape[0], self.dims[0], device=h_index.device, dtype=torch.float)
        	index = h_index.unsqueeze(-1).expand_as(query)
        	# initial (boundary) condition - initialize all node states as zeros
        	boundary = torch.zeros(batch_size, data.num_nodes, self.dims[0], device=h_index.device)
        	#boundary = torch.zeros(data.num_nodes, *query.shape, device=h_index.device)
        	# Indicator function: by the scatter operation we put ones as init features of source (index) nodes
        	boundary.scatter_add_(1, index.unsqueeze(1), query.unsqueeze(1))

      • Or you can take the basic relational features from the graph of relations and append relational features in the entity-level NBFNet:
        

        ULTRA/ultra/models.py

        Lines 182 to 184 in 33c6e6b

        	# initialize relations in each NBFNet layer (with uinque projection internally)
        	for layer in self.layers:
        	layer.relation = relation_representations

      • Add entity-level LLM features to the boundary in the entity-level NBFNet:
        

        ULTRA/ultra/models.py

        Lines 137 to 140 in 33c6e6b

        	# initial (boundary) condition - initialize all node states as zeros
        	boundary = torch.zeros(batch_size, data.num_nodes, self.dims[0], device=h_index.device)
        	# by the scatter operation we put query (relation) embeddings as init features of source (index) nodes
        	boundary.scatter_add_(1, index.unsqueeze(1), query.unsqueeze(1))

   2. Late interaction - you can freeze the main GNN models and only change the final MLP by adding LLM features to the output:
      

      ULTRA/ultra/models.py

      Lines 199 to 200 in 33c6e6b

      	output = self.bellmanford(data, h_index[:, 0], r_index[:, 0]) # (num_nodes, batch_size, feature_dim）
      	feature = output["node_feature"]

   This is a less expressive way but won't require re-training the model from scratch.

   In any case, if your graphs have >10k nodes, I'd recommend projecting down the LLM features (usually 768d or more, depends on the LLM) to smaller dimension (32/64d) in order to fit the full-batch GNN layer onto a GPU.

   🚀React with 🚀3Reacted by Daniel, Zhaocheng Zhu and Claudio Pomo
2. daniel4x commented on Jan 13, 2024

   

   daniel4x

   on Jan 13, 2024

   Last edited by daniel4x

   ContributorAuthor

   More actions

   Just an update, I tested all three suggested methods.
   I'll add a side branch later if you are interested, along with an example of language model integration for future reference.

   Generally, the pre-trained embedding was added to the EntityNBFNet __init__:

   ULTRA/ultra/models.py

   Lines 106 to 127 in 04d5c13

   	def __init__(self, input_dim, hidden_dims, num_relation=1, **kwargs):
   	# dummy num_relation = 1 as we won't use it in the NBFNet layer
   	super().__init__(input_dim, hidden_dims, num_relation, **kwargs)
   	self.layers = nn.ModuleList()
   	for i in range(len(self.dims) - 1):
   	self.layers.append(
   	layers.GeneralizedRelationalConv(
   	self.dims[i], self.dims[i + 1], num_relation,
   	self.dims[0], self.message_func, self.aggregate_func, self.layer_norm,
   	self.activation, dependent=False, project_relations=True)
   	)
   	feature_dim = (sum(hidden_dims) if self.concat_hidden else hidden_dims[-1]) + input_dim
   	self.mlp = nn.Sequential()
   	mlp = []
   	for i in range(self.num_mlp_layers - 1):
   	mlp.append(nn.Linear(feature_dim, feature_dim))
   	mlp.append(nn.ReLU())
   	mlp.append(nn.Linear(feature_dim, 1))
   	self.mlp = nn.Sequential(*mlp)

   Slightly modified the code and added the following:

           if lm_vectors is not None:
               # can decide whether to freeze or not...
               self.lm_vectors = nn.Embedding.from_pretrained(lm_vectors, freeze=True)
               self.merge_linear = nn.Linear(feature_dim, 64)

   Per your 1st suggestion, it seems like training from scratch with the following:

          # .....original code.....
           boundary.scatter_add_(1, index.unsqueeze(1), query.unsqueeze(1))
           
           # interaction of boundary with lm vectors
           if self.lm_vectors is not None:
               lm_vectors = self.lm_vectors(h_index)  # mistake - see @migalkin
               lm_vectors = lm_vectors.unsqueeze(1).expand(-1, data.num_nodes, -1)
               boundary = torch.cat([boundary, lm_vectors], dim=-1)
               boundary = self.merge_linear(boundary)

   The merge_linear may not be the best option and can be modified with any other interaction to fit into the Conv layers.

   👍React with 👍1Reacted by Youjin C.
3. migalkin commented on Jan 13, 2024

   

   migalkin

   on Jan 13, 2024

   Last edited by migalkin

   Collaborator

   More actions

   Those lines

   lm_vectors = self.lm_vectors(h_index)
   lm_vectors = lm_vectors.unsqueeze(1).expand(-1, data.num_nodes, -1)

   would take only lm features of head nodes in the batch and copy them to all nodes in the graph - is this what you want?

   If you want to initialize each node with its own lm feature, then you don't need to call the embedding layer and just take all of its weights (and repeat along the batch dimension) like self.lm_vectors.weight.repeat(bs, 1, 1) or sending the whole number of nodes as the index self.lm_vectors(torch.arange(data.num_nodes)).repeat(bs, 1, 1)

   👍React with 👍1Reacted by Daniel
4. dhall1995 commented on Jan 24, 2024

   

   dhall1995

   on Jan 24, 2024

   More actions

   @daniel4x Have you made this into a separate branch? I would be really interested to see your code and hear which of the integration methods performed best for you.

   In my use case I have a mixture of LLM features (edges) and a set of pre-trained embedding features for each of the node types. Your experience that node features offer significant performance benefits tallies a lot with mine so it would be great to integrate this into my code.

   👍React with 👍2Reacted by Michael A. Davis and Geralt