We thank a lot for the reviewer's quick feedback and generosity in raising the score. As far as we can understand, the reviewer still has two questions:
- whether there is an expectation form in $\nabla f(x)$;
- the result is not full high probability.
We then answer as follows:
- **The loss function $f$ is not necessarily the expected form in Eq. (1).** We indeed only require $f$ to be differentiable and smooth. We will clarify this in the revised version. In this sense, $\nabla f(x)$ also not necessarily includes the expectation.
- **$\\|\nabla f(x)\\|^2 \le \varepsilon$ implies that $x$ is near-stationary, which is a standard measure in non-convex smooth optimization.** Our high probability result shows that $\frac{1}{T}\sum_{t=1}^T \\|\nabla f(x_t)\\|^2\le \varepsilon$, indicating that there exists at least one near-stationary point $x_{\tau},\tau \in [T]$.
- As far as we could understand, (please correct me if I was wrong), the fully high probability you mention is to derive a high probability bound for $\frac{1}{T}\sum_{t=1}^T \\|g(x_t,z_t)\\|^2$. We think that this could be meaningful. However, given the second point, the high probability bound related to $\frac{1}{T}\sum_{t=1}^T \\|\nabla f(x_t)\\|^2$ could implies a near-sationary point while the bound for $\frac{1}{T}\sum_{t=1}^T \\|g(x_t,z_t)\\|^2$ may not, which we will show through an example at the end.
- In addition, since $\frac{1}{T}\sum_{t=1}^T \\|\nabla f(x_t)\\|^2$ is an entirely random variable with respect to $z_1,\cdots,z_{T}$, **our result is fully high probability over random samples $z_1,\cdots,z_T$.** However, when Assumption (A2) holds, we could choose any random sample $z_t'$ that is independent from $z_t$ and obtain that $\nabla f(x_t) = \mathbb{E}_{z_t'}[g(x_t,z_t')]$. Hence, due to the additional expectation on $z_t'$, we agree that the obtained result is not fully high probability without restricting on $z_1,\cdots,z_T$.
We also present a simple example to help better understand the difference of two types of measure. Suppose that the loss function is $f(x)=x^4$, which does not include expectation in its gradient. We now consider the following SGD algorithm (with differences to Adam in step-size):
- Input $x_1$ and step-size $\eta > 0$;
- For $t = 1\cdots T:$
- Draw a random sample $z_t$ such that $\mathbb{E}[z_t \mid x_t]=0$ (e.g., a Gaussian white noise);
- Generate the stochastic gradient $g(x_t,z_t) = 4x_t^3+z_t$ (which is an unbiased estimator);
- Update the sequence as $x_{t+1} =x_{t} - \eta (4x_{t}^3 + z_t)$.
If we obtain that $\frac{1}{T}\sum_{t=1}^T (4x_{t}^3)^2 \le \varepsilon$ with high probability (as in our paper), which is fully high probability over $z_1,\cdots, z_{T}$, then we are able to deduce that at least one gradient $(4x_{\tau}^3)^2,\tau \in [T] $ is small enough and $x_{\tau}$ is near-stationary. However, if we obtain the fully high probability bound $\frac{1}{T}\sum_{t=1}^T (4x_{t}^3 + z_{t})^2 = \frac{1}{T}\sum_{t=1}^T [(4x_{t}^3)^2 + z_t^2 + 8x_t^3z_t]\le \varepsilon$, then we could not derive any stationary point from this inequality since the term $x_t^3 \cdot z_t$ may be non-zero.
We hope that the above answer and example could address your questions. If you have any further inquiries, we are delighted to discuss them with you.
Best regards,
Authors.