SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Context-PIPs: Persistent Independent Particles Demands Spatial Context Features
Weikang Bian $^{1,2*}$  Zhaoyang Huang $^{1*}$  Xiaoyu Shi $^{1}$  Yitong Dong $^{3}$  Yijin Li $^{3}$  Hongsheng Li $^{1,2}$
$^{1}$ CUHK MMLab  $^{2}$ Centre for Perceptual and Interactive Intelligence  $^{3}$ Zhejiang University wkbian@outlook.com, drinkingcoder@link.cuhk.edu.hk, hsli@ee.cuhk.edu.hk
We tackle the problem of Persistent Independent Particles (PIPs), also called Tracking Any Point (TAP), in videos, which specifically aims at estimating persistent long-term trajectories of query points in videos. Previous methods attempted to estimate these trajectories independently to incorporate longer image sequences, therefore, ignoring the potential benefits of incorporating spatial context features. We argue that independent video point tracking also demands spatial context features. To this end, we propose a novel framework Context-PIPs, which effectively improves point trajectory accuracy by aggregating spatial context features in videos. Context-PIPs contains two main modules: 1) a SOurse Feature Enhancement (SOFE) module, and 2) a TTarget Feature Aggregation (TAFA) module. Context-PIPs significantly improves PIPs all-sided, reducing  $11.4\%$  Average Trajectory Error of Occluded Points (ATE-Occ) on CroHD and increasing  $11.8\%$  Average Percentage of Correct Keypoint (A-PCK) on TAP-Vid-Kinetics. Demos are available at https://wkbian.github.io/Projects/Context-PIPs/.
Video particles are a set of sparse point trajectories in a video that originate from the first frame (the source image) and move across the following frames, which are regarded as the target images. In contrast to optical flow estimation that computes pixel-wise correspondences between a pair of adjacent video frames, Persistent Independent Particles (PIPs) [12] or Tracking Any Point (TAP) [5] is interested in tracking the points in the follow-up frames that correspond to the original query points even when they are occluded in some frames. Video particles provide long-term motion information for videos and can support various downstream tasks, such as video editing [16] and Structure-from-Motion [48].
Long-range temporal information is essential for video particles especially when the particles are occluded because the positions of the occluded particles can be inferred from the previous and subsequent frames where they are visible. However, simultaneously encoding long image sequences brings larger computational and memory costs. Previous methods [12, 5] learn to track individual points independently because dense video particles are unnecessary in most scenarios. Inspired by optical flow estimation from visual similarities, they learn to predict point trajectories from the similarities between the query point and the subsequent target images. Specifically, given a query point at the source image, PIPs encode  $T$  feature maps from  $T$  consecutive target images and builds a  $T \times H \times W$  correlation volume by computing the feature similarity between the feature of the query point and the feature maps. The  $T$  particle positions are iteratively refined with the 3D correlation volume through a shared MLP-Mixer [40]. In other words, PIPs trades the spatial context features of
*Weikang Bian and Zhaoyang Huang assert equal contributions.
37th Conference on Neural Information Processing Systems (NeurIPS 2023), Vancouver, Canada.
the particle for longer temporal feature encoding. PIPs achieves great performance on the DAVIS dataset, which contains large movement particles and weak texture images (e.g., fast-moving dogs and black bears).
We argue that independent point tracking still demands spatial context features. Intuitively, although PIPs only accounts for specified query points, spatial context features around them provide informative cues for point trajectory refinement. For example, video particles on the same objects always share similar motions over time. In some video frames where the target particles are occluded, their surrounding particles may be visible and provide guidance for the position estimation of the target particles. However, PIPs only takes correlations and features belonging to the particles while ignoring abundant spatial context features around them. In this work, we propose tracking particles with Context (Context-PIPs) to improve independent point tracking with spatial context features. Context-PIPs contains two key modules for better point trajectory refinement: 1) a source feature enhancement (SOFE) module that learns to adopt more spatial context features in the source image and builds a guidance correlation volume, and 2) a target feature aggregation (TAFA) module that aggregates spatial context features in the target image guided by the correlation information.
In the source image, points that possess similar appearances are supposed to move in similar trajectories in subsequent frames. Such an assumption has also been used in GMA [21] for optical flow estimation. Given a query point, SOFE computes the correlation between the query point and the source image feature map, which is its self-similarity map. With the guidance of the correlation (self-similarity), SOFE predicts  $M$  offsets centered from the query point, and samples at the corresponding  $M$  auxiliary points to collect source context features. During the iterative point trajectory refinement, the correlation information between the  $M$  auxiliary features and  $T$  target feature maps is injected into the MLP-Mixer, which provides strong guidance and shows evident performance improvement.
Existing methods for optical flow and video particle estimation ignore features in target images for iterative refinement. To better utilize the context of target images, in each iteration, our TAFA module collects target features surrounding the previous iteration's resulting movements. TAFA for the first time shows that context features in target images also benefit correspondence estimation and further improve the point tracking accuracy.
Our contributions can be summarized as threefold: 1) We propose a novel framework, Context-PIPs, to improve independent video particle tracking with context features from both source and target features. Context-PIPs ranks 1st on the four benchmarks and shows clear performance superiority. 2) We design a novel Source Feature Enhancement (SOFE) module that builds a guidance correlation volume with spatial context features in the source image, and 3) a novel TTarget Feature Aggregation (TAFA) module that extracts context features from target images.
Optical Flow. Optical flow estimates the dense displacement field between image pairs and has traditionally been modeled as an optimization problem that maximizes the visual similarity between image pairs with regularizations [13, 1, 2, 34]. It serves as the core module of many downstream applications, such as Simultaneously Localization And Mapping (SLAM) [8, 47, 27, 43, 9], video synthesis [46, 14, 15, 42, 41], video restoration [24, 23], etc. Since FlowNet [6, 20], learning optical flow with neural networks presents superior performance over traditional optimization-based methods and is fast progressing with more training data obtained by the renderer and better network architecture [6, 20, 28, 35, 36, 18, 19, 45]. In recent years, iterative refining flow with all-pairs correlation volume presents the best performance. The most successful network designs are RAFT [39] and FlowFormer [17, 32], which achieves state-of-the-art accuracy.
Typical optical flow estimation only takes image pairs but longer image sequences can provide more information that benefits optical flow estimation. Kalman filter [4, 7] had been adopted in dealing with the temporal dynamics of motion and estimating multi-frame optical flow. Recent learning-based methods also attempted to exploit multi-frame information and perform multi-frame optical flow estimation. PWC-Fusion [29] is the first method that learns to estimate optical flow from multiple images. However, it only fuses information from previous frames in U-Net and improves little performance. The "warm start" technique [39, 33, 37] that wrapped the previous flow to initialize the next flow is firstly proposed in RAFT and shows clear accuracy increasing. VideoFlow [31] takes multi-frame cues better, iteratively fusing multi-frame information in a three-frame and five
frame structure, which reveals that longer temporal information benefits pixel tracking. Recently, researchers [25, 10] also tried to learn to estimate optical flow with event cameras.
Tracking Any Point. Optical flow methods merely focus on tracking points between image pairs but ignore tracking points across multiple consecutive frames, which is still challenging. Harley et al. [12] studied pixel tracking in the video as a long-range motion estimation problem inspired by Particle Video [30]. They propose a new dataset FlyingThings++ based on FlyingThings [26] for training and Persistent Independent Particles (PIPs) to learn to track single points in consecutive frames with fixed lengths. Doersch et al. [5] is a parallel work, which formalized the problem as tracking any point(TAP). They also propose a new dataset Kubric [5] for training and a network TAP-Net to learn point tracking. Moreover, they provide the real video benchmarks that are labeled by humans, TAP-Vid-DAVIS [5] and TAP-Vid-Kinetics [5], for evaluation. PIPs and TAP solve the video particle tracking problem in a similar manner, i.e., recurrently refining multi-frame point trajectory via correlation maps. In this paper, our Context-PIPs follows the training paradigm of PIPs and improves the network architecture design of PIPs. We also take the TAP-Vid-DAVIS and TAP-Vid-Kinetics benchmarks from TAP-Net for evaluation.
In contrast to optical flow methods [39, 17] that track dense pixel movement between an image pair, the problem of point tracking takes  $T$  consecutive RGB images with a single query point  $\mathbf{x}_{src} \in \mathbb{R}^2$  at the first frame as input, and estimates  $T$  coordinates  $\mathbf{X} = \{\mathbf{x}_0, \mathbf{x}_1, \dots, \mathbf{x}_{T-1}\}$  at the video frames where every  $\mathbf{x}_t$  indicates the point's corresponding location at time  $t$ . Persistent Independent Particles (PIPs) [12] is the state-of-the-art network architecture for TAP. It iteratively refines the point trajectory by encoding correlation information that measures visual similarities between the query point and the  $T$  video frames. The query points to be tracked are easily lost when the network only looks at them and ignores spatial context features. We propose a novel framework Context-PIPs (Fig. 1) that improves PIPs with a SOource Feature Enhancement (SOFE) module and a TArget Feature Aggregation (TAFA) module. In this section, we first briefly review PIPs and then elaborate on our Context-PIPs.
3.1 A Brief Revisit of PIPs
PIPs [12] processes  $T$  video frames containing  $N$  independent query points simultaneously and then extends the point trajectories to more video frames via chaining rules [12]. Given a source frame with a query point  $\mathbf{x}_{src} \in \mathbb{R}^2$  and  $T - 1$  follow-up target video frames, PIPs first extracts their feature maps  $\mathbf{I}_0, \mathbf{I}_1, \ldots, \mathbf{I}_{T - 1} \in \mathbb{R}^{C \times H \times W}$  through a shallow convolutional neural network and bilinearly samples to obtain the source point feature  $\mathbf{f}_{src} = \mathbf{I}_0(\mathbf{x}_{src})$  from the first feature map at the query point  $\mathbf{x}_{src}$ .  $C, H, W$  are the feature map channels, height, and width. Inspired by RAFT [39], PIPs initialize the point trajectory and point visual features at each frame with the same  $\mathbf{x}_{src}$  and  $\mathbf{f}_{src}$ :
$$
\begin{array}{l}\mathbf {X} ^ {0} = \left\{\mathbf {x} _ {0} ^ {0}, \mathbf {x} _ {1} ^ {0}, \dots , \mathbf {x} _ {T - 1} ^ {0} \mid \mathbf {x} _ {t} ^ {0} = \mathbf {x} _ {s r c}, t = 0, \dots , t = T - 1 \right\},\\\overline {{\mathbf {\Pi}}} ^ {0} = \left\{ \right.\mathbf {\Pi} _ {0} ^ {0}, \mathbf {\Pi} _ {1} ^ {0}, \dots , \mathbf {\Pi} _ {T - 1} ^ {0}, \dots , \mathbf {\Pi} _ {T - 1} ^ {0}, \dots , \mathbf {\Pi} _ {T - 1} ^ {0}, \dots , \mathbf {\Pi} _ {T - 1} ^ {0}, \dots , \mathbf {\Pi} _ {T - 1} ^ {0}, \dots , \mathbf {\Pi} _ {T - T - 1} ^ {0}, \dots , \mathbf {\Pi} _ {T - T - 1} ^ {0}, \dots , \mathbf {\Pi} _ {T - T - 1} ^ {0}, \dots , \mathbf {\Pi} _ {T - T - 1} ^ {0}, \dots , \mathbf {\Pi} _ {T - T - 1} ^ {0}, \dots , T - 1, T - T - 1, T - T - 1, T - T - 1, T - T - 1, T - T - 1, T - T - 1, T - T - 1, T - T - 1, T - T - 1, T - T - 1, T - T - 1, T - T - 1, T - T - 1, T - T - 1, T - T + 1, T - T + 1, T - T + 1, T - T + 1, T - T + 1, T - T + 1, T - T + 1, T - T + 1, T - T + 1, T - T + 1, T - T + 1, T - T + 1, T - T + 1, T - T + 1, T - T + 2, \dots , T - T + n, \dots , T - T + n, \dots , T - T + n, \dots , T - T + n, \dots , T - T + n, \dots , T - T + n, \dots , T - T + n, \dots , T - T + n, \dots , T - T + n, \dots , T - T + n, \dots , T - T + n, \dots , T - T + n.\end{array}\tag {1}
$$
$$
\mathbf {F} ^ {0} = \left\{\mathbf {f} _ {0} ^ {0}, \mathbf {f} _ {1} ^ {0}, \dots , \mathbf {f} _ {T - 1} ^ {0} \mid \mathbf {f} _ {t} ^ {0} = \mathbf {f} _ {s r c}, t = 0, \dots , t = T - 1 \right\},
$$
and iteratively refines them via correlation information.  $\mathbf{x}_t^k$  and  $\mathbf{f}_t^k$  respectively denote the point trajectory and point features in the  $t$ -th frame and  $k$ -th iteration. Intuitively, the point features store the visual feature at the currently estimated query point location in all the  $T$  frames.
Specifically, in each iteration  $k$ , PIPs construct multi-scale correlation maps [39] between the guidance feature  $\{\mathbf{f}_t^k\}_{t=0}^{T-1}$  and the target feature maps  $\{\mathbf{I}_t^k\}_{t=0}^{T-1}$ , which constitutes  $T$  correlation maps  $\mathbf{C}^k = \{\mathbf{c}_0^k, \mathbf{c}_1^k, \dots, \mathbf{c}_{T-1}^k\}$  of size  $T \times H \times W$ , and crops correlation information inside the windows centered at the point trajectory:  $\mathbf{C}^k(\mathbf{X}^k) = \{\mathbf{c}_0^k(\mathbf{x}_0^k), \mathbf{c}_1^k(\mathbf{x}_1^k), \dots, \mathbf{c}_{T-1}^k(\mathbf{x}_{T-1}^k)\}$ , where  $\mathbf{c}_t^k(\mathbf{x}_t^k) \in \mathbb{R}^{D \times D}$  denotes that we crop  $D \times D$  correlations from  $\mathbf{c}_t^k$  inside the window centered at  $\mathbf{x}_t^k$ . The point features  $\mathbf{F}^k$ , point locations  $\mathbf{X}^k$ , and the local correlation information  $\mathbf{C}^k(\mathbf{X}^k)$  are fed into a standard 12-layer MLP-Mixer that produces  $\Delta \mathbf{F}$  and  $\Delta \mathbf{X}$  to update the point feature and the point trajectory:
$$
\Delta \mathbf {F}, \Delta \mathbf {X} = \operatorname {M L P M i x e r} \left(\mathbf {F} ^ {k}, \mathbf {C} ^ {k} \left(\mathbf {X} ^ {k}\right), \operatorname {E n c} \left(\mathbf {X} ^ {k} - \mathbf {x} _ {s r c}\right)\right), \tag {2}
$$
$$
\mathbf {F} ^ {k + 1} = \mathbf {F} ^ {k} + \Delta \mathbf {F}, \mathbf {X} ^ {k + 1} = \mathbf {X} ^ {k} + \Delta \mathbf {X}.
$$
Figure 1: Overview of Context-PIPs Pipeline. Our Context-PIPs improves PIPs [12] with SOurce Feature Enhancement (SOFE) and TTarget Feature Aggregation (TAFA). PIPs iteratively refines the point trajectory  $\mathbf{X}^k$  with an MLP-Mixer with the current point trajectory  $\mathbf{X}^k$ , the correlation features  $\mathbf{C}^k$ , and the point features  $\mathbf{F}^k$ . SOFE and TAFA respectively improves the correlation features and point features, denoted as  $\hat{\mathbf{C}}^k$  and  $\hat{\mathbf{F}}^k$ .
PIPs iterates  $K$  times for updates and the point trajectory in the last iteration  $\mathbf{X}^K$  is the output.
PIPs achieves state-of-the-art accuracy on point tracking by utilizing longer temporal features. However, the previous method ignores informative spatial context features which are beneficial to achieve more accurate point tracking. Context-PIPs keeps all modules in PIPs and is specifically designed to enhance the correlation information  $\mathbf{C}^k$  and point features  $\mathbf{F}^k$  as  $\hat{\mathbf{C}}^k$  and  $\hat{\mathbf{F}}^k$  with the proposed SOFE and TAFA.
3.2 Source Feature Enhancement
Given the query point  $\mathbf{x}_{src}$  and feature map  $\mathbf{I}_0$  of the source image, PIPs simply samples a source feature  $\mathbf{f}_{src}$  at the query point location to obtain the point visual features  $\mathbf{F}^k$ . Although the point features are updated via the iterative refinement, its perceptive field is limited to the single point, easily compromised in harsh scenarios. The correlation maps  $\mathbf{C}^k$  in the  $k$ -th iteration provide vague information when the query point is in a less textured area. Moreover, the correlation map  $\mathbf{c}_t^k$  at timestamp  $k$  is ineffective once the particle is occluded at the  $t$ -th frame. To enhance the source feature, as shown in Fig. 1, we propose SOurse Feature Enhancement (SOFE) that accepts spatial context features in the source image as auxiliary features to guide the point trajectory refinement. The MLP-Mixer can infer the point locations via the auxiliary features even when the points are occluded or on less textured regions, which improves the point tracking accuracy and robustness.
Directly adopting all features in the source image brings large computational costs. SOFE learns to sample a small number of auxiliary features to enhance the source feature. Specifically, SOFE improves the point features in three steps. Firstly, SOFE learns to predict  $M$  offsets  $\delta \mathbf{x}_0, \delta \mathbf{x}_1, \ldots, \delta \mathbf{x}_{M-1} \in \mathbb{R}^2$  with an MLP-based sampler to sample  $M$  auxiliary features  $G = \{\mathbf{g}_0, \mathbf{g}_1, \ldots, \mathbf{g}_{M-1} | \mathbf{g}_m = \mathbf{I}_0(\mathbf{x}_{src} + \delta \mathbf{x}_m)\}$  around the query point  $\mathbf{x}_{src}$  in the source image. Motivated by GMA that aggregates pixel flows from pixels that are likely to belong to the same object through self-similarity, our proposed sampler also learns the locations of the auxiliary features based on local self-similarities  $\mathbf{c}_0^0(\mathbf{x}_{src})$  which store the correlations cropped from the first frame at the query point location. Secondly, we construct the correlation map  $\mathbf{c}_{m,t}' = < \mathbf{g}_m, \mathbf{I}_t > \in \mathbb{R}^{H \times W}$  that measure visual similarities of the  $m$ -th auxiliary feature and the  $t$ -th frame feature map.  $\hat{\mathbf{c}}_{m,t}$  provides additional correlation information to guide the iterative point trajectory refinement. In each iteration  $k$ , we crop the additional correlation information  $\mathbf{c}_m'(\mathbf{x}_t^k)$  according to the point locations
$\mathbf{x}_t^k$  and concatenate them with the original point correlation information  $\mathbf{c}_t^k (\mathbf{x}_t^k)$ , where  $\mathbf{c}_m^\prime (\mathbf{x}_t^k)$  denotes the same cropping operation as  $\mathbf{c}_t^k (\mathbf{x}_t^k)$ . Finally, for each frame  $t$ , we reduce the augmented correlations to a correlation feature vector  $\hat{\mathbf{c}}_t$  of length 196 through a correlation encoder CorrEnc.
$$
\hat {\mathbf {c}} _ {t} ^ {k} = \operatorname {C o r r E n c} \left(\operatorname {C o n c a t} \left(\mathbf {c} _ {0} ^ {\prime} \left(\mathbf {x} _ {t} ^ {k}\right), \mathbf {c} _ {1} ^ {\prime} \left(\mathbf {x} _ {t} ^ {k}\right), \dots , \mathbf {c} _ {M - 1} ^ {\prime} \left(\mathbf {x} _ {t} ^ {k}\right), \mathbf {c} _ {t} ^ {k} \left(\mathbf {x} _ {t} ^ {k}\right)\right)\right), \tag {3}
$$
and inject  $\hat{\mathbf{C}}^k = \{\hat{\mathbf{c}}_0^k,\hat{\mathbf{c}}_1^k,\dots ,\hat{\mathbf{c}}_{T - 1}^k\}$  into the MLP-Mixer. Compared with PIPs that only adopts  $\mathbf{c}_t^k (\mathbf{x}_t^k))$  , SOFE provides more informative correlations to the MLP-Mixer with spatial context features but does not increase its parameters and computations. The additional auxiliary features from the self-similarity map of the source image enhance the original source point features and significantly improves tracking accuracy.
3.3 Target Feature Aggregation
Inspired by existing optical flow methods, PIPs iteratively refines the point trajectory with correlation information and context features and also iteratively updates the point visual feature  $\mathbf{F}^{k + 1} = \mathbf{F}^k +\Delta \mathbf{F}$  after initializing them with the source feature, which presents benefits for point trajectory refinement. We observe that the input for supporting the point feature updating comes from the correlations  $\mathbf{C}^k$  only. However, such correlations  $\mathbf{C}^k$  are calculated as only cosine distances between the source point visual feature  $\mathbf{F}^k$  and the target features around currently estimated point locations  $\mathbf{X}^k$ , which provide limited information on how to conduct visual feature updating. Can we better guide the point feature update with context features in target images?
We, therefore, propose TTarget Feature Aggregation (TAFA) to augment point features with target image features nearby the point trajectory. Specifically, for each target frame  $t$ , a patch of shape  $D \times D$  cropped from the corresponding target feature map  $\mathbf{I}_t$  centered at  $\mathbf{x}_t^k$  to generate key and value. The augmented correlation features  $\hat{\mathbf{C}}$  in Eq. 3 encode abundant visual similarities. Therefore, we generate a query from it to extract target context features and adopt cross-attention with relative positional encoding to obtain the target context feature  $\mathbf{f}_t^{rk}$ , which is added to the original source point feature  $\hat{\mathbf{f}}_t^k = \mathbf{f}_t^k + \mathbf{f}_{t'}^{rk}$ . Finally, such augmented point features  $\hat{\mathbf{F}}^k = \{\hat{\mathbf{f}}_0^k, \hat{\mathbf{f}}_1^k, \dots, \hat{\mathbf{f}}_{T - 1}^k\}$  are injected into the MLP-Mixer. Similar to our proposed SOFE, TAFA also keeps the same parameters and computations of MLP-Mixer as PIPs while providing additional target context features and further improving PIPs. Although context features in the source image are used since RAFT [39], no previous methods adopt context features from target images. TAFA for the first time validates that target images also contain critical context features that benefit point movement refinement. SOFE improves PIPs with auxiliary features in the source image while TAFA absorbs more target image features. Equipping SOFE and TAFA to PIPs constitutes our final model, Context-PIPs.
3.4 Loss Functions
We compute the L1 distance between  $\mathbf{X}^k$  computed in iteration  $k$  and the ground truth  $\mathbf{X}_{gt}$  and constrain with exponentially increasing weights  $\gamma = 0.8$ .
$$
\mathcal {L} _ {T A P} = \sum_ {k = 1} ^ {K} \gamma^ {K - k} \left\| \mathbf {X} ^ {k} - \mathbf {X} _ {g t} \right\| _ {1} \tag {4}
$$
In addition, we will predict the visibility/occlusion  $\mathbf{V}$  by a linear layer according to the  $\hat{\mathbf{F}}^K$  obtained by the iterative update. And the cross entropy loss is used to supervise  $\mathbf{V}$  with the ground truth  $\mathbf{V}_{gt}$ .
$$
\mathcal {L} _ {\text {V i s}} = \mathbf {V} _ {g t} \log \mathbf {V} + (1 - \mathbf {V} _ {g t}) \log (1 - \mathbf {V}). \tag {5}
$$
The final loss is the weighted sum of the two losses:
$$
\mathcal {L} _ {\text {t o t a l}} = w _ {1} \mathcal {L} _ {\text {T A P}} + w _ {2} \mathcal {L} _ {\text {V i s}}. \tag {6}
$$
We use  $w_{1} = 1$  and  $w_{2} = 10$  during training.
Table 1: Experiments on FlyingThings++, CroHD, TAP-Vid-Davis, and TAP-Vid-Kinetics (first). Vis. and Occ. denotes ATE-Vis. and ATE-Occ.
Method	FlyingThings++	CroHD	DAVIS (first)	Kinetics (first)
Vis.	Occ.	Vis.	Occ.	AJ	A-PCK	AJ	A-PCK
DINO	43.05	76.25	22.50	26.06	-	-	-	-
RAFT	16.75	43.21	7.99	13.16	27.1	42.1	35.1	49.7
TAP-Net	-	-	17.00	20.86	33.0	48.6	38.5	54.4
PIPs (Paper)	15.54	36.67	5.16	7.56	-	-	-	-
PIPs (Re-imp.)	7.40	24.40	4.73	7.97	39.2	55.1	33.4	51.0
Context-PIPs (Ours)	6.44	22.22	4.28	7.06	42.7	60.3	40.2	57.0