SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Sébastien Lachapelle*,1
Divyat Mahajan*
Ioannis Mitliagkas
Simon Lacoste-Julien†,1
Mila & DIRO, Université de Montréal
$^{1}$ Samsung - SAIT AI Lab, Montreal
We tackle the problems of latent variables identification and "out-of-support" image generation in representation learning. We show that both are possible for a class of decoders that we call additive, which are reminiscent of decoders used for object-centric representation learning (OCRL) and well suited for images that can be decomposed as a sum of object-specific images. We provide conditions under which exactly solving the reconstruction problem using an additive decoder is guaranteed to identify the blocks of latent variables up to permutation and block-wise invertible transformations. This guarantee relies only on very weak assumptions about the distribution of the latent factors, which might present statistical dependencies and have an almost arbitrarily shaped support. Our result provides a new setting where nonlinear independent component analysis (ICA) is possible and adds to our theoretical understanding of OCRL methods. We also show theoretically that additive decoders can generate novel images by recombining observed factors of variations in novel ways, an ability we refer to as Cartesian-product extrapolation. We show empirically that additivity is crucial for both identifiability and extrapolation on simulated data.
The integration of connectionist and symbolic approaches to artificial intelligence has been proposed as a solution to the lack of robustness, transferability, systematic generalization and interpretability of current deep learning algorithms [53, 4, 13, 25, 21] with justifications rooted in cognitive sciences [20, 28, 43] and causality [57, 63]. However, the problem of extracting meaningful symbols grounded in low-level observations, e.g. images, is still open. This problem is sometime referred to as disentanglement [4, 48] or causal representation learning [63]. The question of identifiability in representation learning, which originated in works on nonlinear independent component analysis (ICA) [65, 31, 33, 36], has been the focus of many recent efforts [49, 66, 26, 47, 3, 9, 41]. The mathematical results of these works provide rigorous explanations for when and why symbolic representations can be extracted from low-level observations. In a similar spirit, Object-centric representation learning (OCRL) aims to learn a representation in which the information about different objects are encoded separately [19, 22, 11, 24, 18, 51, 14]. These approaches have shown impressive results empirically, but the exact reason why they can perform this form of segmentation without any supervision is poorly understood.
* Equal contribution. † Canada CIFAR AI Chair.
Correspondence to: {lachaseb, divyat.mahajan}@mila.quebec
Figure 1: Left: Additive decoders model the additive structure of scenes composed of multiple objects. Right: Additive decoders allow to generate novel images never seen during training via Cartesian-product extrapolation (Corollary 3). Purple regions correspond to latents/observations seen during training. The blue regions correspond to the Cartesian-product extension. The middle set is the manifold of images of balls. In this example, the learner never saw both balls high, but these can be generated nevertheless thanks to the additive nature of the scene. Details in Section 3.2.
Our first contribution is an analysis of the identifiability of a class of decoders we call additive (Definition 1). Essentially, a decoder  $f(z)$  acting on a latent vector  $z \in \mathbb{R}^{d_z}$  to produce an observation  $x$  is said to be additive if it can be written as  $f(z) = \sum_{B \in \mathcal{B}} f^{(B)}(z_B)$  where  $\mathcal{B}$  is a partition of  $\{1, \ldots, d_z\}$ ,  $f^{(B)}(z_B)$  are "block-specific" decoders and the  $z_B$  are non-overlapping subvectors of  $z$ . This class of decoder is particularly well suited for images  $x$  that can be expressed as a sum of images corresponding to different objects (left of Figure 1). Unsurprisingly, this class of decoder bears similarity with the decoding architectures used in OCRL (Section 2), which already showed important successes at disentangling objects without any supervision. Our identifiability results provide conditions under which exactly solving the reconstruction problem with an additive decoder identifies the latent blocks  $z_B$  up to permutation and block-wise transformations (Theorems 1 & 2). We believe these results will be of interest to both the OCRL community, as they partly explain the empirical success of these approaches, and to the nonlinear ICA and disentanglement community, as it provides an important special case where identifiability holds. This result relies on the block-specific decoders being "sufficiently nonlinear" (Assumption 2) and requires only very weak assumptions on the distribution of the ground-truth latent factors of variations. In particular, these factors can be statistically dependent and their support can be (almost) arbitrary.
Our second contribution is to show theoretically that additive decoders can generate images never seen during training by recombining observed factors of variations in novel ways (Corollary 3). To describe this ability, we coin the term "Cartesian-product extrapolation" (right of Figure 1). We believe the type of identifiability analysis laid out in this work to understand "out-of-support" generation is novel and could be applied to other function classes or learning algorithms such as DALLE-2 [59] and Stable Diffusion [61] to understand their apparent creativity and hopefully improve it.
Both latent variables identification and Cartesian-product extrapolation are validated experimentally on simulated data (Section 4). More specifically, we observe that additivity is crucial for both by comparing against a non-additive decoder which fails to disentangle and extrapolate.
Notation. Scalars are denoted in lower-case and vectors in lower-case bold, e.g.  $x \in \mathbb{R}$  and  $\pmb{x} \in \mathbb{R}^n$ . We maintain an analogous notation for scalar-valued and vector-valued functions, e.g.  $f$  and  $\pmb{f}$ . The  $i$ th coordinate of the vector  $\pmb{x}$  is denoted by  $\pmb{x}_i$ . The set containing the first  $n$  integers excluding 0 is denoted by  $[n]$ . Given a subset of indices  $S \subseteq [n]$ ,  $\pmb{x}_S$  denotes the subvector consisting of entries  $\pmb{x}_i$  for  $i \in S$ . Given a function  $\pmb{f}(\pmb{x}_S) \in \mathbb{R}^m$  with input  $\pmb{x}_S$ , the derivative of  $\pmb{f}$  w.r.t.  $\pmb{x}_i$  is denoted by  $D_i\pmb{f}(\pmb{x}_S) \in \mathbb{R}^m$  and the second derivative w.r.t.  $\pmb{x}_i$  and  $\pmb{x}_{i'}$  is  $D_{i,i'}^2\pmb{f}(\pmb{x}_S) \in \mathbb{R}^m$ . See Table 2 in appendix for more.
Code: Our code repository can be found at this link.
2 Background & Literature review
Identifiability of latent variable models. The problem of latent variables identification can be best explained with a simple example. Suppose observations  $\pmb{x} \in \mathbb{R}^{d_x}$  are generated i.i.d. by first sampling a latent vector  $\pmb{z} \in \mathbb{R}^{d_z}$  from a distribution  $\mathbb{P}_{\pmb{z}}$  and feeding it into a decoder function  $\pmb{f}: \mathbb{R}^{d_z} \to \mathbb{R}^{d_x}$ ,
i.e.  $\pmb{x} = \pmb{f}(\pmb{z})$ . By choosing an alternative model defined as  $\hat{\pmb{f}} \coloneqq \pmb{f} \circ \pmb{v}$  and  $\hat{\pmb{z}} \coloneqq \pmb{v}^{-1}(\pmb{z})$  where  $\pmb{v} : \mathbb{R}^{d_z} \to \mathbb{R}^{d_z}$  is some bijective transformation, it is easy to see that the distributions of  $\hat{\pmb{x}} = \hat{\pmb{f}}(\hat{\pmb{z}})$  and  $\pmb{x}$  are the same since  $\hat{\pmb{f}}(\hat{\pmb{z}}) = \pmb{f} \circ \pmb{v}(\pmb{v}^{-1}(\pmb{z})) = \pmb{f}(\pmb{z})$ . The problem of identifiability is that, given only the distribution over  $\pmb{x}$ , it is impossible to distinguish between the two models  $(\pmb{f}, \pmb{z})$  and  $(\hat{\pmb{f}}, \hat{\pmb{z}})$ . This is problematic when one wants to discover interpretable factors of variations since  $\pmb{z}$  and  $\hat{\pmb{z}}$  could be drastically different. There are essentially two strategies to go around this problem: (i) restricting the hypothesis class of decoders  $\hat{\pmb{f}}$  [65, 26, 44, 54, 9, 73], and/or (ii) restricting/adding structure to the distribution of  $\hat{\pmb{z}}$  [33, 50, 42, 47]. By doing so, the hope is that the only bijective mappings  $\pmb{v}$  keeping  $\hat{\pmb{f}}$  and  $\hat{\pmb{z}}$  into their respective hypothesis classes will be trivial indeterminacies such as permutations and element-wise rescalings. Our contribution, which is to restrict the decoder function  $\hat{\pmb{f}}$  to be additive (Definition 1), falls into the first category. Other restricted function classes for  $\pmb{f}$  proposed in the literature include post-nonlinear mixtures [65], local isometries [16, 15, 29], conformal and orthogonal maps [26, 60, 9] as well as various restrictions on the sparsity of  $\pmb{f}$  [54, 73, 7, 71]. Methods that do not restrict the decoder must instead restrict/structure the distribution of the latent factors by assuming, e.g., sparse temporal dependencies [31, 38, 42, 40], conditionally independent latent variables given an observed auxiliary variable [33, 36], that interventions targeting the latent factors are observed [42, 47, 46, 8, 2, 3, 64, 10, 67, 72, 34], or that the support of the latents is a Cartesian-product [68, 62]. In contrast, our result makes very mild assumptions about the distribution of the latent factors, which can present statistical dependencies, have an almost arbitrarily shaped support and does not require any interventions. Additionally, none of these works provide extrapolation guarantees as we do in Section 3.2.
Relation to nonlinear ICA. Hyvarinen and Pajunen [32] showed that the standard nonlinear ICA problem where the decoder  $f$  is nonlinear and the latent factors  $z_{i}$  are statistically independent is unidentifiable. This motivated various extensions of nonlinear ICA where more structure on the factors is assumed [30, 31, 33, 36, 37, 27]. Our approach departs from the standard nonlinear ICA problem along three axes: (i) we restrict the mixing function to be additive, (ii) the factors do not have to be necessarily independent, and (iii) we can identify only the blocks  $z_{B}$  as opposed to each  $z_{i}$  individually up to element-wise transformations, unless  $\mathcal{B} = \{\{1\}, \dots, \{d_z\}\}$  (see Section 3.1).
Object-centric representation learning (OCRL). Lin et al. [45] classified OCRL methods in two categories: scene mixture models [22, 23, 24, 51] & spatial-attention models [19, 12, 11, 18]. Additive decoders can be seen as an approximation to the decoding architectures used in the former category, which typically consist of an object-specific decoder  $f^{(\mathrm{obj})}$  acting on object-specific latent blocks  $z_{B}$  and "mixed" together via a masking mechanism  $m^{(B)}(z)$  which selects which pixel belongs to which object. More precisely,
$$
\boldsymbol {f} (\boldsymbol {z}) = \sum_ {B \in \mathcal {B}} \boldsymbol {m} ^ {(B)} (\boldsymbol {z}) \odot \boldsymbol {f} ^ {(\mathrm {o b j})} \left(\boldsymbol {z} _ {B}\right), \text {w h e r e} \boldsymbol {m} _ {k} ^ {(B)} (\boldsymbol {z}) = \frac {\exp \left(\boldsymbol {a} _ {k} \left(\boldsymbol {z} _ {B}\right)\right)}{\sum_ {B ^ {\prime} \in \mathcal {B}} \exp \left(\boldsymbol {a} _ {k} \left(\boldsymbol {z} _ {B ^ {\prime}}\right)\right)}, \tag {1}
$$
and where  $\mathcal{B}$  is a partition of  $[d_z]$  made of equal-size blocks  $B$  and  $\pmb{a}:\mathbb{R}^{|B|}\to \mathbb{R}^{d_x}$  outputs a score that is normalized via a softmax operation to obtain the masks  $\pmb{m}^{(B)}(\pmb {z})$ . Many of these works also present some mechanism to select dynamically how many objects are present in the scene and thus have a variable-size representation  $\pmb{z}$ , an important technical aspect we omit in our analysis. Empirically, training these decoders based on some form of reconstruction objective, probabilistic or not, yields latent blocks  $\pmb{z}_B$  that represent the information of individual objects separately. We believe our work constitutes a step towards providing a mathematically grounded explanation for why these approaches can perform this form of disentanglement without supervision (Theorems 1 & 2). Many architectural innovations in scene mixture models concern the encoder, but our analysis focuses solely on the structure of the decoder  $\pmb{f}(\pmb {z})$ , which is a shared aspect across multiple methods. Generalization capabilities of object-centric representations were studied empirically by Dittadi et al. [14] but did not cover Cartesian-product extrapolation (Corollary 3) on which we focus here.
Diagonal Hessian penalty [58]. Additive decoders are also closely related to the penalty introduced by Peebles et al. [58] which consists in regularizing the Hessian of the decoder to be diagonal. In Appendix A.2, we show that "additivity" and "diagonal Hessian" are equivalent properties. They showed empirically that this penalty can induce disentanglement on datasets such as CLEVR [35], which is a standard benchmark for OCRL, but did not provide any formal justification. Our work provides a rigorous explanation for these successes and highlights the link between the diagonal Hessian penalty and OCRL.
Compositional decoders [7]. Compositional decoders were recently introduced by Brady et al. [7] as a model for OCRL methods with identifiability guarantees. A decoder  $\pmb{f}$  is said to be compositional when its Jacobian  $D\pmb{f}$  satisfies the following property everywhere: For all  $i \in [d_z]$  and  $B \in \mathcal{B}$ ,  $D_B f_i(\pmb{z}) \neq \mathbf{0} \Rightarrow D_{B^c} f_i(\pmb{z}) = \mathbf{0}$ , where  $B^c \coloneqq [d_z] \setminus B$ . In other words, each  $x_i$  can locally depend solely on one block  $z_B$  (this block can change for different  $z$ ). In Appendix A.3, we show that compositional  $C^2$  decoders are additive. Furthermore, Example 3 shows a decoder that is additive but not compositional, which means that additive  $C^2$  decoders are strictly more expressive than compositional  $C^2$  decoders. Another important distinction with our work is that we consider more general supports for  $z$  and provide a novel extrapolation analysis. That being said, our identifiability result does not supersede theirs since they assume only  $C^1$  decoders while our theory assumes  $C^2$ .
Extrapolation. Du and Mordatch [17] studied empirically how one can combine energy-based models for what they call compositional generalization, which is similar to our notion of Cartesian-product extrapolation, but suppose access to datasets in which only one latent factor varies and do not provide any theory. Webb et al. [70] studied extrapolation empirically and proposed a novel benchmark which does not have an additive structure. Besserve et al. [5] proposed a theoretical framework in which out-of-distribution samples are obtained by applying a transformation to a single hidden layer inside the decoder network. Krueger et al. [39] introduced a domain generalization method which is trained to be robust to tasks falling outside the convex hull of training distributions. Extrapolation in text-conditioned image generation was recently discussed by Wang et al. [69].
3 Additive decoders for disentanglement & extrapolation
Our theoretical results assume the existence of some data-generating process describing how the observations  $\pmb{x}$  are generated and, importantly, what are the "natural" factors of variations.
Assumption 1 (Data-generating process). The set of possible observations is given by a lower dimensional manifold  $\pmb{f}(\mathcal{Z}^{\mathrm{test}})$  embedded in  $\mathbb{R}^{d_x}$  where  $\mathcal{Z}^{\mathrm{test}}$  is an open set of  $\mathbb{R}^{d_z}$  and  $\pmb{f}: \mathcal{Z}^{\mathrm{test}} \to \mathbb{R}^{d_x}$  is a  $C^2$ -diffeomorphism onto its image. We will refer to  $\pmb{f}$  as the ground-truth decoder. At training time, the observations are i.i.d. samples given by  $\pmb{x} = \pmb{f}(\pmb{z})$  where  $\pmb{z}$  is distributed according to the probability measure  $\mathbb{P}_{\pmb{z}}^{\mathrm{train}}$  with support  $\mathcal{Z}^{\mathrm{train}} \subseteq \mathcal{Z}^{\mathrm{test}}$ . Throughout, we assume that  $\mathcal{Z}^{\mathrm{train}}$  is regularly closed (Definition 6).
Intuitively, the ground-truth decoder  $\pmb{f}$  is effectively relating the "natural factors of variations"  $z$  to the observations  $x$  in a one-to-one fashion. The map  $\pmb{f}$  is a  $C^2$ -diffeomorphism onto its image, which means that it is  $C^2$  (has continuous second derivative) and that its inverse (restricted to the image of  $\pmb{f}$ ) is also  $C^2$ . Analogous assumptions are very common in the literature on nonlinear ICA and disentanglement [33, 36, 42, 1]. Mansouri et al. [52] pointed out that the injectivity of  $\pmb{f}$  is violated when images show two objects that are indistinguishable, an important practical case that is not covered by our theory.
We emphasize the distinction between  $\mathcal{Z}^{\mathrm{train}}$ , which corresponds to the observations seen during training, and  $\mathcal{Z}^{\mathrm{test}}$ , which corresponds to the set of all possible images. The case where  $\mathcal{Z}^{\mathrm{train}} \neq \mathcal{Z}^{\mathrm{test}}$  will be of particular interest when discussing extrapolation in Section 3.2. The "regularly closed" condition on  $\mathcal{Z}^{\mathrm{train}}$  is mild, as it is satisfied as soon as the distribution of  $z$  has a density w.r.t. the Lebesgue measure on  $\mathbb{R}^{d_z}$ . It is violated, for example, when  $z$  is a discrete random vector. Figure 2 illustrates this assumption with simple examples.
Objective. Our analysis is based on the simple objective of reconstructing the observations  $\pmb{x}$  by learning an encoder  $\hat{\pmb{g}}: \mathbb{R}^{d_x} \to \mathbb{R}^{d_z}$  and a decoder  $\tilde{\pmb{f}}: \mathbb{R}^{d_z} \to \mathbb{R}^{d_x}$ . Note that we assumed implicitly that the dimensionality of the learned representation matches the dimensionality of the ground-truth. We define the set of latent codes the encoder can output when evaluated on the training distribution:
$$
\hat {\mathcal {Z}} ^ {\text {t r a i n}} := \hat {\boldsymbol {g}} \left(\boldsymbol {f} \left(\mathcal {Z} ^ {\text {t r a i n}}\right)\right). \tag {2}
$$
When the images of the ground-truth and learned decoders match, i.e.  $\pmb{f}(\mathcal{Z}^{\mathrm{train}}) = \hat{\pmb{f}} (\hat{\mathcal{Z}}^{\mathrm{train}})$ , which happens when the reconstruction task is solved exactly, one can define the map  $\pmb {v}:\hat{\mathcal{Z}}^{\mathrm{train}}\to \mathcal{Z}^{\mathrm{train}}$  as
$$
\boldsymbol {v} := \boldsymbol {f} ^ {- 1} \circ \hat {\boldsymbol {f}}. \tag {3}
$$
This function is going to be crucial throughout the work, especially to define  $\mathcal{B}$ -disentanglement (Definition 3), as it relates the learned representation to the ground-truth representation.
Before introducing our formal definition of additive decoders, we introduce the following notation: Given a set  $\mathcal{Z} \subseteq \mathbb{R}^{d_z}$  and a subset of indices  $B \subseteq [d_z]$ , let us define  $\mathcal{Z}_B$  to be the projection of  $\mathcal{Z}$  onto dimensions labelled by the index set  $B$ . More formally,
$$
\mathcal {Z} _ {B} := \left\{\boldsymbol {z} _ {B} \mid \boldsymbol {z} \in \mathcal {Z} \right\} \subseteq \mathbb {R} ^ {| B |}. \tag {4}
$$
Intuitively, we will say that a decoder is additive when its output is the summation of the outputs of "object-specific" decoders that depend only on each latent block  $\mathbf{z}_B$ . This captures the idea that an image can be seen as the juxatoposition of multiple images which individually correspond to objects in the scene or natural factors of variations (left of Figure 1).
Definition 1 (Additive functions). Let  $\mathcal{B}$  be a partition of  $[d_z]^1$ . A function  $f: \mathcal{Z} \to \mathbb{R}^{d_x}$  is said to be additive if there exist functions  $f^{(B)}: \mathcal{Z}_B \to \mathbb{R}^{d_x}$  for all  $B \in \mathcal{B}$  such that
$$
\forall \boldsymbol {z} \in \mathcal {Z}, \boldsymbol {f} (\boldsymbol {z}) = \sum_ {B \in \mathcal {B}} \boldsymbol {f} ^ {(B)} \left(\boldsymbol {z} _ {B}\right). \tag {5}
$$
This additivity property will be central to our analysis as it will be the driving force of identifiability (Theorem 1 & 2) and Cartesian-product extrapolation (Corollary 3).
Remark 1. Suppose we have  $\pmb{x} = \sigma (\sum_{B\in \mathcal{B}}\pmb{f}^{(B)}(\pmb{z}_B))$  where  $\sigma$  is a known bijective function. For example, if  $\sigma (\pmb {y})\coloneqq \exp (\pmb {y})$  (component-wise), the decoder can be thought of as being multiplicative. Our results still apply since we can simply transform the data doing  $\tilde{\pmb{x}}\coloneqq \sigma^{-1}(\pmb {x})$  to recover the additive form  $\tilde{\pmb{x}} = \sum_{B\in \mathcal{B}}\pmb{f}^{(B)}(\pmb{z}_B)$
Differences with OCRL in practice. We point out that, although the additive decoders make intuitive sense for OCRL, they are not expressive enough to represent the "masked decoders" typically used in practice (Equation (1)). The lack of additivity stems from the normalization in the masks  $m^{(B)}(z)$ . We hypothesize that studying the simpler additive decoders might still reveal interesting phenomena present in modern OCRL approaches due to their resemblance. Another difference is that, in practice, the same object-specific decoder  $f^{(\mathrm{obj})}$  is applied to every latent block  $z_{B}$ . Our theory allows for these functions to be different, but also applies when functions are the same. Additionally, this parameter sharing across  $f^{(B)}$  enables modern methods to have a variable number of objects across samples, an important practical point our theory does not cover.
3.1 Identifiability analysis
We now study the identifiability of additive decoders and show how they can yield disentanglement. Our definition of disentanglement will rely on partition-respecting permutations:
Definition 2 (Partition-respecting permutations). Let  $\mathcal{B}$  be a partition of  $\{1, \ldots, d_z\}$ . A permutation  $\pi$  over  $\{1, \ldots, d_z\}$  respects  $\mathcal{B}$  if, for all  $B \in \mathcal{B}$ ,  $\pi(B) \in \mathcal{B}$ .
Essentially, a permutation that respects  $\mathcal{B}$  is one which can permute blocks of  $\mathcal{B}$  and permute elements within a block, but cannot "mix" blocks together. We now introduce  $\mathcal{B}$ -disentanglement.
Definition 3 (B-disentanglement). A learned decoder  $\hat{f}:\mathbb{R}^{d_z}\to \mathbb{R}^{d_x}$  is said to be B-disentangled w.r.t. the ground-truth decoder  $f$  when  $f(\mathcal{Z}^{\mathrm{train}}) = \hat{f} (\hat{\mathcal{Z}}^{\mathrm{train}})$  and the mapping  $\pmb {v}\coloneqq \pmb{f}^{-1}\circ \hat{\pmb{f}}$  is a diffeomorphism from  $\hat{\mathcal{Z}}^{\mathrm{train}}$  to  $\mathcal{Z}^{\mathrm{train}}$  satisfying the following property: there exists a permutation  $\pi$  respecting  $\mathcal{B}$  such that, for all  $B\in \mathcal{B}$ , there exists a function  $\bar{\pmb{v}}_{\pi (B)}:\hat{\mathcal{Z}}_B^{\mathrm{train}}\rightarrow \mathcal{Z}_{\pi (B)}^{\mathrm{train}}$  such that, for all  $\pmb {z}\in \hat{\mathcal{Z}}^{\mathrm{train}}$ ,  $\pmb{v}_{\pi (B)}(\pmb {z}) = \bar{\pmb{v}}_{\pi (B)}(\pmb {z}_B)$ . In other words,  $\pmb{v}_{\pi (B)}(\pmb {z})$  depends only on  $\pmb{z}_B$ .
Thus,  $\mathcal{B}$ -disentanglement means that the blocks of latent dimensions  $z_{B}$  are disentangled from one another, but that variables within a given block might remain entangled. Note that, unless the partition is  $\mathcal{B} = \{\{1\}, \ldots, \{d_z\}\}$ , this corresponds to a weaker form of disentanglement than what is typically seeked in nonlinear ICA, i.e. recovering each variable individually.