Probability Distribution Learning and Its Application in Deep Learning

TL;DR

This paper introduces a probability distribution learning framework to analyze deep learning's optimization and generalization, providing theoretical guarantees and insights into neural network training.

Contribution

It proposes the PD learning framework, establishes the Fenchel-Young loss as optimal, and introduces new concepts to explain SGD effectiveness and generalization bounds.

Findings

Fenchel-Young loss is necessary for PD learning.

Introduces $ ext{H}(\psi)$-convexity and $ ext{H}(\Psi)$-smoothness for DNNs.

Provides bounds on risk and generalization error influenced by training set size and information loss.

Abstract

Despite its empirical success, deep learning still lacks a comprehensive theoretical understanding of model fitting and generalization. This paper proposes the probability distribution (PD) learning framework to analyze the optimization and generalization mechanisms of deep learning. Within this framework, the conditional distribution of labels given features is the primary learning target, with the loss function, prior knowledge, and model properties explicitly characterized. Under these formulations, we establish theoretical guarantees on optimizability, even in non-convex settings, and derive generalization error bounds that provide meaningful explanations for practical performance. Specifically, we first prove theoretically that the Fenchel-Young loss is the natural and necessary choice for solving PD learning problems, thereby justifying the generality of conclusions based on this loss. Second, to capture the characteristics of deep neural networks (DNNs), we introduce the notions of $\mathcal{H}(\psi)$-convexity and $\mathcal{H}(\Psi)$-smoothness, which generalize the classical concepts of strong convexity and Lipschitz smoothness. Based on them, we provide a theoretical explanation for the effectiveness of SGD in training DNNs. Finally, we derive model-independent bounds on the expected risk and generalization error for trained models, revealing the influence of the training set size, regularization term, the mutual information between labels and features, and the information loss caused by model irreversibility on risk and generalization. Based on our theoretical analysis and experimental validation, we believe that the PD learning framework facilitates a deeper and more unified theoretical understanding of deep learning.