Discovering Physics Laws of Dynamical Systems via Invariant Function Learning

We appreciate the reviewer's thoughtful feedback. We will address them point by point below.

Drastic and non-linear changes across function environments

What would happen if changes across function environments were more drastic and non-linear? For example, changing f(theta) = sin(theta) to f(theta) = sin(theta)^2.

Thank you for this insightful question. In short, drastic changes are actually good for invariance discovery, but non-linear changes are much more sophisticated to analyze.

Drastic changes: According to our initial experiments, larger environmental differences will lead to better results. The experimental results shown in the paper indeed come with decent environment differences. For example, in the "powered" pendulum environment, the mass is flung over the top in many trajectories (Figure 7: upper left one).

Non-linear changes: In most optimization problems, expanding from linear to non-linear is challenging. For example, the invariant learning pioneer work IRM can only prove that the invariant principle works under the linear scenario and one SCM. Therefore, we don't expect our method to work directly for all non-linear scenarios. We provide analyses and open discussions around it and hope you find it interesting and useful.

• First, we consider the theory basis, the SCM in Appendix C.1. In the SCM, $f:=g\_{comp}(f\_c, f\_e)$, while $f\_e$ can be different across $e$, $g\_{comp}$ is shared. That means, if we consider $f(theta) = sin(theta)^2=sin(theta) * sin(theta)$ as $f:=f\_c * f\_e$, then the SCM remains valid if all the environments obey $f:=f\_c*f\_e$. If not all environments share the same function composition process, extended SCMs or extended $g\_{comp}$ will be required to be designed.
• Second, if SCM remains valid, the next challenge is the reverse decomposition (red dash arrows in Figure 6). In the current hyper-network implementation, given z_c denotes sin(theta), z_e denotes sin(theta), letting z_c + z_e denote sin(theta)^2 in the function representation space is a learning challenge, which may require much deeper networks or extra techniques. For example, we may choose to learn in the space of the derivative of sin(theta)^2 to alleviate the learning challenge.
• Finally, in the real world, many effects are composed additively or can be transformed to be composed additively like forces, so we believe our method is applicable in many cases.

Baseline implementation details

Please clarify the baselines. It is unclear how the baselines are implemented, what are the architectures, what are the input and outputs, and do they also have separate embeddings z_c and z_e?"

We appreciate this important point about baseline clarification. In the revised paper, we will provide a detailed explanation of the baseline implementations:

Basically, for fair comparisons, we tried our best to make baselines as similar to our architecture as possible. So all the performance gains come from our IFL principle.

Architecture

To better distinguish the baselines, let's denote our DIF framework without $\hat{e}$ predictions (Figure 3 and remove the $g_\phi$ MLP and $\hat{e}$ branch) as DIF-Base.

• MAML: Only Forecastor in the DIF-Base (with MAML optimization) since MAML does meta-learning at the optimization level.
• CoDA: DIF-Base and set the dimension of $\hat{z}_e=2$ (2 is the original paper setting).
• IRM: DIF-Base. IRM regularization is at the loss level so no architecture change.
• VREx: DIF-Base. VREx regularization is also at the loss level.

Inputs and Outputs:

• MAML takes X_p and outputs the forecasting trajectory directly.

• CoDA, IRM, VREx take X_p as input and output function representations $\hat{z}_c$ and $\hat{z}_e$, then finally output the forcasting trajectory. The same as DIF.

Training:

• MAML: MAML optimization on Forcaster.
• CoDA, IRM, VREx the same as DIF.

Inference: (Decoding X^c for the baselines)

• MAML: use the Forcaster with meta-parameters.
• CoDA, IRM, VREx: the same as DIF: use the f_c branch.

Research Questions

At the start of section 5, the authors list their 6 research questions. I think the only 'real' research questions that this paper is addressing is 1 and 2, whereas the other ones are more 'additional evidence'. Maybe it is worth rewriting that to make the contribution of the paper more concise.

We agree with this valuable suggestion. In the revised paper, we will restructure Section 5 to emphasize Research Questions 1 and 2 as the primary contributions of our work.

We will reduce the remaining questions (3-6) to "supplementary analyses" that provide additional evidence for our claims rather than core research questions.

We thank the reviewer for the constructive feedback and hope the questions are addressed. We will incorporate all these improvements in the revised version of the paper.

Thank you for your review. We appreciate your feedback and would like to address your concerns point by point.

Definition and identifiability of invariant function

The definition of an invariant function is vague...

We thank the reviewer for raising this point.

• Rigorous definition: Our invariant function is not defined arbitrarily; it is rigorously grounded in causality and formally specified within our Structural Causal Model (SCM) (see Section 2.2.1, Figure 2, and Appx. C.1). Specifically, the invariant function fc is the function generated by the structural equation and exogenous variable c, which represents the core mechanism shared across all environments. This causal definition distinguishes our approach from ad-hoc solutions.

• Example pitfall: We respectfully argue that the reviewer’s example oversimplifies the problem by focusing on a single realization, however, our treatment considers the invariant function as a random variable, a distribution that must be learned, not a specific instance.

• A more proper example: If we have $\mathbf{f}_1=\bm{\alpha} x + \bm{\beta}$ and $\mathbf{f}_2=\bm{\alpha} (x + 1) + \bm{\rho}=\bm{\alpha} x + (\bm{\alpha} + \bm{\rho})$, then since $f_c \perp e$, we find the decision boundary for the $e$ prediction, which determines by the distribution differences between $\bm{\beta}$ and $\bm{\alpha} + \bm{\rho}$, not by the distribution of $\bm{\alpha} x$ (invariant). Note that the three functions the reviewer provide can be sampled from both $\mathbf{f}_1$ and $\mathbf{f}_2$, illustrating the pitfall of analysing in instance level. This example is not perfect. For rigorous understanding, please resort to our SCM and information theory based proofs. Furthermore, we invite the reviewer to refer to Appendix B (Q2) and related works in invariant learning for additional insights.

• Community acknowledgment: The invariant function is defined rigorously in a statistically significant way using SCM. We like to share that using SCM to define invariant representation/mechanism has been a mature knowledge in invariant representation learning community.

There is no strong theoretical justification for why the invariant function is identifiable in all cases.

• Identifiability: Sufficiency & necessity: In Section C.2 of the appendix, we provide a complete proof of the existence and uniqueness of the solution to our optimization problem. The proof demonstrates that our invariant function learning principle guarantees identifiability under the causal assumptions.

I do not fully understand some theoretical claims in this paper, Theorem 3.1 (and its proof) in particular.

Theorem 3.1 provides the optimization principle to sufficiently and necessarily discover the function random variable defined in our SCM.

Note that our definition and theoretical proof have been acknowledged by other reviewers:

Reviewer pHQv stated:

• "structural causal model drawn by the authors (Figure2) illustrates the concept of invariant functions very well"
• "I read the theoretical proof in Appendix C, and there is no obvious problem with the derivation process."

Reviewer cD7J: remarked on the solid theoretical framework and highlighted the connection to Independent Causal Mechanism, underscoring that invariant learning is a well-established area.

Reviewer T91b confirmed "I have checked the main theoretical claims (3.1, 3.2 and 3.3) and I am happy with their soundness."

Due to space limitations, for further context, we recommend reviewing the fundamentals of invariant learning and the role of SCM in invariant learning listed in our essential related works.

Experimental Design and Generalization

One key concern is whether the method generalizes well to settings where environmental changes are more drastic. ... The "powered" environment would possibly break the typical behavior of the pendulum and fling the mass over the top.

Theoretically: The alterations in parameters (α and ρ) do not undermine the structure of our SCM; hence, our method remains generalizable.

Empirically: Our experiments in the "powered" environment do fling the mass over the top, and the results shown are still robust, suggesting it can handle increasingly drastic changes.

Writing and Organization

We appreciate this feedback and will improve the writing in our revision. We will:

1. Address redundancy and improve grammar
2. Enhance coherence, particularly in Section 2.2
3. Fix notation issues.
4. Clearly define all variables before use
5. Reorganize the explanation of challenges and solutions for improved flow

Related Works

Thank you for highlighting these important references. While we did mention PySR and symbolic regression in the Appendix, we agree that SINDy and genetic programming-based methods deserve more attention in the main text. We will expand Section 4 (Related Work) to include them.

We appreciate your feedback and the acknowledgment of the novelty and effectiveness. Below, we address your concerns point by point.

Limited experimental setup

Thank you for raising this concern. Our work represents the first step in invariant function learning, introducing a new paradigm for scientific discovery. The focus on low-dimensional systems was intentional for three reasons:

1. Conceptual clarity: Low-dimensional systems allow us to clearly demonstrate the core principles of invariant function learning without obscuring them with implementation details for high-dimensional data.
2. Methodological validation: As the first work on invariant function learning, we needed to establish baseline performance on well-understood systems where ground truth is available.
3. Benchmark contribution: We've constructed three multi-environment datasets (ME-Pendulum, ME-Lotka-Volterra, and ME-SIREpidemic) that will serve as important benchmarks for future research.

We fully agree that extending to high-dimensional real-world examples is an important direction, as we acknowledge in our limitations section (Section 6). The current work lays the theoretical and methodological foundation for these future extensions.

Scope limitation to ODE systems

Thank you for this important question. As we discussed in Appendix B, we acknowledge that our method is currently designed for ODE systems. This choice was driven by three key challenges:

1. Lack of multi-environment PDE datasets: Unlike domain adaptation tasks, multi-environment datasets for PDEs aren't readily available and require domain-specific expertise to construct.
2. PDE-specific technical challenges: PDEs introduce multi-scale and multi-resolution problems requiring specialized techniques for stability and computational efficiency.
3. Interpretation challenges: Our current approach uses symbolic regression for interpretability, which doesn't extend naturally to PDE systems.

Our future work will tackle these challenges by: (1) collaborating with domain experts to develop multi-environment PDE benchmarks, (2) incorporating specialized PDE-handling techniques from neural operators literature, and (3) developing new interpretability methods for invariant function learning in PDEs.

Connection to causal mechanisms

We agree that there is indeed a strong connection to the independent causal mechanism (ICM) assumption/principle. Our work can be viewed as discovering ICM specifically to function spaces in dynamical systems.

Note that "SCMs are underpinned by the assumption of independent causal mechanisms, a fundamental premise in causality. Intuitively, this supposition holds that causal correlations in SCMs are stable, independent mechanisms akin to unchanging physical laws, rendering these causal mechanisms generalizable."

ICM is the basis of Jonas's invariant causal predictor and IRM. When we use SCM in Section 2.2.1, we already adopt the ICM assumption, as in the SCM, we're assuming that invariant mechanisms (fc) and environment-specific mechanisms (fe) operate independently and can be disentangled. From the ICM perspective, our optimization goal is to identify the ICM showing in our SCM.

We acknowledge that we could have made this connection more explicit in our related work, and we appreciate the opportunity to clarify it here. In our revised version, we would enhance our discussion of this connection in the related work section.

Novelty limitations and related work

Although the experiments show the effectiveness of the proposed method, the use of information maximization to disentangle causal mechanisms has been used in various contexts before.

We acknowledge that information maximization for disentangling causal mechanisms has precedents in the literature. Our contribution lies in:

1. Novel task formulation: We introduce invariant function learning as a new task for scientific discovery (Section 2.2).
2. Function space application: While previous work has applied similar principles to fixed representations, we extend them to function spaces where invariance must be maintained across all possible states.
3. Theoretical guarantees: Our work provides theoretical guarantees for invariant function discovery (Theorem 3.1) specifically adapted to the dynamical systems context.
4. Empirical validation: Our experimental results (Section 5.4) demonstrate that general invariant learning principles like IRM and VREx do not perform well on our task, highlighting the need for our specialized approach.

These points, we believe, substantiate the novelty and significance of our approach.

Regarding essential references

Thank you for this valuable suggestion. We agree that enhancing our discussion of connections to ICM and ICA literature would strengthen the paper. In the revised versions, we would expand Section 4 (Related Work) to include them.

Thank you for your thoughtful review. We are glad you recognized our theoretical framework and experimental validation. We address your concerns point-by-point as follows.

Practicability in real-world scenarios

Thank you for this question. While we focus on establishing theoretical foundations in this work, we believe our method has significant real-world potential. In principle, the methodology, particularly the disentanglement via adversarial training and mutual information maximization, is applicable to real-world scenarios where environmental variations are present. Invariant function learning can extract fundamental dynamical laws from noisy, environment-influenced data, a critical challenge in scientific discovery. Application areas include:

1. Extracting physical laws from video data where environment factors (e.g., lighting, camera position) complicate analysis
2. Learning generalizable physics for simulation and control across varying conditions
3. Contributing to foundation models in physics by identifying invariant mechanisms

Note that some of these applications might not belong to invariant function learning but fall back to general invariant learning, which is not our target. PDE scenarios that need invariant function learning, are in our future work since much work needs to be done before being able to transfer to PDEs described in Appendix B Q1.

Theoretical claims and noise assumptions

We acknowledge that Appendix C.3 assumes Gaussian white noise. While this is a common starting point in theoretical analysis, we agree this is a limitation. Future work will extend our framework to other noise scenarios or relax this assumption.

High-dimensional systems with graph structure

Thank you for the suggestions and discussions on graph-structured physical mechanisms as in Cranmer et al. (2020) and Shi et al. (2022). We agree these are relevant to extending our work. Our current focus was establishing the basic framework for invariant function learning in ODE systems, but we acknowledge that discussing the approach to high-dimensional or graph-structured settings is worthwhile. In the revised paper, we will discuss these limitations and outline potential directions for adapting our method to such complex scenarios.

Utility Compared to Alternative Approaches

Thank you for your comparison with direct neural network fitting and symbolic regression approaches. Our method targets extracting invariant dynamics. Regarding alternative strategies like finding symbolic expressions in different environments and then obtaining the same part $f_c$ through observing, the most difficult part is observing. The reason is direct: the fitting of f tends to capture both invariant and environment-specific aspects, which leads to spurious correlations so that the final equation forms from different environments seem significantly different. That is to say, after extracting equations from different environments, it is challenging to find the common part directly since the common parts are blended with environment parts and look like not disentangable.

Our DIF method explicitly enforces the separation of invariant dynamics from environment-specific factors, providing a more reliable basis for scientific discovery. This advantage becomes especially important in scenarios with complex environmental effects or when specific invariant mechanisms need to be identified. We will expand on this distinction in the revised manuscript.

Real-world datasets

We agree that real-world validation would strengthen our paper. Our focus in this initial work was establishing theoretical foundations and validating on controlled systems. Testing on real-world data is the next step we discussed in the Appendix and we're actively pursuing this direction.

Question on X_p vs X

Regarding the specific question (Q1): Your understanding is mostly correct. We distinguish between $X_p$ (past observations before time $T_c$) and $X$ (full trajectory) because the forecasting task takes $X_p$ as input $X$ as output. Similar to $X\rightarrow Y$ in general prediction tasks, the forecasting task does $X_p \rightarrow X$. In training, $X$ works as "label $Y$" for loss, while in inference time, $X$ after $T_c$ is not available. Using the whole trajectory $X$ as $X_p$ would blur the line between training and testing data. We will clarify this point further in the revised manuscript.

We hope these responses address your concerns, and we thank you for your constructive feedback.