Title: Improving Compositional Generalization using Iterated Learning and Simplicial Embeddings
Venue: NeurIPS 2023
Accepted: yes
Keywords: compositional generalization

systematic generalization

iterated learning

representation learning

graph neural networks
Abstract: Compositional generalization, the ability of an agent to generalize to unseen combinations of latent factors, is easy for humans but hard for deep neural networks. A line of research in cognitive science has hypothesized a process, "iterated learning," to help explain how human language developed this ability; the theory rests on simultaneous pressures towards compressibility (when an ignorant agent learns from an informed one) and expressivity (when it uses the representation for downstream tasks). Inspired by this process, we propose to improve the compositional generalization of deep networks by using iterated learning on models with simplicial embeddings, which can approximately discretize representations. This approach is further motivated by an analysis of compositionality based on Kolmogorov complexity. We show that this combination of changes improves compositional generalization over other approaches, demonstrating these improvements both on vision tasks with well-understood latent factors and on real molecular graph prediction tasks where the latent structure is unknown.

Improving Compositional Generalization using Iterated Learning and Simplicial Embeddings

Abstract

Compositional generalization, the ability of an agent to generalize to unseen combinations of latent factors, is easy for humans but hard for deep neural networks. A line of research in cognitive science has hypothesized a process, 'iterated learning,' to help explain how human language developed this ability; the theory rests on simultaneous pressures towards compressibility (when an ignorant agent learns from an informed one) and expressivity (when it uses the representation for downstream tasks). Inspired by this process, we propose to improve the compositional generalization of deep networks by using iterated learning on models with simplicial embeddings, which can approximately discretize representations. This approach is further motivated by an analysis of compositionality based on Kolmogorov complexity. We show that this combination of changes improves compositional generalization over other approaches, demonstrating these improvements both on vision tasks with well-understood latent factors and on real molecular graph prediction tasks where the latent structure is unknown.

1 Introduction

Deep neural networks have shown an amazing ability to generalize to new samples on domains where they have been extensively trained, approaching or surpassing human performance on tasks including image classification [62], Go [70], reading comprehension [13], and more. A growing body of literature, however, demonstrates that some tasks that can be easily solved by a human can be hard for deep models. One important such problem is compositional generalization ([18], comp-gen for short). For example, Schott et al. [65] study manually-created vision datasets where the true generating factors are known, and demonstrate that a wide variety of current representation learning methods struggle to learn the underlying mechanism. To achieve true 'artificially intelligent' methods that can succeed at a variety of difficult tasks, it seems necessary to demonstrate compositional generalization. One contribution of this paper is to lay out a framework towards understanding and improving compositional generalization, and argue that most currently-common training methods fall short.

In wondering how deep networks can learn to compositionally generalize, we might naturally ask: how did humans achieve such generalization? Or, as a particular case, how did human languages evolve components (typically, words) that can systematically combine to form new concepts? This has been a long-standing question in cognitive science and evolutionary linguistics. One promising hypothesis is known as iterated learning (IL), a procedure simulating cultural language evolution [41]. Aspects of this proposal are supported by lab experiments [42], a Bayesian model [7], the behavior of neural networks in a simple emergent communication task [60], and real tasks like machine translation [50] and visual question answering [76].

To link the study in cognitive science and deep learning, we first analyze the necessary properties of representations in order to generalize well compositionally. By linking the compositionality and the Kolmogorov complexity, we find iteratively resetting and relearning the representations can introduce compressibility pressure to the representations, which is also the key to the success of iterated learning. To apply iterated learning in a general representation learning problem, we propose to split the network into a backbone and a task head, and discretize the representation at the end of the backbone using simplicial embeddings (SEM, [45]). This scheme is more practical than LSTM [34] encoders previously used for neural iterated learning [60]. We observe in various controlled vision domains that SEM-IL can enhance compositional generalization by aligning learned representations to ground-truth generating factors. The proposed method also enhances downstream performance on molecular graph property prediction tasks, where the generating process is less clear-cut.

2 Compositional Generalization

Generalization is a long-standing topic in machine learning. The traditional notion of (in-distribution) generalization assumes that training and test samples come from the same distribution, but this is insufficient for many tasks: we expect a well-trained model to generalize to some novel scenarios that are unseen during training. One version of this is compositional generalization (comp-gen) [17], which requires the model to perform well on novel combinations of semantic concepts.

2.1 Data-generating assumption and problem definition

Any type of generalization requires some 'shared rules' between training and test distributions. We hence assume a simple data-generating process that both training and test data samples obey. In Figure 1, the semantic generating factors, also known as latent variables, are divided into two groups: the task-relevant factors (or semantic generating factors) G = [ G 1 , ..., G m ] , and task-irrelevant (or noise) factors O . This division depends on our understanding of the task; for example, if we only want to predict the digit identity of an image in the color-MNIST dataset [3], then m = 1 and G 1 represents the digit identity. All the other generating factors such as color, stroke, angle, and possible noise are merged into O . If we want to predict a function that depends on both identity and color, e.g. identifying blue even numbers, we could have G = [ G 1 , G 2 ] with G 1 the identity and G 2 the color.

Each input sample x ∈ X is determined by a deterministic function GenX ( G , O ) . The task label(s) y ∈ Y only depend on the factors G and possible independent noise ϵ , according to the deterministic function GenY ( G , ϵ ) . Note ( x , O ) ⊥ ⊥ ( y , ϵ ) | G , and that O , G , and ϵ are independent. The datagenerating distribution P ( x , y ) is determined by the latent distributions P ( G ) and P ( O ) , along with the GenX and GenY . We assume GenX and GenY are fixed across environments (the 'rules of production' are consistent), while P ( G ) and P ( O ) might change between training and test. 2

For compositional generalization, we wish to model the problem of generalizing to new combinations of previously seen attributes: understanding 'red circle' based on having seen 'red square' and 'blue circle.' Thus, we may assume that the supports of P ( G ) are non-overlapping between train and test. (If this assumption is not true, it only makes the problem easier.) In summary, our goal is to find an algorithm A such that, when trained on a dataset D train ∼ P n train , A achieves small test risk R P test ( A ( D train )) . Here P train and P test should satisfy these conditions:

P train and P test have G , O , ϵ jointly independent, and x = GenX ( G , O ) , y = GenY ( G , ϵ ) .

GenX and GenY are the same deterministic functions for P train and P test .

In challenging cases, we may have supp [ P train ( G )] ∩ supp [ P test ( G )] = ∅ .

Stage IV

Stage III

Stage II

Stage I

•

𝐻 𝐆 𝐳

> 0

𝐳

contains insufficient info. about

𝐆

•

𝐻 𝐆 𝐳

= 0

;

𝐻 𝐳 𝐆

> 0

𝐳

contain irrelevent non-

𝐆

info.

•

𝐻 𝐆 𝐳

= 0

;

𝐻 𝐳 𝐆

= 0

𝐳

only has info. in

𝐆

(bijection)

but their structures can be different

•

𝐳 ⇔ 𝐆

isomorphic bijection

𝐳

maintains the structure of

𝐆

𝐙

𝐎

𝐆

𝐗

𝐘

GenX

(𝐎, 𝐆)

GenY

(𝐆, 𝜖)

𝑔(𝐳)

(a). Data generating assumption

(b). Ladder of compositionality

Train longer

Bigger model

Data augmentation

Compressibility and

expressivity pressures

𝜖

𝐘

ℒ

Y, Y

2.2 Representationl Learning and Ladder of Compositionality

For compositional generalization, we expect that the model must extract atomic semantic features from the training data, and systematically re-combine them in a procedure akin to how the data is generated [41]. We thus consider a typical representation learning framework, which resembles the inverse of the data generation process (Figure 1(a), bottom). We use a backbone h : X → Z to convert the input signal x into a representation z , and a task head g : Z → Y to solve the given task based on that representation z . The prediction of the model is ˆ y = ( g ◦ h )( x ) .

Intuitively, we would like our learned z to uncover the hidden G , and g ( z ) to recover GenY ( G , ϵ ) . We thus analyze how the relationship between z and G influences the model's generalization capability, building off principles such as information bottleneck [74]. Inspired by the 'ladder of causation' [55], we propose a 'ladder of compositionality' in Figure 1(b), which outlining a series of conditions on z and G . We hypothesize that comp-gen roughly requires reaching the highest rung of that ladder:

Hypothesis 1. To generalize compositionally, the learned z should capture exactly the information in G and nothing more ( G to z should be a bijection), and moreover it should preserve the 'structure' of G (i.e. the mapping from G to z should be an isomorphism).

More on this hypothesis, the ladder, and relationship to models of disentanglement [32] are discussed in Appendix A. In short, we find that a model trained using common learning methods relying on mutual information between input x and supervision y cannot reliably reach the final stage of the ladder - it is necessary to seek other inductive biases in order to generalize compositionally.

3 Compressibility pressure and Compositional mapping

From the analysis above, we need to find other inductive biases to obtain compositional mappings. Inspired by how compositionality emerges in human language, 3 we speculate that the compressibility pressure is the key. Note that this pressure does not refer to compressing information from x to z (as in Stage III does), but whether a mapping can be expressed in a compact way by reusing common rules. In this section, we will first link compressibility pressure to Kolmogorov complexity by defining different mappings using group theory. As the Kolmogorov complexity is hard to compute, making explicit regularization dificult, we propose to implicitly regularize via iterated learning , a procedure in cognitive science proposed to increase compositionality in human-like language.

3.1 Compositional mappings have lower Kolmogorov complexity

From Occam's razor, we know efficient and effective mappings are more likely to capture the ground truth generating mechanism of the data, and hence generalize better. The efficiency is determined by how compressed the mapping is, which can also be measured by Kolmogorov complexity [47, 71].

To build a link between compositionality and Kolmogorov complexity, we can first describe different bijections between z and G using group theory, and then use the description length to compare the complexity of a typical element. Specifically, assuming z ∈ Z , G ∈ G and | Z | = | G | , the space of all bijections between z and G is an isomorphism of a symmetric group S | G | . If G = [ G 1 , ..., G m ] and each G m has v different possible values, | G | = v m . For clarity in the analysis, we assume z also has the same shape. Then, any bijection between z and G can be represented by an element in S v m .

The space of compositional mapping, which is a subset of all bijections, has more constraints. Recall how a compositional mapping is generated (see Appendix A.4 for more details): we first select z i for each G j in a non-overlapping way. Such a process can be represented by an element in S m . After that, we will assign different 'words' for each z i , which can be represented by an element in S v . As we have m different z i , this procedure will be repeated m times. In summary, any compositional mapping can be represented by an element in the group S v m ⋊ S v ∈ S v m , where ⋊ is the semidirect product in group theory. The cardinality of S v m is significantly larger than S v m ⋊ S v , and so a randomly selected bijection is unlikely to be compositional. Thus

Proposition 1 (Informal) . For m,v ≥ 2 , among all bijections, any compositional mapping has much lower Kolmogorov complexity than a typical non-compositional mapping.

We prove this by constructing descriptive protocols for each bijection. As a compositional mapping has more reused rules , its description length can be smaller (see Appendix B.1 for more details).

3.2 Compressibility pressure is amplified in iterated learning

Now, our target is finding bijections with higher compositionality and lower Kolmogorov complexity, which are both non-trivial. Because the ground truth G is usually inaccessible and the Kolmogorov complexity is hard to calculate. Fortunately, researchers find that human language also evolved to become more compositional without knowing G . Authors of [42] hypothesize that the compressibility pressure , which exists when an innocent agent (e.g., a child) learns from an informed agent (e.g., an adult), plays an important role. Such pressure is reinforced and amplified when the human community repeats this learning fashion for multiple generations.

However, the aforementioned hypothesis assumes that simplicity bias is inborn in the human cognition system. Will deep neural agents also have similar preferences during training? The answer is yes. By analyzing an overparameterized model on a simple supervised learning problem, we can strictly prove that repeatedly introducing new agents to learn from the old agent (then this informed agent becomes the old agent for the next generation) can exert a non-trivial regularizing effect on the number of 'active bases' of the learned mapping. Restricting the number of active bases encourages the model to reuse the learned rules. In other words, this regularization effect favors mappings with lower Kolmogorov complexity, which is exactly what we expect for compositional generalization. Due to the space limits, we left the formulation and proof of this problem in Appendix B.2.

3.3 Complete the proposed solution

Wethus expect that iteratively resetting and relearning can amplify the compressibility pressure, which helps us to reach the final rung of the ladder from the third. Before that, we need another pressure to reach third rung (i.e., ensure a bijection between z and G ). Expressivity pressure , constraining the learned mapping to be capable enough to accomplish the downstream tasks, is what we need.

The complete iterated learning hypothesis of Kirby et al. [42] claims that the compositional mapping emerges under the interaction between the compressibility pressure (i.e., efficiency) and the expressivity pressure (i.e., effectiveness). Inspired by this, we propose to train a model in generations consisting of two phases. At the t -th generation, we first train the backbone h in an imitation phase , where a student h S t learns to imitate z sampled from a teacher h T t . As analyzed above, iteratively doing so will amplify the compressibility pressure. Then, in the following interaction phase , the model g t ◦ h S t follows standard downstream training to predict y . The task head g t is randomly initialized and fine-tuned together with the backbone in this phase. By accomplishing this phase, the expressivity pressure is introduced. The fine-tuned backbone h S t then becomes the teacher h T t +1 for the next generation, and we repeat, as illustrated in Figure 2 and Algorithm 1.

Another problem with applying iterated learning to deep neural networks is how to create the discrete message, i.e., z . Discretization is not necessary : for example, the imitation phase could use L 2 loss to

...

(a). Iterated learning

Agent [t]

Agent [t-1]

Agent [t]

Learn (

𝐱

,

𝐳

)

Accomplish

downstream

Agent [t]

copy

Agent [t+1]

...

(b). Integration of SEM to iterated learning

(c). The SEM bottleneck

•

Interaction phase:

𝐱

Backbone

𝐡

𝐳

SEM

FC-layer

ℒ𝑑𝑠

•

Imitation phase:

𝐱

Teacher

SEM

ℒ𝑚𝑙(

Logits

, 𝐠

𝟏:𝐋

)

𝐠𝟏:𝐋

, sampled from

𝐳

Student

SEM

Logits

𝐡

𝑧1

𝑧2

𝑧𝐿

...

Softmax,

𝜏

𝑧1

𝑧2

𝑧𝐿

Concat

𝐳

Embed

Out 1

Out 2

𝐳

[1, L*V]

[L,V]

Softmax,

𝜏

Softmax,

𝜏

SEM bottleneck

Out 3

Logits

...

[L,V]

match a student's continuous representations to the teacher's. We find greatly improved performance with our discretization scheme, however, due to much-increased compressibility pressure. It is also possible [60] to use e.g. an LSTM encoder at the end of h ( x ) to produce discrete z , and an LSTM decoder at the start of g ( z ) . The interaction phase is then not directly differentiable; though many estimator options exist [6, 39, 78], training tends to be difficult due to high bias and/or variance.

Instead, we consider a simplicial embedding layer (SEM, [45]), which has proven effective on many self-supervised learning tasks. As illustrated in Figure 2(c), a dense representation h (the output of the original backbone) is linearly transformed into m vectors z i ∈ R v . Then we apply a separate softmax with temperature τ to each z i , yielding ¯ z i which are, if the temperature is not too high, approximately sparse; the ¯ z i are then concatenated to a long vector z . The overall process is

By using an encoder with a final SEM layer, we obtain an approximately-sparse z . In the imitation phase, we generate discrete pseudo-labels by sampling from the categorical distribution defined by each ¯ z i , then use cross-entropy loss so that the student is effectively doing multi-label classification to reconstruct the teacher's representations. In the imitation phase, the task head g operates directly on the long vector z . The full model g ◦ h is differentiable, so we can use any standard task loss. Pseudocode for the proposed method, SEM-IL, is in the appendix (Algorithm 1).

4 Analysis on Controlled Vision Datasets

We will first verify the effectiveness of the proposed SEM-IL method on controlled vision datasets, where the ground truth G is accessible. Thus, we can directly observe how z gradually becomes more similar to G , and how the compressibility and expressivity pressures affect the training process. In this section, we consider a regression task on 3dShapes [9], where recovering and recombining the generating factors is necessary for systematic generalization. The detailed experimental settings and results on additional similar datasets, dSprites [52] and MPI3D-real [23], are given in Appendix C.

4.1 The Effectiveness of SEM-IL

Better comp-gen performance We first show the effectiveness of the proposed method using results on 3dShapes, containing images of objects with various colors, sizes, and orientations against various backgrounds. Here G numerically encodes floor hue, wall hue, object hue, and object scale into discrete values, and the goal is to recover a particular linear function of that G . (Results for a simple nonlinear function were comparable.)

We compare five algorithms:

Baseline: directly train a ResNet18 [31] on the downstream task.

SEM-only: insert an SEM layer to the baseline model.

IL-only: train a baseline model with Algorithm 1, using MSE loss during imitation.

SEM-IL: train an SEM model with Algorithm 1.

Given-G: train an SEM model to reproduce the true G (which would not be known in practice), then fine-tune on the downstream task.

Trainingsetperformance

Testsetperformance

WithSEMbottleneck

NoSEMbottleneck

100

Baseline

100

Baseline

0.30

0.34

10-1

SEM-only

IL-only

10-1

SEM-only

IL-only

similarity

0.25

5

10-3

10-2

SEM-IL

Given-G

10-2

SEM-IL

Given-G

ogical

0.20

0.32

0.30

0.15

10-4

10-3

polog

0.28

10-5

0.10

SEM-only

Baseline

10-4

SEM-IL

0.26

IL-only

0

500

1000

1500

2000

0

500

1000

1500

2000

0

500

1000

1500

2000

0

500

1000

1500

2000

Training epochs

Trainingepochs

Training epochs

Training epochs

In the first panel of Figure 3, we see that the baseline and SEM-only models perform similarly on the training set; IL-based methods periodically increase in error at the beginning of each generation, but are eventually only slightly worse than the baselines on training data. On the test set, however, evaluating compositional generalization by using values of G which did not appear in training, SEM-IL brings significant improvement compared with other methods. Using only SEM or only IL gives no improvement over the baseline, however; it is only their combination which helps, as we will discuss further shortly. The (unrealistic) oracle method Given-G is unsurprisingly the best, since having z similar to G is indeed helpful for this task.

How z evolves during learning To see if better generalization ability is indeed achieved by finding z that resembles the structure of G , we check their topological similarity 4

where d z and d G are distance metrics, z ( i ) is the predicted representation of x ( i ) , and G ( i ) is the corresponding ground-truth generating factors. This measurement is widely applied to evaluate the compositionality of the mappings in cognitive science [8] and emergent communication [60]. Following existing works, we use the Hamming distance for G and discretized z in SEM-based methods, and cosine distance for continuous z in non-SEM methods. We expect h ( x ) to map x with similar G to close z , and dissimilar x to distant z , so that ρ ( z , G ) will be high.

The third panel of Figure 3 shows that the SEM-only model quickly reaches a plateau after 200 epochs and then slowly decreases, while SEM-IL, after briefly stalling at the same point, continues to increase to a notably higher topological similarity. In the last panel, however, the IL-only method doesn't improve ρ over the baseline: it seems both parts are needed.

4.2 Discretized Representation is Beneficial for the Imitation Phase of IL

To explain why SEM and IL cooperate well, we need to look deeper into how the compressibility pressure influences the learning of representations. This pressure induced by iterated learning, which helps us to find mappings with lower Kolmogorov complexity, leads to representations that are more compositional and systematic [42]. However, in prior works, these mappings were only considered in conjunction with some discretized representation [54, 60]. While IL could be used with continuous representation during the imitation phase, similar to born-again networks [19], we found that our algorithm benefits a lot from the discretized representations.

To get a clear picture of why discretized representations are so important, we divide h ( x ) into m submappings h i ( x ) , which map x to ¯ z i ∈ [0 , 1] v . We can understand each ¯ z i as a categorical distribution over v different possible values. As such, during training, the model learns discrete features of the dataset and assigns confidence about each feature for every sample. The neural network will tend to more quickly learn simpler mappings [5, 24], and will assign higher confidence according to the mapping it has learned. In other words, if a mapping does not align well with G , it is more likely to give idiosyncratic learned ¯ z i , and will lead to low confidence for most samples. On the contrary, ¯ z i belonging to compositional mappings will be more general, and on average tend towards higher confidence.

4

6

8

10

Confidence of teacher

10

15

20

25

30

35

40

45

50

Learning speed of student

consistent

different

0

500

1000

1500

2000

Training epochs

0.0

0.2

0.4

0.6

0.8

1.0

Predicting probability

High confidence

0

500

1000

1500

2000

Training epochs

0.0

0.2

0.4

0.6

0.8

1.0

Intermediate confidence

0

500

1000

1500

2000

Training epochs

0.0

0.2

0.4

0.6

0.8

1.0

Low confidence

The imitation phase reinforces this bias when the new student learns from the sampled pseudo labels g i from the teacher's prediction ¯ z i . As such, confident predictions, which are more likely to belong to the compositional mappings, will be learned faster (and harder to forget) by the student. On the contrary, for less confident features where P (¯ z i | x ) is flat, g i could change across epochs. This makes it hard for the student to remember any related ( x , g i ) . For example, a student will be reluctant to build a stable mapping between 'red' and z 1 if the teacher communicates ( 'red square' , g 1 = 0) , ( 'red square' , g 1 = 1) , ( 'red square' , g 1 = 2) in three consecutive epochs.

̸

Furthermore, using the sampled pseudo-labels can help the student to align the learned ( x , g i ) better. Assume during training, the student already remembers some pairs like ( 'blue circle' , g 1 = 0) , ( 'blue square' , g 1 = 0) , ( 'blue star' , g 1 = 0) , but the teacher is not confident in ( 'blue apple' , g 1 ) , perhaps because apples are rarely blue. Following the analysis above, as P (¯ z 1 | blue apple ) is flat, the teacher may generate g 1 = 0 a significant portion of the time. However, if the teacher happens to generate g 1 = 0 at some point, the student would learn ( 'blue apple' , g 1 = 0) faster than those with g 1 = 0 , because it aligns well with the other information stored in the student network. The parameter updates caused by the learning of other ( 'blue [shape]' , g 1 = 0) will also promote the learning of ( 'blue apple' , g 1 = 0) , similar to how 'noisy' labels are fixed as described by [61].

̸

To support the explanations above, we can first observe the correlation between the teacher's confidence and the model's learning speed for ( x , g i ) . Specifically, for each x , the teacher makes m predictions with the corresponding categorical distribution ¯ z i , i ∈ [ m ] . For each ( x , ¯ z i ) , the confidence is measured by the negative logarithm of the teacher's predicted probability, -log[¯ z i ] ˆ ȷ where ˆ ȷ ∈ argmax j [¯ z i ] j . The learning speed of ( x , g i ) is measured by the integral of the student's prediction with training time t , i.e., ∑ t =0 [ˆ z i ( t )] j , where j is the value provided by the teacher and ˆ z i ( t ) is student's prediction at time t . As illustrated in the first panel of Figure 4, the ( x , g i ) with higher confidence are usually learned faster by the student.

We also provide the learning curves of ( x , g i ) with high/intermediate/low confidence (each with 10 samples) in the other three panels of the figure. The curves for high-confidence samples all converge to [ˆ z i ( t )] j = 1 while those for low-confidence predictions could converge to a value less than 0.3. This means the student might make predictions that are different from the teacher's supervision. By highlighting such low-confidence ( x , g i ) pairs in the scatter plot, we find they are all low-confidence samples. Another interesting observation from the high-confidence curves is that some ( x , g i ) pairs are not remembered by the student in the first generation: they emerge at some point and gradually dominate as the training goes on. This phenomenon matches our analysis of how the sampled pseudo-labels help the student align ( x , g i ) to its knowledge well. To further support this explanation, Appendix C shows that performance is substantially harmed by taking pseudo-labels from the argmax , rather than sampling from ¯ z i .

To recap, this subsection provided an explanation (along with some supporting evidence) for why the combination of SEM and IL is so important, based on the perspective of sample difficulty, which we believe to be a significant factor in the success of this algorithm.

5 Application: Molecular Property Prediction

Given the success in controlled vision examples, we now turn to a real problem where the true generative process is unknown. We focus on predicting the properties of molecular graphs, for several

reasons. First, molecular graphs and their labels might follow a (chemical) procedure akin to that in Figure 1: for instance, one G i might be the existence of a specific functional group, or the number of specific atoms. Different molecular properties could then be determined by different subsets of G i , as we desired in the compositional generalization problem. Furthermore, the generating mechanisms ( GenX and GenY ) should be consistent and determined by nature. Second, benchmark datasets in this community contain various types of tasks (e.g., binary classification, multi-label classification, and regression) with similar input signals: performing well on different tasks will broaden the scope of our algorithm. Furthermore, the scaffold split used by most molecular datasets corresponds well to the compositional generalization setup we consider here. (We also try some more challenging splits, using structural information.) Last, learning meaningful representations that uncover the generating mechanisms of molecules is important, of practical significance, and difficult: it can potentially help predict the properties of unknown compounds, or accelerate the discovery of new compounds with specific properties, but scaling based on massive datasets as in recent work on vision or language seems more difficult. We hope our analysis can provide a new perspective on this problem.

5.1 Improvement on the Downstream Performance

We conduct experiments on three common molecular graph property datasets: ogbg-molhiv (1 binary classification task), ogbg-molpcba (128 binary classification tasks), and PCQM4Mv2 (1 regression task); all three come from the Open Graph Benchmark [37]. We choose two types of backbones, standard GCN [40] and GIN [80]. For the baseline experiments, we use the default hyperparameters from [37]. As the linear transform added in SEM-based method gives the model more parameters, we consider 'baseline+' to make a fair comparison: this model has an additional embedding layer, but no softmax operation. Detailed information on these datasets, backbone models, and hyper-parameters is provided in Appendix D.

From Table 1, we see the SEM-IL method almost always gives the best performance. Unlike in the controlled vision experiments (Figure 3), however, SEM alone can bring significant improvements in this setting. We speculate that compressibility pressure might be more significant in the interaction phase (i.e. standard training) when the generating mechanism is complex. This suggests it may be possible to develop a more efficient algorithm to better impose compressibility and expressivity pressures at the same time.

5.2 Probing Learned z by Meaningful Structures

In the controlled vision examples, we know that SEM-IL not only enhances the downstream performance, but also provides z more similar to the ground-truth generating factors, as seen by the improvement in topological similarity. However, as the generating mechanism is usually inaccessible in real problems, we indirectly measure the quality of z using graph probing [2]. Specifically, we first extract some meaningful substructures in a molecule using domain knowledge. For example, we can conclude whether a benzene ring exists in x by directly observing its 2D structure. With the help of the RDKit tool [44], we can generate a sequence of labels for each x , which is usually known as the 'fingerprint' of molecules (denoted FP ( x ) ∈ { 0 , 1 } k , indicating whether each specific structure exists in x ). Then, we add a linear head on top of the fixed z and train it using a generated training set ( x , FP ( x )) , x ∼ D train , and compare the generalization performance on the generated test set

( x , FP ( x )) , x ∼ D test . For fair comparison, we set m = 30 and v = 10 to make z and h be the same width, excluding the influence of the linear head's capacity.

In the experiments, we use the validation split of molhiv as D train and the test split as D test , each of which contain 4,113 distinct molecules unseen during the training of z . The generalization performance of ten different substructures is reported in Table 2. The first block (first two rows) of the table demonstrates the performance of two types of models before training. They behave similarly across all tasks and give a higher AUROC than a random guess. Then, comparing the three algorithms in each block, we see SEM-based methods consistently outperform the baseline, which supports our hypothesis well. SEM-IL outperforms SEM-only on average, but not for every task; this may be because some structures are more important to the downstream task than others.

Comparing the results across the four blocks, we find that the task in the interaction phase also influences the quality of z : the z trained by molpcba is much better than those trained by molhiv . To figure out where this improvement comes from, we first use only 10% of the training samples in molpcba to make the training sizes similar, then make the supervisory signal more similar by using only one task from molpcba . As illustrated in the last two blocks in the table, we can conclude that the complexity of the task in the interaction phase, which introduces the expressivity pressure, plays a more important role in finding better z .

Based on this observation, we can improve SEM-IL by applying more complex interaction tasks. For example, existing works on iterated learning use a referential game or a reconstruction task in the interaction phase, which could introduce stronger expressivity pressure from a different perspective. Furthermore, [45] demonstrates that SEM works well with most contrastive learning tasks. We hope the fundamental analysis provided in this paper can shed light on why SEM and IL collaborate so well and also arouse more efficient and effective algorithms in the future.

6 Related Works

Iterated Learning and its Applications. Iterated learning (IL) is a procedure that simulates cultural language evolution to explain how the compositionality of human language emerges [41]. In IL, the knowledge (i.e., the mapping between the input sample and its representation) is transferred between different generations, during which the compositional mappings gradually emerge and dominate under the interaction between compressibility and expressivity pressures. Inspired by this principle, there are some successful applications in symbolic games [60], visual question answering [76], machine translation [50], multi-label learning [58], reinforcement learning [54], etc.

There are also many algorithms training a neural network for multiple generations, which could possibly support the principles proposed in iterated learning. For example, [19] proposes to iteratively distill the downstream logits from the model in the previous generation, and finally bootstrap all the models to achieve better performance on image classification task; this can be considered as an IL algorithm merging the imitation and interaction phases together. [82] proposes to re-initialize the latter layers of a network and re-train the model for multiple generations, which is similar to an IL algorithm that only re-initializes the task head. [54] extends such a reset-and-relearn training to reinforcement learning and shows that resetting brings benefits that cannot be achieved by other regularization methods such as dropout or weight decay. In the era of large language models, selfrefinement in-context learning [51] and self-training-based reinforcement learning [26] can also

benefit from iteratively learning from the signals generated by agents in the previous generation. We left the discussion and analysis on these more complex real systems in our future work.

Knowledge Distillation and Discrete Bottleneck. Broadly speaking, the imitation phase in SEM-IL, which requires the student network to learn from the teacher, can be considered as a knowledge distillation method [33]. Different from the usual setting, where the student learns from the teacher's prediction on a downstream task, we assume a data-generating mechanism and create a simplex space for the generating factors. By learning from the teacher in this space, we believe the compressibility pressure is stronger and is more beneficial for the compositional generalization ability.

For the discretization, there are also other possible approaches, e.g., [28] uses an LSTM to create a discrete message space, and [48] proposes a method using a vector quantized bottleneck [75]. We choose SEM [45] for its simplicity and universality: it is easy to insert it into a model for different tasks. Besides, SEM has proved to be effective on self-supervised learning tasks; we extend it to classification, regression, and multi-label tasks.

Compressibility, learning dynamics, and Kolmogorov complexity Recently, with the success of large language models, the relationship between compressibility and generalization ability gradually attracted more attention [12]. Authors of [57] propose that how well a model is compressed corresponds to the integral of the training loss curve when negative logarithmic likelihood loss is used. Although this claim assumes the model sees each training sample only once, which might not be consistent with the multiple-epochs training discussed in this paper, the principles behind this claim and our analysis are quite consistent: the mappings generalize better and are usually learned faster by the model. Furthermore, authors of [71] link the generalization ability to Kolmogorov complexity. Our analysis in Appendix B also supports this claim well. Hence we believe the evolution of the human cognition system can provide valuable insights into deep learning systems.

Graph Representation Learning. Chemistry and molecular modeling are some of the main drivers of neural graph representation learning since its emergence [21] and graph neural networks, in particular. The first theoretical and practical advancements [27, 40, 80] in the GNN literature were mostly motivated by molecular use cases. Furthermore, many standard graph benchmarks [15, 16, 37] include molecular tasks on node, edge, and graph-levels, e.g., graph regression in ZINC and PCQM4Mv2 or molecular property prediction in ogbg-molhiv and ogbg-molpcba datasets. Graph Transformers [14, 43, 59] exhibit significant gains over GNNs in molecular prediction tasks. Selfsupervised learning (SSL) on graphs is particularly prominent in the molecular domain highlighted by the works of GNN PreTrain [38], BGRL [73], and Noisy Nodes [22]. We will extend the proposed method to different models and different pretraining strategies in our future work.

7 Conclusion

In this paper, we first define the compositional generalization problem by assuming the samples in the training and test sets share the same generating mechanism while the generating factors of these two sets can have different distributions. Then, by proposing the compositionality ladder, we analyze the desired properties of the representations. By linking the compositionality, compressibility, and Kolmogorov complexity together, we find iterated learning, which is well-studied in cognitive science, is beneficial for our problem. To appropriately apply iterated learning, we attach an SEM layer to the backbone model to discretize the representations. On the datasets where the true generating factors are accessible, we show that the representations learned by SEM-IL can better portray the generation factors and hence lead to better test performance. We then extend the proposed algorithm to molecular property prediction tasks and find it improves the generalization ability.

The main drawback of the current solution is the time-consuming training: we must run multiple generations and some common features might be re-learned multiple times, which is inefficient. Hence a more efficient way of imposing compressibility is desired.

Overall, though, our analysis and experiments show the potential of the SEM-IL framework on compositional generalization problems. We believe a better understanding of where the compressibility bias comes from in the context of deep learning can inspire more efficient and non-trivial IL framework designs. Clearly defining the compositional generalization problem and finding more related practical applications can also promote the development of IL-related algorithms.

References

Kartik Ahuja, Ethan Caballero, Dinghuai Zhang, Jean-Christophe Gagnon-Audet, Yoshua Bengio, Ioannis Mitliagkas, and Irina Rish. 'Invariance principle meets information bottleneck for out-of-distribution generalization.' NeurIPS . Volume 34. 2021, pages 3438-3450.

Mohammad Sadegh Akhondzadeh, Vijay Lingam, and Aleksandar Bojchevski. 'Probing Graph Representation.' AISTATS . 2023.

Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. 'Invariant risk minimization' (2019). arXiv: 1907.02893 .

Hangbo Bao, Li Dong, Songhao Piao, and Furu Wei. 'BEiT: BERT Pre-Training of Image Transformers.' ICLR . 2022.

Yoshua Bengio, Tristan Deleu, Nasim Rahaman, Nan Rosemary Ke, Sebastien Lachapelle, Olexa Bilaniuk, Anirudh Goyal, and Christopher Pal. 'A Meta-Transfer Objective for Learning to Disentangle Causal Mechanisms.' ICLR . 2020.

Yoshua Bengio, Nicholas Léonard, and Aaron Courville. 'Estimating or propagating gradients through stochastic neurons for conditional computation' (2013). arXiv: 1308.3432 .

Aaron Beppu and Thomas Griffiths. 'Iterated learning and the cultural ratchet.' Proceedings of the Annual Meeting of the Cognitive Science Society . Volume 31. 2009.

Henry Brighton and Simon Kirby. 'Understanding linguistic evolution by visualizing the emergence of topographic mappings.' Artificial Life 12.2 (2006), pages 229-242.

Chris Burgess and Hyunjik Kim. 3D Shapes Dataset . 2018. URL: https://github.com/ deepmind/3d-shapes/ .

Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. 'Emerging properties in self-supervised vision transformers.' CVPR . 2021.

Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. 'A simple framework for contrastive learning of visual representations.' ICML . 2020.

Grégoire Delétang, Anian Ruoss, Paul-Ambroise Duquenne, Elliot Catt, Tim Genewein, Christopher Mattern, Jordi Grau-Moya, Li Kevin Wenliang, Matthew Aitchison, Laurent Orseau, Marcus Hutter, and Joel Veness. Language modeling is compression . 2023. arXiv: 2309.10668 .

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 'BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding.' NAACL-HLT . 2019, pages 4171-4186.

Vijay Prakash Dwivedi and Xavier Bresson. 'A generalization of transformer networks to graphs.' AAAI 2021 Workshop on Deep Learning on Graphs: Methods and Applications . 2020. arXiv: 2012.09699 .

Vijay Prakash Dwivedi, Chaitanya K Joshi, Thomas Laurent, Yoshua Bengio, and Xavier Bresson. 'Benchmarking graph neural networks.' Journal of Machine Learning Research 24.43 (2023), pages 1-48.

Vijay Prakash Dwivedi, Ladislav Rampášek, Mikhail Galkin, Ali Parviz, Guy Wolf, Anh Tuan Luu, and Dominique Beaini. 'Long Range Graph Benchmark.' Neural Information Processing Systems (NeurIPS 2022), Track on Datasets and Benchmarks (2022).

Jerry A Fodor and Ernest Lepore. The compositionality papers . Oxford University Press, 2002.

Jerry A Fodor and Zenon W Pylyshyn. 'Connectionism and cognitive architecture: A critical analysis.' Cognition 28.1-2 (1988), pages 3-71.

Tommaso Furlanello, Zachary Lipton, Michael Tschannen, Laurent Itti, and Anima Anandkumar. 'Born again neural networks.' ICML . 2018.

Robert Geirhos, Jörn-Henrik Jacobsen, Claudio Michaelis, Richard Zemel, Wieland Brendel, Matthias Bethge, and Felix A Wichmann. 'Shortcut learning in deep neural networks.' Nature Machine Intelligence 2.11 (2020), pages 665-673.

Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. 'Neural message passing for quantum chemistry.' ICML . 2017.

Jonathan Godwin, Michael Schaarschmidt, Alexander L Gaunt, Alvaro Sanchez-Gonzalez, Yulia Rubanova, Petar Veliˇ ckovi´ c, James Kirkpatrick, and Peter Battaglia. 'Simple GNN Regularisation for 3D Molecular Property Prediction and Beyond.' ICLR . 2022.

Muhammad Waleed Gondal, Manuel Wuthrich, Djordje Miladinovic, Francesco Locatello, Martin Breidt, Valentin Volchkov, Joel Akpo, Olivier Bachem, Bernhard Schölkopf, and Stefan Bauer. 'On the Transfer of Inductive Bias from Simulation to the Real World: a New Disentanglement Dataset.' NeurIPS . 2019.

Anirudh Goyal and Yoshua Bengio. 'Inductive biases for deep learning of higher-level cognition.' Proceedings of the Royal Society A 478.2266 (2022).

Arthur Gretton, Olivier Bousquet, Alex Smola, and Bernhard Schölkopf. 'Measuring Statistical Dependence with Hilbert-Schmidt Norms.' Algorithmic Learning Theory . 2005.

Caglar Gulcehre, Tom Le Paine, Srivatsan Srinivasan, Ksenia Konyushkova, Lotte Weerts, Abhishek Sharma, Aditya Siddhant, Alex Ahern, Miaosen Wang, Chenjie Gu, Wolfgang Macherey, Arnaud Doucet, Orhan Firat, and Nando de Freitas. Reinforced Self-Training (ReST) for Language Modeling . 2023. arXiv: 2308.08998 .

Will Hamilton, Zhitao Ying, and Jure Leskovec. 'Inductive representation learning on large graphs.' NeurIPS . 2017.

Serhii Havrylov and Ivan Titov. 'Emergence of language with multi-agent games: Learning to communicate with sequences of symbols.' NeurIPS . 2017.

Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. 'Masked autoencoders are scalable vision learners.' CVPR . 2022.

Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. 'Momentum contrast for unsupervised visual representation learning.' CVPR . 2020.

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. 'Deep Residual Learning for Image Recognition.' CVPR . 2016.

Irina Higgins, David Amos, David Pfau, Sebastien Racaniere, Loic Matthey, Danilo Rezende, and Alexander Lerchner. 'Towards a definition of disentangled representations' (2018). arXiv: 1812.02230 .

Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. 'Distilling the knowledge in a neural network.' 2.7 (2015). arXiv: 1503.02531 .

Sepp Hochreiter and Jürgen Schmidhuber. 'Long short-term memory.' Neural computation 9.8 (1997), pages 1735-1780.

Charles F Hockett. 'The origin of speech.' Scientific American 203.3 (1960), pages 88-97.

Weihua Hu, Matthias Fey, Hongyu Ren, Maho Nakata, Yuxiao Dong, and Jure Leskovec. 'OGB-LSC: A Large-Scale Challenge for Machine Learning on Graphs.' NeurIPS Datasets and Benchmarks Track . 2021.

Weihua Hu, Matthias Fey, Marinka Zitnik, Yuxiao Dong, Hongyu Ren, Bowen Liu, Michele Catasta, and Jure Leskovec. 'Open graph benchmark: Datasets for machine learning on graphs.' NeurIPS . 2020.

Weihua Hu, Bowen Liu, Joseph Gomes, Marinka Zitnik, Percy Liang, Vijay Pande, and Jure Leskovec. 'Strategies for Pre-training Graph Neural Networks.' ICLR . 2020.

Eric Jang, Shixiang Gu, and Ben Poole. 'Categorical reparameterization with Gumbel-softmax' (2016). arXiv: 1611.01144 .

Thomas N. Kipf and Max Welling. 'Semi-Supervised Classification with Graph Convolutional Networks.' ICLR . 2017.

Simon Kirby, Hannah Cornish, and Kenny Smith. 'Cumulative cultural evolution in the laboratory: An experimental approach to the origins of structure in human language.' PNAS 105.31 (2008), pages 10681-10686.

Simon Kirby, Monica Tamariz, Hannah Cornish, and Kenny Smith. 'Compression and communication in the cultural evolution of linguistic structure.' Cognition 141 (2015), pages 87102.

Devin Kreuzer, Dominique Beaini, William L. Hamilton, Vincent Létourneau, and Prudencio Tossou. 'Rethinking Graph Transformers with Spectral Attention.' NeurIPS . 2021.

Greg Landrum et al. 'Rdkit: Open-source cheminformatics software' (2016).

Samuel Lavoie, Christos Tsirigotis, Max Schwarzer, Ankit Vani, Michael Noukhovitch, Kenji Kawaguchi, and Aaron Courville. 'Simplicial embeddings in self-supervised learning and downstream classification' (2022). arXiv: 2204.00616 .

Bo Li, Yifei Shen, Yezhen Wang, Wenzhen Zhu, Dongsheng Li, Kurt Keutzer, and Han Zhao. 'Invariant information bottleneck for domain generalization.' AAAI Conference on Artificial Intelligence . 2022.

Ming Li and Paul Vitányi. An Introduction to Kolmogorov Complexity and Its Applications . 4th edition. Springer, 2019.

Dianbo Liu, Alex M Lamb, Kenji Kawaguchi, Anirudh Goyal, Chen Sun, Michael Curtis Mozer, and Yoshua Bengio. 'Discrete-valued neural communication.' NeurIPS . 2021.

Ilya Loshchilov and Frank Hutter. 'Decoupled Weight Decay Regularization.' ICLR . 2019.

Yuchen Lu, Soumye Singhal, Florian Strub, Aaron Courville, and Olivier Pietquin. 'Countering language drift with seeded iterated learning.' ICML . 2020.

Aman Madaan, Niket Tandon, Prakhar Gupta, Skyler Hallinan, Luyu Gao, Sarah Wiegreffe, Uri Alon, Nouha Dziri, Shrimai Prabhumoye, Yiming Yang, Shashank Gupta, Bodhisattwa Prasad Majumder, Katherine Hermann, Sean Welleck, Amir Yazdanbakhsh, and Peter Clark. Self-refine: Iterative refinement with self-feedback . 2023. arXiv: 2303.17651 .

Loic Matthey, Irina Higgins, Demis Hassabis, and Alexander Lerchner. dSprites: Disentanglement testing Sprites dataset . 2017. URL: https://github.com/deepmind/dspritesdataset/ .

Hossein Mobahi, Mehrdad Farajtabar, and Peter Bartlett. 'Self-distillation amplifies regularization in Hilbert space.' NeurIPS . Volume 33. 2020, pages 3351-3361.

Evgenii Nikishin, Max Schwarzer, Pierluca D'Oro, Pierre-Luc Bacon, and Aaron Courville. 'The primacy bias in deep reinforcement learning.' ICML . 2022.

Judea Pearl and Dana Mackenzie. The Book of Why: The New Science of Cause and Effect . Basic Books, 2018.

Aahlad Puli, Nitish Joshi, He He, and Rajesh Ranganath. 'Nuisances via Negativa: Adjusting for Spurious Correlations via Data Augmentation' (2022). arXiv: 2210.01302 .

Jack Rae. Compression for AGI . Stanford MLSys Seminar. 2023. URL: https://www. youtube.com/watch?v=dO4TPJkeaaU .

Sai Rajeswar, Pau Rodriguez, Soumye Singhal, David Vazquez, and Aaron Courville. 'Multilabel iterated learning for image classification with label ambiguity.' CVPR . 2022.

Ladislav Rampášek, Mikhail Galkin, Vijay Prakash Dwivedi, Anh Tuan Luu, Guy Wolf, and Dominique Beaini. 'Recipe for a General, Powerful, Scalable Graph Transformer.' NeurIPS . 2022.

Yi Ren, Shangmin Guo, Matthieu Labeau, Shay B. Cohen, and Simon Kirby. 'Compositional languages emerge in a neural iterated learning model.' ICLR . 2020.

Yi Ren, Shangmin Guo, and Danica J. Sutherland. 'Better Supervisory Signals by Observing Learning Paths.' ICLR . 2022.

Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, Alexander C. Berg, and Li Fei-Fei. 'ImageNet Large Scale Visual Recognition Challenge.' IJCV 115.3 (2015), pages 211-252.

Sandro Salsa. Partial differential equations in action: from modeling to theory . Volume 99. Springer, 2016.

Bernhard Schölkopf and Alexander J Smola. Learning with Kernels: Support Vector Machines, Regularization, Optimization, and Beyond . MIT Press, 2002.

Lukas Schott, Julius Von Kügelgen, Frederik Träuble, Peter Vincent Gehler, Chris Russell, Matthias Bethge, Bernhard Schölkopf, Francesco Locatello, and Wieland Brendel. 'Visual Representation Learning Does Not Generalize Strongly Within the Same Domain.' ICLR . 2022.

Dino Sejdinovic, Bharath Sriperumbudur, Arthur Gretton, and Kenji Fukumizu. 'Equivalence of distance-based and RKHS-based statistics in hypothesis testing.' The Annals of Statistics 41.5 (2013), pages 2263-2291.

Yuge Shi, Imant Daunhawer, Julia E Vogt, Philip HS Torr, and Amartya Sanyal. 'How robust are pre-trained models to distribution shift?' (2022). arXiv: 2206.08871 .

Rui Shu, Yining Chen, Abhishek Kumar, Stefano Ermon, and Ben Poole. 'Weakly Supervised Disentanglement with Guarantees.' ICLR . 2020.

Ravid Shwartz-Ziv and Naftali Tishby. 'Opening the black box of deep neural networks via information' (2017). arXiv: 1703.00810 .

David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam, Marc Lanctot, Sander Dieleman, Dominik Grewe, John Nham, Nal Kalchbrenner, Ilya Sutskever, Timothy Lillicrap, Madeleine Leach, Koray Kavukcuoglu, Thore Graepel, and Demis Hassabis. 'Mastering the game of Go with deep neural networks and tree search.' Nature 529.7587 (2016), pages 484-489.

Ilya Sutskever. An observation on Generalization . Simons Institute workshop on Large Language Models and Transformers. 2023. URL: https://simons.berkeley.edu/talks/ ilya-sutskever-openai-2023-08-14 .

Gábor J. Székely, Maria L. Rizzo, and Nail K. Bakirov. 'Measuring and testing dependence by correlation of distances.' The Annals of Statistics 35.6 (2007), pages 2769-2794.

Shantanu Thakoor, Corentin Tallec, Mohammad Gheshlaghi Azar, Mehdi Azabou, Eva L Dyer, Remi Munos, Petar Veliˇ ckovi´ c, and Michal Valko. 'Large-Scale Representation Learning on Graphs via Bootstrapping.' ICLR . 2022.

Naftali Tishby, Fernando C Pereira, and William Bialek. 'The information bottleneck method.' Allerton Conference on Communication, Control, and Computing . 1999. arXiv: physics/ 0004057 .

Aaron van den Oord, Oriol Vinyals, and Koray Kavukcuoglu. 'Neural discrete representation learning.' NeurIPS . 2017.

Ankit Vani, Max Schwarzer, Yuchen Lu, Eeshan Dhekane, and Aaron Courville. 'Iterated learning for emergent systematicity in VQA.' ICLR . 2021.

Pascal Vincent, Hugo Larochelle, Yoshua Bengio, and Pierre-Antoine Manzagol. 'Extracting and composing robust features with denoising autoencoders.' ICML . 2008.

Ronald J Williams. 'Simple statistical gradient-following algorithms for connectionist reinforcement learning.' Reinforcement Learning (1992), pages 5-32.

Zhenqin Wu, Bharath Ramsundar, Evan N Feinberg, Joseph Gomes, Caleb Geniesse, Aneesh S Pappu, Karl Leswing, and Vijay Pande. 'MoleculeNet: a benchmark for molecular machine learning.' Chemical Science 9.2 (2018), pages 513-530.

Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. 'How Powerful are Graph Neural Networks?' ICLR . 2019.

Zhenlin Xu, Marc Niethammer, and Colin A Raffel. 'Compositional generalization in unsupervised compositional representation learning: A study on disentanglement and emergent language.' NeurIPS . 2022.

Hattie Zhou, Ankit Vani, Hugo Larochelle, and Aaron Courville. 'Fortuitous Forgetting in Connectionist Networks.' ICLR . 2022.

Jinghao Zhou, Chen Wei, Huiyu Wang, Wei Shen, Cihang Xie, Alan Yuille, and Tao Kong. 'Image BERT Pre-training with Online Tokenizer.' ICLR . 2022.

Algorithm 1: Proposed IL-SEM algorithm

A The Ladder of Compositionality

To figure out what z we need in order to generalize well compositionally, we propose the 'ladder of compositionality' in Figure 1. To justify our claims, this appendix will discuss how we could climb the ladder step by step by analyzing how the corresponding requirements are generated. A formal definition of the compositional mapping in terms of group theory, which is necessary for reaching the final stage of the ladder, is also provided. In short, we find only relying on the mutual information between z and G (or between x and y ) cannot reach the final rung of the ladder: we need other inductive biases, which is the main motivation of this paper.

A.1 Stage I: z misses some important information in G

The learned representation would have H ( G | z ) > 0 at this stage. As Y = GenY ( G , ϵ ) and ˆ Y = g ( z ) are assumed to be invertible (here Y and ˆ Y are random variables), this condition can be rewritten as I ( Y ; ˆ Y ) < H ( Y ) 6 . Hence following the analysis in [69], a model with such an encoder z = h ( x ) even cannot achieve high enough training performance, let alone generalizing to unseen

test distributions. This condition might occur when the model underfits the training data, e.g., at the beginning of training or the model's capacity is too small. To make an improvement, one can increase the model size or train longer.

A.2 Stage II: z not only contains all information in G but also some in O

At this stage, the learned representation would have H ( G | z ) = 0 and H ( z | G ) > 0 . This means h ( x ) remembers additional information that G doesn't have, e.g., noisy information in O . From H ( G | z ) = 0 , we know H ( Y | ˆ Y ) = 0 and hence I ( Y ; ˆ Y ) = H ( Y ) . Then the model would have perfect training performance and could also be able to generalize well when training and test datasets share the same distribution. However, when facing the out-of-distribution generalization problem, especially when a spurious correlation exists between some factors in O and G , the extra information learned by z can mess the predictions up. Such a phenomenon is named 'short-cut learning' and is quite common in many deep-learning systems [20]. For example, if the background strongly correlates with the object in the training set (e.g., a cow usually co-occurs with the grass while a seagull usually co-occurs with the beach), the DNN then tends to rely more on these 'short-cut' features (e.g., the background rather than the object) during training. If such correlations disappear or reverse in the test set, the models relying on factors in O cannot generalize well to a new distribution.

Making improvement - loss based on information-bottleneck To make an improvement, the model should eliminate the task-irrelevant information as much as possible. Based on this principle, authors of [74] propose to minimize the following information bottleneck equations:

which means the learned z should extract as much as information from Y (or equivalently, G ) and forget as much as irrelevant information about X (i.e., those in O ). This method is also widely applied in other relevant tasks, like domain adaptation [46], invariant risk minimization [1], and etc.

Making improvement - data augmentation Another simple and efficient way to make improvements is data augmentation: one can identify some task-irrelevant factors in O and design specific data augmentation methods to teach the model to be insensitive to them. For example, if we believe that the label of an image should be irrelevant to color jittering, random cropping, rotation, flipping, etc., we can apply random augmentations during training and treat the differently augmented x and x ′ as the same class. Then the model would inherently learn to be insensitive to such factors and hence forget the corresponding information. One interesting thing about data augmentation is that it can be designed and applied in a reverse direction, i.e., we can break some semantic factors and train the model to be insensitive to the broken samples. For example, believing the shape of the image and order of the words are semantic factors in G , the authors of [56] propose to randomly rotate the image patches or words to make negative samples. Those models that perform well on such negative samples are more likely to rely on the factors in O .

Making improvement - auxiliary task design, e.g., SSL Furthermore, one can also consider designing auxiliary tasks in addition to the downstream task, e.g., pretrain using self-supervised learning (SSL) and finetune on the target task. In [67], the authors empirically show that the representations learned via SSL usually generalize better than those learned via supervised learning when facing OOD downstream problems, even though the models are trained using a similar amount of data samples. There are also some works demonstrating that SSL representations encode more semantic information about the input image [10], which is a sign that auxiliary tasks like SSL can introduce extra biases that favor information in G .

̸

Consider the first group of SSL methods, which are usually based on contrastive loss, e.g., SimCLR [11], MoCo [30], etc. These methods usually require h ( x i ) and h ( x ′ i ) to be similar while h ( x i ) and h ( x j ) , i = j to be distinct, where x i is the anchor input, x ′ i is the augmentation of it, and x j is another different image. The carefully designed augmentation can encourage the model to ignore some task-irrelevant factors that belong to O . Imagine x ′ i is generated by deleting the background of x i . As the training enforces d z ( h ( x ) , h ( x ′ )) to be small, the learned model will then become insensitive to the information in the background, and hence avoid relying on this 'short-cut' feature. Note that the contrastive SSL is utilizing the bias from data augmentation in a more aggressive way: the SSL algorithm will tell the model that x and x ′ are the same image while the supervised learning only inform the model that x and x ′ belong to the same class .

Another line of SSL is built on reconstruction tasks, like denoising auto-encoder (DAE [77]) and masked auto-encoder (MAE [29]). The h ( x ) in these methods is usually trained in an auto-encoder fashion: using a reconstruction network r ( z ) : Z → X , we require the reconstructed x recon = ( r ◦ h )( x ) to be similar to the original input x . As the above equation has a trivial solution, i.e., ( r ◦ h )( · ) = identity , which means h ( x ) copies every details about x (including all O and G ), some early works like denoising auto-encoder propose to add noise on x to encourage non-trivial solutions. Depending on the noise we introduce, the model will learn to ignore different factors accordingly, which seems quite similar to the data augmentation mentioned in contrastive SSL methods. To encourage h ( x ) to extract more useful semantic information in G , methods like MAE [29] propose to mask most of the patches of the input image and try to make reconstructions based on the remaining patches. Such methods exhibit amazing reconstruction performance (not in terms of high resolution, but the precise semantic reconstruction), which also implies that the h ( x ) trained in this way is capable of extracting high-level semantic generating factors (those are likely in G ). Furthermore, methods like BEiT [4] and iBOT [83] also patchify and mask the input images, and concurrently, impose extra constraints on z by comparing them with the descriptions generated by the big language model. Such designs also encourages h ( x ) to extract high-level semantic information, as illustrated in [10].

In summary, as the SSL algorithms learn good z by designing loss or tasks on it, we can introduce extra inductive bias via auxiliary tasks. By requiring z to be invariant when adding noise or conducting data augmentation on x , the task-irrelevant information can be ruled out during learning. By requiring z contains the necessary information for reconstruction when only part of x is observable, the taskrelevant semantic information can be highlighted during learning. By combining these principles, it is possible for us to learn good h ( x ) that only extracts information in G . With the help of these methods, our z might learn exactly all information in G , which means the third rung is achieved.

A.3 Stage III: z learns exactly all information in G, i.e., h ( · ) leads to a bijection

̸

Starting from Stage II, if we can design clever training methods (e.g., adding regularization, data augmentation, auxiliary task, etc.) and make our h ( x ) to be insensitive to some O , ideally, we can learn an almost perfect encoder that extracts exactly all information contained in G . In this case, we have H ( G | z ) = 0 and H ( z | G ) = 0 , i.e., h ( · ) leads to a bijection between z -space and G -space, which is denoted as z ↔ G . Ideally, such z can generalize well even when P train = P test , as long as all the concepts in the test set are seen by h ( x ) during training, i.e., supp [ P test ( G )] ⊆ supp [ P train ( G )] . However, we speculate that even h ( x ) on Stage III will struggle in the comp-gen problem, as the problem assumes supp [ P test ( G )] ∩ supp [ P train ( G )] = ∅ . In other words, we need the model to decompose and recombine the learned concepts in a systematic way, which is similar to the groundtruth-generating mechanism. To achieve this goal, we need to consider how to achieve Stage IV.

A.4 Stage IV: h ( · ) leads to a isomorphism bijection between z and G

To generalize well compositionally, we not only need z contains exactly all information in G , the structure of G should also be embodied in z , which means h ( · ) should lead to an isomorphism bijection between z -space and G -space (i.e., z ⇔ G ). Specifically, we need:

Hypothesis 1 in detail. To generalize well compositionally, we need z ⇔ G , which requires:

z ↔ G , i.e., h ( · ) leads to a bijection between z and G ;

P ( z ) and P ( G ) factorize in a similar way;

Each z i maps to some G w i , where w is a permutation vector of length m , and such a mapping is invariant 7 ;

For each z i , the mapping between z i = k and G w i = u k is invariant 8 , where u is a permutation vector of length v w i .

We call such mappings compositional and call the other bijections holistic . To better understand the difference between these two types of mappings, we consider a simple example where the generating factors are G = [ G 1 , G 2 ] , G 1 = { blue, red } and G 2 = { circle, box } , and the representations are

00

01

10

11

Shape

Color

𝑧1

𝑧2

00

01

10

11

(b). Holistic Mapping

z = [ z 1 , z 2 ] , z 1 = { 0 , 1 } and z 2 = { 0 , 1 } . Hence the space of G is { blue circle, blue box, red circle, red box } , and the space of z is { 00 , 01 , 10 , 11 } . To generate a compositional mapping, we can first decide the meaning of z i , e.g., z 1 represents the shape and z 2 represents the color. After that, we assign vocabularies in z i to denote different meanings, e.g., z 1 = 0 → circle, z 1 = 1 → box, z 2 = 0 → blue, and z 2 = 1 → red. Combining these two steps, we generate a compositional mapping, as illustrated in Figure 5(a). Any bijections that cannot be decomposed in this way are holistic mappings, like Figure 5(b). Obviously, the compositional mappings can generalize compositionally while the holistic ones cannot. However, finding a compositional mapping from all possible bijections is a hard problem, as the number of compositional mappings is much smaller than the holistic ones: assuming we have m = 3 different G i and each with v = 3 possible values, there are roughly 10 3 compositional mappings ( m !( v !) m ), and 10 28 holistic mappings ( ( v m )! -m !( v !) m ).

Another important characteristic of compositional mappings is the 'distance-preserving' ability. In other words, they will map two x with similar G to two similar z in the representation space, i.e., z preserve the topological structure of G . We can use topological similarity ( ρ , defined in Equation 2) to quantify how compositional a mapping is. Usually, for all bijections, mappings with larger ρ are more compositional.

To sum up the aforementioned four stages of the learned z , we list the conditions that z must satisfy and their capabilities in Table 3.

A.5 Relationship with disentangled representations

Readers might notice that our problem settings and the requirements for z mentioned in Hypothesis 1 are quite similar to those discussed in disentanglement representation learning [32]. Here we explain the connections and differences between them.

We care more about downstream tasks. In disentanglement representation learning, people focus more on a general property of the learned representation by assuming the existence of independent semantic generating factors. Disentanglement is then considered a desired property of such representations. Works in this line try to formalize this property and believe this property is beneficial for multiple downstream tasks. However, as the downstream task is not incorporated in such analysis, it is hard to conclude whether specific factors are semantic or not (remember the split of O and G highly depends on the task). On the other hand, this paper focuses more on the generalization ability of the learned representations on specific downstream task(s). We directly optimize the loss of the downstream task in the interaction phase, which we believe can regularize the representations more efficiently.

We don't strictly require a disentangled z . Although representations with properties like consistency or restrictiveness mentioned in [68] could be beneficial for the downstream tasks, it seems not necessary to have all of them. In other words, disentanglement is a sufficient but not necessary requirement for a model to generalize well in downstream tasks. That is because the factorization of G could be non-trivial (as nature is not simple). We believe that capturing the hidden structure of G and GenY ( · ) using z and g ( · ) is more important than mapping an involved generating mechanism to a disentangled system. Additionally, implicitly splitting O from G is rather crucial in our settings, which is rarely discussed in the fields of disentanglement representation learning.

We mainly consider discrete factors. Inspired by how human language evolves, it is natural to start from the discrete factors and representation due to the discreteness of human language. Our experimental results also benefit a lot from the discreteness, e.g., using cross-entropy loss to amplify the learning speed advantage, sampling pseudo labels to strengthen the inductive bias, using group theory to formalize the compressibility and Kolmogorov complexity, etc. However, as nature might not be purely discrete, incorporating the continuous latent space is crucial to enlarge the scope of our study. We would leave this in our future work.

B Compositionality, Compressibility, Kolmogorov complexity, and number of active bases

This appendix links several key concepts related to compositional mappings together, i.e., compressibility, Kolmogorov complexity, and number of active bases. The analysis here provides good intuition on why we might expect iterated learning to be helpful in comp-gen.

B.1 Higher compositionality, lower Kolmogorov complexity

We first complete the proof of Proposition 1.

Proposition 1 (Informal) . For m,v ≥ 2 , among all bijections, any compositional mapping has much lower Kolmogorov complexity than a typical non-compositional mapping.

Proof. Recall the fact that any bijection from z to G can be represented by an element in the symmetry group S v m . From the definition of the symmetry group, we know each element in S v m can be represented by a permutation matrix of size v m . As there is only one 1 in each row and column of a permutation matrix, any permutation matrix can be uniquely represented by a permuted sequence of length v m . Specifically, assume we have a sequence of natural numbers { 1 , 2 , ..., v m } , each permuted sequence Perm ( { 1 , 2 , ..., v m } ) represents a distinct permutation matrix, and hence represents a distinct bijection from z to G . In other words, we can encode one bijection from z to G using a sequence of length v m , i.e., Perm ( { 1 , 2 , ..., v m } ) , and bound the corresponding Kolmogorov complexity (in bits) as

As an arbitrary bijection from z to G doesn't have any extra information to improve the coding efficiency, Equation (4) provides an upper bound of the minimal Kolmogorov complexity.

On the contrary, as each compositional mapping can be represented by an element in S v m ⋊ S m , we can encode the mappings more efficiently. Specifically, we need to first use m sequences with length v , i.e., Perm ( { 1 , 2 , ..., v } ) , to represent the assignment of 'words' for each z i . After that, we need one sequence of length m , i.e., Perm ( { 1 , 2 , ..., m } ) to encode the assignment between z i and G j . The corresponding Kolmogorov complexity is then bounded as

Although this is only an upper bound, by a counting argument most such mappings must have a complexity no less than, say, a constant multiple of that bound.

To compare the Kolmogorov complexity, we can define a ratio as γ ≜ K ( bijection ) K ( comp ) . Obviously, when m ≤ v , γ ≥ v m -1 · m 2 , which is larger than 1 as long as m,v ≥ 2 . When m>v , γ ≥ v m log 2 v 2 log 2 m , which is also larger than 1 when m,v ≥ 2 .

Actually, there might be some mappings that are not purely compositional or holistic. For example, we can have a mapping with z i ≤ 10 sharing the reused rules while other z i> 10 doesn't. Then this type of mapping can be represented by an element in S 10 v ⋊ S 10 ⋊ S v m -10 . As a mapping in this subset shares 10 common rules, its Kolmogorov complexity is between K ( bijection ) and K ( comp ) . Intuitively, for all bijections , smaller K ( · ) means higher compressibility and higher compositionality.

B.2 Regularize the Kolmogorov complexity using iterated learning

From the analysis in Appendix A.4 and B.1, we know that finding mappings with lower Kolmogorov complexity is the key to generalizing well compositionally. From existing works in cognitive science, we know iteratively introducing new agents to learn from old agents can impose the compressibility pressure and hence make the dominant mapping more compact after several generations [42]. Although iterated learning reliably prompts the emergence of compositional mapping in lab experiments, directly applying it to deep learning is not trivial: as we are not sure whether the preference for compositionality still exists for the neural agents. Hence in this subsection, we study a simple overparameterized linear model on a 0/1 classification task to show that iterated learning can indeed introduce a non-trivial regularizing effect. Combining with the fact that mappings with lower Kolmogorov complexity are more likely to capture the ground truth generating mechanism and hence generalize better [71], we can conclude that iterated learning is helpful for compositional generalization problems.

Consider a general supervised learning problem, in which we want to learn a mapping f ∈ F : X → Y that could approximate the underlying relationship between random variables X and Y . As we usually have a finite number of training samples and the space of all possible mappings is large, the model could just remember all ( x , y ) pairs in the training set to achieve a perfect training performance. To avoid this trivial solution, we usually expect the optimal f ∗ to have specific properties, e.g., smoothness or Lipschitz continuousness, etc. Hence usually, we want to optimize a problem with a corresponding regularization term:

where R : F → R is regularizing f and ϵ is the training loss tolerance. This regularization term is usually the inner product of f on the functional space, i.e., R ( f ) = ∥ f ∥ H , where H is a reproducing kernel Hilbert space (RKHS) determined by some kernel function κ ( · , · ) . For example, if we consider ∀ ( x, y ) ∈ [0 , 1] 2 and κ ( x, y ) = min( x, y ) , then ∥ f ∥ H = ∥ f ′ ∥ [0 , 1] 2 [64]. In other words, the regularizer will penalize functions with higher first-order derivatives.

In order to make the analysis generalize to other properties of f , we define a linear differential operator L as [ Lf ] ≜ ∫ X u ( x , · ) f ( x ) d x , where u ( · , · ) is a kernel function. Then, the regularization term is:

Substituting this definition back to Equation 6 and then applying the Karush-Kuhn-Tucker (KKT) conditions, the closed-form solution for this optimization problem is (i.e., Proposition 1 in [53]):

where c is a bounded constant, Y = [ y 1 | . . . | y N ] ⊤ is the stacked training labels. The matrix G ∈ R N × N and its vector g x ∈ R N × 1 is defined as:

The g ( x , q ) is the Green's function [63] of this operator and is defined by:

where δ is the Dirac delta function. Following the definition of Green's function, we know G is positive definite and hence decompose it as:

where D = diag ([ d 1 , . . . , d N ]) is determined by its eigenvalues and V contains N corresponding eigenvectors. Now, we can stack the model's prediction for different input samples x n and get the vector form solution of problem 6:

With this solution, following the settings in [53], we can explain where the compressibility pressure (the one that favors mappings with lower Kolmogorov complexity) comes from. Specifically, the optimal model for the first generation is f ∗ 0 = V ⊤ D ( c 0 I + G ) -1 V Y 0 . Then in the following generations, the model in generation t will learn from the predictions of the model in the previous generation. As the problem is identical (the only difference is the labels) for different generalizations, we can have the following recursion formulas:

Solving this yields the expression of the labels in the t -th generation:

where A t ≜ D ( c t I + D ) -1 is a N × N diagonal matrix. Substituting this back to Equation 12, we finally obtain the following expression:

From this solution, the model's prediction at t -th generation can be considered as a weighted combination of transformed Y 0 . The matrix V will first map Y 0 to a space determined by the Green's function. Different dimensions of this space are then rescaled by a diagonal matrix ∏ t i =0 A i . After that, the vector is transformed back to the origin space by multiplying V ⊤ . Among these terms, ∏ t i =0 A i is the only one that depends on t . Recall the definition of A t = D ( c t I + D ) -1 , we can conclude that ∏ t i =0 A i is also a diagonal matrix where each entry has the form like ∏ t d j c t + d j (here d j is the j -th eigenvalue of G ). As c t > 0 , ∀ t , as stated in [53], all the diagonal entries will gradually decrease when t grows. The dimensions with smaller ∏ t d j c t + d j decrease faster, and vice versa. Recall

the role played by this diagonal matrix, we can imagine that the number of active bases in V is decreasing when t grows, which is a strong and unique inductive bias (i.e., compressibility pressure) introduced by this recursive training fashion 9 .

Now, we can link these theoretical analyses to the Kolmogorov complexity and compressibility pressure. The crux is the understanding of 'active bases'. Consider a toy example where G = [ G 1 , G 2 ] , where each G i has 4 possible values (there are N = 16 different objects). Then, to memorize these 16 samples, the model needs 16 bases like 'S 00 → blue circle'. However, for compositional mappings, only 9 bases are enough 10 : because the model reuse some rules.

In summary, when t is small, there is no preference for compositional mappings because the number of active bases is large enough to remember most of the training samples. As t increases, the model then needs to be clever enough to reuse some bases, where the structure of the mapping emerges. If t is too large, where the compressibility pressure is too strong, the model will degenerate into a very naive solution, which is harmful for generalization (hence we need the interaction phase in iterated learning, discussed later).

C Experiments on Controlled Vision dataset

C.1 Experimental Settings

Train split

Test split

Others

3dShapes

dSprites

MPI3D-Real

n

G2

Data generating factors In this paper, we conduct experiments on three vision datasets, i.e., dSprites [52], 3dShapes [9], MPI3D-real [23], where the ground truth G are given. The summary and examples of these datasets are provided in Table 5 and Figure 6.

Here we specify how to split the dataset and generate the downstream labels using 3dShapes as an example. We denote the hue of the floor, wall, and object as G 1 , G 2 , and G 3 , respectively; each has 10 possible values, linearly spaced in [0 , 1] . The object scale, G 4 , has 8 possible values linearly spaced in [0 , 1] . The remaining two factors, object shape (4 possible values) and object orientation (15 possible values), are treated as other factors and merged into O . Data augmentation methods such as adding Gaussian noise, random flipping, and so on, are also merged into O . Under this setting, the universe of G , i.e., G , has 8000 different values, which is further divided into G train and G test . For

the sys-gen problem, we assume G train ∩ G test = ∅ . One measure of the difficulty of the problem is the split ratio, α = |G train | / |G| ; smaller α generally means a more challenging problem.

Data generating mechanisms For the training set, we first select G ∈ G train , and then generate multiple input signals using x = GenX ( G , O ) , where O is uniformly random. The x in the test set is generated in a similar way, but without including data augmentation in O . Labels for all pairs in the dataset are generated by y = GenY ( G , ϵ ) . The main downstream task we study here is regression: GenY ( G , ϵ ) = a ⊤ G + ϵ , where a = [ a 1 , ..., a m ] ⊤ is a column vector and all a i are chosen from [0 , 1] 11 . These examples assume y is generated by a simple combination of different G i , where recovering the generating factors is necessary to generalize well compositionally.

Model and training settings The model structure for this section is illustrated in Figure 2. We consider using a randomly initialized 4-layer CNN for dSprites and ReNet18 [31] for 3dShapes and MPI3D. Unless otherwise specified, we consider a linear head g ( z ) , and a typical L ds , i.e., cross-entropy loss for classification and mean square error loss for regression. The networks are optimized using a standard SGD optimizer with a learning rate of 10 -3 and a weight decay rate of 5 ∗ 10 -4 . Actually, we find the results are insensitive to these settings.

C.2 More Results

C.2.1 Some interesting observations

0

200

400

600

800

1000

Training epochs

10

3

10

2

train

test

0

200

400

600

800

1000

Training epochs

10

3

10

2

sample

argmax

0

2

4

6

8

Generations

1.0

1.2

1.4

1.6

1.8

2.0

2.2

Entropy

Teacher

Student

Uniform

0.2

0.4

0.6

0.8

1.0

Predicting probability

0

100

200

300

400

500

600

700

Counts

Gen 1

Gen 3

Gen 9

See the first panel in Figure 7, which demonstrates the training and testing loss when training a model using SEM-IL. In the first generation, we see the training loss is always smaller than the test one. The test loss then plateaus after some epochs, which matches our expectations. However, in the following generations, the test loss would decrease faster than the training one at the beginning of the interaction phase, which is quite counter-intuitive. We would explore why this happens and whether it is a sign of increased topological similarity in the future.

Another observation is about the sampling mechanism applied in the imitation phase. Remember in Algorithm 1, the pseudo labels used in the imitation phase are sampled from the teacher's prediction ¯ z l . Hence if the teacher is confident in some attributes, the generated labels would be consistent in different epochs, and vice versa. From Figure 4, we provide the scatter plot of the correlation between the teacher's confidence and the student's learning speed. Here we verify this hypothesis through an ablation study. Specifically, we replace the sampling procedure with an argmax function, i.e., the teacher always provides the label with the largest predicting probability regardless of its confidence. As illustrated in the second panel in Figure 7, the test performance of the argmax -case is much worse than the standard SEM-IL method.

The benefits introduced by sampling pseudo labels can also be interpreted as we are making selfadapting τ for different input samples in SEM. In the origin SEM, we only have one τ to control the average entropy of the backbone's prediction (lower τ leads to peakier predicting distributions, and vice versa). However, a model trained using SEM-IL equivalently has different τ for different x .

To understand this, we can compare the entropy of h 1 ( x ) and h 1 ( x ′ ) , where the teacher is confident in ( x , g 1 = 0) and less confident in ( x ′ , g 1 = 0) . Then during imitation, the student will remember ( x , g 1 = 0) much faster than ( x ′ , g 1 = 0) and hence assign ¯ z 1 = 0 higher probability when the input is x . On the contrary, for the input x ′ , as the student might receive different corresponding pseudo labels during imitation, it would assign a lower probability for ¯ z 1 = 0 given x ′ . As a result, the entropy H (¯ z 1 | x ) and H (¯ z 1 | x ′ ) of the student model would be very different as the teacher have different confidence when generating the pseudo labels, which is equivalent as automatically selecting different τ for x and x ′ .

The last two panels in Figure 7 demonstrate how the entropy of the model's prediction on G changes in different generations. (Remember the output of the backbone after SEM, i.e., z = [¯ z 1 , . . . , ¯ z m ] , are m simplicial vectors with length v .) From the figures, we see the entropy gradually decreases as the training goes on, which means the model is becoming more and more confident in its predictions on average. However, the last panel shows that there are still many unconfident predictions even after converging on the downstream task performance: we speculate that these factors contain little information on G as they might have high entropy.

C.2.2 Influence of task difficulty

As the ground-truth generating factors are accessible for these vision datasets, we could explore how the difficulty of the task, i.e., α , influence the performance gap among different methods. As shown in Table 6, SEM-IL brings a significant enhancement when α is not too big nor too small. When α is big, the test loss might be very small and there is no room to make considerable improvement. On another extreme, if α is too small, some attributes might be never observed during training, which is hard for the model to extract correct generating factors. Remember we expect to generalize to 'red circle' by knowing the concept of red and circle from other combinations: if there is only 'blue' and 'box' in the training set, it is impractical for a model to learn such a concept (maybe the model can extrapolate, but that is another topic). We speculate most real tasks are in a relatively small α regime, as the generating factors and their possible values can be very large. Please note that as the train/test split and GenY all depend on the random seeds, the variance of the numbers in these tables could be large. However, we observe the SEM-IL consistently outperforms other methods under each task generated by different random seeds.

C.3 Other baselines from disentangled representation learning

Most related works consider sys-gen as an NLP or emergent language problem rather than a general representation learning problem, so to the best of our knowledge, there are no specific advanced baselines for this concrete problem. The most related works are some VAE-based methods in disentanglement learning. In this part, we re-implement β -VAE and compare them with the baseline method in our setting. Specifically, we first pre-train the encoder of V AE on the same training set and then attach a task head for the downstream task to the ' µ -part' of the encoder's prediction (note that the encoder will output ' µ -part' and ' σ -part' together). We observe that the V AE-based method performs worse than the baseline method, even though they seem to recover some disentangled factors when conducting latent traversal. This observation is consistent with the findings in [65] and [81], where the authors claim that the disentangled representations are incapable of reliably generalizing to new conceptual combinations. We also speculate that the challenging requirement of the comp-gen problem, i.e., G train ∩G test = ∅ , exacerbates this: there will not be enough variations in x to make the VAE model capture the latent vectors precisely. However, as VAE is also an encoder-decoder system, it is possible to combine SEM-IL with it, which is left for our future work.

D Experimental Settings and More Results on Molecular Graph Dataset

Train:

Valid:

Test:

(a). The standard scaffold split

(b). A more challenging comp-gen split

Step2: calculate the similarity score for each x in training set using the fingerprint vector

Step3: order all training x by the score

Step4: prune 𝜂 % training samples with highest scores

Step1:

Fingerprint

[0 1 0 1 0 0 0 1 0 1]

In this part, we provide an overview of the molecular graph learning dataset we used in this paper. Ogbg-molhiv and ogbg-molpcba [37] are molecular property prediction datasets proposed by MoleculeNet and then adopted by open graph benchmark (OGB) project [79]. The molhiv dataset contains roughly 40K samples, and the task is to predict whether a molecule is capable of inhibiting HIV replication (i.e., a binary classification task). The molpcba dataset is more complex, as it contains roughly 400K samples, and the target is to predict 128 different bioassays, which is a multi-task binary classification task. Ogb-PCQM4Mv2 [36] is a large-scale molecular dataset that contains roughly 4000K samples. The task is to predict the HOMO-LUMO gap (i.e., a regression task), which is a quantum physical property that is hard to calculate in traditional methods. As the test split is private, we treat the original validation split as the test split and only report the performance on it in the paper. All of the aforementioned datasets use scaffold splitting, which separates structurally different molecules into different subsets, as illustrated in the left panel in Figure 8. Under such a split, some specific structures in the test set might never occur during training, which makes it a good testbed for systematic generalization ability.

To make the task more challenging, which could simulate the scenario where supp [ P train ( G )] ∩ supp [ P test ( G )] = ∅ , we prune the training set following a procedure demonstrated in the right panel of Figure 8. Specifically, we first calculate the fingerprint of each x in both training, validation, and test sets using RDKit (please refer to Section 5.2 for more details). Similar to the settings used in Table 2, the fingerprint of each x is defined as FP ( x ) ∈ { 0 , 1 } k . In this vector, FP i ( x ) = 1 means the molecule contain the i -th structure. Then, the score for each x in the training set is defined as how many samples in the validation and test sets share the identical FP ( x ) with it. To prune the training samples which are similar to the validation and test set, we delete η % samples with the highest scores (for Table 1, η = 50 ). If we believe these 10 structures are part of G , the remaining training samples

are more likely to have non-overlapping G compared with the test set, which makes the task a better testbed for systematic generalization.

The implementation of the GCN/GIN backbone used in this work is taken from the open-source code released by OGB [37]. We use the default setting of hyperparameters for all experiments (including baseline, baseline+, and interaction phase of SEM-only and SEM-IL). For the backbone structure, the depth of the GCN/GIN is 5, hidden embedding is 300, the pooling method is taking the mean, etc. For the training on downstream tasks, we use the AdamW [49] optimizer with a learning rate of 10 -3 , and use a cosine decay scheduler to stable the training. For the SEM layer, we search L from [10 , 200] and V from [5 , 100] on the validation set. For the IL-related methods, we select the imitation steps from {1,000; 5,000; 10,000; 50,000; 100,000}.

Peer reviews:
comments: SUMMARY:
This paper proposes an iterated learning method with simplicial embeddings (SEM-IL) for systematic generalization.
The method is inspired by iterated learning for humans, encouraging compressibility and expressivity.
The empirical experiments show the improvement in vision tasks (known latent factors) and real molecular graph prediction tasks (unknown latent factors).
It better portrays the generation factors in vision tasks.
It also finds that the success of iterated learning may be caused by discrete messages.

STRENGTHS:
- The analysis and experiments show the potential of the SEM-IL framework on systematic generalization problems.

- It is based on the cognitive science of iterated learning for humans.

WEAKNESSES:
(1) As mentioned in the conclusion section, it is still unclear why the SEM-IL method enables systematic generalization from a theoretical perspective for deep learning.

(2) More experiments will be helpful.
The paper would be more convincing with a strict theoretical explanation or comprehensive experiments.

(3) As mentioned in line 88, the decomposition into G and O is fundamentally unidentifiable because there can be different ways to define factors that work equally well on training data.
So it would be helpful to summarize what additional information SEM-IL provides to identify them.

(4) It seems source codes are not provided.

QUESTIONS:
Please address the points in the weakness section.

LIMITATIONS:
The limitations are mentioned but not in a dedicated section.
It does not have a dedicated section for social impact.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-06-13

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work addresses the challenge deep neural networks face with systematic generalization. Systematic generalization refers to the ability to apply learned concepts to new, unobserved combinations. The authors draw on a cognitive science theory known as \"iterated learning\", a hypothesized process explaining how human language developed this representation ability. The theory is based on simultaneous pressures towards compressibility and expressivity when transmitting messages. The researchers propose using iterated learning, in combination with deep network models that incorporate simplicial embeddings, to generate nearly discrete messages, thereby improving systematic generalization. The effectiveness of this approach is demonstrated in vision tasks with known latent factors and in molecular graph prediction tasks with unknown latent structures, where it outperforms other methods.\n\nSTRENGTHS:\n1. The paper is well-written. The motivation is clear and the problems are stated clearly.\n\n2. The problem addressed in this work is important, and could be particularly interesting to many communities such as causal learning or domain generalization. \n\n3. The empirical performance looks promising, where the gap between the baseline and the proposed approach is non-negligible.\n\nWEAKNESSES:\n1. The method section is not described clearly. In particular, in section 3.1, the authors did not explain why IL can be used to improve the performance. Similarly, in section 3.2, I cannot understand the motivation of adopting SEM. Most of the details are deferred to Appendix B, however, it is still important to briefly explain why IL and SEM can lead to the desired results.\n\n2. The main experiments are done with synthetic dataset such as 3dShapes or dSprites dataset. It would be interesting to see how the proposed problem can address the distribution shift problem in domain generalization benchmark such as DomainBed. \n\n3. The model attempts to learn the minimal feature to perform prediction. This could fail when there is “short-cut” presenting in the dataset. In particular, there might be “easy feature” that is highly correlated with the true target feature, where the model could learn to use those easy features instead in the proposed setting.\n\n4. The main theoretical results in Appendix B.3 are adopted from previous work [56], where they consider a distillation setting.\n\n5. The main motivation is based on compressibility and expressivity. The same concept is also proposed in various previous works such as Variational Information Bottleneck (VIM), supervised contrastive learning, invariant risk minimization (IRM), or works in causal representation learning. Nevertheless, there is no discussion or comparison to previous works.\n\nQUESTIONS:\n1. There are three analysis provided in Appendix B to explain IL. It is a bit confusing why we need so many explanation here. In particular, it is hard to related different explanation, e.g., connecting the Bayesian analysis with the KRR analysis.\n\n2. How do we guarantee the predictor g is a good classifier in IL? Do we also train g in the IL procedure?\n\n3. Why can’t IL alone improve the performance? According to Appendix B, it seems that IL alone can address the proposed problem in representation learning.\n\nLIMITATIONS:\nThere is no obvious potential negative social impact of this work.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper aims to develop methods to improve systematic generalization. The paper proposes several theoretical criteria related to representation learning to improve systematic generalization. Motivated by these principles, the paper then studies whether two methods, iterative learning (IL) and simplicial embeddings (SEM), and their combination can improve systematic generalization. They evaluate these methods on synthetic vision tasks and several tasks related to modeling molecules. They find that the combination of these two methods leads to the highest performance, on average, across these tasks.\n\nSTRENGTHS:\n* I appreciated the formal treatment of the problem in section 2, the intuition offered in section 3, and the empirical analysis connecting the proposed methods to the various hypotheses in these sections. Otherwise it would have been non-obvious how SEM and IL would be expected to improve systematic generalization.\n* The paper is well written and does a nice job connecting various bodies of work. While the empirical results for the specific proposed methods are not overly convincing (see weaknesses), I could see this work potentially inspiring new ideas.\n* Different from prior work on systematic generalization that proposes specialized models with task-specific inductive biases, SEM and IL are relatively general methods that do not necessarily limit the expressiveness of the underlying model.\n\nWEAKNESSES:\n* While I liked the intuition-building and \"ladder of systematicity\", I have two concerns with it. First, Stages 1 to 3 seem to directly follow from prior work, e.g. the information bottleneck principle (to the author's credit, this appears to be acknowledged in the appendix). Second, maybe misunderstanding on my part, but I was not fully satisfied with the definition of stage 4. As informal intuition, it makes sense that we would want to learn representations `z` that are similarly \"structured\" to the latent factors of the data generating procedure `G`. But it wasn't clear to me which operations are expected to be preserved through the proposed \"isomorphism\". Additionally, while `G` informally represents \"semantic factors\" it wasn't clear if this can be translated to a formal constraint on the structure of `G`, which we want `z` to be isomorphic to? I think it's fine to leave this mostly informal, but since the paper seemed to attempt to formalize this, I was left a bit unsatisfied.\n* It seems like most of the intuition offered for SEM and IL corresponds to stages 1 to 3 of the ladder, which per the first point is closely aligned with criteria proposed by prior work. It wasn't clear why these methods should help achieve stage 4, which seemed more unique to the systematic generalization problem setting.\n* The proposed methods seem to assume a `z` with fixed dimensionality. While this seems reasonable for one of the motivating examples (e.g. `z` representing color and shape), it is less clear how this could be applied to systematic generalization in natural language, which is referenced multiple times in the text. Prior work on systematic generalization in natural language has often considered dynamically-sized tree-structured latent variables.\n* The overall effectiveness of SEM and IL is hard to gauge. The experiments show improvements on synthetic vision tasks and 3 tasks related to modeling molecules. However, it's unclear whether the proposed methods are effective on real-world vision or NLP tasks that are more popularly studied within the ML community.\n\nQUESTIONS:\nSee weaknesses above.\n\nLIMITATIONS:\nI more explicit discussion of limitations could be useful.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIterated learning is hypothesized to help human language develop representation ability. The paper proposes to use iterated learning with deep network models containing simplicial embeddings to obtain approximately discrete messages. They show that this combination of changes improves systematic generalization over other approaches, demonstrating these improvements both on vision tasks with well-understood latent factors and on real molecular graph prediction tasks where the latent structure is unknown.\n\nSTRENGTHS:\nThe paper aims to improve the systematic generalization capability of deep neural networks, which is currently an important topic in the field. The writing and figures are clear and quite easy to follow.\n\nWEAKNESSES:\nPlease see in Questions.\n\nQUESTIONS:\n1. A more detailed explanation for iterated learning, which is currently lacking in the paper, will be appreciated.\n\n2. The intuition of using iterated learning in this paper is merely from results in cognitive science. I have not seen convincing reasons why it is appropriate to incorporate iterated learning in deep neural networks, and what limitations of previous works can be tackled by doing so. The introduction should be rewritten to highlight these points.\n\n3. As illustrated in Figure 3 and Table 2 (and the authors also notified), iterated learning alone may not be enough. This again raises a question: Is iterated learning really necessary for this framework?\n\n4. I have not seen a clear connection between iterated learning and simplicial embedding. As far as I understand, the main theme of the paper is about iterated learning, and simplicial embedding empirically improves the performances further (as mentioned in Section 3.2). That is to say, the combination of iterated learning and simplicial embedding in this paper may come from trial and error rather than a thorough theoretical justification. It would be beneficial if the authors can highlight this connection in their paper. \n\n5. The paper mentions compressibility and expressivity pressures throughout the main text. Can we quantify these pressures to show that the proposed method really helps them emerge through iterated learning?\n\nI would happily increase my score if the authors carefully tackle my concerns.\n\nLIMITATIONS:\nThe authors have included some discussions regarding the limitations of their work in the conclusion section.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-07"}