Rethinking Space-Time Networks with Improved Memory Coverage for Efficient Video Object Segmentation

1 Introduction . Video object segmentation ( VOS ) aims to identify and segment target instances in a video sequence . This work focuses on the semi-supervised setting where the first-frame segmentation is given and the algorithm needs to infer the segmentation for the remaining frames . This task is an extension of video object tracking [ 1 , 2 ] , requiring detailed object masks instead of simple bounding boxes . A high-performing algorithm should be able to delineate an object from the background or other distractors ( e.g. , similar instances ) under partial or complete occlusion , appearance changes , and object deformation [ 3 ] . Most current methods either fit a model using the initial segmentation [ 4 , 5 , 6 , 7 , 8 , 9 ] or leverage temporal propagation [ 10 , 11 , 12 , 13 , 14 , 15 , 16 ] , particularly with spatio-temporal matching [ 17 , 18 , 19 , 20 , 21 , 22 , 23 , 24 , 25 , 26 , 27 ] . Space-Time Memory networks [ 18 ] are especially popular recently due to its high performance and simplicity – many variants [ 22 , 16 , 23 , 21 , 24 , 28 , 29 , 30 ] , including competitions ’ winners [ 31 , 32 ] , have been developed to improve the speed , reduce memory usage , or to regularize the memory readout process of STM . In this work , we aim to subtract from STM to arrive at a minimalistic form of matching networks , dubbed Space-Time Correspondence Network ( STCN ) 1 . Specifically , we start from the basic premise †This work was done in The Hong Kong University of Science and Technology . 1Training/inference code and pretrained models : https : //github.com/hkchengrex/STCN 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . that correspondences are target-agnostic . Instead of building a specific memory bank and therefore affinity for every object in the video as in STM , we build a single affinity matrix using only RGB relations . For querying , each target object passes through the same affinity matrix for feature transfer . This is not only more efficient but also more robust – the model is forced to learn all object relations beyond just the labeled ones . With the learned affinity , the algorithm can propagate features from the first frame to the rest of the video sequence , with intermediate features stored as memory . While STCN already reaches state-of-the-art performance and speed in this simple form , we further probe into the inner workings of the construction of affinities . Traditionally , affinities are constructed from dot products followed by a softmax as in attention mechanisms [ 18 , 33 ] . This however implicitly encoded “ confidence ” ( magnitude ) with high-confidence points dominating the affinities all the time , regardless of query features . Some memory nodes will therefore be always suppressed , and the ( large ) memory bank will be underutilized , reducing effective diversity and robustness . We find this to be harmful , and propose using the negative squared Euclidean distance as a similarity measure with an efficient implementation instead . Though simple , this small change ensures that every memory node has a chance to contribute significantly ( given the right query ) , leading to better performance , higher robustness , and more efficient use of memory . Our contribution is three-fold : • We propose STCN with direct image-to-image correspondence that is simpler , more efficient , and more effective than STM . • We examine the affinity in detail , and propose using L2 similarity in place of dot product for a better memory coverage , where every memory node contributes instead of just a few . • The synergy of the above two results in a simple and strong method , which suppresses previous state-of-the-art performance without additional complications while running fast at 20+ FPS . 2 Related Works . Correspondence Learning Finding correspondences is one of the most fundamental problems in computer vision . Local correspondences have been used heavily in optical flow [ 34 , 35 , 36 ] and object tracking [ 37 , 38 , 39 ] with fast running time and high performance . More explicit correspondence learning has also been achieved with deep learning [ 40 , 41 , 42 ] . Few-shots learning can be considered as a matching problem where the query is compared with every element in the support set [ 43 , 44 , 45 , 46 ] . Typical approaches use a Siamese network [ 47 ] and compare the embedded query/support features using a similarity measure such as cosine similarity [ 43 ] , squared Euclidean distance [ 48 ] , or even a learned function [ 49 ] . Our task can also be formulated as a few-shots problem , where our memory bank acts as the support set . This connection helps us with the choice of similarity function , albeit we are dealing with a million times more pointwise comparisons . Video Object Segmentation Early VOS methods [ 4 , 5 , 50 ] employ online first-frame finetuning which is very slow in inference and have been gradually phased out . Faster approaches have been proposed such as a more efficient online learning algorithm [ 8 , 6 , 7 ] , MRF graph inference [ 51 ] , temporal CNN [ 52 ] , capsule routing [ 53 ] , tracking [ 11 , 13 , 15 , 54 , 55 , 56 , 57 ] , embedding learning [ 10 , 58 , 59 ] and space-time matching [ 17 , 18 , 19 , 20 ] . Embedding learning bears a high similarity to space-time matching , both attempting to learn a deep feature representation of an object that remains consistent across a video . Usually embedding learning methods are more constrained [ 10 , 58 ] , adopting local search window and hard one-to-one matching . We are particularly interested in the class of Space-Time Memory networks ( STM ) [ 18 ] which are the backbone for many follow-up state-of-the-art VOS methods . STM constructs a memory bank for each object in the video , and matches every query frame to the memory bank to perform “ memory readout ” . Newly inferred frames can be added to the memory , and then the algorithm propagates forward in time . Derivatives either apply STM at other tasks [ 21 , 60 ] , improve the training data or augmentation policy [ 21 , 22 ] , augment the memory readout process [ 16 , 21 , 22 , 24 , 28 ] , use optical flow [ 29 ] , or reduce the size of the memory bank by limiting its growth [ 23 , 30 ] . MAST [ 61 ] is an adjacent research that focused on unsupervised learning with a photometric reconstruction loss . Without the input mask , they use Siamese networks on RGB images to build the correspondence out of necessity . In this work , we deliberately build such connections and establish that building correspondences between images is a better choice , even when input masks are available , rather than a concession . We propose to overhaul STM into STCN where the construction of affinity is redefined to be between frames only . We also take a close look at the similarity function , which has always been the dot product in all STM variants , make changes and comparisons according to our findings . The resultant framework is both faster and better while still principled . STCN is even fundamentally simpler than STM , and we hope that STCN can be adopted as the new and efficient backbone for future works . 3 Space-Time Correspondence Networks ( STCN ) . Given a video sequence and the first-frame annotation , we process the frames sequentially and maintain a memory bank of features . For each query frame , we extract a key feature which is compared with the keys in the memory bank , and retrieve corresponding value features from memory using key affinities as in STM [ 18 ] . 3.1 Feature Extraction . Figure 1 illustrates the overall flow of STCN . While STM [ 18 ] parameterizes a Query Encoder ( image as input ) and a Memory Encoder ( image and mask as input ) with two ResNet50 [ 62 ] , we instead construct a Key Encoder ( image as input ) and a Value Encoder ( image and mask as input ) with a ResNet50 and a ResNet18 respectively . Thus , unlike in STM [ 18 ] , the key features ( and thus the resultant affinity ) can be extracted independently without the mask , computed only once for each frame , and symmetric between memory and query.2 The rationales are 1 ) Correspondences ( key features ) are more difficult to extract than value , hence a deeper network , and 2 ) Correspondences should exist between frames in a video , and there is little reason to introduce the mask as a distraction . From another perspective , we are using a Siamese structure [ 47 ] which is widely adopted in few-shots learning [ 63 , 49 ] for computing the key features , as if our memory bank is the few-shots support set . As the key features are independent of the mask , we can reuse the “ query key ” later as a “ memory key ” if we decide to turn the query frame into a memory frame during propagation ( strategy to be discussed in Section 3.3 ) . This means the key encoder is used exactly once per image in the entire process , despite the two appearances in Figure 1 ( which is for brevity ) . 2That is , matching between two points does not depend on whether they are query or memory points ( not true in STM [ 18 ] as they are from different encoders ) . Architecture . Following the STM practice [ 18 ] , we take res4 features with stride 16 from the base ResNets as our backbone features and discard res5 . A 3×3 convolutional layer without non-linearity is used as a projection head from the backbone feature to either the key space ( Ck dimensional ) or the value space ( Cv dimensional ) . We set Cv to be 512 following STM and discuss the choice of Ck in Section 4.1 . Feature reuse . As seen from Figure 1 , both the key encoder and the value encoder are processing the same frame , albeit with different inputs . It is natural to reuse features from the key encoder ( with fewer inputs and a deeper network ) at the value encoder . To avoid bloating the feature dimensions and for simplicity , we concatenate the last layer features from both encoders ( before the projection head ) and process them with two ResBlocks [ 62 ] and a CBAM block3 [ 64 ] as the final value output . 3.2 Memory Reading and Decoding . Given T memory frames and a query frame , the feature extraction step would generate the followings : memory key kM ∈ RCk×THW , memory value vM ∈ RCv×THW , and query key kQ ∈ RCk×HW , where H and W are ( stride 16 ) spatial dimensions . Then , for any similarity measure c : RCk×RCk → R , we can compute the pairwise affinity matrix S and the softmax-normalized affinity matrix W , where S , W ∈ RTHW×HW with : Sij = c ( k M i , k Q j ) Wij = exp ( Sij ) ∑ n ( exp ( Snj ) ) , ( 1 ) where ki denotes the feature vector at the i-th position . The similarities are normalized by √ Ck as in standard practice [ 18 , 33 ] and is not shown for brevity . In STM [ 18 ] , the dot product is used as c. Memory reading regularization like KMN [ 22 ] or top-k filtering [ 21 ] can be applied at this step . With the normalized affinity matrix W , the aggregated readout feature vQ ∈ RCv×HW for the query frame can be computed as a weighted sum of the memory features with an efficient matrix multiplication : vQ = vMW , ( 2 ) which is then passed to the decoder for mask generation . In the case of multi-object segmentation , only Equation 2 has to be repeated as W is defined between image features only , and thus is the same for different objects . In the case of STM [ 18 ] , W must be recomputed instead . Detailed running time analysis can be found in Section 6.2 . Decoder . Our decoder structure stays close to that of the STM [ 18 ] as it is not the focus of this paper . Features are processed and upsampled at a scale of two gradually with higher-resolution features from the key encoder incorporated using skip-connections . The final layer of the decoder produces a stride 4 mask which is bilinearly upsampled to the original resolution . In the case of multiple objects , soft aggregation [ 18 ] of the output masks is used .

Summary: This work proposes a space-time correspondence network which improves performance and speed over the standard STM network by finding pixel-level correspondences in individual frames rather than frames+segmentations. Furthermore, this work shows that the dot-product, which is commonly used in STM is suboptimal and they propose an efficient L2 similarity computation which leads to improved correspondences throughout a video. It is evaluated on two standard VOS benchmarks - Youtube-VOS and DAVIS-2017 - and there are extensive analyses and ablations which give insights into the proposed method.