Deep Optical Flow: RAFT & Beyond

"Horn and Schunck handed me one equation and a prayer for smoothness. RAFT handed me a lookup table of every pixel against every other pixel and said, just keep guessing, but guess better each time. Reader, I have never been more confident about where things went."

An Optical Flow Field That Finally Stopped Smearing at the Edges

Every method so far in this chapter has reduced a clip to a single label or a handful of boxes. Now we ask for everything: a motion vector at every one of the hundreds of thousands of pixels in a frame, telling you exactly where each one went. That dense field is optical flow, the displacement of every pixel in frame one to its corresponding location in frame two. You met it classically in Chapter 15, and it returned as the input to the two-stream network of Section 26.2. This section rebuilds it with deep learning. The structure of RAFT (Recurrent All-Pairs Field Transforms) is worth studying in detail not only because it is the dominant flow method but because its three-part design is a template you will see again in stereo depth in Chapter 27. Three words capture the entire recipe and are worth committing to memory: encode, correlate, iterate. Encode both frames into learned features, correlate every pixel against every other to build a similarity table, then iterate a shared-weight update that keeps refining the guess. The three subsections that follow are exactly those three words in order.

1. The Flow Problem and the Classical Baseline Beginner

Optical flow rests on the brightness constancy assumption: a point's intensity does not change as it moves between frames. Writing $I(x, y, t)$ for the image intensity and $(u, v)$ for the per-pixel displacement, a first-order Taylor expansion gives the classical optical flow constraint equation

where $I_x, I_y$ are the spatial gradients and $I_t$ is the temporal gradient. This is one equation in two unknowns per pixel, the famous aperture problem: a single local window can only determine the motion component along the gradient, not the full vector. Classical methods resolve the ambiguity with extra assumptions: Lucas-Kanade assumes flow is constant in a small window (giving a solvable least-squares system), and Horn-Schunck adds a global smoothness penalty. Both break down on large displacements, where the Taylor expansion is invalid, and on textureless regions, where there is no gradient to constrain anything. RAFT keeps the brightness-constancy intuition but replaces the local linearization with a global correlation search.

2. Feature Encoding and the All-Pairs Correlation Volume Intermediate

Here is the move that broke the field open, and at first glance it looks reckless: instead of guessing where each pixel went and checking nearby, RAFT compares every pixel in frame one against every pixel in frame two, all at once, before it commits to a single displacement. That is the "all-pairs" in the name, and it is why RAFT stops smearing at large motions where the classical methods give up. To get there it first runs both frames through a shared convolutional feature encoder that maps each frame to a dense feature map at one-eighth resolution, so each spatial location carries a learned descriptor rather than a raw pixel. This is the learned replacement for the hand-crafted descriptors of Chapter 10; matching learned features is far more robust than matching raw intensities, which is what brightness constancy assumed.

The heart of RAFT is the correlation volume. For every feature in frame one, RAFT computes the dot-product similarity to every feature in frame two, producing a 4D volume $C(i, j, k, l)$ that holds the similarity of pixel $(i, j)$ in frame one to pixel $(k, l)$ in frame two. This is the all-pairs comparison that gives RAFT its name, and it is the global, large-displacement-capable analogue of the local block matching used in classical stereo and flow.

One refinement turns this volume from expensive into practical. RAFT builds a small pyramid of the volume by pooling the second-frame dimensions, so the lookup can read both fine and coarse correspondence. The coarser levels are what let RAFT handle large motion cheaply: a displacement of many pixels at full resolution becomes a displacement of only a few pixels once the second frame is pooled down, so a small fixed-size lookup window still reaches it. Figure 26.4.1 shows the pipeline.

The code below builds the all-pairs correlation volume from two feature maps. It is a single batched matrix multiplication, the same dot-product-as-similarity operation that powered attention in Chapter 22, applied here between two images rather than within one sequence.

import torch

def correlation_volume(fmap1, fmap2):
    """All-pairs correlation between two feature maps.

    fmap1, fmap2: (N, D, H, W) learned features.
    Returns (N, H, W, H, W): similarity of every pixel in 1 to every pixel in 2.
    """
    N, D, H, W = fmap1.shape
    f1 = fmap1.view(N, D, H * W)                 # (N, D, HW)
    f2 = fmap2.view(N, D, H * W)                 # (N, D, HW)
    corr = torch.matmul(f1.transpose(1, 2), f2)  # (N, HW, HW) dot products
    corr = corr / (D ** 0.5)                      # scale, as in attention
    return corr.view(N, H, W, H, W)

f1 = torch.randn(1, 64, 32, 32)                   # 1/8-resolution features
f2 = torch.randn(1, 64, 32, 32)
vol = correlation_volume(f1, f2)
print("correlation volume:", vol.shape)           # correlation volume: torch.Size([1, 32, 32, 32, 32])

3. The Recurrent Update Operator Intermediate

RAFT's third component is what makes it accurate. Rather than predicting flow in one shot, it starts from zero flow and applies a recurrent update operator, a convolutional gated recurrent unit (GRU), that refines the estimate over a fixed number of iterations. A GRU is a small network with a persistent hidden state that it carries from one step to the next: at each step it takes the current input, mixes it with the state through learned gates that decide what to keep and what to overwrite, and emits an updated state. Here that state is a feature map, the same weights act at every iteration, and the running state is what accumulates the flow estimate across steps.

At each iteration the operator does three things. First it looks up the correlation volume at the pixels indicated by the current flow estimate, reading how well the current correspondence matches. Then it combines that lookup with context features from the first frame. Finally it predicts a residual update $\Delta f$ that is added to the current flow. Crucially, all iterations share weights, so the network learns a single update rule applied repeatedly, mimicking the iterations of a classical optimizer but with a learned step. The training loss supervises every intermediate flow, with later iterations weighted more, which teaches the operator to converge. The illustration below captures this encode, correlate, iterate loop as a robot consulting its lookup table and redrawing its arrows.

The simplified update step below shows the structure: lookup, combine, predict residual, add. A real RAFT uses a more elaborate lookup with a local neighborhood and a multi-level correlation pyramid, but the recurrence is exactly this.

import torch
import torch.nn as nn

class FlowUpdate(nn.Module):
    """One shared-weight RAFT-style update: predict a residual flow and add it."""
    def __init__(self, hidden=128, corr_dim=81, ctx_dim=128):
        super().__init__()
        # GRU-like gate over hidden state h, fed correlation lookup + context + flow
        self.enc = nn.Sequential(nn.Conv2d(corr_dim + 2, 64, 3, padding=1), nn.ReLU())
        self.gru = nn.GRUCell(64 + ctx_dim, hidden)
        self.flow_head = nn.Conv2d(hidden, 2, 3, padding=1)   # predicts (du, dv)

    def forward(self, h, flow, corr_lookup, context):
        N, _, H, W = flow.shape
        inp = torch.cat([corr_lookup, flow], dim=1)            # what the lookup said
        inp = self.enc(inp).flatten(2).transpose(1, 2)         # (N, HW, 64)
        ctx = context.flatten(2).transpose(1, 2)               # (N, HW, ctx_dim)
        h_flat = h.flatten(2).transpose(1, 2).reshape(-1, h.shape[1])
        h_new = self.gru(torch.cat([inp, ctx], -1).reshape(-1, inp.shape[-1] + ctx.shape[-1]),
                         h_flat).view(N, H, W, -1).permute(0, 3, 1, 2)
        delta = self.flow_head(h_new)                          # residual update
        return h_new, flow + delta                             # refined flow

# One iteration on toy tensors (corr lookup window of 9x9 = 81 channels):
upd = FlowUpdate()
h = torch.zeros(1, 128, 32, 32)
flow = torch.zeros(1, 2, 32, 32)
corr_lookup = torch.randn(1, 81, 32, 32)
context = torch.randn(1, 128, 32, 32)
h, flow = upd(h, flow, corr_lookup, context)
print("refined flow:", flow.shape)                             # refined flow: torch.Size([1, 2, 32, 32])

4. Using Flow, and the Transformer Successors Intermediate

In practice you rarely implement RAFT; you load it. TorchVision ships pretrained RAFT weights, and computing dense flow between two frames is a few lines. The output is a $(2, H, W)$ field of horizontal and vertical displacements that you can feed to the temporal stream of Section 26.2, use to warp one frame onto another, or visualize as a colour-coded motion image.

# Estimate dense optical flow between two frames with a pretrained RAFT.
# The model emits one flow field per refinement iteration; we keep the last.
import torch
from torchvision.models.optical_flow import raft_large, Raft_Large_Weights

weights = Raft_Large_Weights.DEFAULT
model = raft_large(weights=weights).eval()
transforms = weights.transforms()

img1, img2 = transforms(frame1, frame2)              # resize + normalize a frame pair
with torch.no_grad():
    flow_predictions = model(img1, img2)             # list of flow fields, one per iteration
flow = flow_predictions[-1]                          # the final, most refined estimate
print(flow.shape)                                    # torch.Size([1, 2, H, W])  (du, dv)