combination_slices = tf.unstack(K.reshape(combinations, (B, -1, 2 * self.F_)))
output_features = []
for slice in combination_slices:
dense = Dense(1)(slice) // N x 1 (basically "a(Wh_i, Wh_j)" in the paper)
// TODO: masking
e_i = K.reshape(dense, (1, -1)) // 1 x N (e_i in the paper)
softmax = K.squeeze(K.softmax(e_i)) // N (alpha_i in the paper)
softmax_broadcast = K.transpose(K.reshape(K.tile(softmax, [self.F_]), [self.F_, -1]))
node_features = K.sum(softmax_broadcast * linear_transf, axis=0)
if self.use_bias:
output = K.bias_add(node_features, self.bias)
After Change
def call(self, inputs):
X = inputs[0] // input graph (B x F)
G = inputs[1] // full graph (N x F) (this is necessary in code, but not in theory. Check section 2.2 of the paper)
B = K.shape(X)[0] // Get batch size at runtime
N = K.shape(G)[0] // Get number of nodes in the graph at runtime
outputs = [] // Will store the outputs of each attention head (B x F" or B x KF")
for head in range(self.attention_heads):
kernel = self.kernels[head] // W in the paper (F x F")
attention_kernel = self.attention_kernels[head] // Attention network a in paper (2*F" x 1)
// Compute inputs to attention network
linear_transf_X = K.dot(X, kernel) // B x F"
linear_transf_G = K.dot(G, kernel) // N x F"
// Repeat feature vectors of input: [[1], [2]] becomes [[1], [1], [2], [2]]
repeated = K.reshape(K.tile(linear_transf_X, [1, N]), (-1, self.F_)) // B*N x F"
// Tile feature vectors of full graph: [[1], [2]] becomes [[1], [2], [1], [2]]
tiled = K.tile(linear_transf_G, [B, 1]) // B*N x F"
// Build combinations
combinations = K.concatenate([repeated, tiled]) // N*B x 2F"
combination_slices = K.reshape(combinations, (B, -1, 2 * self.F_)) // B x N x 2F"
// Attention head
dense = K.dot(combination_slices, attention_kernel) // B x N x 1 (a(Wh_i, Wh_j) in the paper)
dense = K.squeeze(dense, -1) // B x N
dense = K.softmax(dense) // B x N
// TODO: masking with Vaswani method (add -inf to masked coefficients)
