Paper: Lion Secretly Solves a Constrained Optimization: As Lyapunov Predicts (ICLR-2024)

Reviewer comment (Theory):
I just have a question regarding the impact of assuming differentiability of $\mathcal{K}$ in the continuous-time analysis (see question below).

Author response:
Dear Reviewer, we thank you for your valuable comments. We have addressed the typos and further refined the draft. Below are answers to your questions on the differentiability of $\mathcal{K}$.

1) The main difference between the continuous and discrete-time results is that $\Delta(x,m)$, which is a measure of stationarity of the problem, decreases with $O(1/t)$ rate in the continuous time, while $O(1/(\epsilon T) + \epsilon)$ in the discrete time, where $T$ denotes the number of steps. Hence, the continuous-time and discrete-time results are parallel and consistent, in the sense that we recover the continuous-time result when taking the step size $\epsilon$ to zero (with $t = \epsilon T$).

Reviewer comment (Theory):
In the discrete-time analysis, $\mathcal{K}$ is allowed not to be differentiable, but in this case the result appears more difficult to interpret: the only result provided is Thm 4.1, which is about $H(x_{t+1}, m_{t+1}) - H(x_{t}, m_t)$ (or, alternatively, about the sum of deltas), and unless mistaken, there is no result analogous to convergence to a stationary point of the constrained objective. Theorem B.9 in the appendix addresses only the first phase of the dynamics.

Author response:
Dear Reviewer, we thank you for your valuable comments. We have addressed the typos and further refined the draft. Below are answers to your questions on the differentiability of $\mathcal{K}$.

1) The main difference between the continuous and discrete-time results is that $\Delta(x,m)$, which is a measure of stationarity of the problem, decreases with $O(1/t)$ rate in the continuous time, while $O(1/(\epsilon T) + \epsilon)$ in the discrete time, where $T$ denotes the number of steps. Hence, the continuous-time and discrete-time results are parallel and consistent, in the sense that we recover the continuous-time result when taking the step size $\epsilon$ to zero (with $t = \epsilon T$).

3) If we assume $\mathcal{K}$ is smooth (with bounded Hessian), it is possible to improve the discrete-time rate to $O(1/(\epsilon T))$. Hence, the impact of the non-differentiability is the $O(\epsilon)$ term, which suggests that the algorithm converges up to an $\epsilon$ accuracy. This is a typical phenomenon in optimization with non-smooth objectives (like sub-gradient descent) or non-smooth update (like signed GD). Because in practice the step size is small or decaying, this may not have a big impact on practical performance.

Reviewer comment (Theory):
It is therefore hard to verify from such theorem whether the conclusions of the discrete-time case (which corresponds to the algorithm used in practice) will actually match the conclusions from the continuous-time case.

Author response:
Dear Reviewer, we thank you for your valuable comments. We have addressed the typos and further refined the draft. Below are answers to your questions on the differentiability of $\mathcal{K}$.

1) The main difference between the continuous and discrete-time results is that $\Delta(x,m)$, which is a measure of stationarity of the problem, decreases with $O(1/t)$ rate in the continuous time, while $O(1/(\epsilon T) + \epsilon)$ in the discrete time, where $T$ denotes the number of steps. Hence, the continuous-time and discrete-time results are parallel and consistent, in the sense that we recover the continuous-time result when taking the step size $\epsilon$ to zero (with $t = \epsilon T$).

3) If we assume $\mathcal{K}$ is smooth (with bounded Hessian), it is possible to improve the discrete-time rate to $O(1/(\epsilon T))$. Hence, the impact of the non-differentiability is the $O(\epsilon)$ term, which suggests that the algorithm converges up to an $\epsilon$ accuracy. This is a typical phenomenon in optimization with non-smooth objectives (like sub-gradient descent) or non-smooth update (like signed GD). Because in practice the step size is small or decaying, this may not have a big impact on practical performance.

Reviewer comment (Theory):
The two theoretical results are interesting in themselves, but it is unclear to what extent they can successfully predict the actual behaviour of the Lion-$\mathcal{K}$ method, since Lyapunov analyses require $\mathcal{K}$ to be differentiable.

Author response:
Dear Reviewer, we thank you for your valuable comments. We have addressed the typos and further refined the draft. Below are answers to your questions on the differentiability of $\mathcal{K}$.

4) The requirement for differentiability in our statement of continuous-time results is mainly for the simplicity of the presentation. We can state a continuous-time result for non-smooth $\mathcal{K}$ using non-smooth versions of Lyapunov theory, but we would need to recall sophisticated technical tools such as Filippov's differential inclusion and Clarke's generalized gradient. See e.g., *Daniel Shevitz and Brad Padent, Lyapunov Stability Theory of Nonsmooth Systems, 1993.*

Reviewer comment (Theory):
When $\mathcal{K}$ is not differentiable (e.g., $\ell_1$ norm considered in the paper), analyses can break, and this may lead to behaviour of algorithms being quite different for differentiable vs. non-differentiable $\mathcal{K}$.

Author response:
Dear Reviewer, we thank you for your valuable comments. We have addressed the typos and further refined the draft. Below are answers to your questions on the differentiability of $\mathcal{K}$.

3) If we assume $\mathcal{K}$ is smooth (with bounded Hessian), it is possible to improve the discrete-time rate to $O(1/(\epsilon T))$. Hence, the impact of the non-differentiability is the $O(\epsilon)$ term, which suggests that the algorithm converges up to an $\epsilon$ accuracy. This is a typical phenomenon in optimization with non-smooth objectives (like sub-gradient descent) or non-smooth update (like signed GD). Because in practice the step size is small or decaying, this may not have a big impact on practical performance.

4) The requirement for differentiability in our statement of continuous-time results is mainly for the simplicity of the presentation. We can state a continuous-time result for non-smooth $\mathcal{K}$ using non-smooth versions of Lyapunov theory, but we would need to recall sophisticated technical tools such as Filippov's differential inclusion and Clarke's generalized gradient. See e.g., *Daniel Shevitz and Brad Padent, Lyapunov Stability Theory of Nonsmooth Systems, 1993.*

Reviewer comment (Theory):
In the simple case of mirror descent (section 3.1), writing the discrete update naively with $\mathcal{K}$ as $\ell_1$ norm gives an update $x_{t+1} \leftarrow \text{sign}(x_t - \nabla f(x_t))$, which typically does not converge to $\min_x f(x)$, as iterates remain restricted to $\\{+1, 0, -1\\}^{d}$; the same problem would happen for the continuous case mirror descent.

Author response:
Dear Reviewer, we thank you for your valuable comments. We have addressed the typos and further refined the draft. Below are answers to your questions on the differentiability of $\mathcal{K}$.

5) The mirror descent with $\ell_1$ norm reduces to Signed GD: $x_{t+1} = x_t - \epsilon sign(\nabla f(x_t))$, which is different from $x_{t+1} = sign(x_t - \nabla f(x_t))$. Hence it is the update (not the value of $x_t$) that is constrained in $[-1,1]$. With standard argument, one can show that $\min_{s\leq T}\|\nabla f(x_s)\|_1 = O(1/(\epsilon T) + \epsilon)$, which is consistent with our result.

Reviewer comment (Theory):
It would be good to elaborate more (and/or give more references) in the continuous case on why the analysis of Lion-$\mathcal{K}$ would still be valid or almost valid for non-differentiable $\mathcal{K}$, and/or in the discrete case to provide more details or a more intuitive reading of the result in Theorem 4.1, to confirm that Lion-$\mathcal{K}$ is indeed performing constrained/penalized minimization without suffering from issues similar to mirror descent.

Author response:
Dear Reviewer, we thank you for your valuable comments. We have addressed the typos and further refined the draft. Below are answers to your questions on the differentiability of $\mathcal{K}$.

3) If we assume $\mathcal{K}$ is smooth (with bounded Hessian), it is possible to improve the discrete-time rate to $O(1/(\epsilon T))$. Hence, the impact of the non-differentiability is the $O(\epsilon)$ term, which suggests that the algorithm converges up to an $\epsilon$ accuracy. This is a typical phenomenon in optimization with non-smooth objectives (like sub-gradient descent) or non-smooth update (like signed GD). Because in practice the step size is small or decaying, this may not have a big impact on practical performance.

4) The requirement for differentiability in our statement of continuous-time results is mainly for the simplicity of the presentation. We can state a continuous-time result for non-smooth $\mathcal{K}$ using non-smooth versions of Lyapunov theory, but we would need to recall sophisticated technical tools such as Filippov's differential inclusion and Clarke's generalized gradient. See e.g., *Daniel Shevitz and Brad Padent, Lyapunov Stability Theory of Nonsmooth Systems, 1993.*

Reviewer comment (Experiments):
This manuscript appears to be incomplete. Notably, there is a lack of discussion on related works, and the empirical results presented in Figures 1-3 are neither interpreted nor mentioned, even though there seems to be sufficient space for such discussions. I strongly recommend that the authors enhance the completeness and rigor of this paper via:  
- Interpretation of the empirical results,  
- Ablation studies on the key components of Lion to justify the theory,  
- Empirical validation of (at least some) algorithms listed in Table 2.

Author response:
### 1. Interpretation of empirical results

**Ans:** We have added further interpretations of the empirical results in the revised draft as mentioned in general response. The following is the interpretation the results in Figure 1-3. 

- **Fig 1:** The trajectories of Lion with weight decay $\lambda = 1.5$ and $\lambda = 0.6$ and corresponding loss curve and $||x||_\infty$ curve. Observation: the converging points of those 2 trajectories are within the feasible region (specifically on the boundary of the feasible region). We use this figure to verify 2 things: 
    1. Weight decay introduces a constraint effect.
    2. 2 phases of training dynamics: (first, exponentially decay to feasible region, then minimize objective function)

- **Fig 2:** We use this figure to show phase I is very fast and the constraint actually exists. The histograms of network parameters trained by Lion with $\lambda = 10$, the initialization of network parameters is Kaiming normal, from the leftmost plot, in the beginning, not all params satisfy the constraint (red vertical lines: $x\leq 1 / \lambda$). After 150 iterations of training, the constraint is almost satisfied, and at the 250-th iteration, the constraint is fully satisfied.

- **Fig 3:** This figure is a stack version of Figure 2 but with different initialization methods and weight decay ($\lambda$) levels. Stacked histograms of network parameters in ResNet on the CIFAR-10 dataset across iterations. We stacked the histogram plots from iteration 0 (back) to 190 (front). With weight decay $\lambda$ equals to 0, there is no constraint effect.

### 2. Empirical validation of algorithms listed in Table 2

**Ans:** As we mentioned in the general response, we have **Experiment II (Toy Example with Different $\mathcal{K}$):**
- **Section & Figure Reference:** Section 5, Page 10 -11, Figure 1 and Figure 5.
- **Summary:** We examine Lion's behavior with different $\mathcal{K}$ values in a similar toy example from Figure 1. This demonstrates that the observed behavior aligns well with our theoretical predictions.

Reviewer comment (Evaluation):
This paper does not connect the superior performance of Lion or other Lion-$\mathcal{K}$ optimizers to their main finding. How does this additional constraint help optimization or generalization?

Author response:
### 3. Connect the superior performance of Lion or other Lion-K optimizers to their main finding.

**Ans:** We have added experiments mentioned in the general response: 

**Experiment I (The Implication of Weight Decay):**
- **Section & Figure Reference:** Page 4-5, Figure 1 and Figure 4.
- **Summary:** We discuss the optimal weight decay $\lambda$ leading to fast convergence without sacrificing training loss. We demonstrate that a larger $\lambda$ generally accelerates learning, but performance drops when $\lambda$ exceeds a critical value $\lambda_0$. This finding underscores the importance of tuning $\lambda$ for Lion.

**Experiment III (Lion-$\ell_p$ on ImageNet and Language Modeling):**
- **Section & Figure Reference:** Section 5, Figure 7.
- **Summary:** We evaluate Lion-$\ell_p$ with $p \in \{1.0, 1.2, 1.4, 1.6, 1.8, 2.0\}$ on ImageNet and Language Modeling tasks using ResNet-50, ViT-16B, and GPT-2. The results are intriguing: Lion-$\ell_p$ with smaller $p$ values performs better for ViT and GPT-2, but the opposite trend is observed with ResNet-50 on ImageNet. This suggests a potential architecture-dependence of the optimal $\mathcal{K}$, marking an interesting avenue for future research.

### 4. How does this additional constraint help optimization or generalization?

**Ans:** There is a trade-off effect. Adding a weight decay term introduces a constraint. Adding the constraint force drives the solution faster towards the region.

### 5. Why Should Lion Decay Weight? 

**Ans:** From the analysis above, the role of weight decay $\lambda$ in Lion is two-fold: 

a) It alternates the solution if $\lambda$ is large and the constraint $||x||_\infty$ $\leq$ $1/\lambda$ 

is strong enough to exclude the unconstrained minimum $x_{unc}^*$ of $f(x)$. This may improve the generalization and stability of the solution while reducing the training loss (Figure 4 and Figure 7).

b) If $\lambda$ is sufficiently small to include the unconstrained minimum $x_{\text{unc}}^*$ in the constrained set, it does not alter the final solution. In this case, the main role of weight decay is to speed up the convergence because the Phase 1 brings the solution into the constrained set with a linear rate. Hence, the ideal choice of $\lambda$ is $\lambda$ $=$ $1/||x_{unc}^*||_\infty$

Reviewer comment (Experiments):
To demonstrate efficiency in the conclusion, authors didn't compare with other methods. At least, it is necessary to include other benchmark algorithms in empirical evaluations.

Author response:
We agree on the importance of comparing our method with benchmark algorithms. In the final manuscript, we will include a comparative analysis with established benchmarks to demonstrate the efficiency and effectiveness of our approach in a broader context.

Reviewer comment (Evaluation):
Can you comment the possible convergence on lion-K vs interesting decomposition (12)? It seems they may both work well.

Author response:
We think that having a Lyapunov function posits a necessary condition for an update rule to be a valid optimizer. It is then up to empirical results and insights to determine whether the algorithm performs well on the different problems of interest.

Reviewer comment (Experiments):
The current analyses apply to only algorithms with full gradient. It would be interesting to see results for Lion with stochastic gradients.

Author response:
The extension to the stochastic gradient case is straightforward given our results with the full gradient. We provide a new theorem in the Appendix (Check Theorem B.16 at page 30).

Reviewer comment (Theory):
It would be helpful if the authors can discuss and comment on the implications of Theorem 4.1. Also, it is worth explaining why it is necessary to use a different implicit scheme.

Author response:
We added comments on Theorem 4.1. The main implication is that we get parallel results on the two-phase convergence for the discrete-time case. The main difference is that the rolling average of $\Delta^1_t$ and $\Delta^2_t$ converges with $O(1/(\epsilon t) + \epsilon)$, where epsilon is the step size, rather than $O(1/t)$ in the continuous case.

   The implicit form fits nicer with our proof as we tried to get a Lyapunov form in the discrete-time. We think it is mainly for technical reasons of the proof and as we explain in Section 4, the implicit form is equivalent to the explicit form with a different step size.

Reviewer comment (Writing):
Typo:
- In the first sentence in the paragraph above Figure 4, "difference $\mathcal{K}$" -> "different $\mathcal{K}$"
- The second paragraph above Lemma 2.1, "a systematic introduce to ..." -> " a systematic introduction to ..."

Author response:
Dear reviewer, thank you very much for your comments. We have addressed the typos and further refined the draft.

Reviewer comment (Theory):
In Section 3.1, the authors discuss the connection with existing algorithms. Do the corresponding convergence results also reduce to those classic ones?

Author response:
One convergence rate result implied from our framework is that $\Delta_1$ and $\Delta_2$ converge with an $O(1/t)$ rate. This is consistent with the results of classical algorithms on non-convex functions that the gradient (and momentum) norm converges to $O(1/t).$

   It is also straightforward to obtain stronger convergence rates by placing stronger assumptions on $f$ (e.g., convexity or strong convexity), but since $f$ is non-convex in deep learning, we defer such derivations to future works.