Weight Space Representation Learning with Neural Fields

TL;DR

The paper investigates whether neural network weights can serve as meaningful data representations by constraining optimization with a pre-trained base model and introducing multiplicative LoRA to create structured, low-dimensional weight spaces for neural fields. It develops a diffusion-based generative model operating directly on these weight representations, augmented by an asymmetric masking strategy to mitigate permutation symmetry, and a hierarchical LoRA-aware diffusion encoder to respect LoRA structure. Across 2D and 3D data, the approach yields improved reconstruction, clear semantic structure in weight space, and state-of-the-art-like generation within the weight-space paradigm, outperforming prior weight-space methods and demonstrating strong discriminative power. The work suggests that carefully parameterized weight spaces can serve as viable, interpretable representations for reconstruction, generation, and analytics, with potential implications for data compression and interpretability in neural representations.

Abstract

In this work, we investigate the potential of weights to serve as effective representations, focusing on neural fields. Our key insight is that constraining the optimization space through a pre-trained base model and low-rank adaptation (LoRA) can induce structure in weight space. Across reconstruction, generation, and analysis tasks on 2D and 3D data, we find that multiplicative LoRA weights achieve high representation quality while exhibiting distinctiveness and semantic structure. When used with latent diffusion models, multiplicative LoRA weights enable higher-quality generation than existing weight-space methods.

Figures (10)

• Figure 1: LoRA based weight space representation with neural fields. Given an input coordinate $\mathbf{p} \in \mathbb{R}^n$, a base neural field is adapted via multiplicative LoRA weights $\boldsymbol{\phi}_i$ to produce signal values $\mathbf{v}_\mathbf{p} \in \mathbb{R}^m$, each weight representing an instance. The LoRA weights themselves form structured representations in weight space, enabling diverse applications including reconstruction, generation, and clustering.
• Figure 2: Diffusion Transformer with hierarchical LoRA layer encoder architecture. For each layer $l$, we treat vector pairs $(\mathbf{a}_l^{(i)}, \mathbf{b}_l^{(i)})$ as tokens. Vector-level positional encodings capture rank dimension indices, followed by multi-head attention that models interactions among the $r$ rank components within the layer. This hierarchical design enables the model to learn both local (within-layer) dependencies among rank components and global (cross-layer) relationships across different layers of the neural field.
• Figure 3: Weight space structure analysis. We measure weight similarity (cosine similarity) and the linear mode connectivity barrier (Chamfer distance) as a function of initialization perturbation strength $\lambda$. Each data point is averaged from $30$ instances, the underlying shades are indicative of standard deviation.
• Figure 4: Qualitative generation results. Generated samples from diffusion models trained on different weight space representations. The top 2 rows show results generated by the Airplane model, followed by 2 rows from the Multi-class model. The bottom rows show 2D FFHQ generations.
• Figure 5: t-SNE visualization of weight spaces. Each point represents one instance from the ShapeNet ten-category dataset, colored by object category. Multiplicative LoRA weight spaces exhibit semantic structure.

Theorems & Definitions (8)

• Theorem 1
• proof
• Remark 1
• Theorem 2
• proof
• Corollary 2.1
• Corollary 2.2
• proof