Tracking Everything Everywhere All at Once

In plain words

Tracking a point through a video, even when it hides behind something for a while, is hard if you only look at pairs of neighbouring frames. This paper takes a different route: for each video it trains a small model from scratch, so that the model knows where every piece of the scene is in one shared reference space. Each frame gets its own mapping into that space, and the mapping can be run in both directions, so a point followed from frame A to frame B and back returns to where it started.

The model is taught with rough motion estimates that some existing tool already produced, plus the requirement that the colours it renders match the video. Because the whole video is fitted together, a point that disappears for a while can still be placed sensibly when it comes back. The model does not separate camera motion from object motion, which is why the authors call it quasi-3D rather than a true 3D reconstruction.

TL;DR

• Test-time optimisation per video: a canonical volume (NeRF-style density and colour) plus one invertible network per frame, conditioned on a per-frame latent code §4 §4.1 §4.2.
• Trained on filtered pairwise optical flow (RAFT by default, TAP-Net as an alternative), with a flow loss, a photometric loss and an acceleration regulariser §5.1 §5.2.
• Author-reported on TAP-Vid: Ours (RAFT) reaches AJ 55.1 on Kinetics, 51.7 on DAVIS and 77.5 on RGB-Stacking, above the compared baselines on AJ tab. 1.
• Costs: all pairwise flows are computed exhaustively, and training runs for 200K iterations §5.4 app. C §7.
• Official code exists (Apache-2.0), but the paper and code disagree on several details; see Paper vs code.

Problem

Dense correspondence methods that look at two frames at a time, such as pairwise optical flow, lose a point once it is occluded, and composing their outputs over many frames produces inconsistent trajectories §4. A full dynamic 3D reconstruction would avoid both problems, but it is an ill-posed problem with camera and scene motion entangled §4. The paper asks whether long-range, occlusion-aware motion can be recovered without an explicit dynamic 3D model §4.

Method

The representation is called OmniMotion. It has two parts §4:

1. Canonical volume. A coordinate network maps each canonical 3D point to a density and a colour, as in NeRF §4.1. The density tells where surfaces are in canonical space, and the colour supports a photometric loss §4.1.
2. Bijections per frame. Each frame has an invertible mapping between its local 3D coordinates and the canonical space. The mapping is a Real-NVP style stack of affine coupling layers, all sharing one network but conditioned on a per-frame latent code §4.2. Composing one frame's mapping with the inverse of another's gives a cycle-consistent correspondence between any two frames §4.2 eq. 1.

The representation does not separate camera motion from scene motion, so the authors call it quasi-3D rather than a physical reconstruction §4. The camera is treated as part of the bijections: rendering uses a fixed orthographic camera §4.3. To compute motion for a query pixel p in frame i and a target frame j §4.3:

• Lift p to K points along a ray at depths sampled between the near and far bounds §4.3 app. C.
• Map each point to canonical space with frame i's bijection and query the volume for density and colour §4.2 §4.1.
• Map each point into frame j with the inverse of frame j's bijection eq. 1.
• Alpha-composite the frame-j points with the weights from the volume eq. 2. Project the result to 2D to obtain the predicted correspondence §4.3.

Training fits the representation to the input flows and the video colours. Input correspondences are computed for all frame pairs, filtered with a cycle check and an appearance check, and optionally extended by chaining §5.1 app. A. Each step samples frame pairs and pixels, then minimises the total loss below §5.2 app. C.

# Correspondence for one query pixel across frames  [§4.3] [eq. 1] [eq. 2]
# Input:  pixel (2,) in frame i; target frame j; K depths
# Output: predicted pixel (2,) in frame j

def predict_pixel(pixel, frame_i, frame_j, depths):
    # Lift the pixel to K points on an orthographic ray  [§4.3]
    points_i = concat(pixel, depths)                    # (K, 3)
    # Map the points into canonical space with frame i's bijection  [§4.2]
    canonical = to_canonical(points_i, latent_code[frame_i])  # (K, 3)
    # Read density and colour from the canonical volume  [§4.1]
    # Contraction of canonical coordinates is applied first  [app. C]
    density, colour = canonical_field(canonical)        # (K,), (K, 3)
    # Alpha values from density  [eq. 2]
    alpha = 1 - exp(-density)                           # (K,)
    # Transmittance T_k = product of (1 - alpha_l) over l < k  [eq. 2]
    transmittance = exclusive_cumprod(1 - alpha)        # (K,)
    weights = transmittance * alpha                     # (K,)
    # Map every sample into frame j with the inverse bijection  [eq. 1]
    points_j = from_canonical(canonical, latent_code[frame_j])  # (K, 3)
    # Composite the mapped samples  [eq. 2]
    composited = sum(weights[:, None] * points_j, axis=0)  # (3,)
    # Orthographic projection keeps the first two coordinates  [§4.3]
    return composited[:2]                               # (2,)

# Total training objective for one batch  [§5.2] [eq. 3] [eq. 4] [eq. 5]
# Input:  sampled pairs with input flow, weight and observed colour
# Output: scalar loss

def training_loss(batch, step):
    # Predict correspondences and rendered colours for sampled pixels  [§4.3]
    predicted_flow = predict_flows(batch)               # (B, P, 2)
    rendered_colour = render_colours(batch)             # (B, P, 3)
    # Flow loss: weighted L1 to the filtered input flow  [eq. 3]
    flow_loss = sum(batch.weights * abs(predicted_flow - batch.flow))
    # Photometric loss: squared error to the observed colour  [eq. 4]
    photo_loss = sum((rendered_colour - batch.colour) ** 2)
    # Acceleration: L1 of the second difference over frames i-1, i, i+1  [eq. 5]
    accel = abs(x_prev + x_next - 2 * x)                # (B, P, 3)
    reg_loss = sum(accel)
    # Photometric weight ramps from 0 to 10 over the first 50k steps  [app. C]
    lambda_pho = min(10.0, 10.0 * step / 50000)
    # Regularisation weight is fixed at 20  [app. C]
    lambda_reg = 20.0
    # Total loss  [eq. 6]
    return flow_loss + lambda_pho * photo_loss + lambda_reg * reg_loss

Other training rules from the paper:

• Adam, 200K iterations, 8 frame pairs per batch with 256 correspondences each, so 1024 correspondences per step §5.4. Learning rates and decay are in the table below app. C.
• Pairs are sampled with a maximum frame interval that starts at 20 and grows by one every 2k steps app. C.
• Hard mining: every 20k steps the model's flow predictions are cached and compared with the input flows. Half of the sampled query pixels are drawn with weights proportional to that error, and half uniformly §5.3 app. C.
• The mapped points are kept inside the predefined depth range to avoid degenerate solutions app. C. A distortion loss from mip-NeRF 360 suppresses floaters app. C.
• An auxiliary gradient loss compares colour differences between random, possibly distant, pixel pairs, and the same gradient loss is applied to the flow app. C.

Inference: the query can be any pixel in any frame. The predicted location uses the single sample with the largest alpha, not the composited average, because the authors found it visually better with similar numbers app. C. Occlusion is detected by checking the transmittance of the mapped 3D location in the target frame §6.2.

Hyperparameters

Values are the paper's unless the Notes column says otherwise. "not stated" means the paper does not give a value.

Results

Author-reported. Setup: TAP-Vid benchmark §6.1. DAVIS is 30 videos from the 2017 validation set. Kinetics is a random subset of 100 videos from 1,189 videos of 250 frames. RGB-Stacking is 50 synthetic videos of 250 frames with 30 tracks each. Baselines: RAFT chained (RAFT-C) and direct (RAFT-D), PIPs, Flow-Walk chained and direct, TAP-Net and Deformable Sprites §6.2. Metrics: average position accuracy over 1, 2, 4, 8 and 16 px (δ_avg), Average Jaccard (AJ), occlusion accuracy (OA) and temporal coherence (TC, lower is better) §6.1. Test-time-optimised methods are Ours and Deformable Sprites; the rest run feed-forward tab. 1 caption.

Ablations on DAVIS, with the full model as reference tab. 2 app. B:

Qualitative results and pseudo-depth maps are described in fig. 3 and fig. 4. The paper's claims are these: the method tracks through long occlusions where PIPs cannot, gives smoother trajectories than direct pairwise methods, and beats chained methods on longer videos §6.3. The pseudo-depth maps show a relative depth ordering, not physical depth §6.4.

Paper vs code

Official code: the project page links https://github.com/qianqianwang68/omnimotion, a repository under the first author's account, with Apache-2.0 licence. The appendix says the code will be released on acceptance app. C. The code is at commit 43868e5d214b795f00c586b790cb832a6f571a4d. Searches made: the paper's own links, the arXiv abstract, the project page, and the first author's GitHub account. Citations below refer to that commit.

1. Pair weight formula. The paper's weight is 1 / cos(Delta / N' * pi / 2), which grows as the interval grows app. C. The code uses cos((Delta - 1) / N' * pi / 2), which falls from 1 towards 0 as the interval grows raft.py. These point in opposite directions. Pick one and say which in your implementation.
2. Loss normalisation. Paper: flow loss as a sum of L1 terms eq. 3, photometric loss as a sum of squared errors eq. 4, regulariser as a sum of L1 terms eq. 5. Code: the flow loss is a weighted mean of L1 with the normalisation switched off, after dropping the single largest element trainer.py criterion.py; the photometric loss is a mean of squared errors trainer.py; the regulariser is a mean-based L1 trainer.py criterion.py. The two differ by constant factors per term, which changes the effective lambda values.
3. Regulariser support. The paper applies the acceleration regulariser to the union of all local 3D samples eq. 5 §5.2. The code applies it to the first quarter of the sampled points in the batch only trainer.py.
4. Inference sample. The paper says the single sample with the largest alpha is used app. C. The code takes the sample with the largest blending weight, alpha times transmittance trainer.py trainer.py.
5. Occlusion. The paper checks the transmittance at the mapped 3D location §6.2. The code sums the blending weights of samples in front of the mapped point and thresholds that sum at 0.99 by default trainer.py config.py. The threshold is not stated in the paper.
6. Input-flow filtering. The paper's cycle threshold (3 px) and DINO check (cosine above 0.5, pairs more than 3 frames apart) match the preprocessing code app. A filter_raft.py filter_raft.py, except the code applies the DINO check at 3 frames or more. The occlusion bypass for pairs closer than 3 frames uses an extra factor of 1.5 on the forward-backward error filter_raft.py, which the paper does not state. In the training loader, pairs one frame apart keep every pixel raft.py. The paper limits the bypass to occluded-region correspondences app. A.
7. Training length. The paper trains for 200K iterations §5.4. The code's default is 100000 config.py, and the README says 100k README.md.
8. Affine pre-transform. The paper does not describe a global affine transform. The code applies a per-frame 2D affine map, predicted from time and depth, before the coupling layers when use_affine is set nvp_simplified.py nvp_simplified.py. The provided config turns it on default.txt; the argument default is off config.py.
9. Losses not in the paper's main equations. The code adds a depth-range loss with weight 10 trainer.py trainer.py, a distortion loss whose weight starts at step 40k and grows by 1/2000 per step up to 10 trainer.py trainer.py, and a flow-gradient loss whose weight grows by 1/500000 per step up to 0.1 trainer.py. The colour-gradient loss is weighted by the photometric schedule trainer.py. The paper gives no weights for these app. C. A canonical-sphere loss exists but its weight is 0 by default trainer.py.
10. Hard mining details. The paper describes pixel-level error sampling §5.3 app. C. The code also draws the base frame with probability proportional to its per-frame flow error raft.py, blurs the error map with a 5 x 5 Gaussian raft.py, and gives pairs one frame apart a weight of 0.5 when choosing the second frame raft.py.
11. Coupling layer details. The paper says the coupling layers alternate their split pattern app. C. The code draws a random permutation of the three coordinates for each group of three layers at construction nvp_simplified.py. The code also clamps the scale output to [-8, 8] nvp_simplified.py nvp_simplified.py; the clamp is not in the paper.
12. Matches. Learning rates, LR decay, photometric ramp, lambda_reg, pairs per batch, sample count, depth range, mapping and volume architectures, contraction and the window schedule all match the paper config.py trainer.py trainer.py train.py.

Pitfalls

• Pair weight. The paper's formula and the code's formula disagree in direction (see Paper vs code 1). Test both on one sequence before choosing.
• Units of the losses. Sums in the paper and means in the code mean the lambda values do not transfer as written. Decide which normalisation you implement and state it.
• Iteration budget. 200K in the paper, 100k in the code and README. Full training takes about 8 to 9 hours on an A100 according to the README README.md.
• Cost scales quadratically. Exhaustive pairwise flow costs grow with the square of the frame count, and the authors name this as a limitation §7. Chaining and appearance checks add preprocessing app. A.
• Sensitivity to initialisation. The optimisation is non-convex. The authors report occasional wrong surface ordering or duplicated objects in canonical space that optimisation can fail to fix §7.
• Fast-moving and thin structures are hard for the input flows, and so for this method §7.
• Occlusion decision. The paper describes a transmittance check §6.2. The code turns an accumulated weight into a binary decision with a threshold of 0.99 (see Paper vs code 5), which the paper does not state.
• Hidden hard-coded device. The fixed positional encoding creates its frequency tensor on cuda nvp_simplified.py, so the code will not run on CPU as written.
• Global SSL change. Importing the criterion module disables HTTPS certificate verification for the whole Python process criterion.py. Remove that line before reusing the module.
• NaN handling. On a NaN the training loop drops into the Python debugger trainer.py trainer.py. Replace these with an error for batch runs.
• Test-time optimisation per video. Every video is a new training run. Results cannot be reused across videos, so the method's output for one video says nothing about another.
• Subset sampling. Kinetics results use a random subset of 100 videos without a stated seed §6.1, so the exact subset cannot be recreated from the paper.
• Version. This card follows arXiv v2 (updated 2023-09-12), the version in the metadata this card was compiled from.

Lineage

• Builds on: NeRF-style volume rendering and coordinate networks §4.1 app. C (mildenhall2021nerf in the paper's bibliography; no card); Real-NVP affine coupling for invertible mappings §4.2 (dinh2016density; no card); RAFT and TAP-Net as input correspondences §5.1 (no card); mip-NeRF 360 contraction and distortion loss app. C (no card).
• Compared against: RAFT, TAP-Net, PIPs, Flow-Walk and Deformable Sprites §6.2. None has a card yet.
• Superseded by: not stated in this card's sources.