TransRefer3D: Entity-and-Relation Aware Transformer
for Fine-Grained 3D Visual Grounding

Visual grounding, also known as phrase grounding, aims to localize descriptive words in images. It has been widely studied in the past several years(Yu et al., 2020; Yang et al., 2019; Zhang et al., 2019; Jing et al., 2020). Karpathy et al. (Karpathy et al., 2014) proposed to map image fragments and sentence fragments into a common space to match images and sentences. Chen et al. (Chen et al., 2017) proposed a Query-guided Regression Network with Context policy (QRC net) and introduce reinforcement learning techniques to achieve better localization. Huang et al. (Huang et al., 2020), proposed a Cross-Modal Progressive Comprehension module (CMPC) and a Text-Guided Feature Exchange module (TGFE) to segment referred entities in the image. Though researchers have made significant progress on 2D visual grounding, the research on 3D visual grounding is still relatively preliminary. In this work, we mainly focus on fine-grained 3D visual grounding task.

3D visual grounding is an emerging research topic. Chen et al. (Chen et al., 2019) released a new benchmark of 3D object localization with natural language descriptions. They proposed ScanRefer to combine visual detection and language encoder for joint inference. Achlioptas et al. (Achlioptas et al., 2020) released another benchmark called ReferIt3D consisting of Nr3D, Sr3D and Sr3D+ datasets, which assumes the segmented object instances are already given. However, its task is more fine-grained where each referred object has distractors of the same category. Text-guided Graph Neural Network (TGNN (Huang et al., 2021)) are proposed for referring instance segmentation on point cloud. Yuan et al. (Yuan et al., 2021) proposed InstanceRefer which utilizes instance attribute, relation and localization perceptions for 3d visual grounding. Different from previous works, we exploit the Transformer architecture and propose Entity-aware Attention (EA) module and Relation-aware Attention (RA) module to build the basic blocks of Transformer to achieve a finer matching between visual and linguistic features.

Transformers have been widely used in natural language processing  (Vaswani et al., 2017; Devlin et al., 2018; Liu et al., 2019). Recently a lot of works also apply them on high-level vision tasks  (Dosovitskiy et al., 2020; Carion et al., 2020), low-level vision tasks  (Chen et al., 2020) and graph tasks (Yun et al., 2019) because of its strong ability to recognize long-term dependency.

A Transformer is composed of several Transformer layers stacked together. Each Transformer layer consists of a multi-head self-attention (MSA) and a feed forward network. MSA is based on the self-attention mechanism. Given three feature matrices: queries $Q\in\mathbb{R}^{n\times d}$, keys and values $K,V\in\mathbb{R}^{m\times d}$, a single-head self-attention is formulated as:

introducing $H$ parallel attention layers as different transformation heads to enlarge the model capacity, we can further obtain a multi-head attention:

where $W_{q}^{(i)},W_{k}^{(i)},W_{v}^{(i)}\in\mathbb{R}^{d_{h}\times d}$ are linear transformation that projects $Q,K,V$ for the $i$-th head and $W_{out}\in\mathbb{R}^{h\times Hd_{h}}$ for the output. After the multi-head attention follows a feed forward network, which is a multi-layer perceptron (MLP). Moreover, the Transformer also utilize Layer Normalization and residual connections.

Multimodal learning requires a full understanding of the contents in each modality and, more importantly, learning the interactions between modalities. Co-attention mechanism (Lu et al., 2016) is proposed to modeling the image attention and question attention in VQA. Yu et al. (Yu et al., 2019) proposed a dense co-attention model to make full interactions between modalities and address the previous issue that each modality learns separate attention distributions. In 3D visual grounding tasks, the context of entity-and-relation has not been fully researched yet. In this paper, we expand the co-attention mechanism to guide context modeling in 3D visual grounding task. This is the first work investigating the co-attention mechanism on the 3D visual-linguistic tasks to the best of our knowledge.

According to the above sections, scenes contain two types of information: entities’ attributes and relations among entities. To better model the multimodal context for fine-grained 3D visual grounding, we propose two cross-modal attention mechanism: Entity-aware Attention (EA) and Relation-aware Attention (RA). Furthermore, we develop an Entity-and-Relation aware Contextual Block (ERCB) that integrate those attention mechanisms as a multimodal encoding unit. By stacking multiple ERCBs, we obtain the proposed TransRefer3D model.

Figure 2 shows the framework of our model. Following previous work  (Achlioptas et al., 2020), we first extract object features with a shared PointNet++ (Qi et al., 2017b) and vectorize the input words with an embedding layer. After that, we fuse position information to both visual and linguistic features and feed them into the proposed TransRefer3D.

We use the same training objective as ReferIt3D  (Achlioptas et al., 2020) consisting of a main task and two auxiliary tasks. For $N$ input objects, the main training task is a $N$-classification supervised by a cross-entropy loss $\mathcal{L}_{\mathrm{main}}$. To obtain the features well embedded to represent $K$ categories, we also adopt $K$-classification on object features and aggregated linguistic features, supervised by cross entropy loss $\mathcal{L}_{\mathrm{obj}}$ and $\mathcal{L}_{\mathrm{lang}}$. The total loss function is $\mathcal{L}=\mathcal{L}_{\mathrm{main}}+\lambda_{\mathrm{obj}}\mathcal{L}_{\mathrm{obj}}+\lambda_{\mathrm{lang}}\mathcal{L}_{\mathrm{lang}}$ where we set $\lambda_{\mathrm{obj}}=\lambda_{\mathrm{lang}}=0.5$ in all of our experiments.

Our goal is to find out the referent according to the referring expression. Usually, the referent and other objects in the scene get mentioned in the referring expression as some keywords. We name all objects appearing in the scene and their corresponding keywords as the entities in each modality. First of all, we establish an Entity-aware Attention (EA) module for finely understanding the multimodal context.

The importance of each entity for locating the referent is different. In a ScanNet indoor scene, often there are more than 20 objects, but most of them make a limited contribution to the visual grounding task. To better understand the referring expression, a strategy is adapt to match the entity objects with corresponding words and suppress relatively useless features. By this, the model can finely identify and establish correlations on those important entities.

Inspired by self-attention mechanism and co-attention widely adopted in 2D vision-language approaches (Huang et al., 2020; Yu et al., 2019), we propose EA module to match those features to model an entity aware multimodal context. In a 3D scene, the $n$ objects are represented by $d$-dimension features. And the referring expression is handled to $m$ word-level features. Given two feature sets in different modality $X\in\mathbb{R}^{n\times d}$ and $Y\in\mathbb{R}^{m\times d}$, the EA module matches $Y$ to $X$:

which is a multimodal generalization of attention. Like self-attention, we further extend EA to a multi-head form by replacing $Attn(\cdot)$ with a multi-head version.

The proposed EA module does symmetrical matchings to visual and word-level linguistic features for better modeling entities in both modals.

The EA module can match entity information in word features and 3D object features. However, only focusing on the entities is insufficient for complicated visual-linguistic reasoning. For example, consider the scene with multiple objects of the same category as the referent. All of them match the corresponding word feature, so their corresponding Entity-aware Attention scores can be similar. Consequently, it can be challenging for the model to distinguish the referent from others.

To address this problem, it is important to let model recognize those relational words (e.g. ”next to”) that appear in the referring expression and enhance entity features according to their relations. We establish a representation of abstract entity-to-entity relation for visual modality and develop a cross-modal Relation-aware Attention (RA) for finer context modeling.

Inspired by DGCNN, we define the relation representation between two entity features $f_{i},f_{j}$ as an asymmetric form:

where $H_{\Theta}$ denotes a nonlinear layer.

Like entity features, each relation features $r_{ij}\in\mathbb{R}^{d}$ can also get matched with corresponding linguistic features. For each entity feature $f_{i}$, we calculate its relation features and introduce co-attention on them to obtain an enhanced cross-modal context of $f_{i}$:

where $L$ is the linguistic feature vector and $\mathrm{Agg}(\cdot)$ denotes a channel-wise aggregation operator gathering the $n$ fused $d$-dimension features in $R^{\prime}_{i}$ to $f^{R}_{i}\in\mathbb{R}^{d}$. Figure 4 shows the diagram of RA. Hence for each $f_{i}$ we can obtain a corresponding $f^{R}_{i}$ as its relation-aware multimodal context representation.

We adopt a binary relation representation calculated on each pair of entities (eq. 4). Assuming that those more complicated relations among multiple entities can be described as a set of binary relations, we stack RA modules in our model to achieve a hierarchical and progressive context comprehension with respect to relation information.

EA and RA can model two essential information in scenes: entities and relations. By composing EA and RA, we propose an Entity-and-relation Aware Contextual Block (ERCB), as shown in Figure 5. ERCB can form a cross-modal context comprehension with respect to both entity and relation information for visual grounding.

In each block, we first feed the visual and linguistic features to self-attention layers. Self-attention layers can effectively extract features in each modality and enhance the performance of cross-modal transformations.

Then, we match the visual features with word-level features with respect to both entity information and relation information. We calculate two types of attention representation $F^{E}$ and $F^{R}$ in parallel and fuse them to obtain the joint features $f^{J}$:

where $f^{J}_{i}\in\mathbb{R}^{d}$. $f^{E}_{i}$ and $f^{R}_{i}$ are the $i$-th feature of $F^{E}$ and $F^{R}$, respectively. $\mathrm{FFN}$ is a feed-forward network and $\mathrm{LN}$ denotes layer normalization (Ba et al., 2016). Further performing an add-and-norm transformation, we can obtain the enhanced multimodal feature $f^{\prime}_{i}$:

We also enhance linguistic features with EA guided by visual features. For each attention layer, a feed-forward network is followed to introduce non-linearity. By stacking those ERCBs, we finally obtain a hierarchical Transformer for fine-grained multimodal context comprehension.

In Section 4.3, we will discuss the ability of each module and the structure of our model. Experiments show that TransRefer3D and its component EACB are much more powerful in modeling multimodal context than other methods.

For a fair comparison, we follow the same settings as ReferIt3D (Achlioptas et al., 2020). We train the model until convergence with Adam optimizer (Kingma and Ba, 2014) and a learning rate of 0.0005. The default feature dimension ($d$) is 128 and the default batch size is 32. We set the depth of ERCBs ($L$) to 4 and use 4-head attentions if not particularly indicated.

We conduct our experiments on three datasets, Nr3D, Sr3D and Sr3D+ (Achlioptas et al., 2020). The details of these datasets are listed as follows:

• •

  Nr3D: Nr3D (Natural Reference in 3D) consists of 41.5K human descriptions collected using a referring game (ReferIt Game). It describes objects in 707 ScanNet scenes. In each scene, there are multiple objects of the same category as the referent.

• •

  Sr3D: Sr3D (Spatial Reference in 3D) contains 83.5K synthetic descriptions. It categorizes spatial relations into 5 types: horizontal proximity, vertical proximity, between, allocentric and support (Achlioptas et al., 2020), and then generates descriptions using language templates. Similar with Nr3D, there are always more than one objects of the same category as the referred object.

• •

  Sr3D+: Sr3D+ is an augmentation of Sr3D. It uses the same method as Sr3D to generate synthetic descriptions. The difference is the referent could be the only object belonging to a fine-grained category in each scene.

Sr3D and Sr3D+ are designed to enhance the performance of the model on Nr3D (Achlioptas et al., 2020). The model is first jointly trained on both Nr3D training set and Sr3D/Sr3D+ and then evaluated on Nr3D test set. Nevertheless, for the completeness of experiments, we also train and test our model only on Sr3D.

In this section, we report the quantitative performance of our model on Nr3D, Sr3D and Sr3D+.

In this subsection, we further discuss the effectiveness of each module in TransRefer3D and the architecture design.

To verify the effectiveness of each module in the proposed model, we train and test TransRefer3D without each of the modules: self-attention (SA), vision-guided linguistic entity aware attention (EA, V$\to$L), language-guided visual entity aware attention (EA, L$\to$V) and relation aware attention (RA). The results are shown in Table 3. The model without SA performs worst in all the experiments. This indicates that the Transformer architecture is suitable for our 3D visual grounding task and that the contextual awareness within a single modal is the basis of cross-modal understanding. Furthermore, removing all of the proposed modules hinders the performance, demonstrating the effectiveness of the proposed EA and RA modules.

Moreover, we train and test our model with different architectures. By default, the proposed ERCB consists of an EA module and a RA module in parallel. We have also tried stacking EA and RA modules, in which the data is first fed into EA module and then RA module. A possible explanation for this is that the entity information and relation information are relatively independent, and a parallel structure might be more reasonable.

As for the depth of our model, we tested TransRefer3D with different depth of ERCB layers on Nr3D dataset. As shown in Figure 6, 4 layers of ERCB achieve the best performance. Shallow models with fewer layers might fail to represent the complicated relations as analyzed in Section 3.3. Simultaneously, deeper Transformers have been observed hard to train because of unbalanced gradients and amplification effect  (Liu et al., 2020a).