Venue: ICLR 2022 Conference

Reviews:

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: This work solves the entropy regularized mean-field model for a two-layer neural network. As the Fenchel–Rockafellar dual of that problem is a finite dimension problem, they consider an SDCA-type scheme to solve the dual problem, which still contains an integral term and they use the Langevin iterate procedure to solve the subproblem approximately. In contrast to the existing analysis for SGD in the mean-field regime, the new algorithm has a fixed length dual step size and has a fast linear convergence rate.
main_review: Comments:
This paper introduces techniques to directly solve the infinite-dimensional mean-field model for two-layer NN. It allows finite step size and improves the convergence rate to linear by solving the dual problem and assuming primal smoothness. My main comments are as follows:

* Proposition 1 needs more careful treatment. The authors claim the Fenchel–Rockafellar strong duality holds without formal justification. Note that, the strong duality holds when certain domain inclusion condition [Eq. 15.33, Bauschke 2011] is verified, which is not trivial for Eq. 2. One should not think the domain of l_i^* would be the whole R as the loss could be a logistic loss, which is not coercive (see [Theorem 11.8(d), R1].

* Intuitively, the update 4+5 in Algorithm 1 is very similar to running one-step SGD for the i-th data point. After computing step 4, the "weights" are updated by step 5. Is there any connection with a primal algorithm given this similarity?

* On stepsize for P-SDCA: it seems the step size length in the Algorithm is n*lambda_2, which is from the lower estimation in Lemma 1. However, on Page 40, additional experimental results, they said "we selected the optimal step size ... for P-SDCA and PDA". If n=100 as on Page 39, then the step sizes for Fig 3(a-c) would be actually different. This should be made clear.

---
[R1] Rockafellar, R.T. and Wets, R.J.B., 2009. Variational analysis (Vol. 317). Springer Science & Business Media.
summary_of_the_review: This work introduces a dual algorithm to solve the infinite-dimensional mean-field training problem. The update step is solved approximation by estimating an integral term in every iteration. The new technique is interesting and the new linear convergence rate improves the prior result.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: This paper presents an optimization alogrithm for mean-field shallow neural networks. The algorithm optimizes the parameter measure using a particle-based implementation of the stochastic dual coordinate ascent algorithm. The authors establish convergence results, namely an exponential convergence rate owing to the convexity of the dual problem.
main_review: This paper concerns optimizing mean-field neural networks. The problem is inherently challenging because the infinite width limit leads to a functional formulation of the representation and, consequently, the optimization must be carried out over the space of probability measures by solving a PDE. Particle-based approaches alleviate some of the challenges, but, as the authors note, the original mean-field papers deal with discretization simply by taking the limit $\Delta t\to 0$. Here, the authors propose to instead work directly with a space-time discretized dynamics.

The authors examine a particular variant of the mean-field objective function in which an entropic regularization term is added.  However, unlike these previous works, the authors take advantage of Fenchel duality to convert the problem into a maximization problem in $\mathbb{R}^n$ as opposed to an infimum over $p\in \mathcal{P}$. 

The proposed algorithm relies on first sampling parameters from a proxy probability measure $\rho^{(t)}(\theta)d\theta$ that is close in total variation distance from the target measure $\rho[g^{(t)}]$. How is it ensured that the proxy is actually close? This strikes me as a challenging thing to verify.

Because the algorithm requires sampling a high-dimensional distribution, this step is performed infrequently and the particles are reweighted iteratively. The sampling with ULA or MALA is the basis of the convergence results, using a LS inequality. Of course, the exponential decay rate will be unfavorable if the target distribution is metastable. There's no real discussion of the nature of the parameter distribution. Is there some reason we should expect that it's easy to sample?
summary_of_the_review: The paper outlines an algorithm and carefully analyzes it. The algorithm appears to be effective in a simple synthetic example. However, the assumption of efficient sampling of the parameter distribution requires substantial additional justification.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: The paper considers optimization of MF two-layer neural nets (the infinite-width version with entropic regularization). Specifically the paper proposes to do dual coordinate ascent, which allows for exponential convergence rate, together with a particle approximation scheme. The paper shows that in the discrete-time setting, an improved convergence rate is obtained (w.r.t. desired error $\epsilon_P$, the number of data points $n$, the regularization $\lambda_2$).
main_review: Thanks for the clarification. I'd like to keep my score.

======================

The paper is clear in its exposition, despite involving many parameters. In the context of MF two-layer neural nets, the idea is new. The proofs seem correct. I also appreciate that the authors have made some honest remark concerning $\lambda_2$, which is clearly a deficiency in this approach.

I do not have much to complain, though a few things could be considered:
- An explanation of the role of the re-weighting stage is needed. It is said to reduce the number of resampling times $\tilde{n}$ down from $n$ and hence brings down the complexity. It would be good to be more quantitative about this. Would the result be $O((n/\lambda_2\gamma)\log(1/\epsilon_P)$ as mentioned in page 7?
- It is said that Mei et al 2018 has the error growing exponentially with time. I think here the same also applies. One can take a different view that is to fix a target error to argue against this exponential dependency, but the same view applies to Mei et al 2018 as well.
summary_of_the_review: The paper proposes an interesting variant of dual coordinate ascent that achieves improvements in the convergence rate. There is not much to criticize since the paper has optimized the complexity to a good extent, except for the inherent difficulty with the approach that depends crucially on $\lambda_2$.
review_title: null

final_score: rating: 8: accept, good paper
review_content: summary_of_the_paper: The problem of training an infinitely wide neural network with one hidden layer can be written as an optimization problem in the space of probability distributions, where each atom in the distribution corresponds to a hidden neuron. The paper proposes to solve instead the dual problem (of an entropy regularized variant), which in the case of finite data is a finite-dimensional optimization problem.  The main contribution is the conception and analysis of an efficient dual coordinate ascent scheme to solve the dual, which in theory and practice requires fewer iterations to find an epsilon-accurate solution than existing methods. 

main_review: ### Strengths 
The paper is well-written, technically sound and I have no doubts about the correctness of the results (though I did not check the proofs in the appendix carefully). 

The presented dual particle optimization scheme appears to be novel and might be useful for other problems on probability density than mean-field neural networks. 

### Weaknesses
A) Motivation
* Since the main contribution is an algorithm (and its analysis) for Eq. (1), it would be great to motivate this problem a bit more for the uninitiated reader. For example, is it expected to be preferable over standard neural network training (eg. due to the regularization or convexity) or is the goal rather to gain a better understanding on standard training through it? In the experiments, the entropic regularization seems to help generalization, so perhaps this could be mentioned as a motivation (with additional explanations). 
 
B) Clarity
* It would be helpful if the paper could provide an explicit definition of the set P of probability density functions on R^{\tilde d} and a proof of Proposition 1. I am confused about this point, since in order for Theorem 15.23 of Bauschke et al. (2011) to be applicable in Proposition 1, it seems we have to be in a Hilbert space setting. But I am unsure whether the set of probability density functions considered here is a subset of a Hilbert space.  As a remark, in case the Hilbert space setting does not suffice, note that the original work by Rockafellar "Duality and stability in extremum problems involving convex functions", Pacific J. Math, 1967 already proves the duality theorem in the more general setting of paired topological vector spaces.

* It was a bit tedious to follow the proof of Lemma 1 (Lemma 3 in the appendix on page 19). I can see that the elements of the Hessian are bounded by one, but it would be nice if the paper could add a clarifying sentence how this yields Eq. (13).  Also, the sentence "Note that the bound (12) of the Hessian in the form of the covariance matrix regarding h_i" did not make sense to me, and Eq. (12) appears not to be a bound but rather an equality. 

C) Experimental evaluation 
* Would it be possible to run the method on a small binary classification example, e.g. on MNIST (2 vs 4)? The current experiment in Figure 1 suggests that the particle method generalizes much better than SGD (mean field / NTK), and it would be interesting to confirm this on a small real world dataset rather than in the teacher/student setup. Perhaps one could even show the learned 28x28 image-filters to illustrate the "feature learning" aspect of the mean-field algorithm (opposed to the NTK one). 

summary_of_the_review: The paper proposes and studies an interesting optimization method for problems in the space of probability density functions. The main application considered are mean-field neural networks, but  I can see how this algorithm will be very useful also for other problems over probability densities which are ubiquitous in ML.  

There are several smaller issues mostly regarding the technical clarity and motivation in the paper (see main review). If these are clarified/addressed in a revision, I am also willing to increase my score.
review_title: null

Meta-review: This is a solid paper and considers the problem of training a wide neural network with a single hidden layer. This can be framed as an optimization problem in the space of probability distributions with a suitable entropy regularization, where each atom in the distribution corresponds to a hidden neuron. The dual of this problem (for finite data) is a finite-dimensional optimization problem and the paper proposes a particle based coordinate ascent scheme.
The paper provides some convergence rate results. After the rebuttal, the authors have also included more experimental/numerical results.

The authors have answered the concerns raised by the reviewers and overall, the paper can be accepted:
The presented approach appears to be sufficiently novel and might be useful in other settings.
The presentation is clear and easy to follow for such a technical paper; the paper is well organized.
The limitations of the approach are clearly stated (dependence on the regularization parameter for entropy term that may be hard to select)
Decision: Accept (Poster)