Random Weight Factorization improves the training of Continuous Neural Representations
Authors: Sifan Wang
Hanwen Wang
Jacob H Seidman
Paris Perdikaris
Venue: Submitted to ICLR 2023 (2023)

Abstract: Continuous neural representations have recently emerged as a powerful and flexible alternative to classical discretized representations of signals. However, training them to capture fine details in multi-scale signals is difficult and computationally expensive. Here we propose random weight factorization as a simple drop-in replacement for parameterizing and initializing conventional linear layers in coordinate-based multi-layer perceptrons (MLPs) that significantly accelerates and improves their training. We show how this factorization alters the underlying loss landscape and effectively enables each neuron in the network to learn using its own self-adaptive learning rate. This not only helps with mitigating spectral bias, but also allows networks to quickly recover from poor initializations and reach better local minima. We demonstrate how random weight factorization can be leveraged to improve the training of neural representations on a variety of tasks, including image regression, shape representation, computed tomography, inverse rendering, solving partial differential equations, and learning operators between function spaces.

Under review as a conference paper at ICLR 2023
RANDOM WEIGHT FACTORIZATION IMPROVES THE TRAINING OF CONTINUOUS NEURAL REPRESENTA- TIONS
Anonymous authors Paper under double-blind review
ABSTRACT
Continuous neural representations have recently emerged powerful and flex- ible alternative to classical discretized representations of signals. However, train- ing them to capture fine details in multi-scale signals is difficult and computa- tionally expensive. Here we propose random weight factorization as simple drop-in replacement for parameterizing and initializing conventional linear layers in coordinate-bass multi-layer perceptrons (MLPs) that significantly accelerates and improves their training. We show how this factorization alters the underlying loss landscape and effectively enables each neuron in the network to learn using its own self-adaptive learning rate. This not only helps with mitigating spectral bias, but also allows networks to quickly recover from poor initializations and reach better local minima. We demonstrate how random weight factorization be leveraged to improve the training of neural representations on variety of tasks, including image regression, shape representation, computed tomography, inverse rendering, solving partial differential equations, and learning operators between function spaces.
INTRODUCTION
Some of the recent advances in machine learning can be attributed to new developments in the de- sign of continuous neural representations, which employ coordinate- based multi-layer perceptrons (MLPs) to parameterize discrete signals (e.g. images, videos, point clouds) across space and time. Such parameterizations are appealing because they are differentiable and much more memory ef- ficient than griplsampled representations, naturally allowing smooth interpolations to unseen input coordinates. As such, they have achieved widespread success in variety of computer vision and graphics tasks, including image representation (Stanley 2007) Nguyen et al. 2015)) shape represen-
tation (Chen & Zhang 2019 Park et al. [2019] Genova et al. 2020), view synthesis (Sitzmann cet al. [019]] Saito et al. [219]] Mildenhall et al.) 2020) Niemeyer et al. (020),, texture generation (Oechsle et al.. [219 Henzler al. (020), etc. Coordinate- based MI MLPs have also been applied to scientific computing applications such as physics-informed neural networks (PINNs) for solving forward and inverse partial differential equations (PDEs) Raissi et al. (2019) 2020): Karniadakis et al. (2021). and Deep Operator networks (DeepONets) for learning operators between infinite-
dimensional function spaces Lu al. 202T
Wang et al. 2022ec)
Despite their flexibility, it has been shown both empirically and theoretically that coordinate- MLPs suffer from "spectral bias" (Rahaman et al., 2019 Cao et al., 2019: Xu al.) 2019). This manifests as a difficulty in learning the high frequency components and fine details of a target func- tion. popular method resolve this issue is embed input coordinates into a higher dimensional space, for example by using Fourier features before the MLP (Mildenhall et al. 20200) Tancik et al.) 2020). Another widely used approach is the use SIREN networks (Sitzmann et al. 2020), which employs MLPs with periodic activations to represent complex natural signals and their derivatives. One main limitation of these methods is that a number of associated hyper parameetrss (e.g. scale factors) need to be carefully tuned in order avoid catastrophic cemeeraizaaiio//neepooooo Unfortunately, the selection of appropriate apper-parametes typically requires some prior knowl- edge about the target signals, which may not be available in some applications.
errors.

---

Under review as a conference paper at ICLR 2023
More general approaches to improve the training and performance of MLPs involve different types
of normalizations, such as Batch Normalization (loffe & Szegedy 2013), Layer Normalization [a
et al., 2016 and Weight Normalization (Salimans & Kingmal 2016). However, despite their markable success in deep learning benchmarks, these techniques are not widely used in MI-P-basdd neural representations. Here draw motivation from the work of (Salimans & Kingma 2016 Wang et al. 2a2 and investigate a simple yet remarkably effective re-parameterization of weight vectors in MLP networks, coined as random weight factorization, which provides generalization of Weight Normalization and demonstrates significant performance gains. Our main contributions are summarized as
We show that random weight factorization alters the loss landscape of a neural representa- tion in a way that can drastically reduce the distance between different parameter configu- rations, and effectively assigns a self-adaptive learning rate to each ncuron in the network.
We empirically illustrate that random weight factorization can effectively mitigate spectral bias, as well as enable coordinate-basee MLP networks to escape from poor intializations and find better local minima.
We demonstrate that random weight factorization can be used as a simple drop-in enhance- ment to conventional linear layers, and yield consistent and robust improvements across wide range of tasks in computer vision, graphics and scientific computing.
WEIGHT FACTORIZATION
Let æ E IRd be the input, g (x)) = x and do = d. We consider a standard multi-layer perceptron
(MLP) fz( recursively defined by
f (x) = w(() (x) + b0) (z) = (x)),, 1 1= 1-1,,.....
(2.1)
with final layer
fo(x) w((+1) ( + +b((+1)
(2.2)
where W Rdixd-1 is the weight matrix in 1-th layer and is element-wise activation function. Here, O (W(1), (()) ..... W(++)) b((+))) represents all trainable parameters in the network.
MLPs are commonly trained by minimizing appropriate loss function (()) via gradient descent. To improve convergence, we propose to factorize the weight parameters associated with each neuron in the network as follows
w (ki)) (k..)) (k:1) = di, 1= 1,,,,...,LL+
(2.3)
where (k.)) E Rd-1 is a weight vector representing the k-th row of the weight matrix W g(k.)) E R is a trainable scale factor assigned to each individual neuron, and v(kkd) E gd Conse- quently, the proposed weight factorization can be written by
( = diag(s V(), 1 1,2, 1.2,..,,444 1+1.
(2.4)
with Rd
2.1 A GEOMETRIC PERSPECTIVE
In this section, we provide geometric motivation for the proposed weight factorization. To this end, we consider the simplest setting of a one-parameter loss function f(w). For this case, the weight factorization is reduced to w s with two scalars 8, v. Note that for a given w o 0 there are infinitely many pairs (s, v) such that w s u. The set of such pairs forms family of hyperbolas in the su- cl-plane (one for each choice of signs for both s and v). As such, the loss function the su-plane constant along these hyperbolas.
2

---

Under review as a conference paper at ICLR 2023
S visual illustration of the difference between the original loss land- scape function of w versus the loss land- scape the factorized se-plane. In the left panel, we plot the original loss function as well as an initial parameter point, the local minimum, and the global minimum. The right panel shows how in the factorized parameter space, each of these three points corresponds to two hyperbolas in the su-plane. Note how the distance between the initialization and the global minima is reduced from the top the bottom panel upon appropriate choice of factorization. The key observation is that the distance between factorizations representing the initial parameter and the global minimum become arbitrarily small in the su-plane for larger values of s. Indeed, we can prove that this holds for any general loss function in arbitrary parameter dimensions (the proof provided in Appendix A..I).
Original {(w)
loss landscape Initialization Local minima Global minima
v Transformed loss landscape
weight factorization:
w v
t Weight factorization transforms loss landscapes and shortens the distance to minima.
Figure 1:
Theorem 1, Suppose that C(0) ) the associated loss function of a neural network defined in equa- tion 2. and equation 2.2 For a given 0, we define Uo as the set containing all possible weight
factorizations
Uo = (o=(100) v0 diag(s!)). v()) W
1= 1=,,... 1-,,..,, L.I.
(2.5)
Then for any 0, 0', we have
disi(Uo, Uor)
min wEU NEU!
||
yl = 0.
(2.6)
2.2 SELF-ADAPTIVE LEARNING RATE FOR EACH NEURON
different way examine the effect of the proposed weight factorization is by studying its associ- ated gradient updates. Recall that standard gradient descent update with a learning rate takes the form
Oc
(k.)) wilk (k.))
(2.7)
Dwn. (k.t)
The following theorem derives the corresponding gradient descent update expressed in the original parameter space for models using the proposed weight factorization.
Theorem 2. Under the weight factorization of equation
2.3 the gradient descent update is given by
(k.1) w a+1 =w
(k:1)
(15 (1.00
11 110011
@w (k.) + O(12),
OL
(2.8)
for 1,2, and l 1.2.....
we observe The proof provided in Appendix A.2 By comparing equation 2.7 and equation 2.8
that the weight factorization w = v re-scales the learning rate of w by a factor of (s² + || 12). Since 8, o are trainable parameters, this analysis suggests that this weight factorization effectively assigns self-adaptive learning rate to each neuron in the network. Similar analysis has been ex- plored in the context of deep linear networks suggesting that the weight factorization can signifi- cantly accelerate the optimization In the following sections, weight factorization (under an appropriate initialization of the scale factors), not only helps with mitigating spectral bias (Rahaman et al.) 2019 Bietti & Mairal 2019 Tancik et al. 2020 Wang et al., 2021c), but also allows networks to quickly move away from a poor initialization and reach better local minima faster.
will demonstrate that the proposed

---

Under review as a conference paper at ICLR 2023
2.3
RELATION TO EXISTING WORKS
The proposed weight factorization is largely motivated by weight normalization (alimans & Kingma 2016), which decouples the norm and the directions of the weights associated with each neuron as
w
||||"
(2.9)
where g ||2]||, and gradient descent updates are applied directly to the new parameters v, Indeed, this can be viewed as a special case of the proposed weight factorization by setting = || in equation 2.3 In contrast weight normalization, our weight factorization scheme allows for more flexibility in the choice of the scale factors 8.
We note that SIREN networks (Sitzmann et al. 2020) also employ a special weight factorization for each hidden layer weight matrix,
W wp * W
(2.10)
where the scale factor we E R is a user-defined hyper-parameter. Although the authors attribute the success of SIREN to the periodic activation functions in conjunction with a tailored initialization scheme, here we will demonstrate that the specific choice of wi is the most crucial element in SIREN's performance, see Appendix G for more details.
worth pointing out that the proposed weight factorization also bears some resemblance the adaptive activation functions introduced in (Jagtap et al.) 2020), which modifies the activation of each neuron by introducing an additional trainable parameter a as
(x) (a (x)))
(2.11)
In practice, the scale factor is generally initialized as a = 1, yielding a trivial weight factorization. As illustrated in the next section, this is fundamentally different from approach the scale factors 8 by a random distribution and re-parameterize the weight matrix accordingly. In Section we demonstrate that, by initializing using an appropriate distribution, can consistently outperform both weight normalization, SIREN, and adaptive activations across a broad range of supervised and selevvuper learning tasks.
initialize
RANDOM WEIGHT FACTORIZATION IN PRACTICE
Here we illustrate the use of weight factorization through the lens of a simple regression task. Specifically, we consider a smooth scalar-valued function I sampled from a Gaussian random field using a square exponential kernel with length scale of 0.02. This generates data-set of N = 256 observation pairs {I,, ( where { lie on a uniform grid in [0, 1] The goal to train network fo learn f by minimizing the mean square error loss L())- = 1/N Σ (J((x,) - [x:))
The proposed random weight factorization is applied as follows. We first initialize the parameters of an ML.P network via the Glorot scheme (Glorot & Bengio 2010). Then, for every weight matrix W, we proceed by initializing a scale vector exp(s) where s sampled from multivariate normal distribution N ((,, ol). Then every weight matrix is factorized by the associated scale factor as W = diag(exp(s) s)) V at initialization. Finally, We train this network by gradient descent the new parameters s, V directly. This procedure is summarized in Appendix B along with simple JAX Flax implementation (Heek et al., 2020) in Appendix [ The use of exponential parameterization is motivated by Weight Normalization (Salimans & Kingmal 2016) to strictly avoid zeros very small values in the scale factors and allow them to span a wide range of different magnitudes. ablation study in Appendix D further validates its benefit.
In Figure 2 we train networks (3 layers, 128 neurons per layer, ReLU activations) to learn the target function using: (a) conventional MLP, (b) an MLP with adaptive activations (AA) (Jagtap et al. 2020). (c) an MLP with weight normalization (WN) (Salimans & Kingma, 2016). and (d) an MLP with the proposed random weight factorization scheme (RWF). Evidently, RWF yields the best pre- dictive accuracy and loss convergence. Moreover, we plot the relative change of the weights in the original (unfactorized) parameter space during training in the bottom middle panel. We observe that
4

---

Under review as a conference paper at ICLR 2023
Plain
AA
WN
RWF
mun
05
10
20
14
00
10
60
05
100
107
107
MSE
3/K3
10'
50
100)
160
x1000
Fletn A
WN
DWP
Figure 2: ID regression: Top: Model predictions using different parameterizations. Plain: Stan- dard MLP; AA: adaptive activation; WN: weight normalization; RWF: random weight factorization. Bottom left: Mean square eтrоr (MSE) during training. Bottom Middle: Relative change of weights during training. The comparison is performed in the original parameter space. Bottom right: Eigen- values (descending order) of the empirical NTK at the end of training.
RWF leads to the largest weight change during training, thereby enabling the network to find better local minima further away from its initialization. To further emphasize the benefit of weight factor- ization, compute the eigenvalues of the resulting empirical Neural Tangent Kernel (NTK) (Jacot
et al. 2018) Ko 8ie 80 (I:), 8fe 80 (x²)
at the last step of training and visualize them in the bot-
tom right panel. As demonstrated in Tancik al. (2020) and Wang et al. (021c), the eigenvalues of the Neural Tangent Kernel (NTK) characterize the convergence rate of the corresponding kernel eigenvectors. Specifically, components of the target function that aligned with kernel eigenvec- tors corresponding to larger eigenvalues will be learned faster. Therefore, the rapid decay of the NTK eigenvalues implies extremely slow convergence function, pathology commonly referred to "spectral bias". Notice how RWF exhibits a flatter NTK spectrum and slower eigenvalue decay than the other methods, indicating training dynamics and less severe spectral bias. To explore the robustness of the proposed RWF, we conduct a systematic study on the effect of µı and a in the initialization of the scale factor s. The re- sults suggest that the choice of μ,, plays an important role. Specifically, too small μl, values may lead to performance that is similar to a conventional MLP, while setting μt, a too large can result in unstable training process. We empirically find that 1, 0.1 consistently improves the loss convergence and model accuracy for the vast majority of tasks considered in this work. Additional details are presented in Appendix D
high frequency components of the target
contitoond
EXPERIMENTS
In this section, demonstrate the effectiveness and robustness of random weight factorization for training continuous neural representations across a range of tasks in computer vision, graphics, and scientific computing. More precisely, we compare the performance of plain MLPs, MLPs with adap- tive activations (AA) (Jagtap et al. (020)., weight normalization (WN) (Salimans & Kingma 2016), and the proposed random weight factorization (RWF). The comparison performed over a collec- tion of MLP architectures, including conventional MLPs, SIREN (Sitzmann et al., 2020). modified MLPs (Wang al., 20211). as well MLPs with positional encodings (Mildenhall et al. 2020) and Gaussian Fourier features (Tancik et al., 2020). In addition, for the Image Regression and Com- puted Tomography benchmarks, we perform a comprehensive
sweep over different well MLPs with
random seeds, lcarning rates, activation functions, scale our experiments are presented in Appendix Table summarizes the results obtained for each benchmark, corresponding the optimal input setup. Overall, the proposed random weight factorization yields the
different depths and widths. The aeper-parameters
mapping and best performance, significantly outperforming the other competing approaches.
is also worth noting that the proposed moderately increases the model sizes by O(L)), where denotes the number of layers and D the width of the network. Besides, the computational overhead
5

---

Under review as a conference paper at ICLR 2023
Task
Metric
Case
Plain
WN
RWF (ours)
Natural
27.35
27.37
27.36
28.08
Image Regression
PSNR (1)
Text
32,09
32.29
32.25
33.13
Dragon
0.980
0.981
0.981
0.984
Shape Representation
IoU (1)
Armadillo
0.978
0.976
0.974
0.982
Shepp
30.09
30.37
30.59
33.73
Computed Tomography
PSNR (1)
ATLAS
21.85
21.93
22.16
23.71
Inverse Rendering
PSNR (1)
Lego
26.11
26.11
26.07
26.21
Advection
25.94%
36.27%
67.40%
3.10%
Solving PDEs
Rel. 1,2 (4)
Navier-Stokes
11.93%
10.95%
14.62%
6.79 %
DR
1.09%
0.95%
0.97%
0.50%
Learning Operators
Rel. 1.3 (4)
Darcy
2.03%
2.06%
2.00%
1.67 %
Burgers 5.11%
4.71%
4.37%
2.46 %
Table I: We compare four different parameterizations over various benchmarks, and demonstrate that Random weight factorization consistently outperforms other parameterizations across all tasks. All comparisons are conducted under exactly the same syper-paramette settings. (t)(L) indicates that higher/lower values are better, respectively. (Plain: conventional MLP; AA: adatpive activation; WN: weight normalization; RWF: random weight factorization).
of our method is marginal, which is further clarified by Table where we report the computational cost of each benchmark. Therefore, RWF can be therefore considered as a drop-in enhancement to any architecture that uses linear layers. All code and data will be made publicly available. A summary of each benchmark study presented below, with more details provided in Appendix.
4.1
2D IMAGE REGRESSION
We train coordinate-based MLPs to learn map from 2D input pixel coordinates the corresponding RGB values of an image, using the benchmarks put forth in (Tancik et al., (0220). We conduct experiments using two data-sets: Natural and Text, each containing 16 images. The Natural data- set is constructed by taking center crops of randomly sampled images from the Div2K data-set (Agustsson & Timofte, 2017). The Text data-set is constructed by placing random text with random font sizes and colors on a white background. The training data is obtained by downsampling each test image by factor of 2. We compare the resulting peak signal-to-noise ratio (PSNR) in the full resolution test images, obtained with different MLP architectures using different input mappings and weight parametrizations.
4.2
3D SHAPE REPRESENTATION
This task follows the original problem setup in (Tancik et al., 22020. The goal is learn an implicit representation of 3D shape using Occupancy networks (Mescheder et al. 2019). which take spatial coordinates as inputs and predict for points outside given shape and for points inside the shape. We use two complex triangle meshes commonly used in computer graphics: Dragon and Armadillo. The training data generated by randomly sampling points inside a bounding box and calculating their labels using the ground truth mesh. We evaluate the trained model performance using the Intersection over Union (IoU) metric on a set of points randomly sampled near the mesh surface to better highlight the different mappings' abilities to resolve fine details.
4.3 2D COMPUTED TOMOGRAPHY (CT)
This task follows the original problem setup in (Tancik et al., 2020). We train an MLP to learn map from 2D pixel coordinates to a corresponding volume density at those locations. Two data- sets considered: procedurally-gennrraee Shepp-Logan phantoms (Shepp & Logan 1974) and 2D brain images from the ATLAS data-set (Liew et al. 2018). Different from the previous tasks, observe integral projections of density field instead of direct measurements. The network is trained an indirect supervised fashion by minimizing a loss between a sparse set of ground-truth integral projections and integral projections computed from the network's output. We use PSNR to

---

Under review as a conference paper at ICLR 2023
GO
Plan
AA
-00
0.5
100
0.0
0.5
1.0
100
3.5
WIN
BWI
E
BLD
000
Flate
1171
R0FF
Figure 3: Adection: Predicted solutions of trained MLPs with different weight parameterizations, along with the evolution of the associated relative L prediction errors during training.
quantify the performance of the trained MLPs with different input mappings and different weight parametrizations.
4.4 3D INVERSE RENDERING FOR VIEW SYNTHESIS
This task follows the original problem setup in (Tancik et al. 2020). We aim learn implicit representation of 3D scene from 2D photographs using Neural Radiance Field (NeRF) (Mildenhall et al. 2020), which is coordinate-basel MLP that takes 3D location input and outputs color and volume density. The network is trained by minimizing a rendering loss between the set of 2D image observations and the same rendered views from the predicted scene representation. In our experiments, we consider a down-sampled NeRF Lego data-set and use a simplified version of the method described in (Mildenhall et al. 2020), where remove hierarchical sampling and view dependence. We compare the PSNR of the trained MLPs with different weight parametrizations over 10 random seeds.
4.5 SOLVING PARTIAL DIFFERENTIAL EQUATIONS (PDEs)
Our goal is to solve partial differential equations (PDEs) using physics-informe neural network (PINNs) (Raissi et al. 2019), which take the coordinates of spatio-temporal domain inputs and predict the corresponding target solution function. PINNs are trained a self-supervised fashion by minimizing composite loss function for fitting given initial and boundary conditions, well as satisfying the underlying PDE constraints. We consider two benchmarks, an advection equation modeling the transport of scalar field, and the Navier-Stokes equation modeling the motion of an incompressible fluid in square cavity. We compare the resulting relative L errors of each param- eterizations over 10 random seeds. Detailed descriptions and implementations of each problem are provided below and in the Appendix K
Figure 3 and Figure 4 present the ground truth against the predicted solutions obtained by training PINNs with different weight parameterizations. It can be observed that the predictions obtained by RWF are in excellent agreement with the ground truth, while the other three parameterizations result in poor or even non-physical approximations, The rapid decrease in the test error further validates the benefit of our method. We attribute these significant performance improvements to the fact that PINN models, due to their self-supervised nature, often suffer from poor initializations. Evidently, RWF can precisely mitigate this by being able to reach better local minima that are located further away from the model initialization neighborhood.

---

Under review as a conference paper at ICLR 2023
GT
Pain
AA
10
10
1.0
1.
1.0
Ga
04
0.0
62
02
0.2
02
022
68
0.0
66g
RWF
1.3
1.0
140
1.0
655
0.5
0.5
0.5
A/3
2
022
600
100
ILS
1.0
Toer (00
Figure 4: Navier-Stokes: Predicted solutions of trained MLPs with different weight parameteriza- tions, along with the evolution of the associated relative 2 errors during training.
Advection equation: The first example is ID advection equation, a linear hyperbolic equation commonly used to model transport phenomena
Du Ou
0, u(x 0) (z)) x (0, 2
+ Or
E 2T). [0, 1],
(4.1)
Ōt
(4.2)
with periodic boundary conditions. This example has been studied in (Krishnapriyan et al. 2021 Daw al.) 2022), exposing some of the limitations that PINNs suffer from as the transport velocity c is increased. In our experiments, we consider = 50 and an initial condition g(z) = sin(z).
Navier-Stokes equation: The second example a classical benchmark problem in computational fluid dynamics, describing the motion of an incompressible fluid in ivo-dimensional lid-driven cavity. The system governed by the Navier-Stokes equations written in non-dimensional form
1 Au 0, (x, v) E 1)?,
u Vu Vp
(4.3)
Re V u 0, (z, v) E 1),
(4.4)
where u = (u, u) denotes the velocity in x and y directions, respectively, and is the scalar pressure field. assume u (1, on the top lid the cavity, and a non-slip boundary condition on the other three walls. All experiments are performed with Reynolds number of Re 1, 000.
4.6 LEARNING OPERATORS
In this task, focus learning the solution operators of parametric PDEs. To describe the problem setup in general, consider parametric PDE of the following form
M(a, u) = O.
(4.5)
where N A - is linear nonlinear differential operator between
function spaces. Moreover, a E A denotes the PDE parameters, and u E U is the corresponding unknown solutions of the PDE system. The solution operator G A 24 is given by
G(a) u(a).
(4.6)
In our experiments, we consider three benchmarks: Diffusion-reaation. Darcy flow and the Burgers' equation. Detailed descriptions of each problem setup are shown below. As shown in Figure 5 we plot the training losses of each model with different weight parameterizations. One can see that random weight factorization yields the best loss convergence for every example, indicating the capability the proposed method to accelerate the convergence of stochastic gradient descent and achieve better local minima.
8

---

Under review as a conference paper at ICLR 2023
DR
Darcy
Burgges
10
10r-3
g
Loo
10-1
10
10--
10-1
-0)
40
Ster x1000)
Ites x 1000
Iter x1000
WN
Figure 5: Learning operators: Loss convergence of training DeepONets with different weight pa- rameterizations for diffusion-reaction equation, Darcy flow and the Burgers' equation.
Diflusion-reaatiin: Our first example involves nonlinear diffusion-reeacionn PDE with source term a (0, + E,
Du = D pu ku² a(x), tx, E (0. x (0,
(4.7)
Ot Oz
with the zero initial and boundary conditions, where D = 0.01 is the diffusion coefficient and k 0.01 is the reaction rate. We train Deep Operator Network (DeepONet) Lu et al. (202I) to learn the solution operator for mapping source terms a(x) to the corresponding PDE solutions u(x). This network takes a PDE parameter and a spatial-temporal coordinate as inputs, and predicts the associated PDE solution evaluated at that location. The model is trained in supervised manner by minimizing a loss between the predicted PDE solutions and the available solution measurements.
The Darcy equation describes steady-state flow through porous medium, taking the following form in two spatial dimensions
Darcy flow:
(a Vu) = f. (I, E (0.
(4.8)
0, 0, (z, E 00,, 1)2
(4.9)
where a (0, R is the diffusion coefficient and f (0, 1) + R is forcing term. We fix [(x, v) and aim to learn a continuous representation of the solution operator G a(x, y)
u(x, v) with DeepONet.
Burgers' equation: As the last example, we consider fundamental nonlinear PDE, the one- dimensional viscous Burgers' equation. It takes the form
Ou + u Or
Ou
82u
= v 2x
(x, E (0, x (0, 1].
(4.10)
Ot u(z,0) 0) = 0, x e (0,1),
(4.11)
(4.12)
with periodic boundary conditions and // = 0.001. Our goal is to learn the solution operator from the initial condition the associated PDE solution with physics-infommen DeepONet (Wang al. 2021ep. Different from the first two examples, physics-ifformed DeepONet trained in a self. supervised manner, i.e. without any paired input-output observations, except for a set of given initial or boundary conditions (see Wang et al. (2021e) for more details).
CONCLUSIONS
In this work, proposed random weight factorization, simple and remarkably effective parameterization the weight matrices in neural networks. Theoretically, we show how this fac- torization alters the geometry of loss landscape by assigning a self-adaptive learning rate each neuron. Empirically, we show that our method can mitigate spectral bias in MLPs and enable net- works search for good local optima further away from their initialization. We validate random weight factorization using six different benchmarks ranging from image regression to learning op- crators, showcasing a consistent and robust improvements across various tasks computer vision, graphics, and scientific computing. These findings provide new insights into the training of continu- ous neural representations and open several exciting avenues for future work, including the applica- tion of our method to deep learning models beyond coordinate-baut based MLPs. such convolutional networks (LeCun et al.) 1998), graph networks (Scarselli et al. 2008). and Transformers (Vaswani al. 2017).

---

Under review as a conference paper at ICLR 2023
REFERENCES
Eirikur Agustsson and Radu Timofte. Ntire 2017 challenge on single image super-resolution: Dataset and study. In Proceedings of the IEEE conference on computer vision and pattern recog- nition workshops, pp. 126-135, 2017.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv:1607.064550 2016.
Alberto Bietti and Julien Mairal. On the inductive bias of neural tangent kernels. Advances in Neural Information Processing Systems, 32, 2019.
Yuan Cao, Zhiying Fang, Yue Wu, Ding- Xuan Zhou, and Quanquan Gu. Towards understanding the spectral bias of deep learning. arXiv preprint arXiv: 1912.01 0180,1198,, 2019.
Zhiqin Chen and Hao Zhang. Learning implicit fields for generative shape modeling. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pp. 5939-5948, 2019.
Steven M Cox and Paul C Matthews. Exponential time differencing for stiff systems. Journal of Computational Physics, 176(2):430-455 2002.
Arka Daw, Jie Bu, Sifan Wang, Paris Perdikaris, and Anuj Karpatne. Rethinking the importance of sampling in Physics-informed neural networks. arXiv preprint vrXiv:2207.023300 2022.
Suchuan Dong and Naxian Ni. A method for representing periodic functions and enforeing exactly periodic boundary conditions with deep neural networks. Journal of Computational Physics, 435: 110242, 2021.
Tobin A Driscoll, Nicholas Hale, and Lloyd N Trefethen. Chebfun guide, 2014.
Kyle Genova, Forrester Cole, Daniel Vlasic, Aaron Sarna, William Freeman, and Thomas Funkhouser. Learning shape templates with structured implicit functions. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 7154-7164, 2019.
Kyle Genova, Forrester Cole, Avneesh Sud, Aaron Sarna, and Thomas Funkhouser. Local deep implicit functions for 3d shape. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 4857-4866, 2020.
Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pp. 249-256, 2010.
Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andreas Steiner, and Mare van Zee. Flax: A neural network library and ecosystem for JAX, 2020. URL
http://github. wtt.ccom/ooooo cot./oogle//llaa......m
Philipp Henzler, Niloy Mitra, and Tobias Ritschel. Learning a neural 3d texture space from 2d exemplars. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recog- nition, pp. 8356-8364, 2020.
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International conference on machine learning, pp. 448-456. PMLR, 2015.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kemel: Convergence and gen- eralization in neural networks. Advances in neural information processing systems, pp. 8571- 8580, 2018.
Ameya D Jagtap, Kenji Kawaguchi, and George Em Karniadakis. Adaptive activation functions accelerate convergence in deep and physies-informed neural networks. Journal of Computational Physics, 404: 109136, 2020.
George Em Karniadakis, Ioannis Kevrekidis, Lu Lu, Paris Perdikaris, Sifan Wang, and Liu Yang. Physics-informed machine learning. Nature Reviews Physics, pp. 1-19, 2021.
10

---

Under review as a conference paper at ICLR 2023
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
Aditi S Krishnapriyan, Amir Gholami, Shandian Zhe, Robert M Kirby, and Michael W Mahoney. Characterizing possible failure modes in physics-infomma informed neural networks. arXiv preprint arXiv:2109.0105g 2021.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11):228-22224, 1998.
Sook-Lei Liew, Julia M Anglin, Nick W Banks, Matt Sondag, Kaori L. Ito, Hosung Kim, Jennifer Chan, Joyce Ito, Connie Jung, Nima Khoshab, et al. A large, open source dataset of stroke anatomical brain images and manual lesion segmentations. Scientific data, 5(1):1-11,
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 229, 2021.
Lars Mescheder, Michael Oechsle, Michael Niemeyer, Sebastian Nowozin, and Andreas Geiger. Oc- cupancy networks: Learning 3d reconstruction in function space. In Proceedings ofe IÉEE/CVF conference on computer vision and pattern recognition, pp.
2019.
Ben Mildenhall, Pratul P Srinivasan, Matthew Tancik, Jonathan T Barron, Ravi Ramamoorthi, and Ren Ng. Nerf: Representing scenes as neural radiance fields for view synthesis. In European conference on computer vision, pp. 405-421. Springer, 2020.
Anh Nguyen, Jason Yosinski, and Jeff Clune. Deep neural networks are easily fooled: High confi- dence predictions for unrecognizable images. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 427 436, 2015.
Michael Niemeyer, Lars Mescheder, Michael Oechsle, and Andreas Geiger. Differentiable volumet- ric rendering: Learning implicit 3d representations without 3d supervision. In Proceedings of the JEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 3504-3515, 2020.
Michael Oechsle, Lars Mescheder, Michael Niemeyer, Thilo Strauss, and Andreas Geiger. Tex- ture fields: Learning texture representations function space. In Proceedings of the IÉEE/CVF International Conference Computer Vision, pp. 4531-4540, 2019.
Jeong Joon Park, Peter Florence, Julian Straub, Richard Newcombe, and Steven Lovegrove. Deepsdf: Learning continuous signed distance functions for shape representation. In Proceed- ings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 165-174, 2019.
Nasim Rahaman, Aristide Baratin, Devansh Arpit, Felix Draxler, Min Lin, Fred Hamprecht, Yoshua Bengio, and Aaron Courville. On the spectral bias of neural networks. In International Confer- on Machine Learning, pp. 2019.
Maziar Raissi, Paris Perdikaris, and George Karniadakis. Physics-informel neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational Physics, 378:686-707, 2019.
Maziar Raissi, Alireza Yazdani, and George Em Karniadakis. Hidden fluid mechanics: Learning velocity and pressure fields from flow visualizations. Science, 367(6481):1026-1030 2020.
Shunsuke Saito, Zeng Huang, Ryota Natsume, Shigeo Morishima, Angjoo Kanazawa, and Hao Li. Pifu: Pixel-aligned implicit function for high-resolution clothed human digitization. In Proceed- ings of the IEEE/CVF International Conference on Computer Vision, pp.
2019.
Tim Salimans and Durk P Kingma. Weight normalization: A simple reparameterization to accelerate training of deep neural networks. Advances in neural information processing systems, 29,2012 2016.
Franco Scarselli, Marco Gori, Ah Chung Tsoi, Markus Hagenbuchner, and Gabriele Monfardini. The graph neural network model. IEFE transactions neural networks, 20(1):61-80, 2008

---

Under review as a conference paper at ICLR 2023
Lawrence A Shepp and Benjamin F Logan. The fourier reconstruction of head section. 1EEE Transactions on nuclear science, 21(3):21-43, 1974.
Vincent Sitzmann, Michael Zollhöfer, and Gordon Wetzstein, Scene representation networks: Con- tinuous 3d-structure-aware neural scene representations. Advances in Neural Information Pro- cessing Systems, 32, 2019.
Vincent Sitzmann, Julien Martel, Alexander Bergman, David Lindell, and Gordon Wetzstein. Im- plicit neural representations with periodic activation functions. Advances in Neural Information Processing Systems, 33:7462-7473, 2020.
Kenneth O Stanley. Compositional pattern producing networks: A novel abstraction of development. Genetic programming and evolvable machines, 8(2):111-12,, 2007.
Matthew Tancik, Pratul P Srinivasan, Ben Mildenhall, Sara Fridovich- Keil, Nithin Raghavan, Utkarsh Singhal, Ravi Ramamoorthi, Jonathan Barron, and Ren Ng. Fourier features let networks learn high frequency functions in low dimensional domains. arXiv:2006. 10739, 2020.
arXiv preprint
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 informa- tion processing systems, 30, 2017.
Hanwen Wang, Isabelle Crawford-Eng, and Paris Perdikaris. Enhancing the trainability and ex- pressivity of deep MIPs with globally orthogonal initialization. In The Symbiosis of Deep Learning and Differential Equations, 2021a. URL id KKMKGG2GTTMM
wwwarreiiewnee//ooo net/forum?
Sifan Wang, Yujun Teng, and Paris Perdikaris. Understanding and mitigating gradient flow patholo- gies in physics-informed neural networks. SIAM Journal on Scientific Computing. A3355: A3081, 2021b.
Sifan Wang. Hanwen Wang, and Paris Perdikaris. On the eigenvector bias of fourier feature net- works: From regression to solving multi-scale PDEs with physics-informed neural networks. Computer Methods Applied Mechanics and Engineering, 384: 113938, 2021c
Sifan Wang, Hanwen Wang, and Paris Perdikaris. Improved architectures and training algorithms for deep operator networks. arXiv preprint arXiv:2110.01654. 2021d.
Sifan Wang, Hanwen Wang, and Paris Perdikaris. Learning the solution operator parametric par- tial differential equations with physics-informed deeponets. Science advances, 7(40):eabi8605, 2021e.
Sifan Wang, Shyam Sankaran, and Paris Perdikaris. Respecting causality is all you need for training physics-informed neural networks. arXiv preprint 22Xi:22033.07044, 2022.
Zhi-Qua Zhi-Qin John Xu, Yaoyu Zhang, Tao Luo, Yanyang Xiao, and Zheng Ma. Frequency principle: Fourier analysis sheds light on deep neural networks. arXiv preprint arXiv: (90110652,2 2019.
12

---

Under review as a conference paper at ICLR 2023
A PROOFS
A.1 PROOF OF THEOREM 1
Proof. Note that for any pair (s!)) v(t) with diag(s()) v()) WW have
| - 0, as 8(1) 00
(A.1)
for = 1,2..... L+1 1. Then for any > 0, there exists M > 0 such that for any (V(),
800)
E
Uo ,nd || > M, we obtain || for 1 1,.,....7 + 1. We define U. by
U. (0,
8 | 8 RR,, = 1,2,..... ................................ L 1,2.1..1++1)
(A.2)
Now can choose 8 such that || M. Then
L+1
dist(Uo,U.)
[1012 +l
(A.3)
Similarly, we can show that
L+1
UA((UO,,U.)
Σ
V 12 S Ic
(A.4)
Therefore,
dist(Uo, Up) +dist(UU, U.) + dist(U,, Upr) S 22L le
(A.5)
Since c is arbitrary, letting - 0 gives
dist(Uo: Upr) = O.
(A.6)
A.2 PROOF OF THBOREM 2
Proof. Suppose that ((e.)) denotes k-th component of f()) E R4. Under the proposed weight
factorization in equation 2.3 differentiating the loss function with respect to wh. and g(h.1), respectively, yields
Oc
dc
n+i
(s.1)
(b.l)
11
dsn
(k.J)
(k.l)
1) (f(k.)
a (k.)) g "1
(A.7)
7L
ss (k.) OC
(k.2)
(k.
(k.l)
g -1)
(A.8)
n+1
(k.1) (k.1)
Of(k.t)
Note that
Oc
OC
11-1)
(A.9)
ow (k.l)
ffik.l)
and the update rule of (k,!) and (k.l)) can be re- re-written as
0c
s sh (To(k.))
(k.)) (k.))
(A.10)
Dw (k.))
v unt1 (k.1)
(k.))
17
(R.))
ac
(A.11)
Ow (k.1)*
Since w (k.1) = (k..)) (k.1), the update rule of w (k.1) is given by
w n+1
(k.I) (k.1) 17 (1s/ + 110 11.3
0c
+ O
(A.12)
dw (k.I)
13

---

Under review as a conference paper at ICLR 2023
B ALGORITHM OF RANDOM WEIGHT FACTORIZATION
Algorithm Random weight factorization (RWF)
1. Initialize neural network fo with 0 (W()) bOO) L+1 (e.g. using the Glorot scheme Gloort & Bengio 2010)).
for | I, 1.2....,4 L do
(a) Initialize each scale factor as g ~ (N,
(b) Construct the factorized weight matrices as W = diag (ex((s())) V
end for
2. Train the network by gradient descent on the factorized parameters {s!1) V (0)00 L+1
The recommended hyper-paramertt are μ 1.0, a 0.1
C A DROP-IN ENHANCEMENT FOR LINEAR LAYERS
alass Dense (nn, Module)
features: int
compatt
def __.calll11 (self, x)
kernel
(' kernel', gloret_normal (),
(x.shapc[-1] self.fcatures))
10
bias self. param ('bias', nn. m.innttializers.zeeee (self. features,)
11
y dot (x, kernel) bias return y
14
Listing 1: JAX Flax implementation (Heek al. 2020 of conventional linear layer.
def facterized glorot normal (mean=1.0, 0, stddev=0.1):
def init (key, shape) :
keyl, key2 random. random.spitt (key)
w glorot_normal () (keyl, shape)
3 = mean + normal (stddev) (key2, (-1),,
np.exp (s)
return 5, return init
10
class FactorizedDense (onoule))))
12
features: int
13
14
enn.compace
def self.param kernel',
x):
5,
17
factorized_
wwt.gg_ooooaaallll_oo normal (),
18
(x.shape(-1), self. features)
19
kernel =
bias self. param bias',
21
nn. it..iitializers.errn (self. features,
22
23
mn.ndddd dot (x, kernel) bias return
Listing 2: JAX Flax implementation Heek et al. 2020 of a linear layer with random weight factorization.
14

---

Under review as a conference paper at ICLR 2023
D ID REGRESSION ABLATION STUDY
We perform a systematic study on the effect of and a used for initializing the scale factor in random weight factorization. To this end, train MLPs with random weight factorization initialized by different and a. Each model (3 layers, 128 channels, ReLU activations) is trained via a full-batch gradient descent for 10 iterations using the Adam optimizer (Kingma Ba 2014) with a starting learning rate of 10 followed by an exponential decay by a factor of 0.9 in every 000 steps. The resulting relative L errors over 10 random seeds are visualized in Figure ( We observe that all the curves are roughly basin-shaped, indicating "sweet spot" for the random weight factorization, i.e., y E [1, 2]. Thus we will mainly use these as the default hyper- parameter of random weight factorization in the majority of the benchmarks presented here (see Table
Furthermore, we perform an ablation study on the different distributions (Normal, LogNor- mal,Uniform, LogUniform) initializing the scale factors and report the average and standard de- viation of mean square errors over 10 random seeds. For fair comparison, we manually tune each distribution such that the sampled values are roughly in the same range. We see that LogNormal is slightly better than the other distributions while overall there is no significant discrepancies in terms of the accuracy for using different distributions.
10
1.0
0.1
10-1
0.01
0.001
MSB
10
10-3
10-4
2
4
H
Figure 6: ID Regression: MSE of training MLPs with random weight factorization initialized by different y and a.
Distribution
MSE
Normal
3.95e 04 2.22e 04
LogNormal
3.76e 04 = 3.3le 04
Uniform
3.94e 04 = 4.89e 04
Uniform
5.40e 04 4.78e 04
Table 2: 1D Regression: MSE training MLPs with random weight factorization initialized by different distributions.
15

---

Under review as a conference paper at ICLR 2023
E HYPER-PARAMETERS
The following tables summarizes the hyper-paramettrs of the different networks architectures em- ployed each benchmark (Table 3). their associated learning rate schedules (Table 4 and the corresponding random weight factorization settings (Table
Task
Case
Backbone
Depth
width
Activation
Natural
Image Regression
MLP
4
256
ReLU
Text
Dragon
Shape Representation
MLP
256
ReLU
Armadillo
Shepp
Computed Tomography
MLP
5
256
ReLU
ATLAS
Inverse Rendering
Lego
MLP
5
256
ReLU
Advection
MLP
5
256
Solving PDEs
MLP
Tanh
Navier-Stokes
128
Modified MLP
DR
5
64
Tanh
DeepONet
Learning Operators
Darcy
4
128
GELU
Burgers
Modified DeepONet
128
Tanh
Table 3: Network architectures for each benchmark Table
Learning Rate Schedule
Task
Case
Iterations
Step Size
Decay Steps
Decay Rate
Warmup Steps
Natural
Image Regression
10-3
x 10a
2 10
Text
Dragon
Shape Representation
5 10
5 x 10a
0.1
103
104
Armudillo
Shepp
Computed Tomography
10-3
102
2 x 10
ATLAS
Inverse Rendering
Lego
10
10a
0.9
x 10a
5 10*
Advection
5 10a
0.9
2 10
Solving PDEs
10-3
Navier-Stokes
2 x 10a
0.9
1.2 10
DR
10a
0.9
5 10a
Learning Operators
Darcy
10-3
Burgers
2 10a
10a
Table 4: Learning Rate Schedules for each benchmark in Table
1
16

---

Under review as a conference paper at ICLR 2023
Task
Case
Initialization of RWF
Natural
(2, 0.01)
Image Regression
Text
8 (1.0..)
Dragon
Shape Representation
8 ~~~ (11,
Armadillo
Shepp
Computed Tomography
8 8~ (~111,)
ATLAS
Inverse Rendering
Lego
NV11
Advection
8 8~1((1,
Solving PDEs
Navier-Stokes
8~ NO.5,
DR
Learning Operators
Darcy
s~1,01)
Burgers
Table 5: Distribution used for initializing the scale factor in random weight factorization for each benchmark.
Hyper- saper-parameters
Values
Random Seed
3, 5, 7, 11
Learning Rate
le-3, 1e-4
Activation
ReLU, GELU, sin
Gaussian scales
5, 10
Width
64, 128, 256
Depth
3, 4, 5
Table 6: Range of hyper-parameters for hyper-parameter sweeps used in Image Regression and Computed Tomography benchmarks.
F COMPUTATIONAL COSTS
Table presents the computational cost in terms of training iterations per second for the networks employed in each benchmark. All timings are reported on single NVIDIA RTX A6000 GPU.
Task
Case
Plain
AA
WN
RWF
Natural
116.36
111.33
114.09
113,80
Image Regression
Text
1 116.86
112.04
113.77
113.99
Shape Representation
Dragon
12.79
12.79
12.76
12.73
Armadillo
13.14
13,06
12.81
12.94
Shepp
30.02
29.21
29.83
29.88
Computed Tomography
ATLAS
29.47
28.74
29.36
29.32
Inverse Rendering
Lego
30.59
29.41
30.68
30.75
Solving PDEs
Advection
853.82
757.19
855.50
789,41
Navier-Stokes
164.02
152.70
160.51
160.45
DR
469.09
450.44
470.34
469.00
Learning Operators
Darcy
63.10
61.21
62.92
62.91
Burgers
27.86
26.25
27.29
27.10
Table 7: Computational cost (training iterations per second) for each benchmark. the computational overhead of random weight factorization is marginal.
We can that
17

---

Under review as a conference paper at ICLR 2023
G 2D IMAGE REGRESSION
two image data-sets: Natural and Natural. All the test images have a 512 x 512 resolution while the training data has a 256 x 256 resolution. For each data-set, train models with different weight parametrizations (Plain, Adaptive Activation, Weight Nor- malization, Random Weight Factorization) over different random seeds, learning rates, activation functions, scale pyper-parameters, well MLP architectures with different depths and widths. The range of hyper-parameters is summarized in Table [ For random weight factorization, we initialize the scale factor s (2, 0.01) and s 0.01) for the Natural and Text data-set, respectively. Each model is trained via a full-batch gradient descent for 2,000 iterations using the Adam optimizer (Kingma & Ba 2014) 4) with default settings and 200 warm-up steps. We compute mean and standard deviations of the resulting PSNRs and visualize our results in 7 Figure 8 respectively. We reach the conclusion that the random weight factorization (RWF) consistently yields the best optimal performance than the alternatives.
As mentioned in Section 4.1
Furthermore, under the best different parameterizations (see 2.3) and the following input embeddings:
setup, we compare the performance of MLPs with
No mapping: MLP with no input feature mapping.
Positional encoding (Tancik et al., 2020): (xx) [.... cos (2mоr//m))) x)) sin (2irad/mm2) x) .... for j = 0....... m 1. where the frequencies are log-linear spaced and the scale o > is a user-specified
Gaussian (Tancik et al., 2020): >(x) = [cos(2n (z) where E Rxdd is sampled from Gaussian distribution NO, 22). The scale 0 is user-specified hyper-parameter.
The resulting test PSNR reported in Table and Table We can see that the weight factorization with Gaussian Fourier features achieves the best PSNR among all the cases. Some visualizations shown in Figure 9 and Figure TO
Natural data-set
Plain
AA
WN
RWF (ours)
No mapping
17.94 + 2.38
18.29 + 2.44
18.28 + 2.44
19.11 + 2.50
Positional Encoding
27.04 + 3.89
26.73 3.67
26.99 = 3.83
27.46 + 3.82
Gaussian
27.35 + 4.05
2.96
27.36 4.03
28.08 4.34
Table 8: 2D Image Regression: Mean and standard deviation of PSNR obtained by training MLPs with different input mappings and weight parameterizations for the Natural data-set.
Text data-set
Plain
AA
WN
RWF (ours)
No mapping
18.42 + 2.42
18.43 = 2.34
18.46 + 2.31
17.85 + 2.36
Positional encoding
31.33 2.71
31.73 + 2.29
31.43 2.57
30.49 + 2.45
Gaussian
32.09 + 1.80
32.29 + 1.99
32.25 1.74
33.13 = 2.03
Table 9: 2D Image Regression: Mean and standard deviation of PSNR obtained by training MLPs with different input mappings and weight parameterizations for the Text data-set.
18

---

Under review as a conference paper at ICLR 2023
25
20
14
Leaming rate
PSNR
10
1c403
None
A.A
W/N
RWF
5
10
20
PSNR
10
ein
Gaussian scale
Activation
gelu
relu
15
20
5
20
PSNR
PSNR
Networkc depth
Network wilth
128
3-
64
15
20
PSNR
PSNR
Figure 7: 2D Image Regression: Mean and standard deviation PSNR obtained by training MLPs with different hyper-parameter setup for the Natural data-set. In particular, top left panel plots settings tested, RWF yields
min/max bar, indicating that over the entire range of the best optimal performance compared to the alternatives.
19

---

Under review as a conference paper at ICLR 2023
1c04
Learnnng rate
PSNR
201
1o-03
0
News
AA
WN
RWF
10
20
30
PSNR
10-
sim
Gameelan soale
Activation
gelu
relu
10
PSNR
PSNR
Networks depth
Network width
7
128
3
64
0
10
20
30
10
20
30
PSNR
PSNR
Figure 8: 2D Image Regression: Mean and standard deviation PSNR obtained by training MLPs with different hyper-paramett setup for the Text data-set. In particular, top left panel plots min/- max bar, indicating that over the entire range of Syper-parameter settings tested, RWF yields the best optimal performance compared to the alternatives.
20

---

Under review as a conference paper at ICLR 2023
Gr
Plan
Wi
WWP
GT
Plan
AA
WN
BWF
Figure 9: 2D Image Regression: Predicted Natural images of trained MLPs with Gaussian Fourier features and with different weight parameterizations.
GT
Plain
AA
WN
RVF
O b &SHDBY
obs &HHDBS
Obs &HDDY
Obs &
ASHDBY
obss S
WolkzuuZ
Wutz
Wouztuu
WizzzuZ
www.
PTeco WoRR
PToce
PToceWHoRx xbA
PToceWHRRx xbA
PTeco xbA
xbA
xbA
IRITZ
gq cc
gqHdCcz
gqHdCcz
gqHdCcz
gqHdCcz
Figure 10: 2D Image Regression: Predicted Text images of trained MLPs with Gaussian Fourier features and with different weight parameterizations.
Comparison with STREN (Sitzmann al. 2020): We find that SIREN also factorizes the weight matrix of every hidden layer as W = wp x W with some scale factor wo. It is indeed special case vary the scale factor wo and train SIREN networks
of our approach. To examine its performance, under the same hyper- parameter setting, and present the test error over the Natural and Text data-set in Figure TT] If we take wy 1, the main difference between SIREN and our baseline is the activation function (sine vs. ReLU). However, SIREN just performs similarly to our baseline (Plain MI.P with no input mapping). This suggests that the periodic activation function does not play an crucial role in improving the model performance. In the same figure, we can see that the SIREN performance is mostly affected by the value of wn- Therefore, we may argue that the success of SIREN can be attributed to that simple weight factorization instead of the sine activations with the associated initialization scheme W U( u(-v6/d. /6/d). Nevertheless, the best PSNR that SIREN achieves still significantly lower than the proposed random weight factorization with positional encodings Gaussian Fourier features.
21

---

Under review as a conference paper at ICLR 2023
Fitting Natural image
Fitting Text image
25
baseline
25
baseline
PNSR
20
PNSR
20
15-
15
10
2
2
Scale w/
Scale chb
Figure 11: 2D Image Regression: PSNR of training SIREN networks with different scale factor for the Natural and Text data-set, respectively. Error bars are plotted over different images in the data- set. The baseline (black dash) represents the result of training a plain MLP with no input mapping.
H 2D COMPUTED TOMOGRAPHY
data-sets: procedurally-generaaee Shepp-Logan phantoms (Shepp Lo- gan 1974) and 2D brain images from the ATLAS data-set (Liew et al. 2018). Each data-set consist of 20 images of 512 x 512 resolution. To generate the training data, we compute 20 and 40 syn- thetic integral projections at evenly-spaced angles for every image of Shepp and ATLAS data-set, respectively.
For this task,
Similar to the previous tasks, we first perform a parametr sweep over different random seeds, learning rates, activation functions, scale hyper-parrmettes, as well as MLP architectures with differ- ent depths and widths. Particularly, we initialize tion. Each model is trained via a full-batch gradient descent for 2,000 iteration using the the Adam optimizer (Kingma & Ba 2014) (Kingma & Ba 2014) with default settings and 200 warm-up steps. Figure and Figure 13 visualizes the the performance MLPs with each weight parameterizations under different hyper-parameter settings.
N (1, 0.1) for using random weight factoriza-
Furthermore, under the best hyper-paramen parameter setup, we compare the performance of MLPs with different input mappings. In experiments, Table [00 summarizes the test PSRN over Shepp and ATLAS data-set, respectively. For different input mappings, random weight factorization yields the best PSNR, consistently outperforming other parameterizations. Besides, we plot some model predictions corresponding to Gaussian input mapping in Figure 14|and and 15]
Shepp data-set
Plain
AA
WN
RWF (ours)
No mapping
22.78 1.35
23.44 1.27
23.37 0.84
24.39 1.56
Positional encoding
29.56 = 1.73
29.71 + 1.80
29.87 1.86
32.43 = 1.51
Gaussian
30.08 = 1.63
30.44 + 1.69
30.55 + 1.72
33.70 + 1.29
ATLAS data-set
Plain
AA
WN
RWF (ours)
No mapping
15.87 = 0.66
16.07 + 0.68
16.17 0.63
16.49 - 0.61
Positional encoding
21.44 + 0.94
21.54 + 0.92
21.70 - 0.63
23.34 - 0.92
Gaussian
22.02 = 0.93
22.02 - 1.05
22.10 1.04
23 61 = 0.83
Table 10: 2D Computed Tomography: Mean and standard deviation of PSNR obtained by training MLPs with different input mappings and weight parameterizations for the Shepp and ATLAS data- set, respectively.
22

---

Under review as a conference paper at ICLR 2023
30
1s04
Leaming rate
PSNR
20
1e-03
None
A.A
WN
RWF
0
10
211
30
PSNR
10-
Canssaan scale
Activation
gelu
relu
30
PSNR
PSNR
256
Networks depth
Network width
128
0
10
20
30
0
20
30
PSNR
PSNR
Figure 12: 2D Computed Tomography: Mean and standard deviation PSNR obtained by training MLPs with different pprramtteee setup for the Shepp data-set. In particular, top left panel plots a min/max bar, indicating that over the entire range of parameter settings tested, RWF yields the best optimal performance compared to the alternatives.
23

---

Under review as a conference paper at ICLR 2023
25
les04
Leaming rale
PSNR
10-
1e-03
None
AAA
WN
RWF
10
15
20
PSNR
10-
Gausela scale
Activation
gelu
melu
10
15
20
PSNR
PSNR
256
Network deptà
Network widti
128
64
0
5
10
15
20
5
10
13
20
25
PSNR
PSNR
Figure 13: 2D Computed Tomography: Mean and standard deviation PSNR obtained by training MLPs with different hyper-parameter setup for the Atlas data-set. In particular, top left panel plots a min/max bar, indicating that over the entire range of paramete settings tested, RWF yields the best optimal performance compared to the alternatives.
24

---

Under review as a conference paper at ICLR 2023
GT
Plain
WF
TAVF
GT
Plain
AA
WN
RWF
Figure 14: 2D Computed Tomography: Predictions of trained MLPs with Gaussian Fourier features and with different weight parameterizations for the Shepp data-set.
WN
Plain
A
WN
RWF
Figure 15: 2D Computed Tomography: Predictions of trained MLPs with Gaussian Fourier features and with different weight parameterizations for the ATLAS data-set.
Comparison with SIREN (Sitzmann et al., 2020): We also test the performance of SIREN for this example. Specifically, we train SIREN network with different scale factor under the same hyper-parameter settings. As shown Figure I6] SIREN using the optimal scale factor just slightly better than our baseline (plain MLP with no input mapping), but still worse than using positional encodings or random Fourier features.
Shepp
ATLAS
20
25
15+
PNSR
20
PNSR
15
10:
10
baseline
baseline
Scale who
Stale u0p
Figure 16: 2D Computed Tomography: PSNR training SIREN networks with different seale factor for the Shepp and ATLAS data-set, respectively. Error bars are plotted over different images in the data-set. The baseline (black dash) represents the result training a plain MLP with no input mapping.
25

---

Under review as a conference paper at ICLR 2023
I 3D SHAPE REPRESENTATION
For this example, we consider two complex triangle meshes Dragon and Armadillo, both of which contain hundreds of thousands of vertices. In experiments, every mesh is rescaled to fit inside the unit cube [0, 1 such that the centroid of the mesh is (0.5, 0.5,
We represent each shape by MLPs with different input mappings and weight parameterizations. For models with input mappings, we use the same hyper-parameters as in Tancik et al. 2020). For models using random weight factorization, we initialize the scale factor using the recommended settings s ~ (1,, 0.1). All networks are trained by minimizing a cross-entropy loss to match the corresponding classification labels (0 for points outside the mesh, for points inside).
We train each model (8 layers, 128 channels, ReLU activations) via mini-batch gradient descent for 10* iterations using the Adam optimizer (Kingma & Ba 2014) with start learning rate 10-4 and exponential decay by factor of 0.1 for every 5, 000 steps. The batch size we use is 8192. To emphasize the learning fine surface details, we calculate the test eroo on set close to the mesh surface, which is generated by randomly choosing mesh vertices that have been perturbed by a random Gaussian vector with a standard deviation of 0.01.
The resulting IoU scores of each model is reported in Table One observe consistent im- provements RWF across different input mappings and data-sets, outperforming the other param- eterizations. Moreover, the learned shape representations depicted in Figure
Dragon data-set
Plain
AA
WN
RWF (ours)
No mapping
0.894
0.894
0.891
0.924
Positional encoding
0.967
0.968
0.970
0.977
Gaussian
0.980
0.981
0.981
0.984
Armadillo data-set
Plain
AA
WN
RWF (ours)
No mapping
0.842
0.846
0.845
0.901
Positional encoding
0.965
0.967
0.967
0.972
Gaussian
0.978
0.976
0.975
0.982
Table 11: 3D Shape Representation: IoU of training MLPs with different input mappings and weight parameterizations for the Dragon and Armadillo data-sets.
GT
Plain
RVF
WX
Figure 17: 3D Shape Representation: Rendered shape representations obtained by training MLPs with Gaussian Fourier features and with different weight parameterizations.
J 3D INVERSE RENDERING FOR VIEW SYNTHESIS
For this task, we use the NeRF Lego data-set 120 images downsampled to 400 x 400 pixel resolution. The data-set is split into 100 training images, validation images, and 13 test images. In
26

---

Under review as a conference paper at ICLR 2023
experiments, only Gaussian Fourier features of scale 10, it has been empirically validated to be the best imput mapping in the previous tasks.
We train MLPs (5 layers, 256 channels, ReLU activations) with different parameterizations and 10 random seeds for x 10 iterations over random seeds. We use the Adam optimizer (Kingma & Ba 2014) with start learning rate of 10-3 and warmup exponential decay by factor of 0.9 for every 000 steps. The batch size is 2048. In Figure 18 we visualize the test PSNR each model during training. In comparison with other three parameterizations, the MLP with RWF achieves the best PSRN. Some visualizations are provided in Figure 19
26.0
PSNR
25.5
Plsin
25.0
AA
WN
RWF
24.5
20
30
40
50
Iter x
Figure 18: 3D Inverse Rendering: Test PSNR of MLPs with different weight parameterizations during training.
Plain
GT
WN
RWF
GT
Plain
WN
R&V
Figure 19: 3D Inverse Rendering: Rendered views of trained MLPs with different weight parame- terizations.
K
SOLVING PDES
In this section, we present the implementation details PINNs for solving advection and Navier- Stokes equation, respectively.
27

---

Under review as a conference paper at ICLR 2023
K.1
ADVECTION EQUATION
Recall
du du σt
c dr
0,
x x (0, 2m), 1 [0,1]
(K.1)
u(r, 0) g(x), x E 2x)
(K.2)
with periodic boundary conditions and 50.
We represent the unknown solution u by an MLP ug. particular, we impose the exact periodic boundary conditions by constructing special Fourier feature embedding of the form (Dong & Ni. 2021
(2, (c)s(2)), sin(
(K.3)
The network can be trained by minimizing the composite loss below
[()) sic 2e (0) (0)
(K.4)
where
No.
1
(i)
(0 x1e) -g(x1e)
(K.5)
N
i=1
(0) 1 N dt (t,, xt) + c
N.
Ouo
Duo
Ox (t) xx)
(K.6)
Here we set No.- 128 and N, 1024, and ( | (Xrr, br (x) - are randomly sampled from the computational domain, respectively, at each iteration gradient descent. In addition, we take Aie 100, A, for better enforcing the initial condition. It is worth pointing out that all the network derivatives are computed via automatic differentiation.
To enhance the model performance, we introduce the curriculum training (Krishnapriyan et al, 2021) and causal training (Wang et al. 2022) in the training process.
Curriculum training starts with simple PDE system and progressively solves the target PDE system. For this example, it is accomplished by minimizing the above PINN loss with a lower advection coefficient 10 first and then gradually increasing to the target value (i.e. 50) during training.
Causal training aims impose temporal causality during the training of a PINNs model by appro- priately re-weighting the PDE residual loss at each iteration of gradient descent. Specifically, we split the temporal domain into M chunks [0, At), [A,, 2Al),.... and assign a weight the corre-
sponding temporal residuals losses (0) as
M
(,) = w; (,C,(
(K.7)
with w1 = 1, and
w; exp
expl - Σ (0)), for 2,...,
(K.8)
Here e is a so-called causal parameter, which is user-specified hyper- parameter that determines the slope the causal weights. We take M 16 and E = 0.1 in this example.
We initialize MIPs (5 layers, 256 channels, tanh activations) with different weight parameterizations and 10 random seeds. We train each model with different strategies for x 10 iterations using the Adam optimizer Kingma & Ba 2014) with start learning rate of 10-3 and exponential decay by a factor of 0.9 for every 5, 000 steps. The resulting relative L errors are presented in Table 12 In contrast to the failure of the other three parameterizations, RWF is the only that enables PINN models to solve the advection equation with a reasonable and stable predictive accuracy. Further improvements can be obtained by combining RWF with curriculum causal training strategies.
28

---

Under review as a conference paper at ICLR 2023
These conclusions further clarified by the visualizations in Figures 20 and 21 where the
predicted solutions corresponding to RWF are in excellent agreement with the ground truth.
Advection
Plain
AA
WN
RWF (ours)
Regular
25.94% + 8.18%
36.27% + 16. 12.
67.40% = 1.72%
3.10% + 0.63%
Curriculum
47.13% + 20.40%
23.79% + 25.86%
24.10% + 18.85%
2.75% + 1.74%
Causal
5.92% + 2.85%
10.22% + 13.50%
37.38% = 15.94%
2.08% = 0.29%
Table 12: Caption
GT
Plain
AA
HLS
0.5
0.5
0.0
3.0
-0.5
0.0
0.5
1.0
0.5
0.0
0,5
1.0
Curriculum
WN
RWF
10F
005
0.5
Bel 12
0.0
3.0
10
Plain
WIN
-0.5
RWV
20
100
200
0.5
0055
1.0
Iter 1000
Figure 20: advection: Predicted solutions of training PINNs with different weight parameterizations using curriculum training, as well as the evolution of the associated relative L errors during training.
GT
Plai
AA
-0.5
-0.5
-1.5
000
10.0
00.5
1.0
Cousuil
WN
TOVF
Bel. 12 stre
0.0
2.0
10-3
Picin
WIN
-0.
-115
RWV
10-4
200
0.5
100
0.5
1.0
110
Figure 21: advection: Predicted solutions of training PINNs with different weight parameterizations using causal training, well as the evolution of the associated relative I errors during training.
29

---

Under review as a conference paper at ICLR 2023
K.2 AAVIER-STOKES EQUATION
The underlying PDE system for this benchmark takes the form
Vu + Vp- Au 0, (z,y) (0.1)?
(K.9)
u = 0, (a, E (0, 10. (1,0), (z, I
(K.10)
(K.11)
= (0,0), (z, on L'2.
(K.12)
Here H1 is the top boundary of a square cavity, while I2 denotes the other three sides of the cavity.
We represent the unknown solution u, U,P using a neural network ug.
[,y vl
us
[uo; vo.Pol.
(K.13)
Then, the PDE residuals defined by
Oue
Oue
Opo
Puo
8uo
RØ
uo
Vo
(K.14)
Oy
8x
Re
Ox22
Ove
due
Ope
0² ue
0
R
u
ue
(K.15)
0x
0y
Oy
Re
Ox?
0y²
oue
Ove
+
(K.16)
Oy
Given these residuals, along with a set of appropriate boundary conditions, we can now formulate loss function for training physics-informed neural network
Cr. + Ar, (r) + Ar (0)
(K.17)
with
Luy
2 a vh) -wh]
(K.18)
No
No.
Co, (0) =
Iv (at, v6) -612
(K.19)
No
N
(0)
Σ
R2 (1t, w/
(K.20)
N,
==1
(0) N,
RR (of v/)]
(K.21)
N: (0) N RR (e/.:))²
(K.22)
where { (4/, 2 and { (#f, vi) /b denote the boundary data for the two velocity com-
N.
ponents at the domain boundaries Ij and 2,, respectively, while (x/. v() ii=1 is set of collocation points for enforcing the PDE constraints. All of them are sampled randomly at each iteration of gra- dient descent. In experiments, we set N₂ 256, N, 1024 and X, X, 100, \r., A
Ar_ 1.
We employ an MLP (5 layers, 128 channels, tanh activations) to represent the latent variables of interest, and train the network with different parameterizations for 10a iterations using the Adam optimizer (Kingma & Ba, 2014) with start learning rate of 10-3 and exponential decay by factor of 0.9 for every 2, 000 traming iterations. Moreover, we use a modified MLP architecure (see definition below) and curriculum training enhance the model stability and performance. For the curriculum training, we minimize the loss with Re = 100 and Re 500 for 10 iterations sequentially and change the Reynolds number to Re 1, 000 for the rest of the training. The
30

---

Under review as a conference paper at ICLR 2023
resulting relative L errors over 10 random seeds are reported in Table 13 We that RWF performs the best among all parameterizations by a large margin. Some visualizations are shown in and Figure 22 We attribute this significant performance improvements the fact that PINN models often suffer from poor initializations, and RWF precisely mitigates this by being able to reach better local minima that are located further away from the model initialization neighhorhood.
Navier-Stokes
Plain
AA
WN
RWF (ours)
MLP
11.93% + 4.59%
10.95% + 2.76%
14.62% + 4.04%
6.79% + 1.25%
Modified MLP
3.66% = 0.51%
3.97% + 0.42%
3.80% = 3.97%
2.61% + 1.92%
Table 13: Navier-Stokes: Relative 2 errors of training conventional MLPs and modified MLPs with different weight parameterizations, respectively.
In Wang et al. 2021b) Wang et al. proposed novel architecture that was demonstrated to outperform conventional MLPs across variety of PINNs benchmarks. Here, we will refer to this architecture as "modified MIP"" The forward pass of a I-layer modified MLP is defined as follows
Modified MLP:
U = 0 (W b1), V (W2 +
(K.23)
H (1) (W x D()),
(K.24)
z = a (W H + b((+1)) 1= 1,...., L-1, (++ (1- ( o U z 3 V, l 1= I....,
(K,25)
(K.26)
u u ( (L) 6 ((++1)
(K.27)
where a denotes a nonlinear activation function, @ denotes point-wise multiplication. All trainable parameters are given by
0 (W1, b1. W2, b1, (W(), b
(K.28)
This architecture is almost the same as a standard MLP network, with the addition of two encoders and minor modification in the forward pass. Specifically, the inputs æ are embedded into a feature space via two encoders U. V, respectively, and merged in each hidden layer of a standard MLP using a point-wise multiplication.
1.0
1.0
1.0
1.0
1.0
1.0
0.8
0.8
0.4
0.4
0,0
0.4
0.4
0.4
0.2
0,2
0.2
0.2
0.2
0.2
0.0
0.5
0.0
1.0
0.0
0.5
1.5
WN
RWF
1.0
1.0
1.0
IP
AA
0.8
0.8
0.5
0.8
Ret. E7 etnee
WN
RVV
0.6
0,2
0.2
0,2
02
0.0
0,81
50
100
0.
1.0
0.
Figure 22: Navier-Stokes: Predicted solutions obtained by modified MLPs with different weight parameterization, as well as the evolution of the associated relative I errors during training.
31

---

Under review as a conference paper at ICLR 2023
L LEARNING OPERATORS
Overview of DeepONets:
DeepONet is supposed to approximate an operator G between func- tional spaces. DeepONet consists of two separate neural networks referred to the "branch net" and "trunk net",, respectively. The branch net takes a vector a as input and returns a features embed- ding [b1, b2. ... bull E R9 output, where [a(ar) aax)), ..... a (xxm) represents a function evaluated at collection of fixed locations { The trunk net takes the continuous coordi- nates y as inputs, and outputs a features embedding [4, 12,...,((, 1a]"" E R2, The DeepONet output is obtained by merging the outputs of the branch and trunk networks together via dot product
Gola)(v) = by (a (2x1) a (x2) a (2m)) te())
(L.1)
branch
trunk
where & denotes the collection of all trainable weight and bias parameters in the branch and trunk networks, These parameters can be optimized by minimizing the following mean square error loss
((0)
1
Go(a (1) ((aaa (3)
(L.2)
NP
j=1
N
1 NP
(x)) a (zm))t -((())))))))))
(L.3)
where ( 1==1 denotes N separate input functions sampled from function space U. For each a (o) 1
{u) 1 are locations in the domain of (() and G(a (v) ()) is the corresponding output
data evaluated at v Contrary to the fixed sensor locations of {æ; ) we remark that the locations of { may vary for different thus allowing construct continuous representation of the output function G(aa)
Remark: All parameterizations (AA, WN, RWF) are applied to every dense layer of the DeepONet architecture (in the cases where such parametrizations employed).
Case
Plain
AA 0.95% 0.46%
WN
RWF
DR Darcy
1.09% = 0.53%
0.97% = 0.46%
0.50% + 0 0.26%
2.03% + 1.54%
2.06% 1.70%
2.05% + 2.00%
1.67% 1.70%
Burgers
5.11% = 3.79%
4.71% + 3.39%
4.37% = 2.63%
2.46% + 2.05%
Table 14:
Learning operators: Relative L errors of trained infooiee)) DeepONets over the test data-set of different examples.
L.1 DIFFUSION-REACTION
The underlying PDE for this benchmark takes the form
Tu 8u kuu + a(x), (a, E (0. x (0,1).
(L.4)
Öt
Oz2
subject to zero initial and boundary conditions.
Data Generation: We sample N 5, 000 input functions a(x) from a GRF with length scale l 0.2 and solve the diffusiin-rraattinn system using second-order implicit finite difference method on a 100 x 100 equispaced grid. To generate the training data, we randomly take / 100 measurements from each solution. The test data-set contains another 100 solutions evaluated the same mesh
We represent the solution operator by DeepONet Ge where the branch and trunk networks two separate MLPs (5 layers, 64 channels, tanh activations). The model is trained for 10 iterations using the Adam optimizer (Kingma & Ba 2014) with start learning rate of 10-3 and
32

---

Under review as a conference paper at ICLR 2023
exponential decay by factor of 0.9 for every 1000 steps. The mean and standard deviation of the relative L errors over the test date-set are shown in Table 14 Figure provides several representative predictions using RWF.
Imput sumple
Exact
Prodicted
1.0
1.0
0.5
0.5
-2
0.0
0.0
0..5
1.0
0.0
0.5
1.0
0.0
0.5
1.0
Inpub sumupl
Exact
Prodicted
1.0
1.0
a(a)
~- 0.5
0.5
0.0
0.0
0.0
0.5
1.0
0.0
0.5
1.0
0.0
0.5
1.0
JF
Input sample
Exact
Prodittd
1.0
0.5
0.0
0.0
0.0
0.5
1.0
0.0
0.5
Figure 23: Diffusion-reaction (DR) PDE: Predicted solutions of trained DeepONet with random
weight factorization, corresponding to randomly chosen input samples in the test data-set.
L.2
DARCY FLOW
The PDE system for this benchmark takes the form
V. (a Vu) = 1, (x.y) E (0, u 0, (z. E 0(0, 1)
(L..5
(L.6)
Data Generation: We sample the coefficient function from a Gaussian random field with length scale / -0.5 and solve the associated Darcy flow using finite element method. The training data contains 2, 000 solutions evaluated at 64 x 64 uniform mesh while the test data contains 100 solutions on the same mesh.
We represent the solution operator by DeepONet Go. where the branch network is convolutional neural network (CNN) for extracting latent feature representation of the input coefficients and the trunk network is 4-layer MLP with GELU activations and 128 neurons per hidden layer. We train each model with different parameterizations for 5 x 104 iterations using the Adam optimizer Kingmaa & Bal 2014) with start learning rate of 10 and an exponential decay by factor of 0.9 for every 1000 steps. The results are summarized Table [4] and some predicted solutions are plotted in Figure 24
33

---

Under review as a conference paper at ICLR 2023
Iuput sumple
Exaet
Predicted
1.0
0.55
0.0
0.0
0.0
0.00
0.5
1.0
0.0
0.5
1.0
0.0
0.5
Input sumple
Exact
Predicted
1.0
1.0
1.0
0.05
0.5
m.0
0.0
1.0
0.5
1.0
0.0
0.5
Input sample
Exact
Predicted
1.0
1.0
1.0
0.5
> 0.5
0,5
0.0
0.0
0.0
0.0
0.5
1.0
0.0
0.5
1.0
0.0
0.5
1.0
2
Figure 24: Darcy PDE: Predicted solutions trained DeepONet with random weight factoriza- tion, corresponding randomly chosen input samples in the test data-set.
L.3 BURGERS' EQUATION
Recall the one-dimensional Burgers' equation is given by
du du duu o, ( x (0,
di + de
(L.7)
dx²
subject to initial and the periodic boundary conditions
u(I, 0) a(x), (0,
(L.8)
Data Generation: The training data only consists of 1, 000 input functions sampled from Gaussian random field ~ / (0, (- 5 To generate the test data-set, we sample an- other 100 input functions from the same Gaussian random field and solve the Burgers' equation using the Chebfun package Driscoll al. (201)) with spectral Fourier discretization and fourth- order stiff time-stepping scheme (ETDRK4) Cox Matthews (2002) with a time-step size of 10-4 Temporal snapshots of the solution saved every At 0.01 to give 101 snapshots in total. Consequently, the test data-set contains 500 realizations evaluated at 100 x 101 spatio-temporal
grid.
Our objective here is to learn the solution operator mapping initial conditions a(x) to the associated full spatio-temporal solution u(x; Z). Here proceed by representing the solution operator by mod- ified DeepONet architecture Wang et al. (2021d) outlined below. To impose the periodic the exact boundary condition, we apply Fourier feature mapping to the input coordinates before passing them through the trunk network
[a., + [cos(2 sin (2 t].
(L.9)
34

---

Under review as a conference paper at ICLR 2023
Then the PDE residual is then defined by
(Ga(a)
Go(a) @Go(a) dr
GG((a)
Ro [a]
(I..10)
dr
Consequently, physies-informed DeepONet can be trained by minimizing the following weighted loss function
((0) - + (0),
(L.II)
where
(0)
Σx Σ NP
Gola ( (z)
-u (=
(I.11)
i=1 j=1
(0) NQ ==1==1
Σ
p( (1)
(9) ?
(L.13)
For this example, we take N -4, 100 and Q 512, which means that we randomly sample N = 64 input functions from the training data-sct and Q 512 collocation points inside the computational domain. In particular, we set Aic = 100, X 1 for better enforcing the initial condition across different input samples. The model with different parameterizations is trained for 102 iterations using the the Adam optimizer KKngmma & Ba 2014) with start learning rate of 33 and exponential decay by factor of 0.9 for every 2, 000 steps. We report the test errors in Table
14 and visualize some predicted solutions in Figure 25
Modified DeepONet: follows
Wang et al. (021d modify the forward pass of an L--ayer DeepONet as
(Wa+bba), V (W,y b,),
(L.14)
a (1) )( a + (()), ( ((W y + b()),
(L.15)
Z (W! H b("), z z()= = =0(W, y H b)),
1= 1= L-1, 1,
(L.16)
(++1) = (1- "U=ZZ"" eV, 1= 1=1,..,L-1,
(L.17)
((++)) = -11 -Z )0 OV, 1= 1=1,..,, L-1,
(L.18)
H (WW!L) H (L-1) +b(H), H = ( (L-)) + ((()).
(L.19)
Go(a)(v) ( (L) H
(L,20)
where C denotes point-wise multiplication, & denotes activation function, and O repre- sents all trainable parameters of the DeepONet model. In particular, (w!") (+1)) L+1 J== and
Jt ) (W(() b (1+1) L+1 are the weights and biases of the branch and trunk networks, respectively. we embed the DeepONet inputs a and y into a high-dimensional feature space via two encoders U, V. respectively. Instead of just merging the propagated information in the output layer of the branch and trunk networks, we merge the embeddings U. V in each hidden layer of these two sub-networks using a posnt-wise multiplication (equation L.17 L.18).. Heuristically, this design may not only help input signals propagate through the DeepONet, but also enhance its capability of representing non-linearity due the extensive of woint-wise multiplications.
35

---

Under review as a conference paper at ICLR 2023
Input sumple
Esartt
Prodieted
0.5
-0.2
0.0
0.0
0.0
0.55
1.0
0.0
0.5
1.0
0.0
0.5
Imput sumple
Prodieted
1.0
1.0
at.)
0.0
0.5
0.5
-0.2
0.0
0.0
0.5
1.0
0.5
1.0
0.0
0.5
Input sample
Esact
Predicted
1.0
1.0
0.25
0(2))
0.00
0.5
0.5
00.25
0.01
0.0
0.0
0.5
1.0
0.0
0.5
1.0
0.0
0.5
1.0
Figure 25: Burgers PDE: Predicted solutions of a trained physics-informedt informed DeepONet with random weight factorization, corresponding randomly chosen input samples in the test data-set.
36


Reviews:
[{"review_id": "sMR22dnAX04", "reviewer": "Reviewer_mJaA", "summary": "", "questions": "", "limitations": "", "rating": 5, "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": "This work proposes random weight factorization as an initialization of the conventional linear layers. Theoretically, it is shown how this factorization alters the underlying loss landscape and effectively enables each neuron in the network to learn using its own self-adaptive learning rate. Empirically, the proposed initialization method can result in improving the training on a variety of tasks, including image regression, shape representation, computed tomography, inverse rendering, solving partial differential equations, and learning operators between function spaces. ", "summary_of_the_review": "I didn’t find much difference in this work compared to the adaptive activation function method. In addition, the random weight factorization method only applies to the initialization. ", "strength_and_weaknesses": "Strength:\nOverall, the paper is well-organized and easy to follow. This paper studies how to better initialize and factorize the weight parameters of the multi-layer perception. Specifically, they suggest using random weight factorization to factorize the original weight parameters. They first theoretically show that the space of this factorization is flexible enough to be arbitrarily close to the global minima. After that, they show that the proposed weight factorization can be recognized as a gradient descent with a self-adaptive learning rate. I also admire the tremendous experimental work in this paper.\nWeakness:\n1. My biggest concern is that I didn’t find much difference between the proposed method and the adaptive activation function method. I check the algorithm in Appendix B. It seems that the random weight factorization is only applied once in the initialization. After that, the model was trained with the standard gradient descent method. If we omit the initialization step, this is almost like training a model with an adaptive activation function. I am also concerned about how to ensure that the initialized scale factor s^{(l)} is better than the vector of ones. Please correct me if I am wrong.\n\n2. The author claims that the proposed method can effectively mitigate spectral bias. But I am not aware of any such analysis in the experiment part. \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": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "gCCuLiWvCBy", "reviewer": "Reviewer_C8fw", "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": "This paper proposes random weight factorization, a new parametrization for linear neural network layers inspired by weight normalization. It argues that the parametrization reduces the distance between different parameter configurations in the loss landscape and reduces spectral bias, leading to performance improvements across a range of tasks.", "summary_of_the_review": "This is a simple, but potentially broadly useful idea. To me this puts the bar on the empirical evaluation rather high. Unfortunately, I believe that the evaluation falls short of that, so I do not recommend acceptance.", "strength_and_weaknesses": "Strengths:\n+ The proposed method is very simple.\n+ The method is scalable and potentially widely applicable.\n+ The paper presents results on a broad range of settings for the experiments.\n\nWeaknesses:\n- My understanding based on Tab. 3 in the appendix is that the optimizer setup is fixed across the parametrizations that are compared in each benchmark. As I read it, the paper argues that random weight factorization outperforms the alternatives, i.e. the best case performance is better. However, then the optimization setup would need to be tuned for each method independently. As it stands, the paper can roughly support the statement \"For a range of benchmarks there exists some setting of the hyperparameter where the proposed method performs best\". However, this is not all too relevant, it would either need to typically perform best across a range of hyperparameter setups (i.e. be more robust) or consistently deliver the best optimal performance. To me, the experimental setup is critically flawed and would need to be addressed for me to recommend acceptance (I'm open to discussion on this point and do not strictly expect all experiments to be rerun with a hyperparameter sweep for each baseline).\n- There are no error bars in most experiments and I did not find any mention of multiple runs with different random seeds, which significantly weakens any reported performance gains. Also as a result, I'm not sure how to interpret e.g. the ablation study in Fig 6, it seems extremely noisy across the board with no observable trend beyond high errors for a mean of 4.\n- The reported performance gains are mostly incremental (although this is probably not unexpected with a relatively small change such as the weight parametrization on something as basic as feedforward neural networks with best practices for training them that have been established over the past decade).\n- The range of architectures is quite narrow, it seems like the parametrization would be straight-forward to plug into CNNs and Transformers if I'm not missing anything.\n- I don't find the argument around the proximity of parametrizations too convincing, I'd expect these to be increasingly far away from the origin, especially in high dimensions. This probably would make them hard to reach by gradient-based optimization, in particular in conjunction with L2 regularization (although I suppose it depends whether that is applied to the weight matrices or the actual parameters). I would appreciate some more discussion around this or some small experiment on an actual neural network architecture.\n\nOther:\n* I haven't seen any of the benchmarks before, so find it difficult to judge the reported performance gains.", "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": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "U-B_88MjSR", "reviewer": "Reviewer_AGKQ", "summary": "", "questions": "", "limitations": "", "rating": 8, "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 proposes a method for improving and accelerating the training of coordinate-based networks by proposing randomized weight factorization. This factorization generalizes weight normalization, but can be applied in the coordinate-based network setting for signal memorization. It is done by factoring each of the weights in the layers into a multiplication between a diagonal scale matrix and a dense weight matrix. The paper demonstrates that this factorization theoretically improves the distance to a local minima, resulting in easier optimization and robustness to poor initializations. This is experimentally validated across a number of signal fitting tasks which coordinate-based networks are applied to, including image and 3D shape memorization, neural radiance field optimization, and partial differential equation solutions.", "summary_of_the_review": "This paper proposes a new factorization for fully connected coordinate-based network architectures, theoretically proves that it has some properties which are amenable to better initialization (resulting in better results), and experimentally validates that this is the case. I find this to be a convincing story, and this is a topic which is very interesting for those who work on applications involving coordinate-based networks. Thus, I believe this paper would have a high impact and recommend it for acceptance.\n\nPOST REBUTTAL UPDATE: After reviewing the authors' response, and the other reviews and responses, I am not inclined to change my score. The authors have empirically addressed the hyperparameter comment, which I believe was valid, and have shown that their factorization leads to improvements. The theoretical section has been improved with respect to prior work. I believe the improvements are noticeable, and I don't believe comments on vastly different architectures are relevant since the paper specifically focuses on coordinate-based networks.", "strength_and_weaknesses": "In my opinion the strengths of the paper are:\n- The paper is extremely clear in description of the proposed method and exposition of theoretical and experimental results. This is a plus, as it will make the impact of the paper significantly higher since it can easily be adapted into application pipelines. \n- The claims of improvement are backed up both theoretically and experimentally, and seem significantly better than the comparable baseline methods. The breadth of experimental experiments for which coordinate-based networks are used is great, and I find it valuable that in all of these situations the randomized weight factorization provides the same improvement.\n\nIn my opinion, there are no glaring weaknesses in the paper. One thing which would be interesting to be directly discussed would be the aspect of computational efficiency. Specifically, does the factorization increase the size of the models, since there are more parameters being optimized after the factorization. Additionally, does the improvement in loss landscape significantly improve the training time of coordinate-based networks across all applications. This is an important question for those who work in neural rendering (ex: with neural radiance fields), and this could be a simple but applicable solution for this.\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": "AMcGrPouVxG", "reviewer": "Reviewer_chW4", "summary": "", "questions": "", "limitations": "", "rating": 6, "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 proposes a simple but well-performing technique of reparameterizing each weight matrix of implicit neural representations (INRs) as a product of diagonal matrix (with diagonal elements being equal to the exponential of a single parameter) and another matrix. Empirically, it turns out that the reparameterized INR converges to a better-fitting solution to the signal being trained.", "summary_of_the_review": "The proposed method is potentially a very useful technique for training implicit neural representations, but I have several concerns regarding the empirical validations of the technique, especially on how the method can be combined with the state-of-the-art architecture, SIREN.", "strength_and_weaknesses": "__Strength__\n- The method is very simple and easy to use, can be used as a drop-in replacement of the weight matrices of INR. I especially like the fact that the method may not incur too much additional compute / memory on top of the vanilla INR training procedure, as the training cost is one of the central issues of INR.\n\n__Weaknesses__\n- The \"exponentiation\" in the main algorithm (in appendix B) is not well-motivated, and is not well discussed in the main text. Could authors give a little more explanation why using $\\exp(s)$ instead of $s$ is a good thing to do? Perhaps a formal ablation study or theoretical justification may help.\n\n- The explanation regarding SIREN is not clear enough to me. In page 4, authors argue that (1) using RWP is better than using the SIREN (2) the benefit of SIREN is mostly from using the reparameterization $\\omega_0 \\cdot \\mathbf{w}$ rather than the periodic activation. I am not really sure about how authors showed this in Appendix G and H. It seems like in Appendix G and H, authors try various values of $\\omega_0$ of SIREN to show that the peak performance of such fixed parameterization is lower than the value achieved by RWP. But it is not clear from the text with which activation function RWP and fixed $\\omega$ are trained with. Did you use periodic activation function, or the ReLU? In either case, I do not see how this experiment leads to a conclusion that having periodic activation is less important than having a scaling factor in SIREN.\n\n- Related to the previous point, I think that perhaps showing that RWP reparameterization boosts the performance of INRs even when the underlying model architecture uses a periodic activation function (like SIREN) is a right way to prove the benefit of RWP. Indeed, in many cases, SIREN is known to achieve way better/faster convergence than the ReLU-based models used in the experiments.\n\n- In fact, I wanted to test the claims of the authors by myself but the code has not been provided (unfortunately). Providing codes will help readers validate the claim.\n\n- I wonder if such reparameterization will help boosting the performance of meta-learned INRs (via MAML or hypernetworks). Is there any related experiment?\n\n- In Theorem 1, I am not sure how meaningful having a small \"distance\" (which has not been properly defined, by the way) is, because the reparameterization also have a critical impact on the scale of the gradient of each component, s and V. In other words, despite having a smaller distance by the reparameterization, the effective number of optimization steps can be similar to the original parameterization.\n\n- Results like Theorem 2 has already appeared in the work of Arora et al. (2018), which should be mentioned in the text. In fact, the theoretical results of Arora et al. may help authors explain the good performance of the proposed method.\n\n---\nArora et al., \"On the optimization of deep networks: Implicit acceleration by overparameterization,\" ICML 2018.", "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": "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": {}}]

Author rebuttal:

Decision: Reject