Weight Space Representation Learning via Neural Field Adaptation

This paper proposes that constraining neural field optimization through low-rank adaptation (LoRA) induces structured, high-quality weight space representations that outperform existing methods in reconstruction, generation, and analysis tasks.

The Gist

Imagine you have a massive library of books. Traditionally, if you wanted to store a book, you'd just keep the physical pages. But what if, instead of the pages, you could store the recipe for writing that book? If you had the perfect recipe, you could recreate the book anytime, anywhere.

In the world of Artificial Intelligence (AI), the "recipe" is the neural network weights. These are the billions of tiny numbers inside an AI that tell it how to think. Usually, scientists treat these numbers as a messy, chaotic byproduct of training—like the dust left over after baking a cake. They are hard to read, hard to compare, and hard to use for anything other than the specific task they were trained for.

This paper, "Weight Space Representation Learning," asks a bold question: What if we could turn that messy dust into a clean, organized library of recipes?

Here is the story of how they did it, explained simply.

1. The Problem: The "Messy Room"

Imagine you ask 100 different people to draw a picture of a cat.

• Person A draws it with a pencil, using a specific style.
• Person B draws it with a marker, using a different style.
• Person C draws it upside down.

Even though they all drew a "cat," the actual drawings (the weights) look completely different. If you tried to put them all in a box and ask a computer to find the "cat-ness" in them, it would be confused. The computer sees 100 different messes, not 100 cats. This is called permutation symmetry—the same result can be achieved in a million different, chaotic ways.

2. The Solution: The "Master Blueprint" (The Base Model)

The authors realized that instead of asking everyone to start from scratch (like a blank piece of paper), they should give everyone the same Master Blueprint.

They took a pre-trained AI (the "Base Model") that already knew how to draw general shapes. Then, instead of training a whole new AI for every single image, they just asked the AI to make tiny adjustments to this Master Blueprint to fit the specific image.

Think of it like a custom suit.

• The Base Model is the tailor's mannequin with a standard suit pattern.
• The Adjustments are the specific measurements (shoulders, waist, length) needed for one specific person.

By only saving the measurements (the adjustments) instead of the whole suit, the data becomes much smaller and much more organized.

3. The Secret Sauce: "Multiplicative LoRA" (mLoRA)

This is the paper's biggest innovation.

Usually, when you adjust a model, you do it by adding numbers (like adding a little more salt to a soup). The authors found that for this specific type of "recipe" (called Neural Fields), adding doesn't work well. It creates a tangled mess where the flavors mix up.

Instead, they used Multiplication (mLoRA).

• Analogy: Imagine you have a dimmer switch for a lightbulb.
  • Additive (Old way): You try to make the light brighter by stacking more lightbulbs on top of each other. It gets messy and inefficient.
  • Multiplicative (New way): You just turn the dimmer switch up or down. You are scaling the existing light, not adding new, confusing parts.

By using this "dimmer switch" approach, the adjustments stay clean and organized. The "recipe" for a cat stays distinct from the "recipe" for a dog, even though they share the same base.

4. Breaking the Symmetry: The "Name Tags"

Even with the dimmer switch, there was still a problem. Imagine you have 5 different dimmer switches. It doesn't matter which one you call "Switch 1" and which you call "Switch 5"; the light is the same. This is the "messy room" problem again.

To fix this, the authors used a trick called Asymmetric Masking.

• Analogy: Imagine you have 5 identical twins. To tell them apart, you give them name tags that say "I am the first," "I am the second," etc.
• In the math, they "froze" (locked) certain parts of the adjustment so that the AI couldn't swap them around. This forced every "recipe" to have a unique, consistent order.

5. The Results: Why This Matters

Once they organized the "recipes" (weights) this way, amazing things happened:

1. Better Reconstruction: They could recreate the original images (faces, 3D chairs) with incredible detail using very little data.
2. Generation (The Magic Trick): They trained a "Generator AI" (a Diffusion Model) to learn the distribution of these organized recipes.
   • The Result: The Generator could create brand new faces and 3D objects it had never seen before, just by mixing and matching these organized recipes.
   • The Breakthrough: Previous methods failed when trying to generate high-quality, complex images (like human faces). This method succeeded where others failed.
3. Understanding: Because the recipes were so organized, a computer could easily tell the difference between a "chair" recipe and a "table" recipe. The "weight space" actually made sense to the AI.

The Big Picture

This paper changes how we view AI weights.

• Old View: Weights are a chaotic, unreadable mess of numbers.
• New View: Weights are structured, semantic representations. They are like a library of unique, organized blueprints.

By using a Master Blueprint and a Dimmer Switch (mLoRA) approach, the authors turned the "dust" of AI training into a powerful new way to store, understand, and create data. It's like realizing that if you organize your recipe cards correctly, you don't just have a cookbook; you have a machine that can invent new dishes on the fly.

Technical Summary

1. Problem Statement

The paper addresses a fundamental question in deep learning: Can neural network weights themselves serve as meaningful, structured representations of data?

Traditionally, weights are viewed as opaque byproducts of optimization. While recent work has explored merging or generating weights, significant challenges remain:

• Ambiguity and Symmetry: Neural networks exhibit permutation symmetry (reordering neurons yields the same function) and scaling invariance. This causes functionally identical models to occupy vastly different, disconnected regions in weight space, creating a multi-modal distribution that is difficult to learn.
• High Dimensionality: Raw weights are high-dimensional vectors, suffering from the "curse of dimensionality," which hinders learning effective distributions for generation.
• Lack of Structure: Standard weight representations (e.g., standalone MLP weights) often lack semantic organization, making them poor inputs for generative models like diffusion models.

The authors aim to transform these "chaotic" parameters into structured, semantic weight space representations that can be used for reconstruction, generation, and discriminative tasks.

2. Methodology

The proposed framework, Weight Space Representation Learning via Neural Field Adaptation, relies on three core components:

A. Base Model and Inductive Bias

Instead of training standalone networks from scratch for every data instance, the authors use a pre-trained base neural field (an Implicit Neural Representation or INR).

• Base Architecture: A coordinate-based neural field with multiplicative weight modulation (similar to StyleGAN or pi-GAN).
• Training: The base model is trained using a Variational Autoencoder (VAE) paradigm (specifically, an autodecoder), where a shared network learns dataset-level priors, and per-instance latent codes are optimized.

B. Multiplicative Low-Rank Adaptation (mLoRA)

The core innovation is the use of Low-Rank Adaptation (LoRA) to represent individual data instances, but with a critical modification:

• Standard LoRA (Additive): Updates weights via $W' = W + BA$. The authors argue this exacerbates feature entanglement in neural fields, which synthesize signals through additive composition.
• Proposed mLoRA (Multiplicative): Updates weights via element-wise multiplication:
  $$W' = W \odot (BA)$$
  where $\odot$ is the Hadamard product.
• Rationale: This multiplicative formulation aligns with the modulation mechanisms in generative neural fields. It scales existing features rather than injecting new signal components, thereby preserving channel structure and avoiding further entanglement. This results in a more organized weight space.

C. Symmetry Breaking via Asymmetric Masking

To address the permutation symmetry inherent in LoRA (where permuting the rank dimensions $r$ yields the same function), the authors apply asymmetric masking:

• For each layer, specific entries in the $A$ matrix are randomly frozen (zeroed out for mLoRA) across all instances and training runs.
• This breaks the $GL(r)$ equivalence class, forcing the optimization to converge to a canonical representation and collapsing the multi-modal distribution into a smoother, learnable manifold.

D. Generative Modeling

To generate new data, the authors train a Diffusion Transformer (DiT) directly on the mLoRA weights:

• Hierarchical Encoder: A specialized encoder treats LoRA vector pairs $(a_i, b_i)$ as tokens. It uses Vector-level positional encodings (for rank indices) and Layer-level positional encodings to capture both intra-layer rank dependencies and inter-layer semantic relationships.
• Process: The model learns the distribution of mLoRA weights. During inference, it generates new weight vectors, which are then used to instantiate new neural fields that decode the data.

3. Key Contributions

1. Proof of Concept: Demonstrates that independently optimized neural network weights, when properly constrained, can serve as effective, semantic data representations.
2. Multiplicative LoRA (mLoRA): Introduces a multiplicative adaptation mechanism for neural fields that outperforms standard additive LoRA and standalone MLPs in terms of representation quality and generation capability.
3. Structured Weight Space: Shows that combining mLoRA with asymmetric masking creates a weight space with high linear mode connectivity and clear semantic clustering.
4. State-of-the-Art Generation: Achieves the first successful high-resolution natural image generation (FFHQ) and high-quality 3D shape generation (ShapeNet) purely via weight-space diffusion, outperforming prior methods like HyperDiffusion.

4. Experimental Results

The authors evaluated their method on 2D images (FFHQ) and 3D shapes (ShapeNet) across three tasks:

• Reconstruction:

  • mLoRA-Asym achieved the highest reconstruction quality (PSNR for 2D, Chamfer Distance for 3D) while maintaining a compact parameter count.
  • It outperformed standalone MLPs and additive LoRA, confirming the benefit of the base model's inductive bias and the multiplicative formulation.

• Weight Space Structure Analysis:

  • Stability: mLoRA-Asym showed high cosine similarity between weights optimized from different random initializations, indicating convergence to a single linear mode.
  • Linear Mode Connectivity: The barrier height between different solutions was significantly lower for mLoRA-Asym compared to additive LoRA, suggesting a smoother, more connected weight manifold.

• Generative Modeling (Diffusion):

  • Metrics: mLoRA-Asym achieved the best scores on Fréchet Distance (FD), Maximum Mean Discrepancy (MMD-G/P), and coverage metrics.
  • Qualitative: It successfully generated diverse, high-frequency details in faces (FFHQ) and 3D objects. In contrast, additive LoRA and standalone MLPs failed to generate recognizable images or shapes, often collapsing or producing noise.
  • Comparison: Outperformed HyperDiffusion (a prior state-of-the-art weight-space method) significantly, particularly on multi-category 3D datasets.

• Discriminative Tasks (Classification/Clustering):

  • Using mLoRA weights as input to a linear classifier yielded 90% accuracy on ShapeNet 10-category classification, far surpassing MLPs and additive LoRA.
  • t-SNE visualizations confirmed that mLoRA weights form distinct, semantically separated clusters corresponding to object categories.

5. Significance and Impact

• Paradigm Shift: The work challenges the view of weights as mere optimization byproducts, establishing them as a viable, high-quality data modality for representation learning.
• Efficiency: By leveraging a pre-trained base model and low-rank adaptation, the method achieves high-quality generation with significantly fewer parameters than training full networks from scratch.
• Modality Agnosticism: The approach is applicable to diverse data types (2D, 3D, potentially 4D) without changing the core architecture, relying only on the coordinate-to-value mapping of INRs.
• Future Directions: The success of mLoRA suggests that multiplicative interactions are crucial for disentangling features in neural fields, offering new insights into how to structure weight spaces for generative AI.

In summary, this paper provides a robust framework for turning the "chaos" of neural weights into a structured, semantic language, enabling high-fidelity generation and analysis directly within the parameter space of neural fields.