TorchDEQ: A Library for Deep Equilibrium Models
This work provides a practical tool for researchers and practitioners working with DEQ models, addressing a domain-specific need by consolidating best practices into an open-source library.
The authors tackled the lack of a unified framework for Deep Equilibrium (DEQ) Models by developing TorchDEQ, a PyTorch library that simplifies their implementation and training, resulting in improved performance, stability, and efficiency across ten datasets and six models.
Deep Equilibrium (DEQ) Models, an emerging class of implicit models that maps inputs to fixed points of neural networks, are of growing interest in the deep learning community. However, training and applying DEQ models is currently done in an ad-hoc fashion, with various techniques spread across the literature. In this work, we systematically revisit DEQs and present TorchDEQ, an out-of-the-box PyTorch-based library that allows users to define, train, and infer using DEQs over multiple domains with minimal code and best practices. Using TorchDEQ, we build a ``DEQ Zoo'' that supports six published implicit models across different domains. By developing a joint framework that incorporates the best practices across all models, we have substantially improved the performance, training stability, and efficiency of DEQs on ten datasets across all six projects in the DEQ Zoo. TorchDEQ and DEQ Zoo are released as \href{https://github.com/locuslab/torchdeq}{open source}.