Simultaneous Weight and Architecture Optimization for Neural Networks

TL;DR

This work introduces a novel neural network training framework that fundamentally transforms the process by learning architecture and parameters simultaneously with gradient descent, and can discover sparse and compact neural networks for given datasets.

Abstract

Neural networks are trained by choosing an architecture and training the parameters. The choice of architecture is often by trial and error or with Neural Architecture Search (NAS) methods. While NAS provides some automation, it often relies on discrete steps that optimize the architecture and then train the parameters. We introduce a novel neural network training framework that fundamentally transforms the process by learning architecture and parameters simultaneously with gradient descent. With the appropriate setting of the loss function, it can discover sparse and compact neural networks for given datasets. Central to our approach is a multi-scale encoder-decoder, in which the encoder embeds pairs of neural networks with similar functionalities close to each other (irrespective of their architectures and weights). To train a neural network with a given dataset, we randomly sample a neural network embedding in the embedding space and then perform gradient descent using our custom loss function, which incorporates a sparsity penalty to encourage compactness. The decoder generates a neural network corresponding to the embedding. Experiments demonstrate that our framework can discover sparse and compact neural networks maintaining a high performance.

Figures (6)

• Figure 1: The overall framework for the multi-scale encoder-decoder.
• Figure 2: Comparison of functionalities between input MLPs and decoded MLPs using the min loss function. The blue graphs on the left show the outputs of the input MLPs (z-axis as outputs). The red graphs on the right represent predictions of the decoded MLPs with smallest loss over 4 decoders (z-axis as predictions). Since the MLP has three input data points, we fixed one input at 0.5, and the other two inputs were assigned to the x and y axes, ranging in [-1 1].
• Figure 3: In each row, the dataset is generated by a 3 hidden-layer MLP with 35 non-zero weights for sigmoid-based, leaky ReLU-based, and linear-based MLPs, respectively. The blue plots represent the ground truth outputs, while the red plots from left to right correspond to the 1, 2, 3, and 4 hidden-layer optimal MLPs found through the search.
• Figure 4: This figure displays three plots generated using datasets from a 2 hidden-layer MLP with 14 non-zero weights and Decoder 3. The non-zero count was adjusted by modifying the weight of the sparsity term in the loss function.
• Figure 5: In each row, the dataset is generated by a 2 hidden-layer MLP with 14 non-zero weights for sigmoid-based, leaky ReLU-based, and linear-based MLPs, respectively.