RAFT

RAFT (Recurrent All-Pairs Field Transforms) is a deep network for optical flow estimation, introduced by Zachary Teed and Jia Deng at ECCV 2020. It is a convolutional network with a recurrent update operator: it compares every pixel of the first frame with every pixel of the second in a 4D correlation volume, then refines a single flow field over many iterations with a convolutional GRU that reads from that volume. This design, matching once and refining iteratively at one resolution, became a template for later learned flow, stereo, and scene flow models.

Problem Addressed

Before RAFT, the strongest flow networks, such as PWC-Net and LiteFlowNet, followed the classical coarse-to-fine scheme. Teed and Deng named three weaknesses of this cascade: coarse-level errors are hard to undo, small fast-moving objects are missed, and multi-stage cascades need long training, often over a million iterations.

Earlier refinement networks also mostly used separate weights for each step, which fixed the number of refinements; the one recurrent exception the authors cite, IRR, was limited to five iterations or to the number of pyramid levels. RAFT instead keeps one flow field at a fixed resolution, looks up both small and large displacements in a precomputed multi-scale correlation volume, and applies a small weight-tied update many times. The authors liken the update to a step of a first-order optimizer whose features and motion priors are learned.

Architecture

Given frames $I_1$ and $I_2$, RAFT extracts features and builds the correlation volume once, then iteratively updates a flow field $\mathbf{f}$ initialized to zero, with all stages trained end to end. The full model has 4.8 million parameters, or 5.3 million with the learned upsampling module; a small version, RAFT-S, has about 1 million.

Feature Encoder

A convolutional encoder $g_\theta$ maps each frame to features at one-eighth resolution, $g_\theta : \mathbb{R}^{H \times W \times 3} \to \mathbb{R}^{H/8 \times W/8 \times D}$ with $D = 256$. It consists of six residual blocks, two each at one-half, one-quarter, and one-eighth resolution, uses instance normalization, and is applied with shared weights to both frames. RAFT-S replaces the residual units with bottleneck units.

Context Encoder

A second network $h_\theta$ with the same architecture, applied only to $I_1$, uses batch normalization in the full model. In the official implementation its output is split in two: one part, passed through $\tanh$, initializes the GRU’s hidden state, and the other is injected as context into every update step.

Correlation Volume

The visual similarity between all pairs of feature vectors forms a 4D volume, where $H$ and $W$ here denote the size of the feature maps:

$$C_{ijkl} = \sum_{h} g_\theta(I_1)_{ijh} \, g_\theta(I_2)_{klh}, \qquad \mathbf{C} \in \mathbb{R}^{H \times W \times H \times W},$$

computed as a single matrix multiplication (the official code also divides by $\sqrt{D}$). RAFT then average-pools the last two dimensions, those of $I_2$, with kernel sizes 1, 2, 4, and 8, giving a four-level correlation pyramid. Because the $I_1$ dimensions are never pooled, every pixel of the first frame keeps its own full-resolution correlation map, which the authors credit for recovering small, fast objects.

A lookup operator turns the volume into features. For each pixel $\mathbf{x}$ of $I_1$, it takes the estimated correspondence $\mathbf{x}' = \mathbf{x} + \mathbf{f}(\mathbf{x})$ and samples each pyramid level bilinearly at integer offsets within radius $r$ of $\mathbf{x}'$, scaled to that level. The paper defines this neighborhood with the $L^1$ distance, $\lVert \mathbf{dx} \rVert_1 \le r$, but the official code samples the full square of $(2r+1)^2$ offsets. With $r = 4$ for the full model, the coarsest level reaches offsets of 256 pixels at input resolution. The samples from all levels are concatenated.

The volume costs $O(N^2)$ for $N$ feature pixels, once per image pair; an equivalent variant pools the features of $I_2$ instead and computes correlations on demand, at $O(NM)$ for $M$ iterations.

Update Operator

The update operator is a convolutional GRU with weights tied across iterations. Its input $x_t$ concatenates encoded correlation features, encoded current flow, and context features:

$$\begin{aligned} z_t &= \sigma\big(\mathrm{Conv}_{3 \times 3}([h_{t-1}, x_t], W_z)\big) \\ r_t &= \sigma\big(\mathrm{Conv}_{3 \times 3}([h_{t-1}, x_t], W_r)\big) \\ \tilde{h}_t &= \tanh\big(\mathrm{Conv}_{3 \times 3}([r_t \odot h_{t-1}, x_t], W_h)\big) \\ h_t &= (1 - z_t) \odot h_{t-1} + z_t \odot \tilde{h}_t \end{aligned}$$

The full model replaces each $3 \times 3$ GRU with two in sequence, one with $1 \times 5$ and one with $5 \times 1$ convolutions, to enlarge the receptive field cheaply; RAFT-S uses a single $3 \times 3$ GRU. Two convolutions on the hidden state predict a residual $\Delta\mathbf{f}$, and the estimate becomes $\mathbf{f}_{k+1} = \mathbf{f}_k + \Delta\mathbf{f}$, always at one-eighth resolution. The update operator has 2.7 million parameters.

Upsampling

Each full-resolution flow vector is a convex combination of the $3 \times 3$ coarse vectors around it. Two convolutions on the hidden state predict an $H/8 \times W/8 \times (8 \times 8 \times 9)$ mask, and a softmax over the 9 neighbors makes the weights convex. The authors found this more accurate than bilinear upsampling, particularly near motion boundaries.

Training

RAFT is trained on ground-truth flow: pretraining on the synthetic FlyingChairs (100k iterations) and FlyingThings3D (100k iterations) datasets, then dataset-specific fine-tuning. For Sintel, the authors fine-tuned for 100k iterations on a mixture of Sintel, FlyingThings3D, KITTI 2015, and HD1K; for KITTI, they fine-tuned that model for another 50k iterations on KITTI 2015. Training used AdamW, photometric, spatial, and occlusion augmentation, and 12 unrolled updates.

The loss is the $L^1$ distance between ground truth $\mathbf{f}_{\mathrm{gt}}$ and every intermediate prediction $\mathbf{f}_1, \dots, \mathbf{f}_N$, upsampled to full resolution, with weights that increase exponentially toward the last iteration:

$$\mathcal{L} = \sum_{i=1}^{N} \gamma^{N-i} \left\lVert \mathbf{f}_{\mathrm{gt}} - \mathbf{f}_i \right\rVert_1, \qquad \gamma = 0.8.$$

Gradients flow through each update $\Delta\mathbf{f}$ but are stopped through the incoming estimate $\mathbf{f}_k$.

Inference

RAFT takes two RGB frames and returns a full-resolution flow field. Because the update is weight-tied, the number of iterations is chosen at test time: the authors evaluated with 32 updates on Sintel and 24 on KITTI. In their ablation, accuracy changes little beyond 32 updates, and even 200 updates do not diverge.

For video, a single-resolution design allows warm starting: the flow from the previous frame pair is forward-projected to initialize the next, with occlusion gaps filled by nearest-neighbor interpolation.

The authors reported 10 frames per second at $1088 \times 436$ on a GTX 1080 Ti with 10 updates (0.10 s per pair, against 0.04 s for PWC-Net+ on the same GPU), and 20 frames per second for RAFT-S. On 1080p video, 12 updates took 550 ms per frame, of which the all-pairs correlation took 95 ms.

Performance

As reported by Teed and Deng (2020), with endpoint error (EPE) and, on the KITTI 2015 test set, Fl-all, the percentage of outlier pixels:

• Generalization from synthetic data. Trained only on FlyingChairs and FlyingThings3D, RAFT reached an EPE of 1.43 on the Sintel training set’s clean pass and 2.71 on its final pass, and 5.04 on the KITTI 2015 training set, compared with 8.36 for the best earlier network trained on the same data.
• Sintel test set (2020). With warm start, RAFT reached 1.61 on the clean pass and 2.86 on the final pass (2.855 in the abstract), a 30% reduction from the best published final-pass result of 4.098. The two-frame model scored 1.94 and 3.18.
• KITTI 2015 test set (2020). Fl-all of 5.10%, against 6.10% for the best published method.
• Efficiency. After the same synthetic training, RAFT-S had a lower Sintel final-pass error than PWC-Net and VCN, which are more than six times larger.

Strengths

• Large and small displacements are handled together, without a coarse-to-fine cascade.
• Accuracy can be traded for speed at test time through the number of iterations.
• Few parameters, about ten times fewer training iterations than earlier networks by the authors’ account, and strong generalization from synthetic data.

Limitations

• The all-pairs volume grows quadratically with the number of pixels, which limits resolution on memory-constrained hardware unless the slower on-demand variant is used.
• The update operator relies on local correlation lookups and convolutions, so occluded pixels, which have no match, receive motion only through spatial propagation; GMA was designed to address this.
• Iterative refinement makes it slower than single-pass networks such as PWC-Net+.
• Like other supervised flow networks, it depends on synthetic pretraining and benchmark-specific fine-tuning.

Historical Importance

RAFT (2020) showed that learned flow does not need the pyramid of warping stages that had dominated it. Its authors carried the pattern beyond optical flow by changing the quantity being updated, to disparity, rigid 3D motion, or camera poses and depth, and other groups built on RAFT directly, as GMA did.

Variants

RAFT-Stereo (Lipson, Teed, and Deng, 2021) applies the design to rectified stereo, building a 3D correlation volume restricted to the same image row and adding GRUs at several resolutions. RAFT-3D (Teed and Deng, 2021) estimates scene flow from stereo or RGB-D frames by iteratively updating a dense field of rigid 3D motions. DROID-SLAM (Teed and Deng, 2021) builds visual SLAM on RAFT’s components, recurrently updating camera poses and depth through a dense bundle adjustment layer over many frames.

Within optical flow, GMA (Jiang et al., 2021) adds a transformer-based global motion aggregation module to RAFT to improve estimates in occluded regions. FlowFormer (Huang et al., 2022) encodes the all-pairs cost volume with transformer layers and decodes it with a recurrent transformer decoder that, as in RAFT, queries costs at positions given by the current flow estimate.

Implementations

The official PyTorch implementation, princeton-vl/RAFT, is released under the BSD 3-Clause license with pretrained models. torchvision includes the full and small models as raft_large and raft_small in torchvision.models.optical_flow, with pretrained weights. The following sketch follows the torchvision 0.29 API; image height and width must be divisible by 8:

import torch
from torchvision.models.optical_flow import raft_large, Raft_Large_Weights

weights = Raft_Large_Weights.DEFAULT
model = raft_large(weights=weights).eval()
preprocess = weights.transforms()  # converts to float and maps to [-1, 1]

# frame1, frame2: uint8 tensors of shape (B, 3, H, W)
img1, img2 = preprocess(frame1, frame2)
with torch.no_grad():
    flows = model(img1, img2, num_flow_updates=12)  # list, one flow per update
flow = flows[-1]  # (B, 2, H, W), in pixels