Exploiting Explainable Metrics for Augmented SGD

TL;DR

This paper introduces explainability metrics based on low-rank factorization to assess layer-wise learning quality in deep neural networks, and uses these metrics to adaptively enhance SGD, resulting in improved generalization with minimal overhead.

Contribution

The paper proposes novel explainability metrics for neural network layers and leverages them to augment SGD with adaptive layer-wise learning rates, improving generalization.

Findings

RMSGD outperforms state-of-the-art methods in generalization.

Metrics strongly correlate with generalization performance.

Minimal additional computational cost.

Abstract

Explaining the generalization characteristics of deep learning is an emerging topic in advanced machine learning. There are several unanswered questions about how learning under stochastic optimization really works and why certain strategies are better than others. In this paper, we address the following question: \textit{can we probe intermediate layers of a deep neural network to identify and quantify the learning quality of each layer?} With this question in mind, we propose new explainability metrics that measure the redundant information in a network's layers using a low-rank factorization framework and quantify a complexity measure that is highly correlated with the generalization performance of a given optimizer, network, and dataset. We subsequently exploit these metrics to augment the Stochastic Gradient Descent (SGD) optimizer by adaptively adjusting the learning rate in each layer to improve in generalization performance. Our augmented SGD -- dubbed RMSGD -- introduces minimal computational overhead compared to SOTA methods and outperforms them by exhibiting strong generalization characteristics across application, architecture, and dataset.