GMFlow: Learning Optical Flow via Global Matching

In plain words

Optical flow answers a simple question: for each pixel in one video frame, where did it move in the next frame? Most learning-based methods answer by looking only at a small neighbourhood around each pixel, then repeating the guess many times to reach far-away matches. That repetition makes them slow when the motion is large.

GMFlow instead treats the problem as matching. It turns each frame into a grid of feature vectors, lets the two grids look at each other so the features become more distinctive, and then compares every feature in one frame with every feature in the other. For each pixel, the network gets a probability over all possible destinations, and the flow is the expected destination minus the starting point. A final smoothing step lets pixels that could not be matched, such as ones hidden by an occluder, borrow the motion of similar-looking neighbours.

The point is that one pass can see the whole image, so large motions do not need many rounds of refinement. The card's results show this trade-off: strong accuracy on synthetic and Sintel-style benchmarks at lower runtime, but weaker results on the real-driving KITTI benchmark.

TL;DR

• GMFlow replaces the iterative regression of RAFT with one global correlation and softmax over all pixel pairs at 1/8 resolution §3.1 eq. 1-4.
• Features are enhanced by six Transformer blocks with self- and cross-attention, computed in 2×2 windows §3.2 fig. 2.
• A self-attention step propagates confident flow to occluded or out-of-boundary pixels §3.3 eq. 6-7.
• An optional refinement reuses the same weights at 1/4 resolution with local 9×9 matching §3.4.
• Author-reported: one refinement beats RAFT with 31 refinements on Things validation and Sintel train, and is faster on an A100; it loses on KITTI tab. 3 tab. 6.

Problem

Regression-based flow networks compute a cost volume over a fixed local window and regress flow from it with convolutions. The search range is therefore local, so large displacements need multi-stage or iterative refinement §1 §2. RAFT reaches a near-global search range by repeating a local update many times, and its inference time grows linearly with the number of refinements §1.

The method assumes two consecutive video frames as input §3.1. Its softmax matching assumes that a pixel's correspondence is visible in both frames, which the paper states does not hold for occluded and out-of-boundary pixels §3.3.

Method

Input: two frames. A weight-shared convolutional encoder produces features at 1/8 resolution, with 128 channels §3.1 §4 backbone.py. The model then runs the pipeline below (fig. 2 gives the overview in words: encoder, Transformer, correlation with softmax, and self-attention propagation).

1. Global correlation. Every feature in frame one is compared with every feature in frame two by a dot product, scaled by 1/√D §3.1 eq. 1.
2. Matching distribution. A softmax over the second frame's positions gives, for each pixel of frame one, a probability over all pixels of frame two §3.1 eq. 2.
3. Correspondence and flow. The expected coordinate under that distribution is the correspondence, and the flow is the correspondence minus the pixel's own coordinate §3.1 eq. 3-4. This step is differentiable and gives sub-pixel values, as the paper states §3.1.
4. Feature enhancement. Sine positional encodings in the style of DETR are added first §3.2 utils.py. Then six blocks run self-attention and cross-attention, with a feed-forward network, in both directions; the two frames' features are updated symmetrically §3.2 eq. 5. Attention is restricted to local windows: the feature map is cut into 2×2 windows, and every second layer shifts the windows by half a window to connect neighbouring windows §3.2 tab. 2b.
5. Flow propagation. The flow from step 3 is refined by a self-attention layer whose weights are the feature self-similarity; the flow itself is the value §3.3 eq. 6-7. The paper reports that this lifts unmatched-pixel accuracy tab. 2d.
6. Upsampling. Flow at 1/8 resolution is upsampled to full resolution with RAFT's convex upsampling, which uses a learned mask over 3×3 neighbours §4 gmflow.py.
7. Optional refinement. The 1/8 flow is upsampled to 1/4 resolution and frame two's 1/4 feature is warped with it. The same Transformer and flow-propagation weights are reused. Matching is local, a 9×9 window per pixel, and propagation uses a 3×3 window §3.4 tab. 2e. Features at the two scales are produced by one weight-shared convolution with strides 1 and 2 §3.4 backbone.py.

Bidirectional flow comes at no extra cost: transposing the global correlation matrix gives the backward flow without a second forward pass §4.1 fig. 3. The paper uses forward-backward consistency for occlusion detection §4.1.

Training uses supervision on every intermediate flow prediction with exponentially increasing weights §3.5 eq. 8.

# GMFlow forward pass at 1/8 resolution  [§3.1-3.3] [fig. 2]
# Input:  image_1, image_2 (B, 3, H, W), H and W divisible by 16
# Output: flow (B, 2, H, W)

def gmflow(image_1, image_2):
    # Weight-shared CNN features at 1/8 resolution  [§3.1]
    feature_1 = backbone(image_1)               # (B, D, H/8, W/8), D = 128
    feature_2 = backbone(image_2)               # (B, D, H/8, W/8)

    # Sine positional encodings, computed per window  [§3.2]
    feature_1, feature_2 = add_position(feature_1, feature_2)

    # Six blocks of self- and cross-attention  [§3.2] [eq. 5]
    feature_1, feature_2 = transformer(feature_1, feature_2)

    # Soft matching over all pixel pairs  [§3.1] [eq. 1-4]
    flow = global_matching(feature_1, feature_2)

    # Self-attention flow propagation to unmatched pixels  [§3.3] [eq. 6-7]
    flow = flow_propagation(feature_1, flow)

    # Learned convex upsampling to full resolution  [§4]
    return convex_upsample(flow, feature_1)

# Global matching: softmax over all pixel pairs  [§3.1] [eq. 1-4]
# Input:  feature_1, feature_2 (B, D, H, W)
# Output: flow (B, 2, H, W)

def global_matching(feature_1, feature_2):
    # Pairwise dot products scaled by 1/sqrt(D)  [eq. 1]
    queries = flatten(feature_1)                # (B, H*W, D)
    keys = flatten(feature_2)                   # (B, D, H*W)
    correlation = queries @ keys / sqrt(D)      # (B, H*W, H*W)

    # Distribution over second-frame positions  [eq. 2]
    match_prob = softmax(correlation, dim=-1)   # (B, H*W, H*W)

    # Expected coordinate of each match  [eq. 3]
    grid = pixel_grid(H, W)                     # (H*W, 2)
    correspondence = match_prob @ grid          # (B, H*W, 2)

    # Flow is correspondence minus start position  [eq. 4]
    flow = correspondence - grid                # (B, H*W, 2)
    return reshape(flow, (B, 2, H, W))

# Self-attention flow propagation  [§3.3] [eq. 6-7]
# Input:  feature (B, D, H, W), flow (B, 2, H, W)
# Output: propagated flow (B, 2, H, W)

def flow_propagation(feature, flow):
    # Flow is a constant for backprop  [code: gmflow/gmflow.py#L155]
    flow = detach(flow)

    # Key is a linear map of the query  [code: transformer.py#L363-L364]
    query = linear_q(flatten(feature))          # (B, H*W, D)
    key = linear_k(query)                       # (B, H*W, D)

    # Self-similarity weights, scaled by 1/sqrt(D)  [eq. 6]
    scores = query @ transpose(key) / sqrt(D)   # (B, H*W, H*W)
    prob = softmax(scores, dim=-1)

    # The flow is the value; no value projection  [eq. 6-7]
    propagated = prob @ flatten(flow)           # (B, H*W, 2)
    return reshape(propagated, (B, 2, H, W))

# One Transformer block: self-attention, then cross-attention with FFN
# [§3.2] [eq. 5] [code: gmflow/transformer.py#L188-L241]
# Input:  source, target (B, L, D); shifted is True on odd layers
# Output: updated source (B, L, D)

def transformer_block(source, target, shifted):
    # Self-attention, no feed-forward network  [code: transformer.py#L201-L207]
    source = attention_layer(source, source, shifted, use_ffn=False)

    # Cross-attention, queries from source, keys and values from target
    # [code: transformer.py#L209-L214]
    source = attention_layer(source, target, shifted, use_ffn=True)
    return source

# Single-head attention inside local windows, with residual
# [code: gmflow/transformer.py#L147-L185]
# Input:  source, target (B, L, D)
# Output: source + message (B, L, D)

def attention_layer(source, target, shifted, use_ffn):
    # Separate query, key and value maps  [code: transformer.py#L128-L130]
    query = linear_q(source)
    key = linear_k(target)
    value = linear_v(target)

    # Attention inside each window; shift windows by half on odd layers
    # [§3.2] [code: transformer.py#L73-L105]
    message = window_attention(query, key, value, shifted)
    message = layer_norm(linear_merge(message))

    # FFN reads the residual stream concatenated with the message
    # [code: transformer.py#L137-L145] [code: transformer.py#L181-L183]
    if use_ffn:
        message = layer_norm(mlp(concat(source, message)))
    return source + message

# Sequence loss over all flow predictions  [§3.5] [eq. 8]
# Input:  flow_preds (list of (B, 2, H, W)), flow_gt (B, 2, H, W),
#         valid (B, H, W)
# Output: scalar loss

def sequence_loss(flow_preds, flow_gt, valid, gamma=0.9):
    # Drop unlabelled pixels and motions of 400 px or more
    # [code: loss.py#L13-L14] [code: main.py#L33-L34]
    magnitude = norm(flow_gt, dim=1)            # (B, H, W)
    valid = valid & (magnitude < 400)

    loss = 0
    for i, flow in enumerate(flow_preds):
        # Later predictions weigh more; the last one has weight 1  [eq. 8]
        # [code: loss.py#L17]
        weight = gamma ** (len(flow_preds) - i - 1)

        # Mean over all elements, invalid ones counted as zero
        # [code: loss.py#L19-L21]
        loss += weight * mean(valid * abs(flow - flow_gt))
    return loss

Hyperparameters

Values from the paper are cited to the paper. Values marked "code" are in the official repository and are not stated in the paper text.

Results

Author-reported. Setup: models trained on FlyingChairs and FlyingThings3D, evaluated on Things validation, Sintel train and KITTI train §4. The inference time is at Sintel resolution (436×1024) on a V100, with the A100 time in parentheses tab. 3.

The Sintel test rows are after the mixed fine-tuning in §4.4, and the paper states they use only two frames at inference, unlike the RAFT† and GMA† rows marked with a dagger tab. 5. Table 5 also lists the matched and unmatched breakdown; GMFlow has the lowest unmatched EPE among the listed methods tab. 5.

Ablations, author-reported, trained on Chairs and Things with 200K fine-tuning steps unless stated §4 tab. 2a-f:

• Transformer without cross-attention raises Things val EPE from 6.67 to 10.84 tab. 2a.
• Flow propagation lowers unmatched-pixel EPE on Sintel clean from 15.54 to 10.39 tab. 2d.
• Local matching with a 3×3 window gives Things val EPE 31.78 vs 6.67 for global matching tab. 2c.
• Sharing both the Transformer and the features between scales gives Things clean EPE 3.24 with 4.7M parameters, vs 3.41 with 8.0M when neither is shared tab. 2e.
• Training longer helps: Things val EPE falls from 3.24 (200K) to 2.80 (800K) tab. 2f.

Paper vs code

Compared at commit b5123431164d01ec14526a1c3d22218aecb62024 of the official repository, which the paper's abstract links and whose owner is one of the authors.

• Loss normalisation and exclusions. The paper's loss is a weighted sum over predictions of the L1 norm of the error, with no normalisation or exclusion stated eq. 8. The code takes the mean over every element of the masked, component-wise absolute error, so the loss is also divided by the number of flow components and invalid pixels are counted as zero. It also excludes pixels with ground-truth magnitude of 400 or more loss.py main.py. The paper does not mention either detail.
• Intermediate upsampling. The paper says all flow predictions are supervised §3.5 and gives no upsampling rule for them. The code upsamples intermediate predictions bilinearly, scaled by the upsampling factor, and applies the learned convex upsampling only to the final prediction gmflow.py.
• Flow propagation. The paper's equation for propagation uses the Transformer features directly eq. 6. The code adds a learned query projection and a second linear map on the query before computing keys transformer.py. The code comment says the "correct" form would differ, and that the existing form was kept so the trained models still work transformer.py. Follow the code if you need to load released weights.
• Gradient stopping. The code detaches the flow before propagation and before warping it into the refinement stage gmflow.py gmflow.py. The paper does not say whether gradients flow through the flow input.
• Positional encoding. The paper describes a fixed 2D sine encoding over the feature map §3.2. When attention is windowed (more than one split, as in the paper's 2×2 setting), the code computes the encoding separately inside each window utils.py, so positions are window-relative.
• Feed-forward network. The paper describes self-attention, cross-attention and FFN §3.2 without detail. The code applies the FFN only after cross-attention, on the concatenation of the residual stream and the attention output (input width 2D), with a LayerNorm on both the attention output and the FFN output transformer.py transformer.py.
• Sintel fine-tuning mixture. The paper says "mixed datasets" §4.4. The code samples the datasets with repetition factors: 100 × Sintel clean, 100 × Sintel final, 200 × KITTI, 5 × HD1K and 1 × FlyingThings3D datasets.py.
• Backbone. The paper says the encoder is identical to RAFT's except for 128 output channels §4. RAFT is not in this repository, so this is not verified. The code's encoder is a 7×7 stride-2 convolution followed by residual blocks with instance normalisation backbone.py.
• Attention heads. The paper does not say how many heads are used; the code is single-head only transformer.py.
• Input size. Training crops and padding are set in the scripts, not in the paper. Padding is 16 for the single-scale model and 32 with refinement train_gmflow.sh train_gmflow_with_refine.sh.

No other discrepancies found in the global matching formula, the 1/√D scaling, the 6-block depth, the 2×2 window split, the 9×9 refinement window, or the training schedules quoted in §4.

Pitfalls

• Quadratic memory. Global matching builds an (H/8·W/8) × (H/8·W/8) correlation matrix per image §3.1 eq. 1. The paper's appendix describes a sequential implementation that computes it in K×K blocks at little speed cost: 52.6 ms with no split and 59.5 ms with 8×8 splits at 448×1024 input app. B tab. B. Without that, plan memory for the full matrix.
• Divisibility. Inputs must be divisible by the padding factor main.py, and window splitting asserts that height and width are divisible by the number of splits utils.py. The paper does not state these constraints.
• Mask constant. The shifted-window mask adds −100, not −∞, to blocked positions transformer.py. Blocked weights are small but not zero.
• Out-of-image matches. In local matching, positions outside the image are masked with −1e9 before the softmax matching.py. This applies to the refinement stage only.
• Occlusions. The paper states the method is weaker in occluded regions, and that the KITTI test gap comes mainly from occluded pixels §4.4 tab. 6. Its broader-impact note warns about unreliable results in occluded areas for safety-critical use §5.
• Synthetic-to-real gap. Trained on synthetic data, GMFlow trails RAFT on KITTI generalisation, and the gap narrows only when Virtual KITTI 2 is added §4.3 tab. 4.
• Bidirectional flow. The paper says training predicts unidirectional flow only; the reverse flow is obtained at inference by transposing the correlation matrix, without a second forward pass §4.1. The code's bidirectional path is enabled only through an inference flag evaluate.py.
• Refinement sharing. The refinement stage reuses the global-stage Transformer and propagation weights §3.4. Changing one stage's depth changes both.
• Reported time. The paper's speed numbers are measured on V100 and A100 GPUs at one resolution tab. 3. The paper argues the advantage grows on newer hardware §4.3 — an author claim, not a benchmark over many devices.

Lineage

• Builds on: RAFT, whose iterative refinement this card's method replaces with a single global matching pass §1 §4.3. Shifted-window attention is taken from Swin Transformer, with the windows split into a fixed count rather than a fixed size §3.2. Sine positional encodings follow DETR §3.2. Convex upsampling follows RAFT §4.
• Compared against: RAFT tab. 3 tab. 5; FlowNet2, PWC-Net+, HD³, VCN, DICL and GMA tab. 5. The paper also discusses sparse matching methods SuperGlue and LoFTR, and STTR for stereo, as related work §2.
• Extended by: the official repository's README names a follow-up that extends the model to stereo and depth estimation README.md. It is not part of this card's content.