How do diffusion models learn and generalize on abstract rules for reasoning?

I think the main weakness of the paper is lack of an interesting and useful conclusion that we previously do not know about diffusion models. The main body of the paper is about experimental results, but I personlly did not find the results particularly surprising. I will follow the abstract to discuss one by one:

1. "diffusion models can synthesize novel samples": this seems not so surprising, because we already know that diffusion models can synthesize novel images (like flying monkeys) not existed before.
2. "memorized and recombined local parts of the training samples": this is an interesting point, because previously we did not know how diffusion models can learn the structure. However, in this paper, they have this claim mainly because of Figure 7, which is a plot showing the model first learns local part, then global parts. It would be better if the authors can provide deeper discussions on this point.
3. "advanced sampling techniques were needed for inpainting": this is also an interesting point, but it seems that here the authors simply mean that Sequential Monte Carlo will give better performance, but did not provide any deep analysis.
4. "pattern completion capability can generalize to rules unseen": this is an interesting point, but it seems that the authors mainly provide some unseen examples in test set, and show that diffusion models can work well on these examples. The fact that models can generalize to unseen cases has been observed by many previous studies, and it would be great if the authors can use Raven's task to provide better intuitions and explanations.
5. "rule representation": this is not in particular surprising, because people previously already observe that deep networks can learn semantic features, or self-supervised learning framework can learn very strong features.

Some of the weaknesses of this work are:

The submitted document needs to be completed.

1. The paper is incomplete. The supplementary material includes sections that are missing content (B.1, B.3, B.4). Some sections are only partially completed, such as B.2.3, which is missing figures and explanations that cannot be linked to experimental results, along with 'XXXXX' entries. Additionally, in section A.2.2, line 748, the authors failed to provide the hyperparameter values. Although these issues are present in the supplementary material, the reviewer believes that the authors should have been more diligent and submitted a complete version of their work.

Clarity and Presentation

The reviewer suggests that the authors revisit their work, as improvements can be made in the presentation of the paper. The figures should be organized more effectively; currently, Figures 2, 3, and others appear to be placeholders with individual images stacked. Additionally, there are other issues with the presentation of the results:

1. In Figure 5.D, not all versions are visible. The label "scale" is unclear in the legend, and it seems different models are referenced. What does "model_class" refer to?
2. In Fig. 2.D, what does the individual held-out rule color correspond to? There are five different cases, but the authors do not mention what each of them represents.

Missing References to Prior Work

Previous research has explored ways to represent Raven's Progressive Matrices (RPMs) in a format suitable for language models. The reviewer suggests that the authors should reference these studies and compare their RPM representation with existing works. It is possible that the proposed representation may not be appropriate for language models. Here are some relevant references: [1], [2], [3]. Additionally, this work introduces a simplified RPM dataset [4], which appears to be quite similar to the proposed representation.

Dataset Concerns

1. The reviewer expresses skepticism regarding the validity and difficulty level of the tested out-of-distribution (OOD) and held-out scenarios. In the proposed simplified dataset, all data is numeric. Therefore, if a model is trained to learn the constancy relation between the numbers on the second channel of the RPM, it may not be very challenging to identify this type of relationship on the first channel. This scenario differs from visual puzzles, where constancy in color is present during training, but testing introduces puzzles with constancy in shape for the first time.

Contribution Concerns

1. The contribution of this work seems to be limited. The reasons are that the authors do not seem to propose a novel component but rather re-train existing models for a different reasoning dataset.

[1] Zhang, Chengru, and Liuyun Wang. "Evaluating Abstract Reasoning and Problem-Solving Abilities of Large Language Models Using Raven's Progressive Matrices." (2024).

[2] Webb, Taylor, Keith J. Holyoak, and Hongjing Lu. "Emergent analogical reasoning in large language models." Nature Human Behaviour 7.9 (2023): 1526-1541.

[3] Ahrabian, K., Sourati, Z., Sun, K., Zhang, J., Jiang, Y., Morstatter, F., & Pujara, J. (2024). The curious case of nonverbal abstract reasoning with multi-modal large language models. arXiv preprint arXiv:2401.12117.

[4] Schug, S., Kobayashi, S., Akram, Y., Sacramento, J., & Pascanu, R. (2024). Attention as a Hypernetwork. arXiv preprint arXiv:2406.05816.

• Formal writing: the authors should carefully check all typos and grammar issues in this paper. Examples include but are not limited to: incorrect openings of quotation marks, incomplete sentences (Section 3.3, around line 180), informal use of abbreviations ("esp."), inconsistent naming (Raven's progressive matrix and Raven's progression matrix).
• Organization and presentation of results: this work includes several crowded figures, when they are referenced in the main text, it is often hard to follow since the text and the reference figure are far away from each other. For example, Section 4.3 involves discussions on Figure 4, Figure 10 and Figure 11, where Figure 10 and 11 are only found in the appendix. These results can potentially be summarized into one or two high-level figures that present the major outcomes, with detailed diagrams or tables included in the appendix. This would help readers locate relevant information easily.
• The claim on comparisons to auto-regressive models: the auto-regressive models used for comparison are GPT2 models, which might be out-dated given that there exist much stronger models. Out-performing GPT2 models does not support the claims on out-performing auto-regressive models.

Literature review misses relevant prior work.

[1] Lachapelle, S., Mahajan, D., Mitliagkas, I. and Lacoste-Julien, S., 2023. Additive decoders for latent variables identification and cartesian-product extrapolation. Advances in Neural Information Processing Systems, 36.

[2] Bansal, A., Schwarzschild, A., Borgnia, E., Emam, Z., Huang, F., Goldblum, M. and Goldstein, T., 2022. End-to-end algorithm synthesis with recurrent networks: Extrapolation without overthinking. Advances in Neural Information Processing Systems, 35, pp.20232-20242.

Both of these papers consider a setting that has certain similarities to the problem that this paper studies. [1] also considers attributes such as location and color.

The experiments of the paper are rather limited as it only evaluates the diffusion models on one specific dataset that is proposed for the first time in the paper. The dataset itself is not large (40 relational rules and 3x3 matrix of panels), and the achieved accuracy is rather high (99.8%) indicating that it is not very challenging. In my view, to draw conclusions about the learning abilities and learning procedures of diffusion models one would need to see more convincing experimental results on a broader set of related tasks.

While the paper contrasts the abilities of diffusion models against that of autoregressive models, and provides some useful empirical evidence about the cases where one performs better than the other, it is hard to be sure that these would generalize beyond the specific task of Raven’s progression matrix. In its current form, I think it would be more appropriate for the paper to change its title and conclusions to reflect that it is only about the Raven’s progression matrix and not the broader topic of abstract reasoning. But, this makes the scope of the paper rather limited. So, I would suggest that the paper expands its experiments to other reasoning tasks, and try to make broader conclusions beyond the RPM task, perhaps for a certain class of reasoning tasks.

Object Tokenization: The study focuses solely on the learning of abstract rules, assuming idealized perception of object features. While this approach isolates rule learning effectively, the paper could benefit from a comparison between pixel tokenization and object/panel tokenization, showing how different tokenization strategies might impact results.

Autoregressive Model Setup: In Section A.3, the authors apply object tokenization in autoregressive modeling by simply combining object features. This factorized approach assumes independence among features within the same object, differing from the diffusion model’s approach which models the joint strictly. It would strengthen the study to model the joint probability of object attributes directly.

RPM for Transfer Learning: The Raven's Progressive Matrices (RPM) test is typically a transfer learning process for humans, where test-takers encounter the task without prior exposure or practice. In contrast, diffusion models in this study train on RPM with extensive sample exposure. It would be valuable to explore how diffusion models trained on natural images could transfer to this task and perform zero-shot rule discovery, likely revealing contrasts with the training approach described in the paper.

Contribution Framing: Some contributions could be framed more effectively. For instance, “showing that unconditional diffusion models can be used for conditional inference, i.e., completing missing panels consistent with the rule” may be perceived as a weaker contribution since converting unconditional diffusion models for inpainting is already well-known. A revised framing of contributions could clarify the study’s unique insights.

• I don't think reasoning is the appropriate term to use, instead, the authors have found a wonderful way to study compositional generalization https://arxiv.org/abs/2307.05596 this has been found to allow diffusion models to even interpolate latent variables https://arxiv.org/abs/2405.19201 and in the specific holdout setting the authors studied (leave-one-out holdout rules), everything is explained by compositional generalization (the semantic debate whether or not this is "true reasoning" is uninteresting to me here). Reasoning might include comp. generalization, but also includes things like multi planning (see e.g. https://www.arxiv.org/abs/2409.13373) so the paper should change its framing of what is being studied
• the paper should repeat the strongest results with increasing subsets of held out rule to see whether the breakdown follows the expectations set in the current interpretation as a sanity check -science is about "kill your babies". this is in particular interesting for the linear probe (how many times must a factor appeat to be visible?)
• I am unsure about the sample sizes chosen being meaningful, please use something like https://cran.r-project.org/web/packages/effectsize/index.html to justify experimental design. where appropriate, statistical tests should be used to justify claims against a concrete null hypothesis
• was only a single seed run for everything (in particular, training the model)? if yes, please rerun/retrain, do a statistical significance test and report uncertainties (I only see this on evaluation right now)