# Fast Nearest Neighbor Machine Translation

Yuxian Meng♣, Xiaoya Li♣, Xiayu Zheng▲, Fei Wu♠
Xiaofei Sun♠♣, Tianwei Zhang◆, Jiwei Li♠♣
♣Shannon.AI, ♠Zhejiang University
▲Peking University, ◆Nanyang Technological University
{yuxian_meng, xiaoya_li, xiaofei_sun<EMAIL_ADDRESS>
<EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS>
###### Abstract

Though nearest neighbor Machine Translation ($k$NN-MT) (Khandelwal et al.,
2020) has proved to introduce significant performance boosts over standard
neural MT systems, it is prohibitively slow since it uses the entire reference
corpus as the datastore for the nearest neighbor search. This means each step
for each beam in the beam search has to search over the entire reference
corpus. $k$NN-MT is thus two-orders slower than vanilla MT models, making it
hard to be applied to real-world applications, especially online services. In
this work, we propose Fast $k$NN-MT to address this issue. Fast $k$NN-MT
constructs a significantly smaller datastore for the nearest neighbor search:
for each word in a source sentence, Fast $k$NN-MT first selects its nearest
token-level neighbors, which is limited to tokens that are the same as the
query token. Then at each decoding step, in contrast to using the entire
corpus as the datastore, the search space is limited to target tokens
corresponding to the previously selected reference source tokens. This
strategy avoids search through the whole datastore for nearest neighbors and
drastically improves decoding efficiency. Without loss of performance, Fast
$k$NN-MT is two-orders faster than $k$NN-MT, and is only two times slower than
the standard NMT model. Fast $k$NN-MT enables the practical use of $k$NN-MT
systems in real-world MT applications. 111The code is available at
https://github.com/ShannonAI/fast-knn-nmt

##  1 Introduction

Machine translation (MT) is a fundamental task in natural language processing
(Brown et al., 1993; Och and Ney, 2003), and the prevalence of deep neural
networks has spurred a diverse array of neural machine translation (NMT)
models to improve translation quality (Sutskever et al., 2014; Bahdanau et
al., 2014; Vaswani et al., 2017). The recently proposed $k$ nearest neighbor
($k$NN) MT model (Khandelwal et al., 2020) has proved to introduce significant
performance boosts over standard neural MT systems. The basic idea behind
$k$NN-MT is that at each decoding step, the model is allowed to refer to
reference target tokens with similar translation contexts in a large datastore
of cached examples. The corresponding reference target tokens provide
important insights on the translation token likely to appear next.

One notable limitation of $k$NN-MT is that it is prohibitively slow: it uses
the entire reference corpus as the datastore for the nearest neighbor search.
This means each step for each beam in the beam search has to search over the
entire reference corpus. $k$NN-MT is thus two-orders slower than vanilla MT
models. The original paper of $k$NN-MT (Khandelwal et al., 2020) suggests
using fewer searching clusters, smaller beams and smaller datastores for
generation speedup, but to achieve satisfactory results, carefully tuning on
these factors under different tasks and datasets is still required according
to analyses in (Khandelwal et al., 2020). The computational overhead
introduced by $k$NN-MT makes it hard to be deployed on real-world online
services, which usually require both model performance and runtime efficiency.

In this work, we propose a fast version of $k$NN-MT – Fast $k$NN-MT, to tackle
the aforementioned issues. Fast $k$NN-MT constructs a significantly smaller
datastore for the nearest neighbor search: for each word in a source sentence,
Fast $k$NN-MT first selects its nearest token-level neighbors, which is
limited to tokens of the same token type. Then at each decoding step, in
contrast to consulting the entire corpus for nearest neighbor search, the
datastore for the currently decoding token is limited within the tokens of
reference targets corresponding to the previously selected reference source
tokens, as shown in Figure 1. The chain of mappings from the target token to
the source token, then to its nearest source reference tokens, and last to
corresponding target reference tokens, can be obtained using FastAlign (Dyer
et al., 2013).

Fast $k$NN-MT provides several important advantages against vanilla $k$NN-MT
in terms of speedup: (1) the datastore in the KNN search is limited to target
tokens corresponding to previously selected reference source tokens, instead
of the entire corpus. This significantly improves decoding efficiency; (2) for
source nearest neighbor retrieval, we propose to restrict the reference
sources tokens that are the same as the query token, which further improves
nearest-neighbor search efficiency. Without loss of performance, Fast $k$NN-MT
is two-orders faster than $k$NN-MT, and is only two times slower than standard
MT model. Under the settings of bilingual translation and domain adaptation,
Fast $k$NN-MT achieves comparable results to $k$NN-MT, leading to a SacreBLEU
score of 39.3 on WMT’19 De-En, 41.7 on WMT’14 En-Fr, and an average score of
41.4 on the domain adaptation task.

##  2 Related Work

#### Neural Machine Translation

Neural machine translation systems (Vaswani et al., 2017; Gehring et al.,
2017; Meng et al., 2019) are often implemented by the sequence-to-sequence
framework (Sutskever et al., 2014) and enhanced with the attention mechanism
(Bahdanau et al., 2014; Luong et al., 2015) which associates the current
decoding token to the most semantically related part in the source side. At
decoding time, beam search and its variants are used to find the optimal
sequence (Sutskever et al., 2014; Li and Jurafsky, 2016). The development of
self-attention (Vaswani et al., 2017) and pretraining (Devlin et al., 2018;
Lewis et al., 2019) has greatly motivated a line of works for more expressive
MT systems. These works include incorporating pretrained models (Zhu et al.,
2020; Guo et al., 2020), designing lightweight model structures (Kasai et al.,
2020; Lioutas and Guo, 2020; Tay et al., 2020; Kasai et al., 2021; Peng et
al., 2021), handling multiple languages (Aharoni et al., 2019; Arivazhagan et
al., 2019; Liu et al., 2020b) and mitigating structural issues in Transformers
(Wang et al., 2019; Nguyen and Salazar, 2019; Liu et al., 2020a; Li et al.,
2020; Xiong et al., 2020b) for more robust and efficient NMT systems.

#### Retrieval-Augmented Models

Retrieving and integrating auxiliary sentences has shown effectiveness in
improving robustness and expressiveness for NMT systems. (Zhang et al., 2018)
up-weighted the output tokens by collecting from the retrieved target
sentences $n$-grams that align with the words in the source sentence, and
(Bapna and Firat, 2019) similarly retrieved $n$-grams but incorporated the
information using gated attention (Cao and Xiong, 2018). (Tu et al., 2018)
updated and stored the hidden representations of recent translation history in
cache for access when new tokens are generated, so that the model can
dynamically adapt to different contexts. (Gu et al., 2018) leveraged an off-
the-shelf search engine to retrieve a small subset of sentence pairs from the
training set and then perform translation given the source sentence along with
the retrieved pairs. (Li et al., 2016; Farajian et al., 2017) proposed to
retrieve similar sentences from the training set for the purpose of adapting
the model to different input sentences. (Bulté and Tezcan, 2019; Jitao et al.,
2020) used fuzzy matches to retrieve similar sentence pairs from translation
memories and augmented the source sentence with the retrieved pairs. Our work
is motivated by $k$NN-MT (Khandelwal et al., 2020) and target improving the
efficiency of $k$NN retrieval while achieving comparable translation
performances.

Apart from machine translation, other NLP tasks have also benefited from
retrieval-augmented models, such as language modeling (Khandelwal et al.,
2019), question answering (Guu et al., 2020; Lewis et al., 2020b, a; Xiong et
al., 2020a) and dialog generation (Weston et al., 2018; Fan et al., 2020;
Thulke et al., 2021). Most of these works perform retrieval at the sentence
level and treat the extracted sentences as additional input for model
generation, whereas fast $k$NN-MT retrieves the most relevant tokens in the
source side and fixes the probability distribution using the aligned target
tokens at each decoding step.

##  3 Background: $k$NN-MT

Given an input sequence of tokens $\bm{x}=\\{x_{1},\cdots,x_{n}\\}$ of length
$n$, an MT model translates it into a target sentence in another language
$\bm{y}=\\{y_{1},\cdots,y_{m}\\}$ of length $m$. A common practice to produce
each token $y_{i}$ on the target side is to obtain a probability distribution
over the vocabulary $p_{\text{MT}}(y_{i}|\bm{x},\bm{y}_{1:i-1})$ from the
decoder and use beam search for generation. The combination of the complete
source sentence and prefix of the target sentence $(\bm{x},\bm{y}_{1:i-1})$ is
called translation context. $k$NN-MT interpolates this probability
distribution with a multinomial distribution
$p_{\text{kNN}}(y_{i}|\bm{x},\bm{y}_{1:i-1})$ derived from the $k$ nearest
neighbors of the current translation context $(\bm{x},\bm{y}_{1:i-1})$ from a
large scale datastore $\mathcal{S}$:

$\displaystyle p(y_{i}|\bm{x},\bm{y}_{1:i-1})$ $\displaystyle=\lambda
p_{\text{kNN}}(y_{i}|\bm{x},\bm{y}_{1:i-1})$ (1)
$\displaystyle+(1-\lambda)p_{\text{MT}}(y_{i}|\bm{x},\bm{y}_{1:i-1})$

More specifically, $k$NN-MT first constructs the datastore $\mathcal{S}$ using
key-value pairs, where the key is the high-dimensional vector of the
translation context produced by a trained MT model $f(\bm{x},\bm{y}_{1:i-1})$,
and the value is the corresponding gold target token $y_{i}$, forming
$\mathcal{S}=\\{(\bm{k},v)\\}=\\{(f(\bm{x},\bm{y}_{1:i-1}),y_{i})\\}$. The
context-target pairs may come from any parallel corpus. Then, using the dense
representation of the current translation context as query
$\bm{q}=f(\bm{x}^{\text{in}},\bm{y}_{1:i-1})$ and $L_{2}$ distance as measure,
$k$NN-MT searches through the entire datastore $\mathcal{S}$ to retrieve $k$
nearest translation contexts along with the corresponding target tokens
$\mathcal{N}=\\{\bm{k}_{j},v_{j}\\}_{j=1}^{k}$. Last, the retrieved set is
transformed to a probability distribution by normalizing and aggregating the
negative $L_{2}$ distances, $-d$, using the softmax operator with temperature
$T$, which can be expressed as follows:

$\displaystyle p_{\text{kNN}}(y_{i}|\bm{x}^{\text{in}},\bm{y}_{1:i-1})$ (2)
$\displaystyle=\frac{\sum_{(\bm{k}_{j},v_{j})\in\mathcal{N}}\mathds{1}_{y_{i}=v_{j}}\exp(-d(\bm{q},\bm{k}_{j})/T)}{Z}$
$\displaystyle
Z=\sum_{(\bm{k}_{j},v_{j})\in\mathcal{N}}\exp(-d(\bm{q},\bm{k}_{j})/T)$

Integrating Eq.(2) into Eq.(1) gives the final probability of generating token
$y_{i}$ for time step $i$. Note that the above $k$NN search-interpolating
process is applied to each decoding step of each beam, and each iteration
needs to run on the full datastore $\mathcal{S}$. This gives a total time
complexity of $\mathcal{O}(|\mathcal{S}|Bm)$, where $B$ is the beam size and
$m$ is the target length. In order for faster nearest neighbor search, $k$NN-
MT leverages FAISS (Johnson et al., 2019), an toolkit for efficient similarity
search and clustering of dense vectors.

##  4 Method: Fast $k$NN-MT

Figure 1: Caching source and target tokens (left, blue): Given a trained NMT
model $f$ and the training corpus $\mathcal{D}_{\text{train}}$, we obtain
representations for all source tokens $\bm{h}$ and target tokens $\bm{z}$ in
the training set, which are the last layer outputs from $f$. Datastore
construction (right, green): Given a test example to translate, which is
{B,C,E} in the example, we first navigate each source token to the tokens of
the same type in the cache, e.g., $x_{12}$, $x_{21}$, $x_{32}$ and $x_{41}$
are identified for token B. Then, the top $c$ nearest neighbors for each
source token are preserved according to the distance between the source token
representation and candidate token representations, e.g., $x_{12},x_{21}$ are
selected for token B. Last, the selected source tokens are aligned to their
target tokens using FastAlign (Dyer et al., 2013). For token B, the aligned
target tokens are $y_{11},y_{24}$. The collection of all aligned target tokens
(along with their representations) constitutes the datastore for the current
input {B,C,E}.

The time complexity of $k$NN-MT before search optimization is
$\mathcal{O}(|\mathcal{S}|Bm)$222Time perplexity can slightly be alleviated
when approximate nearest neighbor search is used, with computation perplexity
being not strictly $\mathcal{O}(|\mathcal{S}|)$. , which is prohibitively slow
when the size of the datastore $\mathcal{S}$ or the beam size $B$ is large. We
propose strategies to address this issue. The same as vanilla $k$NN-MT, fast
$k$NN-MT system is built upon a separately trained MT encoder-decoder model.
To get a better illustration of how Fast $k$NN-MT works, we give a toy
illustration in Figure 1. We use the capitalized characters to denote source
tokens and lower-cased letters to denote target tokens. Given the training
set, which is:

$\displaystyle(\bm{x}^{(1)},\bm{y}^{(1)})$
$\displaystyle=(\\{\text{A,B,C,D}\\},\\{\text{b,c,d,a}\\})$ (3)
$\displaystyle(\bm{x}^{(2)},\bm{y}^{(2)})$
$\displaystyle=(\\{\text{B,C,D}\\},\\{\text{c,d,e,b}\\})$
$\displaystyle(\bm{x}^{(3)},\bm{y}^{(3)})$
$\displaystyle=(\\{\text{A,B,D,E}\\},\\{\text{a,b,c,d,e}\\})$
$\displaystyle(\bm{x}^{(4)},\bm{y}^{(4)})$
$\displaystyle=(\\{\text{B,D,E}\\},\\{\text{b,d,e}\\})$
$\displaystyle(\bm{x}^{(5)},\bm{y}^{(5)})$
$\displaystyle=(\\{\text{D,E,F}\\},\\{\text{d,e,f}\\})$

in the toy example, an encoder-decoder model is trained. Next, we wish to
translate a source string $\\{\text{B,C,E}\\}$ at test time.

###  4.1 Datastore Creation On the Source Side

Given a pretrained encoder-decoder model, and the training corpus, we first
obtain representations for all source tokens and target tokens of the training
set, which are the last layer outputs from the encoder-decoder model. In the
toy example, representations for source tokens {A,B,C,D} in the first training
example $(\\{\text{A,B,C,D}\\},\\{\text{b,c,d,a}\\})$ are respectively
$\bm{h}_{11},\bm{h}_{12},\bm{h}_{13},\bm{h}_{14}$, and for target tokens
{b,c,d,a} are respectively $\bm{z}_{11},\bm{z}_{12},\bm{z}_{13},\bm{z}_{14}$.
Given a test example to translate, which is {B,C,E} in the example, we also
obtain the representation for each of its constituent token, denoted by
$\bm{h}_{\text{B}},\bm{h}_{\text{C}},\bm{h}_{\text{E}}$. Next, we select
nearest neighbor tokens for each source token, i.e., {B, C, E}. The nearest
neighbor tokens are first limited to source tokens of the same token type as
the query token. For token B, tokens of the same token type are $x_{12}$,
$x_{21}$, $x_{32}$, $x_{41}$. Similarly, for the token C in the test example,
tokens of the same type are $\\{x_{13},x_{22}\\}$; for the token E, tokens of
the same type are $\\{x_{34},x_{43},x_{52}\\}$. One issue that stands out is
that, for common words such as “the”, there can be tens of millions of the
same type tokens in the training corpus. We thus need to further limit the
number of nearest neighbors. Let $c$ denote the hyper-parameter that controls
the number of nearest neighbors for each token on the source side, which is
set to 2 in the toy example. We rank all candidates based on the distance
between the source token representation (e.g.,
$\bm{h}_{\text{B}},\bm{h}_{\text{C}},\bm{h}_{\text{E}}$) and candidate token
representations, and select the top $c$. Suppose that in the toy example,
$x_{12},x_{21}$ are selected for token B, $x_{13},x_{23}$ are selected for
token C, $x_{34},x_{52}$ are selected for token E. The concatenation of
selected candidates for all source tokens constitute the datastore on the
source side, which is
$\mathcal{D}_{\text{source}}=\\{x_{12},x_{21},x_{13},x_{23},x_{34},x_{52}\\}$
in the toy example. The datastore creation for source tokens (e.g., {B, C, E})
can be run in parallel.

###  4.2 Datastore Creation On the Target Side

For decoding, the model needs to refer to reference target tokens rather than
source tokens. We thus need to transform $\mathcal{D}_{\text{source}}$ to a
list of target tokens. We use FastAlign (Dyer et al., 2013) toolkit to achieve
this goal. FastAlign maps source tokens to target tokens based on the IBM
model (Och and Ney, 2003). Source tokens in $\mathcal{D}_{\text{source}}$ that
do not have correspondence on the target side are abandoned. Output target
tokens from FastAlign form the datastore on the target side, denoted by
$\mathcal{D}_{\text{target}}$. In the toy example,
$x_{12},x_{21},x_{13},x_{23},x_{34},x_{52}$ are respectively mapped to
$\mathcal{D}_{\text{target}}=\\{y_{11},y_{24},y_{12},y_{21},y_{35},y_{52}\\}$.
The size of $\mathcal{D}_{\text{target}}$ is $c\times n$, where $n$ is the
source length.

In practice, we first iterate over all examples in the training set,
extracting all the source token representations and all the target token
representations. Then, we build a separate token-specific cache
$\mathcal{D}_{v}$ for each $v$ in vocabulary, which consists of (key, value)
pairs where the key is the high-dimensional representation $\bm{h}$ and the
value is a binary tuple containing the corresponding aligned target token
along with its representation $\bm{z}$. Then we could map each source token of
a given input sentence to its corresponding cache $\mathcal{D}_{v}$, and build
the target-side datastore following the steps in Section 4.1 and Section 4.2.
The process of caching source and target tokens is present in Algorithm 1.

###  4.3 Decoding

At the decoding time, the datastore for each decoding step is all limited to
$\mathcal{D}_{\text{target}}$, within which $k$NN search is performed. Since
tokens in $\mathcal{D}_{\text{target}}$ are not all related to the current
decoding, nearest neighbor search is performed to select the top $k$
candidates from $\mathcal{D}_{\text{target}}$ for each decoding step. For the
nearest neighbor search here, we use the current representation $\bm{h}$ at
the decoding time to query target representation $\bm{z}$ for target tokens in
$\mathcal{D}_{\text{target}}$. The selected nearest neighbors and their
representations are used to compute the final word generation probability
based on Eq.(1) and Eq.(2).

Input : All sentence-pairs in training set:
$(\bm{x}^{(1)},\bm{y}^{(1)}),...,(\bm{x}^{(N)},\bm{y}^{(N)})$, vocabulary
$\mathcal{V}$

NMT encoder $f_{e}$, NMT decoder $f_{d}$, word alignment for each sentence-
pair $(A^{(1)},...,A^{(N)})$

Input for test: $\bm{x}$

Output : Target datastore $\mathcal{D}_{\text{target}}(\bm{x})$ for $\bm{x}$

$\mathcal{D}_{\text{target}}(\bm{x})\leftarrow\emptyset$ $\blacktriangleright$
initialize the datastore $\mathcal{D}_{\text{source}}(\bm{x})$ for $\bm{x}$

for  _$v\leftarrow 1$ to $\mathcal{V}$_ do

$\mathcal{D}_{v}\leftarrow\emptyset$ $\blacktriangleright$ initialize the
(key, value) datastore for each word in the vocabulary

end for

% Caching Source and Target Tokens:

for  _$i\leftarrow 1$ to $N$_ do

$n_{i}\leftarrow$ length of $\bm{x}^{(i)}$

$m_{i}\leftarrow$ length of $\bm{y}^{(i)}$

$[\bm{h}_{1},...,\bm{h}_{n_{i}}]\leftarrow f_{e}(\bm{x}^{(i)})$
$\blacktriangleright$ computing representations for each source word

$[\bm{z}_{1},...,\bm{z}_{m_{i}}]\leftarrow f_{d}(\bm{x}^{(i)},\bm{y}^{(i)})$
$\blacktriangleright$ computing representations for each target word

for  _$j\leftarrow 0$ to $n_{i}$_ do

for  _$k\leftarrow 0$ to $m_{i}$_ do

if  _$(j,k)\in A^{(i)}$_ then

add $(\bm{h}_{i},(\bm{z}_{k},y^{(i)}_{k}))$ to $\mathcal{D}_{x^{(i)}_{j}}$

% key of $\mathcal{D}_{x^{(i)}_{j}}$ are all hidden representations
$\bm{h}_{i}$ for tokens of the same token type as $x^{(i)}_{j}$

end if

end for

end for

end for

% Generate Datastore for the test input $\bm{x}$:

$n_{x}\leftarrow$ length of $\bm{x}$

for  _$k\leftarrow 1$ to $N$_ do

obtain hidden representation for the $k$-th token $\bm{h}_{k}$

$\\{\bm{z},y\\}$ $\leftarrow$ top $c$ items from $\mathcal{D}_{x_{k}}$ using
$\bm{h}_{k}$ as the query

add $\\{\bm{z},y\\}$ to $\mathcal{D}_{\text{target}}(x)$

end for

Algorithm 1 Constructing Datastore for a Test Input $\bm{x}$.

###  4.4 Quantization

Although the prohibitive computational cost issue of $k$NN-MT has been
addressed, the intensive memory for datastore remains a problem, as we wish to
cache all source and target representations of the entire training set.
Additionally, frequently accessing Terabytes of data is also extremely time-
intensive. To address this issue, we propose to use product quantization (PQ)
(Jegou et al., 2010) to compress the high-dimensional representations of each
token. Formally, given a vector $\bm{x}\in\mathbb{R}^{D}$, we represent it as
the concatenation of $M$ subvectors:
$\bm{x}=[\bm{x}^{1},\bm{x}^{2},...,\bm{x}^{M}]$, where all subvectors have the
same dimension $d=D/M$. The product quantizer $q$ consists of $M$ sub-
quantizers $q_{1},...,q_{M}$, each maps a subvecor
$\bm{x}^{m}\in\mathbb{R}^{d}$ to a codeword $\bm{c}$ in a subspace codebook
with $n$ codewords:
$\mathcal{C}^{m}=\\{\bm{c}^{m}_{1},\bm{c}^{m}_{2},...,\bm{c}^{m}_{n}\\}$. The
objective function of PQ is:

$\min_{q^{1},...,q^{M}}\sum_{\bm{x}}\sum_{m=1}^{M}\|{\bm{x}^{m}-q_{m}(\bm{x}^{m})}\|^{2}$
(4)

Therefore each $\bm{x}$ is mapped to its nearest codeword in the Cartesian
product space $\mathcal{C}=\mathcal{C}^{1}\times...\times\mathcal{C}^{M}$. If
each subspace codebook $\mathcal{C}^{m}$ has $n$ codewords, then the Cartetian
product space $\mathcal{C}$ could represent $n^{m}$ $D$-dimensional codewords
with only $n\times m$ $d$-dimensional vectors, thus significantly eased the
memory issue. With the quantization technique, we are able to compress each
token representation to 128-bytes. For example, for the WMT19 En-De dataset,
the memory size is reduced from 3.5TB to 108GB.

###  4.5 $k$NN Retrieval Details

In practice, executing exact nearest neighbor search over millions or even
billions of tokens could be time-consuming. Hence we use FAISS (Johnson et
al., 2019) for fast approximate nearest neighbor search. All token
representations are quantized to 128-bytes. Recall that we build a token-
specific datastore $\mathcal{D}_{v}$ for each $v$ in vocabulary. We do brute
force search for tokens whose frequency $n_{v}$ is lower than 30000. For those
tokens whose frequency is larger than 30000, the keys are stored in clusters
to speed up search. The number of clusters for token $v$ is set to
$\min(4\times\sqrt{n_{v}},n_{v}/30)$. To learn the cluster centroids, we use
at most 5M keys for each token $v$. During inference, we query the datastore
for $k=512$ neighbors through searching 32 nearest clusters.

###  4.6 Discussions on Comparisons to Vanilla $k$NN-MT

The speedup of Fast $k$NN-MT lies in the following three aspects:
(1) For nearest neighbor retrieval on the source side, we first restrict the
reference tokens that are the same as the query token. This strategy
significantly narrows down the search space to roughly
$|\mathcal{S}|/\text{mid(F)}$ times, where $|\mathcal{S}|$ denotes the number
of tokens in the corpus, and $\text{mid}(F)$ denotes the medium word frequency
in the corpus.
(2) The nearest neighbor search for all source tokens on the source side can
be run in parallel, which is also a key speedup over $k$NN-MT. For vanilla
$k$NN-MT, $k$NN search is performed on the target side and has to be auto-
regressive: the representation for the current decoding step, which is used
for the $k$NN search over the entire corpus, relies on previously generated
tokens. Therefore, the $k$NN search for the current step has to wait for the
finish of $k$NN searches for all previous generation steps.
(3) On the target side, the datastore in the $k$NN search is limited to target
representations corresponding to selected reference source tokens. Though the
nearest neighbor search in the decoding process is auto-regressive and thus
cannot be run in parallel, the running cost is fairly low: recall that the
size of $\mathcal{D}_{\text{target}}$ is $c\times n$. Across all settings, the
largest value of $c$ is set to 512. The size of $\mathcal{D}_{\text{target}}$
is roughly 15k. Performing nearest neighbor searches among 15k candidates is
relatively cheap for NMT, and is actually cheaper than the softmax operation
for word prediction, where the vocabulary size is usually around 50k.
The combination of all these aspects leads to Fast $k$NN-MT two orders of
magnitude faster than vanilla $k$NN-MT.

##  5 Experiments

###  5.1 Bilingual Machine Translation

We conduct experiments on two bilingual machine translation datasets: WMT’14
English-French and WMT’19 German-English. To create the datastore, we follow
(Ng et al., 2019) to apply language identification filtering, keeping only
sentence pairs with correct languages on both sides. We also remove sentences
longer than 250 tokens and sentence pairs with a source/target ratio exceeding
1.5. For all datasets, we use the standard Transformer-base model provided by
FairSeq (Ott et al., 2019)
library.333https://github.com/pytorch/fairseq/tree/master/examples/translation
The model has 6 encoder layers and 6 decoder layers. The dimensionality of
word representations is 1024, the number of multi-attention heads is 16, and
the inner dimensionality of feedforward layers is 8192. Particularly,
following (Khandelwal et al., 2020), the model for WMT’19 German-English has
also been trained on over 10 billion tokens of extra backtranslation data as
well as fine-tuned on newstest test sets from previous years. We report the
SacreBLEU scores (Post, 2018) for
comparison.444https://github.com/mjpost/sacrebleu Table 1 shows our results on
the two NMT datasets. The proposed Fast $k$NN-MT model is able to achieve
slightly better results to the vanilla $k$NN-MT model on WMT’19 German-
English, and competitive results on WMT’14 English-French, with less $k$NN
search cost.

Model | De-En | En-Fr
---|---|---
base MT | 37.6 | 41.1
+$k$NN-MT | 39.1(+1.5) | 41.8(+0.7)
+fast $k$NN-MT | 39.3(+1.7) | 41.7(+0.6)
Table 1: SacreBLEU scores on WMT’19 De-En and WMT’14 En-Fr.

###  5.2 Domain Adaptation

We also measure the effectiveness of the proposed Fast $k$NN-MT model on the
domain adaptation task. We use the multi-domain datasets which are originally
provided in (Koehn and Knowles, 2017) and further cleaned by (Aharoni and
Goldberg, 2020). These datasets include German-English parallel data for
train/validation/test sets in five domains: Medical, Law, IT, Koran and
Subtitles. We use the trained German-English model introduced in Section 5.1
as our base model, and further build domain-specific datastores to evaluate
the performance of Fast $k$NN-MT on each domain. Table 2 shows that Fast
$k$NN-MT achieves comparable results to vanilla $k$NN-MT on domains of Medical
(54.6 vs. 54.4), IT (45.9 vs. 45.8) and Subtitles (31.9 vs. 31.7), and
outperforms vanilla $k$NN-MT on the domain of Koran (21.2 vs. 19.4). The
average score of Fast $k$NN-MT (41.7) is on par with the result of (Aharoni
and Goldberg, 2020) (41.3), which trains domain-specific models and reports
in-domain results.

Following (Khandelwal et al., 2020), we also carry out experiments under the
out-of-domain and multi-domain settings and report the results on Table 2. “+
WMT19’ datastore” shows the results for retrieving neighbors from 770M tokens
of WMT’19 data that the model has been trained on, and “+ all-domain
datastore” shows the results where the model is trained on the multi-domain
datastore from all six settings. The BLEU improvement is much smaller on the
out-of-domain setup compared to the in-domain setup, illustrating that the
proposed framework relies on in-domain data to retrieve valuable contexts. For
the multi-domain setup, the performance for all six domains generally remains
the same and only a small drop of the average score is witnessed. This shows
that the Fast $k$NN-MT framework is robust to a massive amount of out-of-
domain data and is able to retrieve the context-related information from in-
domain data.

Model | Medical | Law | IT | Koran | Subtitles | Avg.
---|---|---|---|---|---|---
(Aharoni and Goldberg, 2020) | 54.8 | 58.8 | 43.5 | 21.8 | 27.4 | 41.3
base MT | 39.9 | 45.7 | 38.0 | 16.3 | 29.2 | 33.8
+in-domain datastore: |  |  |  |  |  |
$k$NN-MT | 54.4(+14.5) | 61.8(+16.1) | 45.8(+7.8) | 19.4(+3.1) | 31.7(+2.5) | 42.6(+8.8)
fast $k$NN-MT | 53.6(+13.7) | 56.0(+10.3) | 45.5(+7.5) | 21.2(+4.9) | 30.5(+1.3) | 41.4(+7.6)
+WMT19’ datastore |  |  |  |  |  |
$k$NN-MT | 40.2(+0.3) | 46.7(+1.0) | 40.3(+2.3) | 18.0(+1.7) | 29.2(+0.0) | 34.9(+1.1)
fast $k$NN-MT | 41.5(+1.6) | 45.9(+0.2) | 41.0(+3.0) | 17.6(+1.3) | 29.2(+0.0) | 35.0(+1.2)
+all-domains datastore |  |  |  |  |  |
$k$NN-MT | 54.5(+14.6) | 61.1(+15.4) | 48.6(+10.6) | 19.2(+2.9) | 31.7(+2.5) | 43.0(+9.2)
fast $k$NN-MT | 53.2(+13.3) | 54.6(+8.9) | 46.4(+8.4) | 19.5(+3.2) | 29.7(+0.5) | 40.7(+6.9)

Table 2: SacreBLEU results for domain adaptation.

|  Source |  Target
---|---|---
original sentence pair |  Zwei Fisch@@ betriebe haben ihre Tätigkeit eingestellt . |  Two fish establi@@ sh@@ ments have ce@@ ased their activities .
reference of "Zwei" |  Zwei Gemein@@ schaft@@ sher@@ steller in Griechenland , die bei der vor@@ ausge@@ gangenen Untersuchung mit@@ gearbeitet hatten , gaben ihre Tätigkeit auf . |  Fur@@ ther@@ more , two Community producers in Greece who took part in the previous investigation ce@@ ased their activity .
reference of "Fisch@@" |  \- fünf auf die Fisch@@ industrie , |  \- five to representatives of the fi@@ sh@@ ery industries ,
reference of "betriebe" |  ( 4 ) Drei Fleisch@@ betriebe mit Übergangs@@ regelung haben erhebliche Anstrengungen unter@@ nommen , neue Betriebs@@ anlagen zu bauen . |  ( 4 ) Three meat establi@@ sh@@ ments on the list of establi@@ sh@@ ments in transition have made considerable efforts to build new facilities .
reference of "haben", "ihre", "Tätigkeit", "eingestellt" and "." |  Einige Betriebe haben ihre Tätigkeit eingestellt . |  Cer@@ tain establi@@ sh@@ ments have ce@@ ased their activities .

Table 3: A test sentence pair from the Law domain. We show the original
sentence pair for test (the first row), the nearest-neighbor tokens on the
source side along with the sentences that retrieved tokens reside in (the
second column), and the aligned target tokens extracted from FastAlign, along
with sentences in which target tokens reside in (the third column). The
retrieved tokens are in red.

###  5.3 Analysis

Figure 2: First figure: the similarity heatmap between the gold target tokens
and the retrieved target neighbors. Second figure: BLEU scores with respect to
different $c$, the number of nearest neighbors on the source side for each
source token. Third figure: BLEU scores with respect to different $k$, the
number of nearest neighbors on the target side for each target token. Forth
figure: Speed comparison between base MT, $k$NN-MT and fast $k$NN-MT.

#### Examples

To visualize the effectiveness of the proposed Fast $k$NN-MT model, we
randomly choose an example from the test set of the Law domain. Table 3 shows
the test sentence, the retrieved nearest neighbor tokens on the source side,
and the corresponding target tokens. The first figure in Figure 2 demonstrates
the similarity heatmap between the gold target tokens and the selected target
neighbors. We can see that the retrieved target nearest tokens are highly
correlated with the ground-truth target tokens, exhibiting the ability of Fast
$k$NN-MT to accurately extract nearest reference tokens at each decoding step.

#### The Effect of the Number of neighbors per token on the source side

We queried the datastore for nearest $c$ neighbors for each source token.
Intuitively, the larger the $c$ is, the more likely the model could recall the
nearest neighbors on the target side. The second figure in Figure2 verifies
this point: the model performance increases drastically when $c$ increases
from 8 to 64, and then continues increasing as $c$ is up to 512.

#### The Effect of the Number of neighbors per token on the target side

Fast $k$NN-MT selects top $k$ nearest neighbors at each decoding step for
computing the probability $p_{\text{kNN}}$ in Eq.(2). The third figure in
Figure 2 shows that the model performance first increases and then decreases
when we continue enlarging the value of $k$, with $c$ fixed at 512, which is
consistent with the observation in (Khandelwal et al., 2020). This is because
that using neighbors that are too far away from the ground-truth target token
adds noise to the model prediction, and thus hurts the performance.

#### Speed comparison

When the beam size is fixed, the time complexity of Fast $k$NN-MT is mainly
controlled by the number of retrieved neighbors $c$ for each source
token.555$k$ plays a minor role to the overall time complexity because each
search on the target side is performed within a total amount of $cn$ tokens,
which is negligible compared to the time cost spent on the source side. The
last figure in Figure 2 shows the speed comparison between base MT, $k$NN-MT
and fast $k$NN-MT when we vary the value of $c$. Fast $k$NN-MT decoded nearly
as fast as the vanilla MT model when $c$ is small. When $c$ reaches 512,
$k$NN-MT is about two times slower than the vanilla MT model. By contrast,
vanilla $k$NN-MT is two order of magnitude slower than base MT and Fast $k$NN-
MT regarding the decoding speed. This is because Fast $k$NN-MT substantially
restricts the search space during decoding, whereas vanilla $k$NN-MT has to
execute $k$NN search over the entire datastore at each decoding step.

#### Similarity function

We have tried different similarity functions when retrieving $c$ nearest
neighbors on source side and computing the $k$NN distribution. These functions
include cosine similarity, inner product and $L_{2}$ distance, the SacreBLEU
scores for which are respectively 39.2, 39.1 and 38.8 on WMT’19 German-
English, showing that cosine similarity is a better measure for representation
distance than $L_{2}$ distance and inner product.

#### Effect of quantization

Model | Medical | Law | IT | Koran | Subtitles | Avg.
---|---|---|---|---|---|---
fast $k$NN-MT | 53.6 | 61.0 | 45.5 | 21.2 | 31.9 | 41.7
+full-precision | 53.8 | 61.1 | 45.8 | 21.3 | 30.7 | 41.5

Table 4: SacreBLEU scores on domain adaptation when using full precision.

Due to the memory issue, we applied quantization to compress the high-
dimensional representation of each token in the training set. We investigate
how quantization would affect model performances. As shown in Table 4,
quantization has minor side effects in terms of BLEU scores, and when we use
full precision instead of quantization, the average BLEU score only increases
0.1, which suggests that computing similarity using compressed vectors is a
viable trade-off between memory usage and model performance.

##  6 Conclusion

In this work, we propose a fast version of $k$NN-MT – Fast $k$NN-MT – to
address the runtime complexity issue of the vanilla $k$NN-MT. During decoding,
Fast $k$NN-MT constructs a significantly smaller datastore for the nearest
neighbor search: for each word in a source sentence, Fast $k$NN-MT selects its
nearest tokens from a large-scale cache. The selected tokens are the same as
the query token. Then at each decoding step, in contrast to using the entire
datastore, the search space is limited to target tokens corresponding to the
previously selected reference source tokens. Experiments demonstrate that this
strategy drastically improves decoding efficiency while maintaining model
performances compared to vanilla $k$NN-MT under different settings including
bilingual machine translation and domain adaptation. Comprehensive ablation
studies are performed to understand the behavior of each component in Fast
$k$NN-MT. In future work, we plan to further improve the efficiency of Fast
$k$NN-MT by applying clustering techniques to build the datastore.

## Acknowledgement

This work is supported by the Science and Technology Innovation 2030 - “New
Generation Artificial Intelligence” Major Project (No. 2021ZD0110201) and the
Key R & D Projects of the Ministry of Science and Technology (2020YFC0832500).
We would like to thank anonymous reviewers for their comments and suggestions.

## References

  * Aharoni and Goldberg (2020) Roee Aharoni and Yoav Goldberg. 2020.  Unsupervised domain clusters in pretrained language models.  _arXiv preprint arXiv:2004.02105_.
  * Aharoni et al. (2019) Roee Aharoni, Melvin Johnson, and Orhan Firat. 2019.  Massively multilingual neural machine translation.  _arXiv preprint arXiv:1903.00089_.
  * Arivazhagan et al. (2019) Naveen Arivazhagan, Ankur Bapna, Orhan Firat, Dmitry Lepikhin, Melvin Johnson, Maxim Krikun, Mia Xu Chen, Yuan Cao, George Foster, Colin Cherry, et al. 2019\.  Massively multilingual neural machine translation in the wild: Findings and challenges.  _arXiv preprint arXiv:1907.05019_.
  * Bahdanau et al. (2014) Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. 2014.  Neural machine translation by jointly learning to align and translate.
  * Bapna and Firat (2019) Ankur Bapna and Orhan Firat. 2019.  Non-parametric adaptation for neural machine translation.  _arXiv preprint arXiv:1903.00058_.
  * Brown et al. (1993) Peter F Brown, Stephen A Della Pietra, Vincent J Della Pietra, and Robert L Mercer. 1993.  The mathematics of statistical machine translation: Parameter estimation.  _Computational linguistics_ , 19(2):263–311.
  * Bulté and Tezcan (2019) Bram Bulté and Arda Tezcan. 2019.  Neural fuzzy repair: Integrating fuzzy matches into neural machine translation.  In _57th Annual Meeting of the Association-for-Computational-Linguistics (ACL)_ , pages 1800–1809.
  * Cao and Xiong (2018) Qian Cao and Deyi Xiong. 2018.  Encoding gated translation memory into neural machine translation.  In _Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing_ , pages 3042–3047.
  * Devlin et al. (2018) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2018.  Bert: Pre-training of deep bidirectional transformers for language understanding.  _arXiv preprint arXiv:1810.04805_.
  * Dyer et al. (2013) Chris Dyer, Victor Chahuneau, and Noah A Smith. 2013.  A simple, fast, and effective reparameterization of ibm model 2.  In _Proceedings of the 2013 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies_ , pages 644–648.
  * Fan et al. (2020) Angela Fan, Claire Gardent, Chloe Braud, and Antoine Bordes. 2020.  Augmenting transformers with knn-based composite memory for dialogue.  _arXiv preprint arXiv:2004.12744_.
  * Farajian et al. (2017) M Amin Farajian, Marco Turchi, Matteo Negri, and Marcello Federico. 2017.  Multi-domain neural machine translation through unsupervised adaptation.  In _Proceedings of the Second Conference on Machine Translation_ , pages 127–137.
  * Gehring et al. (2017) Jonas Gehring, Michael Auli, David Grangier, Denis Yarats, and Yann N Dauphin. 2017\.  Convolutional sequence to sequence learning.  In _International Conference on Machine Learning_ , pages 1243–1252. PMLR.
  * Gu et al. (2018) Jiatao Gu, Yong Wang, Kyunghyun Cho, and Victor OK Li. 2018.  Search engine guided neural machine translation.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 32.
  * Guo et al. (2020) Junliang Guo, Zhirui Zhang, Linli Xu, Hao-Ran Wei, Boxing Chen, and Enhong Chen. 2020.  Incorporating bert into parallel sequence decoding with adapters.  _arXiv preprint arXiv:2010.06138_.
  * Guu et al. (2020) Kelvin Guu, Kenton Lee, Zora Tung, Panupong Pasupat, and Ming-Wei Chang. 2020.  Realm: Retrieval-augmented language model pre-training.  _arXiv preprint arXiv:2002.08909_.
  * Jegou et al. (2010) Herve Jegou, Matthijs Douze, and Cordelia Schmid. 2010.  Product quantization for nearest neighbor search.  _IEEE transactions on pattern analysis and machine intelligence_ , 33(1):117–128.
  * Jitao et al. (2020) XU Jitao, Josep M Crego, and Jean Senellart. 2020.  Boosting neural machine translation with similar translations.  In _Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics_ , pages 1580–1590.
  * Johnson et al. (2019) Jeff Johnson, Matthijs Douze, and Hervé Jégou. 2019.  Billion-scale similarity search with gpus.  _IEEE Transactions on Big Data_.
  * Kasai et al. (2020) Jungo Kasai, Nikolaos Pappas, Hao Peng, James Cross, and Noah A Smith. 2020.  Deep encoder, shallow decoder: Reevaluating the speed-quality tradeoff in machine translation.  _arXiv preprint arXiv:2006.10369_.
  * Kasai et al. (2021) Jungo Kasai, Hao Peng, Yizhe Zhang, Dani Yogatama, Gabriel Ilharco, Nikolaos Pappas, Yi Mao, Weizhu Chen, and Noah A Smith. 2021.  Finetuning pretrained transformers into rnns.  _arXiv preprint arXiv:2103.13076_.
  * Khandelwal et al. (2020) Urvashi Khandelwal, Angela Fan, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. 2020\.  Nearest neighbor machine translation.  _arXiv preprint arXiv:2010.00710_.
  * Khandelwal et al. (2019) Urvashi Khandelwal, Omer Levy, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. 2019\.  Generalization through memorization: Nearest neighbor language models.  _arXiv preprint arXiv:1911.00172_.
  * Koehn and Knowles (2017) Philipp Koehn and Rebecca Knowles. 2017.  Six challenges for neural machine translation.  _arXiv preprint arXiv:1706.03872_.
  * Lewis et al. (2020a) Mike Lewis, Marjan Ghazvininejad, Gargi Ghosh, Armen Aghajanyan, Sida Wang, and Luke Zettlemoyer. 2020a.  Pre-training via paraphrasing.  _arXiv preprint arXiv:2006.15020_.
  * Lewis et al. (2019) Mike Lewis, Yinhan Liu, Naman Goyal, Marjan Ghazvininejad, Abdelrahman Mohamed, Omer Levy, Ves Stoyanov, and Luke Zettlemoyer. 2019.  Bart: Denoising sequence-to-sequence pre-training for natural language generation, translation, and comprehension.  _arXiv preprint arXiv:1910.13461_.
  * Lewis et al. (2020b) Patrick Lewis, Ethan Perez, Aleksandara Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen-tau Yih, Tim Rocktäschel, et al. 2020b.  Retrieval-augmented generation for knowledge-intensive nlp tasks.  _arXiv preprint arXiv:2005.11401_.
  * Li and Jurafsky (2016) Jiwei Li and Dan Jurafsky. 2016.  Mutual information and diverse decoding improve neural machine translation.  _arXiv preprint arXiv:1601.00372_.
  * Li et al. (2016) Xiaoqing Li, Jiajun Zhang, and Chengqing Zong. 2016.  One sentence one model for neural machine translation.  _arXiv preprint arXiv:1609.06490_.
  * Li et al. (2020) Xiaoya Li, Yuxian Meng, Mingxin Zhou, Qinghong Han, Fei Wu, and Jiwei Li. 2020.  Sac: Accelerating and structuring self-attention via sparse adaptive connection.  _arXiv preprint arXiv:2003.09833_.
  * Lioutas and Guo (2020) Vasileios Lioutas and Yuhong Guo. 2020.  Time-aware large kernel convolutions.  In _International Conference on Machine Learning_ , pages 6172–6183. PMLR.
  * Liu et al. (2020a) Liyuan Liu, Xiaodong Liu, Jianfeng Gao, Weizhu Chen, and Jiawei Han. 2020a.  Understanding the difficulty of training transformers.  _arXiv preprint arXiv:2004.08249_.
  * Liu et al. (2020b) Yinhan Liu, Jiatao Gu, Naman Goyal, Xian Li, Sergey Edunov, Marjan Ghazvininejad, Mike Lewis, and Luke Zettlemoyer. 2020b.  Multilingual denoising pre-training for neural machine translation.  _Transactions of the Association for Computational Linguistics_ , 8:726–742.
  * Luong et al. (2015) Thang Luong, Hieu Pham, and Christopher D. Manning. 2015.  Effective approaches to attention-based neural machine translation.  In _Proceedings of the 2015 Conference on Empirical Methods in Natural Language Processing_ , pages 1412–1421, Lisbon, Portugal. Association for Computational Linguistics.
  * Meng et al. (2019) Yuxian Meng, Xiangyuan Ren, Zijun Sun, Xiaoya Li, Arianna Yuan, Fei Wu, and Jiwei Li. 2019.  Large-scale pretraining for neural machine translation with tens of billions of sentence pairs.  _arXiv preprint arXiv:1909.11861_.
  * Ng et al. (2019) Nathan Ng, Kyra Yee, Alexei Baevski, Myle Ott, Michael Auli, and Sergey Edunov. 2019\.  Facebook fair’s wmt19 news translation task submission.  _arXiv preprint arXiv:1907.06616_.
  * Nguyen and Salazar (2019) Toan Q Nguyen and Julian Salazar. 2019.  Transformers without tears: Improving the normalization of self-attention.  _arXiv preprint arXiv:1910.05895_.
  * Och and Ney (2003) Franz Josef Och and Hermann Ney. 2003.  A systematic comparison of various statistical alignment models.  _Computational linguistics_ , 29(1):19–51.
  * Ott et al. (2019) Myle Ott, Sergey Edunov, Alexei Baevski, Angela Fan, Sam Gross, Nathan Ng, David Grangier, and Michael Auli. 2019.  fairseq: A fast, extensible toolkit for sequence modeling.  _arXiv preprint arXiv:1904.01038_.
  * Peng et al. (2021) Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A Smith, and Lingpeng Kong. 2021.  Random feature attention.  _arXiv preprint arXiv:2103.02143_.
  * Post (2018) Matt Post. 2018.  A call for clarity in reporting BLEU scores.  In _Proceedings of the Third Conference on Machine Translation: Research Papers_ , pages 186–191, Brussels, Belgium. Association for Computational Linguistics.
  * Sutskever et al. (2014) Ilya Sutskever, Oriol Vinyals, and Quoc V Le. 2014.  Sequence to sequence learning with neural networks.  In _Advances in neural information processing systems_ , pages 3104–3112.
  * Tay et al. (2020) Yi Tay, Dara Bahri, Donald Metzler, Da-Cheng Juan, Zhe Zhao, and Che Zheng. 2020\.  Synthesizer: Rethinking self-attention in transformer models.  _arXiv preprint arXiv:2005.00743_.
  * Thulke et al. (2021) David Thulke, Nico Daheim, Christian Dugast, and Hermann Ney. 2021.  Efficient retrieval augmented generation from unstructured knowledge for task-oriented dialog.  _arXiv preprint arXiv:2102.04643_.
  * Tu et al. (2018) Zhaopeng Tu, Yang Liu, Shuming Shi, and Tong Zhang. 2018.  Learning to remember translation history with a continuous cache.  _Transactions of the Association for Computational Linguistics_ , 6:407–420.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Ł ukasz Kaiser, and Illia Polosukhin. 2017.  Attention is all you need.  In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, _Advances in Neural Information Processing Systems 30_ , pages 5998–6008. Curran Associates, Inc.
  * Wang et al. (2019) Qiang Wang, Bei Li, Tong Xiao, Jingbo Zhu, Changliang Li, Derek F Wong, and Lidia S Chao. 2019.  Learning deep transformer models for machine translation.  _arXiv preprint arXiv:1906.01787_.
  * Weston et al. (2018) Jason Weston, Emily Dinan, and Alexander H Miller. 2018.  Retrieve and refine: Improved sequence generation models for dialogue.  _arXiv preprint arXiv:1808.04776_.
  * Xiong et al. (2020a) Lee Xiong, Chenyan Xiong, Ye Li, Kwok-Fung Tang, Jialin Liu, Paul Bennett, Junaid Ahmed, and Arnold Overwijk. 2020a.  Approximate nearest neighbor negative contrastive learning for dense text retrieval.  _arXiv preprint arXiv:2007.00808_.
  * Xiong et al. (2020b) Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. 2020b.  On layer normalization in the transformer architecture.  In _International Conference on Machine Learning_ , pages 10524–10533. PMLR.
  * Zhang et al. (2018) Jingyi Zhang, Masao Utiyama, Eiichro Sumita, Graham Neubig, and Satoshi Nakamura. 2018.  Guiding neural machine translation with retrieved translation pieces.  _arXiv preprint arXiv:1804.02559_.
  * Zhu et al. (2020) Jinhua Zhu, Yingce Xia, Lijun Wu, Di He, Tao Qin, Wengang Zhou, Houqiang Li, and Tie-Yan Liu. 2020.  Incorporating bert into neural machine translation.  _arXiv preprint arXiv:2002.06823_.


# 3D Feature Distillation with Object-Centric Priors

Georgios Tziafas
Department of Artificial Intelligence
University of Groningen, the Neteherlands
<EMAIL_ADDRESS>
&Yucheng Xu
School of Informatics
University of Edinburgh, United Kingdom
<EMAIL_ADDRESS>
&Zhibin Li
Department of Computer Science
University College London, United Kingdom
<EMAIL_ADDRESS>
&Hamidreza Kasaei
Department of Artificial Intelligence
University of Groningen, the Neteherlands
<EMAIL_ADDRESS>

###### Abstract

Grounding natural language to the physical world is a ubiquitous topic with a
wide range of applications in computer vision and robotics. Recently, 2D
vision-language models such as CLIP have been widely popularized, due to their
impressive capabilities for open-vocabulary grounding in 2D images. Recent
works aim to elevate 2D CLIP features to 3D via feature distillation, but
either learn neural fields that are scene-specific and hence lack
generalization, or focus on indoor room scan data that require access to
multiple camera views, which is not practical in robot manipulation scenarios.
Additionally, related methods typically fuse features at pixel-level and
assume that all camera views are equally informative. In this work, we show
that this approach leads to sub-optimal 3D features, both in terms of
grounding accuracy, as well as segmentation crispness. To alleviate this, we
propose a multi-view feature fusion strategy that employs object-centric
priors to eliminate uninformative views based on semantic information, and
fuse features at object-level via instance segmentation masks. To distill our
object-centric 3D features, we generate a large-scale synthetic multi-view
dataset of cluttered tabletop scenes, spawning 15k scenes from over 3300
unique object instances, which we make publicly available. We show that our
method reconstructs 3D CLIP features with improved grounding capacity and
spatial consistency, while doing so from single-view RGB-D, thus departing
from the assumption of multiple camera views at test time. Finally, we show
that our approach can generalize to novel tabletop domains and be re-purposed
for 3D instance segmentation without fine-tuning, and demonstrate its utility
for language-guided robotic grasping in clutter.

Figure 1: Visualization of 3D features (middle), back-projected 2D features
(left) and user query similarity heatmaps (right), for previous SOTA point-
cloud feature distillation method OpenScene and our DROP-CLIP. OpenScene fuses
pixel-wise 2D features across all views with average pooling, leading to
grounding failures, segmentation imprecisions and fuzzy object boundaries. Our
method resolves these issues by employing object-centric priors to fuse
object-level 2D features in 3D instance masks with semantics-informed view
selection.

> Keywords: Open-Vocabulary 3D Segmentation, Multi-view Feature Distillation

##  1 Introduction

Language grounding in 3D environments plays a crucial role in realizing
intelligent systems that can interact naturally with the physical world. In
the robotics field, being able to precisely segment desired objects in 3D
based on open language queries (object semantics, visual attributes,
affordances, etc.) can serve as a powerful proxy for enabling open-ended robot
manipulation. As a result, research focus on 3D segmentation methods has seen
growth in recent years [1, 2, 3, 4, 5, 6]. However, related methods fall in
the closed-vocabulary regime, where only a fixed list of classes can be used
as queries. Inspired by the success of open-vocabulary 2D methods [7, 8, 9,
10], recent efforts elevate 2D representations from pretrained models [7, 11]
to 3D via distillation pipelines [12, 13, 14, 15, 16, 17, 18, 19]. However, we
identify certain limitations of existing distillation approaches. On the one
hand, field-based methods [13, 20, 16, 17, 18] offer continuous 3D feature
fields, but require to be trained online in specific scenes and hence cannot
generalize to novel object instances and compositions, they require a few
minutes to train, and need to collect multiple camera views before training,
all of which hinder their real-time applicability. On the other hand, original
3D feature distillation methods and follow up work [12, 14, 21] use room scan
datasets [22, 23] to learn point-cloud encoders, hence being applicable in
novel scenes with open vocabularies. However, such approaches assume that 2D
features from all views are equally informative, which is not the case in
highly cluttered indoors scenes (e.g. due to partial occlusions from some
view), thus leading in noisy 3D features. 2D features are also usually fused
point-wise from ViT patches [9, 10, 8] or multi-scale crops [13, 6], therefore
leading to the so called “patchyness” issue [24] (see Fig. 1). The latter
issue is especially impactful in robot manipulation, where precise 3D
segmentation is vital for specifying robust actuation goals.

To address such limitations, we revisit 2D $\rightarrow$ 3D point-based
feature distillation but revise the multi-view feature fusion strategy to
enhance the quality of the target 3D features. In particular, we inject both
semantic and spatial object-centric priors into the fusion strategy, in three
ways: (i) We obtain object-level 2D features by isolating object instances in
each camera view from their 2D segmentation masks, (ii) we fuse features only
at corresponding 3D object regions using 3D segmentation masks, (iii) we
leverage object-level semantic information to devise an informativeness
metric, which is used to weight the contribution of views and eliminate
uninformative ones. Extensive ablation studies demonstrate the advantages of
object-centric fusion compared to vanilla approaches. To train our method, we
require a large-scale cluttered indoors dataset with many views per scene,
which is currently not existent. To that end, we build MV-TOD (Multi-View
Tabletop Objects Dataset), consisting of $\sim 15k$ Blender scenes from more
than $3.3k$ unique 3D object models, for which we provide $73$ views per scene
with $360^{\circ}$ coverage, further equipped with 2D/3D segmentations, 6-DoF
grasps and textual object-level annotations. We use MV-TOD to distill our
object-centric 3D CLIP [7] features into a 3D representation, which we call
DROP-CLIP (Distilled Representations with Object-centric Priors from CLIP).
Our 3D encoder operates in partial point-clouds from a single RGB-D view, thus
departing from the requirement of multiple camera images at test time, while
offering real-time inference capabilities. We demonstrate that our learned 3D
features achieve high grounding performance and segmentation crispness, while
significantly outperforming previous 2D open-vocabulary approaches in the
single-view setting. Further, we show that they can be leveraged zero-shot in
novel tabletop domains, as well as be used out-of-the-box for 3D instance
segmentation.

In summary, our contributions are fourfold: (i) we release MV-TOD, a large-
scale synthetic dataset of household objects in cluttered tabletop scenarios,
featuring dense multi-view coverage and semantic/mask/grasp annotations, (ii)
we identify limitations of current multi-view feature fusion approaches and
illustrate how to overcome them by leveraging object-centric priors, (iii) we
release DROP-CLIP, a 3D model that reconstructs view-independent 3D CLIP
features from single-view, and (iv) we conduct extensive ablation studies,
comparative experiments and robot demonstrations to showcase the effectiveness
of the proposed method in terms of 3D segmentation performance, generalization
to novel domains and tasks, and applicability in robot manipulation scenarios.

##  2 Multi-View Tabletop Objects Dataset

Figure 2: (Left:) Example generated Blender scene, multi-view image coverage
and annotations included in our dataset. (Right:) Automatic semantic
annotation generation with large vision-language models.

Dataset | Layout | Multi | Clutter | Vision | Ref.Expr. | Grasp | Num.Obj. | Num. | Num. | Obj.-lvl
---|---|---|---|---|---|---|---|---|---|---
View | Data | Annot. | Annot. | Categories | Scenes | Expr. | Semantics
ScanNet [22] | indoor | ✔ | - | RGB-D,3D | ✗ | ✗ | $17$ | $800$ | $-$ | ✗
S3DIS [25] | indoor | ✔ | - | RGB-D,3D | ✗ | ✗ | $13$ | $6$ | $-$ | ✗
Replica [26] | indoor | ✔ | - | RGB-D,3D | ✗ | ✗ | $88$ | $-$ | $-$ | ✔
STPLS3D [25] | outdoor | ✔ | - | 3D | ✗ | ✗ | $12$ | $18$ | $-$ | ✔
ScanRefer [1] | indoor | ✔ | ✗ | RGB-D,3D | 2D/3D mask | ✗ | $18$ | $800$ | $51.5k$ | ✗
ReferIt-3D [2] | indoor | ✔ | ✗ | RGB-D,3D | 2D/3D mask | ✗ | $18$ | $707$ | $125.5k$ | ✗
ReferIt-RGBD [27] | indoor | ✔ | ✗ | RGB-D | 2D box | ✗ | - | $7.6k$ | $38.4k$ | ✗
SunSpot [28] | indoor | ✗ | ✔ | RGB-D | 2D box | ✗ | 38 | $1.9k$ | $7.0k$ | ✗
GraspNet [29] | tabletop | ✗ | ✔ | 3D | ✗ | 6-DoF | $88$ | $190$ | $-$ | ✗
REGRAD [30] | tabletop | ✔ | ✔ | RGB-D,3D | ✗ | 6-DoF | $55$ | $47k$ | $-$ | ✗
OCID-VLG [31] | tabletop | ✗ | ✔ | RGB-D,3D | 2D mask | 4-DoF | $31$ | $1.7k$ | $89.6k$ | template
Grasp-Anything [32] | tabletop | ✗ | ✗ | RGB | 2D mask | 4-DoF | $236$ | $1M$ | $-$ | open
MV-TOD (ours) | tabletop | ✔ | ✔ | RGB-D,3D | 3D mask | 6-DoF | $149$ | $15k$ | $671.2k$ | open

Table 1: Comparisons between MV-TOD and existing 3D datasets.

Existing 3D datasets mainly focus on indoor scenes in room layouts [33, 22,
26] and related language annotations typically cover closed-set object
categories (e.g. furniture) and spatial relations [1, 2, 27, 34, 28], which
are not practical for robot manipulation tasks, where cluttered tabletop
scenarios and open-vocabulary language are of key importance. On the other
hand, recent grasp-related efforts collect cluttered tabletop scenes, but
either lack language annotations [30, 35, 29] or connect cluttered scenes with
language but only for 4-DoF grasps with RGB data [31, 32]. Further, most of
such datasets lack dense multi-view scene coverage, granting them non
applicable for 2D $\rightarrow$ 3D feature distillation, where we require
multiple images from each scene to extract 2D features with a foundation
model. To cover this gap, we propose MV-TOD, a large-scale synthetic dataset
with cluttered tabletop scenes featuring dense multi-view coverage and rich
language annotations at the object level. We generate a total of $15k$ scenes
in Blender [36], comprising of 3379 unique object models, $99$ collected by us
and the rest filtered from ShapeNet-Sem model set [37]. The dataset features
$149$ object categories, each of which includes multiple instances that vary
in fine-grained details. For each object instance, we leverage GPT-4-Vision
[38] to generate open-set descriptions from various perspectives, including
category, color, material, state, utility, affordance, etc, which spawn over
$670k$ unique referring instance queries (see Fig. 2-right and Appendix A).
For each scene, we provide 2D/3D segmentation masks, 6D object poses, as well
as a set of semantic concepts for each appearing object instance.
Additionally, we include 6-DoF grasp annotations for each object model,
originating from the ACRONYM dataset [35]. To the best of our knowledge, MV-
TOD is the first dataset to combine 3D cluttered tabletop scenes with open-
vocabulary language and 6-DoF grasp annotations, which we hope will accelerate
future research.

##  3 Methodology

Figure 3: Method Overview: Given a 3D scene and multiple camera views, we
employ three object-centric priors (in red) for multi-view feature fusion: (i)
extract CLIP features from 2D masked object crops, (ii) use semantic
annotations to fuse 2D features across views, (iii) apply the fused feature on
all points in the object’s 3D mask. The fused feature-cloud is distilled with
a single-view posed RGB-D encoder and cosine distance loss. During inference,
we compute point-wise cosine similarity scores in CLIP space (higher
similarity towards red).

Our goal is to distill multi-view 2D CLIP features into a 3D representation,
while employing an object-centric feature fusion strategy to ensure high
quality 3D features. Our overall pipeline is illustrated in Fig 3. We first
introduce traditional multi-view feature fusion (Sec. 3.1), present our
variant with object-centric priors (Sec. 3.2) and discuss our feature
distillation method (Sec. 3.3).

###  3.1 Multi-view 2D Feature Fusion

We assume access to a dataset of 3D scenes, where each scene is represented
through a set of $\mathcal{V}$ posed RGB-D views of size $H\times W$:
$\left\\{I_{v}\in\mathbb{R}^{H\times W\times 3},\,D_{v}\in\mathbb{R}^{H\times
W},\,T_{v}\in\mathbb{R}^{4\times 4}\right\\}_{v=1}^{\mathcal{V}}$, where
$T_{v}$ the camera pose from view $v$. For each scene, we first obtain the
full point-cloud $P\in\mathbb{R}^{M\times 3}$ along with a 2D-3D
correspondence map $\mathcal{M}_{v}\in\mathbb{R}^{M\times 2}$, mapping each
point $\mathbf{x}_{i},\,i=1,\dots,M$ to a pixel location
$\mathbf{u}_{v,i}=(u_{x},u_{y})^{T}$ in image view $I_{v}$. RGB images are fed
to a pretrained image model $f^{2D}:\mathbb{R}^{H\times W\times
3}\rightarrow\mathbb{R}^{H\times W\times C}$ [9, 10, 8] to obtain pixel-level
2D features of size C: $Z^{2D}_{v}=f^{2D}(I_{v})$, which are then back-
projected to 3D points via:

$\textbf{z}^{2D}_{v,i}=f^{2D}\left(I_{v}(\textbf{u}_{v,i})\right)=f^{2D}\left(I_{v}(\mathcal{M}_{v}(\mathbf{x}_{i}))\right)$
(1)

To fuse 2D features across views $Z^{3D}\in\mathbb{R}^{M\times C}$, previous
works [12, 6, 14, 21] use average pooling:
$\mathbf{z}_{i}^{3D}=\frac{1}{\mathcal{V}}\sum_{v=1}^{\mathcal{V}}\mathbf{z}_{v,i}^{2D}$
(see Appendix B for a comprehensive overview). In essence, this method assumes
that all views are equally informative for each point, as long as the point is
visible from that view. We suggest that naively average pooling 2D features
for each point leads to sub-optimal 3D features, as noisy, uninformative views
contribute equally, therefore “polluting" the overall representation. We then
propose to instead use a generalized version relying on weighted average:

$\mathbf{z}^{3D}_{i}=\frac{\sum_{v=1}^{\mathcal{V}}\mathbf{z}^{2D}_{v,i}\cdot\omega_{v,i}}{\sum_{v=1}^{\mathcal{V}}\omega_{v,i}}$
(2)

where $\omega_{v,i}\in\mathbb{R}$ is a scalar weight that represents the
informativeness of view $v$ for point $i$. In the next subsection, we describe
how to use text data to dynamically compute an informativeness weight for each
view based on semantic object-level semantic information. Additionally,
vanilla point-wise fusion with pixel-level 2D features leads to non-crisp
segmentations and fuzzy object boundaries. To resolve this, we propose to also
leverage dense spatial information, i.e., instance-wise 2D/3D segmentation
masks, which are used for both: (a) obtaining robust object-level 2D CLIP
features from each view, and (b) fusing features only at the points
corresponding to the 3D object region.

###  3.2 Employing Object-Centric Priors

Let $\left\\{S_{v}^{2D}\in\\{0,1\\}^{N\times H\times
W}\right\\}_{v=1}^{\mathcal{V}}$ be 2D instance-wise segmentation masks for
each scene, where $N$ the total number of scene objects. We aggregate the 2D
masks to obtain $S^{3D}\in\\{0,1\\}^{M\times N}$, such that for each point $i$
we can retrieve the corresponding object instance
$n_{i}=\texttt{argmax}_{\nu}\;S^{3D}_{i,\nu}$.

Semantic informativeness metric Let
$\mathcal{Q}=\left\\{Q_{k}\right\\}_{k=1}^{\mathcal{K}},\;Q_{k}\in\mathbb{R}^{N_{k}\times
C}$ be a set of object-specific textual prompts, where $\mathcal{K}$ the
number of dataset object instances and $N_{k}$ the number of prompts for
object $k$. We use CLIP’s text encoder to embed the textual prompts in
$\mathbb{R}^{C}$ and average them to obtain an object-specific prompt
$\textbf{q}_{k}=1/N_{k}\cdot\sum_{j=1}^{N_{k}}Q_{k,j}$. For each scene, we map
each object instance $n\in[1,N]$ to its positive prompt $\textbf{q}_{n}^{+}$,
as well as a set $Q^{-}_{n}\doteq\mathcal{Q}-\\{\textbf{q}_{n}^{+}\\}$ of
negative prompts corresponding to all other instances. We define our semantic
informativeness metric as:

$G_{v,i}=\texttt{cos}(\textbf{z}^{2D}_{v,i},\textbf{q}^{+}_{n_{i}})-\texttt{max}_{\textbf{q}\sim
Q^{-}_{n_{i}}}\texttt{cos}(\textbf{z}^{2D}_{v,i},\textbf{q})$ (3)

Intuitively, we want a 2D feature from view $v$ to contribute to the overall
3D feature of point $i$ according to how much its similarity with the correct
object instance is higher than the maximum similarity to any of the negative
object instances, hence offering a proxy for semantic informativeness. We clip
this weight to 0 to eliminate views that don’t satisfy the condition
$G_{v,i}\geq 0$. Plugging in our metric in equation (2) already provides
improvements over vanilla average pooling (see Sec. 4.1), however, does not
deal with 3D spatial consistency, for which we employ our spatial priors
below.

Object-level 2D CLIP features For obtaining object-level 2D CLIP features, we
isolate the pixels for each object $n$ from each view $v$ from $S^{2D}_{v,n}$
and crop a bounding box around the mask from $I_{v}$:
$\textbf{z}^{2D}_{v,n}=f^{2D}_{cls}\left(\texttt{cropmask}(I_{v},\,S^{2D}_{v,n})\right)$
(see Appendix C for ablations in CLIP visual prompts). Here we use
$f^{2D}_{cls}:\mathbb{R}^{h_{n}\times w_{n}\times
3}\rightarrow\mathbb{R}^{C}$, i.e., only the [CLS] feature of CLIP’s ViT
encoder, to represent an object crop of size $h_{n}\times w_{n}$. We can now
define our metric from equation (3) also at object-level:

$G_{v,n}=\texttt{cos}(\textbf{z}^{2D}_{v,n},\textbf{q}^{+}_{n})-\texttt{max}_{\textbf{q}\sim
Q^{-}_{n}}\texttt{cos}(\textbf{z}^{2D}_{v,n},\textbf{q})$ (4)

where $G_{v,n}\in\mathbb{R}^{\mathcal{V}\times N}$ now represents the semantic
informativeness of view $v$ for object instance $n$.

Fusing object-wise features A 3D object-level feature can be obtained by
fusing 2D object-level features across views similar to equation (2):

$\mathbf{z}^{3D}_{n}=\frac{\sum_{v=1}^{\mathcal{V}}\mathbf{z}^{2D}_{v,n}\cdot\omega_{v,n}}{\sum_{v=1}^{\mathcal{V}}\omega_{v,n}}=\frac{\sum_{v=1}^{\mathcal{V}}\mathbf{z}^{2D}_{v,n}\cdot\Lambda_{v,n}\cdot
G_{v,n}}{\sum_{v=1}^{\mathcal{V}}\Lambda_{v,n}\cdot G_{v,n}}$ (5)

where each view is weighted by its semantic informativeness metric $G_{v,n}$,
as well as optionally a visibility metric
$\Lambda_{v,n}=\sum_{\textbf{}}S^{2D}_{v,n}$ that measures the number of
pixels from $n$-th object’s mask that are visible from view $v$ [6]. We
finally reconstruct the full feature-cloud $Z^{3D}\in\mathbb{R}^{M\times C}$
by equating each point’s feature to its corresponding 3D object-level one via:
$\mathbf{z}^{3D}_{i}=\mathbf{z}^{3D}_{n_{i}},\;n_{i}=\texttt{argmax}_{\nu}\;S^{3D}_{i,\nu}$.

###  3.3 View-Independent Feature Distillation

Even though the above feature-cloud $Z^{3D}$ could be directly used for open-
vocabulary grounding in 3D, its construction is computationally intensive and
requires a lot of expensive resources, such as access to multiple camera
views, view-aligned 2D instance segmentation masks, as well as a set of text
descriptions to compute informativeness metrics. Such utilities are rarely
available in open-ended scenarios, especially in robotic applications, where
usually only single-view RGB-D images from sensors mounted on the robot are
provided. To tackle this, we wish to distill all the above knowledge from the
feature-cloud $Z^{3D}$ into a 3D encoder that receives only a partial point-
cloud from single-view posed RGB-D. Hence, the only assumption that we make
during inference is access to camera intrinsic and extrinsic parameters, which
is a mild requirement in most robotic works.

In particular, given a partial colored point-cloud from view $v$:
$P_{v}\in\mathbb{R}^{M_{v}\times 6}$ (3D coordinates plus colors), we train a
3D encoder $\mathcal{E}_{\theta}:\mathbb{R}^{M_{v}\times
6}\rightarrow\mathbb{R}^{M_{v}\times C}$ such that
$\mathcal{E}_{\theta}(P_{v})=Z^{3D}$. Notice that the distillation target
$Z^{3D}$ is independent of view $v$. Following [12, 15] we use cosine distance
loss:

$\mathcal{L}(\theta)=1-\texttt{cos}(\mathcal{E_{\theta}}(P_{v}),\,Z^{3D})$ (6)

See Appendix B.2 for training implementation details. With such a setup, we
can obtain 3D features that: (i) are co-embedded in CLIP text space, so they
can be leveraged for 3D segmentation tasks from open-vocabulary queries via
computing cosine similarities between CLIP text embeddings $Q$ and the
predicted feature cloud:
$\hat{S}_{i}=\texttt{argmax}_{i}\;\texttt{cos}(\hat{\mathbf{z}}_{i}^{3D},Q)$,
(ii) are ensured to be optimally informative per object, due to the usage of
the semantic informativeness metric to compute $Z^{3D}$, (iii) maintain 3D
spatial consistency in object boundaries, due to performing object-wise
instead of point-wise fusion when computing $Z^{3D}$, and (iv) are encouraged
to be view-independent, as the same features $Z^{3D}$ are utilized as
distillation targets regardless of the input view $v$. Importantly, no labels,
prompts, or segmentation masks are needed at test-time to reproduce the fused
feature-cloud, while obtaining it amounts to a single forward pass of our 3D
encoder, hence offering real-time performance.

Figure 4: Open-Vocabulary 3D Referring Segmentation. Examples of learned 3D
features and grounding heatmaps from open-ended language queries (class names,
attributes, user affordances, and open instance-specific concepts) in scenes
from MV-TOD dataset. Points are colored based on their query similarity
(higher towards red). We note that table points are excluded from similarity
computation in our visualizations.

##  4 Experiments

In our experiments, we explore the following questions: (i) Sec. 4.1: What are
the contributions of our proposed object-centric priors for multi-view feature
fusion? Does the dense number of views of our proposed dataset also
contribute? (ii) Sec. 4.2: How does our method compare to previous open-
vocabulary approaches for 3D semantic and referring segmentation tasks? Are
the learned features robust to open-ended language? (iii) Sec. 4.3: What are
the generalization capabilities of our learned 3D representation in novel
domains and novel tasks (3D instance segmentation)? (iv) Sec. 4.4: Can we
leverage our 3D learned representation for language-guided 6-DoF robotic
grasping?

###  4.1 Multi-view Feature Fusion Ablation Studies

Fusion | $\mathbf{f^{2D}}$ | $\mathbf{\Lambda_{v,i}}$ | $\mathbf{G_{v,i}}$ |  Ref.Segm (%)
---|---|---|---|---
mIoU | Pr@25 | Pr@50 | Pr@75
point | patch |  |  | 44.2 | 59.9 | 41.4 | 27.0
point | patch | ✓ |  | 37.3 | 55.4 | 33.7 | 16.7
point | patch |  | ✓ | 57.0 | 74.1 | 59.5 | 40.9
point | patch | ✓ | ✓ | 57.4 | 77.0 | 60.9 | 39.9
obj | obj |  |  | 65.6 | 67.0 | 65.4 | 64.1
obj | obj | ✓ |  | 67.3 | 68.7 | 67.1 | 65.8
obj | obj |  | ✓ | 83.1 | 83.9 | 83.1 | 82.4
obj | obj | ✓ | ✓ | 80.9 | 83.1 | 80.2 | 79.7

Table 2: Multi-view feature fusion ablation study for 3D referring
segmentation in MV-TOD.

To evaluate the contributions of our proposed object-centric priors, we
conduct ablation studies on the multi-view feature fusion pipeline, where we
compare 3D referring segmentation results of obtained 3D features in held-out
scenes of MV-TOD. We highlight that here we aim to establish a performance
upper bound that the feature fusion method can provide for distillation, and
not the distilled features themselves.

Figure 5: Referring segmentation accuracy (Pr@25 (%)) vs. number of utilized
views.

We ablate: (i) patch-wise vs. object-wise fusion, (ii) MaskCLIP [8] patch-
level vs. CLIP [7] masked crop features, (iii) inclusion of visibility
($\Lambda_{v,i}$) and semantic informativeness ($G_{v,i}$) metrics for view
selection. Results in Table 2.

Effect of object-centric priors We observe that all components contribute
positively to the quality of the 3D features. Our proposed $G_{v,i}$ metric
boosts mIoU across both point- and object-wise fusion ($57.0\%$ vs. $44.2\%$
and $83.1\%$ vs. $65.6\%$ respectively). Further, we observe that the usage of
spatial priors for object-wise fusion and object-level features leads to both
higher segmentation crispness ($25.7\%$ mIoU delta), as well as higher
grounding precision ($42.5\%$ Pr@75 delta). See qualitative comparisons in
Appendix D.

Effect of the number of views We ablate the 3D referring segmentation
performance based on the number of input views in Fig. 5, where novel
viewpoints are added incrementally. We observe that in both setups (point- and
object-wise) fusing features from more views leads to improvements, with a
small plateauing behavior around 40 views. We believe this is an encouraging
result for leveraging dense multi-view coverage in feature distillation
pipelines, as we propose with the introduction of MV-TOD.

###  4.2 Open-Vocabulary 3D Segmentation Results

Method |  #views |  Ref.Segm. (%) |  Sem.Segm (%)
---|---|---|---
mIoU | Pr@25 | Pr@50 | Pr@75 | mIoU | mAcc
OpenScene† [12] |  73 | 29.32 | 44.00 | 24.51 | 11.26 | 21.79 | 32.14
OpenMask3D∗† [6] |  73 | 65.38 | 73.05 | 63.99 | 57.40 | 59.47 | 66.48
DROP-CLIP† (Ours) |  73 | 82.67 | 86.11 | 82.43 |  79.23 | 75.41 | 80.02
DROP-CLIP (Ours) |  73 | 66.56 | 75.73 | 67.55 |  59.88 | 62.04 | 70.74
OpenSeg→3D [9] |  1 | 12.89 | 17.36 | 2.38 | 0.23 | 12.83 | 17.21
MaskCLIP→3D [8] |  1 | 25.64 | 40.36 | 18.69 | 6.95 | 20.97 | 32.09
DROP-CLIP (Ours) |  1 | 62.31 | 71.96 | 62.75 |  53.85 | 54.48 | 64.41

Table 3: Referring and Semantic segmentation results on MV-TOD test split.
Methods with † denote upper-bound 3D features, whereas DROP-CLIP denotes our
distilled model. Methods with →3D produce 2D predictions that are projected to
3D to compute metrics. Method with $*$ denotes further usage of ground-truth
segmentation masks.

In this section, we compare referring and semantic segmentation performance of
our distilled features vs. previous open-vocabulary approaches, both in multi-
view and in single-view settings. For multi-view, we compare our trained model
with OpenScene [12] and OpenMask3D [6] methods, where the full point-cloud
from all $73$ views is given as input.

Figure 6: Referring segmentation accuracy (Pr@25 (%)) vs. different language
query types.

We note that for these baselines we obtain the upper-bound 3D features as
before, as we observed that our trained model already outperforms them, so we
refrained from also distilling features from baselines (details in Appendix
C2). For single-view, we feed our network with partial point-cloud from
projected RGB-D pair, and compare with 2D baselines MaskCLIP [8] and OpenSeg
[9]. Our model slightly outperforms the OpenMask3D upper bound baseline in the
multi-view setting ($+1.18\%$ in referring and $+2.57\%$ in semantic
segmentation), while significantly outperforming 2D baselines in the single-
view setting ($>30\%$ in both tasks). Importantly, single-view results closely
match the multi-view ones ($\sim-4.0\%$), suggesting that DROP-CLIP indeed
learns view-independent features.

Open-ended queries We evaluate the robustness of our model in different types
of input language queries, organized in 4 families (class name - e.g. “cereal
box", class + attribute \- e.g. “brown cereal box", open \- e.g. “chocolate
Kellogs", and affordance \- e.g. “I want something sweet‘). Comparative
results are presented in Fig. 6 and qualitative in Fig. 4. We observe that our
method achieves high grounding accuracy in all query types, even when using
single-view.

###  4.3 Zero-Shot Transfer to Novel Domains / Tasks

Method | OCID-VLG [31] | REGRAD [30]
---|---|---
IoU | Pr@25 | IoU | Pr@25
MaskCLIP→3D [8] | 24.1 | 30.9 | 33.2 | 39.0
DROP-CLIP (Ours) | 46.2 | 48.9 | 59.1 | 63.0

Table 4: Referring segmentation results in OCID-VLG [31] and REGRAD [30]
datasets

Generalization to Novel Domains We evaluate the 3D referring segmentation
performance of our trained model when applied zero-shot in novel tabletop
domains. We test in 500 scenes from OCID-VLG [31] using the dataset’s
instance-wise open queries, as well as in 1000 scenes from REGRAD [30], using
each model’s class name as a query. Only single-view input is provided for
both datasets.

Method | $mIoU$ | $AP_{25}$ | $AP_{50}$
---|---|---|---
SAM [39] | 70.11 | 95.26 | 79.88
DROP-CLIP (S) | 80.83 | 91.92 | 86.83
Mask3D [40] | 14.41 | 18.65 | 3.41
DROP-CLIP (F) | 88.37 | 93.13 | 91.47

Table 5: Zero-shot 3D instance segmentation results in MV-TOD.

We compare with MaskCLIP [8] as above and report results in Table 4. We note
that test datasets contain both novel object instances (REGRAD) and classes
(OCID-VLG). We observe that our method provides a significant performance
boost across both domains ($22.1\%$ mIoU delta in OCID-VLG and $25.9\%$ in
REGRAD).

Zero-Shot 3D Instance Segmentation Since our method has been distilled from
features with object-level priors, we demonstrate that it can be used out-of-
the-box for 3D instance segmentation, via clustering the 3D features (see
Appendix E for implementation details). We report results in MV-TOD in Table
5, where we compare with SAM [39] with single-view images, as well as Mask3D
[41] with full point-clouds (transferred from ScanRefer [1] with room layout).
Mask3D struggles to generalize to tabletop domains, whereas our method
achieves comparable performance with SAM for segmenting from single-view, even
without being explicitly trained for instance segmentation.

###  4.4 Open-Vocabulary Language-guided Robotic Grasping

Figure 7: Language-guided 6-DoF grasping trial with real robot (left), 3D
features, grounding and grasp proposals (right).

.

In this section, we wish to illustrate the applicability of DROP-CLIP in a
language-guided robotic grasping scenario. We integrate our method with a
6-DoF grasp detection network [42], to segment and then propose gripper poses
for picking a target object indicated verbally. We randomly place 5-12 objects
on a tabletop with different levels of clutter, and query the robot to pick
the target object and place it in a fixed position. The user instruction is
open-vocabulary and can involve open object descriptions, attributes, or
affordances. We conducted 50 trials in Gazebo [43] and 10 with a real robot,
and observed grounding accuracy of 84% and 80% respectively, and a final
success rate of 64% and 60%, where failures were mostly due to grasp proposals
that are outside of the robot’s kinematic range or motion planning that lead
to a collision with other objects and the table. Our setup and example trials
are shown in Fig. 7, while more details and qualitative results are provided
in Appendix E. A video of robot demonstrations is provided as supplementary
material.

##  5 Related work

3D Scene Understanding There’s a long line of works in closed-set 3D scene
understanding [44, 45, 46, 47, 48, 49], applied in 3D classification [50, 51],
localization [52, 1] and segmentation [53, 23, 22], using two-stage pipelines
with instance proposals from point-clouds [54, 55] or RGB-D views [56, 27], or
single-stage methods [3] that leverage 3D-language cross attentions. [57] use
CLIP embeddings for pretraining a 3D segmentation model, but still cannot be
applied open-vocabulary.

Open-Vocabulary Grounding with CLIP Following the impressive results of CLIP
[7] for open-set image recognition, followup works transfer CLIP’s powerful
representations from image- to pixel-level [40, 58, 59, 60, 61, 62, 63, 9, 10,
8], extending to detection / segmentation, but limited to 2D. For 3D
segmentation, the closest work is perhaps OpenMask3D [6] that extracts multi-
view CLIP features from instance proposals from Mask3D [41] to compute
similarities with open text queries.

3D CLIP Feature Distillation Recent works distill features from 2D foundation
models with point-cloud encoders [12, 14, 21] or neural fields [13, 19, 17,
18, 19, 24], with applications in robot manipulation [20, 16] and navigation
[64, 65]. However, associated works extract 2D features from OpenSeg [9], LSeg
[10], MaskCLIP [8] or multi-scale crops from CLIP [7] and fuse point-wise with
average pooling, while our approach leverages semantics-informed view
selection and segmentation masks to do object-wise fusion with object-level
features (see detailed overview in Appendix F).

##  6 Conclusion, Limitations and Future Work

We propose DROP-CLIP, a 2D$\rightarrow$3D CLIP feature distillation framework
that employs object-centric priors to select views based on semantic
informativeness and ensure crisp 3D segmentations, while working with single-
view RGB-D. We also release MV-TOD, a large-scale synthetic dataset of multi-
view tabletop scenes with dense annotations that can be leveraged for several
downstream tasks. We hope our work can benefit the robotics community, both in
terms of released resources as well as illustrating and overcoming theoretical
limitations of existing 3D feature distillation works.

While our spatial object-centric priors lead to improved segmentation quality,
they collapse local features in favor of a global object-level feature, and
hence cannot be applied for segmenting object parts. In the future, we plan to
add object part annotations in our dataset and fuse with both object- and
part-level masks. Second, DROP-CLIP only provides grounding and a two-stage
pipeline is needed for grasping, while our dataset already provides rich 6-DoF
grasp annotations. A next step would be to also distill them, opting for a
joint 3D representation for grounding and grasping.

## References

  * Chen et al. [2020] D. Z. Chen, A. X. Chang, and M. Nießner.  Scanrefer: 3d object localization in rgb-d scans using natural language.  In _Computer Vision–ECCV 2020: 16th European Conference, Glasgow, UK, August 23–28, 2020, Proceedings, Part XX 16_ , pages 202–221. Springer, 2020.
  * Achlioptas et al. [2020] P. Achlioptas, A. Abdelreheem, F. Xia, M. Elhoseiny, and L. Guibas.  Referit3d: Neural listeners for fine-grained 3d object identification in real-world scenes.  _16th European Conference on Computer Vision (ECCV)_ , 2020.
  * Luo et al. [2022] J. Luo, J. Fu, X. Kong, C. Gao, H. Ren, H. Shen, H. Xia, and S. Liu.  3d-sps: Single-stage 3d visual grounding via referred point progressive selection.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 16454–16463, 2022.
  * Huang et al. [2021] P.-H. Huang, H.-H. Lee, H.-T. Chen, and T.-L. Liu.  Text-guided graph neural networks for referring 3d instance segmentation.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 35, pages 1610–1618, 2021.
  * Qian et al. [2024] Z. Qian, Y. Ma, J. Ji, and X. Sun.  X-refseg3d: Enhancing referring 3d instance segmentation via structured cross-modal graph neural networks.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 38, pages 4551–4559, 2024.
  * Takmaz et al. [2023] A. Takmaz, E. Fedele, R. W. Sumner, M. Pollefeys, F. Tombari, and F. Engelmann.  Openmask3d: Open-vocabulary 3d instance segmentation.  _ArXiv_ , abs/2306.13631, 2023.  URL https://api.semanticscholar.org/CorpusID:259243888.
  * Radford et al. [2021] A. Radford, J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, S. Agarwal, G. Sastry, A. Askell, P. Mishkin, J. Clark, G. Krueger, and I. Sutskever.  Learning transferable visual models from natural language supervision.  _CoRR_ , abs/2103.00020, 2021.  URL https://arxiv.org/abs/2103.00020.
  * Dong et al. [2022] X. Dong, Y. Zheng, J. Bao, T. Zhang, D. Chen, H. Yang, M. Zeng, W. Zhang, L. Yuan, D. Chen, F. Wen, and N. Yu.  Maskclip: Masked self-distillation advances contrastive language-image pretraining.  _2023 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , pages 10995–11005, 2022.  URL https://api.semanticscholar.org/CorpusID:251799827.
  * Ghiasi et al. [2021] G. Ghiasi, X. Gu, Y. Cui, and T.-Y. Lin.  Scaling open-vocabulary image segmentation with image-level labels.  In _European Conference on Computer Vision_ , 2021.  URL https://api.semanticscholar.org/CorpusID:250895808.
  * Li et al. [2022] B. Li, K. Q. Weinberger, S. J. Belongie, V. Koltun, and R. Ranftl.  Language-driven semantic segmentation.  _ArXiv_ , abs/2201.03546, 2022.  URL https://api.semanticscholar.org/CorpusID:245836975.
  * Oquab et al. [2023] M. Oquab, T. Darcet, T. Moutakanni, H. Q. Vo, M. Szafraniec, V. Khalidov, P. Fernandez, D. Haziza, F. Massa, A. El-Nouby, M. Assran, N. Ballas, W. Galuba, R. Howes, P.-Y. B. Huang, S.-W. Li, I. Misra, M. G. Rabbat, V. Sharma, G. Synnaeve, H. Xu, H. Jégou, J. Mairal, P. Labatut, A. Joulin, and P. Bojanowski.  Dinov2: Learning robust visual features without supervision.  _ArXiv_ , abs/2304.07193, 2023.  URL https://api.semanticscholar.org/CorpusID:258170077.
  * Peng et al. [2022] S. Peng, K. Genova, ChiyuMaxJiang, A. Tagliasacchi, M. Pollefeys, and T. A. Funkhouser.  Openscene: 3d scene understanding with open vocabularies.  _2023 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , pages 815–824, 2022.  URL https://api.semanticscholar.org/CorpusID:254044069.
  * Kerr et al. [2023] J. Kerr, C. M. Kim, K. Goldberg, A. Kanazawa, and M. Tancik.  Lerf: Language embedded radiance fields.  _2023 IEEE/CVF International Conference on Computer Vision (ICCV)_ , pages 19672–19682, 2023.  URL https://api.semanticscholar.org/CorpusID:257557329.
  * Nguyen et al. [2023] P. D. Nguyen, T. Ngo, C. Gan, E. Kalogerakis, A. D. Tran, C. Pham, and K. Nguyen.  Open3dis: Open-vocabulary 3d instance segmentation with 2d mask guidance.  _ArXiv_ , abs/2312.10671, 2023.  URL https://api.semanticscholar.org/CorpusID:266348609.
  * Koch et al. [2024] S. Koch, N. Vaskevicius, M. Colosi, P. Hermosilla, and T. Ropinski.  Open3dsg: Open-vocabulary 3d scene graphs from point clouds with queryable objects and open-set relationships.  _ArXiv_ , abs/2402.12259, 2024.  URL https://api.semanticscholar.org/CorpusID:267750890.
  * Shen et al. [2023] B. W. Shen, G. Yang, A. Yu, J. R. Wong, L. P. Kaelbling, and P. Isola.  Distilled feature fields enable few-shot language-guided manipulation.  In _Conference on Robot Learning_ , 2023.  URL https://api.semanticscholar.org/CorpusID:260926035.
  * Tschernezki et al. [2022] V. Tschernezki, I. Laina, D. Larlus, and A. Vedaldi.  Neural feature fusion fields: 3d distillation of self-supervised 2d image representations.  _2022 International Conference on 3D Vision (3DV)_ , pages 443–453, 2022.  URL https://api.semanticscholar.org/CorpusID:252118532.
  * Kobayashi et al. [2022] S. Kobayashi, E. Matsumoto, and V. Sitzmann.  Decomposing nerf for editing via feature field distillation.  _ArXiv_ , abs/2205.15585, 2022.  URL https://api.semanticscholar.org/CorpusID:249209811.
  * Engelmann et al. [2024] F. Engelmann, F. Manhardt, M. Niemeyer, K. Tateno, M. Pollefeys, and F. Tombari.  Opennerf: Open set 3d neural scene segmentation with pixel-wise features and rendered novel views, 2024.
  * Rashid et al. [2023] A. Rashid, S. Sharma, C. M. Kim, J. Kerr, L. Y. Chen, A. Kanazawa, and K. Goldberg.  Language embedded radiance fields for zero-shot task-oriented grasping.  In _Conference on Robot Learning_ , 2023.  URL https://api.semanticscholar.org/CorpusID:261882332.
  * Zhang et al. [2023] J. Zhang, R. Dong, and K. Ma.  Clip-fo3d: Learning free open-world 3d scene representations from 2d dense clip.  _2023 IEEE/CVF International Conference on Computer Vision Workshops (ICCVW)_ , pages 2040–2051, 2023.  URL https://api.semanticscholar.org/CorpusID:257404908.
  * Dai et al. [2017] A. Dai, A. X. Chang, M. Savva, M. Halber, T. Funkhouser, and M. Nießner.  Scannet: Richly-annotated 3d reconstructions of indoor scenes.  In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition_ , pages 5828–5839, 2017.
  * Ramakrishnan et al. [2021] S. K. Ramakrishnan, A. Gokaslan, E. Wijmans, O. Maksymets, A. Clegg, J. Turner, E. Undersander, W. Galuba, A. Westbury, A. X. Chang, M. Savva, Y. Zhao, and D. Batra.  Habitat-matterport 3d dataset (hm3d): 1000 large-scale 3d environments for embodied ai.  _ArXiv_ , abs/2109.08238, 2021.  URL https://api.semanticscholar.org/CorpusID:237563216.
  * Qin et al. [2024] M. Qin, W. Li, J. Zhou, H. Wang, and H. Pfister.  Langsplat: 3d language gaussian splatting, 2024.
  * Chen et al. [2022] M. Chen, Q. Hu, Z. Yu, H. Thomas, A. Feng, Y. Hou, K. McCullough, F. Ren, and L. Soibelman.  Stpls3d: A large-scale synthetic and real aerial photogrammetry 3d point cloud dataset.  _arXiv preprint arXiv:2203.09065_ , 2022.
  * Straub et al. [2019] J. Straub, T. Whelan, L. Ma, Y. Chen, E. Wijmans, S. Green, J. J. Engel, R. Mur-Artal, C. Ren, S. Verma, et al.  The replica dataset: A digital replica of indoor spaces.  _arXiv preprint arXiv:1906.05797_ , 2019.
  * Liu et al. [2021] H. Liu, A. Lin, X. Han, L. Yang, Y. Yu, and S. Cui.  Refer-it-in-rgbd: A bottom-up approach for 3d visual grounding in rgbd images.  _2021 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , pages 6028–6037, 2021.
  * Mauceri et al. [2019] C. Mauceri, M. Palmer, and C. Heckman.  Sun-spot: An rgb-d dataset with spatial referring expressions.  _2019 IEEE/CVF International Conference on Computer Vision Workshop (ICCVW)_ , pages 1883–1886, 2019.
  * Fang et al. [2020] H.-S. Fang, C. Wang, M. Gou, and C. Lu.  Graspnet-1billion: A large-scale benchmark for general object grasping.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , pages 11444–11453, 2020.
  * Zhang et al. [2022] H. Zhang, D. Yang, H. Wang, B. Zhao, X. Lan, J. Ding, and N. Zheng.  Regrad: A large-scale relational grasp dataset for safe and object-specific robotic grasping in clutter.  _IEEE Robotics and Automation Letters_ , 7(2):2929–2936, 2022.
  * Tziafas et al. [2023] G. Tziafas, X. Yucheng, A. Goel, M. Kasaei, Z. Li, and H. Kasaei.  Language-guided robot grasping: Clip-based referring grasp synthesis in clutter.  In _7th Annual Conference on Robot Learning_ , 2023.
  * Vuong et al. [2023] A. D. Vuong, M. N. Vu, H. Le, B. Huang, B. P. K. Huynh, T. D. Vo, A. Kugi, and A. Nguyen.  Grasp-anything: Large-scale grasp dataset from foundation models.  _ArXiv_ , abs/2309.09818, 2023.  URL https://api.semanticscholar.org/CorpusID:262045996.
  * Armeni et al. [2016] I. Armeni, O. Sener, A. R. Zamir, H. Jiang, I. Brilakis, M. Fischer, and S. Savarese.  3d semantic parsing of large-scale indoor spaces.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pages 1534–1543, 2016.
  * Rozenberszki et al. [2022] D. Rozenberszki, O. Litany, and A. Dai.  Language-grounded indoor 3d semantic segmentation in the wild.  In _European Conference on Computer Vision_ , pages 125–141. Springer, 2022.
  * Eppner et al. [2020] C. Eppner, A. Mousavian, and D. Fox.  ACRONYM: A large-scale grasp dataset based on simulation.  In _2021 IEEE Int. Conf. on Robotics and Automation, ICRA_ , 2020.
  * Community [2018] B. O. Community.  Blender - a 3d modelling and rendering package.  2018\.  URL http://www.blender.org.
  * Chang et al. [2015] A. X. Chang, T. Funkhouser, L. Guibas, P. Hanrahan, Q. Huang, Z. Li, S. Savarese, M. Savva, S. Song, H. Su, J. Xiao, L. Yi, and F. Yu.  ShapeNet: An Information-Rich 3D Model Repository.  Technical Report arXiv:1512.03012 [cs.GR], Stanford University — Princeton University — Toyota Technological Institute at Chicago, 2015.
  * GPT [2023] Gpt-4v(ision) system card.  2023\.  URL https://api.semanticscholar.org/CorpusID:263218031.
  * Kirillov et al. [2023] A. Kirillov, E. Mintun, N. Ravi, H. Mao, C. Rolland, L. Gustafson, T. Xiao, S. Whitehead, A. C. Berg, W.-Y. Lo, P. Dollár, and R. B. Girshick.  Segment anything.  _2023 IEEE/CVF International Conference on Computer Vision (ICCV)_ , pages 3992–4003, 2023.  URL https://api.semanticscholar.org/CorpusID:257952310.
  * Gu et al. [2021] X. Gu, T.-Y. Lin, W. Kuo, and Y. Cui.  Open-vocabulary object detection via vision and language knowledge distillation.  In _International Conference on Learning Representations_ , 2021.  URL https://api.semanticscholar.org/CorpusID:238744187.
  * Schult et al. [2023] J. Schult, F. Engelmann, A. Hermans, O. Litany, S. Tang, and B. Leibe.  Mask3D: Mask Transformer for 3D Semantic Instance Segmentation.  2023\.
  * Chen et al. [2023] S. Chen, W. N. Tang, P. Xie, W. Yang, and G. Wang.  Efficient heatmap-guided 6-dof grasp detection in cluttered scenes.  _IEEE Robotics and Automation Letters_ , 8:4895–4902, 2023.  URL https://api.semanticscholar.org/CorpusID:259363869.
  * Koenig and Howard [2004] N. P. Koenig and A. Howard.  Design and use paradigms for gazebo, an open-source multi-robot simulator.  _2004 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS) (IEEE Cat. No.04CH37566)_ , 3:2149–2154 vol.3, 2004.
  * Choy et al. [2019] C. B. Choy, J. Gwak, and S. Savarese.  4d spatio-temporal convnets: Minkowski convolutional neural networks.  _2019 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , pages 3070–3079, 2019.  URL https://api.semanticscholar.org/CorpusID:121123422.
  * Han et al. [2020] L. Han, T. Zheng, L. Xu, and L. Fang.  Occuseg: Occupancy-aware 3d instance segmentation.  _2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , pages 2937–2946, 2020.  URL https://api.semanticscholar.org/CorpusID:212725768.
  * Hu et al. [2021a] W. Hu, H. Zhao, L. Jiang, J. Jia, and T.-T. Wong.  Bidirectional projection network for cross dimension scene understanding.  _2021 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , pages 14368–14377, 2021a.  URL https://api.semanticscholar.org/CorpusID:232379958.
  * Hu et al. [2021b] Z. Hu, X. Bai, J. Shang, R. Zhang, J. Dong, X. Wang, G. Sun, H. Fu, and C.-L. Tai.  Vmnet: Voxel-mesh network for geodesic-aware 3d semantic segmentation.  _2021 IEEE/CVF International Conference on Computer Vision (ICCV)_ , pages 15468–15478, 2021b.  URL https://api.semanticscholar.org/CorpusID:236493200.
  * Li et al. [2022] J. Li, X. He, Y. Wen, Y. Gao, X. Cheng, and D. Zhang.  Panoptic-phnet: Towards real-time and high-precision lidar panoptic segmentation via clustering pseudo heatmap.  _2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , pages 11799–11808, 2022.  URL https://api.semanticscholar.org/CorpusID:248811224.
  * Robert et al. [2022] D. Robert, B. Vallet, and L. Landrieu.  Learning multi-view aggregation in the wild for large-scale 3d semantic segmentation.  _2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , pages 5565–5574, 2022.  URL https://api.semanticscholar.org/CorpusID:248218804.
  * Wu et al. [2014] Z. Wu, S. Song, A. Khosla, F. Yu, L. Zhang, X. Tang, and J. Xiao.  3d shapenets: A deep representation for volumetric shapes.  _2015 IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ , pages 1912–1920, 2014.  URL https://api.semanticscholar.org/CorpusID:206592833.
  * Zhang et al. [2021] R. Zhang, Z. Guo, W. Zhang, K. Li, X. Miao, B. Cui, Y. J. Qiao, P. Gao, and H. Li.  Pointclip: Point cloud understanding by clip.  _2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , pages 8542–8552, 2021.  URL https://api.semanticscholar.org/CorpusID:244909021.
  * Caesar et al. [2019] H. Caesar, V. Bankiti, A. H. Lang, S. Vora, V. E. Liong, Q. Xu, A. Krishnan, Y. Pan, G. Baldan, and O. Beijbom.  nuscenes: A multimodal dataset for autonomous driving.  _2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , pages 11618–11628, 2019.  URL https://api.semanticscholar.org/CorpusID:85517967.
  * Behley et al. [2019] J. Behley, M. Garbade, A. Milioto, J. Quenzel, S. Behnke, C. Stachniss, and J. Gall.  Semantickitti: A dataset for semantic scene understanding of lidar sequences.  _2019 IEEE/CVF International Conference on Computer Vision (ICCV)_ , pages 9296–9306, 2019.  URL https://api.semanticscholar.org/CorpusID:199441943.
  * Achlioptas et al. [2020] P. Achlioptas, A. Abdelreheem, F. Xia, M. Elhoseiny, and L. J. Guibas.  Referit3d: Neural listeners for fine-grained 3d object identification in real-world scenes.  In _European Conference on Computer Vision_ , 2020.
  * Zhao et al. [2021] L. Zhao, D. Cai, L. Sheng, and D. Xu.  3dvg-transformer: Relation modeling for visual grounding on point clouds.  _2021 IEEE/CVF International Conference on Computer Vision (ICCV)_ , pages 2908–2917, 2021.
  * Huang et al. [2022] S. Huang, Y. Chen, J. Jia, and L. Wang.  Multi-view transformer for 3d visual grounding.  _2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , pages 15503–15512, 2022.
  * Rozenberszki et al. [2022] D. Rozenberszki, O. Litany, and A. Dai.  Language-grounded indoor 3d semantic segmentation in the wild.  _ArXiv_ , abs/2204.07761, 2022.  URL https://api.semanticscholar.org/CorpusID:248227627.
  * Zhong et al. [2021] Y. Zhong, J. Yang, P. Zhang, C. Li, N. C. F. Codella, L. H. Li, L. Zhou, X. Dai, L. Yuan, Y. Li, and J. Gao.  Regionclip: Region-based language-image pretraining.  _2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , pages 16772–16782, 2021.  URL https://api.semanticscholar.org/CorpusID:245218534.
  * Minderer et al. [2022] M. Minderer, A. A. Gritsenko, A. Stone, M. Neumann, D. Weissenborn, A. Dosovitskiy, A. Mahendran, A. Arnab, M. Dehghani, Z. Shen, X. Wang, X. Zhai, T. Kipf, and N. Houlsby.  Simple open-vocabulary object detection with vision transformers.  _ArXiv_ , abs/2205.06230, 2022.  URL https://api.semanticscholar.org/CorpusID:248721818.
  * Zhou et al. [2022] X. Zhou, R. Girdhar, A. Joulin, P. Krahenbuhl, and I. Misra.  Detecting twenty-thousand classes using image-level supervision.  _ArXiv_ , abs/2201.02605, 2022.  URL https://api.semanticscholar.org/CorpusID:245827815.
  * Minderer et al. [2023] M. Minderer, A. A. Gritsenko, and N. Houlsby.  Scaling open-vocabulary object detection.  _ArXiv_ , abs/2306.09683, 2023.  URL https://api.semanticscholar.org/CorpusID:259187664.
  * Wang et al. [2021] Z. Wang, Y. Lu, Q. Li, X. Tao, Y. Guo, M. Gong, and T. Liu.  Cris: Clip-driven referring image segmentation.  _2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , pages 11676–11685, 2021.  URL https://api.semanticscholar.org/CorpusID:244729320.
  * Lüddecke and Ecker [2021] T. Lüddecke and A. S. Ecker.  Image segmentation using text and image prompts.  _2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , pages 7076–7086, 2021.  URL https://api.semanticscholar.org/CorpusID:247794227.
  * Shafiullah et al. [2022] N. M. M. Shafiullah, C. Paxton, L. Pinto, S. Chintala, and A. Szlam.  Clip-fields: Weakly supervised semantic fields for robotic memory.  _ArXiv_ , abs/2210.05663, 2022.  URL https://api.semanticscholar.org/CorpusID:252815898.
  * Bolte et al. [2023] B. Bolte, A. S. Wang, J. Yang, M. Mukadam, M. Kalakrishnan, and C. Paxton.  Usa-net: Unified semantic and affordance representations for robot memory.  _2023 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS)_ , pages 1–8, 2023.  URL https://api.semanticscholar.org/CorpusID:258298248.