[Paper Review] Inference in Multi-Layer Networks with Matrix-Valued Unknowns

Parthe Pandit, Mojtaba Sahraee-Ardakan+3 more|arXiv (Cornell University)|Jan 26, 2020

Stochastic Gradient Optimization Techniques46 references4 citations

TL;DR

This paper proposes the Multi-Layer Matrix Vector Approximate Message Passing (ML-Mat-VAMP) algorithm for inference in multi-layer networks with matrix-valued unknowns, extending the ML-VAMP framework to handle matrix variables. It establishes exact performance prediction in the large-system limit where input dimensions grow while the number of hidden nodes remains fixed, enabling precise error analysis for deep generative models and two-layer neural network learning.

ABSTRACT

We consider the problem of inferring the input and hidden variables of a stochastic multi-layer neural network from an observation of the output. The hidden variables in each layer are represented as matrices. This problem applies to signal recovery via deep generative prior models, multi-task and mixed regression and learning certain classes of two-layer neural networks. A unified approximation algorithm for both MAP and MMSE inference is proposed by extending a recently-developed Multi-Layer Vector Approximate Message Passing (ML-VAMP) algorithm to handle matrix-valued unknowns. It is shown that the performance of the proposed Multi-Layer Matrix VAMP (ML-Mat-VAMP) algorithm can be exactly predicted in a certain random large-system limit, where the dimensions $N imes d$ of the unknown quantities grow as $N ightarrow\infty$ with $d$ fixed. In the two-layer neural-network learning problem, this scaling corresponds to the case where the number of input features and training samples grow to infinity but the number of hidden nodes stays fixed. The analysis enables a precise prediction of the parameter and test error of the learning.

Motivation & Objective

• To address the inference problem of recovering input and hidden variables in stochastic multi-layer neural networks where hidden variables are represented as matrices.
• To extend the Vector Approximate Message Passing (VAMP) framework to handle matrix-valued unknowns, enabling efficient and scalable inference.
• To provide a unified algorithm for both Maximum A Posteriori (MAP) and Minimum Mean Square Error (MMSE) estimation in such networks.
• To establish a rigorous asymptotic performance analysis that enables exact prediction of parameter and test errors in the large-system limit.
• To support applications in deep generative priors, multi-task regression, and two-layer neural network learning with theoretical guarantees.

Proposed method

• Proposes the ML-Mat-VAMP algorithm as an extension of ML-VAMP to handle matrix-valued variables by processing consecutive layer pairs $(\mathbf{Z}_{\ell}, \mathbf{Z}_{\ell-1})$ sequentially.
• Uses a forward-backward message-passing structure with variables updated in forward ($\widehat{\mathbf{Z}}^{+}_{k\ell}$) and backward ($\widehat{\mathbf{Z}}^{-}_{k\ell}$) passes.
• Employs a state evolution framework that tracks the empirical distribution of estimation errors through a set of deterministic recursion equations.
• Introduces auxiliary variables $\mathbf{p}^{0}_{\ell}, \mathbf{p}^{+}_{k\ell}, \mathbf{q}^{+}_{k\ell}, \mathbf{q}^{-}_{k\ell}$ to represent sufficient statistics of the posterior distributions.
• Derives asymptotic performance predictions by analyzing the convergence of matrix iterates to i.i.d. random variables under the large-system limit $N \to \infty$ with fixed $d$, using tools from random matrix theory.
• Establishes exact prediction of mean-square error and parameter estimation error via a state evolution algorithm (Algorithm 4) that tracks sufficient statistics and correlation structures.

Experimental results

Research questions

• RQ1 Can the performance of inference in multi-layer networks with matrix-valued unknowns be exactly predicted in the large-system limit?
• RQ2 How can the Vector Approximate Message Passing (VAMP) framework be generalized to handle matrix-valued variables rather than vector-valued ones?
• RQ3 What is the asymptotic behavior of the estimation error in deep generative models and two-layer neural network learning under fixed number of hidden nodes and growing input dimensions?
• RQ4 How does the proposed ML-Mat-VAMP algorithm maintain accuracy and convergence in the presence of non-linear activation functions and structured noise?
• RQ5 What are the sufficient statistics and recursion rules that enable exact performance prediction for matrix inference in multi-layer networks?

Key findings

• The ML-Mat-VAMP algorithm enables exact performance prediction in the large-system limit where $N \to \infty$ with fixed $d$, corresponding to growing input features and training samples but fixed number of hidden nodes.
• The asymptotic performance of the algorithm is governed by a state evolution framework (Algorithm 4) that tracks the empirical distribution of estimation errors through deterministic recursion equations.
• The algorithm achieves exact prediction of mean-square error and parameter estimation error in two-layer neural network learning, with convergence of matrix iterates to i.i.d. random variables in the limit.
• The state evolution analysis confirms that the correlations between estimation errors and true variables converge to deterministic limits, enabling precise error prediction.
• The method successfully generalizes VAMP to matrix-valued inference, maintaining theoretical guarantees and scalability in high-dimensional settings.
• The framework supports both MAP and MMSE inference, with the latter enabling optimal estimation under uncertainty in the hidden variables.