PaperlayerView exact revision on arXiv

Source-linked AI summary

Atelier: Learning Local Self-Supervised Features for CryoEM Volumes via Hypernetworks

Phillip Lo, Sudarshan Babu, Dari Kimanius, Aly A. Khan

arXiv:2609.30569v1cs.AIstat.ML

TL;DR

CryoEM maps require substantial modeling and expert validation to obtain biological annotations, while fixed-grid features and independently fitted INRs have limitations for continuous, aligned feature extraction. Atelier uses a transformer hypernetwork to amortize INR generation, and its coordinate-conditioned features improve voxel-level annotation across eight tasks. The results support amortized INRs as a geometry-aware representation for cryoEM analysis.

  • Problem

    CryoEM maps lack annotations such as residue identities and secondary structure, while fixed-grid methods require interpolation off-grid and independently fitted INRs are costly and unaligned.

    Source: §1 · pp. 1–2

  • Method

    Atelier trains a transformer hypernetwork to generate cryoEM volume INRs, whose coordinate-queryable activations are used as auxiliary features for a 3D U-Net.

    Source: §1 · p. 2

  • Results

    Coordinate-conditioned INR activations improved voxel-level annotation accuracy on eight of eight tasks over a volume-only baseline.

    Source: §1 · p. 3, §5 · p. 12

  • Takeaways & Limitations

    The results support amortized implicit neural representations as a geometry-aware feature primitive for cryoEM analysis.

    Source: §5 · pp. 11–12

    Pretraining requires roughly 720 H100-hours, and this cost remains a barrier to iteration and scaling to larger volumes or higher-resolution grids.

    Source: C Limitations and Future Work · p. 22

Abstract

from arXiv · show

CryoEM map interpretation requires features that are spatially localized, consistent across samples, and informative across spatial scales. Most deep learning methods for map annotation extract features from fixed voxel grids. However, implicit neural representations (INRs) are able to model volumetric data as scale-agnostic, coordinate-conditioned functions. INRs are therefore attractive for cryoEM, but fitting a separate INR for each map is too expensive for large-scale feature extraction and produces representations that are not aligned across samples. We introduce Atelier, a self-supervised framework that amortizes INR fitting for reconstructed cryoEM maps. Pretrained on 5,439 Electron Microscopy Data Bank maps, Atelier is a transformer-based hypernetwork that generates high-fidelity reconstructions across a wide range of protein structures, including large multi-subunit assemblies. Beyond reconstruction, the INR generated by the pretrained transformer exposes a continuous, local feature field through its intermediate activations at any spatial query point, a property that voxel grid and patch-tokenizer architectures do not naturally provide. Used as auxiliary channels to a 3D nested U-Net annotation head trained from scratch, these coordinate-conditioned features improve performance on eight voxel-level property prediction tasks over a volume-only baseline. Our results demonstrate that amortized implicit neural representations are an effective primitive for geometry-aware analysis of cryoEM data.

1 Introduction

Atelier addresses the cost and alignment limits of per-map implicit representations by amortizing their generation with a self-supervised hypernetwork. Its generated local features improve voxel-level annotation across eight tasks when added to a 3D U-Net.

1 Introduction

  • CryoEM maps do not specify atomic models, residue identities, secondary structure, or ligand-binding sites, and extracting these annotations remains a computational and expert-validation bottleneck.Interpretation needs both local density detail and broader context, while common voxel-grid methods require interpolation for off-grid predictions.Source: §1 · pp. 1–2
  • Per-map INR fitting requires hundreds to thousands of optimization steps, while independently fitted representations are not aligned for use as shared feature extractors.Hypernetworks can amortize fitting, but prior cryoEM applications focused primarily on reconstruction from particle images.Source: §1 · p. 2
  • Atelier uses a self-supervised transformer hypernetwork to generate per-volume INRs in one forward pass, amortizing INR fitting across structures.On 689 held-out EMDB volumes, it achieves an average normalized area under the Fourier shell correlation curve of 0.634 ± 0.093.Source: §1 · pp. 2–3
  • Annotation accuracy improved on eight of eight voxel-level tasks when coordinate-queryable INR activations augmented a 3D U-Net, compared with a volume-only baseline.The evaluation also uses a capacity-matched transformer-only autoencoder to separate the contributions of pretraining and the INR formulation.Source: §1 · pp. 2–3
  • A capacity-matched transformer-only ablation attributes most downstream gain to self-supervised pretraining; the INR formulation typically adds a smaller improvement and enables coordinate-conditioned, resolution-agnostic inference.The transformer-only model cannot perform that coordinate-conditioned inference.Source: §1 · p. 3

2 Background and Related Work

Prior cryoEM annotation methods use voxel-grid representations, while INRs offer continuous coordinate-conditioned representations but are costly to fit independently. Hypernetworks provide a route to amortized, aligned neural fields whose local activations can serve as annotation features.

Grid-Based Representation Learning in CryoEM

  • 3D U-Nets have supported cryoEM annotation, but their features are tied to voxel grids and off-grid predictions require interpolation.Grid resolution also affects memory and compute costs.Source: §2 · p. 3

Implicit Neural Representations

  • INRs represent signals as continuous coordinate-conditioned functions and can capture sub-voxel geometry, but independent fitting is expensive and yields unaligned features across instances.Earlier cryoEM INR work primarily recovered density from raw particle images rather than learning transferable representations from reconstructed maps.Source: §2 · p. 3

Hypernetworks and Amortized INRs

  • Hypernetworks amortize INR parameter generation with a single forward pass, and their shared mapping induces aligned representations across input volumes.Prior work applied the approach across shapes, images, audio, and structured scientific data, while CryoHype modeled cryoEM heterogeneity.Source: §2 · pp. 3–4

Intermediate Activations as Features

  • INR activations queried at a coordinate encode localized geometry, allowing Atelier to use coordinate-aligned features for voxel-level annotation without task-specific pretraining supervision.This contrasts with voxel encoders that require pooling for localized predictions.Source: §2 · p. 4

3 Method

Atelier preprocesses cryoEM maps, then trains a hypernetwork to generate INRs that reconstruct each map. The resulting INR activations provide spatially localized features queryable at arbitrary coordinates.

3.1 Data preprocessing

  • Maps are denoised, normalized to [0, 1], padded to D = 96, and resampled to 3 Å per voxel; maps with foreground voxels exceeding the padded size are discarded.Background voxels are set to 0.Source: §3.1 · p. 4

3.2 Self-supervised INR pretraining

  • Self-supervised pretraining learns a hypernetwork that maps each discrete volume to INR parameters, with the INR trained to reproduce the volume pointwise.The reconstruction objective combines spatial mean squared error with Fourier-domain error weighted by squared normalized frequencies to emphasize high-frequency components.Source: §3.2 · p. 4

3.3 The Atelier architecture

  • The architecture uses a sparse-volume tokenizer, a transformer that contextualizes volume and weight tokens, and an INR decoder that renders density at queried coordinates.The tokenizer discards background-only patches; the INR is an MLP with ReLU activations and Fourier coordinate encoding.Source: §3.3 · pp. 5–6, Figure 2 · p. 5

3.4 Local Feature Extraction

  • Post-activation INR hidden states at a query coordinate form a feature that is resolution-independent and geometrically localized, including at sub-voxel positions.Features for a set of points can be formed by concatenating their activations.Source: §3.4 · p. 6

4 Experiments

Across four experiments, Atelier reconstructs held-out cryoEM volumes and produces local features whose learned locality and cross-volume consistency are examined. Adding its fine features improves voxel-level property prediction over volume-only inputs across all eight tasks.

4.1 Pretraining

  • 0.634 ± 0.093 nAUFSC is achieved across 689 test volumes; the TCR/CD3/pMHC example reaches 0.621, comparable to the test-set mean.Atelier reconstructs the large asymmetric assembly in a single forward pass.Source: §4.1 · pp. 6–7
  • Pretrained features are approximately invariant to global cube symmetries and often remain similar within protein domains while varying between domains.The visualizations suggest the features can partially segment novel proteins.Source: §4.1 · pp. 7–8, Figure 4 · p. 8

4.2 Demonstrating the Locality of Atelier Fine Features

  • On 16 held-out volumes, local perturbations affect trained-model activations 18×–34× more than distant perturbations, while untrained local and distant effects are nearly equal.Distant perturbations were selected as a mass-matched control.Source: §4.2 · pp. 8–9, Table 1 · p. 9

4.3 Using Local Features for Downstream Property Prediction

  • Fine embeddings pool frozen INR activations within a 3-voxel-radius sphere around annotated voxels, whereas coarse embeddings assign each voxel its patch token.The coarse transformer-only autoencoder is matched to the INR in latent dimension.Source: §4.3 · p. 9
  • Fine-feature augmentation outperforms volume-only EMNUSS on eight of eight tasks and coarse-feature augmentation on seven of eight.The comparison uses test-set metrics averaged over labeled voxels.Source: Table 2 · p. 10, §4.3 · p. 10

4.4 Case Study: Comparing Atelier Features Across Related Structures

  • For 80 aligned corresponding points, Atelier feature cosine similarity tracks structural relatedness across base, same, homologous, and more distant GPCR structures.The four structures were absent from training, and their volume-mean features were subtracted before comparison.Source: §4.4 · pp. 10–11, Table 3 · p. 11

5 Conclusion

Atelier uses a transformer hypernetwork to generate INRs whose intermediate activations provide localized features queryable at any spatial coordinate. These features improve voxel-level property prediction across eight tasks, while the paper identifies interface analysis and related applications as future directions.

5 Conclusion

  • A transformer hypernetwork generates high-resolution cryoEM INRs whose intermediate activations yield localized feature vectors queryable at any spatial coordinate.The paper describes these features as geometrically rich.Source: §5 · p. 11
  • Coordinate-conditioned Atelier activations improve a baseline model consistently on all eight voxel-level property prediction tasks.The features are extracted from generated INRs and can be queried at any spatial coordinate.Source: §5 · pp. 11–12
  • Atelier’s continuous features provide a foundation for future interface prediction, binding-site identification, and structure-guided protein engineering.The discussion highlights interfaces in large protein complexes as a potential application requiring sub-voxel, side-chain-level precision.Source: §5 · p. 12

A Data Preparation

The data preparation pipeline uses curated EMDB maps with fitted atomic structures and applies denoising so the tokenizer can focus on patches containing signal.

A Data Preparation

  • The pipeline starts from curated EMDB IDs with fitted PDB maps and resamples density maps to 3 Å isotropic voxels.Figure 6 illustrates denoising on four uncurated volumes; thresholding removes dust so the tokenizer can focus on patches containing signal.Source: A Data Preparation · p. 16, Figure 6 · p. 17

A.1 Denoising CryoEM Volumes

The appendix describes denoising resampled cryoEM volumes and generating voxel-level labels from fitted atomic structures. The dataset is size-filtered, and the labels cover classification and regression properties.

A.1 Denoising CryoEM Volumes

  • The denoising pipeline thresholds background and removes connected components smaller than 1% of the largest component.It first uses Otsu thresholding; when the tight foreground bounding box reduces volume by less than 10%, it sweeps 50 threshold values and selects the one maximizing the discontinuity in data reduction.Source: §A.1 · p. 17
  • After denoising, voxel values are scaled to [0, 1], volumes are padded to (96, 96, 96), and larger volumes are discarded, leaving 5439/684/689 train/validation/test samples.Splits were generated randomly before size filtering.Source: §A.1 · p. 17

A.2 Label Generation

  • Fitted PDB structures are aligned to cryoEM voxel grids to assign voxel-level labels for secondary structure, amino acid, molecule type and class, hydrophobicity, absolute and relative SASA, and B-factor.Classification labels include three secondary-structure classes and 20 amino-acid labels; regression labels include residue hydrophobicity, solvent-accessible surface area, and mean crystallographic B-factor.Source: §A.2 · pp. 17–18
  • Voxels without alpha-carbons receive background labels, except molecule-type labeling also includes nucleotide and heteroatom voxels.Molecule class is produced by remapping molecule-type predictions into protein, nucleotide, and other.Source: §A.2 · p. 18

B.1 Atelier architectural details

Atelier tokenizes active volume patches and uses a transformer hypernetwork to generate INR weights, producing a continuous reconstruction at spatial query points. The appendix also specifies pretraining settings.

B.1 Atelier architectural details

  • The tokenizer divides padded volumes into 6 × 6 × 6 patches, discards background-only patches, and encodes the remaining active patches as tokens.The median volume has 236 active patches, and 95% of patches are discarded across the dataset.Source: §B.1 · p. 18
  • A 24-layer transformer processes volume and learned weight tokens to generate INR weights, with token rows normalized and scaled by learned per-layer α_i.The architecture has latent dimension 256, eight attention heads, and feedforward hidden dimension 2048; the learned weight tokens are shared across input volumes.Source: §B.1 · pp. 18–19
  • The INR encodes spatial query points with fixed Fourier features and is evaluated at all 96^3 grid points to render the full volume.The model has 38,332,806 trainable parameters and is adapted in part from Trans-INR.Source: §B.1 · pp. 18–19

B.2 Atelier pretraining details.

  • Atelier pretraining used AdamW for 5000 epochs on eight H100 GPUs, with an effective batch size of 64 volumes and approximately 90 hours of wall time.Training used octahedral-group augmentations, gradient clipping at 1.0, and a learning-rate schedule with linear warmup followed by cosine decay.Source: §B.2 · p. 19

B.3 Transformer-only ablation

The transformer-only ablation reconstructs each patch directly from its transformed volume token, without INR weight tokens. It holds the tokenizer, transformer, and training hyperparameters fixed relative to Atelier.

B.3 Transformer-only ablation

  • The ablation removes weight tokens and uses transformed volume tokens to reconstruct their corresponding patches through a learned linear projection.Reconstructed patches are reassembled and evaluated with the same loss as Atelier.Source: §B.3 · p. 19
  • The volume tokenizer and transformer remain unchanged, isolating the change to how the transformed tokens produce reconstructions.The ablated model uses exactly the same hyperparameters as Atelier.Source: §B.3 · p. 19
  • This variant tests a coarse representation: each query point uses its containing patch's transformed volume token, so points in that patch share one feature.Source: §B.3 · p. 19

B.4 EMNUSS property prediction experiment details

The experiment compares volume-only, coarse-feature, and fine-feature inputs to an EMNUSS 3D nested U-Net, alongside a transformer-only reconstruction baseline. Per-volume analyses find significant gains in most comparisons, with exceptions for amino-acid prediction and B-factor prediction.

B.4 EMNUSS property prediction experiment details

  • EMNUSS maps a C-channel volume to per-voxel class scores or one regression value, with the same architecture implementation across all three input variants.The variants are volume-only, coarse-feature augmentation, and fine-feature augmentation.Source: §B.4 · p. 19
  • All three variants use AdamW for 300 epochs, with 10 warmup epochs to a 0.001 learning rate followed by cosine decay to 0.0001.Classification uses cross-entropy and accuracy; regression uses non-background-voxel MSE and is evaluated by relative error.Source: §B.4 · p. 19
  • Fine features are extracted by querying the frozen Atelier INR around each annotated voxel, averaging activations across layers, then Gaussian-weighting nearby voxel features into one 256-dimensional vector.Queries cover a sphere of radius 3 voxels; the Gaussian pooling uses σ = 1.5 voxels.Source: §B.4 · p. 20
  • Coarse features use the frozen transformer-only model's 256-dimensional token for the containing 6×6×6 voxel patch, giving all voxels in that patch identical features.The transformer-only ablation directly maps patch tokens to patch reconstructions.Source: §B.4 · p. 20, Figure 7 · p. 20
  • Table 4 reports Matthews correlation coefficient (MCC) for the classification tasks, averaged over labeled test voxels.Source: §B.4 · p. 20, Table 4 · p. 21
  • For most per-volume comparisons, fine Atelier features significantly outperform baselines; coarse features win for amino acid prediction, while fine features do not significantly beat volume-alone for B-factor.The other comparisons show modest but significant improvements.Source: Table 5 · p. 21, §B.4 · p. 21
  • Per-volume comparisons use Wilcoxon signed-rank tests, Holm-Bonferroni correction, win rates, and Hodges-Lehmann effect sizes with 95% distribution-free confidence intervals.Positive differences and effect sizes indicate better performance from fine features.Source: §B.4 · p. 21

C Limitations and Future Work

The evaluation uses random EMDB-ID splits, which limit claims about generalization to entirely unseen folds. The authors also identify pretraining cost and high-frequency coordinate encoding as areas for improvement.

C Limitations and Future Work

  • Random EMDB-ID splits can place homologous proteins or alternate conformations across partitions, limiting claims of absolute generalization to entirely unseen folds.Sequence-identity- or CATH-topology-stratified splits are proposed as a next step.Source: C Limitations and Future Work · pp. 21–22
  • At roughly 720 H100-hours, pretraining remains a barrier to iteration and scaling to larger volumes or higher-resolution grids, despite amortized inference after training.Better tokenization, sparse attention, or distillation are proposed to reduce the cost.Source: C Limitations and Future Work · p. 22
  • The highest Fourier-coordinate frequencies can exceed the 963-grid Nyquist limit; ReLU spectral bias regularizes in practice, while principled bandlimiting remains a refinement.Source: C Limitations and Future Work · p. 22

Related papers

Loading 2609.30569v1…