Analyzing and Enhancing the Backward-Pass Convergence of Unrolled Optimization
This work addresses a bottleneck in integrating optimization models into deep learning, offering a practical solution for specialized learning tasks, though it is incremental in nature.
The paper tackled the challenge of backpropagation through unrolled optimization solvers in deep networks, showing that the backward pass is asymptotically equivalent to solving a linear system and proposing Folded Optimization to improve efficiency and flexibility, with experiments demonstrating computational advantages across various tasks.
The integration of constrained optimization models as components in deep networks has led to promising advances on many specialized learning tasks. A central challenge in this setting is backpropagation through the solution of an optimization problem, which often lacks a closed form. One typical strategy is algorithm unrolling, which relies on automatic differentiation through the entire chain of operations executed by an iterative optimization solver. This paper provides theoretical insights into the backward pass of unrolled optimization, showing that it is asymptotically equivalent to the solution of a linear system by a particular iterative method. Several practical pitfalls of unrolling are demonstrated in light of these insights, and a system called Folded Optimization is proposed to construct more efficient backpropagation rules from unrolled solver implementations. Experiments over various end-to-end optimization and learning tasks demonstrate the advantages of this system both computationally, and in terms of flexibility over various optimization problem forms.