SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Juliette Bertrand\*1,2 Giorgos Kordopatis-Zilos\*1 Yannis Kalantidis2 Giorgos Tolias
<sup>1</sup>VRG, FEE, Czech Technical University in Prague
$^{2}$ NAVER LABS Europe
Sim-to-real transfer
Frames and ground-truth masks
STCN
STCN + tt-MCC (ours)
Figure 1: Test-time training significantly improves performance in the presence of distribution shifts for the task of video object segmentation (VOS). Left: Example using a model trained on synthetic videos from BL-30K [8] and tested on a real video from DAVIS [38]. Right: Example using a model trained on DAVIS [38] and YouTube-VOS [51] and tested on a cartoonized video from the corrupted DAVIS-C benchmark which we introduce in this work. Second-to-bottom row: Results obtained using the STCN [9] approach. Bottom row: Results after test-time training with the proposed mask cycle consistency loss (tt-MCC) using the single ground-truth mask provided for the first video frame.
Corrupted test examples
The video object segmentation (VOS) task involves the segmentation of an object over time based on a single initial mask. Current state-of-the-art approaches use a memory of previously processed frames and rely on matching to estimate segmentation masks of subsequent frames. Lacking any adaptation mechanism, such methods are prone to test-time distribution shifts. This work focuses on matching-based VOS under distribution shifts such as video corruptions, stylization, and sim-to-real transfer. We explore test-time training strategies that are agnostic to the specific task as well as strategies that are designed specifically for VOS. This includes a variant based on mask cycle consistency tailored to matching-based VOS methods. The experimental results on common benchmarks demonstrate that the proposed test-time training yields significant improvements in performance. In particular for the sim-to-real scenario and despite using only a single test video, our approach manages to recover a substantial portion of the performance gain achieved through training on real videos. Additionally, we introduce DAVIS-C, an augmented version of the popular DAVIS test set, featuring extreme distribution shifts like image-/video-level corruptions and stylizations. Our results illustrate that test-time training enhances performance even in these challenging cases. Project page: https://jbertrand89.github.io/test-time-training-vos/
In one-shot video object segmentation (VOS), we are provided with a single segmentation mask for one or more objects in the first frame of a video, which we need to segment across all video frames. It is a dense prediction task over time and space, and therefore, collecting training data is highly demanding. Early VOS methods design foreground/background segmentation models that operate on single-frame input and require a two-stage training process [4, 29]. During the off-line stage, the model is trained to segment a variety of foreground objects. During the on-line stage for a specific test video, the model is fine-tuned to adapt to the objects of interest indicated in the provided mask.
Recently, matching-based methods like STCN [9] or XMem [6] have shown impressive performance on the common VOS benchmarks. These methods propagate the segmentation mask from a memory of previously predicted masks to the current frame. They only involve an off-line training stage where the model learns how to perform matching and propagation. Lacking any test-time adaptation mechanism, however, such methods are highly prone to test-time distribution shifts.
Our goal is to improve the generalization of matching-based VOS methods under distribution shifts by fine-tuning the model at test time through a single test video. Such one-shot adaptation is a form of test-time training (TTT), a research direction that is lately attracting much attention in classification tasks and a promising way for generalizing to previously unseen distributions. We argue that TTT is a great fit for matching-based VOS, not only because of the additional temporal dimension to exploit, but also because we are given a ground-truth label for a frame of the test example.
In this work, we focus on two cases of test-time distribution shift for VOS: a) when training on synthetic data and testing on real data, also known as sim-to-real transfer, and b) when the test data contain image/video corruption or stylization. We explore three different ways of performing TTT: i) a task-agnostic approach that borrows from the image classification literature and performs entropy minimization [34, 46], ii) an approach tailored to the STCN architecture that employs an auto-encoder loss on the mask provided for the first frame and does not exploit any temporal information, and iii) a mask cycle consistency approach that is tailored to matching-based VOS, and utilizes temporal cycle consistency of segmentation masks to update the model. As we show, these different test-time training strategies do not work equally well, and tailoring TTT to the task and the methods is important and not trivial.
Our experiments demonstrate that some of the proposed approaches constitute a powerful way of generalizing to unseen distributions. When starting from a model trained only on synthetic videos, test-time training based on mask cycle consistency improves STCN by  $+22$  and  $+11$  points in terms of  $\mathcal{J}\&\mathcal{F}$  score on the two most popular VOS benchmarks, YouTube-VOS and DAVIS, respectively. What is even more impressive is that by simply using test-time training, we recover the bulk of the performance gain that training the original model on real videos brings: for YouTube-VOS and DAVIS, TTT recovers  $82\%$  and  $72\%$  of the performance gains, respectively, that training on real videos brings compared to synthetic ones, but without requiring any video annotations during off-line training.
We also study the performance of TTT in the presence of image-level or video-level corruptions and stylization. To that end, we follow a process similar to ImageNet-C [14] and create DAVIS-C, a version of the DAVIS test set with 14 corruptions and style changes at three strength levels. We evaluate models with and without TTT on DAVIS-C and analyze how performance changes as the distribution shift increases. Our results show that TTT significantly improves the performance of STCN by more than 8 points in terms of  $\mathcal{I}\&\mathcal{F}$  score for the hardest case of corruption/stylization. A qualitative example for both types of distribution shift is shown in Figure 1.
Foreground/background models with online fine-tuning for VOS. Early approaches train a segmentation model to separate the foreground object from the background. The offline training stage is followed by an online supervised fine-tuning stage that uses the mask of the first frame for adapting the model [4, 29]. The fine-tuning process is improved by including pseudo-masks derived from highly confident predictions [44] or by additionally including hard-negative examples [48]. Maninis et al. [33] additionally incorporate semantic information about the underlined object categories via an external instance-aware semantic segmentation model. Other methods first produce a coarse segmentation result and then refine it by the guidance of optical flow [15], the appearance cue only [39], or temporal
consistency [32]. Hu et al. [16] utilize an RNN model to exploit the long-term temporal structures, while others additionally incorporate optical flow information [16]. The heavy cost of fine-tuning is reduced by meta-learning in the work of Xiao et al. [49] or Meinhardt and Leal-Taixe [35].
Matching-based VOS. Most recent approaches build upon a matching-based prediction model, which is trained offline and neither requires nor incorporates any fine-tuning stage. The current test frame is matched to either the first annotated frame [52, 18, 5, 45, 3], a propagated mask [20, 17, 28], or a memory containing both [31, 36, 9]. Early methods use simple matching process [18] or are inspired by classical label propagation [60], while different methods improve the design of the matching [1, 9, 40, 53, 55, 54], the type and extent of the memory used [6], or exploit interactions between objects or frames [37, 56], among other aspects.
Hybrid and other VOS methods. Mao et al. [34] borrow from both families of approaches, i.e. matching-based and those that require online fine-tuning, and jointly integrate two models, but, unlike our work, the matching model is not updated at test time. Some methods differ from matching-based in the way they capture the evolution over time, e.g. by RNNs [43] or by spatio-temporal object-level attention [11].
Cycle consistency in VOS. Matching-based methods require dense video annotations for training, which is a limitation. To dispense with the need for mask annotations during training, cycle consistency of pixel or region correspondences over space and time has been successfully used by prior work [19, 24, 61]. This is an unsupervised loss that is also appropriate for test-time training. In our case, we tailor cycle consistency to the task of matching-based VOS, and propose a supervised consistency loss that operates on masks by taking advantage of the provided mask at test-time in order to better adapt to the object of interest. Mask consistency based on the first frame is used by prior work during model training. Khoreva et al. [21] synthesize future frames using appearance and motion priors, while in our case, instead of synthesizing, we use mask predictions of the existing future video frames. Li et al. [26] and HODOR [1] use a temporal cycle consistency loss on a the first frame mask during training that allows learning from videos with a single segmentation mask. In contrast, we utilize the mask frame that is provided at test-time for the first frame of the test video and we are able to outperform HODOR and other state-of-the-art methods for multiple test-time distribution shifts.
Test-time training. A family of approaches adapts the test example to the training domain. Kim et al. [22] chose the appropriate image augmentation that maximizes the loss according to a loss prediction function, while Xiao et al. [50] updates the features of the test sample by energy minimization with stochastic gradient Langevin dynamics. Another family of approaches adapts the model to the test domain. Entropy minimization is a common way to update either batch-normalization statistics only [46] or the whole model [59]. Other self-supervised objectives include rotation prediction [42], contrastive learning [30], or spatial auto-encoding [12]. In our work, we move beyond the image domain and introduce mask cycle consistency as an objective to adapt the model specifically for video object segmentation applications. Azimi et al. [2] evaluate some of the aforementioned TTT techniques on top of self-supervised dense tracking methods under several distribution shifts on videos. Another recent work [57], concurrent to ours, also studies TTT in the video domain but for a classification task; we instead fully tailor the TTT on the VOS task and use temporal mask consistency as our loss.
3 Test-time training for matching-based VOS
In this section, we first formulate the task of Video Object Segmentation (VOS) and briefly describe the basic components of the matching-based VOS methods we build on [9, 6]. We then present three ways for test-time training: a task-agnostic baseline based on entropy minimization that has been highly successful in other tasks and two VOS-specific variants that leverage the fact that we are provided with a ground-truth mask at test time; one using an auto-encoder loss and another a temporal cycle-consistency loss on the masks.
Figure 2: The mask cycle consistency loss for a given frame triplet  $\{x_0, x_i, x_j\}$  and mask  $m_0$ . Frames  $x_0, x_i$  and  $x_j$  are shown with a red, green and blue border, respectively. Function  $f$  represents the overall prediction model. It takes as input the current test frame and a memory  $M$  of predicted masks from the previous frames and outputs the predicted mask for the current frame. For the first frame, the memory contains the ground-truth mask.
The goal of VOS is to segment a set of particular objects in a video provided their ground-truth segmentation mask  $m_0$  in the first video frame  $x_0^1$ . Space-Time Correspondence Networks (STCN) [9] sequentially predict the mask of each frame separately, while the prediction is conditioned on the provided mask and the predicted masks of earlier frames that are stored in a memory queue. XMem [6] is an extension of STCN that features a long-term memory and top performance for long-term prediction. Although we provide here the necessary parts of STCN required for presenting our contribution, we refer the reader to Cheng et al. [9] for a complete description. Our formulation is generic and applicable to a wide range of matching-based VOS methods, including the ones that require only a single forward pass for multiple objects such as AOT [56].
We denote the model that performs the overall process by function  $f$ . The predicted mask of frame  $x_{j}$  is given by  $\hat{m}_{j} = f(x_{j};M)$ , where  $M = \{(x_{i},m_{i})\}$  is the memory of previous predictions consisting of frame-mask tuples  $(x_{i},m_{i})$ . The memory always includes  $(x_0,m_0)$ , while any other tuple for  $i > 0$  is formed by using the corresponding predicted mask  $\hat{m}_{i}$  since no other ground-truth mask is available. Note that ground-truth masks for the first frame are binary, while the predicted ones consist of object probabilities in [0, 1].
The STCN model contains a frame encoder  $e_{x}$  that extracts frame representation  $u_{i} = e_{x}(x_{i}) \in \mathbb{R}^{W \times H \times D_{e}}$  and a mask encoder  $e_{m}$  that extracts mask representation  $v_{i} = e_{m}(m_{i}) \in \mathbb{R}^{W \times H \times D_{m}}$ . The latter can take either ground-truth (binary) or predicted (real-value) masks as input. Frame and mask representations of earlier frames are added to the memory.
During the prediction for the current frame denoted by  $x_{j}$ , representation  $u_{j}$  is matched (each  $D_{e}$ -dimensional vector separately) to the representations of frames in the memory, and the most similar items are identified. In particular, the full similarity matrix is formed between the current frame items and the memory items, but only the top-k values per row (rows correspond to the current frame) are maintained, and the rest are set to 0. By the term item we refer to each  $D_{e}$ -dimensional vector, with each frame having  $W * H$  of them. The corresponding mask representation of these most similar items is aggregated to obtain  $\hat{v}_{j} \in \mathbb{R}^{W \times H \times D_{m}}$ . In this way we compose the mask representation for frame  $x_{j}$  via matching since the mask itself is not available. This process is
In the following, we assume the presence of a single object for simplicity. In the case of multiple objects, the prediction process is repeated for each object independently, similar to [6, 9].
equivalent to a cross-attention operation between the test frame features (queries) and the memory frame features (keys) that aggregates memory mask features (values). Finally, a mask decoder  $d_{m}$  is used to obtain the predicted mask  $\hat{m}_{j} = d_{m}([\hat{v}_{j},u_{j}]$ , where  $[\hat{v}_j,u_j]\in \mathbb{R}^{W\times H\times (D_e + D_m)}$  is the concatenation of the two representations along the depth dimension.
In this work, we assume that we have access to a matching-based VOS model, trained on real or synthetic data with mask annotations. These models do not generalize well to extreme distribution shifts not encountered during supervised training. Our goal is to improve their performance in such cases by fine-tuning the parameters of the model at test time.
3.2 Test-time training
We explore three different losses to perform TTT for the case of matching-based VOS: (i) based on entropy minimization (tt-Ent), which forms a paradigm transferred from the image domain (image classification in particular) and is a task- and method-agnostic loss, (ii) based on an auto-encoder (tt-AE) that operates on segmentation masks, which is an approach tailored for the STCN architecture, (iii) based on mask cycle consistency (tt-MCC), which is more general and appropriate for a variety of matching-based methods. tt-MCC exploits the matching-based nature of STCN and includes mask propagation via a memory of masks, while the tt-AE variant does not include the matching step between multiple frames.
Given a test video, we optimize the parameters of the model using the provided mask and either only the first frame (for  $tt$ -AE) or all its frames ( $tt$ -Ent,  $tt$ -MCC). The test-time loss is optimized, the model parameters are updated, and the model is used to provide predictions for the specific test video. Then, the model is re-initialized back to its original parameters before another test video is processed.
For the two TTT variants that use multiple frames, i.e.  $tt$ -MCC and  $tt$ -Ent, we follow the process used during the off-line training phase of STCN and XMem and operate on randomly sampled frame triplets and octuplets, respectively. In the following, we describe the case of triplets which is extended to larger sequences in a straightforward way. All the triplets include the first video frame, while the second (third) triplet frame is uniformly sampled among the  $s$  frames that follow the first (second) triplet frame in the video. For triplet  $x_0, x_i, x_j$  with  $0 < i < j$ , the first predicted mask is given by  $\hat{m}_i = f(x_i; \{(x_0, m_0)\})$ , and the second by  $\hat{m}_j = f(x_j; \{(x_0, m_0), (x_i, \hat{m}_i)\})$ , i.e. the second frame is added to the memory before predicting the last mask of the triplet. Multiple triplets are sampled during test-time training. The per-pixel losses are averaged to form the frame/mask loss. We denote the cross entropy loss between a ground-truth mask  $m$  and a predicted mask  $\hat{m}$  by  $\ell_{CE}(m, \hat{m})$ .
3.2.1 Entropy (tt-Ent) loss
We start from a task- and method-agnostic loss based on Entropy (tt-Ent) minimization, a common choice for test-time training on image classification [46]. In this case, we force the model to provide confident predictions for the triplets we randomly sample. The entropy of mask  $\hat{m}_i$  is denoted by  $H(\hat{m}_i)$  and the loss for a particular triplet is  $L_{H} = 0.5*(H(\hat{m}_{i}) + H(\hat{m}_{j}))$ , which is minimized over multiple triplets. Prior work on image classification avoids the trivial solution by using a batch of test examples [46] or augmenting the input example and minimizing the entropy over the average prediction of the augmented frames [59]. In the task of semantic segmentation, batches or augmentations are not required since it is a dense prediction task and we have multiple outputs to optimize, i.e. a prediction per pixel. In the case of VOS, we have even more outputs due to the temporal dimension and because we are sampling triplets among many frames. However, similar to Wang et al. [46], we observe that training only the batch normalization parameters is a requirement for tt-Ent to work.
3.2.2 Mask auto-encoder (tt-AE) loss
The STCN model includes a mask encoder and a mask decoder to decode the mask representation coming from the cross-attention process. We repurpose these components for test-time training and compose an auto-encoder that reconstructs the input mask. The auto-encoder input is the provided ground-truth mask, i.e.  $m_0$ , and the loss is given by  $L_{AE} = \ell_{CE}(m_0, d_m([e_m(m_0), e_x(x_0)]))$ . Note that, together with the mask encoder that gets optimized to better represent the specific object shape, we also optimize the frame encoder with the  $tt$ -AE loss, since the mask representation is concatenated with the frame representation in the input of the mask encoder [9].
(b) provided mask
(c) Mask predictions before test-time training
(d) loss before TTT
(e) Gradients of mask predictions before test-time training
$\hat{m}_i$  (step 1)
(f) Mask predictions after test-time training
Figure 3: Unravelling the  $tt$ -MCC loss. (a) a triplet from a video with gaussian noise corruption; (b) the provided mask; (c) mask predictions after each of the four steps of the cycle, before test-time training (TTT); (d) the pixel-wise loss before TTT; (e) gradients of mask predictions for each of the 4 steps of the cycle; (f) predictions after test-time training with the mask cycle consistency loss; (f) the loss after TTT. We see that  $tt$ -MCC significantly improves all predicted masks and reduces the loss. We binarize the predicted masks for visualization purposes; they are otherwise probability maps.
$\hat{m}_j$  (step 2)
$\hat{m}_i^b$  (step 3)
$\hat{m}_0^b$  (step 4)
$\ell_{CE}(m_0,\hat{m}_0^b)$
(g) loss after TTT
3.2.3 Mask cycle consistency (tt-MCC) loss
After sampling a triplet, we make predictions in a time-forward and time-backward manner to form a cycle that allows us to use the provided ground-truth mask as a label with a cross-entropy loss. The cycle consists of four steps. We first obtain masks  $\hat{m}_i$  and  $\hat{m}_j$  with time-forward predictions. Then, we consider the triplet in reverse order going from  $x_j$  to  $x_i$  and then to  $x_0$ , and obtain time-backward predictions. In particular,  $\hat{m}_j$  is treated as the only mask in the memory to enable prediction for frame  $x_i$  which is given by  $\hat{m}_i^b = f(x_i, \{(x_j, \hat{m}_j)\})$ . In the last step, the mask of  $x_0$  is given by  $\hat{m}_0^b = f(x_0, \{(x_j, \hat{m}_j), (x_i, \hat{m}_i^b)\})$ . We compare this prediction to the provided ground-truth mask and the loss to optimize is  $L_{MCC} = \ell_{CE}(\hat{m}_0^b, m_0)$ . The overall process and the four prediction steps are shown in Figure 2. During the optimization, back-propagation is performed through all the predicted masks mentioned above. In this way, earlier predictions are also optimized in an indirect way.
For successful object segmentation in  $\hat{m}_0^b$ , masks  $\hat{m}_i$ ,  $\hat{m}_j$ , and  $\hat{m}_i^b$ , in the whole sequence of predictions, are required to have successful object segmentation too. If for example an object part is not properly segmented in the intermediate predictions, it will be missed in  $\hat{m}_0^b$  too. We further experimented with adding an extra mask consistency loss between the forward and backward predictions for the middle/intermediate frame, i.e.  $\ell_{CE}(\hat{m}_i^b, \hat{m}_i)$  but did not result in any additional gains.
Figure 3 shows the four predictions during a cycle using a particular triplet before and after  $tt$ -MCC. We see that initially, the object is not properly segmented during internal steps of the cycle and is fully missed in the last step, of the cycle resulting in large loss values for most of its pixels. After TTT is performed, not only the mask which is given as input to the loss, i.e.  $\hat{m}_0^b$  gets improved, but so are the predictions during intermediate steps. Note that the prediction of step 4 is not part of the inference process, but is rather a proxy to optimize predictions for steps 1 and 2 that are part of it. The gradients of the internal masks are indicating how the loss is affecting their predictions per pixel.
saturate
gaussian noise
cartoon
Figure 4: The 14 types of corruption in DAVIS-C. Corruptions are depicted at "medium" strength.
defocus blur
style 1
glass blur
style 2
crf compression
style 3
motion blur
style 4
In this section, we first describe the training and implementation details. Then, we introduce the DAVIS-C benchmark and report our experimental results on four datasets and two distribution shifts: sim-to-real transfer and corrupted/stylized test data.
Initial models. We start from two publicly available models for STCN. STCN-BL30K is the model trained on synthetic data from BL-30K [8]. STCN-DY is a model initialized by STCN-BL30K and further trained using the training videos of DAVIS and YouTube-VOS, which include approximately 100K frames annotated with segmentation masks. We use the same two models for XMem, denoted by XMem-BL30K and XMem-DY, respectively.