Machine learning for differential equations paves the way for computationally efficient alternatives to numerical solvers, with potentially broad impacts in science and engineering. Though current algorithms typically require simulated training data tailored to a given setting, one may instead wish to learn useful information from heterogeneous sources, or from real dynamical systems observations that are messy or incomplete. In this work, we learn general-purpose representations of PDEs from heterogeneous data by implementing joint embedding methods for self-supervised learning (SSL), a framework for unsupervised representation learning that has had notable success in computer vision. Our representation outperforms baseline approaches to invariant tasks, such as regressing the coefficients of a PDE, while also improving the time-stepping performance of neural solvers. We hope that our proposed methodology will prove useful in the eventual development of general-purpose foundation models for PDEs. Code: https://github.com/facebookresearch/SSLForPDEs.
Paper
Similar papers
Peer review
Summary
The paper presents an approach for self-supervised learning of PDEs. The main approach uses Lie symmetries (similar e.g. to translational symmetries for images) in the solutions of differential equations not only for data augmentation (as has been done before), but for representation learning for downsteam tasks.
Strengths
The paper is a pleasure to read. There are nice diagrams (perhaps even too fancy?) and a nice verbal presentation everywhere. The appendix comes with an introduction into the theory, which is also well written and correct. The literature is thoroughly cited (Of course, there is always more to cite, but the amount is huge and the choice are reasonable.) Many experiments are conducted. Even limitations and drawbacks are clearly stated, such that after reading the paper a really feel what works easily, what works with problems, and what is not yet known. I feel well-informed after reading the paper. Furthermore, the topic is clearly relevant for NeurIPS: it is a non-trivial machine learning question with relevant application domains. In particular the correct usage of PDE-terms is refreshing to see when reviewing NeurIPS papers about PDEs. Here, I could find no problems. All terms were used correctly, even though more could and should have been said about PDEs, but the paper only has 9 pages; more is saind in the appendix. This quality in mathematical writing is not the norm at NeurIPS! And, at the same time, the paper is not technical.
Weaknesses
I could follow the paper. However, I have a strong math background that includes symmetries of PDEs. The authors try really hard (and are quite good) in making everything clear to the audience. However, I guess that much of the NeurIPS audience might have problems with some techniques used in the paper. After reading the paper, it was not clear to me what kind of data is necessary. In one of your experiments, you used data on a grid and were thus able to apply ResNet18. More details in this regard (what is possible? what is reasonable?) would have been nice. Quite a lot of data is necessary for training. If I have 10000 samples of a PDE, I kind of know the PDE anyway. This, together with the high dependency on several hyperparameters and the additional limitation (see below), is the main problem of this paper which make is "just a strong NeurIPS paper" instead of an outstanding paper.
Questions
I guess for a paper of this level ob abstraction, you do not need to define what a group and what a group operation is? (l. 115)
Rating
8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness
4 excellent
Presentation
4 excellent
Contribution
3 good
Limitations
The approach is certainly interesting. However, where I am working with differential equations, one only knows the symmetries from the differential equations. Hence, I rarely see a need for this approach. Otherwise, limitations and open questions are given both in the paper (e.g. l. 99ff, 130ff, etc.)
Opinion after the rebuttals
I have read all reviews and all rebuttals. Having given the most positive rating, I will try to justify my unchanged opinion. - All reviewers agree that the paper is well-written. - All reviewers agree that the method is (mostly) correct. - The main open point regarding correctness is the boundary conditions. In my corner of differential equations, the boundary conditions are more like additional information rather than an integral part of the differential equations. Hence, at least in my corner of differential equations, the raised problem is less problematic. - The primary point of difference is the amount of novelty. Since I feel that differential equations, especially partial differential equations, pose really challenging questions, the level of novelty I see in this paper is sufficient for me to consider it a valuable contribution to NeurIPS. I acknowledge that this point is highly subjective.
Summary
This paper proposes to learn general-purpose representations of PDEs from heterogeneous data by implementing joint embedding methods for self-supervised learning. Learned representation outperforms baseline approaches for invariant tasks such as regressing the coefficients of a PDE and improve the time-stepping performance of neural solvers.
Strengths
- Propose a general framework for self-supervised learning in PDE by using symmetry transformation in PDE. - Learned representation can be used in several tasks, such as determining unknown parameters and time-stepping. - This paper is well-written and the presentation is clear.
Weaknesses
- The novelty of this paper is quite marginal. The SSL framework adopted in the paper is well-known and not customized for the specific PDE problem. The only contribution is the augmentation of PDE solutions according to symmetry groups, which is also well-studied in previous literature, such as in [1]. - The evaluation on time-stepping is obviously not enough. For example, this paper doesn’t compare with several important baselines, such as MPPDE [2], FNO [3], UNet, PINN, etc. I think the authors need to show the effectiveness of using learned representation to improve at least a few of these models on limited labeled training data. - The regression task of determining external buoyancy force (a constant value) in the NS equation is quite simple in practice. In practice, there are more complicated forcing terms. For example, in FNO [3], they use a forcing term containing cosine and sine waves. Can authors use the forcing term in FNO and regress the parameters in that forcing term? - The viscosity in the NS equation is also an important parameter. Can authors also provide regression results on determining the viscosity? [1] Brandstetter, Johannes, Max Welling, and Daniel E. Worrall. "Lie point symmetry data augmentation for neural PDE solvers." International Conference on Machine Learning. PMLR, 2022. [2] Brandstetter, Johannes, Daniel Worrall, and Max Welling. "Message passing neural PDE solvers." arXiv preprint arXiv:2202.03376 (2022). [3] Li, Zongyi, et al. "Fourier neural operator for parametric partial differential equations." arXiv preprint arXiv:2010.08895 (2020).
Questions
- Why do authors only consider Lie point symmetry for PDE? Is there any other symmetry group that applies to PDE? - If a PDE does not have periodic boundary conditions, does augmentation using symmetry group still valid? - Can elaborate more on why learned representation cannot improve time-stepping performance for the Burgers equation?
Rating
3: Reject: For instance, a paper with technical flaws, weak evaluation, inadequate reproducibility and incompletely addressed ethical considerations.
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness
2 fair
Presentation
3 good
Contribution
1 poor
Limitations
The limitations of this paper are discussed.
Summary
In this paper, the authors propose to use self-supervised learning for obtaining an embedding that can robustly be used for predicting some quantities of interest or for time stepping. Particularly, they use the joint-embedding framework for SSL. They use symmetry groups for training the embedding to be invariant to these symmetry groups and thus capture the underlying behavior of the PDEs.
Strengths
+ novelty in using Lie symmetry groups for SSL training. + Improvement in training physics-based models in much lesser time. + covered a comprehensive list of symmetry groups for different PDEs. This would be very useful to the community.
Weaknesses
- For a person from an engineering background, I do think SSL is quite similar to the ideas of Multi-fidelity modeling (https://arxiv.org/abs/1609.07196, https://arxiv.org/abs/2110.04170, https://arxiv.org/abs/1903.00104). Both help reduce the training data and both are very easy to implement. Unfortunately, SSL for the quit nh - The authors talk about taking steps for building foundation models, I would imagine some kind of amalgamation of all the PDEs while training the SSL. Like you use all the data from all the PDEs to build and train an SSL model. That seems closer to how foundation models would behave. I feel like the paper gives big hopes and doesn't deliver them in the implementation. If they were to not talk about foundation models and just say they build SSL frameworks for PDEs, I would have been very happy. - The authors make a comparison between supervised ones and SSL-trained ones. but, there are a whole host of PDE-residual-based models like DeepONets, FNOsetc. (which the authors acknowledge in the supplement). It would be nice to know if SSL is even needed for PDE-based foundation models. Can PDE-residuals sufficiently help us navigate through foundation models? Is there a need for SSL? Thats not established fairly. (note that I acknowledge the comparisons, but is that sufficient to establish the necessity of SSL for PDEs?)
Questions
See the weaknesses. Further, there are more questions that are not particularly weakening the paper but could make it clearer. 1. Figure 5 and the corresponding text is very interesting to me. It seems that the results from the Lie transformation(0.0038 compared to 0.0078 in supervised) and Crop (0.0052 compared to 0.0078) are not that different. It seems to me that applying some kind of SSL training is more important than using a particular Lie transformation per se. Is it worth the effort to find out the best possible selection of Lie point augmentations? Further, is this selection task dependent? How does one train a generic model for several tasks then? 2. It would be nice to see some additional information on the computational overhead from SSL compared to supervised training. (see weaknesses, ideally, I would like to also know the comparison with DeepONets)
Rating
7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness
3 good
Presentation
3 good
Contribution
3 good
Limitations
N/A
Summary
This paper presents a general framework for self-supervised learning in a PDE context. In a way that is principled and natural, PDE symmetry groups are used to make the requisite augmentations from which self-supervision will learn structure; the augmentations are selected carefully so as to keep the regressed quantity (for the downstream task) constant. The utility of this method is demonstrated by regressing quantities of interest on four downstream PDEs, and the results indicate that the improvement of predicting from self-supervised representations is considerable relative to straightforward supervised learning. This work paves the way for adapting self-supervision to the physical sciences by way of carefully considered symmetry group-oriented data augmentations.
Strengths
1. The high-level idea is relatively novel, in that I do believe this is the first work that aims to make use of self-supervised learning for partial differential equations, with augmentations performed naturally according Lie symmetries. The most closely related work [12] uses Lie point symmetries of PDEs in order to augment PDE datasets, but this happens purely in the context of supervised tasks. 2. On the four PDE datasets presented/four equations considered, the results provide a fairly convincing improvement over just supervised learning alone. The mean squared error of prediction is considerably reduced (Table 1). 3. In terms of writing and presentation, the paper is quite good. The field is well introduced, with relevant literature cited. The idea is clearly presented and figures are given to illustrate the methodology (Figures 1, 2, and 3 were all quite instructive).
Weaknesses
1. Although the idea of using Lie symmetry-based augmentation for self-supervised learning is novel, I was underwhelmed at the rigor with which these augmentations were applied. Symmetries of PDEs are with respect to infinite domain or periodic boundaries. The boundary conditions imposed in the paper violate such symmetries, and hence make the justification for this sort of operation theoretically dubious. The authors admit that, naturally, because they violate these symmetries, they can only implement the group operations with "small strengths." I believe that the work would benefit from further investigating how to preserve these boundary conditions during augmentation, thereby providing a theoretically sound basis for the self-supervised learning proposal. 2. One of the interesting discoveries of this paper is that learning on top of self-supervised representations is considerably better than just immediately employing supervised learning techniques (Table 1); this is in stark contrast to what we see in e.g. computer vision, where both exhibit roughly the same performance. Some cursory post hoc analysis is given to explain this observation, although no further investigation is done in terms of providing a good understanding (be it theoretical or empirical).
Questions
My questions to the authors are listed below: 1. Have you put more thought into preserving boundary conditions during augmentation and/or rectifying the fact that the symmetry transformations applied are no longer the proper symmetries, since the system no longer has infinite domain? 2. Have you done any further investigation into why self-supervised learning gives results that considerably exceed that of supervised learning alone? ### Verdict In this paper, the authors consider the novel idea of applying self-supervised learning to PDEs. Although they lay the basic groundwork (e.g. performing augmentation via Lie symmetries), I feel that more consideration is needed to find a way to preserve boundary conditions during augmentation, such that the work is theoretically principled. The results are promising, in that the proposed approach considerably outperforms supervised learning, but little investigation is done as to why (this is a surprising/interesting observation and I would like to see more of an explanation given). As such, I recommend a borderline reject rating for this paper. ### Additional Comments and Minor Corrections The writing and presentation was, for the most part, very good (as mentioned above). I would recommend clarifying some of the table headings in results figures, for example, changing "Best strength" in Figure 5 (left) to "Best strength ($\epsilon$)." A few minor corrections are given below: L221: "test samples.As shown in" -> "test samples. As shown in" L229: "for to evaluate the models" -> "to evaluate the models" L253: "buyoancy" -> "buoyancy" L254: "which the hardest evaluation setting -> "which is the hardest evaluation setting" ## Post-rebuttal update After reading the authors' rebuttals, overall, I think the paper will make for a positive contribution to the conference. I have increased my score to a 5.
Rating
5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness
4 excellent
Presentation
4 excellent
Contribution
2 fair
Limitations
The authors have adequately addressed limitations of their work in the "Discussion" section (Section 5).
1. Sorry, my previous comment was somehow cut off. I wanted to say that it would be unkind to ignore all the work on multi-fidelity modeling. The authors claim clearly, "SSL attempts to learn a rich yet easy-to-leverage feature space" which is exactly the purpose behind multi-fidelity models. In conclusion, I would like the authors to acknowledge that SSL is another way of creating models with "easy-to-leverage feature spaces" and there are prior works to do so using multi-fidelity models. 2. Given the results that the authors have shared about mixing the differential equations, I think the authors must tone down the claims about foundation models as there are more steps than just scaling the current framework. Nevertheless, this is certainly a step forward toward the goal. I am happy to see this! 3. While I appreciate the clarity of the authors in the response, I do think both SSL and neural operator-type methods would be needed. The example results of FNO and FNO + SSL look good. I would like to see something similar with DeepONets vs. DeepONets+SSL. In other words, a foundation model may not just be SSL, but we may also need ingredients of a neural operator to be incorporated. Therefore, I would like to see the SSL in conjunction with the neural operators.
Response to reviewer clarifications
Thank you for clarifying your question and adding further feedback. We have responded to the additional comments below. *** ### On Multi-Fidelity Modeling and General Discussion of Other Feature Learning Methods Thank you for pointing out the connection between multi-fidelity modeling and our approach using SSL. We appreciate your feedback and agree that recognizing related works is crucial for a holistic understanding of our work's context. We concur that both multi-fidelity modeling and SSL can reduce dependency on extensive training data and have practical implementations. Despite these similarities, the two methods serve somewhat different overall purposes, and exploit data in unique ways. Multi-fidelity modeling primarily combines data and models of varying fidelities. Common goals include training models “using data from different levels of fidelity” [1] or enhancing “accuracy by injecting a small set of high-fidelity observations” into less accurate models [2]. In contrast, SSL aims to harness salient features from diverse data sources without being tailored to specific applications. The techniques we employed capitalize on the inherent structure of our dataset, especially through augmentations and invariances. Looking more broadly, we acknowledge that SSL is not the sole approach to feature learning. There exists a myriad of techniques, including metric learning, kernel design, autoencoders, and others [3-4]. We opted for SSL due to its proven efficacy in fields like computer vision and the direct analogy offered by our data augmentations. Nonetheless, you are correct that there are some high-level commonalities between multi-fidelity modeling and our approach. We appreciate this observation, and will add multi-fidelity modeling to the discussion of related work in the final revision. ### On Foundation Models and Mixing Differential Equations Thank you for your appreciation of the new mixed equation experiment. As stated in our response, we will tone down the language surrounding foundation models and share more details of the mixed equation experiments in the final draft. ### On Experiments with DeepONets Thank you for emphasizing the potential synergy between SSL and various neural operator-type methods. Your insight on the possible complementarity between these approaches is well-taken. To your point on expanding our experiments: While we have showcased the benefits of integrating SSL with the FNO, a representative neural operator architecture, we understand the value of further extending this analysis to other architectures like DeepONets. In preparation for the final paper, we commit to providing experimental results with DeepONets in tandem with SSL. We would like to underscore, however, that while additional architectures will certainly offer a broader perspective, the underlying message of our work regarding the usefulness of SSL remains consistent. We resonate with your viewpoint that the strength of SSL is amplified when used in harmony with other state-of-the-art techniques, thereby underscoring the collaborative nature of innovation in this domain. **References:** [1] Fernández-Godino, M. Giselle, et al. "Review of multi-fidelity models." arXiv preprint arXiv:1609.07196 (2016). [2] Perdikaris, Paris, et al. "Nonlinear information fusion algorithms for data-efficient multi-fidelity modelling." Proceedings of the Royal Society A: Mathematical, Physical and Engineering Sciences 473.2198 (2017): 20160751. [3] Kaya, Mahmut, and Hasan Şakir Bilge. "Deep metric learning: A survey." Symmetry 11.9 (2019): 1066. [4] Williams, Christopher KI, and Carl Edward Rasmussen. Gaussian processes for machine learning. Vol. 2. No. 3. Cambridge, MA: MIT press, 2006.
After reading the author's responses and other reviewers' responses, I have increased the score for this paper to accept; Cheers!
Thank you for the rebuttal
I have read the rebuttal as well as the rebuttals to other authors' concerns. Cumulatively, I think the work lays a reasonable foundation and sets the stage for interesting follow-up work. I have increased my score to a 5.
Decision
Accept (poster)