[Paper Review] Random Weight Factorization Improves the Training of Continuous Neural Representations

Sifan Wang, Hanwen Wang+2 more|arXiv (Cornell University)|Oct 3, 2022

Neural Networks and Applications6 citations

TL;DR

This paper introduces random weight factorization—a simple reparameterization of linear layers in coordinate-based MLPs that improves training by enabling self-adaptive learning rates per neuron. By factorizing weights into scale and direction vectors with random initialization, the method reshapes the loss landscape, mitigating spectral bias and accelerating convergence across diverse tasks in computer vision, graphics, and scientific computing.

ABSTRACT

Continuous neural representations have recently emerged as a powerful and flexible alternative to classical discretized representations of signals. However, training them to capture fine details in multi-scale signals is difficult and computationally expensive. Here we propose random weight factorization as a simple drop-in replacement for parameterizing and initializing conventional linear layers in coordinate-based multi-layer perceptrons (MLPs) that significantly accelerates and improves their training. We show how this factorization alters the underlying loss landscape and effectively enables each neuron in the network to learn using its own self-adaptive learning rate. This not only helps with mitigating spectral bias, but also allows networks to quickly recover from poor initializations and reach better local minima. We demonstrate how random weight factorization can be leveraged to improve the training of neural representations on a variety of tasks, including image regression, shape representation, computed tomography, inverse rendering, solving partial differential equations, and learning operators between function spaces.

Motivation & Objective

• Address the challenge of training continuous neural representations for multi-scale signals, which suffer from slow convergence and poor capture of high-frequency details.
• Overcome spectral bias in coordinate-based MLPs, where networks prioritize learning low-frequency components over fine details.
• Improve training stability and generalization without requiring task-specific hyperparameter tuning or prior signal knowledge.
• Develop a plug-and-play method that enhances standard MLPs with minimal architectural changes and consistent performance gains.

Proposed method

• Propose random weight factorization: decompose each neuron's weight vector as $\bm{w}^{(k,l)} = s^{(k,l)} \cdot \bm{v}^{(k,l)}$, where $s^{(k,l)}$ is a trainable scale and $\bm{v}^{(k,l)}$ is a random direction vector initialized once.
• Use this factorization to reparameterize standard linear layers in MLPs, replacing $\bm{W}^{(l)}$ with $\mathrm{diag}(\bm{s}^{(l)}) \cdot \bm{V}^{(l)}$.
• Train the scale factors $s^{(k,l)}$ and direction vectors $\bm{v}^{(k,l)}$ jointly via gradient descent, enabling each neuron to adapt its own learning rate.
• Demonstrate that this factorization alters the loss landscape, reducing the distance between parameter configurations and enabling faster escape from poor local minima.
• Apply the method as a drop-in replacement in various architectures, including standard MLPs, SIREN, and modified DeepONets for operator learning.
• Integrate the method into physics-informed neural networks (PINNs) and DeepONets by combining it with Fourier feature mappings and residual loss minimization.

Experimental results

Research questions

• RQ1 How does random weight factorization affect the geometry of the loss landscape in continuous neural representations?
• RQ2 To what extent can random weight factorization mitigate spectral bias in coordinate-based MLPs across different signal types?
• RQ3 Can random weight factorization improve training convergence and generalization in tasks involving high-frequency signals, such as PDE solving and inverse rendering?
• RQ4 How does the method perform across diverse domains, including image regression, shape representation, and operator learning, without task-specific tuning?
• RQ5 Does the self-adaptive learning rate induced by the factorization enable faster recovery from poor weight initialization?

Key findings

• Random weight factorization significantly reduces training time and improves convergence speed across all evaluated tasks, including image regression, shape representation, and PDE solving.
• The method enables MLPs to learn high-frequency components more effectively, reducing spectral bias and improving reconstruction accuracy on fine details.
• In the Burgers' PDE task, the physics-informed DeepONet with random weight factorization achieved lower test error compared to standard and SIREN baselines, with improved generalization across unseen initial conditions.
• The approach consistently outperforms standard MLPs and SIREN networks in inverse rendering and computed tomography tasks, demonstrating robustness across signal types.
• The loss landscape becomes flatter and more connected under random weight factorization, enabling faster escape from poor local minima and better final solutions.
• The method requires no additional hyperparameter tuning and delivers consistent improvements across diverse architectures and datasets, including those with complex, multi-scale signals.