Test-time adaptation with slot-centric models
[Figure]
Abstract
We consider the problem of segmenting scenes into constituent objects.
Current supervised visual detectors, though impressive within their training distribution, often fail to segment out-of-distribution scenes.
Recent test-time adaptation methods use auxiliary self-supervised losses to adapt the network parameters to each test example independently and have shown promising results towards generalization outside the training distribution for the task of image classiﬁcation.
In our work, we ﬁnd evidence that these losses can be insufﬁcient for instance segmentation tasks, without also considering architectural inductive biases.
For image segmentation, recent slot-centric generative models break such dependence on supervision by attempting to segment scenes into entities in a self-supervised manner by reconstructing pixels.
Drawing upon these two lines of work, we propose Slot-TTA, a semi-supervised instance segmentation model equipped with a slotcentric image rendering component, that is adapted per scene at test time through gradient descent on reconstruction or novel view synthesis objectives.
We show that test-time adaptation greatly improves segmentation in out-of-distribution scenes.
We evaluate Slot-TTA in scene segmentation benchmarks and show substantial out-of-distribution performance improvements against state-of- the-art supervised feed-forward detectors and self-supervised domain adaptation models. Please ﬁnd the full version of our paper at: https://arxiv.org/abs/2203.
Introduction
While signiﬁcant progress has been made in machine scene perception and segmentation within the last decade, object detectors continue to generalize poorly outside their training distribution [12, 22].
Consider the extremely cluttered scenes shown in Figure 1.
We can intuitively reason about meaningful parts that this scene could be broken into.
Yet, a state-of-the-art 2D detector [6] trained to segment similar objects in less cluttered scenes (5-7 object instances) struggles with this decomposition.
This lack of generalization requires us to build systems that can robustly adapt to such changes in distribution.
Test-time adaptation (TTA) [13, 37, 40] describes a setting where a model adapts to changes in distribution at test-time, at the cost of additional computation.
In recent years, a variety of methods based on TTA have been proposed, focusing on few-shot adaptation [32] where the network is given access to a few labelled examples, or unsupervised domain adaptation (UDA) [43] where the network is given access to many unlabelled examples from the new distribution.
Of particular relevance is a speciﬁc UDA setting where model parameters are adapted independently to each unlabelled example in the test-set.
This setting has been previously referred to as single-example UDA, and here we also refer to it as slow inference since it is similar to a human taking more time to parse a difﬁcult example.
Existing approaches for this setting typically devise a self-supervised loss that aligns well with the task of image classiﬁcation and then optimize this loss during test-time adaptation [37, 11, 1, 20].
However, despite their success for image classiﬁcation, these approaches do not provide adequate support for other scene understanding tasks, and in particular scene segmentation, as we showcase in
Section 3.1.One potentially important aspect to supporting TTA for other scene understanding tasks is the inductive bias of the underlying architecture.
In the context of segmentation, there has been a lot of recent development in building models that segment scenes into entities in an unsupervised way by optimizing a reconstruction objective [10, 18, 38, 15, 8, 28, 44].
These methods differ in details but share the notion of incorporating a ﬁxed set of entities, also known as slots or object ﬁles .
Each slot extracts information about a single entity during encoding, and is “synthesized” back to the input domain during decoding.
Their ability to distinguish visual objects at a representation level makes them a particularly promising candidate for TTA for segmentation tasks.
In light of the above, we propose Slot-centric Test-time adaptation (Slot-TTA), a semi-supervised approach that combines Slot Attention [28] (in the 2D image) or Object Scene Representation Transformer [34] (in multi-view image setting) with a supervised segmentation loss to enable it to leverage instance-level image.
Slot-TTA is trained jointly to synthesize and segment scenes.
At test time, the model adapts without supervision to a single test sample by optimizing the self-supervised objective alone.
Different from fully-unsupervised object-centric generative models, Slot-TTA uses annotations at training time to help it develop the notion of what an object is, which lets it scale to more complex visual settings.
Different from existing TTA methods, Slot-TTA uses a slot-centric architecture and self-supervised synthesis loss that better aligns with the task of segmentation.
Different from state-of-the-art detectors, Slot-TTA is equipped with reconstruction feedback that allows it to adapt at test time without supervision, i.e. without using additional annotated data.
Indeed, we show that test-time adaptation via image synthesis in Slot-TTA enables successfully parsing completely unfamiliar scenes composed of familiar entities.
We test Slot-TTA’s instance segmentation ability on the two datasets: MultiShapeNet-Hard [35] and Multi-Shape.
We evaluate Slot-TTA’s ability to parse out-of-distribution scenes and compare it against state-of-the-art entity-centric generative models [34] and supervised visual detectors [6] trained with labeled data to segment objects.
We show improvements over all baselines in Slot-TTA ability to segment novel scenes.
Additionally, we ablate different design choices of Slot-TTA.
Method
We consider Slot-TTA in two settings: (i) auto-encoding of images, and (ii) novel view synthesis.
For (i), we use the architecture proposed by Slot Attention [28], where a 2D broadcast decoder [41] is used to render the input view.
For (ii), we use OSRT’s architecture which combines the Slot
Attention bottleneck with the geometry-free backbone of SRT [35] to perform object-centric novel view synthesis.
Training for joint segmentation and reconstruction We train all the parameters of our model to jointly optimize image reconstruction or novel view image synthesis objectives and the task segmentation objective over all the n examples in the training set, where x represents the input scene and y the segmentation labels:
[Equation]
For reconstruction, we minimize the mean squared error between the predicted and ground truth RGB images.
For segmentation, we supervise the alpha masks a i of each slot as provided by the decoders.
We use Hungarian matching [27] to associate the ground truth masks with the predicted masks, and upon association we apply a categorical cross-entropy loss l seg .
We weight the segmentation and reconstruction loss by λ s and λ r .
Test-time adaptation In this work, we refer to a single forward pass through our trained model without any test-time adaptation as fast inference (same as regular inference).
We call the process of test-time adapting the model on each example independently slow inference , using only the reconstruction objective of Eq.
We use this terminology to emphasize that the only difference between both settings is the added computation time which results in an effective speed difference between the two inference schemes.
We adapt only the encoder parameters θ enc in our model, which we found to improve results compared to adapting the entire model as shown in our supplementary
Section D.1.
We train for 150 steps per example using the Adam optimizer [25].
For further details on our model architectures and model ﬁgure please refer to supplementary Section B
Experiments
We test Slot-TTA capability for segmenting posed multi-view and single-view RGB images.
We use Adjusted Random Index (ARI) as our evaluation metric for segmentation accuracy [31].
Our experiments aim to answer the following questions: (i) How does Slot-TTA compare against state-ofthe-art 2D segmentation models [6]?
(ii) How does slow inference through reconstruction feedback affect segmentation accuracy in Slot-TTA and its variants?
(iii) How much does supervision during training contribute to segmentation performance?
3.1 Segmenting RGB images in multi-view scenes
Dataset We evaluate Slot-TTA on the MultiShapeNet (MSN) dataset from SRT [35].
The dataset is constructed by rendering 51K ShapeNet objects using Kubric [16] against 382 HDR backgrounds so that there is no overlap of objects between the train and test sets.
Further train and tests sets differ in the number of objects present: scenes with 5-7 object instances are in the training set and scenes with 16-30 objects are in the test set.
For more details please refer to supplementary Section A.
[Table]
Baselines We compare to three baselines: (i) Mask2Former [6], a state-of-the-art 2D image segmentor which adapts detection transformers [4] to image segmentation by using multiscale segmentation decoders with masked attention.
(ii) Mask2Former+BYOL which combines the segmentation model [6] with test time adaptation using BYOL self-supervised loss [1].
(iii) Mask2Former+Recon which combines the segmentation model [6] with rendering submodules and image reconstruction loss for test-time adaptation.
Results We show quantitative segmentation results of our model and baselines on target camera viewpoints in Table 1 and qualitative TTA results in Figure 2.
In Slot-TTA-w/o supervision, instead of training jointly for reconstruction and segmentation, we train using only cross-view image synthesis, similar to OSRT [34].
[Figure]
It can be observed that: (i) Slot-TTA-Slow outperforms the feedforward Mask2Former-Fast, especially for out-of-distribution scenes; (ii) adding self-supervised losses of SOTA image classiﬁcation methods [1] to Mask2Former (eg.
Mask2Former+BYOL) does not sufﬁce to adapt them effectively at test time and (iii) Slot-TTA without supervision, which is identical to OSRT [34] is not competitive with supervised models for object segmentation.
3.2 Segmenting single-view RGB images
As a proof of concept, in this section, we test our model and the baseline Mask2Former in segmenting single RGB images comprised of multiple samples from ﬁve shapes of distinct colors, organized in heavily occluded conﬁgurations, a dataset we create and we call Multi-Shape.
Our training set consists of images with 3-5 object instances, while the test set consists of images with 10-16 object instances.
For this setting, we report the ARI scores for the foreground objects only, since in this dataset the background occupies a large image area and a method that assigns most pixels to background already achieves a very high ARI.
We ﬁnd the performance accuracy ordering of the methods to be the same.
As can be seen in Table 2, before TTA Mask2former and Mask2former+Recon outperform our method.
After TTA, our method signiﬁcantly outperforms the baselines.
Method in-dist (3-5 instances) out-of-dist(10-16 instances)
Fast Infer.
Slow Infer.
Fast Infer.
Slow Infer.
Mask2Former 0.
N/A 0.
N/A
Mask2Former+Recon 0.
[Table]
Please refer to Section D.1 and Section D.2 in the supplementary for qualitative results and ablations in our mult-view RGB and single-view RGB settings respectively.
Conclusion
In this work we show that the architectural choices found in unsupervised object discovery methods such as Slot Attention, could be very helpful in test-time adaptation.
Additionally we show sufﬁcient evidence in our work that a future version of Slot-TTA could potentially compete with state-of-the-art segmentation methods when allowed to do test-time adaptation.
B Method
Slot Attention Current state-of-the-art detectors and segmentors instantiate slots (i.e. the query vectors) from 2D visual feature maps [4].
Most works use iterative cross-attention (features to slots) and self-attention (slot-to-slots) operations [4] to map a set of N input feature vectors to a set of K slot vectors.
Attention-based competition amongst slots and iterative routing popularized in [15, 28] encourages a single location in the input to be assigned to a unique slot vector.
Given a visual scene encoded as a set of feature vectors M ∈ R N × C and K randomly initialized slots sampled from a multivariate Gaussian distribution with a diagonal covariance S ∼ N ( µ, Diag [Equation]
∈ R K × D , where µ, σ ∈ R C are learnable parameters of the Gaussian, Slot Attention [28] computes an attention map a between the feature map M and the slots S :
[Equation]
k , q , and v are learnable linear transformations that map inputs and slots to a common dimension
D .
The softmax normalization over slots ensures competition amongst them to attend to a speciﬁc feature vector in M .
Updates to the slots are computed based on the input features they attend to: updates
[Equation]
where
[Equation]
[Figure]
which are then fed into a GRU [7]: S =
GRU( state = S, input = updates ) .
We iterate 3 times over equations 2 and 3.
For detailed description, please refer to [28].
We ﬁrst describe the encoders and decoders that form the foundation of Slot-TTA for each modality.
Further we detail how we train Slot-TTA and perform test time adaptation through slow inference.
B.0.1 Encoding and Decoding Backbones
Posed multi-view 2D RGB images As shown in Figure 5, Slot-TTA builds upon the architecture of OSRT [34], which is an object-centric, geometry-free novel view synthesis method.
Given a set of multi-view RGB images as input, a CNN encodes each input image
I i into a feature grid, which is then ﬂattened into a set of tokens with camera pose and ray direction information added in each of the tokens, similar to SRT [35].
These are then encoded into a set of latent features using a transformer [39] Enc with multiple self-attention blocks z = Enc(CNN(I i )) .
The latent features z are then mapped into a set of slots S using Slot Attention (Section B).
For decoding, we adopt the spatial broadcast decoder [41] formulation, where a render MLP takes as input the slot vector
S k and the pixel location p parameterized by the camera position and the ray direction pointing to the pixel to be decoded.
It outputs an RGB color c k and an unnormalized alpha score a k for each pixel location c k , a k
= Dec( p, S k ) .
The a k ’s are normalized using a Softmax and used as weights to aggregate the predicted RGB values c k for each slot.
Single-view 2D RGB images For this setting, Slot-TTA uses a ResNet-18 [21] to encode the input
RGB image into a feature grid.
We then add positional vectors to the feature grid and map to a set of slot vectors using Slot Attention.
Similar to the multi-view setting, each slot vector is decoded to the
RGB image and an alpha mask using an MLP renderer.
We parameterize pixel location p as ( x, y ) points on the grid instead of camera position as the above setting.
C Related Work
Entity-centric generative models for scene decomposition Entity-centric (or object-centric ) models use architectural inductive biases to represent perceptual inputs, such as an observation of a visual scene, in terms of separate object variables, often referred to as slots or object ﬁles [19, 33, 26, 9, 14, 23, 3, 17, 42, 30].
Prominent examples of such models include MONet [3],
GENESIS [9], IODINE [17], and Slot Attention (SA) [28], which are trained in a fully-unsupervised setting via a simple auto-encoding objective.
Object representations and scene decomposition emerge via the inductive bias of the model architecture (and in some cases, additional regularizers).
However, without any form of supervision, scene decompositions can be ambiguous, which is particularly challenging for complex real-world scenes or in the presence of complicated textures.
In Slot-TTA, we aid the competition mechanism in SA to address this issue by jointly training with a supervised segmentation loss.
OSRT [34] is a cross-view geometry-free encoder-decoder method, that segments an image into objects through reconstructing novel viewpoints.
OSRT combines SA with SRT [35], a view synthesis model that uses transformer encoder and decoders to fuse information across views, as well as the camera pose, without any explicit 3D information.
Our multi-view RGB Slot-TTA builds upon their architecture.
Test-time adaptation In test-time adaptation
, model parameters are updated at test-time to better generalize to the distribution shift.
In recent years, there has been signiﬁcant development in this direction.
Methods such as pseudo labelling and entropy minimization [36, 40, 2] have demonstrated that supervising the model using its own conﬁdent predictions could help improve its accuracy.
Adaptive BatchNorm methods [24, 5] have shown that updating the BatchNorm parameters using the new examples can help adaptation.
Despite these successes, these methods by deﬁnition are data inefﬁcient as they require conﬁdent predictions or a batch of examples to adapt.
Self-supervised learning (SSL) [37, 1, 11] based methods on the other hand, have empirically shown to be data efﬁcient.
During training, they jointly train using the task and SSL loss, and during test-time, they train only using the SSL loss.
All of the methods in the SSL setting thus far focus on the task of classiﬁcation and mainly differ in terms of the SSL loss used.
For example TTT [37] uses rotation angle prediction as their SSL loss, MT3 [1] uses a BYOL [20] loss and TTT-MAE [11] uses Masked autoencoding loss [29].
In our work, we show that these losses do not generalize to segmentation, and how we might need speciﬁc architectural biases to close the gap.
D Additional Experiments D.1 Segmenting RGB images in multi-view scenes Method
[Table]
Slot-TTA-tta_Norm_param
[Table]
Slot-TTA-tta_Slot_param [Table] Slot-TTA w/
[Table]
Weighted_Sample
[Table]
Slot-TTA [Table] [Table] We conduct various ablations of Slot-TTA in Table 3.
In Figure 6, we show additional qualitative results comparing Slot-TTA-Fast and Slot-TTA-Slow.
(i)
We ablate different decoder choices in the topmost section where instead of using the broadcast decoder we use the Scene representation transformer (SRT) decoder [35] which we refer to as Slot-TTASRT_Decoder or the SlotMixer decoder [34], referred to as Slot-TTA-SlotMixer_Decoder .
(ii) We ablate what parameters to adapt at test time.
As it’s unclear since TENT [40] optimizes BatchNorm or LayerNorm parameters, but TTT [37] optimizes the shared parameters between the SSL and the task-speciﬁc branch, which in our case will be all the parameters in the network.
In Table 3, Slot-TTA-tta_All_param is when we adapt all the network parameters, Slot-TTAtta_Norm_param adapts only the Layer or BatchNorm parameters and Slot-TTA-tta_Slot_param adapts only the learnable slot embeddings.
We ﬁnd that optimizing only the encoder parameters works the best for our setting.
(iii) Further , we ablate error-conditioned pixel sampling where Slot-TTA w/o Weighted_Sample refers to our model that uses uniform sampling instead of the error weighted sampling.
D.2 Segmenting single-view RGB images
In Figure 7 we qualitatively compare Slot-TTA-Slow with Slot-TTA-Fast.
We show that slow inference can help discover objects missed by Slot-TTA.
We also show some failure cases where slow inference could override the object-centric bottleneck to achieve higher reconstruction accuracy.
[Figure]