Optimizing Rank for High-Fidelity Implicit Neural Representations

Abstract

Implicit Neural Representations (INRs) based on vanilla Multi-Layer Perceptrons (MLPs) are widely believed to be incapable of representing high-frequency content. This has directed research efforts towards architectural interventions, such as coordinate embeddings or specialized activation functions, to represent high-frequency signals. In this paper, we challenge the notion that the low-frequency bias of vanilla MLPs is an intrinsic, architectural limitation to learn high-frequency content, but instead a symptom of stable rank degradation during training. We empirically demonstrate that regulating the network's rank during training substantially improves the fidelity of the learned signal, rendering even simple MLP architectures expressive. Extensive experiments show that using optimizers like Muon, with high-rank, near-orthogonal updates, consistently enhances INR architectures even beyond simple ReLU MLPs. These substantial improvements hold across a diverse range of domains, including natural and medical images, and novel view synthesis, with up to 9 dB PSNR improvements over the previous state-of-the-art.

Using Muon for training INRs

Muon (Jordan et al.) is typically used for the hidden layers only. Thus, instead of using a single parameter group, we define two parameter groups and pass them to the SingleDeviceMuonWithAuxAdam. We recommend tryining a range of learning rates for both Muon (hidden weights) and Adam (auxilary, for bias and input/output weights.)

    # define model (this could also be a FF, Siren etc.) 
    model = ReLU_INR(use_ff=use_ff, num_layers=num_layers, hidden_dim=hidden_dim).to(device)
    
    # Conventional optimization (Adam) for comparison
    optim_adam = torch.optim.Adam(model_adam.parameters(), lr=lr)
    
    # New optimization (Muon)
    from muon import SingleDeviceMuonWithAuxAdam
    
    # Optimizing with Muon - split parameters depending on architecture (you might need to adapt this to your class!)
    muon_params = []
    adam_params = [] # first/last layers + biases go here
    
    # 1. Collect 2D weight matrices - we assume the weights to be ordered (!) in the model - sth to double check :)
    all_matrices = []
    for name, p in model.named_parameters():
        if p.ndim == 2 and p.size(0) > 1 and p.size(1) > 1: # check if vector
            all_matrices.append(p)
        else:
            adam_params.append(p) # Biases, vectors, scalars -> Adam
    
    # 2. Distribute: First & Last -> Adam, Hidden -> Muon
    adam_params.append(all_matrices[0])   # First Layer 
    adam_params.append(all_matrices[-1]) # Last Layer
    muon_params.extend(all_matrices[1:-1]) # Hidden Layers 

    # we do not use weight decay, but you might want to experiment with it
    optim_muon = SingleDeviceMuonWithAuxAdam([
    dict(params=muon_params, use_muon=True, lr=muon_lr, weight_decay=0.0),
    dict(params=adam_params, use_muon=False, lr=aux_adam_lr, betas=(0.9, 0.999), weight_decay=0.0)
    ])
    
    # We recommend using a learning rate scheduler
    scheduler_muon = torch.optim.lr_scheduler.CosineAnnealingLR(optim_muon, T_max=steps)

Giving Muon a try

To easily check out Muon and start training Muon-INRs, we provide a reference notebook here. Make a copy to get started.

To replicate the experiments of our paper, we provide the following instructions.

1. Setup and Data

Obtaining the Data

Please use the following links to obtain the necessary datasets:

Experiment	Dataset Source
Image Reconstruction	Kodak Image Dataset
Audio Reconstruction	Audio Data from Sitzmann et al.
CT Reconstruction	CT Data from Saragadam et al.
SISR Image Data	DIV2K Dataset Kaggle
3D Shapes	Stanford Dataset
NeRF Dataset	Synthetic NeRF Dataset

Environment Setup

1. Standard Dependencies You require a standard PyTorch environment (tested on PyTorch 1.13+), along with script-specific libraries for data loading (e.g., librosa for audio, pydicom/mrcfile for CT). These can be installed via uv/pip/conda. Please note that NeRF and Shape Reconstruction are based on other repositories, and include installation instructions.

2. Installing Muon We use the Muon optimizer (Jordan et al.). Please install it via:

pip install git+https://github.com/KellerJordan/Muon

Note on PyTorch Versions: Muon utilizes the .mT syntax established in PyTorch 1.10. If you must use an older PyTorch version (e.g. for torch-ngp), please substitute all .mT with .transpose(-2, -1) in the function zeropower_via_newtonschulz5 in muon.py.

2. Overfitting Experiments

For specific learning rates and hyperparameters, please refer to the tables provided in the Supplementary Material.

Image Reconstruction: To employ Muon, specify --optimizer muon and set the learning rates for Muon and Adam parameter groups respectively.

python3 fit_image.py \
  --image data/kodim01.png \
  --optimizer muon \
  --muon_lr <MUON_LR> \
  --lr <ADAM_LR>

Audio Reconstruction: Specify the audio file (e.g., gt_bach.wav) via --audio.

python3 fit_audio.py \
  --audio data/gt_bach.wav \
  --optimizer muon \
  --muon_lr <MUON_LR> \
  --lr <ADAM_LR>

3D Shapes (SDF): We base our experiments on the FINER/BACON codebase. We will release a cleaned-up fork soon.

3. Inverse Problems

CT Reconstruction: Run the reconstruction script specifying the image file path via --image.

python3 fit_ct.py \
  --image data/ct_scan.dcm \
  --optimizer muon \
  --muon_lr <MUON_LR> \
  --lr <ADAM_LR>

Single Image Super-Resolution (SISR): Run the SISR script specifying the high-res image path via --image.

python3 fit_sisr.py \
  --image data/div2k/0001.png \
  --optimizer muon \
  --muon_lr <MUON_LR> \
  --lr <ADAM_LR>

NeRF: We base our experiments on the FINER/Torch-NGP codebase. We will release a cleaned-up fork soon.

Citation

If you find this work useful, please cite it as follows:

@misc{mcginnis2025optimizingrankhighfidelityimplicit,
      title={Optimizing Rank for High-Fidelity Implicit Neural Representations}, 
      author={Julian McGinnis and Florian A. Hölzl and Suprosanna Shit and Florentin Bieder and Paul Friedrich and Mark Mühlau and Björn Menze and Daniel Rueckert and Benedikt Wiestler},
      year={2025},
      eprint={2512.14366},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={[https://arxiv.org/abs/2512.14366](https://arxiv.org/abs/2512.14366)}, 
}

When using Muon in your work, please cite:

@misc{jordan2024muon,
      title={Muon: An optimizer for hidden layers in neural networks},
      author={Keller Jordan and Yuchen Jin and Vlado Boza and Jiacheng You and Franz Cesista and Laker Newhouse and Jeremy Bernstein},
      year={2024},
      url={[https://kellerjordan.github.io/posts/muon/](https://kellerjordan.github.io/posts/muon/)}
}