Learning Sparse Neural Networks through $L_0$ Regularization

Christos Louizos; Max Welling; Diederik P. Kingma

Source-linked AI summary

Learning Sparse Neural Networks through $L_0$ Regularization

Christos Louizos, Max Welling, Diederik P. Kingma

arXiv:1712.01312v2Dec 4, 2017stat.MLcs.LG

TL;DR

The paper addresses the intractability of direct L0 regularization by introducing differentiable stochastic gates and the hard concrete distribution. Experiments show competitive sparsification with theoretical training speedups.

Problem
Direct L0 regularization explicitly penalizes nonzero parameters but is combinatorial and intractable for large models.
Method
The method smooths the expected L0 objective with continuous distributions, hard-sigmoid transformations, and hard concrete gates while preserving exact parameter zeros.
Results
The proposed L0 minimization produces neural-network sparsification competitive with current approaches while theoretically allowing training speedups.
Takeaways & Limitations
The framework enables principled, efficient gradient-based optimization of sparse model structures and can provide computational benefits with minimal or no performance loss.
Takeaways & Limitations
Efficient training of very large networks with learned sparsity patterns remains identified as future work.

Abstract

from arXiv · show

We propose a practical method for $L_0$ norm regularization for neural networks: pruning the network during training by encouraging weights to become exactly zero. Such regularization is interesting since (1) it can greatly speed up training and inference, and (2) it can improve generalization. AIC and BIC, well-known model selection criteria, are special cases of $L_0$ regularization. However, since the $L_0$ norm of weights is non-differentiable, we cannot incorporate it directly as a regularization term in the objective function. We propose a solution through the inclusion of a collection of non-negative stochastic gates, which collectively determine which weights to set to zero. We show that, somewhat surprisingly, for certain distributions over the gates, the expected $L_0$ norm of the resulting gated weights is differentiable with respect to the distribution parameters. We further propose the \emph{hard concrete} distribution for the gates, which is obtained by "stretching" a binary concrete distribution and then transforming its samples with a hard-sigmoid. The parameters of the distribution over the gates can then be jointly optimized with the original network parameters. As a result our method allows for straightforward and efficient learning of model structures with stochastic gradient descent and allows for conditional computation in a principled way. We perform various experiments to demonstrate the effectiveness of the resulting approach and regularizer.

1 INTRODUCTION

The paper motivates sparsifying neural networks to remove unnecessary computation and reduce overfitting, then proposes a differentiable surrogate for L0 regularization using hard-sigmoid-transformed continuous variables.

Neural networks can be heavily overparameterized, pruned without accuracy loss, and prone to memorizing random patterns without proper regularization.
Sparsification removes irrelevant degrees of freedom, reducing computation and model complexity while alleviating overfitting.
L0 regularization directly penalizes nonzero parameters without additional shrinkage, but its combinatorial optimization is intractable for large models.
The proposed framework smooths the expected L0 objective with continuous distributions while preserving exact zeros and enabling gradient-based optimization.
The hard concrete distribution stretches a binary concrete variable and transforms its samples with a hard-sigmoid for efficient structure learning.

2 MINIMIZING THE L0 NORM OF PARAMETRIC MODELS

The method replaces discrete L0 gates with continuous stochastic variables followed by hard-sigmoid rectification, preserving exact zeros while making the regularized objective differentiable.

L0-regularized empirical risk penalizes the number of nonzero parameters, with λ controlling regularization strength; AIC and BIC arise from particular λ choices.
Direct optimization is intractable because the L0 penalty is nondifferentiable and the parameter vector has 2^|θ| possible states.
Binary gates indicate parameter presence, so the L0 norm equals the number of active gates, but discrete gates obstruct efficient gradient optimization.
2.1 A GENERAL RECIPE FOR EFFICIENTLY MINIMIZING L0 NORMS: Continuous random variables passed through a hard-sigmoid produce gates that can be exactly zero while retaining differentiable probabilities of activation.
2.1 A GENERAL RECIPE FOR EFFICIENTLY MINIMIZING L0 NORMS: The resulting objective is differentiable with respect to gate-distribution parameters, enabling stochastic-gradient optimization while retaining exact parameter zeros.
2.2 THE HARD CONCRETE DISTRIBUTION: The hard concrete stretches a binary concrete distribution to an interval extending beyond [0,1], then applies a hard-sigmoid to obtain binary-like gates.
2.2 THE HARD CONCRETE DISTRIBUTION: Noise smooths the hard-sigmoid into a sigmoid on average, supporting gradient-based optimization of the stochastic gates.

3 RELATED WORK

Prior work used rectified continuous random variables to emulate binary variables, but the paper positions hard concrete as better suited to Bernoulli-like behavior for L0 regularization.

Continuous rectifications of Gaussian or logistic variables have represented exact zeros while retaining tractable continuous optimization.
These unimodal distributions cannot represent the bimodal behavior of Bernoulli variables, which is important for approximating the L0 penalty.

4 EXPERIMENTS

Experiments on MNIST and CIFAR classification show that L0 regularization can learn sparse architectures during training while maintaining competitive accuracy and potentially reducing computation.

MNIST classification and sparsification: Table 1 compares learned architectures by neurons remaining after pruning and test error after 200 epochs.The comparison includes baselines from Louizos et al. (2017) and the proposed L0 minimization under L0hc.
MNIST classification and sparsification: Expected training FLOPs are tracked over iterations for original, dropout, and L0-regularized networks.The FLOP calculation assumes one operation for multiplication and one for addition.
MNIST classification and sparsification: The L0 procedure can potentially yield significant computational benefits with minimal or no performance loss, while stronger regularization favors speedup over compression.For LeNet, larger convolutional-layer λ values in the λ sep. setting produce a larger FLOP reduction; this setting is preferable when speedup is the priority.
CIFAR classification: With λ = 0.001/N, L0-regularized WideResNets improve accuracy over dropout equivalents on both CIFAR 10 and CIFAR 100.The same setting gradually decreases FLOPs without a convergence-speed lag, while increasing λ yields slight error increases and potentially larger speedups.

5 DISCUSSION

The paper presents hard concrete gates as a practical route to optimizing L0 objectives and reports sparsification, competitive performance, and improvement over dropout. It identifies more effective conditional computation and Bayesian extensions as future directions.

Discussion: The proposed recipe smooths the combinatorial L0 problem with continuous distributions followed by a hard-sigmoid while preserving exact zeros and enabling gradient optimization.The hard concrete distribution stretches a binary concrete variable before applying the hard-sigmoid.
Discussion: Experiments show competitive neural-network sparsification and theoretically enabled training speedups, with WideResNet CIFAR experiments improving upon dropout.The paper also characterizes the process as providing an inductive bias and regularizer.
Discussion: Future work includes better exploiting conditional computation in very large networks, Bayesian parameter treatments, learned bit precision, and binary latent-variable applications.These directions are presented as potential extensions rather than completed results.

A RELATION TO VARIATIONAL INFERENCE

The expected L0 objective is connected to variational inference under a spike-and-slab prior, with active parameters incurring a fixed coding cost. The resulting surrogate can also incorporate gate uncertainty and prior information about gate behavior.

Spike-and-slab formulation: A spike-and-slab prior assigns each gate a Bernoulli inclusion variable and each parameter either a zero point mass or a continuous distribution.The prior uses p(z)=Bernoulli(π), a delta spike at zero when z=0, and a standard normal slab when z=1.
Variational formulation: Variational inference factorizes a spike-and-slab approximate posterior over parameter dimensions and yields a variational free energy.The free energy contains KL terms for gate distributions and active-parameter distributions.
Connection to L0 regularization: Assuming optimized parameters and a fixed active-parameter KL cost λ turns the variational objective into an expected L0-style objective.The fixed cost is interpreted as the information required to encode an active parameter under an adaptive hypothetical prior.
Connection to L0 regularization: The expected L0 procedure is a close surrogate to a variational bound with a spike-and-slab parameter distribution and fixed coding cost for active gates.When gate uncertainty matters, the fuller variational objective should be optimized because it penalizes gate entropy and can incorporate gate priors.

B THE HARD CONCRETE DISTRIBUTION

The hard concrete distribution converts a binary concrete variable into a distribution with exact endpoint masses and a continuous interior component. Stretching before hard-sigmoid rectification creates gates that can be exactly zero or one.

Construction: The hard concrete is constructed from a binary concrete random variable by stretching its support to an interval (γ, ζ), with γ<0 and ζ>1.The stretched variable is defined as s̄=s(ζ−γ)+γ.
Construction: Applying a hard-sigmoid to the stretched variable produces gate values z=min(1,max(0,s̄)).The hard-sigmoid rectification maps values below zero to zero and values above one to one.
Distributional form: The resulting gate distribution combines a delta mass at zero, a delta mass at one, and a truncated continuous distribution on (0,1).The endpoint probabilities are Q_s̄(0|φ) for zero and 1−Q_s̄(1|φ) for one.

C NEGATIVE KL-DIVERGENCE FOR HARD CONCRETE DISTRIBUTIONS

For hard concrete gates, the prior must share the approximate posterior’s support so their KL divergence is valid. The KL can then be computed using the mixture structure, with Monte Carlo estimation available when the truncated integral lacks a closed form.

Matching support: A valid prior for hard concrete gates is formed by applying the same hard-sigmoid transformation to an arbitrary continuous distribution.This gives the prior the same three mixture components as the hard concrete posterior.
KL computation: Because the prior and posterior have matching mixture structures, their KL divergence can be computed using the chain rule of relative entropy.The mixture representation includes endpoint masses and a continuous interior component.
Numerical estimation: When the truncated-distribution integral is unavailable in closed form, the KL can be estimated by Monte Carlo sampling through the inverse transform method.Sampling can be performed over an interval such as (γ,ζ).
Caveat: The KL between rectified distributions is not generally equal to the KL between their pre-rectified variables because rectification is non-invertible.The non-invertibility prevents directly transferring the KL through the transformation.