FlowFormer: A Transformer Architecture for Optical Flow

In plain words

Optical flow answers a simple question: for every pixel in one video frame, where did it move to in the next frame? Older methods first build a score table that compares every pixel of the first image with every pixel of the second. The score table is large and hard to reason about, so most methods only look at small neighbourhoods of it.

FlowFormer treats that score table as something a transformer can read. Each pixel's row of scores is squeezed into a handful of learned summary vectors. These summaries then talk to each other, both within the same pixel and across neighbouring pixels, so the model can use the wider context. A decoder then repeatedly asks "where does this pixel's match most likely lie now?", reads the table there, and nudges its motion guess.

The point is accuracy on standard optical-flow benchmarks, where the authors report lower error than the previous best methods. The architecture is heavier than the CNN baselines it is compared with, and the paper's numbers depend on settings described on this card.

TL;DR

• Optical flow needs a dense per-pixel motion field; the paper argues that a transformer must work on a compressed cost volume, because raw all-pairs attention would need thousands of tokens §1.
• The method compresses each source pixel's cost map into K = 8 latent tokens of width D = 128, refines them with three alternate-group transformer layers, and decodes flow with 12 recurrent cross-attention steps (count from code) §3.2 §4.1 default.py.
• Author-reported Sintel test AEPE: 1.16 clean and 2.09 final, the best of the compared methods in the table; the KITTI-15 test F1-all of 4.68 is second to S-Flow tab. 1.
• Several training details (loss weights, optimizer, number of recurrent steps, positional-encoding form) are absent from the paper and are taken from the official code; see Paper vs code.

Problem

• Optical flow is a per-pixel displacement field from a source image to a target image §1.
• Classical formulations assume the appearance of corresponding points is unchanged §1.
• Recent CNN methods such as RAFT build a 4D cost volume and only retrieve costs inside local windows around the current estimate §1 §4.3.
• Applying a transformer directly to the 4D cost volume needs thousands of tokens and is computationally unaffordable §1.
• Perceiver IO applies a transformer to raw pixels for flow. The paper says it needs many parameters and about 80 times more training examples, because it ignores the cost-volume prior §1.
• The paper's question is whether a transformer can use the cost volume efficiently, combining both §1.

Method

FlowFormer has three parts, as in fig. 1: a 4D cost volume built from image features §3.1, a cost-volume encoder that turns it into a cost memory §3.2, and a recurrent cost-memory decoder that predicts flow from that memory and a context feature map §3.3.

Cost volume §3.1. A backbone extracts a feature map of 1/8 the input size, with D_f = 256 channels §4.1. The dot product of every source-pixel feature with every target-pixel feature gives an H×W×H×W volume. The paper states no normalisation of these dot products §3.1; the code applies none either encoder.py.

Cost volume encoder §3.2. For each source pixel, its H×W cost map is treated as an image and processed in three steps:

1. Patchify: zero-pad the right and bottom to multiples of 8, then apply three stride-2 convolutions with D_p/4, D_p/2 and D_p channels (D_p = 64), with ReLU after the first two §3.2 §4.1. Each output cell covers an 8×8 patch of the cost map.
2. Latent summarization: K shared, learned codewords C attend over the patch features, concatenated with a positional embedding, to produce K tokens of width D per source pixel §3.2 eq. 1.
3. Alternate-group transformer (AGT) layers, stacked three times §3.2 §4.1 fig. 2:
   • Intra-cost-map attention: self-attention and an FFN over the K tokens of one source pixel eq. 2.
   • Inter-cost-map attention: for each of the K token indices, spatially separable self-attention (SS-SelfAttention, from Twins) over all H·W source pixels eq. 3 §3.2. Context features from the source image are concatenated to the queries and keys of this step §3.2.

The result is a cost memory of shape H×W×K×D §3.2.

Cost memory decoder §3.3. Flow is estimated at 1/8 resolution and upsampled with a learned convex upsampler §3.3. Let p = x + f(x) be the matched position of pixel x §3.3.

1. Local cost window: a 9×9 window of the pixel's own cost map, centred at p, is cropped §3.3.
2. Dynamic positional cost query: the window is encoded with FFNs and combined with a positional embedding of p to form a query eq. 4. Because p changes each iteration, the query changes too §3.3.
3. Cross-attention: the query attends to the pixel's K cost tokens from the memory eq. 4. Keys and values are computed once and reused in later iterations §3.3.
4. Recurrent update: a ConvGRU receives the cross-attention output with the window, the context feature t_x and the current flow, and outputs a residual Δf eq. 5 §3.3. The flow is updated as f ← f + Δf §3.3.

The flow from every iteration is upsampled and supervised, with later iterations weighted more heavily §3.3. The paper gives no formula for the loss.

Equations (transcribed from the paper; each is numbered there):

\mathbf{K}_x = \mathrm{Conv}_{1\times1}\big(\mathrm{Concat}(\mathbf{F}_x, \mathrm{PE})\big) \qquad [\text{eq. 1}]

\mathbf{V}_x = \mathrm{Conv}_{1\times1}\big(\mathrm{Concat}(\mathbf{F}_x, \mathrm{PE})\big) \qquad [\text{eq. 1}]

\mathbf{T}_x = \mathrm{Attention}(\mathbf{C}, \mathbf{K}_x, \mathbf{V}_x) \qquad [\text{eq. 1}]

\mathbf{T}_x = \mathrm{FFN}\big(\mathrm{Self\text{-}Attention}(\mathbf{T}_x(1), \dots, \mathbf{T}_x(K))\big) \quad \text{for all } \mathbf{x} \in \mathbf{I}_s \qquad [\text{eq. 2}]

\mathbf{T}_i = \mathrm{FFN}\big(\mathrm{SS\text{-}Self\text{-}Attention}(\mathbf{T}_i)\big), \quad i = 1, \dots, K \qquad [\text{eq. 3}]

\mathbf{Q}_x = \mathrm{FFN}\big(\mathrm{FFN}(\mathbf{q}_x) + \mathrm{PE}(\mathbf{p})\big), \quad \mathbf{K}_x = \mathrm{FFN}(\mathbf{T}_x), \quad \mathbf{V}_x = \mathrm{FFN}(\mathbf{T}_x) \qquad [\text{eq. 4}]

\mathbf{c}_x = \mathrm{Attention}(\mathbf{Q}_x, \mathbf{K}_x, \mathbf{V}_x) \qquad [\text{eq. 4}]

\Delta\mathbf{f}(\mathbf{x}) = \mathrm{ConvGRU}\big(\mathrm{Concat}(\mathbf{c}_x, \mathbf{q}_x), \mathbf{t}_x, \mathbf{f}(\mathbf{x})\big) \qquad [\text{eq. 5}]

Unnumbered in the paper: the matched position p = x + f(x) and the window q_x = Crop_{9×9}(M_x, p) §3.3; the update f(x) ← f(x) + Δf(x) §3.3.

Tiling at test time (appendix) blends overlapping tile predictions with weights from a Gaussian over normalised distance to the tile centre app. §2:

\mathbf{f} = \frac{\sum_i w_i \mathbf{f}_i}{\sum_i w_i} \qquad [\text{app. §2, eq. 1}]

d_{u,v} = \left\lVert \left(\tfrac{u}{H_{\mathrm{train}}} - 0.5,\ \tfrac{v}{W_{\mathrm{train}}} - 0.5\right) \right\rVert_2, \qquad w_{u,v} = g(d_{u,v};\ \mu = 0,\ \sigma = 0.05) \qquad [\text{app. §2, eq. 2}]

Appendix citations [app. §N] refer to the numbered sections of the supplementary file in the paper's source archive.

Pseudocode (paper-level structure; each block cites the step it implements):

# FlowFormer forward pass  [§3]  [fig. 1]
# Input:  image1, image2 (B, 3, H_I, W_I)
# Output: one upsampled flow per recurrent step (B, 2, H_I, W_I)

def flowformer(image1, image2):
    # Source-image context features t_x  [§3.3]
    context = context_encoder(image1)            # (B, 256, H, W)
    # Cost memory with H x W x K tokens of width D  [§3.2] [eq. 1-3]
    memory = cost_volume_encoder(image1, image2, context)
    # Recurrent decoder starts from zero flow  [§3.3]
    return cost_memory_decoder(memory, context)

# Cost volume encoder  [§3.2]
# Input:  feature maps (B, 256, H, W) for both images
# Output: cost memory (B, H, W, K, D), with K = 8 and D = 128  [§4.1]

def cost_volume_encoder(feats_src, feats_tgt, context):
    # Dot product of every source pixel with every target pixel  [§3.2]
    cost = dot_products(feats_src, feats_tgt)    # (B, H, W, H, W)
    # Patchify and summarize each source pixel's cost map  [§3.2] [eq. 1]
    tokens = summarize_all_cost_maps(cost)       # (B, H, W, K, D)
    # Three alternate-group transformer layers  [§3.2] [§4.1]
    for layer in range(3):
        tokens = agt_layer(tokens, context)
    return tokens

# Patchify one cost map  [§3.2]
# Input:  cost map (H, W); output: patches (ceil(H/8), ceil(W/8), D_p)

def patchify(cost_map):
    # Zero-pad right and bottom to multiples of 8  [§3.2]
    padded = pad_to_multiple_of_8(cost_map)
    # Three stride-2 convolutions, ReLU after the first two  [§3.2]
    # Channels D_p/4, D_p/2, D_p with D_p = 64  [§4.1]
    patches = relu(conv_stride2(padded, out_channels=D_p // 4))
    patches = relu(conv_stride2(patches, out_channels=D_p // 2))
    return conv_stride2(patches, out_channels=D_p)

# Latent summarization of one cost map  [§3.2] [eq. 1]
# K shared learned codewords attend over the patch features

def summarize(patches, pos_enc, codewords):
    # Keys and values: 1x1 projections of patches with positions  [eq. 1]
    keys = conv1x1(concat(patches, pos_enc))
    values = conv1x1(concat(patches, pos_enc))
    # Dot-product attention; codewords are the queries  [eq. 1]
    return attention(codewords, keys, values)    # (K, D)

# Alternate-group transformer layer  [§3.2] [fig. 2]
# Input and output: tokens (B, H, W, K, D)

def agt_layer(tokens, context):
    # Intra-cost-map: self-attention over the K tokens of each pixel  [eq. 2]
    tokens = self_attention_and_ffn_over_k(tokens)
    # Inter-cost-map: for each codeword index, spatially separable attention
    # over all H*W pixels; context joins queries and keys  [eq. 3] [§3.2]
    tokens = ss_self_attention_and_ffn_over_hw(tokens, context)
    return tokens

# Cost memory decoder  [§3.3] [eq. 4] [eq. 5]
# Input:  memory (B, H, W, K, D), context (B, 256, H, W)
# Output: one upsampled flow per recurrent step

def cost_memory_decoder(memory, context):
    # Keys and values from the memory, computed once and reused  [§3.3]
    keys, values = project_memory(memory)
    # Zero initial flow  [code: LatentCostFormer/decoder.py#L211-L215]
    flow = zeros(B, 2, H, W)
    predictions = []
    # Number of steps is not stated in the paper; the code uses 12
    for step in range(num_steps):
        # Matched position of every pixel  [§3.3]
        position = pixel_grid + flow
        # 9x9 window of each cost map around the match  [§3.3]
        window = sample_window(cost_maps, position, radius=4)
        # Dynamic positional cost query  [eq. 4]
        query = ffn(ffn(window) + positional_encoding(position))
        # Cross-attention from the query to the cost memory  [eq. 4]
        cost_feat = attention(query, keys, values)
        # ConvGRU predicts the flow residual  [eq. 5]
        delta = conv_gru(concat(cost_feat, window), context, flow)
        flow = flow + delta
        # Convex upsampling to input resolution  [§3.3]
        predictions.append(convex_upsample(flow))
    return predictions

Parameter count. Table 3 reports 18.2M parameters for the full model; the paper's own comparison attributes the gain to architecture rather than size §4.3 tab. 3. No runtime or memory measurement appears in the main text §4.

Hyperparameters

The paper's own values are cited to the paper; the code column gives the official default where the paper is silent. Values marked "code" were not stated in the paper.

Results

Author-reported. The paper evaluates on Sintel (clean and final passes) and KITTI-2015 §4 and reports three protocols §4.1:

• Generalization: train on FlyingChairs and FlyingThings (C+T), then test on Sintel and KITTI-2015 training splits §4.1.
• Sintel and KITTI benchmarks: after C+T, finetune on FlyingThings, HD1K, Sintel and KITTI training data §4.1.
• Ablation: RAFT-style schedule, 100k steps on chairs and 60k on things §4.3.

Headline numbers:

The paper states the following relative reductions: 16.5% and 15.5% lower error than GMA* on Sintel test clean and final §4.1; 17.2% and 27.5% against GMA without warm start §4.1; 22.3% and 12.4% on Sintel train clean and final, and 13.9% on KITTI-15 F1-all, for the C+T generalization setting §4.1.

Table 1 (complete, from tab. 1). Columns: Sintel train clean / final; KITTI-15 train F1-epe / F1-all; Sintel test clean / final; KITTI-15 test F1-all. "A" is the AutoFlow dataset. Entries marked † use the tile technique app. §2. The caption does not define the parentheses. They appear on training-split columns of models finetuned on those splits.

Table 2 (complete, from tab. 2). Ablation, trained with the RAFT-style schedule §4.3. Each row changes one component from the RAFT baseline.

Table 3 (complete, from tab. 3). FlowFormer against GMA variants with matched parameter counts. Trained with 120k chairs and 120k things steps §4.3. The percentages in the paper's table are the extra error of each model relative to "Ours".

Appendix AGT ablation (from app. §1): intra-cost-map and inter-cost-map attention replaced by MLP-Mixer and ConvNeXt blocks.

The authors keep the full transformer design because the replacements help Sintel final slightly but hurt KITTI app. §1.

Paper vs code

Official repository: drinkingcoder/FlowFormer-Official at commit 6ba7ea82b45394d3a7a25808399695427c9febd8, Apache-2.0 LICENSE. The owner is the first author's GitHub account; the repository README links the first author's project page and lists the paper README.md. Search: the paper's text and bibliography link no code. A GitHub repository search for "flowformer" returned this repository, whose owner is the first author, and the README confirms it is the release for this paper.

Consistent with the paper: pretrained Twins first two stages encoders.py; D_f = 256, D_p = 64, K = 8, D = 128, three AGT layers default.py; batch sizes and stage lengths default.py things.py kitti.py; learning rates 2.5e-4 and 1.25e-4 default.py things.py; the one-cycle schedule [code: core/optimizer/init.py#L52-L60]; the tile technique with sigma 0.05 evaluate_FlowFormer_tile.py.

Differences, each with both citations:

1. Chairs crop. Paper: 368×498 app. §3. Code: 368×496 default.py. Reproducing the released code exactly needs the code value; the paper value is the one stated in the text.
2. Patch embedding. Paper: three stride-2 convolutions, ReLU after the first two, with positions concatenated before the 1×1 key and value projections §3.2 eq. 1. Code: kernel 6, stride 2, padding 2, ReLU after the first two; the positional channels are concatenated after the convolutions, followed by two 1×1 convolutions with ReLU and a LayerNorm encoder.py.
3. Latent summarization. Paper: the codewords are updated by attention alone eq. 1. Code: the same cross-attention, followed by a residual connection, a LayerNorm and an FFN encoder.py encoder.py.
4. Encoder residual. Paper: not mentioned. Code adds the input latent tokens back to the output of the three AGT layers when cost_encoder_res is true, which the released configs set encoder.py default.py.
5. Inter-cost attention. Paper: spatially separable self-attention (SS-SelfAttention) eq. 3. Code: a Twins window attention (7×7 local window) followed by a global sub-sampled attention with stride 4; both use the context features in queries and keys, and each is followed by an MLP encoder.py twins.py twins.py twins.py. The paper does not state the window size or stride.
6. Decoder query. Paper: Q = FFN(FFN(q) + PE(p)) eq. 4. Code: a two-layer 1×1 convolutional encoder of the window, plus the positional embedding, then a single linear projection; the attention output is concatenated with the query, projected, and passed through a residual FFN decoder.py decoder.py.
7. Decoder keys and values. Paper: K = FFN(T), V = FFN(T) eq. 4. Code: single linear projections, computed once on the first iteration decoder.py decoder.py.
8. Recurrent unit. Paper: a ConvGRU following GMA-RAFT §3.3. Code: the GMA update block, with separable (1×5 then 5×1) ConvGRU, a motion encoder and a GMA-style global aggregator over the context features gru.py gru.py decoder.py. The mask output is scaled by 0.25 gru.py.
9. Gradient stop. Code detaches the matched positions at the start of every iteration, so no gradient flows through the sampling positions decoder.py. Paper: not mentioned.
10. Loss. Paper: the flow of every iteration is supervised, with increasing weights, and no formula is given §3.3. Code: L1 error on each prediction, weighted by gamma^(N-1-i) with gamma 0.8 or 0.85, pixels with ground-truth magnitude at or above 400 excluded loss.py sintel.py.
11. Optimizer. Paper: one-cycle schedule only §4.1. Code: AdamW with the weight decay values in the Hyperparameters table [code: core/optimizer/init.py#L15-L35] default.py. Gradient clipping at 1.0 train_FlowFormer.py.
12. Image noise. Code adds Gaussian noise with sigma drawn uniformly from [0, 5] to both images in every training step (add_noise is true in all configs) train_FlowFormer.py default.py. Paper: not mentioned.
13. Sintel-stage mixture. Paper: "FlyingThings, Sintel, KITTI-2015, and HD1K" §4.1. Code: FlyingThings clean pass only, with weights 100, 100, 200 and 5 for the other sets datasets.py.
14. Positional encoding. Paper: "following COTR" with no formula §3.2. Code: sine and cosine of linearly spaced frequencies with a hard-coded 3.14 and a 1/200 scale attention.py.
15. Attention heads. Paper: not stated. Code: 8 heads by default in the encoder and decoder encoder.py decoder.py.
16. Input range. Code maps pixels to [-1, 1] before the encoders transformer.py. Paper: not stated.

No discrepancies were found for the stage schedule, the token counts, the learning rates or the tile method.

Pitfalls

• Unstated loss details. The paper gives no loss formula, and no weights beyond "increasing". The official code uses L1 with geometric weights (gamma 0.8 or 0.85) loss.py. An implementation written from the paper alone will not match its training dynamics.
• Unstated optimizer and steps. The optimizer, weight decay, gradient clipping and the number of recurrent steps are all missing from the paper. The code uses 12 recurrent steps in training and test default.py.
• Token naming. Section 3 uses K for the number of tokens and D for their width. Section 4.1 writes "N=8 cost tokens of K=128 dimensions" §4.1. The configs and the ablation agree on 8 tokens of width 128 default.py tab. 2.
• Inconsistent "Ours" numbers. Table 1 reports the full model at 1.01 and 2.40 on Sintel train clean and final (C+T). Table 3 reports 1.00 and 2.45 for the same model, and both are described as 120k chairs and 120k things steps §4.1 §4.3. The paper does not explain the difference. Table 2 uses the RAFT-style 100k and 60k schedule §4.3, so its AGT x3 row (1.10 and 2.57) is not directly comparable to either.
• Tiled evaluation. Entries marked † in Table 1 use tiling with Gaussian blending, not a single forward pass app. §2. The code tiles with a 20 px minimum overlap and a 432×960 tile evaluate_FlowFormer_tile.py evaluate_FlowFormer_tile.py. Those numbers cannot be compared one-to-one with single-pass results, and the paper does not report single-pass numbers for them.
• Input size sensitivity. The paper states that positional encodings are sensitive to image size, so test pairs are cropped or tiled to the training size §4.1 app. §2. Feeding a full-resolution image directly is not what the reported numbers use.
• Benchmark protocols. RAFT* and GMA* use warm start from previous frames, which the proposed method does not §4.1. The paper marks the comparison, but a reader comparing numbers across tables should keep the protocol in view.
• Full cost volume in memory. The code materializes the full H×W×H×W cost volume for each pair before summarizing it encoder.py. The paper states the volume's shape §3.1 but gives no memory figures. Keys and values in the decoder are computed once and reused across iterations §3.3; a naive implementation recomputes them each step.
• Coordinate units. The patch positions in the encoder are the cost-map cell centres, eight times the cell index plus four encoder.py. The decoder encodes the flow positions directly, in the same 1/8-resolution units decoder.py. Positional-encoding frequencies depend on these units, so reimplementations must keep them.
• Hard-coded constants. The positional encoding uses 3.14 instead of pi attention.py. Replacing it with math.pi changes the features.
• Pretraining matters. Twins trained from scratch reaches 1.44 and 2.86 on Sintel clean and final, against 1.29 and 2.72 with ImageNet pretraining tab. 2. The paper attributes this to the transformer's larger receptive field and capacity §4.3.
• Cost of the larger model. The full model has 18.2M parameters, against 5.9M for GMA tab. 3. Part of the gain comes from the Twins encoder: GMA-Twins reaches 1.15 on Sintel clean, still 15% above FlowFormer §4.3 tab. 3.

Lineage

• Builds on: RAFT, for the recurrent flow refinement and the 4D cost volume §2 §3.3. The paper's bibliography gives no arXiv identifier for it.
• Builds on: GMA, for the ConvGRU update design and motion aggregation §3.3 [arXiv:2104.02409].
• Builds on: Twins-SVT, for the image encoder and the SS-SelfAttention layer §3.2 §4.1 [arXiv:2104.13840].
• Builds on: COTR, for the positional embedding §3.2 [arXiv:2103.14167]; Perceiver IO, for the tiling procedure §4.1 app. §2 [arXiv:2107.14795].
• Compared against: PWC-Net, RAFT, GMA, S-Flow, FM-RAFT, LiteFlowNet, LiteFlowNet2, HD3, PWC-Net+, VCN and MaskFlowNet tab. 1.
• Superseded by: not determined from this wiki's content.