Title: Can Neural Networks Improve Classical Optimization of Inverse Problems?
Venue: NeurIPS 2023
Accepted: no
Keywords: Inverse problems

neural networks

iterative optimization

chaos

convergence
Abstract: Finding the values of model parameters from data is an essential task in science.
While iterative optimization algorithms like BFGS can find solutions to inverse problems with machine precision for simple problems, their reliance on local information limits their effectiveness for complex problems involving local minima, chaos, or zero-gradient regions.

This study explores the potential for overcoming these limitations by jointly optimizing multiple examples. To achieve this, we employ neural networks to reparameterize the solution space and leverage the training procedure as an alternative to classical optimization.
This approach is as versatile as traditional optimizers and does not require additional information about the inverse problems, meaning it can be added to existing general-purpose optimization libraries.
We evaluate the effectiveness of this approach by comparing it to traditional optimization on various inverse problems involving complex physical systems, such as the incompressible Navier-Stokes equations. Our findings reveal significant improvements in the accuracy of the obtained solutions.

1

Abstract

Finding the values of model parameters from data is an essential task in science. While iterative optimization algorithms like BFGS can find solutions to inverse problems with machine precision for simple problems, their reliance on local information limits their effectiveness for complex problems involving local minima, chaos, or zero-gradient regions. This study explores the potential for overcoming these limitations by jointly optimizing multiple examples. To achieve this, we employ neural networks to reparameterize the solution space and leverage the training procedure as an alternative to classical optimization. This approach is as versatile as traditional optimizers and does not require additional information about the inverse problems, meaning it can be added to existing general-purpose optimization libraries. We evaluate the effectiveness of this approach by comparing it to traditional optimization on various inverse problems involving complex physical systems, such as the incompressible Navier-Stokes equations. Our findings reveal significant improvements in the accuracy of the obtained solutions.

1 Introduction

Estimating model parameters by solving inverse problems [Tar05] is a central task in scientific research, from detecting gravitational waves [GH18] to controlling plasma flows [MLA + 19] to searching for neutrinoless double-beta decay [AAA + 13, AAA + 18]. Iterative optimization algorithms, such as limited-memory BFGS [LN89] or Gauss-Newton [GM78], are often employed for solving unconstrained parameter estimation problems [PTVF07]. These algorithms offer advantages such as ease of use, broad applicability, quick convergence, and high accuracy, typically limited only by noise in the observations and floating point precision. However, they face several fundamental problems that are rooted in the fact that these algorithms rely on local information, i.e., objective values L ( x k ) and derivatives close to the current solution estimate x k , such as the gradient ∂L / ∂x | x k and the Hessian matrix ∂ 2 L / ∂x 2 | x k . Acquiring non-local information can be done in low-dimensional solution spaces, but the curse of dimensionality prevents this approach for high-dimensional problems. These limitations lead to poor performance or failure in various problem settings:

Local optima attract the optimizer in the absence of a counter-acting force. Although using a large step size or adding momentum to the optimizer can help to traverse small local minima, local optimizers are fundamentally unable to avoid this issue.

Flat regions can cause optimizers to become trapped along one or multiple directions. Higher-order solvers can overcome this issue when the Hessian only vanishes proportionally with the gradient, but all local optimizers struggle in zero-gradient regions.

Chaotic regions , characterized by rapidly changing gradients, are extremely hard to optimize. Iterative optimizers typically decrease their step size to compensate, which prevents the optimization from progressing on larger scales.

In many practical cases, a set of observations is available, comprising many individual parameter 37 estimation problems, e.g., when repeating experiments multiple times or collecting data over a time 38 frame [CCC + 19, DJO + 18, GH18, AAA + 13, MAL13] and, even in the absence of many recorded 39 samples, synthetic data can be generated to supplement the data set. Given such a set of inverse 40 problems, we pose the question: Can we find better solutions x i to general inverse problems by 41 optimizing them jointly instead of individually, without requiring additional information about the 42 problems? 43

44

45

46

47

48

49

50

51

To answer this question, we employ neural networks to formulate a joint optimization problem. Neural networks as general function approximators are a natural and straightforward way to enable joint optimization of multiple a priori independent examples. They have been extensively used in the field of machine learning [GBCB16], and a large number of network architectures have been developed, from multilayer perceptrons (MLPs) [Hay94] to convolutional networks (CNNs) [KSH12] to transformers [VSP + 17]. Overparameterized neural network architectures typically smoothly interpolate the training data [BHM18, BPL21], allowing them to generalize, i.e., make predictions about data the network was not trained on.

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

It has recently been shown that this generalization capability or inductive bias benefits the optimization of individual problems with grid-like solution spaces by implicitly adding a prior to the optimization based on the network architecture [UVL18, HSDG19]. However, these effects have yet to be investigated for general inverse problems or in the context of joint optimization. We propose using the training process of a neural network as a drop-in component for traditional optimizers like BFGS without requiring additional data, configuration, or tuning. Instead of making predictions about new data after training, our objective is to solve only the problems that are part of the training set, i.e., the training itself produces the solutions to the inverse problems, and the network is never used for inference. These solutions can also be combined with an iterative optimizer to improve accuracy. Unlike related machine learning applications [KAT + 19, SGGP + 20, SFK + 21, RT21, SHT22, HKT21, SF18, RPM20, ALGS + 22], where a significant goal is accelerating time-intensive computations, we accept a higher computational demand if the resulting solutions are more accurate.

To quantify the gains in accuracy that can be obtained, we compare this approach to classical optimization as well as related techniques on four experiments involving difficult inverse problems: (i) a curve fit with many local minima, (ii) a billiards-inspired rigid body simulation featuring zerogradient areas, (iii) a chaotic system governed by the Kuramoto-Sivashinsky equation and (iv) an incompressible fluid system that is only partially observable. We compare joint optimization to direct iterative methods and related techniques in each experiment.

2 Related work

Neural networks have become popular tools to model physical processes, either completely replacing physics solvers [KAT + 19, SGGP + 20, SFK + 21, RT21] or improving them [TSSP17, UBF + 20, KSA + 21]. This can improve performance since network evaluations and solvers may be run at lower resolution while maintaining stability and accuracy. Additionally, it automatically yields a differentiable forward process which can then be used to solve inverse problems [SF18, RPM20, ALGS + 22], similar to how style transfer optimizes images [GEB16].

Alternatively, neural networks can be used as regularizers to solve inverse problems on sparse tomography data [LSAH20] or employed recurrently for image denoising and super-resolution [PW17]. Recent works have also explored them for predicting solutions to inverse problems [HKT21, SHT22]. In these settings, neural networks are trained offline and then used to infer solutions to new inverse problems, eliminating the iterative optimization process at test time.

Underlying many of these approaches are differentiable simulations required to obtain gradients of 82 the inverse problem. These can be used in iterative optimization or to train neural networks. Many 83 recent software packages have demonstrated this use of differentiable simulations, with general 84 frameworks [HAL + 20, SC19, HKT20] and specialized simulators [TLQL21, LLK19]. 85

Physics-informed neural networks [RPK19] encode solutions to optimization problems in the network 86

of directions [YZWX19, LPY + 21, KGZ + 21]. However, as these approaches rely on loss terms 89 formulated with neural network derivatives, they do not apply to general inverse problems. 90

91

3 Reparameterizing inverse problems with neural networks

92

93

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

We consider a set of n similar inverse problems where we take similar to mean we can express all of them using a function F ( ξ i | x i ) conditioned on a problem-specific vector x i with i = 1 , ..., n . Each inverse problem then consists of finding optimal parameters ξ ∗ i such that a desired or observed output y i is reproduced, i.e.

where L denotes an error measure, such as the squared L 2 norm || · || 2 2 . We assume that F is differentiable and can be approximately simulated, i.e., the observed output y i may not be reproducible exactly using F due to hidden information or stochasticity.

A common approach to finding ξ ∗ i is performing a nonlinear optimization, minimizing L using the gradients ∂ L ∂F ∂F ∂ξ i . In strictly convex optimization, many optimizers guarantee convergence to the global optimum in these circumstances. However, when considering more complex problems, generic optimizers often fail to find the global optimum due to local optima, flat regions, or chaotic regions. Trust region methods [Yua00] can be used on low-dimensional problems but scale poorly to higher-dimensional problems. Without further domain-specific knowledge, these methods are limited to individually optimizing all n inverse problems.

Instead of improving the optimizer itself, we want to investigate whether better solutions can be found by jointly optimizing all problems. However, without domain-specific knowledge, it is unknown which parameters of ξ i are shared among multiple problems. We therefore first reparameterize the full solution vectors ξ i using a set of functions ˆ ξ i , setting ξ i ≡ ˆ ξ i ( θ ) where θ represents a set of shared parameters. With this change, the original parameters ξ i become functions of θ , allowing θ to be jointly optimized over all problems. Here, the different ˆ ξ i can be considered transformation functions mapping θ to the actual solutions ξ i , similar to transforming Cartesian to polar coordinates. Second, we sum the errors of all examples to define the overall objective function L = ∑ n i =1 L i .

For generality, all ˆ ξ i ( θ ) should be able to approximate arbitrary functions. We implement them as an artificial neural network N with weights θ : ˆ ξ i ( θ ) ≡ N ( x i , y i | θ ) . Inserting these changes into Eq. 1 yields the reparameterized optimization problem

We see that the joint optimization with reparameterization strongly resembles standard formulations of neural network training where ( x i , y i ) is the input to the network and F ◦ L represents the effective loss function. However, from the viewpoint of optimizing inverse problems, the network is not primarily a function of ( x i , y i ) but rather a set of transformation functions of θ , each corresponding to a fixed and discrete ( x i , y i ) . Figure 1 shows the computational graph corresponding to Eq. 2.

While the tasks of optimizing inverse problems and learning patterns from data may seem unrelated at first, there is a strong connection between the two. The inductive bias of a chosen network architecture, which enables generalization, also affects the update direction of classical optimizers under reparameterization. This can be seen most clearly if we consider gradient descent steps. There, individual optimization yields the updates ∆ ξ i = -η ∂ L i ∂ξ i with η denoting the step size. After reparameterization,

the updates are ∆ θ = -η ∑ i ∂ L i ∂ξ i ∂ N ∂θ . As we can see, ∂ N ∂θ , which is independent of the specific 132 example, now contributes a large part to the update direction, allowing for cross-talk between the 133 different optimization problems. 134

Net

𝑥𝑖

, 𝑦

𝑖

𝜃

Fwd.

ℒ𝑖

𝑦𝑖

+

𝐿

𝜉𝑖

Despite the similarities to machine learning, the different use case of this setup leads to differences 135 in the training procedure. For example, while overfitting is usually seen as undesirable in machine 136 learning, we want the solutions to our inverse problems to be as accurate as possible, i.e. we want to 137 'overfit' to the data. Consequently, we do not have to worry about the curvature at θ ∗ and will not use 138 mini-batches for training the reparameterization network. 139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

Supervised training. Our main goal is obtaining an optimization scheme that works exactly like classical optimizers, only requiring the forward process F , x i in the form of a numerical simulator, and desired outputs y i . However, if we additionally have a prior on the solution space P ( ξ ) , we can generate synthetic training data { ( x j , y i ) , ξ j } with y j = F ( x j , ξ j ) by sampling ξ i ∼ P ( ξ ) . Using this data set, we can alternatively train N with the supervised objective

Since N has the same inputs and outputs, we can use the same network architecture as above and the solutions to the original inverse problems can be obtained as ξ i = N ( x i , y i ) . While this method requires domain knowledge in the form of the distributions P ( x ) and P ( ξ ) , it has the distinct advantage of being independent of the characteristics of F . For example, if F is chaotic, directly optimizing through F can yield very large and unstable gradients, while the loss landscape of ˜ L can still be smooth. However, we cannot expect the inferred solutions to be highly accurate as the network is not trained on the inverse problems we want to solve and, thus, has to interpolate. Additionally, this method is only suited to unimodal problems, i.e. inverse problems with a unique global minimum. On multimodal problems, the network cannot be prevented from learning an interpolation of possible solutions, which may result in poor accuracy.

Refinement Obtaining a high accuracy on the inverse problems of interest is generally difficult when the training set size is limited, which can result in suboptimal solutions. This is especially problematic when the global minima are narrow and no direct feedback from F is available, as in the case of supervised training. To ensure that all learned methods have the potential to compete with gradient-based optimizers like BFGS, we pass the solution estimates for ξ to a secondary refinement stage where they are used as an initial guess for BFGS. The refinement uses the true gradients of F to find a nearby minimum of L .

4 Experiments

Weperform a series of numerical experiments to test the convergence properties of the reparameterized joint optimization. An overview of the experiments is given in Tab. 1 and additional details of all performed experiments can be found in Appendix B. We run each experiment and method multiple times, varying the neural network initializations and data sets to obtain statistically significant results.

To test the capabilities of the algorithms as a black-box extension of generic optimizers, all experiments use off-the-shelf neural network architectures and only require hyperparameter tuning in terms of decreasing the Adam [KB15] learning rate until stable convergence is reached. We then compare the reparameterized optimization to BFGS [LN89], a popular classical solver for unconstrained optimization problems, and to the neural adjoint method, which has been shown to outperform various other neural-network-based approaches for solving inverse problems [RPM20].

Neural adjoint The neural adjoint method relies on an approximation of the forward process by a surrogate neural network S ( x i , ξ i | θ ) . We first train the surrogate on an independent data set

generated from the same distribution as the inverse problems and contains many examples. We use 175 the same examples as for the supervised approach outlined above but switch the labels to match the 176 network design, { ( x i , ξ i ) , y i } . After training, the weights θ are frozen and BFGS is used to optimize 177 ξ i on the proxy process ˜ F ( ξ i | x i ) = S ( ξ i , x i ) + B ( ξ i ) where B denotes a boundary loss term (see 178 Appendix A). With the loss function L from Eq. 1, this yields the effective objective L ( F ( ξ i | x i ) , y i ) 179 for solving the inverse problems. Like with the other methods, the result of the surrogate optimization 180 is then used as a starting point for the refinement stage described above. 181

182

4.1 Wave packet localization

First, we consider a 1D curve fit. A noisy signal u ( t ) containing a wave packet centered at t 0 is 183 measured, resulting in the observed data u ( t ) = A · sin( t -t 0 ) · exp( -( t -t 0 ) 2 /σ 2 ) + ϵ ( t ) where 184 ϵ ( t ) denotes random noise and t = 1 , ..., 256 . An example waveform is shown in Fig. 2a. For fixed 185 A and σ , the task is to locate the wave packed, i.e. retrieve t 0 . This task is difficult for optimization 186 algorithms because the loss landscape (Fig. 2b) contains many local optima that must be traversed. 187 This results in alternating gradient directions when traversing the parameter space, with maximum 188 magnitude near the correct solution. 189

0

100

200

X

1.0

0.5

0.0

0.5

(a)

Signal

0.2

0.4

0.6

0.8

t

0

20

0

20

(b)

/ 50

10 1

10 2

10 3

Steps

10 1

8×10 0

9×10 0

Loss

(c)

BFGS

Reparameterized

Supervised

Neural Adjoint

2

8

32

128

Dataset Size

0

2

4

6

8

10

12

(d)

Wegenerate the inverse problems by sampling random t 0 and ϵ ( t ) from ground truth prior distributions and simulating the corresponding outputs u ( t ) = Fϵ ( t ) | t 0 ) . Because the noise distribution ϵ ( t ) is not available to any of the optimization methods, a perfect solution with L = 0 is impossible.

Fig. 2c shows the optimization process. Iterative optimizers like BFGS get stuck in local minima quickly on this task. In most examples, BFGS moves a considerable distance in the first iteration and then quickly halts. However, due to the oscillating gradient directions, this initial step is likely to propel the estimate away from the global optimum, leading many solutions to lie further from the actual optimum than the initial guess.

The neural adjoint method finds better solutions than BFGS for about a third of examples for n = 256 (see Tab. 2). In many cases, the optimization progresses towards the boundary and gets stuck once the boundary loss B balances the gradients from the surrogate network.

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

To reparameterize the problem, we create a neural network N that maps the 256 values of the observed signal u ( t ) to the unknown value t 0 . We chose a standard architecture inspired by image classification networks [SZ14] and train it according to Eq. 2. The network consists of five convolutional layers with ReLU activation functions, batch normalization, and max-pooling layers, followed by two fully-connected layers. During the optimization, the estimate of t 0 repeatedly moves from minimum to minimum until settling after around 500 iterations. Like BFGS, most examples do not converge to the global optimum and stop at a local minimum instead. However, the cross-talk between different examples, induced by the shared parameters θ and the summation of the individual loss functions, regularizes the movement in t 0 space, preventing solutions from moving far away from the global optimum. Meanwhile, the feedback from the analytic gradients of F ensures that each example finds a locally optimal solution. Overall, this results in around 80% of examples finding a better solution than BFGS.

For supervised training of N , we use the same training data set as for the neural adjoint method. This 213 approach's much smoother loss landscape lets all solution estimates progress close to the ground 214

215

216

217

218

219

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

243

244

245

246

truth. However, lacking the gradient feedback from the forward process F , the inferred solutions are slightly off from the actual solution and, since the highest loss values are close to the global optimum, this raises the overall loss during training even though the solutions are approaching the global optima. This phenomenon gets resolved with solution refinement using BFGS.

Fig. 2d shows the results for different numbers of inverse problems and training set sizes n . Since BFGS optimizes each example independently, the data set size has no influence on its performance. Variances in the mean final loss indicate that the specific selection of inverse problems may be slightly easier or harder to solve than the average. The neural adjoint method and reparameterized optimization both perform better than BFGS with the reparameterized optimization producing lower loss values. However, both do not scale with n in this example. This feature can only be observed with supervised training whose solution quality noticeably increases with n . This is due to the corresponding increase in training set size, which allows the model to improve generalization and does not depend on the number of tested inverse problems. For n ≥ 32 , supervised training in combination with the above-mentioned solution refinement consistently outperforms all other methods.

A detailed description of the network architecture along with additional learning curves, parameter evolution plots as well as the performance on further data set sizes n can be found in Appendix B.1.

4.2 Billiards

Next, we consider a rigid-body setup inspired by differentiable billiards simulations of previous work [HAL + 20]. The task consists of finding the optimal initial velocity ⃗ v 0 of a cue ball so it hits another ball, imparting momentum in a non-elastic collision to make the second ball come to rest at a fixed target location. This setup is portrayed in Fig. 3a and the corresponding loss landscape for a fixed x velocity in Fig. 3b. A collision only occurs if ⃗ v 0 is large enough and pointed towards the other ball.

Otherwise, the second ball stays motionless, resulting in a constant loss value and

∂

L

0

= 0

.

0.0

0.5

1.0

1.5

2.0

X

0.0

0.5

Y

Cue

Target

(a)

3

2

1

0

4

2

0

2

4

6

(b)

10 1

10 2

10 3

Steps

10

2

10

1

10 0

Loss

(c)

BFGS

Reparameterized

Supervised

Neural Adjoint

4

16

64

256

Dataset Size

0.0

0.1

0.2

0.3

0.4

(d)

This property prevents classical optimizers from converging if they hit such a region in the solution space. The optimization curves are shown in Fig. 3c. BFGS only converges for those examples where the cue ball already hits the correct side of the other ball.

For reparameterization, we employ a fully-connected neural network N with three hidden layers using Sigmoid activation functions and positional encoding. The joint optimization with N drastically improves the solutions. While for n ≤ 32 only small differences to BFGS can be observed, access to more inverse problems lets gradients from some problems steer the optimization of others that get no useful feedback. This results in almost all problems converging to the solution for n ≥ 64 (see Fig. 3d).

In this experiment, the distribution of the solutions P ( ⃗ v 0 ) is not available as hitting the target precisely 247 requires a specific velocity ⃗ v 0 that is unknown a-priori. We can, however, generate training data with 248 varying ⃗ v 0 and observe the final positions of the balls, then train a supervised N as well as a surrogate 249 network for the neural adjoint method on this data set. However, this is less efficient as most of the 250 examples in the data set do not result in an optimal collision. 251

The neural adjoint method fails to approach the true solutions and instead gets stuck on the training 252 data boundary in solution space. Likewise, the supervised model cannot accurately extrapolate the 253 true solution distribution from the sub-par training set. 254

255

256

257

258

259

260

261

262

263

264

265

266

267

268

269

270

271

272

273

274

275

276

277

278

279

280

281

282

283

284

4.3 Kuramoto-Sivashinsky equation

The Kuramoto-Sivashinsky (KS) equation, originally developed to model the unstable behavior of flame fronts [Kur78], models a chaotic one-dimensional system, ˙ u ( t ) = -∂ 2 u ∂x 2 -∂ 4 u ∂x 4 -u · ∇ u . We consider a two-parameter inverse problem involving the forced KS equation with altered advection strength,

where G ( x ) is a fixed time-independent forcing term and α, β ∈ R denote the unknown parameters governing the evolution. Each inverse problem starts from a randomly generated initial state u ( t = 0) and is simulated until t = 25 , by which point the system becomes chaotic but is still smooth enough to allow for gradient-based optimization. We constrain α ∈ [ -1 , 1] , β ∈ [ 1 2 , 3 2 ] to keep the system numerically stable. Fig. 4a shows example trajectories of this setup and the corresponding gradient landscape of ∂ L ∂β ∥ α = α ∗ for the true value of α is shown in Fig. 4b.

0

20

40

Time

0

5

10

15

20

X

(a)

3

2

1

0

0

200

400

(b)

/ 5

10 1

10 2

10 3

Steps

10 2

Loss

(c)

BFGS

Reparameterized

Supervised

Neural Adjoint

4

8

32

128

Dataset Size

0

10

20

30

40

(d)

4

2

0

2

4

Fig. 4c shows the optimization curves for finding α, β . Despite the complex nature of the loss landscape, BFGS manages to find the correct solution in about 60% of cases. The reparameterized optimization, based on a similar network architecture as for the wavepacket experiment but utilizing 2D convolutions, finds the correct solutions in over 80% of cases but, without refinement, the accuracy stagnates far from machine precision. Refining these solutions with BFGS, as described above, sees the accuracy of these cases decrease to machine precision in 4 to 17 iterations, less than the 12 to 22 that BFGS requires when initialized from the distribution mean E [ P ( ξ )] .

Supervised training with refinement produces better solutions in 58% of examples, averaged over the shown n . The unrefined solutions benefit from larger n on this example because of the large number of possible observed outputs that the KS equation can produce for varying α, β . At n = 2 , all unrefined solutions are worse than BFGS while for n ≥ 64 around 20% of problems find better solutions. With refinement, these number jump to 50% and 62%.

This property also makes it hard for a surrogate network, required by the neural adjoint method, to accurately approximate the KS equation, causing the following adjoint optimization to yield inaccurate results that fail to match BFGS even after refinement.

4.4 Incompressible Navier-Stokes

Incompressible Newtonian fluids are described by the Navier-Stokes equations,

with ν ≥ 0 . As they can result in highly complex dynamics [BB67], they represent a particularly challenging test case, which is relevant for a variety of real-world problems [Pop00]. We consider

a setup similar to particle imaging velocimetry [Gra97] in which the velocity in the upper half of 285 a two-dimensional domain with obstacles can be observed. The velocity is randomly initialized in 286 the whole domain and a localized force is applied near the bottom of the domain at t = 0 . The task 287 is to reconstruct the position x 0 and initial velocity ⃗ v 0 of this force region by observing the initial 288 and final velocity field only in the top half of the domain. The initial velocity in the bottom half is 289 unknown and cannot be recovered, making a perfect fit impossible. Fig. 5a,b show an example initial 290 and final state of the system. The final velocity field is measured at t = 56 by which time fast eddies 291 have dissipated significantly. 292

0

25

50

75

100

X

0

20

40

60

80

100

Y

(a)

25

50

75

X

20

40

60

80

Y

(b)

10 1

10 2

10 3

Steps

10 2

Loss

(c)

BFGS

Reparameterized

Supervised

Neural Adjoint

4

16

64

128

Dataset Size

0

20

40

60

80

100

120

(d)

3

2

1

0

1

2

0.4

0.2

0.0

0.2

0.4

Fig. 5c shows the optimization curves. On this problem, BFGS converges to some optimum in all cases, usually within 10 iterations, sometimes requiring up to 40 iterations. However, many examples get stuck in local optima.

For joint optimization, we reparameterize the solution space using a network architecture similar to the previous experiment, featuring four 2D convolutional layers and two fully-connected layers. For all tested n , the reparameterized optimization produces larger mean loss values than BFGS, especially for small n . This results from about 10% of examples seeing higher than average loss values. Nonetheless, 66.7% of the inverse problems are solved more accurately than BFGS on average for n > 4 .

The neural adjoint method nearly always converges to solutions within the training set parameter space, not relying on the boundary loss. With solution refinement, this results in a mean loss that seems largely independent of n and is slightly lower than the results from direct BFGS optimization. However, most of this improvement comes from the secondary refinement stage which runs BFGS on the true F . Without solutions refinement, the neural adjoint method yields inaccurate results, losing to BFGS in 98.2% of cases.

Supervised training does not suffer from examples getting stuck in a local minimum early on. The highest-loss solutions, which contribute the most to L , are about an order of magnitude better than the worst BFGS solutions, leading to a much smaller total loss for n ≥ 16 . With solution refinement, 64%, 73% and 72% of examples yield a better solution than BFGS for n = 16 , 64 , 128 , respectively.

5 Discussion

293

294

295

296

297

298

299

300

301

302

303

304

305

306

307

308

309

310

311

312

313

314

315

316

317

318

319

320

In our experiments, we have focused on relatively small data sets of between 2 and 256 examples to quantify the worst-case for machine learning methods and observe trends. Using off-the-shelf neural network architectures and optimizers with no tuning to the specific problem, joint optimization finds better solutions than BFGS in an average of 69% of tested problems. However, to achieve the best accuracy, the solution estimates must be passed to a classical optimizer for refinement as training the network to this level of accuracy would take an inordinate amount of time and large data sets. Tuning the architectures to the specific examples could lead to further improvements in performance but would make the approach domain-dependent.

When training data including ground truth solutions are available or can be generated, supervised 321 learning can sidestep many difficulties that complex loss landscapes pose, such as local minima, 322

323

324

325

326

327

328

329

330

331

332

333

334

335

336

337

338

339

340

341

342

343

344

345

346

347

348

349

350

351

352

353

354

355

356

357

358

alternating gradient directions, or zero-gradient areas. This makes supervised learning another promising alternative to direct optimization, albeit a more involved one.

The neural adjoint method, on the other hand, yields only very minor improvements over BFGS optimization in our experiments, despite the surrogate network successfully learning to reproduce the training data. This is not surprising as the neural adjoint method tries to approximate the original loss landscape which is often difficult to optimize. Improvements over BFGS must therefore come from regularization effects and exposure to a larger part of the solution space. The fact that the neural adjoint method with solution refinement produces similar results almost independent of the number of data points n shows that the joint optimization has little benefit here. Instead, the refinement stage, which treats all examples independently, dominates the final solution quality. Note that the neural adjoint method is purely data-driven and does not require an explicit form for the forward process F , making it more widely applicable than the setting considered here.

Tab. 2 summarizes the improvements over classical optimizations for all methods. A corresponding table without solution refinement can be found in Appendix B. Considering that reparameterized optimization is the only network-based method that does not require domain-specific information and nevertheless shows the biggest improvement overall, we believe it is the most attractive variant among the three learned versions. Inverse problems for which reparameterized training does not find good solutions are easy to identify by their outlier loss values. In these cases, one could simply compare the solution to a reference solution obtained via direct optimization, and choose the best result.

Limitations We have only considered unconstrained optimization problems in this work, enforcing hard constraints by running bounded parameters through a scaled tanh function which naturally clamps out-of-bounds values in a differentiable manner.

The improved solutions found by joint optimization come with an increased computational cost compared to direct optimization. The time it took to train the reparameterization networks was 3x to 6x longer for the first three experiments and 22x for the fluids experiment.

6 Conclusions and outlook

Wehave investigated the effects of joint optimization of multiple inverse problems by reparameterizing the solution space using a neural network, showing that joint optimization can often find better solutions than classical optimization techniques. Since our reparameterization approach does not require any more information than classical optimizers, it can be used as a drop-in replacement. This could be achieved by adding a function or option to existing optimization libraries that internally sets up a standard neural network with the required number of inputs and outputs and runs the optimization, hiding details of the training process, network architecture, and hyperparameters from the user while making the gains in optimization accuracy conveniently accessible. To facilitate this, we will make the full source code publicly available.

From accelerating matrix multiplications [FBH + 22] to solving systems of linear equations [CHL + 23, 359 SSHR19], it is becoming increasingly clear that machine learning methods can be applied to purely 360 numerical problems outside of typical big data settings, and our results show that this also extends to 361 solving nonlinear inverse problems. 362

363

364

365

366

367

368

369

370

371

372

373

374

375

376

377

378

379

380

381

382

383

384

385

386

387

388

389

390

391

392

393

394

395

396

397

398

399

400

401

402

403

404

405

406

407

408

409

References

[AAA + 13] MAgostini, M Allardt, E Andreotti, AM Bakalyarov, M Balata, I Barabanov, M Barnabé Heider, N Barros, L Baudis, C Bauer, et al. Pulse shape discrimination for gerda phase i data. The European Physical Journal C , 73(10):2583, 2013.

[AAA + 18] Craig E Aalseth, N Abgrall, Estanislao Aguayo, SI Alvis, M Amman, Isaac J Arnquist, FT Avignone III, Henning O Back, Alexander S Barabash, PS Barbeau, et al. Search for neutrinoless doubleβ decay in ge 76 with the majorana demonstrator. Physical review letters , 120(13):132502, 2018.

[ALGS + 22] Kelsey R Allen, Tatiana Lopez-Guevara, Kimberly Stachenfeld, Alvaro SanchezGonzalez, Peter Battaglia, Jessica Hamrick, and Tobias Pfaff. Physical design using differentiable learned simulators. arXiv preprint arXiv:2202.00728 , 2022.

[BB67] Cx K Batchelor and George Keith Batchelor. An introduction to fluid dynamics . Cambridge university press, 1967.

[BHM18] Mikhail Belkin, Daniel J Hsu, and Partha Mitra. Overfitting or perfect fitting? risk bounds for classification and regression rules that interpolate. Advances in neural information processing systems , 31, 2018.

[BPL21] Randall Balestriero, Jerome Pesenti, and Yann LeCun. Learning in high dimension always amounts to extrapolation. arXiv preprint arXiv:2110.09485 , 2021.

[CCC + 19] Giuseppe Carleo, Ignacio Cirac, Kyle Cranmer, Laurent Daudet, Maria Schuld, Naftali Tishby, Leslie Vogt-Maranto, and Lenka Zdeborová. Machine learning and the physical sciences. Reviews of Modern Physics , 91(4):045002, 2019.

[CHL + 23] Salvatore Calì, Daniel C Hackett, Yin Lin, Phiala E Shanahan, and Brian Xiao. Neuralnetwork preconditioners for solving the dirac equation in lattice gauge theory. Physical Review D , 107(3):034508, 2023.

[DJO + 18] S Delaquis, MJ Jewell, I Ostrovskiy, M Weber, T Ziegler, J Dalmasson, LJ Kaufman, T Richards, JB Albert, G Anton, et al. Deep neural networks for energy and position reconstruction in exo-200. Journal of Instrumentation , 13(08):P08023, 2018.

[FBH + 22] Alhussein Fawzi, Matej Balog, Aja Huang, Thomas Hubert, Bernardino RomeraParedes, Mohammadamin Barekatain, Alexander Novikov, Francisco J R Ruiz, Julian Schrittwieser, Grzegorz Swirszcz, et al. Discovering faster matrix multiplication algorithms with reinforcement learning. Nature , 610(7930):47-53, 2022.

[GBCB16] Ian Goodfellow, Yoshua Bengio, Aaron Courville, and Yoshua Bengio. Deep learning . MIT press Cambridge, 2016.

[GEB16] Leon A Gatys, Alexander S Ecker, and Matthias Bethge. Image style transfer using convolutional neural networks. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 2414-2423, 2016.

[GH18] Daniel George and EA Huerta. Deep learning for real-time gravitational wave detection and parameter estimation: Results with advanced ligo data. Physics Letters B , 778:6470, 2018.

[GM78] Philip E Gill and Walter Murray. Algorithms for the solution of the nonlinear leastsquares problem. SIAM Journal on Numerical Analysis , 15(5):977-992, 1978.

[Gra97] Ian Grant. Particle image velocimetry: a review. Proceedings of the Institution of Mechanical Engineers, Part C: Journal of Mechanical Engineering Science , 211(1):5576, 1997.

[HAL + 20] Yuanming Hu, Luke Anderson, Tzu-Mao Li, Qi Sun, Nathan Carr, Jonathan RaganKelley, and Frédo Durand. Difftaichi: Differentiable programming for physical simulation. International Conference on Learning Representations (ICLR) , 2020.

410

411

412

413

414

415

416

417

418

419

420

421

422

423

424

425

426

427

428

429

430

431

432

433

434

435

436

437

438

439

440

441

442

443

444

445

446

447

448

449

450

451

452

453

454

455

456

457

458

459

460

461

462

463

464

465

466

467

468

469

470

471

472

473

474

475

476

477

478

479

480

481

482

483

484

485

486

487

488

489

490

491

492

493

494

495

496

497

498

499

500

501

502

503

504

505

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this paper the authors explore whether better optimization solutions can be found by jointly optimizing several inverse problems together. These inverse problems share a connection as they all can be formulated through a differentiable function $F(\\xi_i | x_i)$. The authors implement the joint optimization through a set of NN parameters $\\theta$ which connect all the different inverse problem variables $\\xi_i = \\hat{\\xi}_i (\\theta)$.\n\nUpon reading the authors' rebuttal to the questions I posed, I am inclined to adjust my score accordingly.\n\nSTRENGTHS:\n* The main problem that the paper is trying to solve is interesting and relevant. As the authors mention in lines 101-105 \"generic optimizers often fail to find the global optimum due to local optima, flat regions, or chaotic regions\".\n* Also the setting expressed in section 3 is not typical which opens up many potential future work leveraging this setting.\n* The experimental results are interesting in how they all improve the solutions by increasing $n$ which does provide evidence of cross-talk between the problems.\n\nWEAKNESSES:\n* The experiments generate inverse problems synthetically. I would like to see a real-life example of a problem that follows the settings exposed in section 3. I understand that several of the equations have practical applications but in this case I'm referring to a real-life problem that has a naturally occurring (not sampled from a known distribution) set of inverse problems that can be connected through a function $F$.\n* The method is not applicable to many problems. I'm not familiar with any ML application that has similar optimization problems that can be pooled together. Put differently, it is unclear to me how restrictive is the setting in section 3 of having a function $F(\\xi_i | x_i)$ that expresses all the inverse problems.\n\nQUESTIONS:\n* I understand the BFGS is used as a high-accuracy optimizer. However, in ML SGD is the workhorse as it seems to avoid getting stock in undesirable parts of the solution space. Could you replicate the experiments in section 3 also showing the loss evolution vs steps for SGD?\n* What are the runtimes of the experiments that you ran on section 4? I would appreciate a runtime comparison between baselines BFGS and Neural adjoint (not only steps, as the steps could be much more expensive for some methods). I know you make mentions of this on lines 346-348 but a break down for each of the problems and baselines would be ideal. How different are the trained NNs from their initialization?\n* Flexible ML models like NNs require several examples to extract the patterns from the data. It is unclear to me how can I learn a NN out of 2 to 256 examples (line 313) as in several problems the neural networks are pretty large too (lines 203-205).\n* Could you provide results without the refinement stage (lines 155-161)? I would like to understand how crucial is this last step is to the results obtained. I understand that this stage is done to achieve high accuracy, however it is unclear to me if this stage is actually doing all the work but simply starting from a better initialization coming from the previous stage or if it is actually the opposite and the second stage only improves marginally the results but the first stage gets relatively close to the optimal solution.\n* \"Supervised learning can sidestep many difficulties that complex loss landscapes pose, such as local minima, alternating gradient directions, or zero-gradient areas\" (lines 321-323). Supervised learning is a paradigm of ML where the model's parameters are trained on a dataset with specified targets. As such supervised learning has nothing to do with complex loss landscapes, etc. Could you elaborate on this line and rephrase in the paper?\n* How you would you ensure reproducibility of the code as in the README.md it reads: \"The code in this directory is preliminary and requires unreleased versions of some libraries\" as it is unclear to me what these \"unreleased versions\" are and whether they would be available as well as your code upon publication (lines 357-358). Could you also add a requirements.txt or setup.py.\n\nLIMITATIONS:\n* The limitations that I see are encapsulated on the questions that I raised above.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper develops a novel approach to gradient-based non-convex optimization. The proposed methodology begins with the reparameterization of the parameter space utilizing neural networks, followed by the application of classical techniques such as BFGS, or alternative Neural Network surrogate models for the forward function, to accomplish the optimization task.\n\nThe efficacy of these methods is verified through four distinctive experiments, which include applications to the Kuramoto-Sivashinsky (K-S) equation and the Incompressible Navier-Stokes equation. The results indicate that the reparameterized optimizer delivers enhanced convergence overall.\n\nSTRENGTHS:\nThe paper presents a compelling concept of reparametrizing the parameter space of the optimization problem using a neural network and accomplishing optimization via a two-step process that employs the trained/optimized neural network as a preconditioner. However, it remains unclear to the me as to why such a reparameterization is likely to benefit the non-convex optimization procedure. Despite this, the experimental results appear to indicate an enhancement, as evidenced by the four case studies investigated by the author.\n\nWEAKNESSES:\nMy primary concern regarding this paper pertains to its lack of rigor. The authors do not clearly define the inverse problem that they are attempting to solve, nor do they provide cogent proofs or insights explaining why the introduced reparameterization would aid the optimization process. It is quite plausible that the limited experimental studies offered in this paper lack generalizability, and it's conceivable that there are counterexamples where optimizers, without the incorporation of reparameterization, achieve superior convergence.\n\nQUESTIONS:\nIn all four cases, the authors do not explicitly define the inverse problem they're attempting to solve, as well as the specific loss function utilized. This lack of clarity greatly complicates the task of comprehending their methodology.\n\nThe paper falls short in normalizing the loss functions - without this normalization, it becomes exceedingly difficult to ascertain whether a loss value, such as \"10\", carries any physical significance.\n\nThe optimization curves presented in the four experiments exhibit some anomalous behaviors. For instance, in Figure 3, the loss of the neural adjoint method escalates with increasing steps, while the supervised method demonstrates oscillation. Could this be attributed to improper tuning of the optimizer's step size?\n\nThe paper doesn't provide a direct comparison between the optimized solutions obtained and the ground truth, except for the reported loss function.\n\nLIMITATIONS:\nWhile the authors recognize that they have not extend the method for constrained optimization problems, I believe there is an inherent limitation in the approach, as it lacks a mechanistic understanding of why such a method would be effective. This comprehension is fundamental for a method to be broadly applicable and reliable.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-06"}

comments: SUMMARY:
The manuscript presents a method to reparameterize and solve multiple inverse problems jointly using neural networks. The manuscript tests the proposed method on multiple inverse problems (including some chaotic problems) and compares against Neural Adjoint and BFGS baselines to show measurable performance improvements.

STRENGTHS:
* The method is simple, and the authors haven't tuned architecture for problems, which would allow their usage as drop-in replacements.
* Comparison against baselines shows the method provides noticeable improvements.

WEAKNESSES:
* The main downside of these methods is the added training cost (which the authors have mentioned in the limitations section)
    * I would recommend adding the training wall clock times + solving times in a table to give potential users of this method a proper estimate.
* Adding benchmarks for the same problems used in the Neural Adjoint paper would strengthen the paper.

QUESTIONS:
* The paper uses BFGS to refine the solution. Would it be possible to use the solver as a predictor-corrector and incorporate BFGS in the training pipeline? Similar to methods used in Deep Equilibrium Networks [1] [2].

[1] [Neural Deep Equilibrium Solvers](https://openreview.net/forum?id=B0oHOwT5ENL)

[2] [Continuous Deep Equilibrium Models: Training Neural ODEs Faster by Integrating Them to Infinity](https://arxiv.org/abs/2201.12240)

LIMITATIONS:
All limitations are clearly stated.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper discusses a novel approach to finding model parameters from data, a crucial task in science. Traditional iterative optimization algorithms like BFGS can accurately solve simple inverse problems, but their reliance on local information can limit their effectiveness in complex situations with local minima, chaos, or zero-gradient regions.\n\nTo overcome these issues, the study proposes the idea of jointly optimizing multiple examples. The authors use neural networks to reparameterize the solution space and utilize the training procedure as an alternative to classical optimization. This method is as versatile as traditional optimizers and does not require additional information about the inverse problems, making it compatible with existing general-purpose optimization libraries.\n\nThe paper evaluates the effectiveness of this novel approach by comparing it to traditional optimization on a variety of complex inverse problems involving physical systems, such as the incompressible Navier-Stokes equations. The findings show significant improvements in the accuracy of the solutions obtained, suggesting that this method could be a powerful tool for tackling complex inverse problems.\n\nSTRENGTHS:\n1. This paper points out a potential new use case of neural networks and deep learning for optimization instead of existing learning to optimize (L2O) methods, that is to use neural networks as part of classic optimization, trying to learn unknown common structures among problem instances of interest. A major difference is that generalization to unseen instances does not matter.\n\n2. The paper is written in crystal clarity with information in every details about the methods, the experiments and the results. The authors discuss about the results of different methods for each setting. Limitations and outlook to future work are also faithfully discussed.\n\n3. Improvements without refinements look impressive on all settings. Improvements after refinements still look great on the first three settings.\n\nWEAKNESSES:\n1. For the 4th setting, Incompressible Navier-Stokes, the proposed reparameterization method gives much higher mean losses than BFGS despite the fact that the majority of problems actually improve over BFGS. Could the authors elaborate more on the potential reasons specific to this experiment setting?\n\n2. Since the mean losses could change with different dataset sizes because of the varying instance difficulty, would it be a better presentation of results with relative error or relative loss? (also a relative improvement could be better for results like Figure 4 in the Appendix.\n\n3. 3~6 times more computational cost for the first three settings and up to 22 times for the fluids could be too high for the benefits achieved. This could be subjective but I hope the authors could provide some discussions or justification.\n\n4. Would the \"similarity\" requirement be too strict to make the proposed method practically useful? For example, in the wave packet localization setting, the parameters A and $\\sigma$ are fixed. Is there a practical application scenario that corresponds to this setting?\n\nQUESTIONS:\nSee the weaknesses part.\n\nLIMITATIONS:\nHigh computation cost and potential lack of practical applicability.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-10"}