RiteeshGrandhi/sciclaimeval-shared-task huggingface dataset:
Neural fields.
Neural fields, also referred to as implicit neural representations (INRs), are coordinate-based neural networks that map coordinates (e.g., pixels on an image or coordinates in 3D Euclidean space) to features (e.g., RGB values or atomic occupancies).
The idea of representing data points implicitly as neural networks dates back to the work of [ 14 ] .
Recently, these representations have been successfully applied to model continuous signals, e.g., 2D images [ 15 , 16 , 17 ] , 3D shapes [ 18 , 19 , 20 , 21 ] , 3D scenes [ 22 , 23 ] , videos [ 24 , 25 ] , physics [ 26 , 27 ] , due to their appealing properties.
Recently, two concurrent seminal work lead to a fast progress of neural fields by overcoming the spectral bias of coordinate-based neural networks [ 28 ] .
Sitzmann et al. [ 29 ] propose SIREN, a neural network that uses periodic activation functions, while Tancik et al. [ 30 ] considers a positional encoding based on Fourier features.
Built on top of those architectures, multiplicative filter networks (MFNs) [ 31 ] represent fields as a simple linear combination over an exponential number of basis functions (e.g. Fourier or Gabor basis).
Due to their simplicity and strong performance, we use MFNs to model the atomic occupancy fields.
Generative models of fields.
Generative models for neural fields were first applied in 3D computer vision problems.
Mescheder et al. [ 19 ] learn the distribution of shape occupancy fields with VAEs [ 32 ] , while [ 18 , 33 ] achieves similar objectives using GANs [ 34 ] . Diffusion models [ 35 ] have also been applied to learn the distribution of neural fields [ 36 , 37 , 38 , 39 ] .
Some work [ 37 , 40 ] parameterize the neural field with the vector of all the corresponding weights.
However, when the signal is complex and the neural fields have large number of parameters (e.g., in the order of millions), it is preferable to parameterize the field with a latent code with much lower dimension [ 36 , 41 , 42 , 43 ] .
Dupont et al. [ 36 ] fit the whole dataset with a shared coordinate-based network and learn a latent modulation code for each field with gradient-based meta learning [ 44 ] .
Similarly to them, we parameterize neural fields with latent modulation codes.
However, instead of applying meta learning, we learn the latent codes through stochastic optimization, either following the “auto-encoding” [ 19 ] or the “auto-decoding” [ 20 ] framework.
D molecule generation.
Most 3D molecule generation approaches represent atoms as points (with coordinates and atom types) and molecules as a set of points. For example, [ 45 , 46 , 47 ] propose autoregressive approaches to sample atoms, while [ 48 , 49 ] use normalizing flows [ 50 ] .
Hoogeboom et al. [ 4 ] propose EDM, a diffusion model [ 35 ] applied to point cloud of atoms with E(3) equivariance [ 6 ] .
Many follow-up works extend EDM [ 51 , 52 , 53 ] .
For example, [ 54 , 55 , 56 ] improve its performance by leveraging extra information during training (such as molecular graph and formal charges).
This contribution is orthogonal to ours and can potentially be incorporated into our generative model.
Other approaches [ 57 , 58 ] map atomic densities on discrete 3D regular grids and leverage computer vision techniques for generation.
Recently, VoxMol [ 5 ] (and its latent version [ 59 ] ), a score-based generative model based on walk-jump sampling [ 1 ] , shows that voxel-based representations can achieve state-of-the-art results on 3D drug-like molecule generation.
However, these methods scale cubically with the volume occupied by molecules, which limits its scope of application.
Neural fields are the continuous generalization of discrete 3D grid representations: they achieve good performance on 3D molecule generation and are more efficient in terms of memory and computation.
Conditional molecule generation.
Voxels and point-clouds have also been used for conditional 3D molecule generation, usually by building upon an unconditional model.
The authors in [ 57 ] condition generation on 3D pharmacophores features, [ 60 , 61 , 62 , 63 ] generate ligands conditioned on protein pockets, [ 64 ] generate molecules conditioned on fragments and [ 65 , 66 ] generate 3D conformations conditioned on molecular graphs.
We are aware of only one other work that uses field-based representation for molecules [ 67 ] . There are several differences between our works: they use different data representation, neural network architecture and noise model. While they consider the problem of generating molecule conformations given a molecular graph, we handle the more general problem of unconditional 3D molecule generation (without access to a molecular graph). Our model can easily be adapted to conformer generation by conditioning the generative model to the molecular graph. Moreover, our approach can also be conditioned to tasks where we do not have access to molecular graphs, such as structure-based drug design or electron density generation.
Neural atomic occupancy fields
We now describe how we represent molecules as continuous occupancy fields, how we approximate them with neural fields and how we decode the neural fields to retrieve molecular conformations. We finish the section by providing some useful properties of our neural field representations.
Molecules as continuous occupancy fields
We represent atoms as continuous Gaussian-like shapes in 3D space, centered around their atomic coordinates.
Molecules are defined as fields mapping every point in the 3D space to the atomic densities of each atom type, v\mathrel{\mathop{\ordinarycolon}}\mathbb{R}^{3}\rightarrow\mathbb{R}^{n} , where n is the number of atom types in the dataset {\mathcal{D}} .
We follow previous work [ 68 , 69 , 70 ] , and compute the occupancy field v_{a} for each atom type a by integrating the occupancy generated by all atoms of this type as follows:

\forall x\in\mathbb{R}^{3},\,v_{a}(x)=1-\prod_{i=1}^{n_{a}}\Big{(}1-exp\Big{(}-\Big{(}\frac{\|x-x_{a_{i}}\|}{.93r}\Big{)}^{2}\Big{)}\Big{)}, (1)

where a_{i} is the i^{\rm th} atom of type a , for a total of n_{a} atoms.
We set the atoms’ radius to be r=.5 Å for all atom types.
Molecular fields are smooth functions taking values between 0 (far away from all atoms) and 1 (at the center of atoms).
Molecular neural fields
Each molecule in the dataset is mapped to a modulation code z\in\mathbb{R}^{d} and we parameterize the molecular occupancy v with a conditional neural field f_{\phi}\mathrel{\mathop{\ordinarycolon}}\mathbb{R}^{3}\times\mathbb{R}^{d}\rightarrow\mathbb{R}^{n} . Our objective is to learn the parameters \phi and the modulation code z such that for any molecular field v and coordinate x\in\mathbb{R}^{3} , f_{\phi}(x,z)\approx v(x) .
We approximate the molecular fields with a linear combination of an exponential large number of parameterized basis functions, such that amplitudes are modulated by the individual codes z .
This parametrization is achieved by modeling the neural field with multiplicative filter networks (MFN) [ 31 ] , a type of coordinate-based network that provides an elegant way to perform this linear combination under some assumptions on the basis functions.
We introduce the parameters associated with these functions in Equation ˜ 4 .

Our conditional MFN is a network with L multiplicative blocks, as illustrated on Figure ˜ 2 (a). We implement conditioning of its parameters with FiLM [ 71 ] . Each multiplicative block is composed of a fully-connected layer, a FiLM modulation layer and an elementwise product with a basis, as illustrated on Figure ˜ 2 (b). The neural field can be expressed by the following recursive expression:

 h^{(0)}(x) =s_{\omega^{(0)}}(x), h^{(l)}(x) =\Big{(}\gamma^{(l-1)}\odot\Big{(}W_{f}^{(l-1)}h^{(l-1)}(x)\Big{)}+\Big{(}b^{(l-1)}+\beta^{(l-1)}\Big{)}\Big{)}\odot s_{\omega^{(l)}}(x),\leavevmode obreak\ \leavevmode obreak\ l\in(1,L-1), f_{\phi}(x,z) \triangleq h^{(L)}(x)=W_{f}^{(L-1)}h^{(L-1)}(x)+b^{(L-1)},

where s_{\omega^{(l)}} is a spatial basis function parameterized by \omega^{(l)} , \odot denotes the Hadamard product and

\beta^{(l-1)}=W^{(l-1)}_{\beta}z, and \gamma^{(l-1)}=W^{(l-1)}_{\gamma}z,

are the bias and scale modulation terms.
We propose two approaches to learn the neural field’s parameters \phi=\{W_{f}^{(l)},b^{(l)},\omega^{(l)},W^{(l)}_{\beta},W^{(l)}_{\gamma}\} and the codes z (one per each molecule in {\mathcal{D}} ).
Auto-decoding.
In this setting, introduced by [ 20 ] , we initialize each code z randomly and directly learn them (together with the parameters of the neural field) with backpropagation. This is achieved by solving the following optimization problem:

\operatorname*{arg\,min}_{\phi,\{z_{v}\}_{v\in{\mathcal{D}}}}\leavevmode obreak\ \sum_{v\in{\mathcal{D}}}\leavevmode obreak\ \int\|f_{\phi}(x,z_{v})-v(x)\|_{2}^{2}\leavevmode obreak\ \mathrm{d}x, (2)

where the integral is approximated by sampling finite sets of points {\mathcal{X}}\subset\mathbb{R}^{3} . While auto-decoding was usually applied in settings with relatively few samples, we were able to scale the training to large datasets of one million samples (see Appendix ˜ B ).
See Algorithm ˜ 1 on Appendix ˜ B for more details.
Auto-encoding.
This approach, introduced by [ 19 ] and illustrated in Appendix ˜ B Figure ˜ 4 , generates the modulation code via an encoder \zeta_{\psi} , parameterized by \psi and decodes the neural field back.
In this work, \zeta_{\psi} is a (trainable) 3D convolutional network encoder that takes (low-resolution) voxel grids \mathcal{G} as inputs.
This approach is flexible and compatible with other encoder architectures and molecule representations (e.g., GNN/point clouds).
The parameters of the encoder and the neural field are learned with the following objective:

\operatorname*{arg\,min}_{\phi,\psi}\leavevmode obreak\ \sum_{v\in{\mathcal{D}}}\leavevmode obreak\ \int\|f_{\phi}(x,\zeta_{\psi}(\mathcal{G}))-v(x)\|_{2}^{2}\leavevmode obreak\ \mathrm{d}x

Once training is done, we generate the code with the trained encoder.
See Algorithm ˜ 2 on Appendix ˜ B for more details.
Instead of learning the codes individually, this approach learns an encoder, which allows to leverage data augmentation more efficiently.
As a result it helps learn a more structured latent space.
These benefits are reflected empirically in our experiments.
From codes to atomic coordinates
By leveraging the modulation codes z and the neural field f_{\phi} , we have access to the (learned) continuous occupancy field, f_{\phi}(\cdot,z) . However, in many useful applications in chemistry and biology, we are more interested in the 3D conformation of molecules.
Next, we describe how we can extract the molecular conformation from a learned (or generated, as we will see next) modulation code.
This is the decoding step outlined in the introduction.

We start by identifying all atoms in the field, their approximate locations and their type.
To this end, we render a discretized voxel grid from the molecular field using a uniform discretization of space and the neural field f_{\phi}(\cdot,z) .
We then apply a peak finding algorithm to infer the number of atoms in the molecule on each channel of the grid (each representing a separate atom type) and their (discretized) coordinates.
Finally, we introduce a new continuous refinement to find the local maximum of the neural field. For each identified atoms a , we refine its coordinates around the neighborhood of the coordinates found with the peak detector x_{a}^{0} :

x_{a}=\operatorname*{arg\,max}_{x\in\mathbb{R}^{3}\mathrel{\mathop{\ordinarycolon}}\|x-x_{a}^{0}\|\leq r}[f_{\phi}(x,z)]_{a},

where [f_{\phi}(x,z)]_{a} denotes the field restricted to the channel corresponding to the atom type.
This continuous refinement finds atomic coordinates that lie beyond the initial coarse uniform discretization.
In practice, we batch the refinement process across molecules and use L-BFGS.
We demonstrate in Section ˜ E.2 its efficiency compared to prior non-continuous refinement approaches from [ 72 , 5 ] .
Molecular neural fields properties
The proposed conditional neural field enjoys many properties that make it a natural choice for handling large 3D molecules represented as continuous fields.
Flexibility w.r.t. basis.
Conditioning MFNs gives the flexibility to choose any type of spatial basis that satisfies a multiplicative-sum property (see the definition in [ 31 ] ).
In our preliminary experiments, we observed that setting the spatial basis to Gabor filters performed better than Fourier filters as they account for the sparse nature of occupancy fields.
For each layer l , we consider the following Gabor parameterization, also used in [ 27 ] :

s_{\omega^{(l)}}(x)=\exp\Bigl{(}-\frac{ u^{(l)}}{2}\|x-\mu^{(l)}\|_{2}^{2}\Bigr{)}(\cos(\Omega^{(l)}x),\sin(\Omega^{(l)}x)), (4)

where \mu^{(l)} is the mean of the Gabor filter, \nu^{(l)} is the scale, \Omega^{(l)} is the frequency and (\cdot,\cdot) refers to the concatenation operator. Equation ˜ 4 combines both real and imaginary parts of the complex Gabor filter. This allows to remove phase parameters and reduce the overall parameter count of MFNs [ 31 ] .
Other choices of basis are also possible and are left for future work.
Parameter efficiency.
Our overall conditional MFN formulation is parameter efficient and shares parameters across molecules and channels (i.e. atom types).
As [ 27 ] , we excluded the basis functions parameters \omega from FiLM to further decrease the parameter count.
Memory efficiency.
Our conditional neural field can be trained on any free-form discretization of the input field.
Occupancy values are computed on the fly.
This allowed to train FuncMol with large batch size even on large 3D molecules.
We found that training the neural field by up-sampling points close to the atoms’ center improved training time as further detailed in Appendix ˜ B .
Alternative approaches like VoxMol [ 5 ] cannot be trained efficiently in this setting: for reference, on the macro-cyclic peptide generation task of Section ˜ 5.4 , on 4 A100 GPUs VoxMol’s training cost per epoch was 10 hours while our neural field’s training cost was less than 12 minutes.
Reconstruction quality and robustness to noise.
Finally our neural field reconstructs accurately the input data as demonstrated in Section ˜ E.1 .
Moreover, operating on these latent codes makes our model extremely robust to noise in code space.
We demonstrate this property in Section ˜ E.3 by reporting the sampling metrics when perturbing the codes z by a Gaussian noise.
Sampling efficiency.
We use the latent codes for generative modeling as explained in Section ˜ 4 .
Most sampling operations are done on a small dimensional latent space, while decoding into a full molecular field is done only after sampling.
As we show in Section ˜ 5 , our approach (which involves sampling latent code followed decoding them into molecules) achieves at least one order magnitude faster molecule sampling time than previous methods.
Score-based generative modeling
We use our latent modulation representations for a downstream generative modeling task. Section ˜ 4.1 describes the neural empirical Bayes (NEB) formalism used in our method and Section ˜ 4.2 explains how we perform sampling.
Neural empirical Bayes
Let p(z) be the distribution of codes and p(v) be the (unknown) distribution of molecular fields, defined more formally as the pushforward of p(z) via the mapping z\mapsto f_{\phi}(\cdot,z) .
NEB estimates the score function of a smoothed density of the codes p(y) , g_{\theta}(y)\approx\nabla\log p(y) .
Indeed sampling from a smoothed density p(y) benefits from faster mixing than on the original density p(z) [ 1 , 73 , 74 ] .
This smoothed distribution is defined by transforming the random variable Z with an additive isotropic Gaussian noise with a known noise level \sigma , Y=Z+N , where N\sim\mathcal{N}(0,\sigma^{2}I_{d}) .
The noise level \sigma plays a key role, trading-off simplicity of the denoising objective and the sampling quality.

NEB is based on an empirical Bayes view of (denoising) score-based models that relates the estimator of clean data (denoiser) and the score function of the smoothed density at a fixed noise level [ 75 , 76 , 1 ] .
The denoiser is taken to be the least-square estimator of Z given Y=y which is the Bayes estimator, i.e. \hat{z}(y)=\mathbb{E}[Z|Y=y] . Under Gaussian noise, denoiser and smoothed score function are related by the Tweedie-Miyasawa formula:

\hat{z}(y)=y+\sigma^{2} abla\log p(y)

The denoiser is parameterized by a neural network and learned by minimizing the following objective:

{\mathcal{L}}(\theta)=\mathbb{E}_{z\sim p(z),\varepsilon\sim\mathcal{N}(0,\sigma^{2}I_{d})}\mathopen{}\mathclose{{\left\|z-\hat{z}_{\theta}(z+\varepsilon)}}\right\|_{2}^{2}

The score function is recovered from a learned denoiser via Equation ˜ 5 and is used for sampling smoothed codes (see Section ˜ 4.2 ).
In practice, we optimize the empirical loss based on the latent codes inferred from a set of molecular fields {\mathcal{D}} .
See pseudo-code in Appendix ˜ B , Algorithm ˜ 3 .
Walk-jump sampling
We use the score function g_{\theta} to sample codes using the walk-jump sampling (WJS) scheme [ 1 , 77 , 73 , 78 ] .
This approach samples molecules from p(z) using the learned score function of noisy data instead of clean data.
It consists of two main steps: walking and jumping as detailed in Appendix ˜ B , Algorithm ˜ 4 . Figure ˜ 1 (b) illustrates these two main steps in a WJS chain: walking consists in generating noisy codes while jumping consists in generating clean codes z .

(initialization) To improve mixing, as [ 77 ] , we initialize the chains by adding uniform noise to Gaussian noise (with the same \sigma used when training the denoiser).
In practice we sample the uniform noise over the range of code values, \mu\sim\mathcal{U}_{d}(\min_{z\in{\mathcal{D}}_{z},i\in\{1\cdots d\}}z_{i},\max_{z\in{\mathcal{D}}_{z},i\in\{1\cdots d\}}z_{i}) , where {\mathcal{D}}_{z} is the training dataset of codes, arriving at y_{0}=\mu+\varepsilon , \varepsilon\sim\mathcal{N}(0,\sigma^{2}I_{d}) .

(walk) Noisy codes are sampled from p(y) with Langevin MCMC algorithms that discretize the underdamped Langevin diffusion [ 79 ] starting from y_{0} and u_{0}=0 :

\mathrm{d}u_{t}=-\gamma u_{t}\mathrm{d}t\;+\;g_{\theta}(y_{t})\mathrm{d}t\;+\;\sqrt{2\gamma}\leavevmode obreak\ \mathrm{d}{B}_{t},\quad\quad\mathrm{d}y_{t}=u_{t}\mathrm{d}t\;, (7)

where B_{t} is the standard Brownian motion in \mathbb{R}^{d} and \gamma is the friction (the “mass” is set to 1).
We discretize this SDE using the ABOBA scheme from Sachs et al. [ 80 ] , given a discretization step \delta and a fixed number of walk steps K .
We analyze the impact of K in Section ˜ E.4 .

(jump) At a given time step K , clean samples are estimated by denoising the smooth code, i.e., z_{K}=\hat{z}_{\theta}(y_{K}) .
These codes are then used to obtain the atomic coordinates as detailed in Section ˜ 3.3 .
Experiments
We now evaluate our model for unconditional generation.
We start with a description of our experimental setup ( Section ˜ 5.1 ), then present our results on two popular small molecule datasets ( Sections ˜ 5.2 and 5.3 ) and a recent macro-cyclic peptide dataset ( Section ˜ 5.4 ).
Datasets.
We evaluate FuncMol on three datasets: QM9 [ 81 ] , GEOM-drugs [ 82 ] and CREMP [ 12 ] .
QM9 contains an enumeration of all possible molecules up to 9 heavy atoms (29 including hydrogens) satisfying some constraints [ 83 ] .
GEOM-drugs contains multiple conformations for 430K drug-sized molecules (computed with semi-empirical density functional theory), with an average of 44 heavy atoms per molecule.
CREMP is a recent dataset that contains multiple conformations of macrocyclic peptides 4-6 residue long, with an average of 74 heavy atoms per molecule.
We model hydrogen explicitly and consider 5 chemical elements for QM9 (C, H, O, N, F), 6 for CREMP (C, H, O, N, F, S) and 8 for GEOM-drugs (C, H, O, N, F, S, Cl and Br), ignoring the P, I and B elements that occur extremely rarely.
We use a split of 100K/20K/13K molecules for QM9, 1.1M/146K/146K on GEOM-drugs and 409K/10K/9K on CREMP for train, validation and test, respectively. We use the same pre-processing and splits in [ 54 ] for QM9 and GEOM-drugs and in [ 84 ] for CREMP.
Implementation details.
Our main model, FuncMol , follows the auto-encoding approach described in Section ˜ 3.2 .
The codes z are computed with an encoder that takes as input a low-resolution voxelized representation of the molecular field with grid dimension of 16 \times 16 \times 16.
The encoder is a 3D CNN containing 4 residual blocks, where each block contains 3 convolutional layers followed by BatchNorm, ReLU and pooling (max pooling on the first three blocks and average pooling on the last one) layers.
We consider modulation codes with dimension 1024 on QM9 and 2048 on GEOM-drugs and CREMP.
We use the same neural field network for all datasets: a conditional MFN with Gabor filters and 6 FiLM-modulated layers, where each fully-connected layer has 2048 hidden units.
We augment the training set by applying random rotations on the three Euler angles.
The weights of the latent code encoder and neural field decoder are trained jointly.

We also show results for the auto-decoding based model, FuncMol dec .
In this setting, we initialize the codes randomly and optimize them together with the neural field weights.
This approach is less fit for performing large amounts of augmentation as it solves a costly per-sample optimization problem; thus we did not apply data augmentation.
As a consequence, we observed that this model is more prone to memorization than the auto-encoding approach (e.g., on GEOM-drugs, around 33% of the generated molecules are copies from the training set).

We normalize the codes to have zero mean and unit variance.
We choose a noise level in normalized space of \sigma=1.2 for GEOM-drugs and CREMP, \sigma=2.0 for QM9.
Our code denoiser is a modified version of the denoiser used in [ 36 ] : a fully-connected network with 18 residual blocks (each with two linear layers with 6144 hidden units) and skip connections.
We remove the bias of all layers and use ReLU activations as in [ 85 ] .
To limit memorization in FuncMol dec , we add dropout (ratio 0.3) between the fully-connected layers in each residual block.
For QM9, we consider a smaller network (6 residual blocks and 4096 hidden units).
We initialize the MCMC chains with noise and use the following sampling hyperparameters \gamma=1.0 and \delta=\sigma/2 as in [ 5 , 78 ] .
For evaluation purposes, we generate one sample per chain.
We consider 1000 steps for QM9 and GEOM-drugs and 10000 steps for CREMP.
See Appendix ˜ B for more details on the implementation.