SGTR: End-to-end Scene Graph Generation with Transformer

TLDR

SGTR在笔者看来是一个经典地利用Query将Two stage model复杂度高的问题转换成Sparse的One stage model的方法，并且很大程度上保留了Two stage model的优势。

优点如下:

• 用$Q_{init}$去建模不同Entity之间的关系，从而避免了$O(N_e^2)$个数的Entity-Pair。
• 对Predicate Node Generator的设计使得$Q_p,Q_s,Q_o$可以通过不同的Decoder更新Feature，并且交互信息，从而得到$Q_p$对应的最终的Entity Feature: $Q_e^o\space and\space Q_e^s$。
• 最后的Bipartite Graph Assembling直接利用每个predicate($Q_p$)对应的$Q_e^o\space and\space Q_e^s$和Entity Node之间的距离建立邻接矩阵，再利用TopK得到最终的预测三元组。

缺点如下:

• $Q_{init}$不够灵活
• 二分图边权设计有待挖掘

细节解释

Predicate Query的设计

文中特意将Predicate Node Generator分为三个部分，其中Predicate Query Initialization可能看起来很不起眼，但是整个模型的核心思想笔者认为是在这里的。

从DETR出现以来，Query在Object Detection里被重新设计了很多次，最早DETR发现每个Query预测到的目标分布有一定的规律性，大致分布在图片某一块区域，它们在decode的过程中寻找对应区域的object。之后人们发现这种通过训练学出来pattern的Query很难收敛，不如直接先验地均匀撒在全图，也就是AnchorDETR的思想。然而在Object Detection里，预测目标是框，所以将Anchor从点变成框，也就是将(x,y,w,h)作为query的先验，这是DAB-DETR的思想。因此Query的设计对每个任务来说都有所区别。在SGG任务中，很自然的想法是将Predicate作为Query，直接去解码关系，这个方法出现在Hotr: End-to-end human-object interaction detection with transformers中，它的问题在于，最终的目标不只是预测关系，还需要对应的Object和Subject，如果直接解码Predicate Query，很难建立对应的Triplet，并且没有主客体的先验，关系的语义学习会非常的困难。因此SGTR设计了一种Entity-aware Predicate Query，这种Query先验的建立了和Entity的关系，在学习过程中，会主动寻找和对应Entity相关的关系，这是符合Two stage model的逻辑的。我们可以先学习一下，作者是如何将Entity Feature和Predicate Query联系起来的。

首先SGTR设计了一个Learnable Embedding $Q_{init}$，并且通过公式(1)对Entity Node Generator得到的Entity Feature和Box进行融合，得到Geometric-aware entity representation $R_G^{e}$，并作为$K_{init}$和$V_{init}$，和$Q_{init}$组成MHA的三个输入(q,k,v)。通过公式(2)，得到初始的Predicate Query的表示$Q_{p}^{e} = [Q_{is}, Q_{io}, Q_p]$，很显然这个Query的初始特征实际上就是$R_G^e$的线性组合，那么就可以理解为何说 用$Q_{init}$去建模不同Entity之间的关系 。基于Predicate Query中的Entity Feature，在解码关系的时候就可以并行地解码对应的Subject, Object，这种Alignment是SGG任务中非常StraightForward的想法，因此笔者才认为这里是SGTR最大的贡献之处。

$$\mathbf {K}_{init}=\mathbf {V}_{init}=(\mathbf {H}_e+\mathbf {G}_e), \mathbf {G}_e=\text {ReLU}(\mathbf {B}_e\mathbf {W}_g) \tag{1}$$

$$Q^{e}_{p} = A(Q_{init}, K_{init}, V_{init})W_e, where \space W_e ∈ R^{d×3d} = [W^{is}_{e},W^{is}_{e},W^{p}_{e}] \tag{2}$$

Bipartite Graph Assembling

SGTR中另一个值得我们学习和借鉴的就是Bipartite Graph Assembling边权的设计。有向二分图确实是一种很好的关系图的表示方法，大部分SGG任务最后生成图的时候虽然没有把问题formula成二分图，但是实际上都是去计算Entity和Predicate的距离，然后去连接边，例如上文提到的HOTR，它解码完Predicate Feature之后，直接用feature similarity-based matching function去建立Triplet。但是有一个很自然的问题是，Predicate Query的特征和Entity Query的特征并不在同一个Embedding Space里，很难准确地刻画相似度。在SGTR中，由于并行解码了对应的Object Query和Subject Query，直接利用这两个Query和第一步生成的Entity Query之间计算类别距离和空间距离即可，当然Predicate Query也可以预测对应的Object和Subject的中心点作为辅助信息。

总结

一路分析下来，SGTR的缺点也比较明显。由于$Q_{init}$在inference的时候是Fix住的，所以它很可能会被SGG普遍的长尾问题所影响，只学到了某些Entity-Pair的关系，因此一种动态的$Q_{init}$的设计应该更有作用，比如说设计多组$Q_{init}$，让Predicate Query更加多样性。二分图边权的设计应该还是有改进空间的，看Table 3的实验结果，这种直接Entity之间的距离比基于特征相似性或者空间距离的只优了一点点，因此可能可以考虑结合特征相似性和目前的边权去做一个更好的刻画。SGTR还是一个很典型的例子，让我们看到Query的设计应该是越白越好、越有语义性越好。