Documentation • Colab Tutorial • DEQ Zoo • Roadmap • Citation
Deep Equilibrium Models, or DEQs, a recently developed class of implicit neural networks, merge the concepts of fixed point systems with modern deep learning. Fundamentally, DEQ models establish their output based on the equilibrium of nonlinear systems. This can be represented as:
Here,
Enter TorchDEQ - a fully featured, out-of-the-box, and PyTorch-based library tailored for the design and deployment of DEQs. It provides intuitive, decoupled, and modular interfaces to customize general-purpose DEQs for arbitrary tasks, all with just a handful of code lines.
Dive into the world of DEQ with TorchDEQ! Craft your own DEQ effortlessly in just a single line of code. Kickstart your journey with our Colab Tutorial — best enjoyed with a comforting cup of tea!
-
Through pip.
pip install torchdeq
-
From source.
git clone https://github.com/locuslab/torchdeq.git && cd torchdeq pip install -e .
- Automatic arg parser decorator. You can call this function to add commonly used DEQ args to your program.
add_deq_args(parser)
- Automatic DEQ instantiation. Call
get_deq
to get your DEQ layer in a single line! It's highly decoupled implementation agnostic to your model design.
deq = get_deq(args)
- Easy DEQ forward. Even for a multi-equilibria system, you can execute your DEQ forward in a single line!
# Assume f is a function of three tensors a, b, c.
def fn(a, b, c):
# Do something here...
# Having the same input and output tensor shapes.
return a, b, c
# A callable object (`fn` here) that defines your fixed point system.
# `fn` can be a functor defined in your Pytorch forward function.
# A functor can take your input injection from the local variables.
# You can also pass a Pytorch Module into the DEQ class.
z_out, info = deq(fn, (a0, b0, c0))
- Automatic DEQ backward. Gradients (both exact and inexact grad) are tracked automatically! Working with TorchDEQ is the same as other standard PyTorch operators. Just post-process
z_out
as normal tensors!
We warmly welcome contributions to TorchDEQ from the community! If you have suggestions for improving the library, introducing new features, or identifying and fixing bugs, please open an issue to discuss with us! Once a direction has been discussed, we can proceed to build, test, and submit a pull request (PR) to TorchDEQ together. Keep a PR clean, well-tested, and have a single focus! While numerical errors and stability may seem minor initially, they can culminate in significant effects over time.
We have provided a preliminary roadmap for the development of this library and are always open to fresh perspectives. Feel free to reach out for questions, discussions, or library developments! Here is my email.
The logo we’ve chosen draws inspiration from the ancient symbol, Ouroboros, a powerful emblem depicting a serpent or dragon eternally consuming its own tail. Unearthed in the tomb of Tutankhamun, the Ouroboros symbolizes the cyclicality of time, embodying both creation and destruction, inception and conclusion. It’s a profound representation of infinity and wholeness, transcending various mythologies and philosophies across time.
For DEQ models, our choice of logo bears a metaphorical weight. The dragon, denoting
@misc{torchdeq,
author = {Zhengyang Geng and J. Zico Kolter},
title = {TorchDEQ: A Library for Deep Equilibrium Models},
year = {2023},
publisher = {GitHub},
journal = {GitHub repository},
howpublished = {\url{https://github.com/locuslab/torchdeq}},
}
This codebase is largely inspired by remarkable projects from the community. We would like to sincerely thank DEQ, DEQ-Flow, PyTorch, and scipy for their awesome open source.