🤖 machine learning

What Cohort INRs Encode and Where to Freeze Them

This paper investigates cohort-trained Implicit Neural Representations (INRs) by identifying that the optimal layer for freezing corresponds to the highest weight stable rank and revealing, through sparse autoencoder analysis, that SIREN and Fourier-feature MLPs encode fundamentally different representations—localized coordinate tiles versus image-spanning contours—thereby providing the first mechanistic understanding of transferable representations in INRs.

Original authors: Vasiliki Sideri-Lampretsa, Sophie Starck, Robbie Holland, Julian McGinnis, Daniel Rueckert

Published 2026-05-12

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Vasiliki Sideri-Lampretsa, Sophie Starck, Robbie Holland, Julian McGinnis, Daniel Rueckert

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: The "Universal Translator" Problem

Imagine you have a team of expert artists (called INRs or Implicit Neural Representations). Each artist is hired to draw a single, specific picture. To draw a face, they spend hours learning exactly how that one face looks. To draw a brain scan, they start from scratch and learn that all over again.

This is slow and expensive. The paper asks: Can we teach these artists a "universal skill" first? If we train them on a whole group (a "cohort") of faces, can we then reuse that training to draw a new face much faster?

The researchers found that yes, you can reuse the training, but you have to be very careful about which parts of the artist's brain you freeze (lock in place) and which parts you let them relearn.

Discovery 1: The "Sweet Spot" for Freezing

When you train an artist on a group of 10 faces, they develop a shared "encoder" (the part of the brain that understands coordinates and shapes) and a unique "decoder" (the part that draws the specific face).

The researchers asked: If we want to draw a new face, how much of the old training should we keep frozen?

The Old Way: The standard method (called STRAINER) tried to update every single layer of the network when drawing a new picture.
The New Finding: The researchers discovered a "sweet spot." There is a specific layer in the middle of the network where the "weight" of the information is most stable and useful.
- The Analogy: Imagine the network is a multi-story building. The ground floor is the foundation, the top floor is the roof, and the middle floors are the living spaces. The researchers found that if you freeze the building up to the floor with the strongest, most complex structural beams (measured by something called "stable rank"), you get the best results.
- The Result: By freezing the network exactly at this "strongest floor" and letting the rest adapt, the new drawing is often better than the old method. It's like keeping the solid foundation and walls of a house but letting the interior decorator rearrange the furniture for the new owner.

Discovery 2: Two Different Types of "Muscle Memory"

The paper studied two different types of artists (network architectures): SIREN and FFMLP. Even though they both learned to draw faces equally well, they learned in completely different ways. The researchers used a special tool called a Sparse Autoencoder (SAE) to peek inside their brains and see what they were actually "thinking."

Think of the SAE as a translator that breaks down the artist's thoughts into individual "atoms" or building blocks.

1. SIREN: The "Tile Makers"

How they think: SIREN artists learned to tile the coordinate plane.
The Analogy: Imagine a floor covered in a grid of small, distinct tiles. Each tile covers a specific, small patch of the floor. If you want to draw a nose, the artist only uses the tiles in the "nose region." If you want to draw a chin, they use the "chin region."
The Behavior: These "atoms" are localized. They don't care what the picture is (a face or a brain); they just care about where they are on the grid.
The Test: If you remove one of these "tiles" (ablation), the damage is tiny and stays only in that specific patch. The rest of the picture remains perfect.

2. FFMLP: The "Ghost Tracers"

How they think: FFMLP artists learned to trace the contours of the specific images they memorized.
The Analogy: Imagine these artists don't use a grid. Instead, they have "ghost outlines" of the 10 faces they memorized floating in their minds. Every time they draw, they are tracing these specific ghost outlines.
The Behavior: These "atoms" are image-spanning. They cover the whole picture and trace the specific curves of the memorized faces.
The Test: If you remove just one of these "ghost outlines" (even if there are 4,000 of them), the whole picture falls apart. The quality drops drastically (by up to 10.6 dB) because the artist lost a crucial part of the global shape.

Why This Matters (According to the Paper)

We know where to freeze: We don't need to guess. We can calculate the "stable rank" of the layers to find the exact point where the network has learned the most reusable structure.
We know what is being learned:
- SIREN learns a general, content-independent map of space (good for generalization).
- FFMLP learns specific, memorized shapes (good for fitting, but fragile).
Causality: The paper proved that these "atoms" aren't just random noise; they are the actual cause of the image. Removing a SIREN atom hurts a small spot; removing an FFMLP atom hurts the whole image.

Summary

The paper reveals that when training AI to represent signals (like images), there is a specific "middle layer" that holds the most transferable knowledge. Furthermore, different AI architectures learn this knowledge in opposite ways: one builds a flexible, local grid (SIREN), while the other memorizes global, specific shapes (FFMLP). Understanding this helps us build better systems that generalize to new tasks rather than just memorizing old ones.

Technical Summary: What Cohort INRs Encode and Where to Freeze Them

Problem Statement

Implicit Neural Representations (INRs) are coordinate-based multi-layer perceptrons (MLPs) that learn continuous functional approximations of signals. A primary limitation of standard INRs is that each network fits a single signal from scratch, with features finely tuned to that specific signal's morphology. This results in high per-signal optimization costs. Recent work, specifically STRAINER, attempts to mitigate this by training a "cohort" of INRs that share early encoder layers while maintaining per-signal decoders. The shared encoder is then used as an initialization for new signals.

However, two critical questions remain unanswered regarding this transfer learning paradigm:

Which layers of the shared encoder actually carry transferable representations?
What do those specific layers encode?

Current STRAINER protocols update every layer during test-time fitting, preventing the isolation of per-layer transferability. Furthermore, while transfer learning in vision and language often localizes transferable features to intermediate layers, it is unclear if this holds for INRs, particularly given that dominant backbones like SIREN and Fourier-feature MLPs (FFMLP) encode coordinate positions through fundamentally different mechanisms.

Methodology

The authors propose a two-pronged diagnostic approach to analyze cohort-trained INRs on three datasets (CelebA-HQ, OASIS-MRI, Kodak) using two backbones (SIREN and FFMLP).

1. Locating Transferable Representations via Weight Stable Rank

To identify which layers transfer best, the authors introduce a test-time freeze sweep protocol. Instead of updating all layers, they freeze the encoder up to a specific depth $\tau$ and reinitialize the remaining layers and the decoder head. They evaluate the performance (PSNR) of new signals across all possible freeze depths.

To predict the optimal freeze depth a priori, they utilize the weight stable rank ( $sr(W_\ell)$ ) of the trained encoder layers. Unlike activation rank, which depends on the specific input distribution, weight stable rank is an intrinsic property of the learned weights, measuring the effective dimensionality of the subspace spanned by the weight matrix. The hypothesis is that the layer with the highest weight stable rank represents the "coordinate embedding" layer, which learns the most reusable structural information across the cohort.

2. Characterizing Encoded Representations via Sparse Autoencoders (SAEs)

To determine what the transferable layers encode, the authors apply Sparse Autoencoders (SAEs) to the hidden activations of the trained encoders. SAEs decompose dense activations into an overcomplete dictionary of sparse "atoms."

Architecture: A TopK SAE with $n=4096$ atoms and sparsity $k=32$ is fitted to the activations of each hidden layer.
Diagnostics:
- Atom Maps: Visualizing the spatial firing patterns of atoms across the coordinate domain.
- Single-Atom Ablations: Systematically removing individual atoms (zeroing their contribution) to measure the causal impact on the final signal reconstruction (PSNR drop).

Key Results

1. Optimal Freeze Depth Coincides with Peak Weight Stable Rank

The experiments reveal that the optimal freeze depth $\tau^\star$ consistently aligns with the layer exhibiting the peak weight stable rank in the cohort-trained encoder.

SIREN: The optimal freeze depth is $\tau^\star = 1$ . The first trainable layer ( $W_1$ ) acts as an emergent positional embedding, achieving the highest weight rank.
FFMLP: The optimal freeze depth is $\tau^\star = 0$ . Since FFMLP uses a fixed Fourier projection before the first trainable layer, the embedding occurs externally, and the first trainable layer ( $W_0$ ) operates on this high-rank input.
Performance: Freezing at this specific depth matches or significantly improves upon the standard STRAINER recipe (which updates all layers). Notably, on the high-resolution Kodak dataset, the freeze protocol yields gains of +3.1 dB for SIREN and +1.3 dB for FFMLP on average, as it prevents the model from wasting the limited test-time optimization budget on relearning the coordinate embedding.

2. Qualitatively Different Dictionaries: Localization vs. Memorization

Despite achieving comparable cohort-fitting quality, SIREN and FFMLP learn fundamentally different internal representations:

SIREN (Localized): The SAE atoms are spatially localized, tiling the coordinate plane. Each atom fires in a confined region independent of the specific cohort content. This suggests SIREN learns a content-independent positional code.
FFMLP (Image-Spanning/Memorized): The atoms are image-spanning and diffuse. They trace the contours of memorized cohort signals (e.g., facial features in CelebA, brain anatomy in OASIS). The dictionary concentrates on a few high-magnitude atoms that cover large portions of the image.

3. Causal Validation via Ablation

Single-atom ablations confirm the causal role of these dictionaries:

SIREN: Ablating a single atom causes PSNR drops (2–3 dB) confined strictly to the spatial region where that atom fires. Other atoms compensate for the rest of the image.
FFMLP: Ablating a single atom causes global degradation, with PSNR drops of up to 10.6 dB across the entire image. This indicates that FFMLP relies on a few "super-atoms" to reconstruct global signal structure, effectively memorizing the cohort content.

Significance and Claims

The paper claims to provide the first mechanistic account of what transfers in cohort-trained INRs.

Diagnostic Tool: The weight stable rank metric offers a computable, target-data-free method to identify the optimal transferable layer depth, eliminating the need for hyperparameter sweeps.
Interpretability: By decomposing INR activations into sparse atoms, the authors turn "black box" activations into inspectable dictionary atoms, revealing that SIREN encodes reusable positional structure while FFMLP encodes memorized content.
Architectural Insight: The findings suggest that the architectural choice (SIREN vs. FFMLP) dictates whether the network learns generalizable structural priors or overfits to memorized cohort signals.
Future Direction: These tools open a path toward characterizing INR encodings and designing architectures specifically for generalization rather than memorization.

The authors note limitations, including that the analysis is currently restricted to SIREN and FFMLP backbones and does not yet characterize the semantic meaning of individual atoms (e.g., edges, colors) or explore steering reconstructions beyond single-atom ablations.