A Novel Convergence Analysis for the Stochastic Proximal Point Algorithm

Abstract: In this paper, we study the stochastic proximal point algorithm (SPPA) for general empirical risk minimization (ERM) problems as well as deep learning problems. We present an efficient implementation of SPPA with minor modification for different problem definitions and we observe that efficiently implemented SPPA has faster and more stable convergence than the celebrated stochastic gradient descent (SGD) algorithm, and its many variations, for both convex and non-convex problems. Due to the fact that per-iteration update of SPPA is defined abstractly and has long been considered expensive, its convergence proof has not been well-studied until recently. In this paper, we close the theoretical gap by providing its convergence for convex problems. Our proof technique is different from some of the recent attempts. As a result, we present a surprising result that SPPA for convex problems may converge \emph{arbitrarily fast}, depending on how the step sizes are chosen. As a second contribution, we also show that for some of the canonical ERM problems and deep learning problems, each iteration of SPPA can be efficiently calculated either in closed form or closed to closed form via bisection---the resulting complexity is exactly the same as that of SGD. Real data experiments showcase its effectiveness in terms of convergence compared to SGD and its variants.

Reviews and author responses:

# === Update ===

I have read the authors reviews as well as the response by the authors.

I'd like to thank the authors for engaging in a discussion around Lemma 1. It is unfortunate that the result does not appear to hold even in a restricted setting.

I will maintain my score. I wish the authors the best in their future work.

# ======

This submission investigates the convergence rate of the stochastic proximal point algorithm (SPPA) for convex, finite-sum functions.
The authors derive new convergence rates for SPPA under the assumption that each function in the finite sum is convex (or strongly convex).
Key to these derivations is a lemma characterizing the probabilities of sampling a sub-function form a time-reversible Markov chain. 
In addition to these results, the authors derive efficient implementations of SPPA for a variety of loss functions, generally by reducing the problem to one-dimensional root-finding.
The paper concludes with experiments comparing SPPA to SGD, Adam, and other baseline stochastic optimization algorithms for (i) convex optimization and (ii) training CNNs and ResNets on MNIST and CIFAR-10, respectively.
The experiments show that SPPA converge faster than stochastic gradient methods and have a similar computational cost. ## Writing

The writing quality is very good and the paper is polished.

## Theory

The key lemma (Lemma 1) underpinning the theoretical contributions of the submission is incorrect.
As a result, Proposition 1, Theorem 1, Theorem 2, Corollary 1, Corollary 2, and Proposition 2 do not hold as stated. 
Note that I did not check the other proofs for separate/additional bugs.

### Lemma 1:

The proof of Lemma 1 is incorrect and the Lemma cannot hold as stated.

**Counter Example**: To see that Lemma 1 cannot hold, consider the following scenario: $n = 2$, $d = 1$, $f_1(x) = \\frac{1}{2} (x-1)^2$, and 
$$ f_2(x) = 
\\begin{cases} 
    \\frac{1}{2} x^2 & \\mbox{if $x \leq 0$}\\\\
    \\infty          & \\mbox{if $x > 0$.}
\\end{cases}
$$
Note that both functions are convex, closed, and coercive.
Now, suppose that $\\lambda_t = \\lambda = 1$ for all $t$ and consider when $x_{t+1} = \frac{1}{2}$.
We have that

$$
x_{t+1} = \\text{argmin}_{x} f_1(x) + \\frac{1}{2} (x - 0)^2,
$$

implying that $x_{t} = 0$ and sampling $i_t = 1$ gives $x_{t+1}$ with non-zero probability.
On the other hand, $x_{t+1}$ cannot be obtained for any value of $x_t$ when $i_t = 2$. 
We deduce $P(i_t = 1 \\mid x_{t+1} = 1)$, which is not equal to $P(i_t = 1 \\mid i_{t+1}) = P(i_t = 1)$, where the equality here follows from independence of $i_t$ and $i_{t+1}$
This shows that Lemma 1 cannot hold in the general sampling model for convex $f_i$. 

**Flaw in the Proof**: the mistake in the proof is the claim that $P(i_t \mid x_{t+1}) = P(i_t \mid i_{t+1})$.
Yes, it is true that $P(i_t \mid w_t) = P(i_t \mid i_{t-1})$ --- this is because $i_t$, $w_t$, and $i_{t_1}$ are all _independent_.
The relationship cannot be generalized to $w_{t+1}$ because it is _dependent_ on $i_t$, while $i_t$ and $i_{t+1}$ are independent.

I have used $f_i$ with nested effective domains in my counter example because the construction is easy and this makes the issue very clear;
however, I am certain more general constructions with domain $\\mathbb{R}^d$ can be constructed where the dependence between $x_{t+1}$ and $i_t$ breaks Lemma 1. 


## Experiments

The experiments in this work cannot be evaluated properly because key experimental details are not reported. 

**Step-sizes**: the authors report a grid of step-sizes which they considered for SPPA, but I did not find a discussion of how step-sizes were chosen for the baseline optimization methods.
The performance of stochastic gradient methods like SGD and Adam is highly dependent on the choice of step-size; 
a fair comparison to SPPA requires conducting a hyper-parameter search to find "reasonable" step-sizes for all baseline methods. 
It is not clear that the authors have done this. 
Figures 1 and 3 show surprisingly poor performance for Adam/AdaBelief, indicating the step-size has not been tuned properly. 
Similarly, Figures 5, 9, and 7 show what looks like insufficient tuning of SGD's step-size. 
As a result, it is not clear that the experiments in this submission meet a minimum threshold of "fairness".

**Repeats**: it is critical that sufficient repeats be performed for stochastic optimization algorithms and that spread information (e.g. interquartile range, standard deviation) be reported in figures.
The authors do an admirable job of this task for the convex optimization problems considered, but don't report any spread information in Figures 5-10, which consider training neural networks.
This information must be included for readers to judge whether or not the experimental results are statistically meaningful. 

## Minor Comments

Section 1, Paragraph 2:
- Please cite the original work by Robbins and Monro [1] on stochastic gradient descent.
- Note that the first stochastic methods to use some form of variance reduction to speed-up convergence were SAG [2] and SDCA [3]. SVRG and SAGA are follow-up works. 

Section 1, Paragraph 3:
- Incremental methods for finite-sums are (and should be) distinguished from stochastic methods by the order of updates. Incremental methods use a deterministic, cyclic order of updates. See e.g. Gurbuzbalaban et al. [4]. 
- "SPPA has an abstract per-iteration update rule and it acquires more information from the problem than solely the first order derivatives" --- I don't agree with this statement. The proximal-point iteration solves the inclusion 
$$0 \\in \\partial f(x_{k+1}) + \\frac{1}{\\lambda_k}(x_{k+1} - x_k) \\iff x_{k+1} = x_{k} - \\lambda_k v_k, $$
where $v_k \in \\partial f(x_{k+1})$.
This is an "implicit" gradient step which requires only first order information.
Solving the inclusion (e.g. computing $v_k$) can be achieved with higher order information, but the iteration itself requires only (sub)-gradients. 
- "While some ‘accelerated’ versions of SGD demand additional time or space overhead to go over the entire data set..." --- you should cite Katyusha (Allen-Zhu, 2017) or related methods for this claim. 
 
All Figures:
- The lines in the plots should be thicker. All font-sizes should be larger to improve readability. 
- Line markers are useful for distinguishing lines and are helpful for color-blind readers. 


## References

[1] Herbert Robbins, Sutton Monro "A Stochastic Approximation Method," The Annals of Mathematical Statistics, Ann. Math. Statist. 22(3), 400-407, (September, 1951)

[2] Schmidt, Mark, Nicolas Le Roux, and Francis Bach. "Minimizing finite sums with the stochastic average gradient." Mathematical Programming 162.1-2 (2017): 83-112. 

[3] Shalev-Shwartz, Shai, and Tong Zhang. "Stochastic dual coordinate ascent methods for regularized loss minimization." Journal of Machine Learning Research 14.2 (2013).

[4] Gurbuzbalaban, M., Asu Ozdaglar, and Pablo A. Parrilo. "Convergence rate of incremental gradient and incremental newton methods." SIAM Journal on Optimization 29.4 (2019): 2542-2565. Stochastic proximal-point methods are an interesting class of algorithms that have recently gained attention as a robust alternative to stochastic gradient methods.
This topic is highly relevant to ICLR and new convergence rates (with minimal assumptions) are of interest to the broader "optimization for machine learning" community. 

Unfortunately, I believe the theoretical derivations in this submission are not sound as given and it is not clear that the convergence results hold as stated.
The issue is that a key lemma (Lemma 1) does not hold (see "Theory" above for a counter-example).
Moreover, I do not think that proofs can be corrected easily, as the entire approach is invalid without Lemma 1.

The experiments in the submission, while interesting, are also flawed.
In particular, insufficient information is provided to verify that the comparisons to baseline stochastic optimizers were conducted fairly.
I am concerned that baseline methods were not tuned properly and so do not represent a meaningful comparison.

Given the above, I recommend that the paper be rejected.

This paper provides a novel elegant proof of convergence for the Stochastic Proximal Point Algorithm, under the assumption of convexity. It proceeds by surprisingly conditioning on future iterates, reversing usual convergence proofs. Experiments show that the SPPA is competitive, including in nonconvex settings such as training neural networks for image classification. The theoretical contributions are solid, to the best of my knowledge. The technique of reversing the conditioning in the expectations allows for a proof which is similar to the deterministic proximal point algorithm. 

EDIT: After reading the other reviews, I'm convinced that Lemma 1 is incorrect, therefore invalidating the proposed proof.

Comments: 
- What is the point of section 2.3 ? This method is not used in the rest of the paper, and no convergence theory is given. I suggest that the authors develop this theory in a separate paper, since it would definitely be interesting.
- In numerical experiments, some parameters are non standard. Why use such small batch sizes for CIFAR-10? 
- For more than 1 sample, can a closed form be found for the least squares/logistic loss on linear models? Does the algorithm converge in theory when the proximal problem is solved approximately? 
- The appendix says that 10 steps of L-BFGS are used in the inner loop. How good an approximation is that? Could the authors plot the suboptimality for the inner problem wrt number of L-BFGS iterations at different stages of model training?
- The plots for neural networks have number of iterations in the x-axis. Could the authors provide plots with time in the x-axis? Since the iterations are much more costly (10 iterations of L-BFGS), and would be difficult to make more efficient, this might significantly alter presented results in the nonconvex setting.
- Could the authors provide experiments with step size schemes which, in theory, would provide acceleration? 
- Could the authors provide log-log suboptimality plots to verify the 1/t^k rates, in the convex setting?
- It seems like the models are not yet converged, even for convex problems. Can the authors plot suboptimalities, rather than loss functions? Perhaps the norm of the gradient along iterations?
- There are many typos; attentive proof-reading would help readability.
- The second contribution is almost exclusively provided in appendix. Replacing section 2.3 with a discussion of how to compute the proximal update would be helpful. I'm on the fence. I need to see better empirical validation to recommend acceptance, but am willing to change my score if the authors provide it.

The submitted paper proposes a new approach to study the convergence of stochastic proximal point algorithm, which requires also no stronger assumptions than existing results. Numerical experiments also demonstrate the effectiveness of the algorithms.  The main result of the paper is Lemma 1, which shows backward conditional probability. The rest of the proof is not surprising since it is quite close to the standard deterministic proximal point algorithm. The extension to the Douglas--Rachford splitting case is also interesting, however, i'm curious about what type of convergence guarantees can be obtained for this method. For example, under what conditions can we have convergence rate for the objective function value?

edit: After reading the other reviews, I think the result needs further polish.  The submitted work proposed a new convergence analysis for the stochastic proximal point algorithm, numerical experiments from various problems demonstrated the advantage of the algorithm. 

 Assume $x_0 \\sim \\mathcal{N}(0,1)$ and take $f_1$, $f_2$ as I gave above.
Then, 
$$P(x_1 \\leq -1 \\mid i_0 = 1) = P(x_0 \\leq -3),$$ 
since the proximal update with $i_0 = 1$ ensures $x_1 > -1$ whenever $x_0 > -3$. 
Similarly, 
$$P(x_1 \\leq -1 \\mid i_1 = 2) = P(x_0 \\leq 0),$$ 
also holds by definition of the proximal update.

Applying Bayes' rule gives
$$P(i_1 = 1 \\mid x_1 \\leq -1) = \\frac{P(x_0 \\leq -3)}{P(x_0 \\leq -3) + P(x_0 \\leq 0)},$$
and
$$P(i_1 = 2 \\mid x_1 \\leq -1) = \\frac{P(x_0 \\leq 0)}{P(x_0 \\leq -3) + P(x_0 \\leq 0)}.$$

Independence does not hold has you have claimed. If you are convinced otherwise, I ask that you please provide a rigorous proof. This is the last construction I will give to convince you that Lemma 1 cannot be valid in a general probability model for stochastic proximal point.


 I can see the argument if we know $x_0=0$. The statement that "this argument is easily extended to randomized initialization" is bit of a jump for me. In fact, our understanding is that if we know nothing about the past iterations (absolutely nothing, including the initialization), then there is no useful information to be used to deduce which $f_i$ is more likely to be picked in the previous iteration.

Again, we really appreciate the discussion as it greatly strengthen's our understanding about the subject, and we look forward to your reply.

 No, Lemma 1 can't be fixed. The result simply does not hold. Try writing down the directed graphical model for the dynamical system --- you will see that $w_{t+1}$ and $i_t$ are dependent even without observing $w_t$. 

Here is a counter example which does not rely on $f_2$ with effective domain strictly included in $\\mathbb{R}$:

Let $f_1 = x^2$ and $f_2 = x$. Take $\\lambda_k = \\lambda = 1$. The stochastic update for $f_2$ is

$$
x_{k+1} = \\text{argmin}_x \\{ x + \\frac{1}{2} (x - x_k)^2 \\}
= x_k - 1.
$$

Now, suppose we initialize deterministically at $x_0 = 0$ and that $x_{k+1} = -(k+1)$. Then we must have
$$
P(i_k = 2 \\mid x_{k+1} = -(k + 1)) = 1,
$$
since $P(x_{k+1} = -(k + 1) \\mid i_{k} = 1) = 0$ by construction. 
In fact, $i_l = 2$ almost surely for every $l \\in [k]$, otherwise $x_{k+1} = -(k+1)$ cannot be "reached" from the initialization.

Note that this argument is easily extended to randomized initialization, higher dimensions, etc. The important part is creating a "reachable" ball that requires sampling a specific function repeatedly for the iterates to escape.

 Thank you very much for the counter-example. This is very interesting.

Essentially, the domain of $f_2$ is $(-\infty,0]$, not $\mathbb{R}$. This is why $\Pr[i_t=2|x_t=1/2]=0$, because $x_t=1/2$ is not in the domain of $f_2$. Do you think Lemma 1 could be fixed by adding the assumption that the domain of all $f_i$ are $\mathbb{R}^d$, meaning no function can take value $+\infty$ anywhere in $\mathbb{R}^d$? 

''The relationship cannot be generalized to $w_{t+1}$ because it is dependent on $i_t$, while $i_t$ and $i_{t+1}$ are independent.''

Our reasoning is that the dependency between $w_{t+1}$ and $i_t$ is through $w_t$, so without knowing the value $w_t$, $w_{t+1}$ and $i_t$ are independent (maybe with some additional assumptions on the functions, now that we've seen the nice counter-example you gave).


This paper shows the arbitrary convergence rate of the stochastic proximal point algorithm(SPPA), under the convex assumption of the objective function. The authors use a novel approach to analyze SPPA. Numerical results show the efficiency of SPPA on some realistic data-sets.

 Major comments:
1. Theorem 1 suggests the arbitrary convergence rate of SPPA. Rockafellar (1976) showed this result for the deterministic PPA. Some discussions are necessary.	
2. In the proof of Lemma 1, I don't understand why $p(i_t|w_{t+1})=p(i_t |i_{t+1})$. More details should be given.
3. In the numerical experiments, the authors use the default settings of SGD, Adam, etc, but select  the parameter $\lambda$   for SPPA. This would be unfair. Moreover, though the stepsize of SGD or the adaptive variants of SGD is difficult to tune, there are a lot of works studying how to fine-tune the stepsize of SGD, e.g., [Wilson, et. al., 2017].   Especially when training deep neural networks(DNNs), people will not use a constant stepsize nor a diminishing stepsize like $1/t.$ The experiments suggest the SPPA outperforms SGD when training DNNs, but the settings of SGD or Adam should be modified. Otherwise, the experiments will mislead the readers and weaken the presentation in this paper. 

4. In SPPA, the subproblem is solved by bisection in this paper.
    But the complexity of the bisection method is not discussed.
    Though   Theorem 1 shows the arbitrary rate, when $\lambda$ increases to infinity, the subproblem will be  extremely ill-conditioned, which results in slower convergence of the bisection method. Hence, the authors should also consider the complexity of the subproblem.

Detailed comments:

1.Page 3:  The PPA is not $w_{t+1} = w_t - \lambda_t \nabla F(w_{t+1})$, where the last term should be the subgradient, since this paper also includes nonsmooth functions.

2. The reference of the equation number should be consistent, e.g., in the proof of Theorem 1, the equation 2 should be equation (2).

3. In the figures, the labels are not explained, e.g., SPPA-1/t, SPPA-AGD.

4. In Page 5: "As we will see in
the experiment section, a constant step size rule still gives the best performance in most cases." I cannot find the direct comparison between different $\lambda_t$ for SPPA.  

[Reference]:
Wilson A C, Roelofs R, Stern M, et al. The marginal value of adaptive gradient methods in machine learning[C]//Proceedings of the 31st International Conference on Neural Information Processing Systems. 2017: 4151-4161. This convergence result of SPPA is interesting. But the numerical experiments weaken the presentation.

Meta-review: The paper considers the stochastic proximal point algorithm, the main contribution is a convergence proof (in addition to arguing that it is a practical algorithm for ERM problems).  However, reviewer EvJ6 and xJB8 pointed out a fatal flaw in Lemma 1 that affects all the downstream results, and the other two reviewers and myself agree that this invalidates the core results. Reviewer EvJ6 gives extensive examples of how the lemma cannot be simply fixed. So while this is an interesting stochastic extension of the famous deterministic method and may warrant further study, this paper doesn't yet provide a rigorous convergence proof.

Decision: Reject