Training Transformers with Enforced Lipschitz Constants

2025-07-17

AI Summary

Neural networks—particularly Transformers—are highly sensitive to input and weight perturbations, leading to adversarial vulnerability, training divergence, and overfitting. Existing Lipschitz-constrained networks struggle to maintain strict Lipschitz continuity throughout full training, especially in modern architectures, and lack guarantees beyond initialization. This paper introduces a全程 (end-to-end) enforced Lipschitz training framework: it jointly designs spectrally normalized matrix optimization strategies, integrating spectral normalization, weight decay, and the Muon optimizer to enforce strict Lipschitz constraints on weight matrices in both MLPs and Transformers. Notably, it yields the first Lipschitz Transformer trained without LayerNorm or other stabilization mechanisms. On the Shakespeare dataset (2M parameters), it achieves 60% accuracy; on internet text (145M parameters), it reaches 21%—marking substantial improvements in the robustness–performance trade-off.

Technology Category

• Machine Learning: Adversarial Learning & Robustness
• Natural Language Processing: Safety and Robustness
• Computer Vision: Adversarial Attacks & Robustness

Application Category

• Security and Privacy: Security and privacy of machine learning and AI applications
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications
• Search and Retrieval-Augmented AI: Large language models for search

Abstract

Neural networks are often highly sensitive to input and weight perturbations. This sensitivity has been linked to pathologies such as vulnerability to adversarial examples, divergent training, and overfitting. To combat these problems, past research has looked at building neural networks entirely from Lipschitz components. However, these techniques have not matured to the point where researchers have trained a modern architecture such as a transformer with a Lipschitz certificate enforced beyond initialization. To explore this gap, we begin by developing and benchmarking novel, computationally-efficient tools for maintaining norm-constrained weight matrices. Applying these tools, we are able to train transformer models with Lipschitz bounds enforced throughout training. We find that optimizer dynamics matter: switching from AdamW to Muon improves standard methods -- weight decay and spectral normalization -- allowing models to reach equal performance with a lower Lipschitz bound. Inspired by Muon's update having a fixed spectral norm, we co-design a weight constraint method that improves the Lipschitz vs. performance tradeoff on MLPs and 2M parameter transformers. Our 2-Lipschitz transformer on Shakespeare text reaches validation accuracy 60%. Scaling to 145M parameters, our 10-Lipschitz transformer reaches 21% accuracy on internet text. However, to match the NanoGPT baseline validation accuracy of 39.4%, our Lipschitz upper bound increases to 10^264. Nonetheless, our Lipschitz transformers train without stability measures such as layer norm, QK norm, and logit tanh softcapping.

Problem

• Address sensitivity of Transformers to input perturbations
• Enforce Lipschitz constraints during Transformer training
• Improve Lipschitz-performance tradeoff in large neural networks

Innovation

• Developed efficient norm-constrained weight matrices
• Trained transformers with enforced Lipschitz bounds
• Co-designed weight constraint method with Muon