Multi-Scale Fusion for Object Representation

Assume there is an input image or video frame $\bm{T}\in\mathbb{R}^{h_{0}\times w_{0}\times c_{0}}$, where typically $h_{0}=w_{0}=128$ and $c_{0}=3$.

Initially, the primary encoder encodes the input into a feature map. $\bm{\phi}_{\mathrm{e}}:\bm{T}\rightarrow\bm{F}$, where $\bm{F}\in\mathbb{R}^{h\times w\times c}$, and $\bm{\phi}_{e}$ is a CNN (He et al., 2016) or ViT (Caron et al., 2021) module. Typically, $h=w=32$ and $c=256$.

Then, the Slot Attention (SA) (Locatello et al., 2020) is used to aggregate the dense feature map $\bm{F}$ into sparse object-level feature vectors, i.e., slots $\bm{S}\in\mathbb{R}^{s\times c}$, along with the byproduct, i.e., segmentation masks $\bm{M}\in\mathbb{R}^{s\times h\times w}$, under query vectors $\bm{S}_{0}$. $\bm{\phi}_{\mathrm{sa}}:(\bm{S}_{0},\bm{F})\rightarrow(\bm{S},\bm{M})$, which is basically an iterative Query-Key-Value attention (Bahdanau et al., 2015) mechanism with $\bm{S}_{0}$ as the query and $\bm{F}$ as the key and value. Typically, $s$ is a predefined number, usually the maximum number of objects plus the background in an image or video sample of a dataset.

Meanwhile, pretrained VAE (Im Im et al., 2017; Van Den Oord et al., 2017) converts $\bm{T}$ into continuous intermediate representation. $\bm{\phi}_{\mathrm{e}}^{\mathrm{v}}:\bm{T}\rightarrow\bm{Z}$, where $\bm{Z}\in\mathbb{R}^{h\times w\times c}$ and typically $h=w=32$ and $c=256$. This continuous representation is quantized into discrete representation by matching continuous representation with codebook and selecting matched codes to form discrete representation. $f_{\mathrm{q}}=f_{\mathrm{m}}\circ f_{\mathrm{s}}:(\bm{Z},\bm{C})\rightarrow(\bm{P},\bm{I},\bm{X})$, where codebook $\bm{C}\in\mathbb{R}^{m\times c}$, matching probabilities $\bm{P}\in\mathbb{R}^{h\times w\times m}$, matched indexes $\bm{I}\in\mathbb{R}^{h\times w}$ and discrete representation $\bm{X}\in\mathbb{R}^{h\times w\times c}$. The matching $f_{\mathrm{m}}$ and selection $f_{\mathrm{s}}$ are implemented as below, respectively

where each super-pixel $\bm{P}^{(i,j,:)}\in\mathbb{R}^{m}$ means the matching probabilities of the $m$ template features in $C$. So the index of the maximum element in $\bm{P}^{(i,j,:)}$, i.e., $\bm{I}^{(i,j)}$ is the index of the most matched code, which can quantize $\bm{Z}^{(i,j,:)}$ into $\bm{X}^{(i,j,:)}:=\bm{C}^{(\bm{I}^{(i,j)},:)}$.

Lastly, with $\bm{S}$ as the condition, the primary decoder reconstructs the input guided by VAE discrete intermediate representation conditioned on $\bm{S}$, which drives every slot to aggregate as much object information as possible. $\bm{\phi}_{\mathrm{d}}:(\bm{X},\bm{S})\rightarrow\bm{X}^{\prime}$, where $\bm{X}^{\prime}\in\mathbb{R}^{h\times w\times c}$. For transformer-based OCL, a transformer decoder (Vaswani et al., 2017) with internal masking performs the reconstruction as classification. For diffusion-based OCL, a diffusion model, typically a conditional UNet (Rombach et al., 2022), carries out the reconstruction as regression of the noise added to $\bm{X}$.

For VAE pretraining, there is also VAE decoding from discrete representation to the input $\bm{\phi}_{\mathrm{d}}^{\mathrm{v}}:\bm{X}\rightarrow\bm{T}^{\prime}$, where $\bm{T}^{\prime}\in\mathbb{R}^{h_{0}\times w_{0}\times c_{0}}$. And the supervision signal comes from Mean-Squared Error (MSE) between $\bm{T}$ and $\bm{T}^{\prime}$. For OCL training, the pretrained VAE module is frozen, and no VAE decoding occurs. Remaining parts are trained under the supervision signal of either Cross Entropy (CE) between $\bm{X}$ and $\bm{X}^{\prime}$, or Mean-Squared Error (MSE) between the ground-truth noise added to $\bm{X}$ and the reconstructed noise.

There can be objects of different sizes in images or videos, while models usually have best-at patterns scales due to non-uniform data distributions. This makes image pyramid (Adelson et al., 1984) an effective technique, where the input is resized into different sizes and processed in parallel. Such multi-scale representations make the pattern scale distribution more uniform for training, and ensures different sized objects to fall within models’ comfort zone for testing.

Input image or video frame $\bm{T}_{1}:=\bm{T}$ is down-sampled with typical bi-linear interpolation into $\bm{T}_{n}\in\mathbb{R}^{\frac{h_{0}}{2^{n-1}}\times\frac{w_{0}}{2^{n-1}}\times c_{0}}$, where $n=1...N$ means different spatial scales and we use $N=3$.

Multi-scale VAE encoding is $\bm{\phi}_{\mathrm{e}}^{\mathrm{v}}:\bm{T}_{n}\rightarrow\bm{Z}_{n}$, where $\bm{Z}_{n}\in\mathbb{R}^{\frac{h}{2^{n-1}}\times\frac{w}{2^{n-1}}\times c}$ and typically $h=w=32$ and $c=256$.

Multi-scale VAE quantization is $f_{\mathrm{q}}=f_{\mathrm{m}}\circ f_{\mathrm{s}}:(\bm{Z}_{n},\bm{C})\rightarrow(\bm{P}_{n},\bm{I}_{n},\bm{X}_{n})$, where $\bm{P}_{n}\in\mathbb{R}^{\frac{h}{2^{n-1}}\times\frac{w}{2^{n-1}}\times m}$, $\bm{I}_{n}\in\mathbb{R}^{\frac{h}{2^{n-1}}\times\frac{w}{2^{n-1}}}$ and $\bm{X}_{n}\in\mathbb{R}^{\frac{h}{2^{n-1}}\times\frac{w}{2^{n-1}}\times c}$. Here codebook $\bm{C}$ is shared over all $N$ scales so as to facilitate scale-invariant pattern learning that differentiates among different object types. But this strong inductive bias can lose scale-variant information that differentiates among different object instances of the same type; Besides, no information exchange among the scales is a kind of information waste – The object pattern that is well represented in one scale should augment its poorly represented patterns in the other scales. We address this in the next subsection.

Multi-scale VAE decoding is $\bm{\phi}_{\mathrm{d}}^{\mathrm{v}}:\bm{X}_{n}\rightarrow\bm{T}_{n}^{\prime}$, where $\bm{T}_{n}^{\prime}\in\mathbb{R}^{\frac{h_{0}}{2^{n-1}}\times\frac{w_{0}}{2^{n-1}}\times c}$.

During VAE pretraining, there are $N$ sets of reconstruction and quantization losses. During OCL training, only scale $n=1$ is used as guidance, thus the supervision is no different from the former subsection. Although multi-scale guidance can be used, it is always harmful in practice.

By following such recipe, hopefully at least one of the $N$ scales of representation provides high-quality representation to some sized objects.

Note: The original VQ-VAE has only $n=1$ scale, which is a special case of our method.

Based on the above multi-scale quantization, we further devise our novel fusion among the scales. Inter-scale fusion utilizes scale-invariance to augment poorly represented features in some scales with the well-represented features in other scales; Intra-scale fusion supplements with scale-variant details in the corresponding scale. Scale-invariant patterns differentiate among different object types, while scale-variance distinguishes among different object instances of the same type.

Firstly, we project VAE continuous representation of a certain scale $\bm{Z}_{n}$ into two copies, $\bm{Z}_{n}^{\mathrm{si}}$ for scale-invariance and $\bm{Z}_{n}^{\mathrm{sv}}$ for scale-variance, respectively. These two will finally be combined back, thus we employ the invertible project-up-project-down design proposed in Zhao et al. (2024) to maintain the inductive biases added in between.

where $\mathrm{pinv}(W)$ is pseudo-inverse of the project-down matrix in Eq. 5 for project-up here; $\mathrm{chunk}_{c,2}(\cdot)$ splits the tensor into two along the channel dimension.

Secondly, for the scale-variant continuous representation, we conduct quantization by a codebook specific to the current scale. $f_{\mathrm{q}}=f_{\mathrm{m}}\circ f_{\mathrm{s}}:(\bm{Z}_{n}^{\mathrm{sv}},\bm{C}_{n}^{\mathrm{sv}})\rightarrow(\bm{P}_{n}^{\mathrm{sv}},\bm{I}_{n}^{\mathrm{sv}},\bm{X}_{n}^{\mathrm{sv}})$. These non-sharing codebooks enforces scale-variant representation learning.

Meanwhile, for the scale-invariant continuous representation, we conduct inter-scale fusion over scale-invariant quantization of all scales. $f_{\mathrm{rf}}:\{(\bm{Z}_{n}^{\mathrm{si}},\bm{C}^{\mathrm{si}})\}\rightarrow\{(\dot{\bm{P}}_{n}^{\mathrm{si}},\dot{\bm{I}}_{n}^{\mathrm{si}},\dot{\bm{X}}_{n}^{\mathrm{si}})\}$, where $n=1...N$. The shared codebook enforces scale-invariant representation learning.

Now that all scales $\{\bm{X}_{n}^{\mathrm{si}}\}$ consist of codes in a shared codebook $\bm{C}^{\mathrm{si}}$, matching probabilities $\{\bm{P}_{n}^{\mathrm{si}}\}$ for the same (super-)pixel region under different scales indicate how confident the model is about the representation quality. We utilize this to augment the low-quality representation in some scales with the high-quality in other scales. Take scales $n-1$, $n$ and $n+1$ as an example:

1. 1.

   Given any super-pixel $\bm{z}_{n}^{\mathrm{si}}:=\bm{Z}_{n}^{\mathrm{si},(i,j,:)}\in\mathbb{R}^{c}$ in any scale $n$, find the corresponding region/super-pixel(s) in other scales, and mean them along all spatial dimensions as voting. For scale $n-1$, we choose $\bm{Z}_{n-1}^{\mathrm{si},(2i:2i+2,2j:2j+2,:)}\in\mathbb{R}^{2\times 2\times c}$ and mean it into a vector $\bm{z}_{n-1}^{\mathrm{si}}\in\mathbb{R}^{c}$; For scale $n+1$, we choose $\bm{z}_{n+1}^{\mathrm{si}}:=\bm{Z}_{n+1}^{\mathrm{si},(\lfloor\frac{i}{2}\rfloor,\lfloor\frac{j}{2}\rfloor,:)}\in\mathbb{R}^{c}$.

2. 2.

   Match $\{\bm{z}_{n-1}^{\mathrm{si}},\bm{z}_{n}^{\mathrm{si}},\bm{z}_{n+1}^{\mathrm{si}}\}$ respectively with codebook $\bm{C}^{\mathrm{si}}$ as in Eq. 1, then we have corresponding matching probabilities $\{\bm{p}_{n-1}^{\mathrm{si}},\bm{p}_{n}^{\mathrm{si}},\bm{p}_{n+1}^{\mathrm{si}}\}\in\{\mathbb{R}^{m}\}$. These probabilities can be taken as the super-pixel quality of the same region in scales $\{n-1,n,n+1\}$.

3. 3.

   Select the most matched code with the highest probability among all scales as the quantization $\dot{\bm{X}}_{n}^{\mathrm{si},(i,j,:)}$ of the super-pixel $\bm{Z}_{n}^{\mathrm{si},(i,j,:)}$. Specifically, for scales $\{n-1,n,n+1\}$, the most matched code index is

   $$\dot{\bm{I}}_{n}^{\mathrm{si},(i,j)}=\mathrm{argmax}_{m}(\mathrm{max}_{n}(\{\bm{p}_{n-1}^{\mathrm{si}},\bm{p}_{n}^{\mathrm{si}},\bm{p}_{n+1}^{\mathrm{si}}\}))$$

   (4)

   with which we can determine $\dot{\bm{X}}_{n}^{\mathrm{si},(i,j,:)}$ as in Eq. 2.

Lastly, for each scale $n$, conduct the intra-scale fusion between the fused scale-invariant representation and scale-variant representation. $f_{\mathrm{af}}:(\dot{\bm{X}}_{n}^{\mathrm{si}},\bm{X}_{n}^{\mathrm{sv}})\rightarrow\bm{X}_{n}$. This combines the augmented scale-invariant and scale-variant copies back into the final discrete intermediate representation.

where $\mathrm{concat}_{c}(\cdot)$ is channel concatenation; project-down matrix $W\in\mathbb{R}^{2c\times c}$ is trainable.

The above processes can be parallelized using PyTorch, as shown in Algo. 1 pseudo code.

Codebook-related parameters. The basis methods use a codebook of size $(m,c)=4096\times 256$, which amounts to $1.048$M. We use (1) invertiable projection to project the continuous representation into two copies (and project them back), up to $c\times 2c=2^{17}$; (2) four groups of codebooks of size $(m^{0.5},c)=(64,256)$, one shared for $N=3$ scale-invariant scales, three specifically for $N=3$ scale-variant scales, up to $m^{0.5}\times c\times(1+N)=2^{16}$. Thus our codebook-related parameters are $0.196$M – we have $81.25\%$ fewer codebook parameters than the basis methods.

Computation and memory consumption. The exact calculation relates with the VAE encoder/decoder layers, like Conv2d, GroupNorm, Mish, etc., thus has too many details. But due to the $O(n^{2})$ complexity nature of convolution layers, we can quickly get the estimation in ratio. Denote the computation of the original VAE as unit, then our $N=3$ scales’ computations would be $1$, $\frac{1}{4}$ and $\frac{1}{16}$, since the latter two has $2\times$ and $4\times$ down sampling rates. Thus in VAE pretraining, our technique requires roughly $31.25\%$ more computation. The memory consumption increase is similar.

To evaluate the quality of OCL object representations, i.e., slots, we use the accuracy of OCL’s byproduct (unsupervised) segmentation as an intuitive measurement. Metrics including ARI (Adjusted Rand Index)¹¹1https://scikit-learn.org/stable/modules/generated/sklearn.metrics.adjusted_rand_score.html, ARI_fg (foreground), mIoU (mean Intersection-over-Union)²²2https://scikit-learn.org/stable/modules/generated/sklearn.metrics.jaccard_score.html and mBO (mean Best Overlap) (Caron et al., 2021) are used to measure OCL’s byproduct segmentation accuracy. Datasets are either synthetic images, i.e., ClevrTex³³3https://www.robots.ox.ac.uk/~vgg/data/clevrtex/; or real-world images, i.e., COCO⁴⁴4https://cocodataset.org/#panoptic-2020 and VOC⁵⁵5http://host.robots.ox.ac.uk/pascal/VOC/voc2012/index.html; or synthetic videos, i.e., MOVi-C/D/E⁶⁶6https://github.com/google-research/kubric/tree/main/challenges/movi.

To evaluate how MSF boosts OCL, we use transformer-based classics, SLATE (Singh et al., 2022a) for images and STEVE (Singh et al., 2022c) for videos; use diffusion-based state-of-the-arts, SlotDiffusion (Wu et al., 2023b) for images. The foundation-based DINOSAUR (Seitzer et al., 2023) and competitors SysBinder (Singh et al., 2022b) and GDR (Zhao et al., 2024) are compared too. All primary encoders are unified as DINO (Caron et al., 2021) to form strong baselines.

Our MSF is a general improver to both transformer-based classics and diffusion-based state-of-the-arts, as shown in Tab. 2, 2 and 4. On complex synthetic dataset ClevrTex, MSF improves SLATE significantly, up to 12% in ARI_fg. On challenging real-world dataset COCO, MSF still manages to improve SLATE by 6% in ARI_fg. MSF even boosts SlotDiffusion, which is state-of-the-art, by 3% in ARI on VOC. Compared with other improvers, i.e., SysBinder@$g$2 and GDR@$g$2, our MSF also beats or at least ties them on SLATE, while beats them consistently on STEVE. On SlotDiffusion, our MSF beats the strong competitor GDR@$g$2, and defeats DINOSUAR.

By the way, our MSF boosts OCL performance on COCO and VOC less than on other datasets. We attribute this to those object sizes are often out of the coverage of scales 1~3. Namely, many humans, trees and grass lands in those real-world images take up the whole scene, and either down sampling 4$\times$ (scale 1) or 16$\times$ (scale 3) cannot produce scale-invariant patterns for inter-scale fusion.

To investigate how MSF fuses multi-scale information into VAE guidance, we analyze object separability (Stanic et al., 2024; Lowe et al., 2024). Empirically, better object separability in VAE guidance contributes to better OCL performance. We cluster on the VAE guidance $\bm{X}_{n=1}$ then calculate inter-/intra-cluster distances; We also visualize (Caron et al., 2021) object separability of scale-invariant representations $\{\dot{\bm{X}}_{n}^{\mathrm{si}}\}$ and scale-variant representations $\{\bm{X}_{n}^{\mathrm{sv}}\}$, as well as representations before inter-scale fusion $\{\bm{X}_{n}^{\mathrm{si}}\}$ and representations after intra-scale fusion $\{\bm{X}_{n}\}$.

As shown in Fig. 3 and Tab. 4, our MSF contributes to significant object separability boosts both qualitatively and quantitatively upon basis methods. In contrast, the object separability of basis methods is consistently inferior to ours.

We use SLATE on ClevrTex as an example to evaluate how different designs contribute to our MSF. We use ARI+ARI_fg as the metrics because ARI is mostly dominated by background segmentation performance while ARI_fg only measures the foreground. We employ naive CNN (Kipf et al., 2022) as the primary encoder to reduce experiment time.

Inter/intra-scale fusion: with vs without. As shown in Tab. 5 upper left, both types of fusion are crucial to performance. Yet using inter-scale fusion alone can be even worse than disabling both, meaning that inter-scale fusion enforces scale-invariance at the cost of losing too much information, wheras intra-scale fusion can recover it. Therefore, enabling both types of fusion is the best setting.

Scale-invariant/variant codebook: shared vs specified. Firstly, using different codebooks as the scale-invariant codebook is meaningless because the corresponding matching probabilities for the same super-pixel cannot be compared. As shown in Tab. 5 upper right, sharing one codebook over different scales for the scale-variance is harmful to the performance. Hence different sets of codes are necessary to foster diverse patterns that are different among the scales.

Number of scales: 2, 3 or 4. As shown in Tab. 5 lower left, the default number of scales $N=3$ is the best setting. By contrast, $N=2$ includes too less explicit scales; As for $N=4$, its last scale, i.e., scale 4, down samples too much and loses too much spatial information to support the inter-scale fusion of our MSF technique.

Number of guidance: 1 or 2. After MSF in VAE, we have three augmented scales. But as shown in Tab. 5 lower right, using the augmented scale 1 as VAE guidance is the best choice, compared with either using both augmented scale 1 and 2, or using augmented scale 2 solely. This is because the augmented scale 1 representation has the largest resolution, i.e., the most spatial details, and also multiple scales of information. Although the other two scales are also multi-scale fused, the information incompleteness due to their low resolution is harmful to guide OCL training.