Unveiling the Training Dynamics of ReLU Networks through a Linear Lens

TL;DR

Problem: understanding the training dynamics of multi-layer ReLU networks is challenging due to depth and non-linearity. Approach: recast a bias-free L-layer ReLU network as a per-sample linear operator with an input-dependent $W_{\\text{eff}}(x)$, where $W_{\\text{eff}}(x) = \\prod_{l=L}^{1} (D_l(x) W_l)$ and $f(x) = W_{\\text{eff}}(x) x$. The study tracks the evolution of $W_{\\text{eff}}(x)$ via snapshots and visualizes the resulting manifold with $t$-SNE, on MNIST. Key findings: intra-class $W_{\\text{eff}}(x)$ converge while inter-class $W_{\\text{eff}}(x)$ diverge, and training yields class-specific, adaptive transformations with sub-class structure. Significance: provides an interpretable bridge between deep non-linear learning and dynamic linear models, offering a tool for explainability and guidance for extending to architectures with biases, residuals, and normalization layers.

Abstract

Deep neural networks, particularly those employing Rectified Linear Units (ReLU), are often perceived as complex, high-dimensional, non-linear systems. This complexity poses a significant challenge to understanding their internal learning mechanisms. In this work, we propose a novel analytical framework that recasts a multi-layer ReLU network into an equivalent single-layer linear model with input-dependent "effective weights". For any given input sample, the activation pattern of ReLU units creates a unique computational path, effectively zeroing out a subset of weights in the network. By composing the active weights across all layers, we can derive an effective weight matrix, $W_{\text{eff}}(x)$, that maps the input directly to the output for that specific sample. We posit that the evolution of these effective weights reveals fundamental principles of representation learning. Our work demonstrates that as training progresses, the effective weights corresponding to samples from the same class converge, while those from different classes diverge. By tracking the trajectories of these sample-wise effective weights, we provide a new lens through which to interpret the formation of class-specific decision boundaries and the emergence of semantic representations within the network.

Figures (1)

• Figure 1: t-SNE visualization of the effective weight vectors $W_{\text{eff}}(x)$ for 1,000 MNIST test samples. Each point represents the effective weight for a single sample, colored by its true class label. (a) Initial state. (b) After training.