Hypergradient-free Training for Deep Equilibrium Models

Abstract

Deep Equilibrium (DEQ) models hold great promise for expanding the capabilities of neural networks and showcasing competitive performances with less memory cost. However, the backward propagation steps in their training process typically involve expensive Jacobian-inverse calculations, resulting in significantly higher computational costs than training conventional neural networks. Despite previous attempts to alleviate the burden, the trade-off between performance and efficiency remains unsatisfactory. In this study, we recast the training problem as a specific bilevel optimization problem. Then we propose BiDEQ, an efficient training algorithm based on penalty method to address the problem without calculating or approximating the hypergradient and matrix inverse in the backward propagation. Numerical experiments on various datasets demonstrate the superiority of our method over state-of-the-art methods for training DEQ models.