Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Adaptive Width Neural Networks

Federico Errica, Henrik Christiansen, Viktor Zaverkin, Mathias Niepert, Francesco Alesiani

TL;DR

The paper tackles the entrenched practice of fixing layer widths and costly hyper-parameter searches by introducing Adaptive Width Neural Networks (AWNN), which learn unbounded layer widths during training. AWNN jointly optimizes per-layer widths and weights via backpropagation, using a variational objective over latent widths $\boldsymbol{\lambda}$ and neuron weights $\boldsymbol{\theta}$ and a finite truncation $D_\ell$ defined as the $k$-quantile of a monotonically decreasing distribution $f_\ell$. A soft ordering of neuron importance is enforced through a decaying function $f_\ell$ to enable gradual growth and safe post-hoc truncation, while a rescaled initialization (Kaiming+) stabilizes deep AWNN Training. The approach demonstrates comparable or improved performance across tabular, image, text, sequence, and graph tasks, with practical benefits including online compression and the ability to reduce transformer parameter counts by large factors without sacrificing accuracy. These results suggest AWNN as a viable path to reduce hyper-parameter tuning costs for large foundation models and to enable dynamic, task-dependent width adaptation.

Abstract

For almost 70 years, researchers have mostly relied on hyper-parameter tuning to select the width of neural networks' layers. This paper challenges the status quo by introducing an easy-to-use technique to learn an unbounded width of a neural network's layer during training. The technique does not rely on alternate optimization nor hand-crafted gradient heuristics; rather, it jointly optimizes the width and the parameters of each layer via simple backpropagation. We apply the technique to a broad range of data domains such as tables, images, text, sequences, and graphs, showing how the width adapts to the task's difficulty. The method imposes a soft ordering of importance among neurons, by which it also is possible to truncate the trained network at virtually zero cost, achieving a smooth trade-off between performance and compute resources in a structured way. Alternatively, one can dynamically compress the network with no performance degradation. In light of recent foundation models trained on large datasets, believed to require billions of parameters and where hyper-parameter tuning is unfeasible due to humongous training costs, our approach stands as a viable alternative for width learning.

Adaptive Width Neural Networks

TL;DR

The paper tackles the entrenched practice of fixing layer widths and costly hyper-parameter searches by introducing Adaptive Width Neural Networks (AWNN), which learn unbounded layer widths during training. AWNN jointly optimizes per-layer widths and weights via backpropagation, using a variational objective over latent widths and neuron weights and a finite truncation defined as the -quantile of a monotonically decreasing distribution . A soft ordering of neuron importance is enforced through a decaying function to enable gradual growth and safe post-hoc truncation, while a rescaled initialization (Kaiming+) stabilizes deep AWNN Training. The approach demonstrates comparable or improved performance across tabular, image, text, sequence, and graph tasks, with practical benefits including online compression and the ability to reduce transformer parameter counts by large factors without sacrificing accuracy. These results suggest AWNN as a viable path to reduce hyper-parameter tuning costs for large foundation models and to enable dynamic, task-dependent width adaptation.

Abstract

For almost 70 years, researchers have mostly relied on hyper-parameter tuning to select the width of neural networks' layers. This paper challenges the status quo by introducing an easy-to-use technique to learn an unbounded width of a neural network's layer during training. The technique does not rely on alternate optimization nor hand-crafted gradient heuristics; rather, it jointly optimizes the width and the parameters of each layer via simple backpropagation. We apply the technique to a broad range of data domains such as tables, images, text, sequences, and graphs, showing how the width adapts to the task's difficulty. The method imposes a soft ordering of importance among neurons, by which it also is possible to truncate the trained network at virtually zero cost, achieving a smooth trade-off between performance and compute resources in a structured way. Alternatively, one can dynamically compress the network with no performance degradation. In light of recent foundation models trained on large datasets, believed to require billions of parameters and where hyper-parameter tuning is unfeasible due to humongous training costs, our approach stands as a viable alternative for width learning.
Paper Structure (34 sections, 3 theorems, 39 equations, 12 figures, 6 tables, 1 algorithm)

This paper contains 34 sections, 3 theorems, 39 equations, 12 figures, 6 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Orthogonal Methods
  4. Adaptive Width Learning
  5. Imposing a Soft Ordering on Neurons' Importance
  6. Rescaled Weight Initialization for Deep AWNN
  7. Future Directions and Limitations
  8. Experiments and Setup
  9. Results
  10. Adaptation to Task Difficulty and Convergence
  11. Training Stability Analysis
  12. Online Network Compression via Regularization
  13. Post-hoc Truncation Achieves a Trade-off between Performance and Compute Resources
  14. Conclusions
  15. Background Notions
  16. ...and 19 more sections

Key Result

Theorem 3.1

Consider an MLP with activations as in Equation eq:unit-rescale and ReLU nonlinearity. At initialization, given $\alpha_j^\ell=\sigma\left(\sum_{k=1}^{D_{\ell-1}} w^\ell_{jk} h^{\ell-1}_k\right), \mathrm{Var}[w_{j*}^\ell] = \frac{2}{\sum_{j=1}^{D_{\ell-1}} f^2_{\ell}(j)} \Rightarrow \mathrm{Var}[\a

Figures (12)

  • Figure 1: (Left) The graphical model of AWNN, with dark observable random variables and white latent ones. (Middle) The distribution $f_\ell$ over hidden units' importance at layer $\ell$ is parametrized by $\lambda_\ell$. The width of layer $\ell$ is chosen as the quantile function of the distribution $f_\ell$ evaluated at $k$ and denoted by $D_\ell$. (Right) The hidden units' activations at layer $\ell$ are rescaled by their importance.
  • Figure 2: (Left) The learned width adapts to the increasing difficulty of the task, from the DoubleMoon to SpiralHard. (Right) AWNN reaches perfect test accuracy with a comparable amount of epochs on DoubleMoon and Spiral, while it converges faster on SpiralHard.
  • Figure 3: Training converges to similar widths on SpiralHard for different batch sizes (left) and starting rates $\nu$, but the latter seems to require a bounded nonlinearity such as ReLU6 to converge in a reasonable amount of epochs (right).
  • Figure 4: It is possible to regularize the width at training time by increasing the magnitude of the loss term $\log\frac{p(\boldsymbol{\nu})}{q(\boldsymbol{\nu})}$. The total width is reduced by more than 50% (left) while preserving accuracy (right). The inset plot refers to the loss term that AWNN tries to maximize.
  • Figure 5: (Left) Thanks to the soft ordering imposed on the neurons, one can also truncate the neural network after training by simply removing the last neurons. (Right) The distribution of neurons' activations for all Spiral test samples follows an exponential-like curve.
  • ...and 7 more figures

Theorems & Definitions (9)

  • Theorem 3.1
  • proof
  • Definition A.1
  • Definition A.2
  • Proposition D.1: AWNN ELBO Lipschitz continuity
  • proof
  • proof
  • Lemma E.1
  • proof