∇-Prox: Differentiable Proximal Algorithm Modeling for Large-Scale Optimization
∇-Prox: Differentiable Proximal Algorithm Modeling for Large-Scale Optimization
复制标题
DOI：
10.1145/3592144
复制
发表时间：
2023-07
期刊：
影响因子：
--
通讯作者：
Zeqiang Lai;Kaixuan Wei;Ying Fu;P. Härtel;Felix Heide
中科院分区：
文献类型：
--
作者：
Zeqiang Lai;Kaixuan Wei;Ying Fu;P. Härtel;Felix Heide
跨不同应用领域的任务可以构成大规模优化问题，这些问题包括图形、视觉、机器学习、成像、健康、调度、计划和能源系统预测。独立于应用领域，近端算法已经成为一种正式的优化方法，它成功地解决了大量存在的问题，通常在优化中利用特定于问题的结构。尽管基于模型的形式化优化为问题建模提供了一种有原则的方法，并保证了收敛性，但乍一看，这似乎与黑盒深度学习方法不一致。最近的一项研究表明，当与基于学习的成分相结合时，基于模型的优化方法是有效的、可解释的，并且允许在很少或没有额外训练数据的情况下将其推广到广泛的应用程序。然而，对不同任务进行这种混合方法的手工实验需要在近端优化和深度学习方面的专业知识，这通常容易出错且耗时。此外，天真地展开这些迭代方法会产生冗长的计算图，当通过自梯度技术进行区分时，会导致内存消耗激增，使基于批处理的训练具有挑战性。在这项工作中，我们引入了∇-Prox，这是一种特定于领域的建模语言和编译器，用于使用可微近邻算法进行大规模优化问题。△-Prox可以在高层次上简洁地指定未知的优化目标函数，并智能地将问题编译为计算效率和内存效率高的可微分求解器。∇- prox的核心特征之一是其完全可微性，它支持基于模型和学习的混合求解器，将近端优化与神经网络管道集成在一起。该方法的示例应用包括基于学习的先验和/或样本依赖的内循环优化调度程序，通过深度平衡学习或深度强化学习学习。通过几行代码，我们展示了∇-Prox可以为一系列图像优化问题生成高性能的求解器，包括端到端计算光学、图像脱除和压缩磁共振成像。我们还证明了∇-Prox可以用于能源系统规划的完全正交应用领域，这是能源危机和清洁能源转型的重要任务，在这方面，它优于最先进的CVXPY和商用Gurobi求解器。
Tasks across diverse application domains can be posed as large-scale optimization problems, these include graphics, vision, machine learning, imaging, health, scheduling, planning, and energy system forecasting. Independently of the application domain, proximal algorithms have emerged as a formal optimization method that successfully solves a wide array of existing problems, often exploiting problem-specific structures in the optimization. Although model-based formal optimization provides a principled approach to problem modeling with convergence guarantees, at first glance, this seems to be at odds with black-box deep learning methods. A recent line of work shows that, when combined with learning-based ingredients, model-based optimization methods are effective, interpretable, and allow for generalization to a wide spectrum of applications with little or no extra training data. However, experimenting with such hybrid approaches for different tasks by hand requires domain expertise in both proximal optimization and deep learning, which is often error-prone and time-consuming. Moreover, naively unrolling these iterative methods produces lengthy compute graphs, which when differentiated via autograd techniques results in exploding memory consumption, making batch-based training challenging. In this work, we introduce ∇-Prox, a domain-specific modeling language and compiler for large-scale optimization problems using differentiable proximal algorithms. ∇-Prox allows users to specify optimization objective functions of unknowns concisely at a high level, and intelligently compiles the problem into compute and memory-efficient differentiable solvers. One of the core features of ∇-Prox is its full differentiability, which supports hybrid model- and learning-based solvers integrating proximal optimization with neural network pipelines. Example applications of this methodology include learning-based priors and/or sample-dependent inner-loop optimization schedulers, learned with deep equilibrium learning or deep reinforcement learning. With a few lines of code, we show ∇-Prox can generate performant solvers for a range of image optimization problems, including end-to-end computational optics, image deraining, and compressive magnetic resonance imaging. We also demonstrate ∇-Prox can be used in a completely orthogonal application domain of energy system planning, an essential task in the energy crisis and the clean energy transition, where it outperforms state-of-the-art CVXPY and commercial Gurobi solvers.