SWAT-NN: Simultaneous Weights and Architecture Training for Neural Networks in a Latent Space

TL;DR

SWAT-NN introduces a novel method that simultaneously optimizes neural network architecture and weights by embedding both into a continuous latent space, enabling efficient discovery of compact, high-performing models.

Contribution

It proposes a universal autoencoder to embed architectures and weights into a latent space, allowing joint optimization via gradient descent for neural network design.

Findings

Effectively discovers sparse, compact neural networks.

Achieves strong performance on synthetic regression tasks.

Demonstrates the viability of continuous latent space optimization.

Abstract

Designing neural networks typically relies on manual trial and error or a neural architecture search (NAS) followed by weight training. The former is time-consuming and labor-intensive, while the latter often discretizes architecture search and weight optimization. In this paper, we propose a fundamentally different approach that simultaneously optimizes both the architecture and the weights of a neural network. Our framework first trains a universal multi-scale autoencoder that embeds both architectural and parametric information into a continuous latent space, where functionally similar neural networks are mapped closer together. Given a dataset, we then randomly initialize a point in the embedding space and update it via gradient descent to obtain the optimal neural network, jointly optimizing its structure and weights. The optimization process incorporates sparsity and compactness penalties to promote efficient models. Experiments on synthetic regression tasks demonstrate that our method effectively discovers sparse and compact neural networks with strong performance.