Technical Summary: Optimizing Rank for High-Fidelity Implicit Neural Representations

1. Problem Statement

Implicit Neural Representations (INRs) have become a versatile framework for modeling diverse signal modalities, including shapes, images, scenes, and medical data. However, vanilla Multi-Layer Perceptrons (MLPs) with standard ReLU activations are widely believed to suffer from "spectral bias," an intrinsic architectural limitation that prevents them from representing high-frequency content effectively.

Consequently, the field has largely focused on architectural interventions to overcome this bias, such as coordinate embeddings (e.g., Fourier Features) and specialized activation functions (e.g., SIREN, WIRE). While these methods improve expressiveness, they introduce new challenges: highly expressive architectures often lack implicit regularization, leading to oscillatory patterns in ill-posed inverse problems, while low-frequency biased architectures may fail to fit fine details.

This paper challenges the prevailing notion that low-frequency bias is an inherent architectural flaw. Instead, the authors argue it is a symptom of stable rank degradation during training. They posit that the inability of vanilla MLPs to fit high-frequency signals stems from the collapse of the effective dimensionality of weight updates, rather than the network architecture itself.

2. Methodology: A Stable Rank Perspective

The authors shift the focus from architectural design to optimization dynamics, utilizing the concept of stable rank as a unifying metric.

2.1 Theoretical Framework

The paper establishes a linear-algebraic link between activations and weight updates. For a layer $l$, the gradient with respect to weights is defined as the outer product of the upstream gradient and the input activations: $\nabla W_l \mathcal{L} = G_{l+1} H_l^\top$.

• Rank Constraint: The row space of the weight update is contained within the span of the batch activations ($rowspace(\nabla W_l \mathcal{L}) \subseteq colspace(H_l)$).
• Degradation: In coordinate-based networks, input activations often lie in low-dimensional subspaces. As signals propagate through layers, nonlinearities can locally increase rank, but without specific interventions, the network tends to operate on progressively lower-dimensional subspaces, leading to "rank collapse."
• Stable Rank: The authors use stable rank, $s(A) = \frac{\|A\|_F^2}{\|A\|_2^2}$, to quantify the effective dimensionality of these updates. A low stable rank indicates that the update is concentrated in a few singular directions, limiting the network's ability to learn high-frequency details.

2.2 Reinterpreting Existing Interventions

The authors reinterpret common architectural modifications through this lens:

• Input-Level (e.g., Fourier Features): These expand the input batch matrix $H_0$, increasing the upper bound on the stable rank of the first-layer update.
• Activation-Level (e.g., SIREN, BatchNorm): These aim to maintain the span of activations $H_l$ throughout the network, preventing the rank from diminishing with depth.

2.3 Proposed Solution: The Muon Optimizer

The paper proposes that the most direct remedy is to intervene at the optimizer level. They introduce Muon, an orthogonalizing optimizer designed to explicitly enforce isotropy in gradient updates.

• Mechanism: For a weight update matrix $U$, Muon maps it to the nearest semi-orthogonal matrix $Ortho(U)$ (via Newton-Schulz iteration, approximating SVD). This forces the update to be distributed evenly across all singular directions, maximizing the stable rank of the update to $\min(d_l, d_{l+1})$.
• Implementation: Muon is applied to hidden-layer weight matrices, while standard Adam is retained for biases and input/output parameters to maintain stability and efficiency. This approach is model-agnostic and operates on the actual per-step gradients.

3. Key Contributions

1. Stable Rank as a Unifying Metric: The paper proposes stable rank as a key measure of expressiveness in INRs, arguing that rank degradation during training, not architectural limitations, is the primary cause of low-frequency bias in vanilla MLPs.
2. Unifying Framework: It provides a framework that explains the effectiveness of diverse architectural modifications (coordinate embeddings, alternative activations) as methods to increase or preserve stable rank.
3. Muon Optimizer: The authors propose Muon, an orthogonalizing optimizer that explicitly prevents rank collapse, serving as a natural remedy for high-fidelity signal learning without requiring architectural changes.
4. Empirical Validation: Extensive experiments demonstrate that Muon significantly improves performance across all tested architectures (ReLU MLPs, SIREN, FINER, etc.) and modalities (images, audio, 3D shapes, medical CT, NeRFs).

4. Experimental Results

The authors evaluated Muon against Adam (and Lion) across four domains:

• Image Overfitting (Kodak Dataset):

  • Muon improved the PSNR of vanilla ReLU MLPs by ~7–8 dB (from ~22.8 to ~30.0 dB), bringing them to parity with Fourier Feature Networks (FFNs) trained with Adam.
  • Significant gains (2–4 dB) were also observed for Gaussian MLPs, SIREN, and FINER.
  • Stable rank analysis showed Muon maintains high rank in hidden layers, whereas Adam and Lion exhibit significant rank collapse.

• Audio Overfitting:

  • Muon achieved substantially higher Signal-to-Noise Ratio (SNR) across all architectures. For a ReLU FFN, SNR improved by +12.5 dB.
  • Vanilla ReLU MLPs failed to reconstruct long audio clips with Adam but succeeded with Muon (on shorter segments), suggesting the failure was due to optimization instability rather than capacity limits.

• Signed Distance Fields (SDF):

  • On the Armadillo and Thai Statue datasets, Muon reduced Chamfer Distance and improved IoU, particularly for small architectures (8x32).

• Inverse Problems:

  • CT Reconstruction: Muon improved PSNR by up to +7 dB for ReLU MLPs in sparse-view CT reconstruction, reducing artifacts in heterogeneous structures.
  • Single Image Super-Resolution (SISR): Muon made ReLU FFNs particularly effective, outperforming architectures with strong high-frequency biases (like FINER) in some metrics, likely due to the beneficial low-frequency bias of ReLU combined with orthogonal updates.
  • Neural Radiance Fields (NeRF): Muon reduced reconstruction artifacts and improved PSNR by up to +9 dB on synthetic scenes (e.g., Chair, Hotdog), yielding more faithful scene reconstructions.

5. Significance and Claims

The paper claims that optimization is a third, underexplored axis of INR design, complementary to input and activation-level interventions.

• Challenging Architectural Determinism: The results challenge the perception that architectural changes are the primary driver of INR performance. The authors demonstrate that optimization alone can significantly impact model fidelity, enabling simple ReLU MLPs to represent high-frequency signals with high fidelity.
• Additive Effects: The paper argues that input-level, activation-level, and optimization-level interventions act on distinct, complementary axes. Combining them (e.g., using Muon with Fourier Features) yields the most faithful reconstructions.
• Practical Impact: By using Muon, practitioners can achieve high-fidelity results with simpler, more robust architectures (like ReLU MLPs) that are less prone to the oscillatory patterns often seen in highly expressive, specialized networks.

6. Limitations

The authors acknowledge several limitations:

• Target Rank: The stable rank framework serves as a diagnostic but does not predict an optimal target rank. In some ill-posed inverse problems, a lower induced stable rank may act as a beneficial implicit regularizer.
• Approximation: Muon uses the Newton-Schulz approximation for orthogonalization. The paper does not deeply study the implications of this approximation, viewing Muon as a point in a broader design space rather than a final optimal solution.
• Tuning Complexity: Muon introduces additional tuning complexity regarding learning rate schedules (separate rates for Muon and auxiliary Adam parameters).
• Expressiveness Gap: While Muon narrows the gap, a ReLU MLP trained with Muon remains less expressive than specialized architectures like SIREN or FINER in all settings.