SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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(\mathbf{x})$  into  $m$  sub-mappings  $h_i(\mathbf{x})$ , which map  $\mathbf{x}$  to  $\bar{z}_i \in [0,1]^v$ . We can understand each  $\bar{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  $\mathbf{G}$ , it is more likely to give idiosyncratic learned  $\bar{z}_i$ , and will lead to low confidence for most samples. On the contrary,  $\bar{z}_i$  belonging to compositional mappings will be more general, and on average tend towards higher confidence.
This measure is also known as the distance correlation [72]; it is a special case of the Hilbert-Schmidt Independence Criterion (HSIC, [25]) for a particular choice of kernel based on  $d_{z}$  and  $d_{G}$  [66].
Figure 4: First panel: correlation between teacher's confidence and student's learning speed for each  $(\mathbf{x},\bar{z}_i)$ ;  $\bar{z}_i$  is the prediction of the  $l$ -th attribute in imitation phase. "Consistent" means the student makes the same prediction as the teacher. Other panels: learning curves of the student's predictions.
The imitation phase reinforces this bias when the new student learns from the sampled pseudo labels  $\mathbf{g}_i$  from the teacher's prediction  $\bar{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(\bar{z}_i \mid \mathbf{x})$  is flat,  $\mathbf{g}_i$  could change across epochs. This makes it hard for the student to remember any related  $(\mathbf{x}, \mathbf{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",  $\mathbf{g}_1 = 0$ ), ("red square",  $\mathbf{g}_1 = 1$ ), ("red square",  $\mathbf{g}_1 = 2$ ) in three consecutive epochs.
Furthermore, using the sampled pseudo-labels can help the student to align the learned  $(\mathbf{x},\mathbf{g}_i)$  better. Assume during training, the student already remembers some pairs like ("blue circle",  $\mathbf{g}_1 = 0$ ), ("blue square",  $\mathbf{g}_1 = 0$ ), ("blue star",  $\mathbf{g}_1 = 0$ ), but the teacher is not confident in ("blue apple",  $\mathbf{g}_1$ ), perhaps because apples are rarely blue. Following the analysis above, as  $P(\bar{z}_1|$  blue apple) is flat, the teacher may generate  $\mathbf{g}_1 \neq 0$  a significant portion of the time. However, if the teacher happens to generate  $\mathbf{g}_1 = 0$  at some point, the student would learn ("blue apple",  $\mathbf{g}_1 = 0$ ) faster than those with  $\mathbf{g}_1 \neq 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]",  $\mathbf{g}_1 = 0$ ) will also promote the learning of ("blue apple",  $\mathbf{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  $(\mathbf{x},\mathbf{g}_i)$ . Specifically, for each  $\mathbf{x}$ , the teacher makes  $m$  predictions with the corresponding categorical distribution  $\bar{z}_i$ ,  $i\in [m]$ . For each  $(\mathbf{x},\bar{z}_i)$ , the confidence is measured by the negative logarithm of the teacher's predicted probability,  $-\log [\bar{z}_i]_j$  where  $\hat{j}\in \mathrm{argmax}_{j}[\bar{z}_{i}]_{j}$ . The learning speed of  $(\mathbf{x},\mathbf{g}_i)$  is measured by the integral of the student's prediction with training time  $t$ , i.e.,  $\sum_{t = 0}[\hat{z}_i(t)]_j$ , where  $j$  is the value provided by the teacher and  $\hat{z}_i(t)$  is student's prediction at time  $t$ . As illustrated in the first panel of Figure 4, the  $(\mathbf{x},\mathbf{g}_i)$  with higher confidence are usually learned faster by the student.
We also provide the learning curves of  $(\mathbf{x},\mathbf{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  $[\hat{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  $(\mathbf{x},\mathbf{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  $(\mathbf{x},\mathbf{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  $(\mathbf{x},\mathbf{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  $\bar{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
Table 1: Downstream performance on different tasks. The numbers of AUROC and average precision are in percent form. For PCQM, we report the validation performance, as the test set is private and inaccessible. Means and standard deviations of 5 seeds are given. Valid/test-full means the standard train/val/test split provided by the dataset. Valid/test-half means we train the model on half of the training data which is less similar to the validation and test sets. See Appendix D for more.
Model and Algorithm	molhiv (AUROC ↑)	molpca (Avg.Precision ↑)	PCQM (MAE ↓)
Valid-full	Test-full	Valid-half	Test-half	Valid-full	Test-full	Valid-half	Test-half	Valid
GCN	Baseline	82.41±1.14	76.25±0.38	75.65±0.91	72.31±1.86	21.44±0.25	22.13±0.46	21.13±0.38	20.78±0.62	0.125±0.002
Baseline+	81.61±0.63	75.58±1.00	73.23±0.75	72.17±1.02	22.31±0.34	22.68±0.30	21.01±0.45	20.60±0.37	0.118±0.004
SEM-only	84.00±1.10	78.40±0.67	74.84±1.57	72.81±2.32	26.39±0.66	25.89±0.71	22.79±0.91	22.09±1.02	0.106±0.002
SEM-IL	84.89±0.68	79.09±0.67	78.48±0.67	74.02±0.78	28.81±0.72	27.15±0.74	22.59±0.84	21.90±0.81	0.102±0.005
GIN	Baseline	81.76±1.04	76.99±1.42	76.95±1.40	71.63±2.21	23.09±0.32	22.64±0.49	20.52±0.39	20.15±0.42	0.109±0.003
Baseline+	81.55±0.72	77.01±0.94	74.77±1.62	69.75±3.10	23.85±0.29	22.91±0.40	21.71±0.12	20.98±0.27	0.108±0.003
SEM-only	83.05±0.90	78.21±0.78	76.29±2.06	72.70±4.94	26.01±0.52	25.66±0.47	22.26±0.39	21.50±0.48	0.106±0.004
SEM-IL	83.32±1.51	78.61±0.73	78.06±1.24	72.89±0.48	29.30±0.48	28.02±0.61	24.41±0.47	23.89±0.77	0.098±0.005
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  $\mathbf{z}$  by Meaningful Structures
In the controlled vision examples, we know that SEM-IL not only enhances the downstream performance, but also provides  $\mathbf{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  $\mathbf{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  $\mathbf{x}$  by directly observing its 2D structure. With the help of the RDKit tool [44], we can generate a sequence of labels for each  $\mathbf{x}$ , which is usually known as the "fingerprint" of molecules (denoted  $\mathrm{FP}(\mathbf{x})\in \{0,1\}^k$ , indicating whether each specific structure exists in  $\mathbf{x}$ ). Then, we add a linear head on top of the fixed  $\mathbf{z}$  and train it using a generated training set  $(\mathbf{x},\mathrm{FP}(\mathbf{x})),\mathbf{x}\sim \mathcal{D}_{train}$ , and compare the generalization performance on the generated test set
$(\mathbf{x}, \mathrm{FP}(\mathbf{x})), \mathbf{x} \sim \mathcal{D}_{test}$ . For fair comparison, we set  $m = 30$  and  $v = 10$  to make  $\mathbf{z}$  and  $\mathbf{h}$  be the same width, excluding the influence of the linear head's capacity.
Table 2: AUROC for graph probing based on different  $\mathbf{z}$  ; random guessing would be  $\approx  {0.5}$  .
Sat.Ring	Aro.Ring	Aro.Cycle	Aniline	Ketone	Bicyc.	Methoxy	ParaHydrox.	Pyridine	Benzene	Avg.
Init. base	0.870	0.958	0.811	0.629	0.595	0.615	0.627	0.706	0.692	0.812	0.732
Init. SEM	0.872	0.958	0.812	0.635	0.597	0.638	0.613	0.692	0.683	0.815	0.731
Train on Molhiv	Baseline	0.874	0.948	0.916	0.700	0.717	0.694	0.804	0.740	0.703	0.913	0.801
SEM-only	0.893	0.989	0.938	0.722	0.751	0.779	0.823	0.763	0.763	0.938	0.836
SEM-IL	0.907	0.980	0.967	0.781	0.801	0.794	0.903	0.815	0.869	0.965	0.878
Train on Molpba	Baseline	0.921	0.988	0.968	0.866	0.875	0.835	0.875	0.855	0.856	0.968	0.901
SEM-only	0.942	0.991	0.981	0.888	0.916	0.854	0.921	0.888	0.897	0.980	0.926
SEM-IL	0.940	0.988	0.982	0.910	0.931	0.849	0.912	0.910	0.912	0.981	0.931
Train on 10% pcba	Baseline	0.923	0.980	0.962	0.863	0.857	0.832	0.870	0.833	0.864	0.962	0.895
SEM-only	0.943	0.993	0.989	0.872	0.906	0.835	0.913	0.876	0.900	0.989	0.922
SEM-IL	0.938	0.994	0.985	0.891	0.918	0.847	0.927	0.874	0.907	0.985	0.927
Train on pcba-1task	Baseline	0.892	0.974	0.948	0.723	0.750	0.689	0.845	0.758	0.782	0.947	0.831
SEM-only	0.906	0.989	0.958	0.772	0.809	0.735	0.876	0.770	0.835	0.957	0.861
SEM-IL	0.906	0.988	0.963	0.741	0.851	0.744	0.887	0.765	0.869	0.962	0.867
In the experiments, we use the validation split of molhiv as  $\mathcal{D}_{train}$  and the test split as  $\mathcal{D}_{test}$ , each of which contain 4,113 distinct molecules unseen during the training of  $\mathbf{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  $\mathbf{z}$ : the  $\mathbf{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  $\mathbf{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.
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, self-refinement 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. Self-supervised 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.
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.