KDD '22 | Nimble GNN Embedding with Tensor-Train Decomposition

jasperzhong commented 10 months ago

https://dl.acm.org/doi/pdf/10.1145/3534678.3539423

jasperzhong commented 10 months ago

和 #371 的background一样，也是想在解决缺失node features的图上降低node embedding table size的工作.

这篇使用了一个叫Tensor Train (TT) decomposition的技术. 之前听说过矩阵分解，而TT这是一种对张量的分解方法. https://zhuanlan.zhihu.com/p/269804845 这篇paper基本上是把这样的张量分解技术应用到了N x d维的node embedding table上. 不过他们也做了一些创新，利用到图的homophily.

jasperzhong commented 10 months ago

TTM就是把一个M x N的矩阵拆成d个4D的tensors，叫TT cores. 每个TT core的shape是R_{k-1} x m_k x n_k x R_k，这些R_k叫做TT ranks，是一个input参数，可以控制最后拆解的tensors大小，相当于控制压缩率. 得到了这些TT cores，如果要lookup原始的node embedding table的某一行，也就是一个vector，可以根据TT cores计算出要lookup的那一行vector. TT cores其实就成了model parameters，会接受反向传播的更新的.

然后TT cores的初始化很重要, 生成的TT-cores构建出来的matrix，最好是正交的，之前的工作发现这样对于converge有帮助.

原始的TTM其实是给定了一个正交的matrix，然后decompose出TT cores. 这样是非常costly的，里面还有SVD分解.

为了解决这个问题，就提出了一个比较lightweight的Algorithm 2，直接生成TT cores. 但这个方法就有一些限制，就是TT ranks不能太大.

最后效果都不错看上去. 因为TT cores很小，就几百兆，可以直接放到GPU memory上，虽然为了lookup node embedding需要做一些矩阵乘法计算，但还是比起从full embedding需要很多CPU-to-GPU的transer反而快了很多. 他们有一个breakdown.

值得一提的是他们有一个优化，大致的idea就是想让相近的节点share相同的TT cores，所以他们就relabel了一些node ID，通过graph partition来group nodes.

jasperzhong / read-papers-and-code

KDD '22 | Nimble GNN Embedding with Tensor-Train Decomposition #372