google-research

1
# coding=utf-8
2
# Copyright 2024 The Google Research Authors.
3
#
4
# Licensed under the Apache License, Version 2.0 (the "License");
5
# you may not use this file except in compliance with the License.
6
# You may obtain a copy of the License at
7
#
8
#     http://www.apache.org/licenses/LICENSE-2.0
9
#
10
# Unless required by applicable law or agreed to in writing, software
11
# distributed under the License is distributed on an "AS IS" BASIS,
12
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
# See the License for the specific language governing permissions and
14
# limitations under the License.
15

16
"""EdgeScorer layer of the GSL Model.
17

18
The edge scorer contains multiple functions tried in existing models.
19
"""
20
import tensorflow as tf
21

22

23
@tf.keras.utils.register_keras_serializable(package="GSL")
24
class EdgeScorer(tf.keras.layers.Layer):
25
  """Wraps edge scorers to be used for graph structure learning."""
26

27
  def compute_cosine_similarities(self, node_embeddings):
28
    node_embeddings = tf.math.l2_normalize(node_embeddings, dim=1)
29
    return self.compute_dot_product_similarities(node_embeddings)
30

31
  def compute_dot_product_similarities(self, node_embeddings):
32
    similarities = tf.matmul(node_embeddings, node_embeddings, transpose_b=True)
33
    return similarities
34

35

36
@tf.keras.utils.register_keras_serializable(package="GSL")
37
class Attentive(EdgeScorer):
38
  """Generates a fully connected adjacency using an attentive approach.
39

40
  This edge scorer, first uses a vector to project the node features and then
41
  creates a fully connected graph of their similarities.
42
  """
43

44
  def __init__(
45
      self,
46
      initialization,
47
      nheads,
48
      seed = 133337,
49
      **kwargs,
50
  ):
51
    super().__init__()
52
    initialization_dict = {"method1": "ones", "method2": "random_uniform"}
53
    self._initialization = initialization_dict[initialization]
54
    if nheads <= 0:
55
      raise ValueError("Number of heads should be greater than zero.")
56
    self._nheads = nheads
57
    self._seed = seed
58

59
  def build(self, input_shape):
60
    node_embedding_dim = input_shape[-1]
61
    self._attention_vectors = []
62
    for _ in range(self._nheads):
63
      # Multiple instances of initializer should be created to give different
64
      # outputs.
65
      if self._initialization == "ones":
66
        cfg = {"class_name": self._initialization, "config": {}}
67
      else:
68
        cfg = {
69
            "class_name": self._initialization,
70
            "config": {"seed": self._seed},
71
        }
72
        self._seed += 1
73
      initializer = tf.keras.initializers.get(cfg)
74
      self._attention_vectors.append(
75
          tf.Variable(
76
              initial_value=initializer(shape=(node_embedding_dim,)),
77
              trainable=True,
78
          )
79
      )
80

81
  def compute_one_head(
82
      self, features, attention_vector
83
  ):
84
    node_embeddings = tf.multiply(attention_vector, features)
85
    return self.compute_cosine_similarities(node_embeddings)
86

87
  def call(self, inputs):
88
    similarities = 0
89
    for vector in self._attention_vectors:
90
      similarities += self.compute_one_head(inputs, vector)
91
    return similarities / self._nheads
92

93
  def get_config(self):
94
    return dict(
95
        initialization=self._initialization,
96
        nheads=self._nheads,
97
        seed=self._seed,
98
        **super().get_config(),
99
    )
100

101

102
@tf.keras.utils.register_keras_serializable(package="GSL")
103
class FP(EdgeScorer):
104
  """Generates a fully connected adjacency with all values as parameters."""
105

106
  def __init__(
107
      self,
108
      node_features,
109
      initialization,
110
      **kwargs,
111
  ):
112
    super().__init__()
113
    initialization_dict = {"method1": "similarity", "method2": "glorot_uniform"}
114
    self._initialization = initialization_dict[initialization]
115
    self._node_features = node_features
116

117
  def build(self, input_shape=None):
118
    number_of_nodes = input_shape[-2]
119
    if self._initialization == "similarity":
120
      self._similarities = tf.Variable(
121
          initial_value=self.compute_cosine_similarities(self._node_features),
122
          trainable=True,
123
      )
124
    else:
125
      initializer = tf.keras.initializers.get(self._initialization)
126
      self._similarities = tf.Variable(
127
          initial_value=initializer(shape=(number_of_nodes, number_of_nodes)),
128
          trainable=True,
129
      )
130

131
  def call(self, inputs):
132
    # FP edge scorer only returns the adjacecy as is.
133
    return self._similarities
134

135
  def get_config(self):
136
    return dict(
137
        node_features=self._node_features,
138
        initialization=self._initialization,
139
        **super().get_config(),
140
    )
141

142

143
@tf.keras.utils.register_keras_serializable(package="GSL")
144
class MLP(EdgeScorer):
145
  """Generates a fully connected adjacency using an MLP model.
146

147
  This edge scorer, first uses an MLP to project the node features and then
148
  creates a fully connected graph of their similarities.
149
  """
150

151
  def __init__(
152
      self,
153
      hidden_size,
154
      output_size,
155
      nlayers,
156
      activation,
157
      initialization,
158
      dropout_rate,
159
      **kwargs,
160
  ):
161
    super().__init__()
162
    initialization_dict = {"method1": "identity", "method2": "glorot_uniform"}
163
    self._initialization = initialization_dict[initialization]
164
    self._hidden_size = hidden_size
165
    self._output_size = output_size
166
    self._nlayers = nlayers
167
    self._activation = activation
168
    self._dropout_rate = dropout_rate
169

170
  def build(self, input_shape=None):
171
    layers = []
172
    for i in range(self._nlayers):
173
      layers.append(
174
          tf.keras.layers.Dense(
175
              units=self._hidden_size
176
              if i < (self._nlayers - 1)
177
              else self._output_size,
178
              activation=self._activation if i < (self._nlayers - 1) else None,
179
              kernel_initializer=tf.keras.initializers.get(
180
                  self._initialization
181
              ),
182
              use_bias=False,
183
          )
184
      )
185
      if i < (self._nlayers - 1):
186
        layers.append(tf.keras.layers.Dropout(rate=self._dropout_rate))
187
    self._model = tf.keras.Sequential(layers)
188

189
  def call(self, inputs):
190
    node_embeddings = self._model(inputs)
191
    similarities = self.compute_cosine_similarities(node_embeddings)
192
    return similarities
193

194
  def get_config(self):
195
    return dict(
196
        hidden_size=self._hidden_size,
197
        output_size=self._output_size,
198
        nlayers=self._nlayers,
199
        activation=self._activation,
200
        initialization=self._initialization,
201
        dropout_rate=self._dropout_rate,
202
        **super().get_config(),
203
    )
204

205

206
def get_edge_scorer(
207
    name, node_features, **kwargs
208
):
209
  if name == "mlp":
210
    return MLP(**kwargs)
211
  elif name == "attentive":
212
    return Attentive(**kwargs)
213
  elif name == "fp":
214
    return FP(node_features, **kwargs)
215
  else:
216
    raise ValueError(f"Edge scorer {name} is not defined.")
217