[Paper Review] Stable and Interpretable Unrolled Dictionary Learning

Bahareh Tolooshams, Demba Ba|arXiv (Cornell University)|May 31, 2021

Sparse and Compressive Sensing Techniques4 citations

TL;DR

This paper proposes PUDLE, a provably stable and interpretable unrolled dictionary learning method that addresses bias in code and dictionary estimates from backpropagation through unrolled sparse coding networks. By modifying network initialization, loss functions, and gradient updates, PUDLE ensures convergence to the ground-truth dictionary and enables mathematical interpretability through explicit links between network weights, outputs, and training data.

ABSTRACT

The dictionary learning problem, representing data as a combination of a few atoms, has long stood as a popular method for learning representations in statistics and signal processing. The most popular dictionary learning algorithm alternates between sparse coding and dictionary update steps, and a rich literature has studied its theoretical convergence. The success of dictionary learning relies on access to a "good" initial estimate of the dictionary and the ability of the sparse coding step to provide an unbiased estimate of the code. The growing popularity of unrolled sparse coding networks has led to the empirical finding that backpropagation through such networks performs dictionary learning. We offer the theoretical analysis of these empirical results through PUDLE, a Provable Unrolled Dictionary LEarning method. We provide conditions on the network initialization and data distribution sufficient to recover and preserve the support of the latent code. Additionally, we address two challenges; first, the vanilla unrolled sparse coding computes a biased code estimate, and second, gradients during backpropagated learning can become unstable. We show approaches to reduce the bias of the code estimate in the forward pass, and that of the dictionary estimate in the backward pass. We propose strategies to resolve the learning instability by tuning network parameters and modifying the loss function. Overall, we highlight the impact of loss, unrolling, and backpropagation on convergence. We complement our findings through synthetic and image denoising experiments. Finally, we demonstrate PUDLE's interpretability, a driving factor in designing deep networks based on iterative optimizations, by building a mathematical relation between network weights, its output, and the training set.

Motivation & Objective

• To address the bias in code and dictionary estimates arising from unrolled sparse coding networks during backpropagation.
• To ensure stable convergence of dictionary learning via theoretical conditions on initialization and data distribution.
• To enhance interpretability of deep networks by establishing a mathematical link between network weights, outputs, and training data.
• To resolve instability in gradient updates during backpropagation through unrolled networks.
• To provide a theoretically grounded framework for unrolled dictionary learning that preserves support recovery of sparse codes.

Proposed method

• PUDLE employs a unrolled ISTA-based network with T iterations of sparse coding, where each layer learns a differentiable approximation of the lasso solution.
• It introduces a modified loss function that reduces bias in the forward pass code estimate and stabilizes the backward pass gradient updates.
• The method uses a tailored network initialization strategy based on the pairwise method (Arora et al., 2015) to improve initial support recovery.
• Gradient updates are stabilized by tuning step size α and introducing a decay mechanism for the regularization parameter ν.
• Interpretability is achieved by constructing a Gram matrix G that quantifies the contribution of each training sample to the final dictionary and code outputs.
• The framework uses automatic differentiation via backpropagation through the unrolled network to learn the dictionary, with normalization applied after each update.

Experimental results

Research questions

• RQ1 Under what conditions on network initialization and data distribution can unrolled dictionary learning preserve the true support of the sparse code?
• RQ2 How can the bias in the code estimate from the forward pass of an unrolled network be reduced to improve dictionary recovery?
• RQ3 What modifications to the loss function and network parameters stabilize gradient updates during backpropagation in unrolled dictionary learning?
• RQ4 How can the learned dictionary and code be made interpretable by linking them to training data via a mathematical framework?
• RQ5 What is the impact of unrolling depth, learning rate, and batch size on convergence and stability in unrolled dictionary learning?

Key findings

• PUDLE achieves stable convergence to a neighborhood of the ground-truth dictionary across various learning rates, with the best convergence observed using the ae-ls gradient.
• The method successfully recovers the support of the true sparse code under proper initialization and data distribution assumptions, as validated in synthetic experiments.
• Dictionary learning with PUDLE outperforms NOODL and SPORCO in image denoising, achieving a PSNR of 31.4 dB on the BSD100 dataset with a 9×9 convolutional dictionary.
• The use of the pairwise initialization method results in faster convergence and improved support recovery, even with limited unrolling depth (T=25).
• The Gram matrix construction enables interpretable analysis, showing that the top 200 contributing training images account for over 90% of the reconstruction energy in MNIST experiments.
• Experiments demonstrate robustness to learning rate variations, with convergence maintained across 10−4, 10−3, and 10−2 learning rates, though smaller rates require more epochs.