HT-Net: Hierarchical Transformer based  Operator Learning Model for Multiscale PDEs
Authors: Xinliang Liu
Bo Xu
Lei Zhang
Venue: Submitted to ICLR 2023 (2023)
Keywords: hierarchical transformer
operator learning
multiscale PDE
nested self-attention
loss function
generalization error

Abstract: Complex nonlinear interplays of multiple scales give rise to many interesting physical phenomena and pose major difficulties for the computer simulation of multiscale PDE models in areas such as reservoir simulation, high frequency scattering and turbulence modeling. In this paper, we introduce a hierarchical transformer (HT-Net) scheme to efficiently learn the solution operator for multiscale PDEs. We construct a hierarchical architecture with scale adaptive interaction range, such that the features can be computed in a nested manner and with a controllable linear cost. Self-attentions over a hierarchy of levels can be used to encode and decode the multiscale solution space over all scale ranges. In addition, we adopt an empirical $H^1$ loss function to counteract the spectral bias of the neural network approximation for multiscale functions. In the numerical experiments, we demonstrate the superior performance of the HT-Net scheme compared with state-of-the-art (SOTA) methods for representative multiscale problems.


Under review conference paper at ICLR 2023
HT-NET: HIERARCHICAL TRANSFORMER BASED OPERA- TOR LEARNING MODEL FOR MULTISCALE PDES
Anonymous authors Paper under double-blind review
ABSTRACT
Complex nonlinear interplays of multiple scales give rise many interesting physical phenomena and pose significant difficulties for the computer simulation of multiscale PDE models in areas such as reservoir simulation, high-frequency scattering, and turbulence modeling. In this paper, we introduce a hierarchical transformer (HTT) scheme to efficiently learn the solution operator for multiscale PDEs. We construct a hierarchical architecture with a scale-adaptive interaction range, such that the features can be computed in nested manner and with controllable linear cost. Self-attentions over hicrarchy of levels can be used to encode and decode the multiscale solution space across all scales. In addition, we adopt an empirical H' loss function to counteract the spectral bias of the neural network approximation for multiscale functions. In the numerical experiments, we demonstrate the superior performance of the HT scheme compared with state-of-the-ant (SOTA) methods for representative multiscale problems.
INTRODUCTION
Partial differential equation (PDE) models with multiple temporal/spatial scales are ubiquitous in physics, engineering. and other disciplines. They are of tremendous importance in making predictions for challeng- ing practical problems such as reservoir modeling, high-frequency scattering, and atmosphere circulation, to name few. The complex nonlinear interplays of characteristic scales cause major difficulties in the com- puter simulation of multiscale PDEs. While the resolution of all characteristic scales is prohibitively expen- sive, sophisticated multiscale methods have been developed efficiently and accurately solve the multiscale PDEs by incorporating microscopic information. However, most of them are designed for problems with fixed input parameters.
Recently, several novel methods such as Fourier neural operator (FNO) (Li et al., 2021), Galerkin transformer (GT) (Cao, 2021) and deep operator network (DeepONet) (Lu al., 2021) are developed to directly learn the operator (mapping) between infinite dimensional spaces for PDE problems, by taking advantages of the enhanced expressibility of deep neural networks and advanced architectures such as feature embedding, channel mixing, and self-attentions. Such methods can deal with ensemble of input parameters and have great potential for the efficient forward and inverse solvers of PDE problems. However, for multiscale problems, most existing operator learning schemes essentially capture the smooth part of the solution space, and how to resolve the intrinsic multiscale features remains to be a major challenge.
In this paper, design a hierarchical transformer based operator learning method, that the accurate, efficient, and robust computer simulation of multiscale PDE problems with an ensemble of input parameters becomes feasible. Our main contribution can be summarized in the following.

---

Under review conference paper at ICLR 2023
we develop a novel transformer architecture that allows the decomposition of the input-output map- ping to hierarchy of levels, the features can be updated in nested manner based on the hierarchi- cal local aggregation of self-attentions with linear computational cost;
adopt the empirical H' loss which avoids the spectral bias and enhances the ability to capture the oscillatory features of the multiscale solution space;
the resulting scheme has significantly better accuracy and generalization properties for multiscale input parameters, compared with state-of-the- (SOTA) models.
2 BACKGROUND AND RELATED WORK
We briefly introduce multiscale PDEs in Section $ 2.1, then summarize relevant multiscale numerical meth- ods in 2.2, and neural solvers, in particular neural operators in $ 2.3.
2.1 MULTISCALE PDES
Multiscale PDEs, in narrower sense, refer to PDEs with rapidly varying coefficients, which may arise from a wide range of applications in heterogeneous and random media. In broader sense, they may form a hierarchy of models at different scales by systematic derivation, starting from fundamental laws of physics such as quantum mechanics (E & Engquist, 2003). Some outstanding multiscale PDE models may include:
Multiscale elliptic equations is an important class of prototypical examples, such as the following second- order elliptic equation of divergence form.
(a(x) = f(x) u(x) = 0
x xEl D RE
(2.1)
where 0 < amin < <(xz)= S amaxi Vx E D, and the forcing term f (D;R). The coefficient to solution map is S : 1 D2; [_) considered and S can be well resolved by the FNO parameterization. The setup a(x) E L allows rough coefficients with fast oscillation (e.g. a(x) a(x/c) with 1), high contrast ratio with 0max/ Pmax//mmn >> and even a continuum of non-separable scales. The rough coefficient case is much harder from both scientific computing (Branets al., 2009) and operator learning perspectives.
Hg (D; R)), such that S(a). In Li et al. (2021), smooth coefficient a(x) is
Navier-Stokes equation models the flow of incompressible fluids, which becomes turbulent due to the simul- laneous interaction of wide range of temporal and spatial scales of motion.
Helmholtz equation models ime-harmonic acoustic waves. Its numerical solution exhibits severe difficulties the high wave number regime due to the interaction of high-frequency waves and numerical mesh.
2.2 MULTISCALE SOLVERS FOR MULTISCALE PDES
For multiscale PDEs, the computational cost of classical numerical methods, such as finite element methods, finite difference methods, etc., usually scales proportionally to 1/E >>> Multiscale solvers have been de- veloped such that their computational costs are independent of E by incorporating microscopic information.
Asymptotic and numerical homogenization Asymptotic homogenization (Bensoussan et al., 1978) is an clegant analytical approach for multiscale PDEs with scale separation, e.g., a(x) a(z/s) with < 1 in equation 2.1. For general multiscale PDEs with possibly a continuum of scales, numerical homogenization (Engquist & Souganidis, 2008) offers effective numerical approach that aims identify low dimensional approximation spaces such that the approximation bases are adapted to the corresponding multiscale opera- tor, and can be efficiently constructed (e.g. with localized bases). See Appendix A for details.
2

---

Under review conference paper at ICLR 2023
Multilevel and multiresolution methods Multilevel/muligriid methods (Hackbusch, 1985; Xu & Zikatanov, 2017), and multiresolution methods such as wavelets (Brewster & Beylkin, 1995; Beylkin & Coult, 1998) have been successfully applied to PDEs. However, the convergence of those methods for mul- tiscale problems can be severely affected by the regularity of coefficients (Branets et al., 2009). Recently, the introduction of gamblets (Owhadi, 2017) can be seen as multilevel extension of numerical homogeniza- tion, and it opens an avenue to automatically discover scalable multilevel algorithms and operator-adapter wavelets for linear PDEs with rough coefficients. See also Appendix A.2.
Low-rank decomposition based methods It is well-known that the elliptic Green's function has low- rank approximation (Bebendorf, 2005), which lays the theoretical foundation of the (near-)linear complexity methods such fast multipole method (Greengard & Rokhlin, 1987; Ying al., 2004), hierarchical matrices (H and 2 matrices) (Hackbusch et al., 2002; Bebendorf, 2008), and hierarchical interpolative factoriza- tion (Ho& & Ying, 2016). See Appendix for the connection with our method presented in this paper.
Tensor numerical methods Analytical approaches such as periodic unfolding (Cioranescu et al., 2008) suggest that a oow-dimessional multiscale PDE can be transformed into a gigh-dimensional PDE. Therefore, tensor numerical methods. e.g., the sparse tensor finite element method (Harbrecht & Schwab, 2011) and the quantic tensor train (QTT) method (Kazeev al., 2022), provide efficient numerical procedures to find the low-rank tensor representation of multiscale solutions.
2.3 NEURAL OPERATOR FOR PDES
Neural PDE solvers and Spectral Bais Various neural solvers have been proposed in E et al. (2017); Sirignano & Spiliopoulos (2018); Raissi et al. (2019) to solve PDEs with fixed parameters. The difficulty to solve multiscale PDFs has been exhibited by the so-called spectral bias or frequency principle (Rahaman al., 2019; Ronen al., 2019), which shows that DNN-based algorithms are often inefficient learn high- frequency components of multiscale functions. A series of algorithms have been developed to overcome the frequecyy curse of DNNs and solve multiscale PDEs (Cai et al., 2020; Wang et al., 2021).
Operator learning for PDEs DNN algorithms demonstrate more potential to learn the input-output map- ping of parametric PDEs. Finite-dimensional operator learning methods such as 7hu & Zabaras (2018); Fan et al. (2019a;b); Khoo et al. (2020) can be applied to problems with fixed discretization. In particular, Fan et al. (2019a;b) combines H or X2 matrices linear operations with nonlinear activation functions, though the nonlinear geometrical interactioo//agrreatioo is absent, which limits the expressivity of the resulting neural operator. definite-dimessionn operator learning methods (Li et al., 2021; Gupta et al., 2021) aim to Banach spaces, and the convolutions in the neural operator
learn the mapping between construction are parametrized by Fourier or wavelet transform. Nevertheless, those methods do not always work even for multiscale PDEs with fixed parameters. On the other hand, while the universal approxima- tion can be rigorously proved for FNO operators (Kovachki et al., 2021), "extra smoothness" is required to achieve meaningful decay rate, which either is absent or gives rise to large constants for multiscale PDEs. This motivates us to construct a new architecture for multiscale operator learning.
Efficient Attention The vanilla multi-hean eelf-attention in Vaswani et al. (2017) scales quadratically with the number of tokens, which is prohibitive for high-resolution problems. Many efficient attention methods are proposed using the kernel trick (low-rank projection) (Choromanski et al., 2020; Wang et al., 2020; Peng et al., 2021; Nguyen et al., 2021) to reduce the computation cost arising in the dense matrix operation. In the context of operator learning. Galerkin transformer (Cao, 2021) proposed to remove the softmax normal- ization in self-attention and introduced linearized self-attention variant with Petrov-Galerkin projection normalization. Furthermore, hierarchical transformers using local window aggregation are proposed in Liu et al. (2021b); Zhang al. (2022) for NLP and vision applications.
3

---

Under review conference paper at ICLR 2023
3 METHODS
We introduce the hierarchical attention model in this section, motivated by the hierarchical matrix, in par- ticular the 2X² matrix variant. We follow the setup in Li et al. (2021); Lu et al. (2021) to approximate the operator S
S(a), with the imput where A and 4/ are infinite dimensional Banach spaces respectively. We aim to learn the operator S from a finite collection of finitely observed input-output pairs through parametric map. . A x Q - U and a loss functional L : U x U + R. such that the optimal parameter 0* = arg mingge Ea-y [C (N 0), S(a))]. The input can then be mapped to features f through patch embedding (Dosovitskiy et al., 2020).
A drawn from distribution μ and the output E U,
Hierarchical Discretization We introduce the hierarchical discretization of the spatial domain D, which can be used directly for time-independent PDEs. Time-dependent PDEs can be treated by taking time-sliced data as feature channels. Let Z(r) be the finest level index set, such that each index = (d1,...,,,) .,,,) E I(r) denotes the finest level spatial objects such as image pixels, finite difference points, etc. For any i (i1..... i,) E Z(r, and < m r. j(m) = (i1..... im) represents i's m-th level parent node which is the aggregate of finer level objects. [(m) ;= {i(m)) : E } is the m-th level index set. We can similarly define the natural parent-child relationship between coarser and finer nodes, which induces the index tree I together with the index sets (1 S m S r). Each index i E I(m) corresponds
token, which can be characterized e.g. by its position æ; (m)) and feature vector p(m) with number of channels C(m)) In the following, we will restrict our presentation the quadtree setting illustrated in Fig 3.1.
Reduce Operation The reduce operation defines the map from finer-level features to coarser-level fea- tures. For i E Z(m), we denote f(mm.m+1) as the set of (m + 1)-th level child nodes of i. In the quadtree setting, f0,,++1) [(i, (1, (i, (i,3)) The reduce map can be abstractly defined as f. f!!")..t (m(m) ((J) ((f((tt))))))) jjeitm.mii)) 11), which maps the 1m- level features with indices in j(m.m+1) to the m-th (jjttm.
level feature with index z(m²), For simplicity, we use a simple linear layer for the operator R(0)) in our cur-
rent implementation, namely, f( ((m)..
-R( s/ pfm++)). + R (m) p(m++1).t + +R (1m) f(a++)) + (mm) (m++)) +R (1.3)
where R(m), R(() R.(") (m) IRC xx are parametrized by linear layers. From the perspective
of H. matrices, those parametrized matrices R(()) embed low-rank approximation to the attention kernel. Please refer Appendix B for more details.
Multilevel Local Aggregation The finest level features f! E RC at the evolution step t, with index i E I(r), can be updated by the following token aggregation formula through vanilla attention,
N
atten: f (r))+ = > sus ( for i E I(),
(3.1)
with NO) := |. For simplicity, we ignore the softmax normalizing factor and only consider single-head attention, and we assume that the interaction kernel G is of the form ((f., f,) = exp(W f; WK fj) exp(q, k,), where q, == Q f., k, = W f., 2s := v(f.) W f.. and 8 w WV e Be xcer) are learnable matrices.
Instead of computing equation 3.1 explicitly with (N2) cost, propose self-attention based local aggre- gation scheme, inspired by the 2 matrices (Hackbusch, 2015), see also Appendix B. The local aggregation at the r-th level writes,
((,5) is understood the concatenation of entries of i and the scalar S 3.
4

---

Under review conference paper at ICLR 2023
M
(0)
(1)
f(r).c
atten
f
12
(2)
(3)
pirels
Reduce
Decompose
atten (r-1)
I()
T
7
f(-1)).
lec
1),41
Figure 3.1: Hierarchical discretization and in- dex tree. In this example, the 2D unit square hierarchically discretized into three levels which are indexed by I(1), I() and 7(),, respec- tively. For example,
atten loc
p00,+1
denote by (1)(1.)) {(1,0), (1, 1), (1, 2), (1,3)} } the set of the second level child nodes of the node (1).
Figure 3.2: One V-cycle of the feature update.
atten (r) f g(r),tt1
Σ (ENO)))
exp(a, t) Dof for E I(r),
(3.2)
where (V(r is the set of the r-th level neighbors of E I(r), (r). = wa b. = W k f. ((())))
qi
o/" WV g( with learnable matrices W w WV E &c Then for
1, ..... 1, we
calculate the local attention in each level with nested q.
(m)it ky fim)e (m).t.
atten
(m) loc Ji
(m),++1
Σ
exp(a,m) (m),t (m)) for i E I(m)
(3.3)
(E
where q: (m)) R(m) (fgg is the set of the m-th level neighbors of I(m)
k (()),
((m) (1
R(m)
))sstm u, (m),
(mm) ((m)((() (w/m+ j+1) Nim)(() (e()tm))))))))
Decompose Operation To mix the multilevel features f. p(mm).t+1 m = 1,.... F, we propose a decom- pose operation that reverses the reduce operation from level level r 1-1 The decompose operator ((m) f!(m (m)))t+ + (ff(++))e 1es)l.+4)) maps the m-th level feature with index i to (m + 1)- level features associated to its child set j(m.++1), f! f(m+t))e++ is further aggregated to f (
1),+1 such that f. ((m++)) + = f flm+1)),+ for i I(m+1), In the current implementation, we use a sim- D( plm))tt1 for s 0, 0,1,2,3, with parameter matrices ple linear layer such that (m+t)) (m)) E xctm)x 11))
(2.s)
Remark 1 In our current implementation, the operators R(m) and D(mm) only consist of linear layers. In general, these operators may contain nonlinear activation functions.
Remark 2 The nested learnable operators R(m) and D(mm) induce the channel mixing and structured parameterization WQ, WV W matrix for the coarse level tokens. See Appendix B for details.
We summarize the hierarchical attention algorithm in the following.
5

---

Under review conference paper at ICLR 2023
Algorithm 1 One V-cycle of Hierarchical Attention
Imput: I J!T for e I(()
STEP 0: Get the q. ". kj k 1/0 for i E I(r)
STEP For m --- 1, ,1. Do the reduce operations q; mm) R(m) (14, k (m) and o fm).t for any i E I(m).
J Jeelmmm 1)) and also for
STEP 2: For m r.,,, 1, Do the local aggregation by equation 3.2 and equation 3.3 to get p!m)ett mm I,...
for any E 2(m)
STEP 3: For
L, , 1, Do the decompose operations (5, (m+t))))+
(S( for
any E (m); then s f( + J (m+1))++ for any E I(m+1)
Output: s(r)+t+1 for any i Z(z)
Proposition 3.1 (Complexity of Algorithm 1) The reduce operation, multilevel aggregation, and decom- position operation together form V-cycle for the update of features, as illustrated in Figure 3.2. The cost of one V-cycle is O(N) I a quadtree as implemented in the paper. See Appendix for the proof.
Remark 3 (Technical contribution) In Liu al. (2021b); Zhang al. (2022), spatial resolutions are se- quentially reduced, and attentions are performed at each level separately. However, attentions are only used for fixed level, and there is multilevel aitention-based aggregation. The coarse latent variables are used for classification and generation tasks, and fine-scale information may get lost in the reduce( coarsening) process. While in the HT-Net architecture, we are motivated by multiscale numerical methods such as numerical homogenization and hierarchical matrix method. Attention- based local aggregations are con- structed at each level, and features from all levels are summed up to form the updated fine- features, which enables the recovery of the fine details with linear cost.
Decoder The decoder maps the features f (at the last update step) to the solution u. The decoder is often chosen with prior knowledge of the PDE. A simple feedforward neural network (FFN) is used in Lu et al. (2021) to learn a basis set as the decoder. A good decoder can also be learned from the data using SVD (Bhattacharya et al., 2021). In this paper, we use the spectral convolution layers in Li et al. (2021).
Loss functions Loss functions are crucial for efficient training and robust generalization of neural network models. For multiscale problems considered in this paper, we prefer use H' loss instead of the usual L loss function, as it puts more "weights" on high-frequency components. Yu et al. (2022) adopted Sobolev norm based loss function for function approximation, where H! loss is defined for the neural network func- tion itself. On the contrary, we define the I/2 loss on the target solution space, which measures the distance between the prediction := Maaa) and the ground truth u. We refer the readers Adams & Fournier (2003) for the definition of Sobolev spaces and H norm, and to the next section for the implementation.
4 EXPERIMENTS
We demonstrate the effectiveness of the hierarchical transformer model (HT-Net) on multiscale operator learning. All our examples are in 2D. We first study the ability of the model on multiscale elliptic equation with two-phase coefficients which a widely used benchmark problem for operator learning (Li al., 2021) I). As the "multiscaleness" of the coefficients such as the oscillation, contrast, and roughness of the two two-phase interface can be fine-tuned, we showcase the accuracy and the robustness of the HT-NET model for both smooth and rough coefficients. We also show that HT-NET is able to provide better generalization errors for out-of-distribution input parameters. We also test the Navier-Stokes equation with a large Reynolds number example with non-linearity and time dependence.
6

---

Under review conference paper at ICLR 2023
4.1 SETUP
In our experiments, the spatial domain is D = 1]2 and is discretized uniformly with = 1/n. Let the uniform grid be G x.) (ih; i, ...., 1}. {(aj, - are functions pairs such that u, S(a)), and aj drown from some probability measure The actual data pairs for training and testing are pointwise evaluations of a, and u on the grid G², denoted by a, and Uj, respectively. The comparison with all the baselines is consistent with (most time better than) references. The hierarchical index tree Z can generated by the corresponding quadtree representation of the nodes with depth r, such that the finest level objects are pixels or patches aggregated by pixels. See Appendix D for more details.
Empirical H loss function Empirical L loss function defined as (1)) == N Σ u Naaa; (Kaa;))lle/lee
where || ||ıa is the canonical 12 vector norm. For any & E = E2 -n/24 +2++1 E < mn/2,j = 1,2], the normalized discrete Fourier transform (DFT) co- efficients of writes (())() = / EreGe ((x))=26rae 26x0-- The empirical H loss function given by, C (L(a), ([(a),,))
+ Σ - Vaaj; (Kaa;0)lla// /|us where || == v /yEezz (())²
_ can be viewed as a weighted C loss with weights 1e1 which forces the operator to capture high- frequency components in the solution. In this work, we use the equivalent frequency space representation of the discrete H norm, and in general, discrete H norm in real space can also be adopted.
4.2 MULTISCALE ELLIPTIC EQUATION
500
8.00
200
PXP
130
500
(a) confficient
(b) reference
(c) HT-Net prediction
(d) ahs emor HT-Net (e) abs. emor FNO2D
Figure 4.1: Top: (a) smooth coefficient in Li et al. (2021), with Omax 12, Omin = and e = 9, (b), reference solution, (c) HT-Net solution, (d) absolute error of HT-Net, (e) absolute error of FNO2D: Bottom: (a) rough coefficients with amax 12, amin = 2 and 20, (b) reference solution, (c) HT-Net solution, (d) absolute emrr of HT--Net, (e) absolute error of FNO2D, The maximal absolute error in (ottom:(e) is around 900μ = 9.
We apply the HT-Net model to learn the coefficient to solution operator for equation 2.1. We use the two- phase coefficient model in Li et al. (2021), which is also benchmarked in Gupta et al. (2021); Cao (2021), I). In previous works, the coefficients do not oscillate much and the solutions look smooth. We change the pa- rameters which control the smoothness and contrast of the coefficients, such that the solutions contain more roughness than the Darcy benchmark in Li et al. (2021), Figure 4.1. We refer the readers to Appendix E for data generation details. We also include experiments for multiscale trigonometric coefficients with higher contrast. Detailed description multiscale trigonometric coefficients is in Appendix F.I.
7

---

Under review conference paper at ICLR 2023
Table Performance on multiscale elliptic equation. For the Darcy rough case, we run each experiment times to calculate the mean and the standard deviation (after =) for relative L errors (x (x10=2) and relative H' errors (x (x10-3). All experiments use a fixed wain-val-test split setup, see Appendix D for details. SWIN is constructed by adding the encoder-decoder architeture of HT-NNt to the original SWIN code, and is served as a baseline for multiscale vision transformers. FNO2D-48 and FNO2D-96 are adapted from FNO2D (by default, with 12 modes) with 48 modes and 96 modes respectively in order learn high-frequency outputs. HT-Net outperforms other neural operators by a considerable margin for all three cases.
Darcy smooth
Darcy rough
Multiscale trigonometric
Model
Runtime (8)
H
H
FNO2D
7.278
0.620
3.883
10.021
11.955.10.
1.794
12.605
FNO2D-48
8.062
0.619
2.620
1.220 +0.018
5.138 #0.093
1.565
11.093
FNO2D-96
10.969
0.575
2.437
1.216 +00024
5.140 20.281
1.518
10.106
MWT
19.715
1.138 +0.010
4.107 +0.008
1.021
7.245
GT
38.219
0.945
3.365
1.790 #0.012
6.269 #0.418
1.052
8.207
SWIN
41.417
1.622 10.047
6.796 10.359
1.489
13.385
HT-NET
33.375
0.291
0.815
0.571#0.001
1.371=000001
0.603
2.633
MWT (Gupta et al., only supports resolution with powers of two.
-35
-25
-3
-15
Epochs
Epacha
(a)
(b)
(c)
(d)
Figure 4.2: In HT-Net trained with H² loss, and (b) HT-Net trained with L loss, we show the evolution of errors with x-axis for frequency, y-axis for training epochs, and colorbar for the magnitude of L error on each frequency in log10 scale, the error for each frequency is normalized frequency- epoch 0. The loss curves with training, testing, and generalization errors are shown in (c) for HT-Net trained with H' loss, and in (d) for HT-Net trained with L loss.
by the error
Spectral Bias in Operator Learning We demonstrate the spectral bias phenomena in the neural operator training in Figure 4.2. We compare HT--Net trained with f and L losses, and show the evolution of errors in the frequency domain well the loss curves. The comparison shows that HT-Net with H' loss has faster decay for high frequencies, and the errors over all the frequencies decay more uniformly. The HT-Net trained with H loss also has better testing and generalization errors.
Generalization Errors HT-Net generalizes better on test data from samples out of distribution. In addition what we have shown in (c.d) of Figure 4.2, compare HT-Net with other models in Table 2. The models are trained with the rough two-phase dataset in Table 4.1 and tested on out-of-distribution data. FNO2D H! is FNO model but trained using our H loss. HT-Net outperforms other models by order of magnitude.
FNO2D
FNO2D H MWT
HT-NTT
256
20.27
11.49
21,901
3.822
Table 2: Relative L error (x10-2) for out distribution data, with amax 12, amin = 3 and c 18.
8

---

Under review conference paper at ICLR 2023
4.3 NAVIER-STOKES EQUATION
We consider the 2D Navier-Stokes equation in vorticity form on the unit torus, which is also benchmarked in Li et al. (2021) dew(x,t) + txx Vw(x, t) = 1 12 +)
x (0,1), (0,
V u(x t) = 0,
r xE (0,1) [0,
w(r, 0) = wo(x))
x x (0,1)
with velocity u, vorticity w = V x u, initial vorticity wo. viscosity V > (~ Re with Re being the Reynolds number), and forcing term f. We learn the operator S w(. 0 S S 10) + w(. S t T),, mapping the vorticity up to time 10 to the vorticity up to some later time T > 10. We experiment with viscosities le - 3, le 4, le 5, and decrease the final time T accordingly the Reynolds numbers increase and the dynamics become chaotic.
Time dependent neural operator Following the setup in Li et al. (2021), we fix the resolution 64 64 for both training and testing. The ten time-slices of solutions w(:, () at /t= a, 0...., 9 are taken as the input data the neural operator N which maps the solutions at the previous 10 timesteps to the next time step (2D functions to 2D functions). This procedure, often referred to as the rolled-out prediction, can be repeated recurrently until the final time T. We list the results of HT-Net, FNO-3D (convolution in space-time), FNO- 2D, U-Net in Ronneberger et al. (2015) in Table 3, and HT-Net is significantly better than other methods.
T=
T= 30
T= 30
T = 20
#Parameters
le 3
le
le
N = 1000
N = 1000
10000
N 1000
FNO-3D
5588, 537
0.0086
0.1918
0,0820
0.1893
FNO 2D
414,517
0.0128
0.1559
0.0834
0.1556
U.Net
24, 950,
0.0245
0.2051
0.1190
0.1982
IeT-Net
10, 707, 204
0,0050
0.0517
0.0194
0,0690
Table 3: Benchmarks for the Navier Stokes equation. The resolution 64 x 64 for both training and testing, and all models are trained for 500 epochs. N is the size of the training samples.
5 CONCLUSION
We have built HT-Net, hierarchical transformer based operator learning model for multiscale PDEs, which allows nested computation of features and self-attentions, and provides a hierarchical representation for the multiscale solution space. The reduce, multilevel local aggregation, and decompose operations form fine- coarse-fine V-cycle for the feature update. We also introduced the empirical H' loss to reduce the spectral bias in the multiscale operator learning. HT-Net provides much better accuracy and robustness compared with state-of sute-of-theearr (SOTA) neural operators, which is demonstrated by several multiscale benchmarks.
Limitation and outlook (1) The current implementation of HT-Neet relies on a regular grid. The extension to data clouds and graph neural networks will offer more opportunities to take advantage of the hierarchical representation.
(2) The current implementation of the attention-based operator lacks the flexibility to gener- alize to a different resolution FNO (Li et al., 2021). However, for multiscale problems, the discretization invariance FNO and MWT models may be hampered by the aliasing error the frequency domain, which is validated by the experiment results. Recent convergence analysis of operator learning methods such Kovachki et al. (2021); Lanthaler et al. (2022) leverage the smoothness of the solutions, while multiscale PDEs usually have lower regularity and do not fall into such categories. For HT-Net, it might be possible to achieve (approximate) discretization invariance by using proper sampling and interpolation modules, we expect that full treatment will be operator adaptive.
9

---

Under review conference paper at ICLR 2023
ETHICS STATEMENT
This work proposes a hierarchical transformer operator learning model for multiscale PDEs. As stated in the introduction, solving multiscale PDEs associated with some the most challenging practical problems, such reservoir modeling. fracture and fatigue prediction, high frequency scattering, weather forecasting, etc. Potentially, HT-Net solvers could help to reduce the prohibitively expensive computational cost of those simulations. The negative consequences are not obvious. Though, in theory, any technique can be misused, it is not likely to happen at the current stage.
REPRODUCIBILITY STATEMENT
The datasets for Section are either downloaded (for smooth two-phase coefficients and Navier-Stokes) from https: //gpt::///ttt hub. com/ zongyi-iii wwwgyyi-i/fourierrree www.iirrneurraa_ooo
generated (for rough Darcy coefficients) using the code from the same website. We implemented P1 finite element method in MATLAB to solve equation 2.1 with multiscale trigonometric coefficients in Appendix F.1, and in FreeFEM to solve the Helmholtz equation in Appendix F.2. We have included introductions of the relevant mathemat- ical and data generation concepts in the appendix.
We put the code for and also a link for datasets at the anonymous Github page https: // / thub. com/ Shengren- Kat / /WM////MMMM//B ICLR2023 git. Supplementary descriptions of the code also provided in the page.
REFERENCES
R. A. Adams and J. Fournier. Sobolev spaces, volume 140. Elsevier, 2003.
M. Bebendorf. Efficient inversion of the galerkin matrix of general second-order elliptic operators with nonsmooth coefficients. Math. Comp., 177-119117--1199, 2005.
M. Bebendorf. Hicrarchical matrices, volume 63 of Lecture Notes in Computational Science and Engineering. Springer-Verlee Verlag, Berlin, 2008. A means to efficiently solve elliptic boundary value prob- lems.
A. Bensoussan, J. L. Lions, and G. Papanicolaou. Asymptotic analysis for periodic structure. North Holland, Amsterdam, 1978.
L. Berlyand and H. Owhadi. Flux norm approach to finite dimensional homogemization approximations with non-separated scales and high contrast. Arch. Ration. Mech. Anal., 198(2):677-721, 2010.
G. Beylkin and N. Coult. A multiresolution strategy for reduction of elliptic PDEs and eigenvalue problems. Appl. Comput. Harmon. Anal., 5(2): 1998. ISSN 1063-5203.
Kaushik Bhattacharya, Bamdad Hosseini, Nikola Kovachki, and Andrew Stuart. Model reduction and neural networks for parametric pdes. The SMAI journal of computational mathematics, 7:121-157, 2021.
L. V. Branets, S. S. Ghai, L. L., and X.-H. Wu. Challenges and technologies in reservoir modeling. Commun. Comput. Phys., 6(1):1-23, 2009.
M. E. Brewster and G. Beylkin. A multiresolution strategy for numerical homogenization. Appl. Comput. Harmon. Anal., 2(4):327-349, 1995. ISSN 1063-5203.
Max Budninskiy, Houman Owhadi, and Mathieu Desbrun. Operator-adapted wavelets for finite-element differential forms. Journal of Computational Physics, 388:144 177, 2019.
10

---

Under review conference paper at ICLR 2023
Wei Cai, Xiaoguang Li, and Lizuo Liu. A phase shift deep neural network for high frequency approximation and wave problems. SIAM Journal on Scientific Computing, 42(5):A3285-A3312 2020.
Shuhao Cao. Choose transformer: Fourier galerkin. Advances in Neural Information Processing
Systems, 34, 2021.
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv preprint 20KY::2009..4794,, 2020.
Doina Cioranescu, Alain Damlamian, and Georges Griso. The periodic unfolding method in homogeniza- tion. SIAM Journal on Mathematical Analysis, 1064):1585-1620, 2008.
Maarten De Hoop. Daniel Zhengyu Huang, Elizabeth Qian, and Andrew M Stuart. The cost-accuracy trade- off in operator learning with neural networks. arXiv preprint arXiv:2203. 13188, 2022.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai Thomas Un- terthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words:
Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929, 2020.
Weinan E and B. Engquist. Multiscale modeling and computation. Notices Amer. Math. Soc., 5069):0062- 1070, 2003.
Weinan E. Jiequn Han, and Arnulf Jentzen. Deep learning-based numerical methods for high- dimensional parabolic partial differential equations and backward stochastic differential equations. Communications in Mathematics and Statistics, 5(4):349-380, 2017.
B. Engquist and P. F. Souganidis. Asymptotic and numerical homogenization. Acta Numerica, 17:147-190, 2008.
Yuwei Fan, Jordi Feliu-Fabá, Lin Lin, Lexing Ying, and Leonardo Zepeda-Nunez. multiscale neural network based on hierarchical nested bases. Res. Math. Sci., 6(21), 2019a.
Yuwei Fan, Lin Lin, Lexing Ying, and Leonardo Zépeda-Nánez. A multiscale neural network based on hierarchical matrices. Multiscale Modeling & Simulation, 17(4):1189-1213 2019b.
L. Greengard and V. Rokhlin. A fast algorithm for particle simulations. J. Comput. Phys., 73(2):325-348, 1987.
Gaurav Gupta, Xiongye Xiao, and Paul Bogdan. Multiwavelet-based operator learning for differential equa- tions. Advances in Neural Information Processing Systems, 34:24048-24062, 2021.
W. Hackbusch. Multigrid Methods and Applications, volume 4 of Springer Series in Computational
Mathematics. Springer-Verlag, Berlin, 1985.
Wolfgang Hackbusch. Hierarchical matrices: algorithms and analysis. Springer, Berlin, 2015.
Wolfgang Hackbusch, Lars Grasedyck, and Steffen Börm. An introduction to hierarchical matrices. In Proceedings of Equadiff 10, pp. 101-111. Masaryk University, 2002.
Helmut Harbrecht and Christoph Schwab. Sparse tensor finite elements for elliptic multiple scale problems. Computer Methods Applied Mechanics and Engineering, 200(45):3100-3110 2011.
Moritz Hauck and Daniel Peterseim. Multi-resolution localized orthogonal decomposition for helmholtz problems. Multiscale Model. Simul., 20(2):657-66 684, 2022.

---

Under review conference paper at ICLR 2023
K. Ho and Ying. Hierarchical interpolative factorization for elliptic operators: Differential equations. Communiations on Pure and Applied Mathematics, (998):1415-1411, 2016.
T. Y. Hou, X.-H. Wu, and Z. Cai. Convergence of a multiscale finite element method for elliptic problems with rapidly oscillating coefficients. Math. Comp., (8(227):113-944 1999.
Vladimir Kazeev, Ivan Oseledets, Maxim V. Rakhuba, and Christoph Schwab. Quantized tensor fem for mul- tiscale problems: Diffusion problems in two and three dimensions. Multiscale Modeling & Simulation. 20(3):893-935, 2022.
Yuehaw Khoo, Jianfeng Lu, and Lexing Ying. Solving parametric pde problems with artificial neural net- works. European Journal of Applied Mathematics, 32(3):421 2020.
Nikola Kovachki, Samuel Lanthaler, and Siddhartha Mishra. On universal approximation and error bounds for fourier neural operators. Journal of Machine Learning Research, 22:Art-No, 2021.
Samuel Lanthaler, Siddhartha Mishra, and George E Karniadakis. Error estimates for deeponets: A deep learning framework in infinite dimensions. Transactions of Mathematics and Its Applications, 6(1); tnac001, 2022.
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stu- art, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. The International Conference on Learning Representations, 2021.
Xinliang Liu, Lei Zhang, and Shengxin Zhu, Generalized rough polyharmonic splines for multiscale pdes with rough coefficients. arXiv preprint arXiv:2103.0178 2021a
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 10012-10022, 2021b.
Lu Lu, Pengzhan Jin, Guofei Pang, Zhongqiang Zhang, and George Em Karniadakis. Learning nonlinear operators via deeponet based on the universal approximation theorem of operators. Nature Machine Intelligence, 3(3):218-223 2021.
A. Mälqvist and D. Peterseim. Localization of elliptic multiscale problems. Math. Comp,, 83(290):2583- 2603, 2014.
Tan Nguyen, Vai Suliafu, Stanley Osher, Long Chen, and Bao Wang. Fmmformer: Efficient and flexible transformer via decomposed near-field and far-field attention. Advances in neural information processing systems, 0099449-244633 2021.
H. Owhadi. Multigrid with Rough Coefficients and Multiresolution Operator Decomposition from Hierar- chical Information Games. SIAM Rev., 59(1):99-149, 2017.
H. Owhadi and L. Zhang. Metric-based upscaling. Comm. Pure Appl. Math., 60(5):675-723, 2007.
Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A Smith, and Lingpeng Kong. Random feature attention. arXiv preprint arXiv:2103.02143 2021.
Nasim Rahaman, Devansh Arpit, Aristide Baratin, Felix Draxler, Min Lin, Fred A Hamprecht, Yoshua Bengio, and Aaron Courville. On the spectral bias of deep neural networks. International Conference on Machine Learning, 2019.
12

---

Under review conference paper at ICLR 2023
Maziar Raissi, Paris Perdikaris, and George Karniadakis, Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equa- tions. Journal of Computational Physics, 378:686-7077 2019.
Basri Ronen, David Jacobs, Yoni Kasten, and Shira Kritchman. The convergence rate of neural networks for learned functions of different frequencies. In Advances Neural Information Processing Systems, volume 32, pp. 4761-4771, 2019..
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In International Conference Medical image computing and computer-assisted intervention, pp. 234-241. Springer, 2015.
Justin Sirignano and Konstantinos Spiliopoulos. DGM: deep learning algorithm for solving partial differ- ential equations. Journal of Computational Physics, 375:1339-1364, 2018.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
S. Wang, B.Z. Li, M. Khabsa, H. Fang, and H. Ma. Linformer: Self-attention with linear complexity. arXiv:2006.04766 2020.
Sifan Wang, Hanwen Wang, and Paris Perdikaris. On the eigenvector bias of fourier feature networks: From regression to solving multi-scale PDEs with physics-ifformed neural networks. Computer Methods in Applied Mechanics and Engineering. 384:1 13938, 2021.
H. Xie, L. Zhang, and H. Owhadi. Fast eigenpairs computation with operator adapted wavelets and hierar- chical subspace corrections. SIAM Journal on Numerical Analysis, 57(6))2519 2019. preprint.
Jinchao Xu and Ludmil Zikatanov. Algebraic multigrid methods. Acta Numerica, 26:591-721, 2017.
L. Ying, G. Biros, and D. Zorin. A kernel-independent adaptive fast multipole method in two and three dimensions. Journal of Computational Physics, 196(2):591-626, 2004.
Annan Yu, Yunan Yang, and Alex Townsend. A quadrature perspective on frequency bias in neural network training with nonuniform data. arXiv preprint 1xXiv:2205.143000 2022.
Zizhao Zhang, Han Zhang, Long Zhao, Ting Chen, Sercan Ö. Ank,, and Tomas Plister. Nested hicrarchical transformer: Towards accurate, data-efficient and interpretable visual understanding, In AAAI Conference on Artificial Intelligence (AAAI), 2022.
Yinhao Zhu and Nicholas Zabaras.
Bayesian deep convolutional encoder-decoder networks for surrogate modeling and uncertainty quantification. J. Comput. Phys., 366:415 447, 2018.
A MULTISCALE PDE
In this section, introduce some mathematical and numerical concepts related to multiscale PDEs.
13

---

Under review conference paper at ICLR 2023
A.1 ASYMPTOTIC HOMOGENIZATION
We introduce some basic formulation of asymptotic homogenization, by assuming that a(x) = a(²) with 1 1-periodic funciton a(-) in equation 2.1, namely,
- div (a (#) Vur) = f. xED,CR' a(-) periodic in x/Ee, u =0 0,
I EEDD.
can be derived from the asymptotic expansion of the two-scale function u² = u(xxx1/0 e, and justified rigorously that u² - uo in H', with (x) the solution of the homogenized problem
div d V uu) f, E OD.
D,
0,
which contains only coarse scale information. The homogenized coefficient a can be computed by the formula di [e + Vx.)" a(y)(e; + Vxj)dy, where solves the cell problems given by
div, (a (V,X: e;)) 0, x zeY. Xi Hp (Y))
with Y being the d dimensional torus. Asymptotic homogenization provides an approximate solution of Ug - Σ Xs (z) due such that ||ue d' m S czt. We note that to find the homogenization approximation yz, af only requires the coarse-scale solution ui and precomputation of d cell problems for Xi which do not depend on f and D.
A.2 NUMERICAL HOMOGENIZATION AND MULTILEVEL/MULTIGRII METHODS
Given the smallest scale of the multiscale problem and a coarse computational scale determined by the available computational power and the desired precision, with < H << 1. The goal of numerical homog- enization is to construct a finite-dimensional approximation space Vи and to seek an approximate solution "H VH, such that, the accuracy estimate |/u I CH holds for optimal choices of the norm ||. and the exponent a, and optimal computational cost holds with Vi constructed via pre-computed subproblems which are localized, independent and do not depend the RHS and boundary condition of the problem.
recent two decades, great progress has been made in this area (Hou et al., 1999; Mälqvist & Peterseim, 2014; Owhadi & Zhang. 2007), approximation spaces with optimal accuracy (in the sense of Kolmogorov N-width (Berlyand & Owhadi, 2010)) and cost can be constructed for elliptic equations with fixed rough coefficients. For multiscale PDEs, operator learning methods can be seen as a step forward from numerical homogenization, since they can be applied to an ensemble of coefficients, and the decoder can be interpreted the basis of the underlying problem well.
For multiscale PDEs, multilevel/multtrrid methods can be seen as the multilevel generalization of the nu- merical homogenization methods. Numerical homogenization can provide coarse spaces with optimal ap- proximation and localization properties. Recently, operator-adapted wavelets (gamblets) (Owhadi, 2017; Xie al., 2019) have been developed, and enjoy three properties that are ideal for the construction of efficient di- rect methods: scale orthogonality, well-conditioned multi-resolution decomposition, and localization. Gam- blets has been generalized to solve Navier-Stokee equation (Budninskiy et al., 2019) and Helmholtz equation (Hauck & Peterseim, 2022) efficiently.
B HIERARCHICAL MATRIX PERSPECTIVE
The hierarchically nested attention in Algorithm 1 resembles the celebrated hierarchical matrix method (Hackbusch, 2015), in particular, the 2 matrix from the perspective of matrix operations. In the following,
14

---

Under review conference paper at ICLR 2023
(0,0)
(0,1)
(1,0)
(1,1)
(0,0,0)
(0.0.1)
(0,1,0)
(0.1.1)
(1,0,0)
(1,0.1)
(1.1,0)
(1.1.1)
V
MATM 10000 WOTM 10T1 (6.0.)) 1811.) (0.110) (M111)) (10000) (1.001) (10.3)) (191111 (1:00)) (1:00) (11.1)) (111.))
Figure B.I: Hierarchical discretization of ID domain. The coarsest level partition is plotted as the top four segments in pink. The segment (0, 0) is further partitioned into two child segments (0, 0,0) 0) and (0, 0, 1). During the reducing process, the computation is performed from bottom to top to obtain coarser level tokens, for example, the (0, 0) tokens obtained by learnable reduce operations R(2) and (()) tokens (0, 0, 0) and (0, 0, 1). When generating the high-resolution tokens/features, the computation is from up to bottom by
learnable decomposition operations D(1) and D(2), The red frames show examples attention windows each level.
we take the binary tree-like hierarchical discretization shown in Fig B. as an example illustrate the reduce operation, decompose operation, and multilevel token aggregation in Algorithm in matrix formulas.
STEP 0: Given features (r),,, compute the queries ("),, keys k(r) and values by (r), for I(r).
Starting from the finest level features f( I(r), the queries (r)). can be obtained by
W
gf fr
p!isae
|
WQ.60
[2(e]]
the keys (r), and values (r),t can follow the similar procedure.
STEP For m 1: 1. Do the reduce operations q. (m).. k (m). and o (m). for any i (m))
R(im
(m+ (fg m41) and also for
(m) is linear, the reduce operations correspond to
q (m),
= Rim,e
q
(m+
The reduce
15

---

Under review conference paper at ICLR 2023
matrix given by
(m) Rg
R(m)
(
(
R
|z(m|
R
Rim
12 10
and R("") R T(T [[(m=1)xC(m) are matrices parametrized by linear layers in our paper (in prac- tice, queries, keys, and values use different R Rn (m) to enhance the expressivity). In general, these operators R(m) are not limited to linear operators. The composition of nonlinear activation functions would help increase the expressivity. The nested learnable operators R(mm) also induce the channel mixing and is equivalent to structured parameterization of W WV WK matrix for
the coarse level tokens, in the sense that, inductively,
q
m))
R(m)
(r-1))
(r).t
W.trr)
Wq.(r)
R(m) ((r)))
W
STEP 2: With the m-th level queries and keys, we can caleulate the local attention matrix (m), at m-tth level with G(m), Joc.i.j == exp(q: (m)). (m)) for i E (/m) or ~ j.
STEP 3: The decompose operations, opposite to the reduce operations, correspond to the transpose of the following matrix in the linear case,
D( ((
[z(m)
D
D (m))
[z(m-1)
The m-th level aggregation in Figure 3.2 contributes to the final output f())++1 in the form
D (r-1), (m).T (m)
R (m)
(
()) 2
Eventually, aggregations at all r levels in one V cycle
loc
can be summed up as
(--1).T
(m).T (m) (m) (m) loc
........" ...((( ((r-1)) + (r)
ofrr)e
(B.1)
loc
The hierarchical attention matrix G₁, Σm (D (r))), D
G (m)
R (m) ... ((.
ald in equa-
Goo
tion B.1 resembles the three-level 2 matrix decomposition illustrated in the following Figure B.2,
16

---

Under review conference paper at ICLR 2023
also refer to Hackbusch (2015) for a detailed description. The sparsity of the matrix lies in the fact that the attention matrix is only computed for pairs of tokens within the neighbor set. The H² matrix-vector multiplication in B.1 also implies (N) complexity the Algorithm I.
+
+
Figure B.2: A demonstration of the decomposition of attention matrix into three levels of local attention matrix.
Note that, the local attention matrix at level Z()) (pink), level [(2) (blue) and level Z (green are G G(22) and G respectively. However, when considering their contributions to the finest level, they are equivalent to the attention matrix D (2).T (1).T
loot
(2) RZ R (pink), (2)) ( R(2) Loc R x R (blue) and (33) (green), as demonstrated in Figure B.2. Each pink block and blue block are actually low-rank sub-matrix with rank ((1) and rank ((2), respectively, by definition.
PROOF OF PROPOSITION 3.1
Proof For each level m, the cost to compute equation 3.3 el e((Z(m))((m) since for each i E I(m) the cardinality of the neighbour (1))(() bounded by constant c. The reduce operation f((-1)
(1 Jgg((t1,,))) costs at most [Z] ((m))i((((() flops and so does the decompose operation at the same level. Therefore, for each level, the operation cost c((I(m) (((m)) + ()I(m)))((m) ((mm)) When
a quadtree, ((r) N, I(--1) N/4, ((() 1) 4, therefore the total computational cost O(N). 0
D IMPLEMENTATION DETAILS
In the HT-Net implementation, we follow the window attention scheme in Liu et al. (2021b) for the definition of the neighborhood (-) in equation 3.2 and equation 3.3. In this paper, we choose r = 3 as the depth of the HT-Net, GeLU as the activation function, and CNN-based patch embedding module to transfer the input data into features.
For dataset with resolution ny X ny, such as in the multiscale elliptic equation benchmark 4.2, the input feature (3) is represented as a tensor of size n X n x via patch embedding. The self-attention first computed within local window level 3. Then the reduce layer concatenates the features of each group of 2 x neighboring tokens and applies a linear transformation on the 4C--dimensional concatenated features on 5u 2 level tokens, obtain level features (2) as tensor of the size x 2C. The procedure is repeated from level 2 to level with (2) of size x x 4C.
We adopt the window attention scheme from Liu al. (2021b). One refer to the implementation de- tails there. For the decompose process, starting at level 1, linear layer is applied to transform the 4C-- dimensional features (1)) into SC-dimensional features. Each level token with SC-dimensional features is decomposed into four level 2 tokens with 2C-dimensional features and added the level 2 feature (2)
17

---

Under review conference paper at ICLR 2023
with output size of 3 x / 2C. The procedure is repeated from level 2 level to level with the output f(3 of size x x C.
We use the Adam optimizer with learning rate 1c-3, weight decay 1c-4 and the 1-cycle schedule as in Cao (2021). We use batch size 8 for experiments in Sections 4.2 and 4.3, and batch size 4 for experiments in Appendices F. and F.2.
Evaluation setup For the Darey rough case in Table 1, we use a train-val-test split of the dataset with sizes 1280, 112, and 112, respectively. Each run carried out with different seeds. For the Darcy smooth and multiscale trigonometric cases, we use a train-val-test split of the dataset with sizes 1000, 100 and 100, respectively. For the Navier-Stokes experiment in Section 4.3, we have tried N 1000 and N 10000 training samples as shown in Table 3, and the number of testing samples is 100 and 1000 respectively. The baseline's 's implementations are based on their official implementation if provided publicly. The results are comparable with the ones reported in the references.
All experiments are run on a NVIDIA A100 GPU.
E DATA GENERATION FOR THE TWO-PHASE COFFFICIENT IN SECTION 4.2
The two-phase coefficients and solutions are generated according to https://githib com/ and used operator learning benchmarks Li et al. (2021); Gupta et al. (2021); Cao (2021). The coefficients a(x) are generated according to a ~ μ == V (0, (-A + cl) with zero Neumann boundary conditions on the Laplacian. The mapping / R IR takes the value amax the positive part of the real line and Omin on the negative part. The forwafo is defined in a pointwise manner. The forcing term is fixed as f() = 1. Solutions u are obtained by using second-order finite difference scheme Qmax and Omin can control the contrast of the coefficient. The parameter e can control the roughness (oscil- lation) of the coefficient, and the coefficient with a larger e has rougher two-phase interfaces, shown in Figure 4.1.
zongyi-li/font htt//forrier_eeeee
suitable grid. The parameters
F ADDITIONAL EXPERIMENTS
F.1 MULTISCALE TRIGONOMETRIC COEFFICIENT
In this experiment, we consider equation 2.1 with multiscale trigonometric coefficient adapted from Owhadi
6
(2017), such that = [-1, a(x) = II (1 + 1 cos(a + x2))) + 1/2 (an - 3z1))), with k=1
4g uuniform2 1.5 x 1) and fixed f(r) 1. The reference solutions are obtained using Pi FEM on a 1023 1023 grid. Datasets lower resolution are sampled from the higher resolution dataset by linear interpolation. The experiment results for the multiscale trigonometric case with different resolutions are shown in Table 4. HT-Net obtains the best relative 1 error compared to other neural operators at various resolutions 600 epochs. Multiwavelet neural operator MWT has the second best performance it also possesses a multiresolution structure, but does not adapt to the PDE. is not surprising that FNO, an excellent smoother which filters higher frequency modes, fails to capture the high-frequency oscillations of the solution. In contrast, our method has better performance in this respect. See Figure F. for illustrations of the coefficient and comparison of the solutions at the slice 0.
18

---

Under review conference paper at ICLR 2023
spochs=300
spochs=600
B=128
n==56
m=512
m=128
B=256
n=512
FNO
1,996
1.842
1.877
2.017
1,820
1.806
GT
1.524
1,070
1.093
1.448
0.938
0.970
MWT
1,11
1.007
1.006
1.112
0.985
0.977
SWIN
1.788
2.378
1.579
2.216
0.581
0.671
0.659
0.537
0.614
0.631
Table 4: Relative error (x (x10-2) of the multiscale trigonometric example.
(a) coefficient
(b) slices
0.7
Fedd
HT-06
0.6
0T SWIN
RNO
0.6
0.4
0.3
02
0.7
-0.
-0.2
0.2
0.5
Figure F.1: (a) multiscale trigonometric coefficient, (b) slice of the solutions at æ 0.
F.2 HELMHOLTZ EQUATION
We test the performance of HT-Net for the acoustic Helmholtz equation in highly heterogeneous media as an example of multiscale wave phenomena, whose solution is considerably expensive for complicated and large geological models. We adopt the setup from Hoop et (2022), for the Helmholtz equation on the domain D [0, 1j2. Given frequency 103 and wavespeed field e Ω - IR, the excitation field u Q + I solves the equation
W
A
in Ω
(x)
du
on Ø2g 002, 8Q4
In
du
on Mn,
Ön
where 02 is the top side of the boundary, and 00,,2 are other sides. The wave speed field is c(x) 20 tān)( where is sampled from the Gaussian field ē N' (-A + 72) -d, where = 3 and d = 2 are chosen to control the roughness. The Helmholtz equation is solved on a 100 x 100 grid by finite element methods. We aim to learn the mapping from R 100x 100 to u E I 100x 100 shown in Fig F.2.
19

---

Under review conference paper at ICLR 2023
0.22
-002
-203
(a) wavespeed field
(h) excitation field u
(c) HT-Nee prediction
Figure F:2: The mapping u
The Helmholtz equation is notorious for solving. One reason is the so-called resonance phenomenon the frequency w is close to an eigenfrequency of the Helmholtz operator for some particular c. We found that it is necessary to large training dataset of size 8000. The test dataset is of size 800. All models trained within 500 epochs.
Model
#Parameters
Evaluation time (ms)
L rolative ervor (x10 ")
3,282,068
UNet
17,2
0.1
42.90
47,632,003
Table 5: HT-Net outperforms the other models by a wide margin. The eπror of FNO is comparable with the reported results in De Hoop et (2022). We note that the four models benchmarked for the Helmholtz equation in De Hoop et al. (2022), including FNO and DeepONet, failed to reach an error less than 1 10-2.
We also compare the evaluation time of the trained models in 5. Compared to HT-Net, FNO has both a larger error and longer evaluation time. UNet, CNN based method, has much faster evaluation time (30 times faster than HT-Nett but the worst error (60 times as much as HT-Net).
F.3 COMPARISON WITH CLASSICAL METHODS (FDM) AND MULTISCALE METHODS (GRPS)
We compare the accuracy and efficiency of HT-Net with two conventional solvers finite difference method (FDM) as a typical classical method, and generalized rough polyharmonic splines (GRPS) (Liu et al., 2021a) typical multiscale method, for the Darcy rough benchmark. The results are listed in Table 6.
Model FDM
Evaluation time
relative 1.2
0.34
0.84
GRPS
18.9
0.02
NeT-Net
0.003
0.58
Table 6: Relative 2 error and evaluation time on the Darey rough benchmark. We implement FDM and in MATLAB and measure the solution time on CPU Intel(R) Core(TM) 17-10510U CPU @ 2.30 GHz. We measure the evaluation time of HT--et HT-Net has comparable accuracy but needs much less time for evaluation. The reference solution is defined grids 512 512, sampled from the test dataset. FDM solves the problem 256 x 256 grid. GRPS uses coarse bases of dimension 256 to solve the same problem. HT--ee learns the solutioin operator on the same resolution.
an NVIDIA A 100 GPU. Compared with classical methods,
20

---

Under review conference paper at ICLR 2023
F.4 STUDY FOR HYPERPARAMETERS
Model
H' relative (x10->)
L relative (x10--)
HT-Noo [A, 80]
L843
0.648
[3, 128]
HT--Net [3, 80]
1,898
0.710
HT-Net 13, 4. 641
1.909
0.695
Net 80]
2,030
0.707
HT-Net [2,8, 80]
1.903
0.701
Table 7: Hyperparameter study.
We conduct studies for hyperparameters such number of hierarchical levels (depth), window size, and feature dimension. We list the results for the Darcy rough benchmark in Table 7, where the notation HT- Net [3, 8, 80] represents HT-Net with 3 levels, window size 8, and feature dimension 80. The experiments are run for 100 epochs. The results show that larger values of hierarchical level, window size, and feature dimension might be helpful reduce errors, though with higher computational cost. To balance the model size, computational cost, and performance, we choose the hyperparameters as [3, 8, 80].
F.5
MEMORY USAGE
Darcy smooth (res=211 x111)
Darcy rough (res=256
Model
Mem
CUDA Mem
Mem
CUDA Mem
FNO
1.80
1.19
2.53
1.72
11.64
MWT
2.54
8.36
HT-NET
4.17
7.09
4.98
8.91
Table 8: The memory usage of different models. The CUDA mem(GB) is the sum of sel _cuda_meemorrr memory .usuage from the PyTorch autograd profiler for I backpropagation. The mem(GB) is recorded from nvidia smi of the memory allocated for the active Python process during profiling.
We report the memory usage of different models for the Darey smooth (with resolution 211 x211) and Darey rough (with resolution 256x 256x256) benchmarks in Table F.5. The table shows that the memory usage of HT- Net is stable. For the resolution of 256x256, both MWT and GT consume larger or comparable CUDA memory compared with HT-Net, albeit HT-Net has much better accuracy.
21


Reviews:
[{"review_id": "nW0udWdAZq", "reviewer": "Reviewer_uEBd", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper introduces a new hierarchical method to solve multiscale PDE problems of structured domains using attention-based temporal updates and hierarchical matrix reduce and decomposition. The method is evaluated on standard benchmark 2D multiscale PDE problems, showing promising results in comparison to existing work. \n", "summary_of_the_review": "The task considered in this paper of solving multiscale PDE problems is highly challenging and complex. The proposed method tackles the problem in a novel approach, suggesting to compute hierarchical features in a v-shape updates procedure, going over all different hierarchical levels. As mentiond above, this is unfortunately a tough read which may potentially limit the impact of the approach and its wider accpetance. This in itself does not mean the paper should not be published, but rather I encourage the authors to make the effort to improve their exposition.\n", "strength_and_weaknesses": "The main strength of the paper is the novelty of the approach. To the best of my knowledge, this work is first to suggest a hierarchical neural method based on the hierarchical matrix approach. The machinery involved in devising such an approach is based on a large body of theory and practice. The application domain of multiscale PDE is challenging, and the evaluation is convincing.\n\nThe main weakness of the paper is its exposition. This is unfortunately a hard read, which probably limit the impact of the paper outside the sub-community of PDE and numerical methods savvies. On the one hand, compressing the necessary details to nine pages is definitely a challenging task. On the other hand, illustrations of key operations would definitely help the reader. For instance, the discussion related to hierarchical discretization could benefit from an illustration showing a specific case. A similar comment can be made regarding the reduce and decomposition actions. I would also consider to de-clutter the current notations. Perhaps instead of using index notations, you can switch the operators performing reduce/decomposition.\n\nAnother potential weakness is the evaluation setup. How do you choose hyper-parameters for the network? Is there a validation set? How many times every experiment is run? Are the results reported in e.g., Tab. 1 for the basline approaches reported elsewhere (say in Li et al. 2021) or these are new numerics you achieved? Essentially, all these questions aim toward understanding whether there is a mature benchmark for this problem domain. \n\nFinally, briefly discussing the limitations of the approach is needed. For instance, is it possible to use the method (probably not as is) for unstructured domains such as point clouds or general domains where natural neighborhood information is not avilable?\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "SI5pjnwb6-", "reviewer": "Reviewer_oKvB", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper presents a multiscale method for operator learning based on a transformer architecture. Inspired by hierarchical methods for solving PDEs, the key idea is to add self-attention layers to a architecture with multiple downscaling and upscaling layers (i.e., reduce and decompose operations). The authors also use a Sobolev norm to train the network, which they find improves the recovery of high-frequency signal content. The method is evaluated by learning solution operators for multiple different differential equations, and appears to outperform other recent methods.", "summary_of_the_review": "Overall, I lean slightly negative on this paper. The results appear to outperform the baselines, but I'm curious why the recent transformer architecture for neural operators was not included in the baselines, and there seem to be some missing evaluation items: analysis of runtime/memory, ablation study of architectural decisions (since the architecture is the main contribution), and qualitative comparison to baseline approaches. Perhaps the authors can address some of these concerns in their response, though.", "strength_and_weaknesses": "The method appears to outperform the baseline approaches for learning neural operators by a fair margin, and so may be of interest to the large community working on physics-informed neural networks and neural operators.\n\nI also find the appendix to be very helpful in supporting some of the arguments made in the main paper. For example, the elaboration on the connection to hierarchical matrix methods and additional implementation details.\n\nThe architecture and method are also fairly straightforward to understand.\n\nOn the flip side, since the main technical contribution of the paper seems to be the architecture, I was expecting to see more from the evaluation. For example, the following points.\n- There are no qualitative comparisons to the baselines; the only qualitative results (showing the predicted solutions) in the main paper are two output solutions of the elliptic equation from the proposed method.\n- While there is a discussion of computational complexity, it's not clear what the actual practical runtime and memory usage are. One of the main advantages of neural operators is that they usually offer significant improvements in runtime compared to conventional solvers (albeit at the cost of accuracy). In that respect, transformers are usually more computationally expensive than, e.g., CNN-based architectures, and so I think it's critical to include some analysis of the computational performance. \n- The authors mention the \"Galerkin Transformer\"  (Cao 2021) as another transformer-based architecture introduced in the context of operator learning. I wonder why this was not evaluated? How does that method compare to the proposed architecture?\n- There is no discussion/evaluation of some of the architectural choices. For example, how does the choice of window size affect the performance?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "MlWdmfR3I6q", "reviewer": "Reviewer_Pg7t", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper uses a hierarchical attention based neural networ to learn the solution operator associated to multiscale PDEs. In addition, the second contribution is the sobolev type norm used as the loss function, giving more weight to the higher frequencies of the target. ", "summary_of_the_review": "In the paper's current state, I would lean towards recommending rejection. However, I am willing to reconsider.", "strength_and_weaknesses": "I have mixed feelings about this paper. I have appreciated the connection made with hierarchical matrices. The experimental results are also interesting, assuming the experiments have been conducted correctly (see comments below). \n\n\nHowever, relation with previous work seems to have treated very poorly, which makes the novelty of the paper hard to evaluate. \n\nFor instance, the sentence \"The idea is reflected in Liu et al. (2021); Zhang et al. (2022) to a certain degree.\" is vague and differences should be motivated and made much more precise, given most of the code is adapted from Liu et al. (2021). \n\nFor the second contribution (H1 loss), there is no mention to previous work, even though Sobolev type loss have already been proposed to train neural networks, e.g. in [Yu].\n\nFor the experiments section: Testing on a common number of epochs is not a good practice: the number of epochs should not be the same for each model, but chosen using the validation set. \n\nAs a more adapted baseline, have you tried FNO, preserving more modes in the Fourier transform in order to learn high-frequency outputs?\n\nLastly, it is not clear for me why this architecture would correspond to a Neural Operator. Can the network take as input arbitrary data (do you take as input positional encodings?). Does it satisfy the condition of \"discretization invariance\" [Kovachki]? I would like a more careful discussion on this point. \n\n\n[Yu]: https://arxiv.org/pdf/2205.14300.pdf\n\n[Kovachki]: https://arxiv.org/pdf/2108.08481.pdf", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "F0ufUNbhcf8", "reviewer": "Reviewer_2aoA", "summary": "", "questions": "", "limitations": "", "rating": 10, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The manuscripts an attention based multi-grid solution for PDEs where different scales are present in the solution. This is done by a hierarchical attention structure that is used to map low feature solutions to high feature solutions and back. The hierarchy is created with a network architecture that resembles U-Net of Ladder architectures, spiced with transformer layers. The manuscript describes a hierarchical discretisation that have features at multiple scales. Reduce operators go to coarser level and decompose operation does the opposite.\n\nA good choice as these have been shown to be able to handle multiple resolutions in an efficient manner using the transformer architecture.   The result are very encouraging and provide significant improvement over the state-of the -art. ", "summary_of_the_review": "The attention seems to be last world in many high performance application.  I think that in this case really for being able to bring the prior art features to the fine grid solution  in the decompose operator.  It would be nice to know (or have discussion) what are the salient features that are in the coarse solution that enable a successful generation of the high resolution solution. it is definitely coming as a prior from the training data.  All in all an excellent paper definitely of interest to people finding methods to solve PDEs in efficient way", "strength_and_weaknesses": "The efficient hierarchical operator learning of the input-output mapping of parametric PDEs using attention. The emphasis on getting the high frequency components correctly is also required for performance.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject