TorchDEQ: A Library for Deep Equilibrium Models
Zhengyang Geng J. Zico Kolter
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}.