Package: torchdeq
TorchDEQ: A Library for Deep Equilibrium ModelsDocumentation•Colab Tutorial•DEQ Zoo•Roadmap•CitationIntroductionDeep 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:$$\mathbf{z}^\star=f_\theta(\mathbf{z}^\star, \mathbf{x})$$Here, $\mathbf{x}$ is the input fed into the network, while $\mathbf{z}^\star$ stands as its output.EnterTorchDEQ- 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 ourColab Tutorial— best enjoyed with a comforting cup of tea!InstallationThrough pip.pipinstalltorchdeqFrom source.gitclonehttps://github.com/locuslab/torchdeq.git&&cdtorchdeq
pipinstall-e.Quick StartAutomatic arg parser decorator. You can call this function to add commonly used DEQ args to your program.add_deq_args(parser)Automatic DEQ instantiation. Callget_deqto 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 functioin of three tensors a, b, c.deffn(a,b,c):# Do something here...# Having the same input and output tensor shapes.returna,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-processz_outas normal tensors!ContributionsWe 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 preliminaryroadmapfor 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 myemail.Logo ExplainedThe 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 $f(\mathbf{x})$, biting its tail, representing $\mathbf{x}$, paints a vivid picture of a function attaining a fixed point. It's a metaphor layered with meaning, visualizing the attainment of stability, illustrated by the dragon completing its circle by biting its tail. This symbol is not just a snapshot of equilibrium; it's a dynamic representation of the infinite nature inherent in DEQ models.Citation@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}},}AcknowledgementsThis codebase is largely inspired by remarkable projects from the community.
We would like to sincerely thankDEQ,DEQ-Flow,PyTorch, andscipyfor their awesome open source.