AutoBasisEncoder: Pre-trained Neural Field Basis via Autoencoding for Operator Learning

Authors: thomas x wang
nicolas  baskiotis
patrick  gallinari
Venue: ICLR_cc_2024_Workshop_AI4DiffEqtnsInSci
Published: 2024-02-11 12:58:55

Abstract: We introduce AutoBasisEncoder, a novel framework designed for operator learn-
ing – the task of learning to map from one function to another. This approach au-
tonomously discovers a basis of functions optimized for the target function space
and utilizes this pre-trained basis for efficient operator learning. By introducing
an intermediary auto-encoding task to the popular DeepONet framework, AutoBa-
sisEncoder disentangles the learning of the basis functions and of the coefficients,
simplifying the operator learning process. Initially, the framework learns basis
functions through auto-encoding, followed by leveraging this basis to predict the
coefficients of the target function. Preliminary experiments indicate that Auto-
BasisEncoder’s basis functions exhibit superior suitability for operator learning
and function reconstruction compared to DeepONet. These findings underscore
the potential of AutoBasisEncoder to enhance the landscape of operator learning
frameworks

Main paper: 
 introduction: 
 Traditional approaches for solving PDEs often rely on predefined basis functions, limiting their adaptability to complex spatial patterns. Central to the success of these methods is the choice of basis, so that it can effectively represent the solution space. While traditional basis functions like Fourier series and polynomials have been widely employed, they usually have limited flexibility, due to being problem-specific (Meuris et al., 2023) . Additionally, handcrafting an appropriate basis for a given problem often requires extensive domain knowledge.Recently, Lu et al. (2021) introduced DeepONet, a framework that learns a basis of functions directly from data. They represent the target function as a linear combination of basis functions, which are learned jointly with a network that predicts the corresponding coefficients. To address the complexity of this complex high-dimensional optimization problem (Wang et al., 2022b) , we explore decomposing the DeepONet scheme into two distinct, simpler sub-problems: optimizing the basis network and coefficient network separately. This idea has been explored by POD-DeepONet (Lu et al., 2022) , who apply a Proper Orthogonal Decomposition (POD) on discretized target functions, learning a basis that is then used to optimize the branch network. Notably, the outcomes of this approach are discretizations of functions on a pre-defined grid, thus requiring the incorporation of additional interpolation methods to obtain predictions outside of this grid (Witman et al., 2022; Demo et al., 2023; Meuris et al., 2023) .In this context, we propose AutoBasisEncoder, a two-step framework that breaks down DeepONet's optimization into two parts: the learning of basis functions and the learning of their coefficients. By introducing an intermediary auto-encoding objective, we first focus on learning a basis of functions tailored to our data; and leverage this pre-trained basis to learn the operator itself, by learning the coefficient network. The auto-encoding task helps to learn a basis of functions with better representation power.Preliminary experiments show that this disentanglement of the DeepONet's optimization leads to the learning of a basis that captures the target function space better, and that this basis can improve the operator learning.To summarize, AutoBasisEncoder simplifies the optimization of DeepONet through an additional auto-encoding task. The resulting pre-trained basis of functions is more expressive and easier to optimize for than the basis of DeepONet. 
problem setting: 
 Our goal is to learn a mapping between two functional spaces.Let U ⊂ L 2 (Ω U , R p ) and V ⊂ L 2 (Ω V , R q ) be two infinite-dimensional function spaces, with Ω U , Ω V compact subsets of R m and R n .Our goal is to learn an operator F that maps functions from space U to space V:F : U → V, u → vIn practice, we are given N pairs of functions(u i | Xi , v i | Yi ) i∈ 1,N discretized on grids X i ⊂ Ω U and Y i ⊂ Ω V .AutoBasisEncoder's architecture is based on DeepONet (Lu et al., 2021) . In the DeepONet framework, the target function v is modeled as a linear combination λ = (λ 1 , . . . , λ d ) of a set of basis functions:Φ = {Φ i } i , so that vλ,Φ (x) = d i=1 λ i Φ i (x).The network responsible for predicting the coefficients (λ i ) i is referred to as the branch net, while the network handling the function bases is known as the trunk net.As illustrated in Figure 1 , DeepONet approximates the operator F with: F = L Φ • U , where L Φ is a linear combination with a basis of function B = (Φ i ) i , and U is the coefficient predictor:U : U → R d , u → (λ 1 , . . . , λ d ), L Φ : R d → V, (λ 1 , . . . , λ d ) → d i=1 λ i Φ i ,Challenges in DeepONet While DeepONet aims to learn a basis for representing the output space V, the primary focus of the operator learning task is not centered on learning an expressive basis, but rather on efficiently processing u to accurately map it to v. Consequently, DeepONet is compelled to strike a delicate balance between basis learning and coefficient learning, which makes its optimization highly complex, potentially resulting in a compromise of basis expressivity. As introduced better, previous work such as POD-DeepONet (Lu et al., 2022) have shown that explicitly stating these two sub-tasks (basis learning and coefficient learning) as separate objectives can enhance the operator learning process.Building upon this understanding, we introduce AutoBasisEncoder -a novel framework specifically designed to address these two sub-tasks separately. While AutoBasisEncoder shares the same network architecture as DeepONet, AutoBasisEncoder distinguishes itself through the incorporation of an intermediary auto-encoding task (see Figures 1 and 2 ). The goal of this task is to promote a basis of functions that better represents the output space V. Consequently, AutoBasisEncoder disentangles the optimization process, breaking down the complex task of operator learning into two simpler tasks: basis learning and coefficient learning.U V R d U F L Φ Figure 1: DeepONet decomposition. L Φ is trained concurrently with U . U V R d U F V L Φ Figure 2: AutoBasisEncoder's Two-Step De- composition: 1. Training L Φ with V . 2.Freezing L Φ and training U .Φ ! Φ " Φ # ⋮ Coefficient network 𝑉 𝜆 ! " 𝜆 # " ⋮ 𝜆 $ " 𝑣 Φ ! (𝑥) Φ " (𝑥) Φ # (𝑥) ⋮ 𝑥 X & $%! # 𝜆 $ & * Φ $ (𝑥) = * 𝑣 ' ! ,𝚽 (𝑥)Step 1 -Basis learning via autoencoding We reconstruct 𝐼𝑑: 𝑣 ↦ 𝑣 by learning the basis functions (Φ % ) and the coefficient network𝑉 ℒ &'()*+,)$* = 𝑀𝑆𝐸(𝑣, 3 𝑣 -! ,𝚽 ) Φ ! Φ " Φ # ⋮ Coefficient network 𝑈 𝜆 ! 0 𝜆 # 0 ⋮ 𝜆 $ 0 𝑢 Φ ! (𝑥) Φ " (𝑥) Φ # (𝑥) ⋮ 𝑥 X & $%! # 𝜆 $ * * Φ $ (𝑥) = - 𝑣 ' " ,𝚽 (𝑥)Step 2 -Coefficient learningWe train 5 ℱ: 𝑢 ↦ 𝑣 by learning the new coefficient network 𝑈 The basis functions (Φ % ) from Step 1 are frozenℒ )1*2&()2 = 𝑀𝑆𝐸 (𝑣, 8 𝑣 -" ,𝚽 )Figure 3 : Detailed AutoBasisEncoder architecture. On top, the basis pre-training step: through the task of autoencoding V, we learn the set of functions. L Φ is trained concurrently with U . On the bottom, the coefficient learning step: the networks {Φ i } i are frozen, and only the coefficient network U is trained.Pre-training a neural field basis of the output space We propose to discover a basis of functions via auto-encoding the target space V. We constrain our latent space to represent the coefficients corresponding to the decomposition on this basis.During this phase, both the basis of functions Φ and the coefficient predictor V : V → R d are learned concurrently, so that:L Φ • V ≈ Id V .In other terms, for any function v ∈ V, we autoencode it into a vector λV = (λ V 1 , . . . , λ V d ) ∈ R d so that: ∀x ∈ Ω V , vλ V ,Φ (x) = d i=1 λ V i Φ i (x) ≈ v(x), where:λ V = V (v).For supervision, we use a MSE loss between the groundtruth v and its reconstruction:L autoencode = M SE(v, vλ V ,Φ ).A detailed view of this step is presented in the top part of Figure 3 .Coefficient Learning In its second phase, AutoBasisEncoder focuses on acquiring the coefficient networkU : U → R d , where u maps to λ U = (λ U 1 , . . . , λ U d ) ∈ R d .The objective is to ensure that the composition of the basis learning operator L Φ (which was pre-trained during the previous step) with the coefficient network U approximates the true operator F, expressed as L Φ • U ≈ F, as shown in the bottom part of Figure 3 . We denote vλ U ,Φ (x) = d i=1 λ U i Φ i (x). This step's goal is to learn the operator, therefore we supervise this step via a MSE loss on the target function v: L operator = M SE(v, vλ U ,Φ ). Unlike other operator methods (Li et 2021; Lu et al., 2022) , which often require specific sampling constraints on the input functions, we do not assume specific architectures for processing the input function u, allowing all kinds of architectures as encoders. 
experiments: 
 We first compare the performance of AutoBasisEncoder on standard operator learning tasks, and then compare the bases learned by the different methods.We evaluate the performance of AutoBasisEncoder and DeepONet over two standard datasets for operator learning: the Darcy flow equation and the Navier-Stokes equation dataset with viscosity 1e -3 from Li et al. (2021) . Further details are presented in Appendix A.2.1.Our model is compared to vanilla DeepONet (Lu et al., 2021) , and POD-DeepONet (Lu et al., 2022) . For fairness, these methods all share the same architecture for the coefficient network, use 64 basis elements and are trained during 20k epochs in total, allowing convergence.We introduce the two following tasks:T autoencode : Given an input function v ∈ V, we aim to reconstruct v itself. T operator : Given an input function u ∈ U, we aim to predict the corresponding v = F(u) ∈ V.As we can see in Table 1 , AutoBasisEncoder improves vanilla DeepONet performance on both datasets, showcasing the importance of our pre-training step. We can see that our framework performs well on the Operator learning task, even though it was first trained on the auto-encoding task. It also performs better than POD-DeepONet on the Darcy dataset. POD-DeepONet outperforms AutoBasisEncoder on the Navier-Stokes dataset, but AutoBasisEncoder can provide continuous predictions all over the space V, while POD-DeepONet can only produce predictions on points of its training data grid. In this section, we evaluate the efficacy of the acquired function bases by examining their ability to comprehensively span the target function space V.We evaluate the effectiveness of the bases learned by DeepONet, AutoBasisEncoder, and POD-DeepONet by assessing their ability to represent target function spaces V. To do this, we first pre-train the bases using each respective algorithm. Then, we use gradient descent to learn the optimal coefficients for linearly decomposing the target function space V in order to auto-encode it. These bases are trained on regular grids, with a resolution of 64 for the Navier-Stokes dataset and 29 for Darcy equation. We then use these bases to represent functions sampled at different resolutions: 32, 64, and 128. The results of this evaluation are summarized in Table 2 . Notably, the basis produced by AutoBasisEncoder consistently outperforms that of DeepONet. Furthermore, AutoBasisEncoder also facilitates faster convergence of the gradient descent algorithm to the optimal coefficients, as demonstrated in Figure 5 . While POD-DeepONet's basis enables better reconstruction of the function space V, it is constrained by its discrete nature and is unable to accurately reconstruct functions with different discretizations from those used during training. Conversely, DeepONet and Auto-BasisEncoder exhibit greater flexibility, as they employ functions as components of their bases, allowing them to reconstruct functions sampled differently from the training sampling. DeepONet 1.32×10 -4 2.48×10 -5 3.79×10 -2 3.83×10 -2 POD-DeepONet 8.60×10 -5 3.50×10 -6 N/A N/A AutoBasisEncoder 1.07×10 -4 1.36×10 -5 3.08×10 -2 3.16×10 -2 Therefore, compared to DeepONet, AutoBasisEncoder's learned basis demonstrates better suitability for capturing the underlying structure of the output functions in space V , and while also improving the effectiveness in learning the operator F.However, a discernible gap exists between the MSE of operator learning and the MSE of the V space reconstruction task using AutoBasisEncoder's basis, indicating room for further improvement. 
conclusion: 
 We propose AutoBasisEncoder, a novel framework that disentangles the operator learning task into two sequential tasks: the learning of an appropriate basis of neural networks that can represent the output function space more accurately than DeepONet, and the prediction of the target function via its corresponding basis coefficients. Through an intermediary auto-encoding task, this disentangling makes it possible to learn a basis of functions that is more suitable for multiple tasks, and in particular operator learning. This framework shows potential for operator learning towards better basis functions learning. 
 
 Appendix: A.1 RELATED WORK Auto-encoding in Operator learning Works such as Oommen et al. (2022) and Li et al. (2023) utilize CNN encoder-decoder architectures for auto-encoding functions in a latent space. They then learn temporal evolution by fitting a DeepONet in this latent space. BelNet (Zhang et al., 2023) encode their input function via a POD-like decomposition, while Zhang et al. (2022) uses an autoencoding framework based on DeepONet, specifically tailored for Stochastic Differential Equations.One common challenge in these methods is the interpretability of the latent spaces. By adopting DeepONet as a general framework, we can view the predicted coefficients of the linear decomposition as a latent vector, which significantly enhances the interpretability of the latent space. This approach provides a more intuitive understanding of the latent space's content and facilitates insights into the learned representations.Among existing methods, BasisONet (Hua & Lu, 2023) stands out as the closest to AutoBasisEncoder: BasisONet auto-encodes both the input and output spaces onto orthonormal bases of functions, and uses these auto-encoding constraints as additional regularizing loss terms. On the other hand, AutoBasisEncoder adopts a distinct approach by concentrating solely on auto-encoding the output function space, offering greater flexibility in processing the input function, enabling it to effectively handle non-continuous functions. -∇ • (a(x)∇u(x)) = f (x) for x ∈ (0, 1) 2 subject to the Dirichlet boundary condition u(x) = 0 for x ∈ ∂(0, 1) 2 , where a ∈ L ∞ ((0, 1) 2 ; R + ) denotes the diffusion coefficient, and f ∈ L 2 ((0, 1) 2 ; R) represents the forcing function.We aim to learn the operator F : L ∞ ((0, 1) 2 ; R + ) → H 1 0 ((0, 1) 2 ; R + ), which maps the diffusion coefficient a to the solution u. It describes a flow fluid through a porous medium.We use the dataset from Li et al. (2021) , consisting of 1000 couples of functions, discretized on a 29x29 regular grid.We investigate the behavior of a 2D Navier-Stokes equation on the unit torus. We consider the flow of a viscous, incompressible fluid. The equation involves the evolution of vorticity w(x, t) over time, subject to certain boundary conditions and initial vorticity w 0 (x). Our objective is to learn an operator F that maps vorticity values w(x, t) from a given time t = 0 to a later time T : F : w(x, t = 0) → w(x, t = T ).The 2D Navier-Stokes equation is describes as follows:∂w(x, t) ∂t + u(x, t) • ∇w(x, t) = ν∆w(x, t) + f (x), x ∈ (0, 1) 2 , t ∈ (0, T ] ∇ • u(x, t) = 0, x ∈ (0, 1) 2 , t ∈ [0, T ] w(x, 0) = w 0 (x), x ∈ (0, 1) 2 (1)where u is the velocity field and w = ∇ × u is the vorticity. w 0 is the initial vorticity, ν ∈ R + is the viscosity coefficient of the fluid. A forcing function f ∈ L 2 per ((0, 1) 2 ; R) is used. We use the Navier-Stokes equation dataset with viscosity ν = 1e-3 from Li et al. (2021) . It consists of 1000 couples of vorticity functions at times 0 and T : (w(x, t = 0), w(x, t = T )), generated on a 256x256 regular grid, and then downsampled at a 64x64 resolution. We also use the 32 and 128 resolution data to assess the capabilities of our framework to model the distribution of vorticities w(x, t = T ) at time T .In practice, we instantiate our neural basis using neural networks with sinusoidal activation functions, capitalizing on their excellent high-frequency representation capabilities (Tancik et al., 2020; Sitzmann et al., 2020; Wang et al., 2022a) . They indeed yield better performance than more traditional activation functions (ReLU, tanh...). Unlike the sparse Mixture of Experts employed by Wang et al. (2022a) , we empirically found that a non-sparse mixture better suits our operator learning tasks. We use the L 2 norm for supervision during the basis learning step and during the coefficient learning step.AutoBasisEncoder residualGroundtruth  


Decision: This paper introduces a novel framework designed for operator learning, mapping from one function to another. All reviewers agree unanimously. The concerns raised by reviewers can be address by enhancing the clarification of the paper. Therefore, I support the acceptance. However, author needs to address the concerns in the camera ready version

Reviews: review: title: title: This work presents an AutoBasisEncoder which is a framework designed for operator learning based on of set of functions optimized for the target function space and utilizes this pre-trained basis for efficient learning process.
main_review: main_review: This work presents an AutoBasisEncoder, which is a framework designed for operator learning based on of set of functions optimized for the target function space and utilizes this pre-trained basis for an efficient learning process.
The paper is well written and some conclusions underscore the potential of AutoBasisEncoder to enhance the landscape of operator learning frameworks
title: title: Official Review
main_review: main_review: DeepONet is a framework for operator learning that learns a basis of functions from data and represents functions as linear combinations of the bases. Both the basis functions and the network that predicts the coefficients of the linear combination are learned jointly. Following recent work (POD-DeepONet, BasisONet), this paper proposes AutoBasisEncoder that separates these into two steps: learning the basis function using an autoencoding objective and the learning the coefficients. Empirical results show improved performance on Darcy flow equation and the Navier-Stokes equation datasets over DeepONet and the ability to interpolate over different resolutions of data.

### Strengths
- The paper is easy to follow (although some discussion on related work can be improved; see weaknesses).
- The proposed approach improves over DeepONet on two datasets, highlighting the the effectiveness of disentangling basis learning from coefficient learning.

### Weaknesses
- The proposed method is simplistic in nature and the empirical analysis is also limited. This makes the overall contribution marginal, especially in the context of works such as BasisONet.
- Discussion on related work needs improvement. In particular, the differences from POD-DeepONet and BasisONet need to be highlighted more to better position the work.

### Questions
- Why did the authors not compare with BasisONet given its similarity to the proposed method? 
- Why did the authors not use any interpolation techniques for POD-DeepONet in Table 2? Can they not be used?
- Why is POD-DeepONet missing from Fig. 3?
title: title: Review of AutoBasisEncoder: Pre-trained Neural Field Basis via Autoencoding for Operator Learning
main_review: main_review: The paper introduces a new framework for operator learning by separating the optimization of the basis coefficients from the basis functions in a pre-existing architecture, DeepONet. The authors first use an autoencoder to find the basis functions from the target space, which are then fixed to find the coefficients using the input functions, and show that this approach leads to the discovery of more expressive basis functions. The paper, for the most part, is clear - the idea is simple and is relevant to the field. A major weakness is that the autoencoding step to determine the basis functions requires knowledge of the target space at points (unlike DeepONet and POD-DeepONet which just uses locations to encode the basis functions). This limits its use for test examples where there is zero or limited access to the target space. There are also many typos in the paper.