Gradient Flossing: Improving Gradient Descent through Dynamic Control of Jacobians

Training recurrent neural networks (RNNs) remains a challenge due to the instability of gradients across long time horizons, which can lead to exploding and vanishing gradients. Recent research has linked these problems to the values of Lyapunov exponents for the forward-dynamics, which describe the growth or shrinkage of infinitesimal perturbations. Here, we propose gradient flossing, a novel approach to tackling gradient instability by pushing Lyapunov exponents of the forward dynamics toward zero during learning. We achieve this by regularizing Lyapunov exponents through backpropagation using differentiable linear algebra. This enables us to"floss"the gradients, stabilizing them and thus improving network training. We demonstrate that gradient flossing controls not only the gradient norm but also the condition number of the long-term Jacobian, facilitating multidimensional error feedback propagation. We find that applying gradient flossing prior to training enhances both the success rate and convergence speed for tasks involving long time horizons. For challenging tasks, we show that gradient flossing during training can further increase the time horizon that can be bridged by backpropagation through time. Moreover, we demonstrate the effectiveness of our approach on various RNN architectures and tasks of variable temporal complexity. Additionally, we provide a simple implementation of our gradient flossing algorithm that can be used in practice. Our results indicate that gradient flossing via regularizing Lyapunov exponents can significantly enhance the effectiveness of RNN training and mitigate the exploding and vanishing gradient problem.

Paper

Similar papers

Peer review

Reviewer 5Z2M4/10 · confidence 4/52023-07-04

Summary

The authors propose a novel approach called "gradient flossing" to tackle the instability of gradients in the training of recurrent neural networks (RNNs). The process of pushing the Lyapunov exponents toward zero is referred to as "flossing" the gradients. This stabilizes the gradients and improves network training.

Strengths

+ The authors show improvements in several toy sequence learning tasks: Delayed Copy Task, Temporal XOR Task, etc. + Idea of using Lyapunov exponents towards zero to mitigate vanishing and exploding gradients is well justified and works as expected on toy tasks.

Weaknesses

The lack of any experiments on non-simulated data is problematic, and the explanation given by the authors is not convincing. The paper's authors explicitly state that they deliberately did not use sequential MNIST or similar non-toy tasks commonly used to probe exploding/vanishing gradients. > We deliberately do not use sequential MNIST or similar toy tasks commonly used to probe exploding/vanishing gradients, because we want a task where the structure of long-range dependencies in the data is transparent and can be varied ad libitum While toy tasks are more transparent, real-world data is important to understand the proposed methods' potential failure modes and expected characteristics in *uncontrolled* settings. Avoiding running experiments on the basis that the results are more difficult to interpret than simulated data is not an acceptable excuse. Transparency in the work itself is hindered when experiments are consciously avoided.

Questions

Can the authors comment that such a technique of pushing the Lyapunov exponents to zero could limit the capacity to learn certain tasks? Maybe enabling the network to amplify or forget things is important for some task contexts, which has been a motivation for specialized architectures like the LSTM.

Rating

4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

Experimental limitation mentioned above.

Reviewer DGA66/10 · confidence 4/52023-07-06

Summary

The authors propose a new technique for handling gradient instability (vanishing/exploding) in neural network models, especially sequence models such as RNNs. Specifically, by leveraging results from recent works that establish a connection between Lyapunov growth exponents (from dynamical systems) and the singular values of the model-Jacobian itself, the authors propose a (regularization) approach that they coin ‘gradient flossing’ to control the evolution of these singular values desirably. They demonstrate that, as a result, the gradient signal remains stable, thereby allowing for successful training of the model itself. Additionally, the authors provide insights on the steps (schedule, pre/during training) to perform such a regularization leading to optimal gradient stabilization & computational benefits.

Strengths

The authors visit a longstanding problem in training NN models, esp. sequence models with long temporal dependencies. By leveraging recently established connections between Lyapunov exponents and the singular values of the (long-term) Jacobian, the authors propose a new technique in the form of: - Defining a loss function out of these Lyapunov exponents - Determine a clever strategy to evaluate this loss function through a QR decomposition to constrain these exponents directly to stable target values. This, in turn, allows control over the evolution of the Jacobian spectra itself. The authors provide experimental evidence supporting this technique on two specific tasks that, by design, require long-time gradient signals to propagate efficiently. Furthermore, the authors provide practical insights into the (expensive) computational needs of the exact implementation of such a scheme and the place (pre/during training) of the same. They then propose a modified, intermittent regularization routine with preserved benefits yet having lower compute necessities. The authors do an overall good job of satisfactorily articulating the problem statement, the proposed solution, observed benefits, and implementation limitations.

Weaknesses

While I believe this is a good contribution, I feel the work can undergo a fair amount of extra polishing in terms of clarity. Please find below (Questions) my comments/questions on areas where this work can be further improved.

Questions

- Eq. (2) does not necessarily need to refer to the ‘long-term Jacobian’ unless $\Delta{t} \equiv t - \tau \to \infty$. That should be clarified. - Is the loss function from Eq. (4) used as an additional or the only loss to train? - Eq. (4) does not seem to contain a reference to the target $\lambda$ values that the authors eventually aim to constrain them to - While the analytical $t\to\infty$ limit is understandable, can you provide an estimate of a reasonable sequence length for *practical data* that can serve as a finite size equivalent of this limit? - Is Eq. (5) subject to the assumption that parameter updates are gradient-based? - I believe the ’s’-index in Eq. (6) is supposed to run from $\tau$ to be exact (?) - Clearly, in arriving at Eq. (6), the $Q$’s from *subsequent* time-steps $s$ and $s+1$ act on each other to reduce to unity ($Q_{s+1}^T Q_{s}$), leaving with the $R$’s only. Can the authors explain why that might be the case in general? Is there an assumption about a small learning rate here? - Since the authors claim the QR decomposition need not be applied at every time step, will the equivalent Eq. (6) thus arrived at still remain to be factually correct? The expression of $T_t$ contains an iterative product of $ D_s$, of which if one chooses to QR decompose a handful only, then that does not seem to yield $\prod_s R^s$ for some period $t_{ons}$. - In Lines 140-141, you refer to using vanilla gradient descent for the experiments of Fig. 1 when you specify in Lines 132-133 that experiments were performed using ADAM. - Isn’t a target $\lambda$ value of -1 indicative of vanishing gradients as per Eq. (3)? How is the model still reliably attaining that value regardless, then? - In Fig. 2 (A), the ‘theory’ and ‘simulation’ legends seem to be interchanged. - In Fig. 2 (A) & (C), the numerical and theoretical results should overlap in the long time-step regime, yet as time increases, the two appear to diverge in the former and maintain a uniform scaling (~10^2) in the latter. Can you explain why? - Can you explain what you mean by ‘flossing before training’? That does not seem to be well clarified in the text.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

To the best of my knowledge, the authors act commendably in attributing to the various limitations of their approach.

Reviewer ZKKk5/10 · confidence 3/52023-07-06

Summary

Authors propose gradient flossing for RNNs, which adds an additional regularization term that keeps the sum of Lyapunov exponents close to zero. This essentially encourages the singular values of the long-term Jacobian to be close to 1, hence addressing the vanishing/exploding gradient problem in RNNs. The authors point out that minimizing this regularization term again has the same gradient instability problem, and they purpose to apply QR decomposition once in a while to avoid the ill-conditioned long-term Jacobian pass. Finally, the authors show gradient flossing achieves faster convergence and a higher success rate on simple long-term tasks up to 500 steps.

Strengths

* The idea of regularizing the Lyapunov exponent and using QR decomposition is novel and interesting. * The problem is well-motivated.

Weaknesses

* Empirical evaluations are toy. Did not include results on standard long-time horizon benchmarks like Longe Range Arena [1]. * The paper could have studied gradient flossing with architectures that are more suited for long horizon tasks such as state space models [2] and LRUs [3], instead of Vanilla RNNs and LSTMs. * Authors should compare to more recent methods as baselines such as HiPPO Matrices[2] and initialization methods used in LRU [3]. Those methods tend to be more competitive than a simple orthogonal initialization. * Clarity of the paper can be improved. For example, $\tau$ in equation 1 is never defined, and equation 4 could be better explained. In addition, I did not quite understand the technical parts of the paper explaining how equation 4 is calculated in practice (see questions below). [1] Tay, Yi, et al. "Long range arena: A benchmark for efficient transformers." arXiv preprint arXiv:2011.04006 (2020). [2] Gu, Albert, Karan Goel, and Christopher Ré. "Efficiently modelling long sequences with structured state spaces." arXiv preprint arXiv:2111.00396 (2021). [3] Orvieto, Antonio, et al. "Resurrecting recurrent neural networks for long sequences." arXiv preprint arXiv:2303.06349 (2023).

Questions

* I’m still unclear how QR decomposition is able to avoid the long-term Jacobian being ill-conditioned. * Why can we not use QR decomposition directly on the original RNN objective in equation 1 to avoid the long-term gradient instability problem?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

2 fair

Contribution

3 good

Limitations

See above.

Reviewer tH3Y7/10 · confidence 3/52023-07-07

Summary

The presented paper proposes a new method for tackling numerical instabilities during training for recurrent neural networks. The proposed method exploits a theoretical link between Lyapunov exponents and the singular values of the long-term Jacobian. The set of experiments is well-chosen to showcase the extent to which gradient flossing helps stabilize training for tasks involving long-range dependencies. There is an in-depth coverage of the practical aspects of the methods in the appendix, which is very useful for both machine learning practitioners and researchers. By steering the Lyapunov exponents to 0, either before and/or during training, the authors are able to show consistent improvements on a variety of synthetic learning tasks involving long-range dependencies in temporal sequences.

Strengths

- There are experiments validating basic hypotheses, such as the fact that the proposed algorithm is able to control the Lyapunov exponents, and that estimating the Lyapunov exponents via the QR decomposition matches the theory. - Practical aspects such as when to perform gradient flossing, how often should it be used during training, and also how often should the QR decomposition should be computed when unrolling the neural network dynamics.

Weaknesses

- There are no evaluation on real world complex benchmark to estimate the impact of the proposed contribution. However, the paper is already quite dense. - Minor remark: I don't understand the "the finite network size fluctuations of the Lyapunov exponents" in figure 1.

Questions

- (line 197 - 201) How does fig 2.C shows that the estimated condition number based on Lyapunov exponents can predict differences in condition number originating from finite network size $N$ ? The authors refer to appendix F for this matter, but it doesn't seem related to this claim. - Can the author explain in more details what are the "the finite network size fluctuations of the Lyapunov exponents" in figure 1 ? - Figure 11 in the appendix shows that too much gradient flossing during training can be detrimental, which means that the local minima of the gradient flossing procedure hardly coincides with the one corresponding to the actual learning task. It is also said in the introduction that appropriate initialization schemes can ensure well-behaved gradient at initialization, a property that can be lost during training. I suspect that the gradient flossing procedure prevents the weights from falling into an "ill-conditioned" region. Thus, how often should gradient flossing be used should depend on, given a fixed learning rate, how many iterations are needed to go from a "well-behave" region to an ill-condition one. Can the author elaborate on this specific intuition ? I am asking this question because I suspect that you already found regimes where the regularization objective conflicts with the actual learning task.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

4 excellent

Contribution

3 good

Limitations

As emphasized by the author, the Lyapunov exponent characterizes the singularity of the long-term Jacobian, and thus cannot characterize local-in-time behavior.

Reviewer 5Z2M2023-08-15

Having reviewed the rebuttals of the authors to my questions and weaknesses and also reviewing rebuttals to other reviewers, I will keep my score. There is still the importance of demonstrating this method on data that the overall research community is familiar with, and the authors do not seem willing to demonstrate this while still claiming to solve a well-known problem in the community of vanishing/exploding gradients. I still consider this a severe weakness of this work. The request is not to see this method as performant or superior to other work. It is to understand the proposed methods' potential failure modes and expected characteristics in uncontrolled settings.

Authorsrebuttal2023-08-21

Addressing Reviewer Concerns: Failure modes, Optimized GradientFlossing Methodology and sequentialMNIST

We thank the reviewer for their feedback and the emphasis on understanding the potential failure modes and expected characteristics of GradientFlossing in uncontrolled settings. 1. **Addressing Real-World Data Concerns**: We acknowledge the importance of evaluating GradientFlossing on datasets familiar to the research community. We are currently undergoing numerical evaluations using SequentialMNIST, with preliminary results appearing promising. These findings will be included in the camera-ready version of the manuscript. 2. **Mitigating Potential Failure Modes**: The potential conflict of persistent gradient flossing with certain task goals is a concern we have also observed (as depicted in Supplementary Figure 11). However, we have now devised an automatic solution to this challenge. By simultaneously conducting gradient flossing and training, we employ a combined loss function: $L_{total} = L_{task} + \alpha L_{gradientflossing}$ Crucially, the strength of gradient flossing, represented by $\alpha$, is trained alongside the network, dynamically adjusting its value throughout the learning process. This approach ensures that gradient flossing is optimized in relation to the task goal, effectively circumventing the described failure mode. We are also exploring a similar automated strategy to determine the optimal number of flossed Lyapunov exponents. 3. **Further Investigations and Future Work**: Exploding and vanishing gradients remain a prevalent issue in RNNs, and our GradientFlossing technique is also theoretically applicable to feedforward networks, offering a promising avenue for future research. While there exist other methods and architectures addressing these challenges, we believe GradientFlossing presents a complementary approach. Preliminary experiments also indicate its potential in further enhancing LSTM performance on tasks with long time-horizons to bridge. In conclusion, we are dedicated to ensuring our contribution is well-understood, robust, and applicable to real-world scenarios. We appreciate the reviewer's insights, which drive our ongoing efforts to enhance and clarify the impact and scope of GradientFlossing.

Reviewer ZKKk2023-08-17

Thank you authors for the rebuttal

Thank you authors' for the detailed response, as well as the more intuitive explanation for my clarification questions. I really appreciate it. I am still inclined to retain my rating because the authors did not touch upon the weaknesses I raised in the original review. In addition to agreeing with Reviewer 5Z2M's point on the method should be evaluated on standard benchmarks to investigate its failure modes, it should also be evaluated against more competitive baselines than simple vanilla RNNs/LSTMs (as suggested in the weakness section of the review).

Authorsrebuttal2023-08-21

Adressing weaknesses

We thank the reviewer for their feedback, both the constructive critique and the recognition that "The idea of regularizing the Lyapunov exponent and using QR decomposition is novel and interesting" and that "The problem is well-motivated." We apologize for not thoroughly addressing the weaknesses in our initial rebuttal. * **Standard Benchmarks**: Evaluations on SequentialMNIST are underway; results will be in the final manuscript. * **Competitive Baselines**: We're extending evaluations to advanced architectures and recent methods such as HiPPO Matrices and LRU initialization methods. * **Clarity**: We will clarify definitions and explanations in the manuscript, especially for equations 1 and 4. * **Mitigating Failure Modes**: Our manuscript has been updated to emphasize a combined loss function: $L_{\text{total}} = L_{\text{task}} + \alpha L_{\text{gradientflossing}}$ Here, the strength of gradient flossing, represented by $\alpha$, is trained alongside the network with backpropagation, dynamically adjusting its value throughout the learning process. This approach ensures that gradient flossing is optimized in relation to the task goal, effectively circumventing the described failure mode. We are also exploring a similar automated strategy to determine the optimal number of flossed Lyapunov exponents. * **References and Evaluations**: New references have been incorporated, and additional results will be shared on our anonymous GitHub repository. We are committed to enhancing our work based on your valuable insights.

Reviewer tH3Y2023-08-20

Satisfied by the answer

Dear authors, I apologize for the late response, but I did read your answer. I am satisfied with it. Likewise, I will keep my score as is, and I encourage the authors to pursue their work with further validation on real-world scenarios. I indeed think that understanding the possible mode of failure in an uncontrolled scenario would be very informative for practitioners, as pointed out by other reviewers. I would however like to point out that I highly appreciated the motivation behind the use of synthetic tasks.

Authorsrebuttal2023-08-21

Further improvements to GradientFlossing

We thank the reviewer for their positive feedback and the emphasis on understanding the potential failure modes and expected characteristics of GradientFlossing in uncontrolled settings. We'd like to highlight key updates to our manuscript: * **Mitigating Failure Modes**: We appreciate the reviewer's concerns about potential failure modes. Our recent advancements, as highlighted, involve a combined loss function: $L_{total} = L_{task} + \alpha L_{gradientflossing}$ With the parameter $\alpha$ also trained using backpropagraion, this approach aims to strike a balance, ensuring gradient flossing optimally complements the primary task goal. * **Evaluations of GradientFlossing beyond Synthetic Tasks**: We acknowledge the significance of real-world scenario validation. As mentioned, evaluations using SequentialMNIST are in progress, preliminary results look promising, and the results will be integrated into the final manuscript. * **Synthetic Tasks**: We are grateful for the reviewer's recognition of our motivation behind using synthetic tasks. They serve as controlled environments to rigorously test our method's fundamentals. We thank the reviewer again to dedicate their time and effort towards this submission.

Authorsrebuttal2023-08-21

Reviewer questions addressed?

As the discussion period draws to a close, we'd like to inquire if our responses have adequately addressed the reviewer's concerns. For clarity, here's a summary of our main updates: 1. Elucidated when $T$ is considered the long-term Jacobian. 2. Defined the application of the gradient flossing loss from Eq. (4). 3. Noted that the target Lyapunov exponents in Eq. (4) are set to $0$. 4. Showcased the Lyapunov spectrum's convergence over time steps. 5. Explained the context and assumptions of Eqs. (5) and (6). 6. Delved into the nuances of infrequent QR decomposition. 7. Rectified inconsistencies in optimization methods and figure descriptions. 8. Addressed discrepancies in Figure 2. 9. Detailed the 'flossing before training' approach. We trust that our clarifications meet the expectations set by the review. If there are any outstanding concerns, we're more than willing to provide further clarity. Otherwise, we extend our gratitude to the reviewer for their invaluable feedback, which has undeniably enhanced the quality of our manuscript.

Program Chairsdecision2023-09-21

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC