SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Aleksandar Stanić<sup>1*</sup> Anand Gopalakrishnan<sup>1*</sup> Kazuki Irie<sup>2†</sup> Jürgen Schmidhuber<sup>1,3</sup>
<sup>1</sup>The Swiss AI Lab, IDSIA, USI & SUPSI, Lugano, Switzerland
$^{2}$ Center for Brain Science, Harvard University, Cambridge, USA
$^{3}$ AI Initiative, KAUST, Thuwal, Saudi Arabia
{aleksandar, anand, juergen}@idsia.ch kirie@fas.harvard.edu
Current state-of-the-art object-centric models use slots and attention-based routing for binding. However, this class of models has several conceptual limitations: the number of slots is hardwired; all slots have equal capacity; training has high computational cost; there are no object-level relational factors within slots. Synchrony-based models in principle can address these limitations by using complex-valued activations which store binding information in their phase components. However, working examples of such synchrony-based models have been developed only very recently, and are still limited to toy grayscale datasets and simultaneous storage of less than three objects in practice. Here we introduce architectural modifications and a novel contrastive learning method that greatly improve the state-of-the-art synchrony-based model. For the first time, we obtain a class of synchrony-based models capable of discovering objects in an unsupervised manner in multi-object color datasets and simultaneously representing more than three objects.<sup>1</sup>
The visual binding problem [1, 2] is of great importance to human perception [3] and cognition [4-6]. It is fundamental to our visual capabilities to integrate several features together such as color, shape, texture, etc., into a unified whole [7]. In recent years, there has been a growing interest in deep learning models [8-19] capable of grouping visual inputs into a set of 'meaningful' entities in a fully unsupervised fashion (often called object-centric learning). Such compositional object-centric representations facilitate relational reasoning and generalization, thereby leading to better performance on downstream tasks such as visual question-answering [20, 21], video game playing [22-26], and robotics [27-29], compared to other monolithic representations of the visual input.
The current mainstream approach to implement such object binding mechanisms in artificial neural networks is to maintain a set of separate activation vectors (so-called slots) [15]. Various segregation mechanisms [30] are then used to route requisite information from the inputs to infer each slot in a iterative fashion. Such slot-based approaches have several conceptual limitations. First, the binding information (i.e., addresses) about object instances are maintained only by the constant number of slots—a hard-wired component which cannot be adapted through learning. This restricts the ability of slot-based models to flexibly represent varying number of objects with variable precision without tuning the slot size, number of slots, number of iterations, etc. Second, the inductive bias used by the grouping module strongly enforces independence among all pairs of slots. This restricts individual
†Work done at IDSIA.
Official code repository: https://github.com/agopal42/ctcae
slot to store relational features at the object-level, and requires additional processing of slots using a relational module, e.g., Graph Neural Networks [31, 32] or Transformer models [33-35]. Third, binding based on iterative attention is in general computationally very demanding to train [36]. Additionally, the spatial broadcast decoder [37] (a necessary component in these models) requires multiple forward/backward passes to render the slot-wise reconstruction and alpha masks, resulting in a large memory overhead as well.
Recently, Löwe et al. [36] revived another class of neural object binding models [38-40] (synchrony-based models) which are based on complex-valued neural networks. Synchrony-based models are conceptually very promising. In principle they address most of the conceptual challenges faced by slot-based models. The binding mechanism is implemented via constructive or destructive phase interference caused by addition of complex-valued activations. They store and process information about object instances in the phases of complex activations which are more amenable to adaptation through gradient-based learning. Further, they can in principle store a variable number of objects with variable precision by partitioning the phase components of complex activations at varying levels of granularity. Additionally, synchrony-based models can represent relational information directly in their distributed representation, i.e., distance in phase space yields an implicit relational metric between object instances (e.g., inferring part-whole hierarchy from distance in "tag" space [39]). Lastly, the training of synchrony-based models is computationally more efficient by two orders of magnitude [36].
However, the true potential of synchrony-based models for object binding is yet to be explored; the current state-of-the-art synchrony-based model, the Complex-valued AutoEncoder (CAE) [36], still has several limitations. First, it is yet to be benchmarked on any multi-object datasets [41] with color images (even simplistic ones like Tetrominoes) due to limitations in the evaluation method to extract discrete object identities from continuous phase maps [36]. Second, we empirically observe that it shows low separability (Table 2) in the phase space, thereby leading to very poor (near chance-level) grouping performance on dSprites and CLEVR. Lastly, CAE can simultaneously represent at most 3 objects [36], making it infeasible for harder benchmark datasets [41, 42].
Our goal is to improve the state-of-art synchrony models by addressing these limitations of CAE [36]. First, we propose a few simple architectural changes to the CAE: i) remove the 1x1 convolution kernel as well as the sigmoid activation in the output layer of decoder, and ii) use convolution and upsample layers instead of transposed convolution in the decoder. These changes enable our improved CAE, which we call  $\mathrm{CAE}++$ , to achieve good grouping performance on the Tetrominoes dataset—a task on which the original CAE completely fails. Further, we introduce a novel contrastive learning method to increase separability in phase values of pixels (regions) belonging to two different objects. The resulting model, which we call Contrastively Trained Complex-valued AutoEncoders (CtCAE), is the first kind of synchrony-based object binding models to achieve good grouping performance on multi-object color datasets with more than three objects (Figure 2). Our contrastive learning method yields significant gains in grouping performance over  $\mathrm{CAE}++$ , consistently across three multi-object color datasets (Tetrominoes, dSprites and CLEVR). Finally, we qualitatively and quantitatively evaluate the separability in phase space and generalization of CtCAE w.r.t. number of objects seen at train/test time.
We briefly overview the CAE architecture [36] which forms the basis of our proposed models. CAE performs binding through complex-valued activations which transmit two types of messages: magnitudes of complex activations to represent the strength of a feature and phases to represent which features must be processed together. The constructive or destructive interference through addition of complex activations in every layer pressurizes the network to use similar phase values for all patches belonging to the same object while separating those associated with different objects. Patches of the same object contain a high amount of pointwise mutual information so their destructive interference would degrade its reconstruction.
The CAE is an autoencoder with real-valued weights that manipulate complex-valued activations. Let  $h$  and  $w$  denote positive integers. The input is a positive real-valued image  $\mathbf{x}' \in \mathbb{R}^{h \times w \times 3}$  (height  $h$  and width  $w$ , with 3 channels for color images). An artificial initial phase of zero is added to each pixel of  $\mathbf{x}'$  (i.e.,  $\phi = \mathbf{0} \in \mathbb{R}^{h \times w \times 3}$ ) to obtain a complex-valued input  $\mathbf{x} \in \mathbb{C}^{h \times w \times 3}$ :
$$
\mathbf {x} = \mathbf {x} ^ {\prime} \odot e ^ {i \phi}, \quad \text {w h e r e} \odot \text {d e n o t e s a H a d a m a r d p r o d u c t} \tag {1}
$$
Let  $d_{\mathrm{in}}$ ,  $d_{\mathrm{out}}$  and  $p$  denote positive integers. Every layer in the CAE transforms complex-valued input  $\mathbf{x} \in \mathbb{C}^{d_{\mathrm{in}}}$  to complex-valued output  $\mathbf{z} \in \mathbb{C}^{d_{\mathrm{out}}}$  (where we simply denote input/output sizes as  $d_{\mathrm{in}}$  and  $d_{\mathrm{out}}$  which typically have multiple dimensions, e.g.,  $h \times w \times 3$  for the input layer), using a function  $f_{\mathbf{w}}: \mathbb{R}^{d_{\mathrm{in}}} \to \mathbb{R}^{d_{\mathrm{out}}}$  with real-valued trainable parameters  $\mathbf{w} \in \mathbb{R}^p$ .  $f_{\mathbf{w}}$  is typically a convolutional or linear layer. First,  $f_{\mathbf{w}}$  is applied separately to the real and imaginary components of the input:
$$
\psi = f _ {\mathbf {w}} \left(\operatorname {R e} (\mathbf {x})\right) + f _ {\mathbf {w}} \left(\operatorname {I m} (\mathbf {x})\right) i \quad \in \mathbb {C} ^ {d _ {\text {o u t}}} \tag {2}
$$
Note that both  $\operatorname{Re}(\mathbf{x}), \operatorname{Im}(\mathbf{x}) \in \mathbb{R}^{d_{\mathrm{in}}}$ . Second, separate trainable bias vectors  $\mathbf{b}_{\mathbf{m}}, \mathbf{b}_{\phi} \in \mathbb{R}^{d_{\mathrm{out}}}$  are applied to the magnitude and phase components of  $\psi \in \mathbb{C}^{d_{\mathrm{out}}}$ :
$$
\mathbf {m} _ {\psi} = | \psi | + \mathbf {b} _ {\mathbf {m}} \quad \in \mathbb {R} ^ {d _ {\text {o u t}}} \quad ; \quad \phi_ {\psi} = \arg (\psi) + \mathbf {b} _ {\phi} \quad \in \mathbb {R} ^ {d _ {\text {o u t}}} \tag {3}
$$
Third, the CAE uses an additional gating function proposed by Reichert and Serre [40] to further transform this "intermediate" magnitude  $\mathbf{m}_{\psi} \in \mathbb{R}^{d_{\mathrm{out}}}$ . This gating function dampens the response of an output unit as a function of the phase difference between two inputs. It is designed such that the corresponding response curve approximates experimental recordings of the analogous curve from a Hodgkin-Huxley model of a biological neuron [40]. Concretely, an intermediate activation vector  $\chi \in \mathbb{R}^{d_{\mathrm{out}}} ( \text{called classic term} [ 40 ] )$  is computed by applying  $f_{\mathbf{w}}$  to the magnitude of the input  $\mathbf{x} \in \mathbb{C}^{d_{\mathrm{in}}}$ , and a convex combination of this classic term and the magnitude  $\mathbf{m}_{\psi}$  (called synchrony term [40]) from Eq. 3 is computed to yield "gated magnitudes"  $\mathbf{m}_{\mathbf{z}} \in \mathbb{R}^{d_{\mathrm{out}}}$  as follows:
$$
\chi = f _ {\mathbf {w}} (| \mathbf {x} |) + \mathbf {b} _ {\mathbf {m}} \quad \in \mathbb {R} ^ {d _ {\text {o u t}}} \quad ; \quad \mathbf {m} _ {\mathbf {z}} = \frac {1}{2} \mathbf {m} _ {\boldsymbol {\psi}} + \frac {1}{2} \chi \quad \in \mathbb {R} ^ {d _ {\text {o u t}}} \tag {4}
$$
Finally, the output of the layer  $\mathbf{z} \in \mathbb{C}^{d_{\mathrm{out}}}$  is obtained by applying non-linearities to this magnitude  $\mathbf{m}_{\mathbf{z}}$  (Eq. 4) while leaving the phase values  $\phi_{\psi}$  (Eq. 3) untouched:
$$
\mathbf {z} = \operatorname {R e L U} \left(\operatorname {B a t c h N o r m} \left(\mathbf {m} _ {\mathbf {z}}\right)\right) \odot e ^ {i \phi_ {\psi}} \quad \in \mathbb {C} ^ {d _ {\text {o u t}}} \tag {5}
$$
The ReLU activation ensures that the magnitude of  $\mathbf{z}$  is positive, and any phase flips are prevented by its application solely to the magnitude component  $\mathbf{m}_{\mathbf{z}}$ . For more details of the CAE [36] and gating function [40] we refer the readers to the respective papers. The final object grouping in CAE is obtained through K-means clustering based on the phases at the output of the decoder; each pixel is assigned to a cluster corresponding to an object [36].
We describe the architectural and contrastive training details used by our proposed CAE++ and CtCAE models respectively below.
$\mathbf{CAE} + +$  . We first propose some simple but crucial architectural modifications that enable the vanilla CAE [36] to achieve good grouping performance on multi-object datasets such as Tetrominoes with color images. These architectural modifications include - i) Remove the 1x1 convolution kernel and associated sigmoid activation in the output layer ("  $f_{\mathrm{out}}$  " in Löwe et al. [36]) of the decoder, ii) Use convolution and upsample layers in place of transposed convolution layers in the decoder (cf. "  $f_{\mathrm{dec}}$  architecture in Table 3 from Löwe et al. [36]). We term this improved CAE variant that adopts these architectural modifications as  $\mathbf{CAE} + +$  . As we will show below in Table 1, these modifications allow our  $\mathbf{CAE} + +$  to consistently outperform the CAE across all 3 multi-object datasets with color images.
Contrastive Training of CAEs. Despite the improved grouping of  $\mathrm{CAE}++$  compared to CAE, we still empirically observe that  $\mathrm{CAE}++$  shows poor separability $^2$  (we also illustrate this in Section 4). This motivates us to introduce an auxiliary training objective that explicitly encourages higher separability in phase space. For that, we propose a contrastive learning method [43, 44] that modulates the distance between pairs of distributed representations based on some notion of (dis)similarity between them (which we define below). This design reflects the desired behavior to drive the phase separation process between two different objects thereby facilitating better grouping performance.
Before describing our contrastive learning method mathematically below, here we explain its essential ingredients (illustrated in Figure 1). For setting up the contrastive objective, we first (randomly)
2We define separability as the minimum angular distance of phase values between a pair of prototypical points (centroids) that belong to different objects
Figure 1: Sampling process of positive (green) and negative (red) pairs for one anchor (purple) in the CtCAE model. The sampling process here is visualized only for the decoder output. Note that we contrast the encoder output in an identical manner. Anchor address (the purple box marked with an X) corresponds to the patch of magnitude values and the feature corresponds to the phase values. See Contrastive Training of CAEs in Section 3 for more details.
Algorithm 1 Mining positive and negative pairs for a single anchor for the contrastive objective. Scalar parameters are  $k_{\mathrm{top}}$  — the number of candidates from which to sample one positive pair and  $m_{\mathrm{bottom}}$  — the number of candidates from which to sample  $(M - 1)$  negative pairs.
Inputs: anchor_address  $\in \mathbb{R}^{d_f}$ , addresses  $\in \mathbb{R}^{h' \times w' \times d_f}$ , features  $\in \mathbb{R}^{h' \times w' \times d_f}$ .
Params: Scalars  $k_{\mathrm{top}}$ ,  $m_{\mathrm{bottom}}$ ,  $M$ .
1: addresses  $\leftarrow$  Flatten(addresses, dims=(0,1))
2: features  $\leftarrow$  Flatten(features, dims=(0,1))
3: distances  $\leftarrow$  CosineDistance(addresses, anchor_address)
4: top_k_features  $\leftarrow$  features[argsort(distance, dim=0)[:ktop]]
5: bottom_m_features  $\leftarrow$  features[argsort(distance, dim=0) [-mbottom:]]
6:pos_pairidx  $\sim$  Uniform [0,  $k_{\mathrm{top}}]$
7: neg_pair_idxs  $\sim$  Uniform [0,  $m_{\mathrm{bottom}}]\times (M - 1)\triangleright M - 1$  samples without replacement
8: positive_pair  $\leftarrow$  top_k_features[pos_pairidx]
9: negative_pairs  $\leftarrow$  bottom_m_features[neg_pair_idxs]
10: return positive_pair, negative_pairs
sample "anchors" from a set of "datapoints". The main idea of contrastive learning is to "contrast" these anchors to their respective "positive" and "negative" examples in a certain representation space. This requires us to define two representation spaces: one associated with the similarity measure to define positive/negative examples given an anchor, and another one on which we effectively apply the contrastive objective, itself defined as a certain distance function to be minimized. We use the term addresses to refer to the representations used to measure similarity, and consequently extract positive and negative pairs w.r.t. to the anchor. We use the term features to refer to the representations that are contrasted. As outlined earlier, the goal of the contrastive objective is to facilitate separability of phase values. It is then a natural choice to use the phase components of complex-valued outputs as features and the magnitude components as addresses in the contrastive loss. This results in angular distance between phases (features) being modulated by the contrastive objective based on how (dis)similar their corresponding magnitude components (addresses) are. Since the magnitude components of complex-valued activations are used to reconstruct the image (Equation (7)), they capture requisite visual properties of objects. In short, the contrastive objective increases or decreases the angular distance of phase components of points (pixels/image regions) in relation to how (dis)similar their visual properties are.
Our contrastive learning method works as follows. Let  $h'$ ,  $w'$ ,  $d_{\mathrm{f}}$ ,  $N_{A}$ , and  $M$  denote positive integers. Here we generically denote the dimension of the output of any CAE layer as  $h' \times w' \times d_{\mathrm{f}}$ . This results in a set of  $h' \times w'$  "datapoints" of dimension  $d_{\mathrm{f}}$  for our contrastive learning. From this set of datapoints, we randomly sample  $N_{A}$  anchors. We denote this set of anchors as a matrix  $\mathbf{A} \in \mathbb{R}^{N_{A} \times d_{\mathrm{f}}}$ ; each anchor is thus denoted as  $\mathbf{A}_k \in \mathbb{R}^{d_{\mathrm{f}}}$  for all  $k \in \{1, \dots, N_{A}\}$ . Now by using
Figure 2: Unsupervised object discovery on Tetrominoes, dSprites and CLEVR with CtCAE. "Phase Rad." (col. 5) is the radial plot with the phase values from  $-\pi$  to  $\pi$  radians. "Phase" (col. 6), are phase values (in radians) averaged over the 3 output channels as a heatmap (colors correspond to those from "Phase Rad.") and "Magnitude" (col. 7) is the magnitude component of the outputs.
Algorithm 1, we extract 1 positive and  $M - 1$  negative examples for each anchor. We denote these examples by a matrix  $\mathbf{P}^k\in \mathbb{R}^{M\times d_{\mathrm{f}}}$  for each anchor  $k\in \{1,\dots,N_A\}$  arranged such that  $\mathbf{P}_1^k\in \mathbb{R}^{d_{\mathrm{f}}}$  is the positive example and all other rows  $\mathbf{P}_j^k\in \mathbb{R}^{d_{\mathrm{f}}}$  for  $j\in \{2,\dots,M\}$  are negative ones. Finally, our contrastive loss is an adaptation of the standard InfoNCE loss [44] which is defined as follows:
$$
\mathcal {L} _ {\mathrm {c t}} = \frac {1}{N _ {A}} \sum_ {k = 1} ^ {N _ {A}} \log \left(\frac {\exp \left(d \left(\mathbf {A} _ {k} ; \mathbf {P} _ {1} ^ {k}\right) / \tau\right)}{\sum_ {j = 1} ^ {M} \exp \left(d \left(\mathbf {A} _ {k} ; \mathbf {P} _ {j} ^ {k}\right) / \tau\right)}\right) \tag {6}
$$
where  $N_{A}$  is the number of anchors sampled for each input,  $d(\mathbf{x}_k;\mathbf{x}_l)$  refers to the cosine distance between a pair of vectors  $\mathbf{x}_k,\mathbf{x}_l\in \mathbb{R}^{d_f}$  and  $\tau \in \mathbb{R}_{>0}$  is the softmax temperature.
We empirically observe that applying the contrastive loss on outputs of both encoder and decoder is better than applying it on only either one (Table 5). We hypothesize that this is the case because it utilizes both high-level, abstract and global features (on the encoder-side) as well as low-level and local visual cues (on the decoder-side) that better capture visual (dis)similarity between positive and negative pairs. We also observe that using magnitude components of complex outputs of both the encoder and decoder as addresses for mining positive and negative pairs while using the phase components of complex-valued outputs as the features for the contrastive loss performs the best among all the other possible alternatives (Table 13). These ablations also support our initial intuitions (described above) while designing the contrastive objective for improving separability in phase space.
Finally, the complete training objective function of CtCAE is:
$$
\mathcal {L} = \mathcal {L} _ {\mathrm {m s e}} + \beta \cdot \mathcal {L} _ {\mathrm {c t}}; \quad \mathcal {L} _ {\mathrm {m s e}} = \left\| \mathbf {x} ^ {\prime} - \hat {\mathbf {x}} \right\| _ {2} ^ {2}; \quad \hat {\mathbf {x}} = | \mathbf {y} | \tag {7}
$$
where  $\mathcal{L}$  defines the loss for a single input image  $\mathbf{x}' \in \mathbb{R}^{h \times w \times 3}$ , and  $\mathcal{L}_{\mathrm{mse}}$  is the standard reconstruction loss used by the CAE [36]. The reconstructed image  $\hat{\mathbf{x}} \in \mathbb{R}^{h \times w \times 3}$  is generated from the complex-valued outputs of the decoder  $\mathbf{y} \in \mathbb{C}^{h \times w \times 3}$  by using its magnitude component. In practice, we train all models by minimizing the training loss  $\mathcal{L}$  over a batch of images. The CAE baseline model and our proposed CAE++ variant are trained using only the reconstruction objective (i.e.  $\beta = 0$ ) whereas our proposed CtCAE model is trained using the complete training objective.
Here we provide our experimental results. We first describe details of the datasets, baseline models, training procedure and evaluation metrics. We then show results (always across 5 seeds) on grouping of our CtCAE model compared to the baselines (CAE and our variant CAE++) separation in phase space, generalization capabilities w.r.t to number of objects seen at train/test time and ablation studies for each of our design choices. Finally, we comment on the limitations of our proposed method.
Table 1: MSE and ARI scores (mean ± standard deviation across 5 seeds) for CAE, CAE++ and CtCAE models for Tetrominoes, dSprites and CLEVR on their respective full resolutions. For all datasets, CtCAE vastly outperforms CAE++ which in turn outperforms the CAE baseline. Results for 32x32 dSprites and CLEVR are also provided, these follow closely the scores on the full resolutions. SlotAttention results are from Emami et al. [45].
Dataset	Model	MSE ↓	ARI-FG ↑	ARI-FULL ↑
Tetrominoes (32x32)	CAE	4.57e-2 ± 1.08e-3	0.00 ± 0.00	0.12 ± 0.02
CAE++	5.07e-5 ± 2.80e-5	0.78 ± 0.07	0.84 ± 0.01
CtCAE	9.73e-5 ± 4.64e-5	0.84 ± 0.09	0.85 ± 0.01
SlotAttention	-	0.99 ± 0.00	-
dSprites (64x64)	CAE	8.16e-3 ± 2.54e-5	0.05 ± 0.02	0.10 ± 0.02
CAE++	1.60e-3 ± 1.33e-3	0.51 ± 0.08	0.54 ± 0.14
CtCAE	1.56e-3 ± 1.58e-4	0.56 ± 0.11	0.90 ± 0.03
SlotAttention	-	0.91 ± 0.01	-
CLEVR (96x96)	CAE	1.50e-3 ± 4.53e-4	0.04 ± 0.03	0.18 ± 0.06
CAE++	2.41e-4 ± 3.45e-5	0.27 ± 0.13	0.31 ± 0.07
CtCAE	3.39e-4 ± 3.65e-5	0.54 ± 0.02	0.68 ± 0.08
SlotAttention	-	0.99 ± 0.01	-
dSprites (32x32)	CAE	7.24e-3 ± 8.45e-5	0.01 ± 0.00	0.05 ± 0.00
CAE++	8.67e-4 ± 1.92e-4	0.38 ± 0.05	0.49 ± 0.15
CtCAE	1.10e-3 ± 2.59e-4	0.48 ± 0.03	0.68 ± 0.13
CLEVR (32x32)	CAE	1.84e-3 ± 5.68e-4	0.11 ± 0.07	0.12 ± 0.11
CAE++	4.04e-4 ± 4.04e-4	0.22 ± 0.10	0.30 ± 0.18
CtCAE	9.88e-4 ± 1.42e-3	0.50 ± 0.05	0.69 ± 0.25