Molecular Representations in Implicit Functional Space via Hyper-Networks

TL;DR

MolField reframes molecular representation learning as learning in the space of continuous functions over $\mathbb{R}^3$, enforcing $SE(3)$ invariance through a canonical coordinate system. It introduces a three-part architecture—Canonical Implicit Neural Representation (C-INR), Structured Weight Tokenization (SWT), and a Function Space Hyper-Network (FSHN)—to generate distributions over molecular functions conditioned on context, enabling end-to-end training across molecular dynamics and property prediction. Empirical results show improved molecular-dynamics surface reconstruction and state-of-the-art or competitive performance on QM9 properties, along with robust ablations and analysis of discretization robustness. The work offers a principled, transferable representation for molecules that unifies dynamics, properties, and generation under function-space learning, with implications for generalization, data efficiency, and cross-task transfer.

Abstract

Molecular representations fundamentally shape how machine learning systems reason about molecular structure and physical properties. Most existing approaches adopt a discrete pipeline: molecules are encoded as sequences, graphs, or point clouds, mapped to fixed-dimensional embeddings, and then used for task-specific prediction. This paradigm treats molecules as discrete objects, despite their intrinsically continuous and field-like physical nature. We argue that molecular learning can instead be formulated as learning in function space. Specifically, we model each molecule as a continuous function over three-dimensional (3D) space and treat this molecular field as the primary object of representation. From this perspective, conventional molecular representations arise as particular sampling schemes of an underlying continuous object. We instantiate this formulation with MolField, a hyper-network-based framework that learns distributions over molecular fields. To ensure physical consistency, these functions are defined over canonicalized coordinates, yielding invariance to global SE(3) transformations. To enable learning directly over functions, we introduce a structured weight tokenization and train a sequence-based hyper-network to model a shared prior over molecular fields. We evaluate MolField on molecular dynamics and property prediction. Our results show that treating molecules as continuous functions fundamentally changes how molecular representations generalize across tasks and yields downstream behavior that is stable to how molecules are discretized or queried.

Figures (7)

• Figure 1: Molecular representations: discrete versus continuous function-space modeling. (a) A molecule exhibits inherently continuous and smooth spatial structure. (b) Conventional molecular representations discretize molecules into sequences, graphs, or point clouds, which only partially capture this continuity. (c) In contrast, we model molecules directly in function space by learning a continuous molecular field that maps 3D coordinates to physical quantities (e.g., density or potential), enabling a unified representation for diverse downstream tasks.
• Figure 2: Overview of MolField.Top: A molecule is mapped into a canonical coordinate system to construct an implicit neural representation (INR) that is invariant to SE(3) transformations. The INR weights are tokenized in a structured, layer-wise manner to enable transformer-based processing, while the model is trained as a hyper-network that learns to generate task-conditioned INRs. Bottom: The resulting INR supports diverse molecular learning tasks, including dynamic surface modeling and property prediction, and can be naturally extended to molecular generation (Appendix \ref{['sec:generation']}).
• Figure 3: Long-horizon prediction on molecular dynamics. Quantitative comparison of future surface prediction over increasing time horizons. MolField (ours) yields more accurate and stable long-term predictions, particularly for extended horizons.
• Figure 4: Data efficiency on molecular dynamics. Performance of different methods under varying training data ratios. MolField (ours) consistently achieves higher IoU and NC across all data regimes, demonstrating improved data efficiency.
• Figure 5: Correlation between INR reconstruction loss and property prediction error (MAE) of MolField. Strong Pearson and Spearman correlations indicate that more accurate function-space reconstructions lead to improved property prediction. Pretraining the INR on molecular generation further strengthens this relationship and consistently reduces prediction error.

Theorems & Definitions (6)

• Definition 2.1: Molecular Functions
• Definition 2.2: SE(3) Symmetry
• Proposition 3.1: Frame equivariance
• Lemma 3.2: SE(3)-invariant coordinate mapping
• proof
• proof