Robust Stochastically-Descending Unrolled Networks

Deep unrolling, or unfolding, is an emerging learning-to-optimize method that unrolls a truncated iterative algorithm in the layers of a trainable neural network. However, the convergence guarantees and generalizability of the unrolled networks are still open theoretical problems. To tackle these problems, we provide deep unrolled architectures with a stochastic descent nature by imposing descending constraints during training. The descending constraints are forced layer by layer to ensure that each unrolled layer takes, on average, a descent step toward the optimum during training. We theoretically prove that the sequence constructed by the outputs of the unrolled layers is then guaranteed to converge for in-distribution problems. We then analyze the generalizability to certain out-of-distribution (OOD) shifts in the optimization problems being solved. Our analysis shows that the descending nature imposed by the proposed constraints is transferable under these distribution shifts, subject to a generalization error, thereby providing the unrolled networks with OOD robustness. We numerically assess unrolled architectures trained with the proposed constraints in two different applications, including the sparse coding using learnable iterative shrinkage and thresholding algorithm (LISTA) and image inpainting using proximal generative flow (GLOW-Prox), and demonstrate the performance and robustness advantages of the proposed method.