# GMFlow: Learning Optical Flow via Global Matching

Haofei Xu1 Jing Zhang2 Jianfei Cai1 Hamid Rezatofighi1 Dacheng Tao3,2
1Department of Data Science and AI, Monash University, Australia
2The University of Sydney, Australia 3JD Explore Academy, China
<EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS>Work done while Haofei was at JD
Explore Academy

###### Abstract

Learning-based optical flow estimation has been dominated with the pipeline of
cost volume with convolutions for flow regression, which is inherently limited
to local correlations and thus is hard to address the long-standing challenge
of large displacements. To alleviate this, the state-of-the-art framework RAFT
gradually improves its prediction quality by using a large number of iterative
refinements, achieving remarkable performance but introducing linearly
increasing inference time. To enable both high accuracy and efficiency, we
completely revamp the dominant flow regression pipeline by reformulating
optical flow as a _global matching_ problem, which identifies the
correspondences by directly comparing feature similarities. Specifically, we
propose a GMFlow framework, which consists of three main components: a
customized Transformer for feature enhancement, a correlation and softmax
layer for global feature matching, and a self-attention layer for flow
propagation. We further introduce a refinement step that reuses GMFlow at
higher feature resolution for residual flow prediction. Our new framework
outperforms 31-refinements RAFT on the challenging Sintel benchmark, while
using only one refinement and running faster, suggesting a new paradigm for
accurate and efficient optical flow estimation. Code is available at
https://github.com/haofeixu/gmflow.

##  1 Introduction

Since the pioneering learning-based work, FlowNet [10], optical flow has been
regressed with convolutions for a long time [17, 33, 38, 16, 41, 48, 19, 53].
To encode the matching information into the network, the cost volume (_i.e_.,
correlation) [13] was shown to be an effective component and thus has been
extensively used in popular frameworks [17, 38, 41]. However, such regression-
based approaches have one major intrinsic limitation. That is, the cost volume
requires a predefined size, as the search space is viewed as the channel
dimension for subsequent regression with convolutions. This requirement
restricts the search space to a _local_ range, making it hard to handle large
displacements.

Figure 1: Conceptual comparison of flow estimation approaches. Most previous
methods _regress_ optical flow from a _local_ cost volume (_i.e_.,
correlation) with convolutions, while we perform _global matching_ with a
Transformer and differentiable matching layer (_i.e_., correlation and
softmax).

To alleviate the large displacements issue, RAFT [41] proposes an iterative
framework with a large number of iterative refinements where convolutions are
applied to different local cost volumes at different iteration stages so as to
gradually reach near-global search space, achieving outstanding performance on
standard benchmarks. The current state-of-the-art methods [19, 48, 53, 41] are
all based on such an iterative architecture. Despite the excellent
performance, such a large number of sequential refinements introduce linearly
increasing inference time, which makes it hard for speed optimization and
hinders its deployment in real-world applications. This raises a question: Is
it possible to achieve both high accuracy and efficiency without requiring
such a large number of iterative refinements?

We notice that another visual correspondence problem, _i.e_., sparse matching
between an image pair [36, 45, 40], which usually features a _large viewpoint
change_ , moves into a different track. The top-performing sparse methods
(_e.g_., SuperGlue [36] and LoFTR [40]) adopt _Transformers_ [43] to reason
about the mutual relationship between feature descriptors, and the
correspondences are extracted with an explicit _matching_ layer (_e.g_.,
softmax layer [45]).

Inspired by sparse matching, we propose to completely remove the additional
convolutional layers operating on a predefined local cost volume, and
reformulate optical flow as a _global matching_ problem, which is distinct
from all previous learning-based optical flow methods. Fig. 1 provides a
conceptual comparison of these two flow estimation approaches. Our flow
prediction is obtained with a differentiable matching layer, _i.e_.,
correlation and softmax layer, by comparing feature similarities. Such a
formulation calls for more discriminative feature representations, for which
the Transformer [43] becomes a natural choice.

We would like to point out that although our pipeline shares the conceptual
major components (_i.e_., Transfromer and the softmax matching layer) with
sparse matching [36, 40], our motivation is originated from the development of
optical flow methods and the challenges associated with formulating optical
flow as a global matching problem are quite different. Optical flow focuses on
dense correspondences for every pixel and modern learning-based architectures
are mostly designed by regression from a local cost volume. The scale and
complexity of optical flow are much higher. Moreover, optical flow needs to
deal with occluded and out-of-boundary pixels, for which the simple softmax-
based matching layer will not be effective.

In this paper, we propose a GMFlow framework to realize the global matching
formulation for optical flow. Specifically, the dense features extracted from
a convolutional backbone network are fed into a Transformer that consists of
self-, cross-attentions and feed-forward network to obtain more discriminative
features. We then compare the feature similarities by correlating all pair-
wise features. After that, the flow prediction is obtained with a
differentiable softmax matching layer. To address occluded and out-of-boundary
pixels, we incorporate an additional self-attention layer to propagate the
high-quality flow prediction from matched pixels to unmatched ones by
exploiting the feature self-similarity [15, 19]. We further introduce a
refinement step that reuses GMFlow at higher feature resolution for residual
flow prediction. Our full framework achieves competitive performance and
higher efficiency compared with the state-of-the-art methods. Specifically,
with only one refinement, GMFlow outperforms 31-refinements RAFT on the
challenging Sintel [4] dataset, while running faster.

Our major contributions can be summarized as follows:

  * •

We completely revamp the dominant flow regression pipeline by reformulating
optical flow as a global matching problem, which effectively addresses the
long-standing challenge of large displacements.

  * •

We propose a GMFlow framework to realize the global matching formulation,
which consists of three main components: a Transformer for feature
enhancement, a correlation and softmax layer for global feature matching, and
a self-attention layer for flow propagation.

  * •

We further propose a refinement step to exploit higher resolution feature,
which allows us to reuse the same GMFlow framework for residual flow
estimation.

  * •

GMFlow outperforms 31-refinements RAFT on the challenging Sintel benchmark,
while using only one refinement and running faster, suggesting a new paradigm
for accurate and efficient flow estimation.

Figure 2: Overview of GMFlow framework. We first extract $8\times$ downsampled
dense features from two input video frames with a weight-sharing convolutional
network. Then the features are fed into a Transformer for feature enhancement.
Next we compare feature similarities by correlating all pair-wise features and
the optical flow is obtained with a softmax matching layer. An additional
self-attention layer is introduced to propagate the high-quality flow
predictions in matched pixels to unmatched ones by considering the feature
self-similarity.

##  2 Related Work

Flow estimation approach. The flow estimation approach is fundamental to
existing popular optical flow frameworks [17, 38, 16, 41, 48, 19, 53], notably
the coarse-to-fine method PWC-Net [38] and iterative refinement method RAFT
[41]. They both perform some sort of multi-stage refinements, either at
multiple scales [38] or a single resolution [41]. For flow prediction at each
stage, their pipeline is conceptual similar, _i.e_., regressing optical flow
from a local cost volume with convolutions. However, such an approach is hard
to handle large displacements. Thus multi-stage refinements are required to
estimate large motion incrementally. The success of RAFT largely lies in the
large number of iterative refinements it can perform. Two exceptions to these
local regression approaches are DICL [44] and GLU-Net [42]. DICL performs
_local matching_ with _convolutions_ , while GLU-Net _regresses_ flow from
_global correlation_ with _convolutions_ , which also makes itself restricted
to fixed image resolution and thus an additional sub-network is required to
handle this issue. Distinct from these approaches, we perform _global
matching_ with a _Transformer_ and we show it is indeed possible to achieve
highly accurate results without relying on a large number of refinements.

Large displacements. Large displacements have been a long-standing challenge
for optical flow [2, 47, 3, 7, 1]. One popular strategy is to use the coarse-
to-fine approach [14, 3, 38] that estimates large motion incrementally.
However, coarse-to-fine methods tend to miss fast-moving small objects [34] if
the resolution is too coarse. To alleviate this issue, RAFT [41] proposes to
maintain a single high resolution and gradually improve the initial prediction
with a large number of iterative refinements, achieving remarkable performance
on standard benchmarks. However, such a large number of sequential refinements
introduce linearly increasing inference time, which makes it hard for speed
optimization and hinders its integration into real-world systems. In contrast,
our new framework streamlines the optical flow pipeline and estimates large
displacements with both high accuracy and efficiency, achieved by a
reformulation of the optical flow problem and a strong Transformer.

Transformer for correspondences. SuperGlue [36] pioneered the use of
Transformer for sparse feature matching. LoFTR [40] further improves its
performance by removing the feature detection step in the typical pipelines.
COTR [20] formulates the correspondence problem as a functional mapping by
querying the interest point and uses a Transformer as the function. These
frameworks are mainly designed for sparse matching problems, and it is non-
trivial to directly adopt them to dense correspondence tasks. Although COTR in
principle can also predict dense correspondences by querying every pixel
location, the inference speed will be significantly slower. For dense
correspondences, STTR [24] is a Transformer-based method for stereo matching,
which can be viewed as a special case of optical flow. Besides, STTR relies on
a complex optimal transport matching layer and doesn’t produce predictions for
occluded pixels, while we use a much simpler softmax operation and a simple
flow propagation layer to handle occlusion. Another related work Perceiver IO
[18] targets general problems and uses optical flow as an application in
experiments. Although both method are based on Transformers, Perceiver IO
directly regresses optical flow via self-attention, while we aim at learning
strong feature representations (in particular with cross-attention) for
matching.

##  3 Methodology

Optical flow is intuitively a matching problem that aims at finding
corresponding pixels. To achieve this, one can compare the similarity of the
features for each pixel, and identify the corresponding pixel that gives the
highest similarity. Such a process requires the features to be discriminative
enough to stand out. Aggregating the spatial contexts within the image itself
and information from another image can intuitively alleviate ambiguities and
improve their distinctiveness. Such design philosophies have enabled great
achievement in sparse feature matching frameworks [36, 40]. The success of
sparse matching, which usually features a large viewpoint change, motivates us
to formulate optical flow as an explicit global matching problem in order to
address the challenge of large displacements.

In the following, we first provide a general description of our global
matching formulation, and then present a Transformer-based framework to
realize it.

###  3.1 Formulation

Given two consecutive video frames ${\bm{I}}_{1}$ and ${\bm{I}}_{2}$, we first
extract downsampled dense features
${\bm{F}}_{1},{\bm{F}}_{2}\in\mathbb{R}^{H\times W\times D}$ with a weight-
sharing convolutional network, where $H,W$ and $D$ denote height, width and
feature dimension, respectively. Considering the correspondences in the two
frames should share high similarity, we first compare the feature similarity
for each pixel in ${\bm{F}}_{1}$ with respect to all pixels in ${\bm{F}}_{2}$
by computing their correlations [45]. This can be implemented efficiently with
a simple matrix multiplication:

${\bm{C}}=\frac{{\bm{F}}_{1}{\bm{F}}_{2}^{T}}{\sqrt{D}}\in\mathbb{R}^{H\times
W\times H\times W},$ (1)

where each element in the correlation matrix ${\bm{C}}$ represents the
correlation value between coordinates ${\bm{p}}_{1}=(i,j)$ in ${\bm{F}}_{1}$
and ${\bm{p}}_{2}=(k,l)$ in ${\bm{F}}_{2}$, and $\frac{1}{\sqrt{D}}$ is a
normalization factor to avoid large values after the dot-product operation
[43].

To identify the correspondence, one naïve approach is to directly take the
location that gives the highest correlation. However, this operation is
unfortunately non-differentiable, which prevents end-to-end training. To
tackle this issue, we use a differentiable matching layer [45, 21, 49].
Specifically, we normalize the last two dimensions of ${\bm{C}}$ with the
softmax operation, which gives us a matching distribution

${\bm{M}}=\mathrm{softmax}({\bm{C}})\in\mathbb{R}^{H\times W\times H\times W}$
(2)

for each location in ${\bm{F}}_{1}$ with respect to all locations in
${\bm{F}}_{2}$. Then, the correspondence $\hat{{\bm{G}}}$ can be obtained by
taking a weighted average of the 2D coordinates of pixel grid
${\bm{G}}\in\mathbb{R}^{H\times W\times 2}$ with the matching distribution
${\bm{M}}$:

$\hat{{\bm{G}}}={\bm{M}}{\bm{G}}\in\mathbb{R}^{H\times W\times 2}.$ (3)

Finally, the optical flow ${\bm{V}}$ can be obtained by computing the
difference between the corresponding pixel coordinates

${\bm{V}}=\hat{{\bm{G}}}-{\bm{G}}\in\mathbb{R}^{H\times W\times 2}.$ (4)

Such a softmax-based approach can not only enable end-to-end training but also
provide sub-pixel accuracy.

###  3.2 Feature Enhancement

Key to our formulation lies in obtaining high-quality discriminative features
for matching. Recall that the features ${\bm{F}}_{1}$ and ${\bm{F}}_{2}$ in
Sec. 3.1 are extracted _independently_ from a weight-sharing convolutional
network. To further consider their mutual dependencies, a natural choice is
Transformer [43], which is particularly suitable for modeling the mutual
relationship between two sets with the attention mechanism, as demonstrated in
sparse matching methods [36, 40]. Since ${\bm{F}}_{1}$ and ${\bm{F}}_{2}$ are
only two sets of features, they have no notion of the spatial position, we
first add the fixed 2D sine and cosine positional encodings (following DETR
[6]) to the features. Adding the position information also makes the matching
process consider not only the feature similarity but also their spatial
distance, which can help resolve ambiguities and improve the performance
(Table LABEL:tab:transformer).

After adding the position information, we perform six stacked self-, cross-
attentions and feed-forward network (FFN) [43] to improve the quality of the
initial features. Specifically, for self-attention, the query, key and value
in the attention mechanism [43] are the same feature. For cross-attention, the
key and value are same but different from the query to introduce their mutual
dependencies. This process is performed for both ${\bm{F}}_{1}$ and
${\bm{F}}_{2}$ symmetrically, _i.e_.,

$\hat{\bm{F}}_{1}=\mathcal{T}({\bm{F}}_{1}+{\bm{P}},{\bm{F}}_{2}+{\bm{P}}),\hskip
6.00006pt\hat{\bm{F}}_{2}=\mathcal{T}({\bm{F}}_{2}+{\bm{P}},{\bm{F}}_{1}+{\bm{P}}),$
(5)

where $\mathcal{T}$ is a Transformer, ${\bm{P}}$ is the positional encoding,
the first input of $\mathcal{T}$ is query and the second is key and value.

One issue in the standard Transformer architecture [43] is the quadratic
computational complexity due to the pair-wise attention operation. To improve
the efficiency, we adopt the shifted local window attention strategy from Swin
Transformer [26]. However, unlike Swin that uses _fixed window size_ , we
split the feature to _fixed number of local windows_ to make the window size
adaptive with the feature size. Specifically, we split the input feature of
size $H\times W$ to $K\times K$ windows (each with size
$\frac{H}{K}\times\frac{W}{K}$), and perform self- and cross-attentions within
each local window independently. For every two consecutive local windows, we
shift the window partition by $(\frac{H}{2K},\frac{W}{2K})$ to introduce
cross-window connections. In our framework, we split to $2\times 2$ windows
(each with size $\frac{H}{2}\times\frac{W}{2}$), which represents a good
speed-accuracy trade-off (Table LABEL:tab:split_attn).

Figure 3: GMFlow also simplifies backward flow computation by transposing the
global correlation matrix without requiring to forward the network twice. The
bidirectional flow can be used for occlusion detection with forward-backward
consistency check.

###  3.3 Flow Propagation

Our softmax-based flow estimation method implicitly assumes that the
corresponding pixels are visible in both images and thus they can be matched
by comparing their similarities. However, this assumption will be invalid for
occluded and out-of-boundary pixels. To remedy this, by observing that the
optical flow field and the image itself share high structure similarity [15,
19], we propose to propagate the high-quality flow predictions in matched
pixels to unmatched ones by measuring the feature self-similarity. This
operation can be implemented efficiently with a simple self-attention layer
(illustrated in Fig. 2):

$\tilde{\bm{V}}=\mathrm{softmax}\left(\frac{\hat{\bm{F}}_{1}\hat{\bm{F}_{1}}^{T}}{\sqrt{D}}\right)\hat{\bm{V}}\in\mathbb{R}^{H\times
W\times 2},$ (6)

where

$\hat{\bm{V}}=\mathrm{softmax}\left(\frac{\hat{\bm{F}}_{1}\hat{\bm{F}_{2}}^{T}}{\sqrt{D}}\right){\bm{G}}-{\bm{G}},$
(7)

is the optical flow prediction from the softmax layer, which is obtained by
substituting the stronger features in Eq. (5) into the softmax matching layer
in Sec. 3.1. Fig. 2 provides an overview of our GMFlow framework.

Method | #blocks | Things (val, clean) | Sintel (train, clean) | Sintel (train, final) |  Param (M)
---|---|---|---|---|---
|  | EPE | $s_{0-10}$ | $s_{10-40}$ | $s_{40+}$ | EPE | $s_{0-10}$ | $s_{10-40}$ | $s_{40+}$ | EPE | $s_{0-10}$ | $s_{10-40}$ | $s_{40+}$ |
cost volume + conv | 0 | 18.83 | 3.42 | 6.49 | 49.65 | 6.45 | 1.75 | 7.17 | 38.19 | 7.75 | 2.10 | 8.88 | 45.29 | 1.8
4 | 10.99 | 1.70 | 3.41 | 29.78 | 3.32 | 0.73 | 3.84 | 20.58 | 4.93 | 0.99 | 5.71 | 31.16 | 4.6
8 | 9.59 | 1.44 | 2.96 | 26.04 | 2.89 | 0.65 | 3.36 | 17.75 | 4.32 | 0.88 | 4.95 | 27.33 | 8.0
12 | 9.04 | 1.37 | 2.84 | 24.46 | 2.78 | 0.65 | 3.32 | 16.69 | 4.07 | 0.84 | 4.76 | 25.44 | 11.5
| 18 | 8.67 | 1.33 | 2.74 | 23.43 | 2.61 | 0.59 | 3.07 | 15.91 | 3.94 | 0.82 | 4.62 | 24.58 | 15.7
Transformer + softmax | 0 | 22.93 | 8.57 | 11.13 | 52.07 | 8.44 | 2.71 | 11.60 | 42.10 | 10.28 | 3.11 | 13.83 | 53.34 | 1.0
1 | 11.45 | 2.98 | 4.68 | 28.35 | 4.12 | 1.27 | 5.08 | 22.25 | 6.11 | 1.70 | 7.89 | 33.52 | 1.6
2 | 8.59 | 1.80 | 3.28 | 21.99 | 3.09 | 0.90 | 3.66 | 17.37 | 4.54 | 1.24 | 5.44 | 26.00 | 2.1
4 | 7.19 | 1.40 | 2.62 | 18.66 | 2.43 | 0.67 | 2.73 | 14.23 | 3.78 | 1.01 | 4.27 | 22.37 | 3.1
| 6 | 6.67 | 1.26 | 2.40 | 17.37 | 2.28 | 0.58 | 2.49 | 13.89 | 3.44 | 0.80 | 3.97 | 21.02 | 4.2
Table 1: Methodology comparison. We stack different number of convolutional
residual blocks or Transformer blocks to see how the performance varies. All
models are trained on Chairs and Things training sets. We report the
performance on Things validation clean set and the cross-dataset
generalization results on Sintel training clean and final sets.

###  3.4 Refinement

The framework presented so far (based on $1/8$ features) can already achieve
competitive performance (Table 3). It can be further improved by introducing
additional higher resolution ($1/4$) feature for refinement. Specifically, we
first upsample the previous $1/8$ flow prediction to $1/4$ resolution, and
warp the second feature with the current flow prediction. Then the refinement
task is reduced to the residual flow learning, where the same GMFlow framework
depicted in Fig. 2 can be used but in a local range. Specifically, we split to
$8\times 8$ local windows (each with $1/32$ of the original image resolution)
in the Transformer and perform a $9\times 9$ local window matching for each
pixel. After obtaining the flow prediction from the softmax layer, we perform
a $3\times 3$ local window self-attention operation for flow propagation.

Note that here we share the Transformer and self-attention weights in the
refinement step with the global matching stage, which not only reduces
parameters but also improves the generalization (Table LABEL:tab:refine). To
generate the $1/4$ and $1/8$ features, we also share the backbone features.
Specifically, we take a similar approach to TridentNet [23] but use a weight-
sharing convolution with strides $1$ and $2$, respectively. Such a weight-
sharing design also leads to better performance than the feature pyramid
network [25].

###  3.5 Training Loss

We supervise all flow predictions using $\ell_{1}$ loss between the ground
truth:

$L=\sum_{i=1}^{N}\gamma^{N-i}\|{\bm{V}}_{i}-{\bm{V}}_{\mathrm{gt}}\|_{1},$ (8)

where $N$ is the number of flow predictions including the intermediate and
final ones, and $\gamma$ (set to 0.9) is the weight that is exponentially
increasing to give higher weights for later predictions following RAFT [41].

setup | Things (val) | Sintel (train) |  Param (M) |
---|---|---|---|---
| clean | clean | final |  |
full | 6.67 | 2.28 | 3.44 | 4.2 |
w/o cross | 10.84 | 4.48 | 6.32 | 3.8 |
w/o pos | 8.38 | 2.85 | 4.28 | 4.2 |
w/o FFN | 8.71 | 3.10 | 4.43 | 1.8 |
w/o self | 7.04 | 2.49 | 3.69 | 3.8 |
(a)

#splits | Things (val, clean) |  Time (ms)
---|---|---
| EPE | $s_{0-10}$ | $s_{10-40}$ | $s_{40+}$ |
$1\times 1$ | 6.34 | 1.26 | 2.37 | 16.36 | 105
$2\times 2$ | 6.67 | 1.26 | 2.40 | 17.37 | 53
$4\times 4$ | 7.32 | 1.29 | 2.58 | 19.26 | 35
(b)

matching space | Things (val, clean) |  Time (ms)
---|---|---
| EPE | $s_{0-10}$ | $s_{10-40}$ | $s_{40+}$ |
global | 6.67 | 1.26 | 2.40 | 17.37 | 52.6
local $3\times 3$ | 31.78 | 1.19 | 12.40 | 85.39 | 51.2
local $5\times 5$ | 26.51 | 0.89 | 6.67 | 76.76 | 51.5
local $9\times 9$ | 19.88 | 1.01 | 2.44 | 61.06 | 52.9
(c)

prop. | Sintel (clean) | Sintel (final)
---|---|---
| all | matched | unmatched | all | matched | unmatched
w/o | 2.28 | 1.06 | 15.54 | 3.44 | 1.95 | 19.50
w/ | 1.89 | 1.10 | 10.39 | 3.13 | 1.98 | 15.52
(d)

setup |  share Trans? |  share feature? | Things | Sintel (train) |  Param (M) |
---|---|---|---|---|---|---
|  |  | clean | clean | final |  |
w/o refine | - | - | 3.98 | 1.65 | 2.94 | 4.7 |
w/ refine |  |  | 3.41 | 1.28 | 2.75 | 8.0 |
✓ |  | 3.21 | 1.27 | 2.60 | 4.9 |
| ✓ | 3.26 | 1.21 | 2.70 | 7.9 |
✓ | ✓ | 3.24 | 1.20 | 2.59 | 4.7 |
(e)

training iterations | Things (val) | Sintel (train) | KITTI (train)
---|---|---|---
| clean | clean | final | EPE | F1-all
200K | 3.24 | 1.20 | 2.59 | 9.32 | 26.95
400K | 3.01 | 1.16 | 2.51 | 8.96 | 26.18
600K | 2.93 | 1.13 | 2.42 | 8.34 | 24.77
800K | 2.80 | 1.08 | 2.48 | 7.77 | 23.40
(f)

Table 2: GMFlow ablations. All models are trained on Chairs and Things
training sets.

##  4 Experiments

Datasets and evaluation setup. Following previous methods [17, 38, 41], we
first train on the FlyingChairs (Chairs) [10] and FlyingThings3D (Things) [29]
datasets, and then evaluate Sintel [4] and KITTI [31] training sets. We also
evaluate on the Things validation set to see how the model performs on the
same-domain data. Finally, we perform additional fine-tuning on Sintel and
KITTI training sets and report the performance on the online benchmarks.

Metrics. We adopt the commonly used metric in optical flow, _i.e_., the end-
point-error (EPE), which is the average $\ell_{2}$ distance between the
prediction and ground truth. For KITTI dataset, we also use _F1-all_ , which
reflects the percentage of outliers. To better understand the performance
gains, we also report the EPE in different motion magnitudes. Specifically, we
use $s_{0-10},s_{10-40}$ and $s_{40+}$, to denote the EPE over pixels with
ground truth flow motion magnitude falling to $0-10$, $10-40$ and more than
$40$ pixels.

Implementation details. We implement our framework in PyTorch. Our
convolutional backbone network is identical to RAFT’s model, except that our
final feature dimension is 128, while RAFT’s is 256. We stack 6 Transformer
blocks. To upsample the low-resolution flow prediction to the original image
resolution, we use RAFT’s convex upsampling [41] method. We use AdamW [27] as
the optimizer. We first train the model on Chairs dataset for 100K iterations,
with a batch size of 16 and a learning rate of 4e-4. We then fine-tune it on
Things dataset for 200K iterations for ablation experiments, with a batch size
of 8 and a learning rate of 2e-4. Our best model is fine-tuned on Things for
800K iterations. For the final fine-tuning process on Sintel and KITTI
datasets, we report the details in Sec. 4.4. Further details are provided in
the _supplementary material_.

###  4.1 Methodology Comparison

Flow estimation approach. We compare our Transformer and softmax-based flow
estimation method with the cost volume and convolution-based approach.
Specifically, we adopt the state-of-the-art cost volume construction method in
RAFT [41] that concatenates 4 local cost volumes at 4 scales, where each cost
volume has a dimension of $H\times W\times(2R+1)^{2}$, where $H$ and $W$ are
image height and width, respectively, and the search range $R$ is set to 4
following RAFT. To regress flow, we stack different number of convolutional
residual blocks [12] to see how the performance varies. The final optical flow
is obtained with a $3\times 3$ convolution with 2 output channels. For our
proposed framework, we stack different number of Transformer blocks for
feature enhancement, where one Transformer block consists of self-, cross-
attentions and a feed-forward network (FFN). The final optical flow is
obtained with a global correlation and softmax layer. Both methods use
bilinear upsampling in this comparison. Table 1 shows that the performance
improvement of our method is more significant compared to the cost volume and
convolution-based approach. For instance, our method with 2 Transformer blocks
can already outperform 8 convolution blocks, especially for large motion
($s_{40+}$). The performance can be further improved by stacking more layers,
surpassing the cost volume and convolution-based approach by a large margin.
We present more comparisons with all the possible combinations of different
flow estimation approaches in the _supplementary material_ , where our method
is consistently better and has less parameters than other variants.

Bidirectional flow prediction. Our framework also simplifies backward optical
flow computation by directly transposing the global correlation matrix in Eq.
(1). Note that during training we only predict unidirectional flow while at
inference we can obtain bidirectional flow for free, without requiring to
forward the network twice, unlike previous regression-based methods [30, 16].
The bidirectional flow can be used for occlusion detection with forward-
backward consistency check (following [30]), as shown in Fig. 3.

Method | #refine. | Things (val, clean) | Sintel (train, clean) | Sintel (train, final) |  Param (M) |  Time (ms)
---|---|---|---|---|---|---
|  | EPE | $s_{0-10}$ | $s_{10-40}$ | $s_{40+}$ | EPE | $s_{0-10}$ | $s_{10-40}$ | $s_{40+}$ | EPE | $s_{0-10}$ | $s_{10-40}$ | $s_{40+}$ |  |
RAFT [41] | 0 | 14.28 | 1.47 | 3.62 | 40.48 | 4.04 | 0.77 | 4.30 | 26.66 | 5.45 | 0.99 | 6.30 | 35.19 | 5.3 | 25 (14)
3 | 6.27 | 0.69 | 1.67 | 17.63 | 1.92 | 0.47 | 2.32 | 11.37 | 3.25 | 0.65 | 4.00 | 20.04 | 39 (21)
7 | 4.66 | 0.55 | 1.38 | 12.87 | 1.61 | 0.39 | 1.90 | 9.61 | 2.80 | 0.53 | 3.30 | 17.76 | 58 (31)
11 | 4.31 | 0.53 | 1.33 | 11.79 | 1.55 | 0.41 | 1.73 | 9.19 | 2.72 | 0.52 | 3.12 | 17.43 | 78 (41)
23 | 4.22 | 0.53 | 1.32 | 11.52 | 1.47 | 0.36 | 1.63 | 9.00 | 2.69 | 0.52 | 3.05 | 17.28 | 133 (71)
31 | 4.25 | 0.53 | 1.31 | 11.63 | 1.41 | 0.32 | 1.55 | 8.83 | 2.69 | 0.52 | 3.00 | 17.45 | 170 (91)
GMFlow | 0 | 3.48 | 0.67 | 1.31 | 8.97 | 1.50 | 0.46 | 1.77 | 8.26 | 2.96 | 0.72 | 3.45 | 17.70 | 4.7 | 57 (26)
1 | 2.80 | 0.53 | 1.01 | 7.31 | 1.08 | 0.30 | 1.25 | 6.26 | 2.48 | 0.51 | 2.81 | 15.67 | 4.7 | 151 (66)
Table 3: RAFT’s iterative refinement framework vs. our GMFlow framework. The
models are trained on Chairs and Things training sets. We use RAFT’s
officially released model for evaluation. The inference time is measured on a
single V100 and A100 (in parentheses) GPU at Sintel resolution ($436\times
1024$). Our framework gains more speedup than RAFT ($2.29\times$ _vs_.
$1.87\times$, _i.e_., ours: $151\to 66$, RAFT: $170\to 91$) on the high-end
A100 GPU since our method doesn’t require a large number of sequential
computation.

###  4.2 Ablations

Transformer components. We ablate different Transformer components in Table
LABEL:tab:transformer. The cross-attention contributes most, since it models
the mutual relationship between two features, which is missing in the features
extracted from the convolutional backbone network. Also, the position
information makes the matching process position-dependent, and can be
conducive to alleviate the ambiguities in pure feature similarity-based
matching. Removing the feed-forward network (FFN) reduces a large number of
parameters, while also leading to a moderate performance drop. The self-
attention aggregates contextual cues within the same feature, bringing
additional performance gains.

Figure 4: Optical flow end-point-error _vs_. number of refinements at
inference time. This figure shows the generalization results on Sintel (clean)
training set after training on Chairs and Things datasets. Our method
outperforms 31-refinements RAFT’s performance while using only one refinement
and running faster.

Local window attention. We compare the speed-accuracy trade-off of splitting
to different numbers of local windows for attention computation in Table
LABEL:tab:split_attn. Recall that the extracted features from backbone is
$1/8$ resolution, further splitting to $H/2\times W/2$ local windows (_i.e_.,
$1/16$ of the original resolution) represents a good trade-off between
accuracy and speed, and thus is used in our framework.

Matching Space. We replace our global matching with local matching in Table
LABEL:tab:global_local_match and observe a significant performance drop,
especially for large motion ($s_{40+}$). Besides, the global matching can be
computed efficiently with a simple matrix multiplication, while larger size
for local matching will be slower due to the excessive sampling operation.

Flow propagation. Our flow propagation strategy results in significant
performance gains in unmatched regions (including occluded and out-of-boundary
pixels), as shown in Table LABEL:tab:prop. The structural correlation between
the feature and flow provides a valuable clue to improve the performance of
pixels that are challenging to match. Visual comparisons are provided in the
_supplementary material_.

Refinement. We compare whether to share the backbone when extracting $1/8$ and
$1/4$ features, and whether to share the Transformer for matching at $1/8$ and
$1/4$ resolutions. The results are shown in Table LABEL:tab:refine. Sharing
Transformer and multi-scale features better regularizes the learning process,
leading to better performance and less parameters.

Training length. Our ablations thus far are based on 200K-iterations fine-
tuning on Things dataset. In Table LABEL:tab:train_schedule, we show our
framework benefits from more training iterations, consistent with the
observations in previous vision Transformer works [9, 26]. We use
800K-iterations model as our final model in subsequent comparisons.

We analyze the computational complexities of core components in our framework
in the _supplementary material_.

Training data | Method | EPE | F1-all | $s_{0-10}$ | $s_{10-40}$ | $s_{40+}$ |
---|---|---|---|---|---|---|---
C + T | RAFT | 5.32 | 17.46 | 0.67 | 1.58 | 13.68 |
GMFlow | 7.77 | 23.40 | 0.74 | 2.19 | 20.34 |
C + T + VK | RAFT | 2.45 | 7.90 | 0.43 | 1.18 | 5.70 |
GMFlow | 2.85 | 10.77 | 0.49 | 1.16 | 6.87 |
Table 4: Generalization on KITTI after training on synthetic Chairs (C),
Things (T) and Virtual KITTI 2 (VK) datasets.  Figure 5: Visual comparisons on
Sintel test set.

###  4.3 Comparison with RAFT

Sintel. Table 3 shows the results on Things validation set and Sintel clean
and final training sets after training on Chairs and Things training sets.
Without using any refinement, our method achieves better performance on Things
and Sintel (clean) than RAFT with 11 refinements. By using an additional
refinement, our method outperforms RAFT with 31 refinements, especially on
large motion ($s_{40+}$). Fig. 4 visualizes the results. Furthermore, our
framework enjoys faster inference speed compared to RAFT and also does not
require a large number of sequential processing. On the high-end A100 GPU, our
framework gains more speedup compared with RAFT’s sequential framework
($2.29\times$ _vs_. $1.87\times$, _i.e_., ours: $151\to 66$, RAFT: $170\to
91$), reflecting that our framework can benefit more from advanced hardware
acceleration and showing its potential for further speed optimization.

KITTI. Table 4 shows the generalization results on KITTI training set after
training on Chairs and Things training sets. In this evaluation setting, our
framework doesn’t outperform RAFT, which is mainly caused by the gap between
the synthetic training sets and the real-world testing dataset. One key reason
behind our inferior performance is that RAFT, relying on fully convolutional
neural networks, benefits from the inductive biases in convolution layers,
which requires a relatively smaller size training data to generalize to a new
dataset in comparison with Transformers [9, 8, 50, 54]. To substantiate this
claim, we fine-tune both RAFT and our method on the additional Virtual KITTI 2
[5] dataset. We can see from Table 4 that the performance gap becomes smaller
when more data is available.

Method | Sintel (clean) | Sintel (final)
---|---|---
| all | matched | unmatched | all | matched | unmatched
FlowNet2 [17] | 4.16 | 1.56 | 25.40 | 5.74 | 2.75 | 30.11
PWC-Net+ [39] | 3.45 | 1.41 | 20.12 | 4.60 | 2.25 | 23.70
HD3 [52] | 4.79 | 1.62 | 30.63 | 4.67 | 2.17 | 24.99
VCN [51] | 2.81 | 1.11 | 16.68 | 4.40 | 2.22 | 22.24
DICL [44] | 2.63 | 0.97 | 16.24 | 3.60 | 1.66 | 19.44
RAFT [41] | 1.94 | - | - | 3.18 | - | -
GMFlow | 1.74 | 0.65 | 10.56 | 2.90 | 1.32 | 15.80
RAFT† [41] | 1.61 | 0.62 | 9.65 | 2.86 | 1.41 | 14.68
GMA† [19] | 1.39 | 0.58 | 7.96 | 2.47 | 1.24 | 12.50
Table 5: Comparisons on Sintel test test. † represents the method uses last
frame’s flow prediction as initialization for subsequent refinement, while
other methods all use two frames only.

###  4.4 Comparison on Benchmarks

Sintel. Following previous works [41, 48], we further fine-tune our Things
model on several mixed datasets that consist of KITTI [31], HD1K [22],
FlyingThings3D [29] and Sintel [4] training sets. We perform fine-tuning for
200K iterations with a batch size of 8 and a learning rate of 2e-4. The
results on Sintel test set are shown in Table 5. We achieve best performance
among all the competitive state-of-the-art methods that only use two frames at
inference. Although RAFT can also use last frame’s prediction as
initialization for subsequent refinement, we still achieve comparable
performance in matched regions even without using multi-frame information.
Qualitative comparisons between GMFlow and other state-of-the-art approaches
on Sintel test set are shown in Fig. 5. More visual results on DAVIS [32]
dataset are given in the _supplementary material_.

KITTI. We perform additional fine-tuning on the KITTI 2015 training set from
the model trained on Sintel. We train GMFlow for 100K iterations with a batch
size of 8 and a learning rate of 2e-4. Table 6 shows the evaluation results.
For non-occluded (Noc) pixels, our performance is slightly inferior to RAFT.
For all pixels that contain both non-occluded and occluded pixels, the
performance gap becomes larger, which indicates that our inferior performance
largely lies in occluded regions.

Method |  FlowNet2 [17] |  PWC-Net+ [39] |  RAFT [41] | GMFlow
---|---|---|---|---
All | 11.48 | 7.72 | 5.10 | 9.32
Noc | 6.94 | 4.91 | 3.07 | 3.80
Table 6: Comparisons on KITTI test set. The metric is F1-all. “All” denotes
the evaluation results on all pixels with ground truth, and “Noc” denotes non-
occluded pixels only.

###  4.5 Limitation and Discussion

Our framework still has room for future improvement in occluded regions, as
can be seen from the KITTI results in Table 6. Besides, our framework may not
generalize very well when the training data has significantly large gap with
the test data (_e.g_., synthetic Things to real-world KITTI). Fortunately,
there are many large-scale datasets available currently, _e.g_., Virtual KITTI
[11, 5], VIPER [35], REFRESH [28], AutoFlow [37] and TartanAir [46], they can
be used to enhance Transformer’s generalization ability.

##  5 Conclusion

We have presented a new global matching formulation for optical flow and
demonstrated its strong performance. We hope our new perspective will pave a
way towards a new paradigm for accurate and efficient optical flow estimation.

Broader impact. Our proposed method might produce unreliable results in
occluded regions, thus care should be taken when using the prediction results
from our model, especially for safety-critical scenarios like self-driving
cars.

Acknowledgement. This research is supported in part by Monash FIT Start-up
Grant. Dr. Jing Zhang is supported by ARC FL-170100117.

## References

  * [1] Christian Bailer, Bertram Taetz, and Didier Stricker.  Flow fields: Dense correspondence fields for highly accurate large displacement optical flow estimation.  In ICCV, pages 4015–4023, 2015.
  * [2] Thomas Brox, Christoph Bregler, and Jitendra Malik.  Large displacement optical flow.  In ICCV, pages 41–48. IEEE, 2009.
  * [3] Thomas Brox, Andrés Bruhn, Nils Papenberg, and Joachim Weickert.  High accuracy optical flow estimation based on a theory for warping.  In ECCV, pages 25–36. Springer, 2004.
  * [4] Daniel J Butler, Jonas Wulff, Garrett B Stanley, and Michael J Black.  A naturalistic open source movie for optical flow evaluation.  In ECCV, pages 611–625. Springer, 2012.
  * [5] Yohann Cabon, Naila Murray, and Martin Humenberger.  Virtual kitti 2.  arXiv preprint arXiv:2001.10773, 2020.
  * [6] Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko.  End-to-end object detection with transformers.  In ECCV, pages 213–229. Springer, 2020.
  * [7] Zhuoyuan Chen, Hailin Jin, Zhe Lin, Scott Cohen, and Ying Wu.  Large displacement optical flow from nearest neighbor fields.  In CVPR, pages 2443–2450, 2013.
  * [8] Stéphane d’Ascoli, Hugo Touvron, Matthew Leavitt, Ari Morcos, Giulio Biroli, and Levent Sagun.  Convit: Improving vision transformers with soft convolutional inductive biases.  ICML, 2021.
  * [9] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby.  An image is worth 16x16 words: Transformers for image recognition at scale.  ICLR, 2021.
  * [10] Alexey Dosovitskiy, Philipp Fischer, Eddy Ilg, Philip Hausser, Caner Hazirbas, Vladimir Golkov, Patrick Van Der Smagt, Daniel Cremers, and Thomas Brox.  Flownet: Learning optical flow with convolutional networks.  In ICCV, pages 2758–2766, 2015.
  * [11] Adrien Gaidon, Qiao Wang, Yohann Cabon, and Eleonora Vig.  Virtual worlds as proxy for multi-object tracking analysis.  In CVPR, pages 4340–4349, 2016.
  * [12] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  In CVPR, pages 770–778, 2016.
  * [13] Asmaa Hosni, Christoph Rhemann, Michael Bleyer, Carsten Rother, and Margrit Gelautz.  Fast cost-volume filtering for visual correspondence and beyond.  TPAMI, 35(2):504–511, 2012.
  * [14] Yinlin Hu, Rui Song, and Yunsong Li.  Efficient coarse-to-fine patchmatch for large displacement optical flow.  In CVPR, pages 5704–5712, 2016.
  * [15] Tak-Wai Hui and Chen Change Loy.  Liteflownet3: Resolving correspondence ambiguity for more accurate optical flow estimation.  In ECCV, pages 169–184. Springer, 2020.
  * [16] Junhwa Hur and Stefan Roth.  Iterative residual refinement for joint optical flow and occlusion estimation.  In CVPR, pages 5754–5763, 2019.
  * [17] Eddy Ilg, Nikolaus Mayer, Tonmoy Saikia, Margret Keuper, Alexey Dosovitskiy, and Thomas Brox.  Flownet 2.0: Evolution of optical flow estimation with deep networks.  In CVPR, pages 2462–2470, 2017.
  * [18] Andrew Jaegle, Sebastian Borgeaud, Jean-Baptiste Alayrac, Carl Doersch, Catalin Ionescu, David Ding, Skanda Koppula, Daniel Zoran, Andrew Brock, Evan Shelhamer, et al.  Perceiver io: A general architecture for structured inputs & outputs.  In ICLR, 2022.
  * [19] Shihao Jiang, Dylan Campbell, Yao Lu, Hongdong Li, and Richard Hartley.  Learning to estimate hidden motions with global motion aggregation.  In ICCV, pages 9772–9781, October 2021.
  * [20] Wei Jiang, Eduard Trulls, Jan Hosang, Andrea Tagliasacchi, and Kwang Moo Yi.  Cotr: Correspondence transformer for matching across images.  ICCV, 2021.
  * [21] Alex Kendall, Hayk Martirosyan, Saumitro Dasgupta, Peter Henry, Ryan Kennedy, Abraham Bachrach, and Adam Bry.  End-to-end learning of geometry and context for deep stereo regression.  In ICCV, pages 66–75, 2017.
  * [22] Daniel Kondermann, Rahul Nair, Katrin Honauer, Karsten Krispin, Jonas Andrulis, Alexander Brock, Burkhard Gussefeld, Mohsen Rahimimoghaddam, Sabine Hofmann, Claus Brenner, et al.  The hci benchmark suite: Stereo and flow ground truth with uncertainties for urban autonomous driving.  In CVPR Workshops, pages 19–28, 2016.
  * [23] Yanghao Li, Yuntao Chen, Naiyan Wang, and Zhaoxiang Zhang.  Scale-aware trident networks for object detection.  In ICCV, pages 6054–6063, 2019.
  * [24] Zhaoshuo Li, Xingtong Liu, Nathan Drenkow, Andy Ding, Francis X Creighton, Russell H Taylor, and Mathias Unberath.  Revisiting stereo depth estimation from a sequence-to-sequence perspective with transformers.  In ICCV, pages 6197–6206, 2021.
  * [25] Tsung-Yi Lin, Piotr Dollár, Ross Girshick, Kaiming He, Bharath Hariharan, and Serge Belongie.  Feature pyramid networks for object detection.  In CVPR, pages 2117–2125, 2017.
  * [26] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo.  Swin transformer: Hierarchical vision transformer using shifted windows.  ICCV, 2021.
  * [27] Ilya Loshchilov and Frank Hutter.  Decoupled weight decay regularization.  arXiv preprint arXiv:1711.05101, 2017.
  * [28] Zhaoyang Lv, Kihwan Kim, Alejandro Troccoli, Deqing Sun, James M Rehg, and Jan Kautz.  Learning rigidity in dynamic scenes with a moving camera for 3d motion field estimation.  In ECCV, pages 468–484, 2018.
  * [29] Nikolaus Mayer, Eddy Ilg, Philip Hausser, Philipp Fischer, Daniel Cremers, Alexey Dosovitskiy, and Thomas Brox.  A large dataset to train convolutional networks for disparity, optical flow, and scene flow estimation.  In CVPR, pages 4040–4048, 2016.
  * [30] Simon Meister, Junhwa Hur, and Stefan Roth.  Unflow: Unsupervised learning of optical flow with a bidirectional census loss.  In AAAI, 2018.
  * [31] Moritz Menze and Andreas Geiger.  Object scene flow for autonomous vehicles.  In CVPR, pages 3061–3070, 2015.
  * [32] Federico Perazzi, Jordi Pont-Tuset, Brian McWilliams, Luc Van Gool, Markus Gross, and Alexander Sorkine-Hornung.  A benchmark dataset and evaluation methodology for video object segmentation.  In CVPR, pages 724–732, 2016.
  * [33] Anurag Ranjan and Michael J Black.  Optical flow estimation using a spatial pyramid network.  In CVPR, pages 4161–4170, 2017.
  * [34] Jerome Revaud, Philippe Weinzaepfel, Zaid Harchaoui, and Cordelia Schmid.  Epicflow: Edge-preserving interpolation of correspondences for optical flow.  In CVPR, pages 1164–1172, 2015.
  * [35] Stephan R Richter, Zeeshan Hayder, and Vladlen Koltun.  Playing for benchmarks.  In ICCV, pages 2213–2222, 2017.
  * [36] Paul-Edouard Sarlin, Daniel DeTone, Tomasz Malisiewicz, and Andrew Rabinovich.  Superglue: Learning feature matching with graph neural networks.  In CVPR, pages 4938–4947, 2020.
  * [37] Deqing Sun, Daniel Vlasic, Charles Herrmann, Varun Jampani, Michael Krainin, Huiwen Chang, Ramin Zabih, William T Freeman, and Ce Liu.  Autoflow: Learning a better training set for optical flow.  In CVPR, pages 10093–10102, 2021.
  * [38] Deqing Sun, Xiaodong Yang, Ming-Yu Liu, and Jan Kautz.  Pwc-net: Cnns for optical flow using pyramid, warping, and cost volume.  In CVPR, pages 8934–8943, 2018.
  * [39] Deqing Sun, Xiaodong Yang, Ming-Yu Liu, and Jan Kautz.  Models matter, so does training: An empirical study of cnns for optical flow estimation.  TPAMI, 42(6):1408–1423, 2019.
  * [40] Jiaming Sun, Zehong Shen, Yuang Wang, Hujun Bao, and Xiaowei Zhou.  Loftr: Detector-free local feature matching with transformers.  In CVPR, pages 8922–8931, 2021.
  * [41] Zachary Teed and Jia Deng.  Raft: Recurrent all-pairs field transforms for optical flow.  In ECCV, pages 402–419. Springer, 2020.
  * [42] Prune Truong, Martin Danelljan, and Radu Timofte.  Glu-net: Global-local universal network for dense flow and correspondences.  In CVPR, pages 6258–6268, 2020.
  * [43] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  In NeurIPS, pages 5998–6008, 2017.
  * [44] Jianyuan Wang, Yiran Zhong, Yuchao Dai, Kaihao Zhang, Pan Ji, and Hongdong Li.  Displacement-invariant matching cost learning for accurate optical flow estimation.  NeurIPS, 33, 2020.
  * [45] Qianqian Wang, Xiaowei Zhou, Bharath Hariharan, and Noah Snavely.  Learning feature descriptors using camera pose supervision.  In ECCV, pages 757–774. Springer, 2020.
  * [46] Wenshan Wang, Delong Zhu, Xiangwei Wang, Yaoyu Hu, Yuheng Qiu, Chen Wang, Yafei Hu, Ashish Kapoor, and Sebastian Scherer.  Tartanair: A dataset to push the limits of visual slam.  2020\.
  * [47] Philippe Weinzaepfel, Jerome Revaud, Zaid Harchaoui, and Cordelia Schmid.  Deepflow: Large displacement optical flow with deep matching.  In ICCV, pages 1385–1392, 2013.
  * [48] Haofei Xu, Jiaolong Yang, Jianfei Cai, Juyong Zhang, and Xin Tong.  High-resolution optical flow from 1d attention and correlation.  In ICCV, pages 10498–10507, 2021.
  * [49] Haofei Xu and Juyong Zhang.  Aanet: Adaptive aggregation network for efficient stereo matching.  In CVPR, pages 1959–1968, 2020.
  * [50] Yufei Xu, Qiming Zhang, Jing Zhang, and Dacheng Tao.  Vitae: Vision transformer advanced by exploring intrinsic inductive bias.  NeurIPS, 2021.
  * [51] Gengshan Yang and Deva Ramanan.  Volumetric correspondence networks for optical flow.  NeurIPS, 32:794–805, 2019.
  * [52] Zhichao Yin, Trevor Darrell, and Fisher Yu.  Hierarchical discrete distribution decomposition for match density estimation.  In CVPR, pages 6044–6053, 2019.
  * [53] Feihu Zhang, Oliver J. Woodford, Victor Adrian Prisacariu, and Philip H.S. Torr.  Separable flow: Learning motion cost volumes for optical flow estimation.  In ICCV, pages 10807–10817, October 2021.
  * [54] Qiming Zhang, Yufei Xu, Jing Zhang, and Dacheng Tao.  Vitaev2: Vision transformer advanced by exploring inductive bias for image recognition and beyond.  arXiv preprint arXiv:2202.10108, 2022.

## Appendix

##  A More Comparisons

We present more comprehensive comparisons (as a supplement of Table 1 in the
main paper) with all the possible combinations of flow estimation approaches
in Table A. Our Transformer and softmax-based method is consistently better
and has less parameters than other variants.

Method |  feature enhancement |  flow prediction | #convs | Sintel (train, clean) | Sintel (train, final) |  Param (M)
---|---|---|---|---|---|---
all | matched | unmatched | all | matched | unmatched
variants | - | cost + conv | 14 | 3.32 | 1.56 | 22.34 | 4.93 | 2.82 | 27.73 | 4.64
Transformer | cost + conv | 2 | 3.41 | 2.40 | 14.32 | 4.57 | 3.27 | 18.67 | 4.95
Transformer | cost + conv | 14 | 2.04 | 1.09 | 12.34 | 3.37 | 2.03 | 17.80 | 7.79
conv | softmax | 14 | 6.36 | 3.22 | 40.30 | 8.00 | 4.80 | 42.58 | 5.12
GMFlow (w/o prop.) | Transformer | softmax | 0 | 2.28 | 1.06 | 15.54 | 3.44 | 1.95 | 19.50 | 4.20
GMFlow (w/ prop.) | Transformer | softmax | 0 | 1.89 | 1.10 | 10.39 | 3.13 | 1.98 | 15.52 | 4.23
Table A: Comparisons on different variants of flow estimation approaches.
Although the Transformer can also be used for feature enhancement in the cost
volume and convolution-based approach (cost + conv), its performance heavily
relies on a deep convolutional regressor (_e.g_., 14 layers to catch up). In
contrast, our softmax-based method is _parameter-free_ (4.20M _vs_. 7.79M).
The flow propagation (prop.) layer further improves ours performance in
unmatched regions, while only introducing additional 0.03M parameters.
Replacing the Transformer with convolutions for feature enhancement leads to
significantly large performance drop, since convolutions are not able to model
the mutual relationship between two features.

##  B Computational Complexity

We analyze the computational complexities of core components in our framework
below.

Global Matching. In our global matching formulation, we build a 4D correlation
matrix $H\times W\times H\times W$ to model all pair-wise similarities between
two features (with size $H\times W$, $1/8$ of the original image resolution).
There exists an equivalent implementation should it become a bottleneck for
high-resolution images. Note that the pixels in the first feature are
_independent_ and thus their flow predictions can be computed _sequentially_.
Specifically, we can sequentially compute $K\times K$ correlation matrices
(each with size $H/K\times W/K\times H\times W$), and finally merge the
results for all pixels. Such a sequential implementation can save the memory
consumption while having little influence on the overall inference time (see
Table B), since the global matching operation only needs to compute once, and
it’s not a significant speed bottleneck in the full framework.

#splits | $1\times 1$ | $2\times 2$ | $4\times 4$ | $8\times 8$
---|---|---|---|---
Time (ms) | 52.57 | 52.64 | 52.90 | 59.45
Table B: Inference time _vs_. number of splits for sequential global matching
implementation. The input image resolution is $448\times 1024$, and the
features are downsampled by $8\times$.

Transformer. We use shifted local window attention [26] in the Transformer
implementation, where each local window size is $1/16$ of the original image
resolution by default. The computational cost is usually acceptable for
regular image resolutions (_e.g_., $448\times 1024$). Note that we can always
switch to smaller windows size (_e.g_., $1/32$, see Table LABEL:tab:split_attn
of the main paper) should it become a bottleneck.

Flow Propagation. Our default flow propagation scheme computes a global self-
attention. The sequential implementation in global matching can also be
adopted here. It’s also possible to compute a local window self-attention only
for less memory consumption by trading some accuracy in large motion (Table
C). Such a local attention operation can be implemented efficiently with
PyTorch’s unfold function.

self-attn. | Sintel (train, final) |
---|---|---
| EPE | $s_{0-10}$ | $s_{10-40}$ | $s_{40+}$ |
global | 3.13 | 0.80 | 3.87 | 18.04 |
local $3\times 3$ | 3.31 | 0.79 | 3.75 | 20.22 |
local $5\times 5$ | 3.21 | 0.75 | 3.66 | 19.69 |
Table C: Global _vs_. local self-attention for flow propagation.

Refinement. Although the feature resolution of our refinement architecture is
higher ($1/4$), it is not a significant bottleneck since smaller local window
($1/32$ of the original image resolution) attention is used in the Transformer
and matching is performed within a local window.

Overall, our GMFlow framework is general and flexible, and many concrete
implementations are possible to meet specific needs.

##  C More Visual Results

Flow Propagation. Our flow propagation scheme with self-attention is quite
effective for handling occluded and out-of-boundary pixels, as can be seen
from Fig. A.

Prediction on DAVIS dataset. We test our pre-trained Sintel model on the DAVIS
[32] dataset, the results on diverse scenes are shown in Fig. B.

Figure A: Our flow propagation (prop.) scheme significantly improves the
performance of occluded and out-of-boundary pixels. Figure B: Visual results
on DAVIS dataset.

##  D More Implementation Details

Network Architectures. The Transformer feature dimension is 128, and the
intermediate feed-forward network expands the dimension by $4\times$. We only
use a single head in all the attention computations, since we observe that
multi-head attention slows down the speed without bringing obvious performance
gains. Our refinement architecture uses exactly the same Transformer for
feature enhancement, except that the attentions are performed within smaller
local windows. The self-attention layer in the flow propagation step is also
shared for $1/8$ and $1/4$ resolutions, where we perform global attention at
$1/8$ resolution and local $3\times 3$ window attention at $1/4$ resolution.

Training Details. Our data augmentation strategy mostly follows RAFT [41]
except that we didn’t use occlusion augmentation, since no obvious improvement
is observed in our experiments. During training, we perform random cropping
following previous works. The crop size for FlyingChairs is $384\times 512$,
FlyingThings3D is $384\times 768$, Sintel is $320\times 896$ and KITTI is
$320\times 1152$. Our framework without refinement is trained on 4 V100 (16GB)
GPUs. The full framework with refinement is trained on 4 A100 (40GB) GPUs. We
are also able to reproduce the results on 4 V100 (16GB) GPUs by halving the
batch size and doubling the training iterations.


11institutetext: Vanderbilt University

# Toro-SAM: Towards robust 3D interactive medical image segmentation with
visual prompts

Hao Li  Han Liu  Dewei Hu  Jiacheng Wang  Ipek Oguz

# PIM: Towards robust 3D interactive medical image segmentation with visual
prompts

Hao Li  Han Liu  Dewei Hu  Jiacheng Wang  Ipek Oguz

# PRISM: Towards robust 3D interactive medical image segmentation with visual
prompts

Hao Li  Han Liu  Dewei Hu  Jiacheng Wang  Ipek Oguz

# PRISM: A Promptable and Robust Interactive Segmentation Model with Visual
Prompts

Hao Li  Han Liu  Dewei Hu  Jiacheng Wang  Ipek Oguz

###### Abstract

In this paper, we present PRISM, a Promptable and Robust Interactive
Segmentation Model, aiming for precise segmentation of 3D medical images.
PRISM accepts various visual inputs, including points, boxes, and scribbles as
sparse prompts, as well as masks as dense prompts. Specifically, PRISM is
designed with four principles to achieve robustness: (1) Iterative learning.
The model produces segmentations by using visual prompts from previous
iterations to achieve progressive improvement. (2) Confidence learning. PRISM
employs multiple segmentation heads per input image, each generating a
continuous map and a confidence score to optimize predictions. (3) Corrective
learning. Following each segmentation iteration, PRISM employs a shallow
corrective refinement network to reassign mislabeled voxels. (4) Hybrid
design. PRISM integrates hybrid encoders to better capture both the local and
global information. Comprehensive validation of PRISM is conducted using four
public datasets for tumor segmentation in the colon, pancreas, liver, and
kidney, highlighting challenges caused by anatomical variations and ambiguous
boundaries in accurate tumor identification. Compared to state-of-the-art
methods, both with and without prompt engineering, PRISM significantly
improves performance, achieving results that are close to human levels. The
code is publicly available at https://github.com/MedICL-VU/PRISM.

###### Keywords:

Interactive medical image segmentation Iterative correction Tumor segmentation
Visual prompts Segment Anything Model (SAM).

##  1 Introduction

In recent years, deep learning-based methods have achieved state-of-the-art
performance in various segmentation tasks [21, 12]. However, achieving robust
outcomes remains a challenge due to significant anatomical variations among
individuals and ambiguous boundaries in medical images. Alternatively,
interactive segmentation models, such as the Segment Anything Model (SAM)
[14], offer a solution by involving humans in the loop and utilizing their
expertise to achieve precise segmentation. This requires users to indicate the
target region by providing visual prompts, such as points [22, 14, 23, 9, 18,
17, 30, 8], boxes [32, 35, 7, 29, 14, 24, 3, 34], scribbles [30, 4, 33, 6],
and masks [14, 27, 31, 5, 22].

Importantly, a robust interactive segmentation model should be effective in
responding to visual prompts given by users with minimal interactions.
Inspired by SAM [14], many interactive segmentation methods have been proposed
in medical imaging. However, most of these SAM-based methods are limited to
using a single type of prompt [32, 9, 34, 8, 18, 7, 24, 3, 35]. This limits
the available information, impacting the effectiveness of the model. Due to
the heterogeneity in anatomy and appearance of medical images, the models
should leverage the advantages of different prompts to achieve optimal
performance across various medical applications. Moreover, these interactive
methods do not involve humans in the loop [9, 8, 35, 18, 7, 32, 24, 3, 34],
i.e., the visual prompts are applied to the automatic segmentation model only
once, which may not be sufficient to achieve robust outcomes with the initial
prompts provided. In contrast, practical applications often necessitate
iterative corrections based on new prompts from the user until the outcome
meets their criteria [26]. Unfortunately, few studies have explored such
human-in-loop approach for medical interactive segmentation [5, 31], and their
performance has been suboptimal.

An additional concern is the efficiency of the model with respect to training
and user interaction requirements. Training a model from scratch with
extensive datasets [5, 33, 24, 31] improves overall representational
capabilities, but this is time-consuming and typically requires substantial
computational resources for optimal performance. Furthermore, models may lack
specificity for particular applications due to existing domain gaps in medical
imaging [31]. This requires an increased number of user prompts to reach
adequate performance. Moreover, 2D models [5, 33, 24] are not considered
“efficient” for 3D images as they require significant user effort to provide
visual prompts for each slice.

In this work, we propose a robust method for interactive segmentation in
medical imaging. We strive for human-level performance, as a human-in-loop
interactive segmentation model with prompts should gradually refine its
outcomes until they closely match inter-rater variability [10]. We present
PRISM, a Promptable and Robust Interactive Segmentation Model for 3D medical
images, which accepts both sparse (points, boxes, scribbles) and dense (masks)
visual prompts. We leverage an iterative learning [14, 27] and sampling
strategy [26] to train PRISM to achieve continual improvements across
iterations. The sparse visual prompts are sampled based on the erroneous
region of the prediction from previous iteration, which simulates the human
correcting behavior. For each iteration, multiple segmentation masks are
generated [19] along with regressed confidence scores. The output with the
highest score is selected, increasing the robustness of the model. This
approach is similar to model ensembling [20]. Moreover, inspired from [25,
13], the output is fed into a shallow corrective refinement network to correct
the mislabeled voxels and refine the final segmentation, adhering to the goal
of being effectiveness. A hybrid image encoder with parallel convolutional and
transformer paths [16] is used as the backbone to better capture the local and
global information from a medical image. Our main contributions are summarized
as:

  * •

Interaction: PRISM accepts various visual prompts, offering versatility that
effectively addresses the diverse challenges of medical imaging segmentation.
This approach ensures precise outcomes through user-friendly interaction.

  * •

Human-in-loop: PRISM is trained with iterative and confidence learning,
allowing for continuous improvements and robust performance. With each user
input, the performance progressively approaches expert-level precision.

  * •

Network architecture: PRISM employs a corrective refinement network and a
hybrid encoder to produce precise segmentation for challenging conditions.

We evaluate PRISM on four public tumor datasets, including tumors in colon
[1], pancreas [1], liver [2] and kidney [11], where anatomical differences
among individuals and ambiguous boundaries are present. Comprehensive
validation is performed against state-of-the-art automatic and interactive
methods, and PRISM significantly outperforms all of them with substantial
improvements.

Figure 1: (a) PRISM takes an image ($x$) and visual prompts ($v$) to produce a
segmentation ($y^{\prime})$. The user then provides prompts for next
iteration. (b) Interaction between image and visual features in the latent
space to produce image ($Z_{x}$) and visual ($Z_{v}$) embeddings with self-
and cross-attention mechanisms.

##  2 Methods

PRISM employs a generic encoder-decoder architecture as displayed in Fig.
1(a). For a given input image $x$, along with visual prompts $v$, the image
and prompt encoders generate the image and visual features, respectively.
These latent features then interact through self- and cross-attention
mechanisms (Fig. 1(b)). The resulting embeddings, $Z_{x}$ and $Z_{v}$, are fed
into the decoder to produce the final segmentation $y^{\prime}$. The expert
then gives sparse prompt based on the erroneous regions in $y^{\prime}$ for
next iteration, in human-in-loop manner.

Fig. 2(a) details PRISM, where a hybrid encoder combines parallel CNN and
vision transformer paths to better extract image features. The last feature
map of decoder $f_{d}$ is integrated with the visual embedding to generate
multiple mask predictions and associated confidence scores (Fig. 2(b)). A
selector is employed to pick the candidate prediction with the highest
confidence score, which is subsequently fed into the corrective refinement
network to generate the $y^{\prime}$. In our experiments, we use a sampler to
mimic expert corrective actions by identifying point and scribble prompts from
the false positive and negative areas within $y^{\prime}$.

Iterative learning. PRISM is designed for clinical applications, which require
improvement through successive iterations. We define the loss for a single
iteration as $L_{i}$, and the total loss as $L_{total}=\sum_{i=1}^{N}L_{i}$,
where $N$ denotes the total number of iterations. The input visual prompts
come only from the current iteration, i.e., they are not cumulative. Since the
dense prompt $y^{\prime}_{i}$ is generated from the previous iteration and
retains the gradient, PRISM can learn the relationship between iterations. For
efficiency, the image features are generated only once at the initial
iteration. The following prompt types are allowed (Fig. 2(b)):

  * •

At each iteration $i$, point prompts are randomly sampled with uniform
distribution from the false negative (FN) and positive (FP) regions of
$y^{\prime}_{i-1}$.

  * •

The scribbles are generated similar to Scribbleprompt [33] by first extracting
the skeleton of the FN and FP regions of $y^{\prime}_{i-1}$ from the previous
iteration. Next, a random mask is created by pixels randomly sampled from a
normal distribution, applying Gaussian blur, and thresholding the image based
on its mean value. This random mask is applied to divide the skeleton of
$y^{\prime}_{i-1}$ into separate, smaller parts. Finally, a random deformation
field and random Gaussian filtering are used to change the scribble curvature
and thickness.

  * •

The 3D bounding box (BB) is determined based on the ground truth $y$ at the
initial iteration ($i=1$) and remains unchanged for the subsequent iterations.

  * •

The logits map at iteration $i-1$ is used to provide additional information
instead of a binary mask as the dense prompt for iteration $i$, for $i>1$.

Figure 2: (a) Details of PRISM. Green highlights the corrective refinement
network. (b) Top row shows the multi-mask prediction with labeled confidence
scores. The selector would then pick Mask 3 as the dense prompt. Possible
visual prompts given this dense prompt are shown in the bottom row.

Confidence learning. Unlike traditional methods that output a single mask,
PRISM increases robustness by generating multiple outputs, each with a
confidence score. The last feature map of the decoder, denoted as
$f_{d}=Decoder(Z_{x})$, interacts with the visual embedding to produce
continuous maps $m_{j}$ and confidence scores $s_{j}$ via $m_{j}=f_{d}\times
MLP^{m}_{j}(Z_{v})$ and $s_{j}=MLP^{s}_{j}(Z_{v})$, where $MLP$ indicates the
multi-layer perceptrons, and $j$ is an index over the multiple outputs. Thus,
$m_{j}$ interacts with visual prompts at both latent- and image-level. The
process of confidence supervision for each iteration is defined by:
$L_{con}=\sum_{j=1}^{M}(\lambda_{s}L_{s}(m_{j},y)+\lambda_{b}L_{b}(m_{j},y)+\lambda_{r}L_{r}(s_{j},1-L_{Dice}(m_{j},y)))$,
where $L_{s}$, $L_{b}$, $L_{r}$ and $L_{Dice}$ indicate structural, boundary,
regression and Dice loss, respectively. More precisely, the structural
information is captured through a combination of Dice loss and cross-entropy
loss: $L_{s}=L_{Dice}+L_{CE}$. For a given logits map or binary mask, the
smooth boundary is derived as: $B(m)=|m-P_{ave}(m)|$, where $P_{ave}$ denotes
the average pooling layer. The boundary loss ($L_{b}$) is evaluated by
$L_{MSE}(B(m),B(y))$, with $L_{MSE}$ representing the mean square error loss.
We also use $L_{MSE}$ for the regression loss.

Corrective learning. As shown in Fig. 2(a), a selector is employed to choose
the candidate mask $\hat{m}$ corresponding to the highest confidence score
$s_{j}$, which is then fed into the corrective refinement network. We design
this network to be both shallow and efficient, consisting of two cascaded
residual blocks and two convolution operations in two different branches. It
takes the input $x_{c}$ at original image resolution. Importantly, $x_{c}$ is
the concatenation of 4 images: (1) the input image $x$, (2) the binary mask
$\hat{b}$ derived from $\hat{m}$, (3) the binary positive prompt maps
accumulated from iteration 0 to i, and (4) the similar binary cumulative
negative prompt map. For efficiency, the $x_{c}$ is downsampled and then
upsampled back for $y^{\prime}_{i}$. The final segmentation $y^{\prime}_{i}$
is supervised by:
$L_{cor}=\lambda_{s}L_{s}(y^{\prime}_{i},y)+\lambda_{b}L_{b}(y^{\prime}_{i},y)$.
Notably, the corrective network does not backpropagate to the network weights
for candidate mask generation and focuses exclusively on refining any given
logit map. The $y^{\prime}_{i}$ is fed to sampler for next iteration $i+1$.

Implementation details. For training, the total objective function is:
$L_{total}=\sum_{i=1}^{N}(L_{{con}_{i}}+L_{{cor}_{i}})$, where the ratio of
$\lambda_{s}:\lambda_{b}:\lambda_{r}$ as $1:10:1$. The batch size was set to
2, and the initial learning rate was $4e^{-5}$, which was reduced by $2e^{-6}$
after every epoch. We employed the AdamW optimizer across a maximum of 200
epochs. Our study was conducted on an NVIDIA A6000. The input image was
cropped to a 3D $128\times 128\times 128$ patch centered on a foreground
pixel. The number of iterations ($N$) was fixed at 11, with 1 to 50 point
prompts randomly chosen per iteration during training. The Dice and normalized
surface Dice (NSD) are used as evaluation metrics. We compared to both fully
automated [12, 28, 15] and interactive [9, 18, 31, 14] methods. We used
$y^{\prime}_{i=11}$ as final segmentation for the iterative methods [31] and
PRISM. We retrained using the published code of each model, and the pretrained
weights were also employed if publicly available. We present two versions of
our method. PRISM-plain only uses point prompts, while PRISM-ultra can handle
other sparse visual prompts such as boxes and scribbles. All results are
generated with seeds.

##  3 Experiments

Datasets. We use four public 3D CT datasets, including colon (N=126) [1],
pancreas (N=281) [1], liver (N=118) [2] and kidney (N=300) [11] tumors.
Challenges include large anatomical differences among individuals, varying
shapes of target structures, and ambiguous boundaries. Additionally, multiple
tumors [11, 2] may occur in a single subject. We adopted the same data split
as a prior study [9] for each task with a training/validation/testing split of
0.7/0.1/0.2, and we only use tumor labels to focus on binary segmentation. We
resampled to a $1mm$ isotropic resolution, performed intensity clipping based
on the $0.5$ and $99.5$ percentiles of the foreground, and applied Z-score
normalization based on foreground voxels. Random zoom and intensity shift are
used as data augmentations.

Table 1: Quantitative results. Bold indicates best performance. PRISM-plain is a simplified model for a fair comparison to methods not allowing BB prompts [9, 18, 31], and methods[14, 9, 18, 31] only using 1 point prompt per input. For liver tumor, 10 points are used for all methods to accommodate multiple tumors. |  Public datastes (Dice $\%$ / NSD $\%$)
---|---
Methods | Colon tumor | Pancreas tumor | Liver tumor | Kidney tumor
nnU-Net [12] | 43.91 / 52.52 | 41.65 / 62.54 | 60.10 / 75.41 | 73.07 / 77.47
3D UX-Net [15] | 28.50 / 32.73 | 34.83 / 52.56 | 45.54 / 60.67 | 57.59 / 58.55
Swin-UNETR [28] | 35.21 / 42.94 | 40.57 / 60.05 | 50.26 / 64.32 | 65.54 / 72.04
SAM [14] | 28.83 / 33.63 | 24.01 / 26.74 | 8.56 / 5.97 | 36.30 / 29.86
3DSAM-adapter [9] | 57.32 / 73.65 | 54.41 / 77.88 | 56.61 / 69.52 | 73.78 / 83.86
ProMISe [18] | 66.81 / 81.24 | 57.46 / 79.76 | 58.78 / 71.52 | 75.70 / 80.08
SAM-Med3D [31] | 54.34 / 78.58 | 65.61 / 92.40 | 23.64 / 26.97 | 76.50 / 88.41
SAM-Med3D-organ | 70.75 / 91.03 | 76.40 / 97.75 | 66.52 / 77.97 | 88.20 / 97.80
SAM-Med3D-turbo | 73.77 / 94.95 | 74.87 / 96.43 | 69.36 / 81.70 | 89.26 / 98.40
PRISM-plain | 67.18 / 85.28 | 65.73 / 89.51 | 79.70 / 91.60 | 85.29 / 93.55
PRISM-ultra |  93.79 / 99.96 |  94.48 / 99.99 |  94.18 / 99.99 |  96.58 / 99.80
| “PRISM-plain” only uses 1 point and no BB. “-ultra” adds the BB, and
scribbles.
---

Comparison with state-of-the-art. Table 1 compares quantitative segmentation
results (See qualitative results in Fig. S0..1). As expected, fully automated
methods [12, 28, 15] are unable to deliver satisfactory outcomes for these
challenging datasets. Even with visual prompts, SAM [14] struggles with the
substantial domain shift between natural and medical images. Furthermore, the
interactive segmentation models [9, 18, 31] trained from scratch fail to
produce adequate results. Notably, “-organ” and “-turbo” denote the fine-
tuning of the pretrained SAM-Med3D model [31]. “-organ” focuses on organ-
specific datasets, while “-turbo” includes these organ datasets (including, in
some cases, the test data in their pretraining) alongside a broader range of
other medical imaging data. In contrast, excluding the pretrained models, the
proposed PRISM-plain demonstrates superior overall Dice performance with only
a single point as sparse prompt. Moreover, as more visual prompts (points, BB,
and scribbles) are provided, PRISM-ultra delivers outcomes close to human
level performance, demonstrating its strength as an effective human-in-loop
algorithm.

Figure 3: Dice score of proposed PRISM on four tumor datasets, where the mean
values (lines) and their 95% confidence intervals (shades) are presented.
Figure 4: Qualitative results of PRISM-ultra for colon tumors characterized by
irregular shapes and ambiguous boundaries. The orange arrows indicate the
major defects which are corrected in the subsequent iteration. The initial
output has noticeable errors that rapidly get corrected in the first few
iterations. More qualitative results can be viewed in Fig. S0..2.

Analysis of iterative learning performance. Fig. 3 shows the performance of
PRISM across the iterative learning process. Although PRISM-plain delivers
superior outcomes (Table 1), it does not have monotonically increasing
performance across iterations and can suffer from a decrease in performance at
later iterations, except for liver tumor. In practice this would be
frustrating to the human user, as more input prompts paradoxically lowers the
segmentation accuracy. However, PRISM-ultra presents not just substantial
improvements in overall performance, but also a monotonically increasing
accuracy and narrower 95% confidence intervals. This highlights the PRISM-
ultra as a robust model, and its qualitative results in Fig. 4 and Fig. S0..2
illustrate the iterative correction. Initially, the results may not meet
expectations due to ambiguous boundaries and anatomical variations. Yet, as
iterations advance, the outcomes progressively align more closely with the
human label. Major corrections occur at the early stages which indicate the
efficiency and effectiveness of PRISM-ultra.

Table 2: Ablation study on the colon tumor. ConL: confidence learning. PRISM-
plain-b: PRISM-plain plus BB. Model configurations are detailed in Table
S0..2.

| ViT encoder | CNN encoder | Hybrid | PRISM-plain | PRISM-plain-b | no ConL
---|---|---|---|---|---|---
Dice | 58.38 | 63.58 | 66.40 | 67.18 | 76.40 | 73.65

Table 3: Prompt analysis for colon tumors. “-1”, “-50”: number of points used
during training. PRISM-basic: PRISM-plain-b with random number of points.
“-erode” and “-dilate”: modified BB by 5 voxel radius. CorL: corrective
learning. PRISM-ultra: PRISM-basic plus test-time scribbles. PRISM-ultra+:
PRISM-ultra plus train-time scribbles. Prompt configurations are detailed in
Table S0..2.

Test Points | -plain-b-1 | -plain-b-50 | -basic | no CorL |  -erode |  -dilate |  -ultra |  -ultra+
---|---|---|---|---|---|---|---|---
1 | 76.40 | 61.33 | 76.66 | 77.65 | 72.99 | 75.47 | 93.79 | 96.47
10 | 77.71 | 82.42 | 83.85 | 80.12 | 82.85 | 83.72 | 93.95 | 96.54
50 | 72.52 | 89.91 | 89.36 | 81.60 | 89.99 | 89.42 | 94.07 | 96.73
100 | 69.60 | 91.30 | 90.26 | 81.20 | 91.06 | 90.70 | 94.28 | 96.93

Analysis of framework. An ablation study for various frameworks on the colon
dataset are presented in Table 3, and the details of model configurations in
Table S0..2. The CNN encoder outperforms the use of a Vision Transformer (ViT)
alone. Using a hybrid encoder improves results, while proposed confidence and
corrective learning further increase PRISM-plain performance. PRISM-plain is
advanced to PRISM-plain-b through the incorporation of a 3D BB, a favored
prompt in medical segmentation due to its strong image prior, leading to
improved outcomes. For efficiency, drawing a box does not require
significantly more effort than point prompt. For effectiveness, points can
address the limited flexibility of boxes in making corrections in subsequent
iterations. The proposed confidence learning approach also contributes to the
performance of the model.

Analysis of number of prompts. We analyze the performance with respect to the
number of point prompts in Table 3. The compared model configurations are
detailed in Table S0..2. We first evaluate the number of prompts used for
training. The model plain-b-1 uses 1 prompt for training, plain-b-50 uses 50,
and basic uses a random number between 1 and 50. Too few or too many training
points lowers performance, while the random selection is the most robust. We
also note that corrective learning has a more pronounced contribution when
more prompts are used. We next test the sensitivity to BB precision by
dilating and eroding the true BB. Either eroding or dilating BB decreases
performance compared to PRISM-basic when few points are used, but using more
points resolves this issue. Adding test-time scribbles (ultra) and train-time
scribbles (ultra+) both improves the performance of PRISM-basic. However, this
improvement increases training time, making scribbles less practical for large
datasets.

##  4 Conclusion

In this study, we propose PRISM, a Promptable Robust Interactive Segmentation
Model for 3D medical images. It supports various visual prompts and applies
diverse learning strategies to achieve performance comparable to that of human
raters. This is demonstrated across four challenging tumor segmentation tasks,
where existing state-of-the-art automatic and interactive segmentation
approaches do not meet the expected standards. Future work will include
adaptation of pretrained models, and learning from datasets with sparse
annotations.

Acknowledgments. This work was supported, in part, by NIH U01-NS106845, and
NSF grant 2220401.

## References

  * [1] Antonelli, M., Reinke, A., Bakas, S., Farahani, K., Kopp-Schneider, A., Landman, B.A., Litjens, G., Menze, B., Ronneberger, O., Summers, R.M., et al.: The medical segmentation decathlon. Nature communications 13(1), 4128 (2022)
  * [2] Bilic, P., Christ, P., Li, H.B., Vorontsov, E., Ben-Cohen, A., Kaissis, G., Szeskin, A., Jacobs, C., Mamani, G.E.H., Chartrand, G., et al.: The liver tumor segmentation benchmark (lits). Medical Image Analysis 84, 102680 (2023)
  * [3] Chen, C., Miao, J., Wu, D., Yan, Z., Kim, S., Hu, J., Zhong, A., Liu, Z., Sun, L., Li, X., et al.: Ma-sam: Modality-agnostic sam adaptation for 3d medical image segmentation. arXiv preprint arXiv:2309.08842 (2023)
  * [4] Chen, X., Cheung, Y.S.J., Lim, S.N., Zhao, H.: Scribbleseg: Scribble-based interactive image segmentation. arXiv preprint arXiv:2303.11320 (2023)
  * [5] Cheng, J., Ye, J., Deng, Z., Chen, J., Li, T., Wang, H., Su, Y., Huang, Z., Chen, J., Jiang, L., et al.: Sam-med2d. arXiv preprint arXiv:2308.16184 (2023)
  * [6] Cho, S., Jang, H., Tan, J.W., Jeong, W.K.: Deepscribble: interactive pathology image segmentation using deep neural networks with scribbles. In: 2021 IEEE 18th international symposium on biomedical imaging (ISBI). pp. 761–765. IEEE (2021)
  * [7] Deng, G., Zou, K., Ren, K., Wang, M., Yuan, X., Ying, S., Fu, H.: Sam-u: Multi-box prompts triggered uncertainty estimation for reliable sam in medical image. In: International Conference on Medical Image Computing and Computer-Assisted Intervention. pp. 368–377. Springer (2023)
  * [8] Gao, Y., Xia, W., Hu, D., Gao, X.: Desam: Decoupling segment anything model for generalizable medical image segmentation. arXiv preprint arXiv:2306.00499 (2023)
  * [9] Gong, S., Zhong, Y., Ma, W., Li, J., Wang, Z., Zhang, J., Heng, P.A., Dou, Q.: 3dsam-adapter: Holistic adaptation of sam from 2d to 3d for promptable medical image segmentation. arXiv preprint arXiv:2306.13465 (2023)
  * [10] Haarburger, C., Müller-Franzes, G., Weninger, L., Kuhl, C., Truhn, D., Merhof, D.: Radiomics feature reproducibility under inter-rater variability in segmentations of ct images. Scientific reports 10(1), 12688 (2020)
  * [11] Heller, N., Isensee, F., Maier-Hein, K.H., Hou, X., Xie, C., Li, F., Nan, Y., Mu, G., Lin, Z., Han, M., et al.: The state of the art in kidney and kidney tumor segmentation in contrast-enhanced ct imaging: Results of the kits19 challenge. Medical image analysis 67, 101821 (2021)
  * [12] Isensee, F., Jaeger, P.F., Kohl, S.A., Petersen, J., Maier-Hein, K.H.: nnu-net: a self-configuring method for deep learning-based biomedical image segmentation. Nature methods 18(2), 203–211 (2021)
  * [13] Jang, W.D., Kim, C.S.: Interactive image segmentation via backpropagating refinement scheme. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pp. 5297–5306 (2019)
  * [14] Kirillov, A., Mintun, E., Ravi, N., Mao, H., Rolland, C., Gustafson, L., Xiao, T., Whitehead, S., Berg, A.C., Lo, W.Y., et al.: Segment anything. arXiv preprint arXiv:2304.02643 (2023)
  * [15] Lee, H.H., Bao, S., Huo, Y., Landman, B.A.: 3d UX-net: A large kernel volumetric convnet modernizing hierarchical transformer for medical image segmentation. In: The Eleventh International Conference on Learning Representations (2023)
  * [16] Li, H., Hu, D., Liu, H., Wang, J., Oguz, I.: Cats: Complementary cnn and transformer encoders for segmentation. In: 2022 IEEE 19th International Symposium on Biomedical Imaging (ISBI). pp. 1–5. IEEE (2022)
  * [17] Li, H., Liu, H., Hu, D., Wang, J., Oguz, I.: Assessing test-time variability for interactive 3d medical image segmentation with diverse point prompts. arXiv preprint arXiv:2311.07806 (2023)
  * [18] Li, H., Liu, H., Hu, D., Wang, J., Oguz, I.: Promise: Prompt-driven 3d medical image segmentation using pretrained image foundation models. arXiv preprint arXiv:2310.19721 (2023)
  * [19] Li, Z., Chen, Q., Koltun, V.: Interactive image segmentation with latent diversity. In: Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition. pp. 577–585 (2018)
  * [20] Linmans, J., Elfwing, S., van der Laak, J., Litjens, G.: Predictive uncertainty estimation for out-of-distribution detection in digital pathology. Medical Image Analysis 83, 102655 (2023)
  * [21] Liu, H., Hu, D., Li, H., Oguz, I.: Medical image segmentation using deep learning. Machine Learning for Brain Disorders pp. 391–434 (2023)
  * [22] Liu, Q., Xu, Z., Bertasius, G., Niethammer, M.: Simpleclick: Interactive image segmentation with simple vision transformers. In: Proceedings of the IEEE/CVF International Conference on Computer Vision. pp. 22290–22300 (2023)
  * [23] Luo, X., Wang, G., Song, T., Zhang, J., Aertsen, M., Deprest, J., Ourselin, S., Vercauteren, T., Zhang, S.: Mideepseg: Minimally interactive segmentation of unseen objects from medical images using deep learning. Medical image analysis (2021)
  * [24] Ma, J., He, Y., Li, F., Han, L., You, C., Wang, B.: Segment anything in medical images. Nature Communications 15(1), 654 (2024)
  * [25] Oguz, B.U., Shinohara, R.T., Yushkevich, P.A., Oguz, I.: Gradient boosted trees for corrective learning. In: International Workshop on Machine Learning in Medical Imaging (2017)
  * [26] Sofiiuk, K., Petrov, I.A., Konushin, A.: Reviving iterative training with mask guidance for interactive segmentation (2021)
  * [27] Sun, S., Xian, M., Xu, F., Yao, T., Capriotti, L.: Cfr-icl: Cascade-forward refinement with iterative click loss for interactive image segmentation. arXiv preprint arXiv:2303.05620 (2023)
  * [28] Tang, Y., Yang, D., Li, W., Roth, H., Landman, B., Xu, D., Nath, V., Hatamizadeh, A.: Self-supervised pre-training of swin transformers for 3d medical image analysis (2022)
  * [29] Wang, G., Li, W., Zuluaga, M.A., Pratt, R., Patel, P.A., Aertsen, M., Doel, T., David, A.L., Deprest, J., Ourselin, S., et al.: Interactive medical image segmentation using deep learning with image-specific fine tuning. IEEE transactions on medical imaging 37(7), 1562–1573 (2018)
  * [30] Wang, G., Zuluaga, M.A., Li, W., Pratt, R., Patel, P.A., Aertsen, M., Doel, T., David, A.L., Deprest, J., Ourselin, S., et al.: Deepigeos: a deep interactive geodesic framework for medical image segmentation. IEEE transactions on pattern analysis and machine intelligence 41(7), 1559–1572 (2018)
  * [31] Wang, H., Guo, S., Ye, J., Deng, Z., Cheng, J., Li, T., Chen, J., Su, Y., Huang, Z., Shen, Y., et al.: Sam-med3d. arXiv preprint arXiv:2310.15161 (2023)
  * [32] Wei, X., Cao, J., Jin, Y., Lu, M., Wang, G., Zhang, S.: I-medsam: Implicit medical image segmentation with segment anything. arXiv preprint (2023)
  * [33] Wong, H.E., Rakic, M., Guttag, J., Dalca, A.V.: Scribbleprompt: Fast and flexible interactive segmentation for any medical image. arXiv preprint (2023)
  * [34] Yao, X., Liu, H., Hu, D., Lu, D., Lou, A., Li, H., Deng, R., Arenas, G., Oguz, B., Schwartz, N., et al.: False negative/positive control for sam on noisy medical images. arXiv preprint (2023)
  * [35] Zhang, Y., Hu, S., Jiang, C., Cheng, Y., Qi, Y.: Segment anything model with uncertainty rectification for auto-prompting medical image segmentation. arXiv preprint arXiv:2311.10529 (2023)

Supplementary Materials

Figure S0..1: Qualitative results of four different tumor segmentation tasks.
The orange arrows indicate the major defects. Figure S0..2: Qualitative
results of PRISM-ultra across iterations. The orange arrows indicate the major
defects which are corrected in the subsequent iteration. Table S0..1: Detailed
settings for architecture and learning strategy in the proposed ablation
studies.

| Encoder type | Learning strategy
---|---|---
Variants | ViT | CNN | ViT + CNN | Confidence learning | Corrective learning
ViT encoder | ✓ |  |  |  |
CNN encoder |  | ✓ |  |  |
Hybrid |  |  | ✓ |  |
PRISM-plain |  |  | ✓ | ✓ | ✓
PRISM-plain-b |  |  | ✓ | ✓ | ✓
no ConL |  |  | ✓ |  | ✓
PRISM-basic |  |  | ✓ | ✓ | ✓
no CorL |  |  | ✓ | ✓ |
PRISM-ultra |  |  | ✓ | ✓ | ✓

Table S0..2: Detailed prompt setting for the proposed ablation studies. Tr.
and T. represent training and test, respectively. The notation [1, 50]
indicates the range from which the number of point prompts is randomly sampled
during training. The term “varies” refers to the number of points used in
tests, which include 1, 10, 50, and 100. “-erode” and “-dilate” denote box
sizes that are 5 voxels smaller or larger in each dimension, respectively.

| Detailed prompt setting
---|---
Variants | point | point num. (Tr. / T.) | box | box type in T. | scribble usage
ViT encoder | ✓ | 1 / 1 |  |  |
CNN encoder | ✓ | 1 / 1 |  |  |
PRISM-plain | ✓ | 1 / 1 |  |  |
PRISM-plain-b | ✓ | 1 / 1 | ✓ | tight |
no ConL | ✓ | 1 / 1 | ✓ | tight |
-plain-b-1 | ✓ | 1 / 1 | ✓ | tight |
-plain-b-50 | ✓ | 50 / varies | ✓ | tight |
PRISM-basic | ✓ | [1, 50] / 1 | ✓ | tight |
no CorL | ✓ | [1, 50] / 1 | ✓ | tight |
-erode | ✓ | [1, 50] / varies | ✓ | undersized |
-dilate | ✓ | [1, 50] / varies | ✓ | oversized |
PRISM-ultra | ✓ | [1, 50] / varies | ✓ | tight | T.
PRISM-ultra+ | ✓ | [1, 50] / varies | ✓ | tight | Tr. and T.