A Bregman Proximal Viewpoint on Neural Operators

1. Explain the architecture in a more direct way. Is there a bilevel optimization? This was much easier to understand in Frecon et al. (2022) "Bregman neural networks"

   We would like to start with a few preliminary facts.
   First, our work is a generalization of the paper "Bregman neural networks" to infinite dimensional function spaces which additionally provides universal approximation results, that were not present in the original work by Frecon et al. (2022).
   Second, similar to Frecon et al. (2022), once each layer is viewed as an optimization problem, the overall training of the architecture can be seen as a bilevel optimization problem, where at the lower level appear the optimization problems defining the action of each layer. Therefore, such bilevel formulation of the training phase is also shared by our proposal.
   Third, extending the work by Frecon et al. (2022) to handle neural operators between (infinite dimensional) functional spaces is not an easy task and some level of mathematical formalism is unavoidable, since infinite dimensional problems, when properly addressed, inherently require a more complex and delicate analysis. Our work has the same flavour of the work by Kovachki et al. (2023) who indeed extends the classical feed forward neural network architecture to function spaces and has required an entire work published in JMLR.
   Anyhow, we have tried to answer the reviewer's concern by extensively improving the accessibility of multiple parts of the paper (as indicated in blue). In addition, we have also added multiple figures which, we hope, clarify both our architecture and our contribution.

2. "Previous methods often lack a general formalism for characterizing their architecture." But there is a well-known paper that does: "Neural Operator: Learning Maps Between Function Spaces. How does the contribution of this paper relate to that one?

   Indeed, we have cited this influential work multiple times in the submitted version. In this work, the authors have shown that each operator layer maps each hidden representation to the next via the action of the sum of a local linear operator, a non-local integral kernel operator, and a bias function, composing the sum with a pointwise activation function. In particular, the FNO can be covered by a peculiar instance of the integral kernel operator. Our contribution differs by viewing the action of the operator layer as the minimizer of a regularized optimization problem, comparing the current hidden representation and the next. In our case, the choice of the regularization is implicitly linked to the resulting activation function. We also show that we can cover classical neural operators such as FNOs. To summarize, while in the mentioned work the authors have proposed and studied a parametric form of operator layers, our work takes a step back and shows how these parametric forms can notably be seen as minimizers of optimization problems. This perspective allows using the extensive literature on proximal optimization and monotone operator theory to further study the properties of neural operators (see our answers below).

3. "In this work, we propose a novel expressive framework for neural operators by conceptualizing the action of operator layers as the minimizers of Bregman regularized optimization problems over Banach function spaces." This sentence does not make sense, please clarify.

   We are sorry if our sentence is not clear. We wanted to highlight in part the elements provided in the previous answer. Put simply, we would like to emphasize that, in our paper, we demonstrate that applying one operator layer to an input function is equivalent to finding the minimizer of a regularized optimization problem over functions. Our setting requires working in Banach space, and the term expressive refers to the fact that we can cover classic neural operators and multiple activation functions. Please note that, in the revised version (see parts in blue), we have reworked the Contributions paragraph to better clarify the present work. In addition, we have added Figure 3 to help the reader to better understand our framework.

1. I think the authors should compare the performance of BFNO to that of F-FNO [...].
   We sincerely thank the reviewer for their kind comment. Indeed, both approaches exhibit similarities. We highlight that BFNO can also be viewed from the ODE point of view through a change of variables. More precisely, and without loss of generality, let us consider a simplified update of the form $v_{t+1} = \sigma( \sigma^{-1}(v_{t}) + W_t v_{t})$. Then, for $z_{t} = \sigma^{-1}(v_{t})$, it follows that $z_{t+1} = z_{t} + W_t \sigma (z_{t})$ which can be seen as a discretization of $\frac{\mathrm{d}z(t)}{\mathrm{d}t} = W(t)\sigma(z(t))$. In contrast, a residual-based architecture of the form $v_{t+1} = v_{t}+\sigma(W_t v_{t})$, such as F-FNO, would lead to the following ODE $\frac{\mathrm{d}v(t)}{\mathrm{d}t} = \sigma(W(t)v(t))$ on $v$ itself. The fact that, in BFNO, the weight matrix is placed outside the activation function can be seen as a different way of mitigating vanishing gradients compared to residual architectures like F-FNO. Concerning the comments on the performance comparison and on the BatchNormalization, we point the reviewer to our general answer.

2. The paper conducted theoretical analysis but were not able to show in theory why BFNO performs better [...]
   We thank the reviewer for highlighting the importance of a theoretical analysis to explain why BFNO outperforms FNO in practice. We agree that such insights would strengthen the paper. However, to the best of our knowledge, providing a formal proof of architectural superiority is a highly challenging task in the current landscape of theoretical deep learning research. One promising avenue for future work could involve analyzing the expressivity and trainability of BFNO using frameworks such as mean-field theory (e.g., [Ref. 1, Ref. 2]). However, this would require an entirely new work.
   That said, we have incorporated the reviewer’s suggestion to compare BFNO with F-FNO experimentally (see our general answer).

3. It is not clear to me if BatchNormalization [...].
   To the best of our knowledge, both BatchNormalization and LayerNormalization cannot be seen through the prism of proximal optimization techniques. Therefore, such architectures cannot be put in the framework of monotone operator theory. Nonetheless, this does not preclude the practitioner from using such layers both for FNO and BFNO. In the experiments described in our general answer, we have tried this and have found a marginal improvement.

4. Another weakness is that the activation function needs to be monotonic [...].
   Concerning the Bregman variant (e.g., BFNO), our framework requires the activation function to be strictly monotone (and so invertible). However, we would like to stress that our framework is more general. For reviewer's convenience, we report here the answer we gave to Reviewer JYsk.
   In section 3.1, we presented the general framework (6) for a layer operator. Looking at such general form, one can note that there is an outer operation (which is the $\mathrm{prox}$) and an inner invertible operation (which is $\tilde{\nabla} \Phi_t$). Therefore, by playing with the choices of the functions $\Phi_t$ and $g_t$ one can cover a number of situations of the form $\sigma_2(\sigma_1^{-1}(v)+\mathcal{K}_t(v)+b_t)$ with $\sigma_1$ being strictly monotone and $\sigma_2$ only monotone. Classical and Bregman neural operators emerge as special cases, where i) $\sigma_1 = \mathrm{Id}$ and $\sigma_2$ is any monotone function, for the former, and ii) $\sigma_1 = \sigma_2$ is strictly monotone, for the latter. We have made this clearer in the revised Section 3.1 (see the parts in blue).

5. I wonder if the authors use BatchNormalization when implementing FNO in their experiments because the original FNO paper used it.
   For our experiments, we relied on the latest available version of the FNO implementation, which does not include BatchNormalization. We note that the original FNO code was removed from the GitHub repository by its author (the 'master' branch was deleted). While we retrieved an earlier version of the code https://github.com/zeng-hq/fourier_neural_operator/commit/6f1d6fe8ac0b26a6941135dbc5c77b3c1cfcaf19, we observed that BatchNormalization was implemented in the initial commit but was subsequently removed in a later commit (commit number 6f1d6fe) titled "remove unnecessary batchnorm". We hope this clarifies our implementation choices and are happy to provide further details if needed. We have added an experiment in our general answer to support why BatchNormalization is not necessary.

[Ref. 1] T. Koshizuka et al. Understanding the expressivity and trainability of Fourier neural operator: a mean-field perspective, NeurIPS 2024.

[Ref. 2] G. Yang and S. Schoenholz. Mean field residual networks: On the edge of chaos. NeurIPS 2017.

We thank the reviewer for their prompt response and would like to provide additional context to clarify the reported results, helping them make a more informed decision on the scoring of our paper.

1. BFNO permits to build deep models

• The key takeaway for practitioners is that BFNO scales effectively with an increasing number of layers, ensuring that performance either improves or plateaus. This behavior is driven by two factors. First, the architecture supports stability, as $v_{t+1} = \sigma(\sigma^{-1}(v_t) + W_t v_t) = v_t$ when the weights are zero ($W_t=0$), reducing the layer to an identity mapping. Second, empirical evidence (see Figure 7 on page 29 of the new revision or here https://imgur.com/a/K5PvjOj) shows that many weights converge to (or remain near) zero in deeper models. As a result, when additional layers are unnecessary, they effectively default to the identity, maintaining the model's performance without degradation.
• It could be argued that, in residual-like networks, the relation $v_{t+1} = v_t + \sigma(W_t v_t) = v_t$ holds when $W_t$ is zero. However, we have not observed cases where the learned $W_t \approx 0$ across all configurations of layers tested. We believe that this discrepancy arises from the underlying dynamics explained by the ODE perspective (see the previous discussion), which makes this behavior challenging to achieve during training in residual networks, unlike in BFNOs.
• F-FNO was initially designed to support the construction of deeper models (up to 24 layers, whereas BFNO scales to at least 64 layers). However, as demonstrated by our experiments, this claim does not hold for the tasks we are considering in our paper: predicting of the final output from the initial condition. We believe their ability to reach 24 layers in their task—predicting all maps from $t$ to $t+1$—relies heavily on additional techniques, such as the Markov assumption and teacher forcing, which are specifically tailored to that type of task. To emphasize this, we present an extended table of results for models with up to 32 layers, clearly showing that BFNO increasingly outperforms F-FNO as model depth grows.
• In practice, cross-validating the number of layers is unnecessary for BFNO. One can simply initialize the model with as many layers as their GPU can handle and proceed with training. This approach guarantees achieving the best possible performance without additional tuning.

2. A more fair comparison with F-FNO

• In the previously reported results, F-FNO operated under a different setup compared to its competitors, primarily due to two factors: 1. the use of a different lifting layer and 2. operator layers with larger widths. These choices were made to align with the original implementation by its authors. To provide a more aligned comparison and isolate the impact of the operator layers, we present results for a modified model, F-FNO, which uses the same lifting layer and operator width as the other models (i.e., the setting originally considered in the FNO paper). In this setting, BFNO systematically matches or outperforms F-FNO.
• Our initial was to make a comparison with the best F-FNO configuration "out-of-the-box" (see F-FNObase). To achieve the reported performance (see F-FNO), we invested significant effort in carefully tuning the optimizer and cross-validating hyperparameters. In contrast, BFNO is far simpler to train, requiring none of the extensive adjustments and optimization tricks necessary for F-FNO.
• We stress that the main takeaway of our work is not merely to outperform F-FNO, but to demonstrate the robustness and scalability of BFNO in terms of number of layers.

Table: Comparison of Test relative ℓ² error across different architectures and number of layers. The table in our general answer has been changed accordingly. Changes in the text are indicated in italic

If any further explanation or additional experiments are needed for a positive reassessment of our work, we would be happy to provide them. Once again, we appreciate the reviewer’s time and valuable feedback.

1. The writing is poor, and I recommend that the authors carefully revise the paper from start to finish, especially regarding newly defined matrices or functions. For example, in Eq. 4, the definitions of $M_t$ and $K_t$ are not clearly stated when they first appear in the paper.
   We thank the reviewer for their feedback and for identifying the omission regarding the definitions of the matrices in Eq. 4. We have carefully revised the paper to check that all concepts are properly introduced without any missing definition. In addition, we have also improved the writing of multiple parts of the paper and moved the more technical parts to the appendix, as indicated in blue. We would be happy to incorporate any other recommendation that the reviewer deems useful.

2. The experiments are too simplistic. The authors only compared BFNO with FNO. I suggest including other FNO improvements as baselines.
   We appreciate the reviewer’s suggestion to include additional FNO improvements as baselines. Our primary objective was to focus on evaluating the added benefits of the $\sigma^{-1}$ term, which is why we restricted our comparison to the original FNO, a widely used and well-established baseline. That said, based on Reviewer HY9Q's comment, we are now conducting experiments to include comparisons with F-FNOs in the revised version. Preliminary results are provided in our general answer. Additionally, we would like to highlight that, in the submitted version, our experiments also included (in the appendix) multiple datasets and a comparison between WNOs and Bregman WNOs. Finally, following (Takamoto et al., 2022) we also provided a detailed frequency analysis showing high error rates on low frequencies and boundary conditions for FNOs.

3. Which $\psi$ in Table 1 do you use in the experiments? Do you compare the BFNOs obtained from different $\psi$?
   The choice of $\psi$ is directly linked to the choice of the activation function. Since we have resorted to the softplus activation function, this is equivalent to choosing the last $\psi$ of Table 1. The choice of the softplus activation function is motivated by the fact that it serves as a smooth approximation of ReLU, commonly used in FNOs.

1. It seems that there is still a gap between the universal approximation result and the numerical results in the article as the theoretical assumption about the non-linearity sigma (sigmoid type) does not hold in the numerical models ($\sigma=\text{softplus}$ is not sigmoid type). Therefore it would be good to mention this gap in the conclusion.
   We agree with the reviewer. We have tried to be careful by mentioning a "preliminary positive result concerning the universal approximation properties of Bregman neural operators" in the outline paragraph and at the beginning of Section 4. We propose to reinforce this gap both in the contributions paragraph by adding "Its applicability is grounded by universal approximation results proven for sigmoidal-type activation operators" and in the conclusion with "As part of our theoretical advancements, we have also established universal approximation results for Bregman neural architectures with sigmoidal-type activation functions. However, it must be acknowledged that a gap exists between this result and common practices, which predominantly rely on ReLU-like activations, as in our work, opening the door to new theoretical developments".

2. There is some notation inconsistency in the definition of the kernel $K_t^{\rm ac}$ in eq. 3 and the $K_t$ in Section 3.1. Are you talking about the same type of kernel in these 2 places? Why do you use $k^{(t)}$ as the kernel density, rather than the $k_t$ as before (below eq. 3)?
   The reviewer is right, we have now made the notation uniform and consistent along the paper, as shown in blue.

3. Section 3.1, it is unclear what the $\sigma_1$ and $\sigma_2$ after Remark 6 come from, do they depend on $g_t$?
   The reviewer is correct. To recall, in section 3.1, we presented the general framework (6) for a layer operator. Looking at such general form, one can note that there is an outer operation (which is the $\mathrm{prox}$) and an inner invertible operation (which is $\tilde{\nabla} \Phi_t$). Therefore, by playing with the choices of the functions $\Phi_t$ and $g_t$ one can cover a number of situations of the form $\sigma_2(\sigma_1^{-1}(v)+\mathcal{K}_t(v)+b_t)$ with $\sigma_1$ being strictly monotone and $\sigma_2$ only monotone. Classical and Bregman neural operators emerge as special cases, where i) $\sigma_1 = \mathrm{Id}$ and $\sigma_2$ is any monotone function, for the former, and ii) $\sigma_1 = \sigma_2$ is strictly monotone, for the latter. We have made this clearer in the revised Section 3.1 (see the parts in blue).

4. Is $\bar{D}$ the closure of the set $D$ in the definition of A in Section 4? Why do you consider the space $C$ with $\bar{D}$ rather than with $D$?
   The reviewer is right, it indeed corresponds to its closure. This has been made explicit in the revised version. The reason to consider the closure is that you need to evaluate the functions on the boundary of the domain when boundary conditions are imposed. This is pretty standard in the literature on PDEs (see, for instance, the book Introduction to Partial Differential Equations by G.B. Folland). By the way, this choice is also made in (Kovachki et al. 2023, Section 9.3).

5. Some type in Remark 9: "no in"?
   Since $\Psi$ is a strongly convex Legendre function, there is a unique solution and the equality holds. Note that in the revised version, we decided to remove it from the main body to gain space for additional comments since this was not a critical part for understanding our contribution.

3. Activation Function Analysis.

Reviewer HY9Q is right to stress that a limitation of our framework is the fact that it requires Bregman variants (such as BFNO) to have a strictly monotonic activation function, which excludes a few functions such as ReLU. This justifies why in our experiments we used Softplus as a surrogate of ReLU.

However, as a thought experiment, we also implemented BFNO with ReLU. The table below shows that BFNO with ReLU achieves comparable or better performance than Softplus for the same number of layers. However, the best results are the same (i.e., 12.2% for 16 layers).

We believe these extensive results, both from our original submission and the new additional experiments, validate our theoretical framework while demonstrating its practical benefits. We have carefully designed and conducted all experiments to ensure their rigor and reliability, and we have incorporated these findings in the appendix to address the reviewers’ concerns thoroughly.

We hope these efforts will clarify any remaining uncertainties and provide a stronger basis for a reassessment of our work. We are also happy to provide additional experimental details or conduct further specific comparisons if the reviewers find it helpful.

Respectfully,

The authors.