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 paper presents the Space-Time Correspondence Network (STCN) to tackle the semi-supervised video object segmentation. STCN improves previous STM by building affinity based on RGB-only input and reused key features. Moreover, the author further investigated the computation of similarity and find that negative squared Euclidean distance leads to better performance. Experiments on multiple datasets are conducted to demonstrate the effectiveness of the proposed method.