Community Spotlight: DWSNets
DWSNets is a public repo demonstrating a novel network architecture for learning in deep weight spaces. Here's what you need to know
Created on May 11|Last edited on May 18
Comment
This post is a piece in our series highlighting some of our favorite community repos that have instrumented W&B. If you'd like your repo to be featured, please reach out to editor@wandb.com
💡
Introduction
Deep nets are everywhere now and for good reason. These complex systems are the primary model for a range of applications, from image recognition to natural language processing. Recently, deep nets also became a primary model for representing data samples, for example, INRs for representing images, 3D objects, or scenes.
In all of these use cases, there is a need to perform some sort of transformation to the model weights. But what is the best way to carry out that process? That question is precisely the focus of the project DWSNets, a repo that is the official implementation of the paper Equivariant Architectures for Learning in Deep Weight Spaces, recently accepted at the International Conference on Machine Learning 2023.
What is DWSNets?
Many modern techniques for neural network training are based on their inherent symmetrical structures. While such techniques yield impressive results in a variety of applications, it makes adjusting weights a challenge.
DWSNets demonstrates a new architecture for effectively processing weights of neural networks while preserving equivariance to their permutation symmetries. This approach has a wide potential and could enable a wide range of interesting tasks, such as adapting a pre-trained network to a new domain.
A figure of MLP permutation symmetries
Why was DWSNets built?
Learning over the weights of neural nets is a relatively novel concept. There are few papers to draw from, and even so, it's unclear if the approaches proposed apply to other types of neural networks. Developing DWSNets was a way of furthering the research and exploring new perspectives on how to best process the weights of neural nets.
Why DWSNets uses Weights & Biases
As a practitioner, having a more transparent and efficient approach to model training often means building better models.
Integrating DWSNets with W&B provided visibility into how the model was converging and performing. The wandb.log function made it easy to track losses, metrics, gradients, parameters, and sample predictions during model development. "We use W&B to monitor all the metrics we care about," said Aviv Navon, co-creator of the repo.
Past that, leveraging W&B also helped Aviv and his team in debugging their optimization process. By including wandb.watch(model,log='all') before their training loop, W&B displayed histograms of their model’s parameters and their gradients at each update step. This helped catch any optimization difficulties that would come from vanishing gradients.
MNIST-INR-classification Experiments
9
What's next for DWSNets?
As with many repositories, keeping the project useful and relevant is key. For DWSNets, that means continually iterating to uncover new ways of expanding the repo. At the same time, the team also wants to revisit previous design choices and assess if that's the best approach or if a better method is needed. Building upon the findings of their current research and addressing its limitations creates opportunities for them or other researchers to improve the quality and validity of any future studies.
How can you get involved with DWSNets?
Reading the paper DWSNets is based on is a great place to start learning more about the project. To get a better sense of the repo, run their introductory Colab. If you're interested in collaborating, reach out to the team on their GitHub page.
Add a comment
Iterate on AI agents and models faster. Try Weights & Biases today.