Original: O BJECT - CENTRIC C OMPOSITIONAL I MAGINATION FOR V ISUAL A BSTRACT R EASONING
Revised: O BJECT - CENTRIC C OMPOSITIONAL I MAGINATION FOR V ISUAL A BSTRACT R EASONING

Original: Anonymous authors

Original: Paper under double-blind review

Original: Imagination will often carry us to worlds that never were.
Revised: Imagination will often carry us to worlds that never were.

Original: But without it we go nowhere — Carl Sagan.
Revised: But without it we go nowhere — Carl Sagan.

Original: A BSTRACT
Revised: A BSTRACT

Original: Like humans devoid of imagination, current machine learning systems lack the ability to adapt to new, unexpected situations by foreseeing them, which makes them unable to solve new tasks by analogical reasoning.
Revised: Like humans devoid of imagination, current machine learning systems lack the ability to adapt to new, unexpected situations by foreseeing them, which makes them unable to solve new tasks by analogical reasoning.

Original: In this work, we introduce a new compositional imagination framework that improves a model’s ability to generalize.
Revised: In this work, we introduce a new compositional imagination framework that improves a model’s ability to generalize out-of-distribution.

Original: One of the key components of our framework is object-centric inductive biases that enables models to perceive the environment as a series of objects, properties, and transformations.
Revised: One of the key components of our framework is object-centric inductive biases that enables models to perceive the environment as a series of objects, properties, and transformations.

Original: By composing these key ingredients, it is possible to generate new unseen tasks that, when used to train the model, improve  generalization.
Revised: By composing these key ingredients, it is possible to generate new unseen tasks that, when used to train the model, improve systematic generalization.

Original: Experiments on a simplified version of the Abstraction and Reasoning Corpus (ARC) demonstrate the effectiveness of our framework.
Revised: Experiments on a simplified version of the Abstraction and Reasoning Corpus (ARC) demonstrate the effectiveness of our framework.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Using simple concepts as building blocks, humans have the capacity to compose and enhance knowledge by relating it to previous experiences.
Revised: Using simple concepts as building blocks, humans have the capacity to compose and enhance knowledge by relating it to previous experiences.

Original: Such analogical reasoning [23, 21] leverages the compositional structure of the world to make sense of new experiences and to imagine previously unseen scenarios.
Revised: Such analogical reasoning [23, 21] leverages the compositional structure of the world to make sense of new experiences and to imagine previously unseen scenarios.

Original: Many machine learning models are also able to acquire knowledge from data and use it successfully to perform a given task.
Revised: Many machine learning models are also able to acquire knowledge from data and use it successfully to perform a given task.

Original: However, their ability to adapt this knowledge to new domains remains unsatisfactory [4].
Revised: However, their ability to adapt this knowledge to new domains remains unsatisfactory [4].

Original: Objects are one of the core abstractions of the human brain when applying analogical reasoning [23].
Revised: Objects are one of the core abstractions of the human brain when applying analogical reasoning [23].

Original: For instance, we can infer the properties of a new object by transferring our knowledge of these properties from similar objects [21].
Revised: For instance, we can infer the properties of a new object by transferring our knowledge of these properties from similar objects [21].

Original: This realization has inspired a recent body of work that focuses on learning models that discover objects in a visual scene without supervision [8, 17, 10, 27, 11, 3, 28, 20].
Revised: This realization has inspired a recent body of work that focuses on learning models that discover objects in a visual scene without supervision [8, 17, 10, 27, 11, 3, 28, 20].

Original: Many of these works propose several inductive biases that lead to a visual scene decomposition in terms of its constituting objects.
Revised: Many of these works propose several inductive biases that lead to a visual scene decomposition in terms of its constituting objects.

Original: The expectation is that such objectcentric decomposition would lead to better generalization since it better represents the underlying structure of the physical world [22].
Revised: The expectation is that such objectcentric decomposition would lead to better generalization since it better represents the underlying structure of the physical world [22].

Original: To the best of our knowledge, the effect of such object-centric representations for systematic generalization in visual reasoning tasks remains largely unexplored.
Revised: To the best of our knowledge, the effect of such object-centric representations for systematic generalization in visual reasoning tasks remains largely unexplored.

Original: In an attempt to better measure the gap between machine and human learning, Chollet [4] introduced the Abstraction and Reasoning Corpus (ARC) .
Revised: In an attempt to better measure the gap between machine and human learning, Chollet [4] introduced the Abstraction and Reasoning Corpus (ARC) .

Original: ARC consists of visual analogy tasks made of few input-output pairs.
Revised: ARC consists of visual analogy tasks made of few input-output pairs.

Original: The goal of a model is to learn to infer the program applied to the inputs from few demonstrations and test it on novel samples.
Revised: The goal of a model is to learn to infer the program applied to the inputs from few demonstrations and test it on novel samples.

Original: The ARC challenge remains insurmountable by current
Revised: The ARC challenge remains insurmountable by current

Original: ML systems.
Revised: ML systems.

Original: Chollet [4] hypothesizes that a successful algorithm should be able to synthetize new solutions by recombining a set of subprograms or primitives based on core knowledge [23].
Revised: Chollet [4] hypothesizes that a successful algorithm should be able to synthetize new solutions by recombining a set of subprograms or primitives based on core knowledge [23].

Original: Dreamcoder [7] is one of the methods that are closest to Chollet’s desiderata for solving ARC.
Revised: Dreamcoder [7] is one of the methods that are closest to Chollet’s desiderata for solving ARC.

Original: In their work, Ellis et al.
Revised: In their work, Ellis et al.

Original: [7] propose imagining new program instances by composing primitives from a domain-specific language (DSL) and training the model to recognize the resulting solutions.
Revised: [7] propose imagining new program instances by composing primitives from a domain-specific language (DSL) and training the model to recognize the resulting solutions.

Original: In their work, the authors show that augmenting the training data with imagined programs is critical in a low data regime.
Revised: In their work, the authors show that augmenting the training data with imagined programs is critical in a low data regime.

Original: They also test their model on a suite of program induction domains where the developer can anticipate which primitives it needs to provide in the initial DSL.
Revised: They also test their model on a suite of program induction domains where the developer can anticipate which primitives it needs to provide in the initial DSL.

Original: In a visual domain such as ARC, however, such primitives require more high-level notions related to core knowledge and are difficult to anticipate.
Revised: In a visual domain such as ARC, however, such primitives require more high-level notions related to core knowledge and are difficult to anticipate.

Original: In this work, we propose a way of using object-centric inductive biases to derive a new compositional imagination framework in which a model can learn a series of concepts without the need for specifying a DSL.
Revised: In this work, we propose a way of using object-centric inductive biases to derive a new compositional imagination framework in which a model can learn a series of concepts without the need for specifying a DSL.

Original: We show that by imagining new scenarios that are composed of learned concepts, and by learning to predict back those concepts, the model better generalizes to compositions unseen during training.
Revised: We show that by imagining new scenarios that are composed of learned concepts, and by learning to predict back those concepts, the model better generalizes to compositions unseen during training.

Original: Finally, since exploring new methods directly on ARC is a daunting task, we introduce Sort-of-ARC , a dataset where the programs operate on objects and are conditioned on some of their attributes in a controlled environment.
Revised: Finally, since exploring new methods directly on ARC is a daunting task, we introduce Sort-of-ARC , a dataset where the programs operate on objects and are conditioned on some of their attributes in a controlled environment.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Visual Reasoning.
Revised: Visual Reasoning.

Original: Both Hoshen & Werman [13] and Barrett et al.
Revised: Both Hoshen & Werman [13] and Barrett et al.

Original: [2] propose to assess the reasoning ability of a neural network by applying standardized intelligence tests such as Raven’s Progressive Matrices (RPM) [15, 35].
Revised: [2] propose to assess the reasoning ability of a neural network by applying standardized intelligence tests such as Raven’s Progressive Matrices (RPM) [15, 35].

Original: These kinds of problems have inspired a large number of deep learning methods [12, 14, 24, 38, 31, 34, 36].
Revised: These kinds of problems have inspired a large number of deep learning methods [12, 14, 24, 38, 31, 34, 36].

Original: In a recent survey, Mitchell [21] recommends evaluating models on generative tasks that focus on human core knowledge [23].
Revised: In a recent survey, Mitchell [21] recommends evaluating models on generative tasks that focus on human core knowledge [23].

Original: ARC presents one such task, where rather than choosing an answer, the solver has to generate its own solution, and hence is more resistant to learning shortcuts (unlike in the case of RPM).
Revised: ARC presents one such task, where rather than choosing an answer, the solver has to generate its own solution, and hence is more resistant to learning shortcuts (unlike in the case of RPM).

Original: However, ARC remains unapproachable by current deep learning methods.
Revised: However, ARC remains unapproachable by current deep learning methods.

Original: Since its release, the neural abstract reasoner [16] is the only deep learning method that succeeds in a reduced subset of ARC’s problems.
Revised: Since its release, the neural abstract reasoner [16] is the only deep learning method that succeeds in a reduced subset of ARC’s problems.

Original: Object-centric Representations.
Revised: Object-centric Representations.

Original: The objective of this line of work is to represent a visual scene in terms of the objects that compose it.
Revised: The objective of this line of work is to represent a visual scene in terms of the objects that compose it.

Original: Spatial mixture models [20, 3, 11] define a Gaussian mixture model weighted by the slot masks decoded from a set of learned structured latents.
Revised: Spatial mixture models [20, 3, 11] define a Gaussian mixture model weighted by the slot masks decoded from a set of learned structured latents.

Original: Spatialtransformer-based models further disentangle each slot’s latent representation into several variables (e.g. content, location, presence, depth) [8, 17, 5, 25, 19].
Revised: Spatialtransformer-based models further disentangle each slot’s latent representation into several variables (e.g. content, location, presence, depth) [8, 17, 5, 25, 19].

Original: All these approaches focus on the generative abilities of the models; in our case, we study the impact of object-centric inductive biases on systematic generalization of the models in a visual reasoning task.
Revised: All these approaches focus on the generative abilities of the models; in our case, we study the impact of object-centric inductive biases on systematic generalization of the models in a visual reasoning task.

Original: We observe that modularity of representations is as important as the mechanisms [9] that operate on them.
Revised: We observe that modularity of representations is as important as the mechanisms [9] that operate on them.

Original: Additionally, we show that object-centric inductive biases of both representations and mechanisms allow us to derive an imagination framework that leads to better systematic generalization.
Revised: Additionally, we show that object-centric inductive biases of both representations and mechanisms allow us to derive an imagination framework that leads to better systematic generalization.

Original: Imagination in Deep Learning.
Revised: Imagination in Deep Learning.

Original: It is well known that data generation techniques such as data augmentation [1, 6, 32], domain randomization [26], and mixup [37] can be used to improve the generalization capacity of a model.
Revised: It is well known that data generation techniques such as data augmentation [1, 6, 32], domain randomization [26], and mixup [37] can be used to improve the generalization capacity of a model.

Original: These methods tend to generate samples that lie in the training set manifold.
Revised: These methods tend to generate samples that lie in the training set manifold.

Original: Thus, they are not adequate for algorithmic reasoning tasks that that evaluate models on completely novel situations.
Revised: Thus, they are not adequate for algorithmic reasoning tasks that that evaluate models on completely novel situations.

Original: Another body of work leverages the compositional structure of the world and of algorithms to generate samples that could not be obtained by interpolation [18, 33, 7].
Revised: Another body of work leverages the compositional structure of the world and of algorithms to generate samples that could not be obtained by interpolation [18, 33, 7].

Original: However, the latter works rely on a set of defined primitives or DSL, which is not possible to obtain in more realistic or challenging settings.
Revised: However, the latter works rely on a set of defined primitives or DSL, which is not possible to obtain in more realistic or challenging settings.

Original:  O BJECT -C ENTRIC C OMPOSITIONAL I MAGINATION
Revised:  O BJECT -C ENTRIC C OMPOSITIONAL I MAGINATION

Original: We propose a model that leverages compositional inductive biases in the architecture to generate new data points.
Revised: We propose a model that leverages compositional inductive biases in the architecture to generate new data points.

Original: Our model learns to compose object-centric representations with a series of learned operations, and generate new imagined scenarios by re-arranging these components in novel ways.
Revised: Our model learns to compose object-centric representations with a series of learned operations, and generate new imagined scenarios by re-arranging these components in novel ways.

Original: Training models on imagined scenarios helps these models to generalize more systematically.
Revised: Training models on imagined scenarios helps these models to generalize more systematically.

Original: 3. Sort-of-ARC
Revised: 3. Sort-of-ARC

Original: While the performance of a model on ARC is an interesting metric for comparing artificial reasoning systems, ARC’s level of difficulty makes it impenetrable from a research perspective.
Revised: While the performance of a model on ARC is an interesting metric for comparing artificial reasoning systems, ARC’s level of difficulty makes it impenetrable from a research perspective.

Original: We thus introduce Sort-of-ARC , a toy version of ARC that shares the same input space and presents a series of simpler problems.
Revised: We thus introduce Sort-of-ARC , a toy version of ARC that shares the same input space and presents a series of simpler problems.

Original: These problems are generated by systematically applying a series of operations that test different aspects of reasoning algorithms, and which might eventually be extended and scale to the full benchmark.
Revised: These problems are generated by systematically applying a series of operations that test different aspects of reasoning algorithms, and which might eventually be extended and scale to the full benchmark.

Original: Some of these aspects include the ability to make abstractions such as objects, retrieving these objects based on their properties, applying operations to objects, and generalizing to new unseen combinations of objects, properties, and operations.
Revised: Some of these aspects include the ability to make abstractions such as objects, retrieving these objects based on their properties, applying operations to objects, and generalizing to new unseen combinations of objects, properties, and operations.

Original: For each task, the model is given a support set indexed by S composed of visual I/O pairs { x i , y i } i ∈ S with x i , y i ∈
Revised: For each task, the model is given a support set indexed by S composed of visual I/O pairs { x i , y i } i ∈ S with x i , y i ∈

Original: R
Revised: R

Original: d × d × c .
Revised: d × d × c .

Original: A task transforms each input of the same support set according to a program that is defined by a condition and a transformation .
Revised: A task transforms each input of the same support set according to a program that is defined by a condition and a transformation .

Original: The condition selects the object properties to which we will apply a particular transformation (e.g., the condition/transformation pair red/right translation would consist in translating all the red objects to the right).
Revised: The condition selects the object properties to which we will apply a particular transformation (e.g., the condition/transformation pair red/right translation would consist in translating all the red objects to the right).

Original: During training, the model is presented with a support set of N I/O pairs { x i , y i } to which we apply the same condition/transformation rule and train it to reconstruct the corresponding output y q of query input x q .
Revised: During training, the model is presented with a support set of N I/O pairs { x i , y i } to which we apply the same condition/transformation rule and train it to reconstruct the corresponding output y q of query input x q .

Original: The model has to first infer the underlying program that generated the support set and then apply it to a query input to generate the right output.
Revised: The model has to first infer the underlying program that generated the support set and then apply it to a query input to generate the right output.

Original: To explicitly measure the systematic generalization abilities of object-centric models, we consider an out-of-distribution (OOD) train/test split setting where all the condition and transformation primitives are seen during training individually, but some pairs are held-out in the test set.
Revised: To explicitly measure the systematic generalization abilities of object-centric models, we consider an out-of-distribution (OOD) train/test split setting where all the condition and transformation primitives are seen during training individually, but some pairs are held-out in the test set.

Original: More details can be found in the Appendix A.
Revised: More details can be found in the Appendix A.

Original: 3. O UR M ODEL
Revised: 3. O UR M ODEL

Original: Our model is composed of two parts, a controller and an executor .
Revised: Our model is composed of two parts, a controller and an executor .

Original: The controller encodes the support set and outputs an embedding containing instructions ; the executor then modifies an input query conditioned on these instructions.
Revised: The controller encodes the support set and outputs an embedding containing instructions ; the executor then modifies an input query conditioned on these instructions.

Original: All our model variants and baselines differ from one another by the executor architecture.
Revised: All our model variants and baselines differ from one another by the executor architecture.

Original: We provide a schematic of our model in Figure 5 and more details in
Revised: We provide a schematic of our model in Figure 5 and more details in

Original: Algorithm 1.
Revised: Algorithm 1.

Original: 3.2. C ONTROLLER
Revised: 3.2. C ONTROLLER

Original: The controller encodes a set of pairs { x i , o i } i =1 ..N of inputs X and outputs O, and outputs an instruction embedding z that will be given to the executor.
Revised: The controller encodes a set of pairs { x i , o i } i =1 ..N of inputs X and outputs O, and outputs an instruction embedding z that will be given to the executor.

Original: We consider different architectures for the controller: The perception backbone can either be slot-attention or CNN-based that would encode each image in the support set.
Revised: We consider different architectures for the controller: The perception backbone can either be slot-attention or CNN-based that would encode each image in the support set.

Original: The backbone is then followed by a transformer encoder [29] to output the instruction embedding z ∈ R p given to the executor.
Revised: The backbone is then followed by a transformer encoder [29] to output the instruction embedding z ∈ R p given to the executor.

Original: 3.2. E XECUTOR
Revised: 3.2. E XECUTOR

Original: The executor takes as input a query x q ∈ R d × d × c and an instruction embedding z ∈ R p provided by the controller, and outputs the resulting query modification o q ∈ R d × d × c .
Revised: The executor takes as input a query x q ∈ R d × d × c and an instruction embedding z ∈ R p provided by the controller, and outputs the resulting query modification o q ∈ R d × d × c .

Original: This process is composed of four main steps: (1) Decomposing the visual query input x q into a set of K entity-centric latent representations H
Revised: This process is composed of four main steps: (1) Decomposing the visual query input x q into a set of K entity-centric latent representations H

Original: = { h k } k
Revised: = { h k } k

Original: =1 ..K .
Revised: =1 ..K .

Original: (2) Translating the instruction embedding z into a neural program via a selection bottleneck (details below).
Revised: (2) Translating the instruction embedding z into a neural program via a selection bottleneck (details below).

Original: (3) Structured update of the the entity-centric latent query set H according to the selected neural program .
Revised: (3) Structured update of the the entity-centric latent query set H according to the selected neural program .

Original: (4) Entity-centric rendering of the updated latent set.
Revised: (4) Entity-centric rendering of the updated latent set.

Original: Steps (1) and (4) are parametrized by a slot attention module [20] and a spatial broadcast decoder [30], respectively.
Revised: Steps (1) and (4) are parametrized by a slot attention module [20] and a spatial broadcast decoder [30], respectively.

Original: The novelty of our approach lies in Steps (2) and (3), which we describe in detail below.
Revised: The novelty of our approach lies in Steps (2) and (3), which we describe in detail below.

Original: Neural Program Selection Bottleneck.
Revised: Neural Program Selection Bottleneck.

Original: The executor translates an instruction embedding z into a neural program via a key-query attention mechanism.
Revised: The executor translates an instruction embedding z into a neural program via a key-query attention mechanism.

Original: This mechanism selects a (1) condition c , which is used to select which latent slots the program should update; and a (2) transformation p , which defines how the latent slots should be modified.
Revised: This mechanism selects a (1) condition c , which is used to select which latent slots the program should update; and a (2) transformation p , which defines how the latent slots should be modified.

Original: We let ( c , c prob ) , ( p , p prob ) =
Revised: We let ( c , c prob ) , ( p , p prob ) =

Original: Selection ( z ) be the output of this selection mechanism, with c prob , p prob the selection probabilities obtained by the soft attention mechanism described below.
Revised: Selection ( z ) be the output of this selection mechanism, with c prob , p prob the selection probabilities obtained by the soft attention mechanism described below.

Original: The queries of the two soft attention mechanisms are extracted from the instruction embedding z with two different parametric operations, f c query and f p query , such that Q c = f c query ( z ) ∈ R 1 × p and
Revised: The queries of the two soft attention mechanisms are extracted from the instruction embedding z with two different parametric operations, f c query and f p query , such that Q c = f c query ( z ) ∈ R 1 × p and

Original: Q p = f p query ( z ) ∈ R 1 × p where c and p are used to index the condition and the transformation selection mechanisms, respectively.
Revised: Q p = f p query ( z ) ∈ R 1 × p where c and p are used to index the condition and the transformation selection mechanisms, respectively.

Original: The keys and values are partitioned and learned embeddings as part of the model and we denote them as K c , V c ∈ R N c × p and K p ,
Revised: The keys and values are partitioned and learned embeddings as part of the model and we denote them as K c , V c ∈ R N c × p and K p ,

Original: V p ∈
Revised: V p ∈

Original: R N p × p , where N c (resp.
Revised: R N p × p , where N c (resp.

Original: N p ) denotes the number of learned condition (resp.
Revised: N p ) denotes the number of learned condition (resp.

Original: transformation) embeddings.
Revised: transformation) embeddings.

Original: The resulting neural program is thus described by the softly selected condition c and transformation p embeddings obtained by the usual key-query attention mechanism:
Revised: The resulting neural program is thus described by the softly selected condition c and transformation p embeddings obtained by the usual key-query attention mechanism:

Original: [Equation]
Revised: [Equation]

Original: Entity-Centric Update.
Revised: Entity-Centric Update.

Original: Given a condition c and a transformation p , the latent slots H extracted from a query image are updated in an entity-centric way.
Revised: Given a condition c and a transformation p , the latent slots H extracted from a query image are updated in an entity-centric way.

Original: This update mechanism produces a latent output H new
Revised: This update mechanism produces a latent output H new

Original: = Update ( H , c , p ) .
Revised: = Update ( H , c , p ) .

Original: The condition c selects which slot to update by defining an update gate for each slot; for a given slot k , the update gate α k is obtained by comparing it to the condition such that α k
Revised: The condition c selects which slot to update by defining an update gate for each slot; for a given slot k , the update gate α k is obtained by comparing it to the condition such that α k

Original: = σ ( MLP pres ([ h k , c ])) .
Revised: = σ ( MLP pres ([ h k , c ])) .

Original: The transformation p defines a slot-wise update ˜ h k
Revised: The transformation p defines a slot-wise update ˜ h k

Original: = MLP up ([ h k , p ]) such that updated slot representation for slot k is defined by h new k 
Revised: = MLP up ([ h k , p ]) such that updated slot representation for slot k is defined by h new k = h k

Original: [Equation]
Revised: + α k ∗ ˜ h

Original: k .
Revised: k .

Original: The output image is then rendered using a Spatial Broadcast Decoder [30] on the new slots H new .
Revised: The output image is then rendered using a Spatial Broadcast Decoder [30] on the new slots H new .

Original: We let o =
Revised: We let o =

Original: Decode ( H new ) be the output image obtained by the structured entity-centric update to the query image x , conditioned on condition and transformation embeddings c and p .
Revised: Decode ( H new ) be the output image obtained by the structured entity-centric update to the query image x , conditioned on condition and transformation embeddings c and p .

Original: 3. C OMPOSITIONAL I MAGINATION
Revised: 3. C OMPOSITIONAL I MAGINATION

Original: Our main contribution consists of showing how object-centric inductive biases enable a learning paradigm in which new tasks are imagined and composed of primitive concepts learned during training.
Revised: Our main contribution consists of showing how object-centric inductive biases enable a learning paradigm in which new tasks are imagined and composed of primitive concepts learned during training.

Original: Similar to the work by Ellis et al.
Revised: Similar to the work by Ellis et al.

Original: [7], this paradigm is loosely inspired by the wake-sleep algorithm in which a recognition model (e.g., our controller) is trained on imagined data produced by a generative model (e.g., our executor).
Revised: [7], this paradigm is loosely inspired by the wake-sleep algorithm in which a recognition model (e.g., our controller) is trained on imagined data produced by a generative model (e.g., our executor).

Original: We propose to imagine new scenarios by selecting, in a random and uniform way, a condition i c and a transformation i p .
Revised: We propose to imagine new scenarios by selecting, in a random and uniform way, a condition i c and a transformation i p .

Original: We then train the model to predict the proper condition and transformation indices given the imagined I/O set.
Revised: We then train the model to predict the proper condition and transformation indices given the imagined I/O set.

Original: During the imagination phase, we reuse the episodes seen during the ”wake” phase X = { x i } i =1 ..N .
Revised: During the imagination phase, we reuse the episodes seen during the ”wake” phase X = { x i } i =1 ..N .

Original: The imagination phase is thus composed of the following steps:
Revised: The imagination phase is thus composed of the following steps:

Original: [Equation]
Revised: [Equation]

Original:  The model is trained to predict which condition/transformation were sampled to construct the imagines episodes and thus minimize the following cross-entropy loss:[Equation]
Revised: [Equation] The model is trained to predict which condition/transformation were sampled to construct the imagines episodes and thus minimize the following cross-entropy loss:

Original: [Equation]
Revised: [Equation]

Original:  where ( c i c ,i p prob )
Revised: ) where ( c i c ,i p prob )

Original: i c is the i c
Revised: i c is the i c

Original: (and i p respectively) position of predicted probability vector for the inferred condition c i c .
Revised: (and i p respectively) position of predicted probability vector for the inferred condition c i c .

Original: The overall training objective is thus composed of the query output prediction loss and the imagination loss.
Revised: The overall training objective is thus composed of the query output prediction loss and the imagination loss.

Original: Details about the coefficients and different warm-up schedules are given in the Appendix C.
Revised: Details about the coefficients and different warm-up schedules are given in the Appendix C.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: Our main goal is to showcase that object-centric inductive biases allow us to derive a compositional imagination training framework that help with systematic generalization.
Revised: Our main goal is to showcase that object-centric inductive biases allow us to derive a compositional imagination training framework that help with systematic generalization.

Original: Thus, we address the following question: Does training on imagined situations improve generalization on new unseen situations composed of known training concepts?
Revised: Thus, we address the following question: Does training on imagined situations improve generalization on new unseen situations composed of known training concepts?

Original: To do so, we propose a number of baselines and model variations that specifically measure the importance of, on the one hand, object-centric inductive biases in the representations and/or the execution and, on the other hand, the compositional imagination framework that we propose.
Revised: To do so, we propose a number of baselines and model variations that specifically measure the importance of, on the one hand, object-centric inductive biases in the representations and/or the execution and, on the other hand, the compositional imagination framework that we propose.

Original: Baselines.
Revised: Baselines.

Original: In order to verify whether object-centric representations are enough for systematic generalization to emerge we consider a baseline where we remove the selection bottleneck and the mechanism that operate on slots is thus monolithic.
Revised: In order to verify whether object-centric representations are enough for systematic generalization to emerge we consider a baseline where we remove the selection bottleneck and the mechanism that operate on slots is thus monolithic.

Original: This corresponds to the No Selection baseline in Table 1.
Revised: This corresponds to the No Selection baseline in Table 1.

Original: We additionally consider variations where the controller backbone to encode images in the support set is either CNN-based or the same slot attention network than the executor.
Revised: We additionally consider variations where the controller backbone to encode images in the support set is either CNN-based or the same slot attention network than the executor.

Original: Split
Revised: Split

Original: No
Revised: No

Original: [Table]
Revised: [Table]

Original: No Imagination
Revised: No Imagination

Original: With Imagination
Revised: With Imagination

Revised: Full Slot Att.

Original: [Table] CNN
Revised: CNN/Slot Att.

Original: [Table]
Revised: Full Slot Att.

Original: CNN/Slot Att
Revised: CNN/Slot Att.

Original: [Table]

Original: Full [Table] Att.
Revised: Full Slot Att.

Original: [Table]
Revised: [Table]

Revised: Results are averaged over five random seeds.

Original: Results
Revised: Results

Original: The results in Table 1 and Figure 2 indicate that (1) object-centric representations are not enough to obtain systematic generalization in a simple visual reasoning task and that the modularity of the mechanisms that operate on slots is important; (2) our proposed compositional imagination framework helps to generate scenarios that were not seen during training and leads to better systematic generalization of the model.
Revised: The results in Table 1 and Figure 2 indicate that (1) object-centric representations are not enough to obtain systematic generalization in a simple visual reasoning task and that the modularity of the mechanisms that operate on slots is important; (2) our proposed compositional imagination framework helps to generate scenarios that were not seen during training and leads to better systematic generalization of the model.

Original: On the other hand, we do observe a large variance in terms of OOD performance.
Revised: On the other hand, we do observe a large variance in terms of OOD performance.

Original: One hypothesis is that during imagination programs are sampled in a uniform way and are not conditioned on the support set at hand.
Revised: One hypothesis is that during imagination programs are sampled in a uniform way and are not conditioned on the support set at hand.

Original: This unconditioned sampling may lead to invalid programs that the controller cannot infer.
Revised: This unconditioned sampling may lead to invalid programs that the controller cannot infer.

Original: How to sample valid program constitutes an interesting future research direction.
Revised: How to sample valid program constitutes an interesting future research direction.

Original: We provide qualitative results on the selectivity patterns of the proposed method in Appendix A.
Revised: We provide qualitative results on the selectivity patterns of the proposed method in Appendix A.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: We have introduced a new compositional imagination framework to improve the performance of machine learning models in analogical reasoning tasks.
Revised: We have introduced a new compositional imagination framework to improve the performance of machine learning models in analogical reasoning tasks.

Original: We leverage object-centric representations to decompose problems into a series of objects, properties, and transformations.
Revised: We leverage object-centric representations to decompose problems into a series of objects, properties, and transformations.

Original: Once learned, we re-assemble these factors in novel ways, generating new unseen problems on which the model is further trained.
Revised: Once learned, we re-assemble these factors in novel ways, generating new unseen problems on which the model is further trained.

Original: We obtain encouraging results on a simplified version of ARC that we will make publicly available to facilitate further research on this challenging benchmark.
Revised: We obtain encouraging results on a simplified version of ARC that we will make publicly available to facilitate further research on this challenging benchmark.

Original: Future research will focus on multi-step reasoning tasks as well as adapting the model to the original ARC benchmark.
Revised: Future research will focus on multi-step reasoning tasks as well as adapting the model to the original ARC benchmark

Original: A S ORT - OF -A RC
Revised: A S ORT - OF -A RC

Original: We introduce Sort-of-ARC , a toy version of ARC that shares the same input space and presents a series of simpler problems.
Revised: We introduce Sort-of-ARC , a toy version of ARC that shares the same input space and presents a series of simpler problems.

Original: Currently, all the images in Sort-of-ARC are of size 20 × 20 , each containing 3 objects of size 3 × 3 .
Revised: Currently, all the images in Sort-of-ARC are of size 20 × 20 , each containing 3 objects of size 3 × 3 .

Original: We list them in Figure 3.
Revised: We list them in Figure 3.

Original: Each object is defined by a shape, a color and position.
Revised: Each object is defined by a shape, a color and position.

Original: There are 16 different shapes and 10 different colors in addition to the background color.
Revised: There are 16 different shapes and 10 different colors in addition to the background color.

Original: The positions are sampled so that none of the objects initially overlap with each other.
Revised: The positions are sampled so that none of the objects initially overlap with each other.

Original: Each support set is composed of 5 input/output pairs and the model is trained to predict the output of a single query image.
Revised: Each support set is composed of 5 input/output pairs and the model is trained to predict the output of a single query image.

Original: Each episode is constructed by sampling at random a condition and a transformation.
Revised: Each episode is constructed by sampling at random a condition and a transformation.

Original: The conditions can be either on the shape index or on the color whereas the transformations correspond to simple translations in one of the four cardinal directions.
Revised: The conditions can be either on the shape index or on the color whereas the transformations correspond to simple translations in one of the four cardinal directions.

Original: To avoid data points where the output is the same as the input, when generating an episode, we make sure that the condition is satisfied for at least one object of each input image of the support set.
Revised: To avoid data points where the output is the same as the input, when generating an episode, we make sure that the condition is satisfied for at least one object of each input image of the support set.

Original: B Q UALITATIVE R ESULTS
Revised: B Q UALITATIVE R ESULTS

Original: Figure 4 shows the attention maps on the learned conditions and transformations given different input programs.
Revised: Figure 4 shows the attention maps on the learned conditions and transformations given different input programs.

Original: For each ground-truth program (defined by a ground truth condition and transformation) we report, for each partitioned condition
Revised: For each ground-truth program (defined by a ground truth condition and transformation) we report, for each partitioned condition

Original: (resp. transformation) their proportion of having the maximum probability of being selected.
Revised: (resp. transformation) their proportion of having the maximum probability of being selected.

Original: The sparse pattern indicates a clear selectivity for each of the different conditions and transforms (both in distribution and out-of-distribution) depending on the ground-truth program despite some overlaps.
Revised: The sparse pattern indicates a clear selectivity for each of the different conditions and transforms (both in distribution and out-of-distribution) depending on the ground-truth program despite some overlaps.

Original: C M ODEL AND T RAINING D ETAILS
Revised: C M ODEL AND T RAINING D ETAILS

Original: The model is trained on 10000 episodes.
Revised: The model is trained on 10000 episodes.

Original: We distinguish between two phases during training: during the first phase the model is trained on episodes coming from the training set only and on during the second phase, imagined episodes are added and the model is trained to predict the condition/transformation that sampled to produce these episodes.
Revised: We distinguish between two phases during training: during the first phase the model is trained on episodes coming from the training set only and on during the second phase, imagined episodes are added and the model is trained to predict the condition/transformation that sampled to produce these episodes.

Original: In order for the imagination loss ( L im ) to be informative enough, we observed that it helped to wait for a few epochs before introducing it, hence the two phases.
Revised: In order for the imagination loss ( L im ) to be informative enough, we observed that it helped to wait for a few epochs before introducing it, hence the two phases.

Original: The imagination loss is introduced after 300 epochs and the corresponding coefficient in increased linearly from 0 to 10 during 200 epochs.
Revised: The imagination loss is introduced after 300 epochs and the corresponding coefficient in increased linearly from 0 to 10 during 200 epochs.

Original: During both phases the model is trained to reconstruct both the query outputs ( L query ) and the outputs of each input of the support set ( L support ).
Revised: During both phases the model is trained to reconstruct both the query outputs ( L query ) and the outputs of each input of the support set ( L support ).

Original: We also include some auxiliary loss L rec to train the perception part of the model to reconstruct images X from encoded slots H before any update such that L rec =
Revised: We also include some auxiliary loss L rec to train the perception part of the model to reconstruct images X from encoded slots H before any update such that L rec =

Original: CrossEntropyLoss ( X , decode ( encode ( X ))) with H = encode ( X ) .
Revised: CrossEntropyLoss ( X , decode ( encode ( X ))) with H = encode ( X ) .

Original: The total loss is thus expressed as :
Revised: The total loss is thus expressed as :

Original: [Equation]
Revised: [Equation]

Original: I D
Revised: I D

Original: G r o un d - t r u t h p
Revised: G r o un d - t r u t h p

Original: r
Revised: r

Original: o g r a m 0 3 6 

Revised: o g r a m 0 3 6

Original: Inferred condition
Revised: Inferred condition

Original: OO
Revised: OO

Original: D
Revised: D

Original: Inferred transform    0. Inferred transform    0.
Revised: Inferred transform    0. Inferred transform    0.

Original: Following recommendations to train the slot attention model [20], we use a linear learning rate warmup from 1 − 6 to 1 − 4 over the first 10 epochs.
Revised: Following recommendations to train the slot attention model [20], we use a linear learning rate warmup from 1 − 6 to 1 − 4 over the first 10 epochs.

Original: The overall pseudo-code is described in Figureand a sketch of the model is described in Figure 5.
Revised: The overall pseudo-code is described in Figureand a sketch of the model is described in Figure 5.

Original: C. M ODEL V ARIATIONS
Revised: C. M ODEL V ARIATIONS

Original: Controller The controller is composed of a perception backbone than can either be CNN or slot attention-based [20].
Revised: Controller The controller is composed of a perception backbone than can either be CNN or slot attention-based [20].

Original: The perception backbone is then followed by a transformer encoder [29] and a mean average to aggregate the information from all the elements of the support set.
Revised: The perception backbone is then followed by a transformer encoder [29] and a mean average to aggregate the information from all the elements of the support set.

Original: For the CNN-based backbone, each image (input and output) is independently encoded using a simple CNN.
Revised: For the CNN-based backbone, each image (input and output) is independently encoded using a simple CNN.

Original: For each pair in a support set we also encode the difference between the output and the input.
Revised: For each pair in a support set we also encode the difference between the output and the input.

Original: The input, output and difference embeddings are then concatenated to form the representation of the I/O that will then be processed by the transformer encoder.
Revised: The input, output and difference embeddings are then concatenated to form the representation of the I/O that will then be processed by the transformer encoder.

Original: For the slot-based backbone, the slot attention module is shared with the perception part of the executor.
Revised: For the slot-based backbone, the slot attention module is shared with the perception part of the executor.

Original: The input of each I/O pair of a support set is first encoded into a set of slots H in (starting from random slots) and the output is encoded into a set of output slots H out using the same slot attention module but starting from the input slots H in .
Revised: The input of each I/O pair of a support set is first encoded into a set of slots H in (starting from random slots) and the output is encoded into a set of output slots H out using the same slot attention module but starting from the input slots H in .

Original: The input and output slots are then concatenated in a slot-wise manner H pair =
Revised: The input and output slots are then concatenated in a slot-wise manner H pair =

Original: [ H in , H out ] and an I/O pair is represented by a vector z s using a weighted sum of each slot contribution such that:
Revised: [ H in , H out ] and an I/O pair is represented by a vector z s using a weighted sum of each slot contribution such that:

Original: [Equation]
Revised: [Equation]

Original: with the slot-wise importance weight w i = f importance ( H pair i ) and the slot-wise contribution h si = f contribution ( H pair i ) where both f importance and f contribution are parametrized with simple MLPs.
Revised: with the slot-wise importance weight w i = f importance ( H pair i ) and the slot-wise contribution h si = f contribution ( H pair i ) where both f importance and f contribution are parametrized with simple MLPs.

Original: Each
Revised: Each

Original: I/O representation of the support set will also be processed by the transformer encoder to get a final support episode representation z that will be given to the Executor.
Revised: I/O representation of the support set will also be processed by the transformer encoder to get a final support episode representation z that will be given to the Executor.

Original: Algorithm
Revised: Algorithm

Original: 1: Initialize all parameters 2: Set max iters and warmup iters 3: for it ≤ max iters do 4: # Controller Block 5: { X , Y } , { x q , y q } ∼ D ▷ Sample support set and a query from the dataset 6: z = Controller ( X , Y )
Revised: 1: Initialize all parameters 2: Set max iters and warmup iters 3: for it ≤ max iters do 4: # Controller Block 5: { X , Y } , { x q , y q } ∼ D ▷ Sample support set and a query from the dataset 6: z = Controller ( X , Y )

Original: ▷ Obtain program instructions 7:
Revised: ▷ Obtain program instructions 7:

Original: H =
Revised: H =

Original: Slot Attention ( x q ) ▷ Extract slots 8: # Executor block 9: ( c , c prob ) , ( p , p prob ) =
Revised: Slot Attention ( x q ) ▷ Extract slots 8: # Executor block 9: ( c , c prob ) , ( p , p prob ) =

Original: Selection ( z ) ▷ Infer neural program 10: function U PDATE ( H , c , p ) 11: α k = MLP pres ([ h k , c ])
Revised: Selection ( z ) ▷ Infer neural program 10: function U PDATE ( H , c , p ) 11: α k = MLP pres ([ h k , c ])

Original: ▷ Predict the presence of each slot for input update 12: ˜ h
Revised: ▷ Predict the presence of each slot for input update 12: ˜ h

Original: k =
Revised: k =

Original: MLP update ([ h k , p ])
Revised: MLP update ([ h k , p ])

Original: ▷ Obtain update for each slot 13: h newk =
Revised: ▷ Obtain update for each slot 13: h newk =

Original: h k + α k ∗ ˜ h
Revised: h k + α k ∗ ˜ h

Original: k ▷ Update slots according to presence 14: return H new 15: end function 16:
Revised: k ▷ Update slots according to presence 14: return H new 15: end function 16:

Original: H newk =
Revised: H newk =

Original: Update ( H , z ) 17: o = Decode ( H new )
Revised: Update ( H , z ) 17: o = Decode ( H new )

Original: ▷ Decode updated slots 18:
Revised: ▷ Decode updated slots 18:

Original: L rec = cross entropy ( o , y q )
Revised: L rec = cross entropy ( o , y q )

Original: ▷ Compute reconstruction loss 19: # Imagination 20: if it ≥ warmup iters then 21: i c , i p ∼ U (1 , N c ) , U (1 , N p )
Revised: ▷ Compute reconstruction loss 19: # Imagination 20: if it ≥ warmup iters then 21: i c , i p ∼ U (1 , N c ) , U (1 , N p )

Original: ▷ Randomly sample a neural program 22:
Revised: ▷ Randomly sample a neural program 22:

Original: o i c ,i p i =
Revised: o i c ,i p i =

Original: Update ( H i , c i c , p i p )
Revised: Update ( H i , c i c , p i p )

Original: ▷ Imagine new support outputs by applying the randomly sampled program to the original inputs 23: { X , ˆ Y } =
Revised: ▷ Imagine new support outputs by applying the randomly sampled program to the original inputs 23: { X , ˆ Y } =

Original: { x i , o i c ,i p i } i =1 ..
Revised: { x i , o i c ,i p i } i =1 ..

Original: N ▷ Build a new support set with imagined outputs 24: z i c ,i p =
Revised: N ▷ Build a new support set with imagined outputs 24: z i c ,i p =

Original: Controller ( X , ˆ Y )
Revised: Controller ( X , ˆ Y )

Original: ▷ Obtain program instructions 25: c i c,i p prob ,p i c ,i p prob =
Revised: ▷ Obtain program instructions 25: c i c , i c , prob =

Original: Selection ( zi c ,i p )
Revised: Selection ( z c ,i p )

Original: ▷ Infer randomly sampled neural program 26:
Revised: ▷ Infer randomly sampled neural program 26:

Original: [Equation]
Revised: [Equation]

Original:  Compute imagination loss 27:
Revised: ▷ Compute imagination loss 27:

Original: [Equation]
Revised: [Equation]

Original: 28: else 29:
Revised: 28: else 29:

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original:  30: end if 31: end for
Revised: rec 30: end if 31: end for

Original: No selection For the baseline without any selection bottleneck the slots are still updated in an entity-centric manner but this time the output of the controller z is directly given to the updating module such that, following the same notations introduced previously, a slot h i is updated as :
Revised: No selection For the baseline without any selection bottleneck the slots are still updated in an entity-centric manner but this time the output of the controller z is directly given to the updating module such that, following the same notations introduced previously, a slot h i is updated as :

Original: [Equation]
Revised: [Equation]

Original: with g parametrized by a simple MLP.
Revised: with g parametrized by a simple MLP.

Original: Controller
Revised: Controller

Original: Support set p key-query selection
Revised: Support set p key-query selection

Original: Query image argument & transformation selection
Revised: Query image argument & transformation selection

Original: Predicted query output
Revised: Predicted query output

Original: Conditions
Revised: Conditions

Original: Slot attention
Revised: Slot attention

Original: Transformations
Revised: Transformations

Original: Imagined
Revised: Imagined

Original: Executor
Revised: Executor

Original: Decoder z
Revised: Decoder z

Original: H c
Revised: H c