Can HGT implementation be optimized?

Hi, I understand that the code purpose might be to implement the Heterogeneous Graph Transformer (HGT), but not specifically in the most efficient way.

At first glance, it seems to me that that at least K and Q are being calculated multiple times unnecessarily, for example every time the src or dst node is repeated. It also seems that possibly some K, Q, V features are being stored even after they were used to calculate the attention and message, consuming unnecessary memory.

Am I understanding it correctly? Thanks!

def forward(self, G, h):
        with G.local_scope():
            node_dict, edge_dict = self.node_dict, self.edge_dict
            for srctype, etype, dsttype in G.canonical_etypes:
                sub_graph = G[srctype, etype, dsttype]

                #If srctype or dsttype is the same, this calculation is performed again without reason
                k_linear = self.k_linears[node_dict[srctype]]
                v_linear = self.v_linears[node_dict[srctype]]
                q_linear = self.q_linears[node_dict[dsttype]]
                
                k = k_linear(h[srctype]).view(-1, self.n_heads, self.d_k)
                v = v_linear(h[srctype]).view(-1, self.n_heads, self.d_k)
                q = q_linear(h[dsttype]).view(-1, self.n_heads, self.d_k)

                e_id = self.edge_dict[etype]

                relation_att = self.relation_att[e_id]
                relation_pri = self.relation_pri[e_id]
                relation_msg = self.relation_msg[e_id]

                k = torch.einsum("bij,ijk->bik", k, relation_att)
                v = torch.einsum("bij,ijk->bik", v, relation_msg)

                sub_graph.srcdata['k'] = k
                sub_graph.dstdata['q'] = q
                sub_graph.srcdata['v_%d' % e_id] = v
               
                #After the attention and message are calculated, the K,V,Q are not deleted, and will stay in memory, even if they wont be used again during this for loop
                sub_graph.apply_edges(fn.v_dot_u('q', 'k', 't'))
                attn_score = sub_graph.edata.pop('t').sum(-1) * relation_pri / self.sqrt_dk
                attn_score = edge_softmax(sub_graph, attn_score, norm_by='dst')

                sub_graph.edata['t'] = attn_score.unsqueeze(-1)

            G.multi_update_all({etype : (fn.u_mul_e('v_%d' % e_id, 't', 'm'), fn.sum('m', 't')) \
                                for etype, e_id in edge_dict.items()}, cross_reducer = 'mean')

[zheng-da]

The new heterogeneous graph API will be available in our next release. However, the new API may be incomplete. We’ll try our best to complete the new API as soon as possible.

However, the API isn’t HGT specific. It is designed to speed up HGT computation. You can still choose how to do aggregation/normalization, etc, with the heterogeneous message passing API.

[zheng-da]

We are in the process of optimizing HGT. Please stay tuned.