Generative modeling enables molecular structure retrieval from Coulomb explosion imaging

Domain: chemistry
Published: 2026-03-03
Editorial outcome: minor_revision

Abstract
Capturing the structural changes that molecules undergo during chemical reactionsinrealspaceandtimeisalong-standingdreamandanessentialprerequisite forunderstandingandultimatelycontrollingfemtochemistry.Akeyapproachto tackle this challenging task is Coulomb explosion imaging, which has benefited decisively from recently emerging high-repetition-rate X-ray free-electron laser sources. With this technique, information on the molecular structure is inferred from the momentum distributions of the ions produced by the rapid Coulomb explosion of molecules. Retrieving molecular structures from these distributions poses a highly nonlinear inverse problem that remains unsolved for molecules consisting of more than a few atoms. Here, we address this challenge using a diffusion-based Transformer neural network. We show that the network reconstructsunknownmoleculargeometriesfromion-momentumdistributionswitha meanabsoluteerrorbelowoneBohrradius,whichishalfthelengthofatypical chemical bond.

Introduction
Imaging molecular structure, and, in particular, its temporal evolution is fundamental to understanding and steering ultrafast processes, including chemical reactions1. Severalexperimentaltechniqueshavebeendevelopedduringthelastdecadestostudy theevolutionofmolecularstructureonpicosecondandfemStosecondtimescales.Rely- ing on a variety of measurement concepts, these techniques probe different aspects of molecular structure and dynamics with different levels of fidelity. Static molecular structures can be captured with the highest position-space resolution using electron microscopy2.Time-resolvedmeasurementsofevRolvingmoleculargeometriesoftenrely on X-rays or high-energy electrons to infer molecular structure from recorded diffrac- tionpatterns3,4.Coulombexplosionimaging(CEI),whichtakesadvantageofCoulomb repulsion of nuclei within molecules thNat are rapidly stripped of their electrons5–7, is a less mature technique that can also provide time-resolved information if combined with short laser pulses8. Since atom I ic motion typically unfolds on femtosecond time scales(asdeterminedbymolecEularvibrations),CEIwithintensefemtosecondlaseror X-ray pulses has been exploited for studying molecular structural changes9–18. These pulses rapidly ionize the target molecules, causing their atomic constituents to repel and fragment as a result of Coulombic forces. The resulting ion-momentum distribu- tions contain informaItion about the initial geometric configuration of the molecule before ionization1T9–24.

In all of tRhe imaging methods discussed above, extracting molecular structure from experimental data requires computational algorithms of varying complexity. For diffraction-based imaging techniques, reliable inversion methods are available. In contrast, a corresponding inversion of measured momentum-space data to molecular geometry is not routinely available in the case of CEI. For CEI, geometry retrieval requires solving a highly nonlinear inverse problem, which is extremely challenging whendealingwithmoleculescontainingmorethan3-4atoms.Ingeneral,inverseprob- lems involve the reconstruction of hidden causal factors from observable data25,26, which are connected by a forward process. If the forward process is trivial to cal- culate and the noise distribution of this process is known, inverse problems can be solvedwiththemaximumlikelihoodestimationorthemaximumaposterioriapproach. Both approaches are typically implemented with an iterative solver, which requires the forward process to be calculated at each step of the iterations. This makes them unfeasible for solving the CEI inverse problem because, in this case, the forward pro- cess is driven by the time-dependent many-body interactions governed by quantum mechanics, which is computationally prohibitive to be integrated into an iterative solver. Consequently, direct reconstruction of molecular geometry from CEI has only beendemonstratedinafewcasesusingaclassicalimplementationoftheforwardpro- cess27,28.MostCEIstudies5,10,11,13–15,17–21,23,24havereliedonasingle-passsimulation oftheforwardprocesstocomparewithexperimentalmeasurements,leavingaccurate, general reconstruction of molecular geometries an open and unresolved problem.

In this work, we address the molecular structure retrieval problem in CEI with a deep generative neural network designed to reconstruct molecular geometries from ion momentum measurements, which we termed MOLEXA (molecular structure reconstructionfromCoulombexplosionimaging).ItisbuiltontheTransformerarchi- tecture29 andthediffusiongenerativemodelingframework30–35,withanovelmemory mechanism implemented in between the Transformer blocks. The complex forward process of CEI not only renders the classical iterative solvers inapplicable, but also poses a severe challenge for deep learning techniques because it is computationally too demanding to generate adequate data for neural network training. To address the issueofdatascarcity,MOLEXAusesatwo-stagetrainingapproach.Stage1trainson a large dataset generated using a computationally inexpensiSve, approximate forward model, while stage 2 fine-tunes the model on a smaller,Ehigh-quality dataset derived fromabinitiosimulations.Thedual-phasestrategyreducesthemeanabsolutepredic- tion error to less than one atomic unit, or half the length of a typical chemical bond. Our present work focuses on the reconstruction of the molecular structure from CEI measurementsusingX-raypulses,butthedemon stratedgenerativemodelingapproach can also be applied to building reconstruction models for CEI measurements using optical lasers19,36,37 and highly charged ioIn beams38.

Results
The MOLEXA netwoCrk

The MOLEXA model takes the measurable quantities (i.e., the three-dimensional ion momenta measured in coincidence) from CEI as an input and predicts the ini- tial structure of a molecule before its interaction with the X-ray pulse (Fig. 1a). It comprises four mAodules for input embedding, dynamics extraction, structure denois- ing, and uncertainty estimation, which will be briefly described in the following. Full network details can be found in Supplementary Note 1.

TheinputtotheEmbeddingModule(Fig.1a)containstheatomicnumber,charge state, and molecular-frame momentum of each atomic fragment. The embeddings of the atomic number and charge state are concatenated with the linear projection of the momentum to form atom-wise features. The atomic features are concatenated to create pairwise features, which are then processed by a residual block before being sent as input to the Dynamics Extraction Module.

The Dynamics Extraction Module (Fig. 1b) generates conditioning information usedintheStructureDenoisingModule.ThebasicTransformerblock,whichincludes multi-head self-attention, is implemented and accounts for the majority of the com- putationalloadforboththisandsubsequentmodules.Insteadofdirectlystackingthe Transformer blocks on top of one another, we found that adding memory operations at the end of each block enhances model performance. Similar to the long short-term memory mechanism39, the memory operations (on the right side of Fig. 1b) include a forget gate that regulates what information to discard from the previous state, an updategatethatdecideswhichinformationintheTransformeroutputshouldbeadded to the memory, and an output gate that selectively sends the current state of the memory to the next Transformer block. In comparison to using skip connections, the memory mechanism is found to suppress the mean (maximum) atomic distance and angle errorsby3.6% (4%) and1.3% (4.4%), respectively. We referto the combination of the Transformer and memory operations as the ”Transformer with Memory” (TM) block. There are six TM blocks in the Dynamics Extraction Module.

TheStructureDenoisingModule,illustratedinFig.1c,reconstructsthemolecular structure using a reverse diffusion process. It starts with a noisy molecular structure. Its atomic positions are encoded on the basis of the output of the Dynamics Extrac- tion Module and the current noise level. Pairwise features derived from this encoding S areprocessedbytwoTMblocks.Theoutputisprojectedtoobtainatom-wisefeatures that are further processed through a self-attention blocSk. The Position Decoder takes the transformed atomic features and outputs a lessEnoisy molecular structure. During inference,thisstructureisiterativelyrefinedbyadiffusionsamplerwithanoisesched- R ule adapted from Ref.34. As shown in Fig. 1c, five iterations, corresponding to four P intermediate structures, are performed to obtain the final molecular structure. The smaller atom sizes shown for the earlier ite rations reflect larger interatomic distances.

TheUncertaintyEstimationModuleistrainedtomatchthepredicteduncertainty with the absolute error between thIe predicted and ground-truth structures. It can provide uncertainty estimations fo r the structure predictions. Using pairwise features E from the Dynamics Extraction Module and the predicted molecular structure, the UncertaintyEstimationMoLduleestimatestheerrorsofthepredictedatomicpositions usingtwoTMblocks,folClowedbyanuncertaintydecoder.Itpre-definestheuncertainty bins r=[0,0.05,0.1,...,9.95] and estimates the probability that the prediction error I falls within each oTf these bins. For each coordinate of the predicted atomic positions, the uncertainty is then calculated as the probability-weighted sum of the bin values.

Training

Unlike text or image generation models, for which there exists an enormous amount oftrainingdata,deeplearningmodelsinphysicalsciencesoftenfacethedatascarcity issue, which is one of the main obstacles preventing the widespread adoption of deep learning techniques for solving physics-related problems. For the molecular structure retrieval problem, we created two training datasets by performing Coulomb explosion simulations at two levels of theory. One level involves the ab initio calculation of the XFEL-induced Coulomb explosion of molecules, which tracks the quantum transition probabilitiesacrossallparticipatingelectronicconfigurationswhiletreatingthenuclei as moving in a classical force field40. It has been shown to produce results that agree with experiments18. A similar level of theory was used, for example, in Refs.21,23. These high-level simulations are computationally expensive. Thus we created only a small dataset containing 76000 samples, using a thousand CPUs for more than a month. A portion of this dataset was kept for validation (10%) and testing (10%) purposes. Since the computation time scales roughly exponentially with the number of atoms, these simulations were limited to molecules with fewer than ten atoms, which was a compromise to balance the dataset size with computational constraints.

Thesecondleveloftheoryisamuchcheaper,classicalCoulombexplosionmodelwith crude approximations22. It was used to generate a dataset that is about a hundred times larger. MOLEXA was first trained on this large but inaccurate dataset and then on the small dataset, which is more accurate and best reflects the reality of a CEIexperiment.Wefoundthatthetwo-stagetrainingapproachreducedthestructure predictionerrorbyafactoroftwocomparedtotrainingsolelyonthesmallerbutmore accurate dataset.

Before training in each stage, both the ion-momentum and ground-truth position distributions were centered and aligned to a common molecular frame. Specifically, for each molecule, the emission direction of the heaviest ion fragment is defined as the x axis. Then the cosine similarities of the momenta of the other ion fragments S relativetothex directionarecalculated.Theionmomentumwiththesmallestcosine similarity is chosen to define the y axis through Gram-SchSmidt orthonormalization so that this momentum vector would fall within the firsEt quadrant of the x-y plane. Boththemomentaofallionfragmentsandtheinitialmolecularstructure,afterbeing R centered at the origin, are transformed into this coordinate frame. For diatomics, the P coordinate system is fully determined once the emission direction of the heavier ion fragment is defined as the x axis. This pre-align ment procedure, applied consistently N across all molecules, transforms the structure retrieval task from one with arbitrary coordinateframestoonewithinafixedmIolecularframe,therebyeliminatingtheneed to explicitly incorporate translational a nd rotational invariance into the model.

The loss function consists of two parts: a weighted mean squared error for the predicted molecular structures, aLnd a cross-entropy loss for uncertainty estimations. Both parts were used througChout the two training stages. The Uncertainty Estima- tion Module was further fine-tuned using the validation dataset while keeping the I othermodulesfrozen,duTringwhichonlythesecondpartofthelossfunctionwasused. All reported results were generated from the test dataset. The training and valida- R tion datasets only contain molecules with fewer than eight atoms, while the eight- or A nine-atom molecules were set aside to test the generalization capability of the model.

Model Performance

Using the test dataset of molecules with less than eight atoms, the mean absolute error (MAE) is 0.52 a.u. (atomic units), the mean (maximum) distance error (DE) of allatomicpairsinamoleculeis0.98(2.11)a.u.,andthemean(maximum)angleerror (AE)ofalldirectionaltripletsinamoleculeis13.97(38.39)degrees.MOLEXA,which was trained on only molecules with up to seven atoms, is also capable of reconstruct- ing the structure of molecules containing eight or nine atoms. For these molecules larger than those included in the training dataset, the MAE is 0.66 a.u., the mean (maximum) DE is 1.16 (3.43) a.u., and the mean (maximum) AE is 14.12 (57.91) degrees. The inference time distribution evaluated on all test molecules is plotted in SupplementaryFig.1andhasameanof59.8ms.Eachofthereconstructedmolecular structures is from a single model prediction. Based on the standard deviations of the modelpredictionsplottedinSupplementaryFig.2,theexpecteddiscrepancybetween the corresponding atomic coordinates of two independent reconstructions for a single molecule is about 0.07 a.u.

Figure 2 provides an overview of the structure retrieval performance of the model. In each column of the figure the results for molecules consisting of N atoms are pre- sented, showing exemplary structures predicted with low (top row) and high (second row) reconstruction uncertainties. The two bottom rows indicate how the general performance of MOLEXA (in terms of mean DE and AE) behaves as a function of the predicted uncertainty for molecules of different sizes. The Uncertainty Estima- tionModuleproducesuncertainties forallcoordinatesofall atomsina molecule.The overall uncertainty displayed here was obtained by taking an average of the uncer- tainties of all atomic positions of a molecule. The heatmaps in Fig. 2, as well as those in Supplementary Figs. 4-5 for MAE, maximum DE and AE, show that there S is a strong correlation between the predicted uncertainty and the errors of the recon- structed molecular structures. This indicates that theSformer can serve as a reliable metric for assessing whether a MOLEXA reconstruEction is trustworthy or not. It can be observed from these heatmaps that the predicted uncertainties can underestimate R the reconstruction errors, especially for larger molecules with larger prediction errors. P We attribute this to a training bias caused by the Uncertainty Estimation Module being trained on reconstructions most of which have relatively low errors. Such bias N canpotentiallybemitigatedwithaweightedcross-entropylossinfuturedevelopments.

With MOLEXA being a diffusioIn model, an uncertainty quantification for each sample is also obtainable by calc ulating the standard deviation of an ensemble of E its structure predictions. The uncertainties calculated with these two approaches are shown in Supplementary FiLg. 2 to have an approximately linear relationship. The for- mermethodisusedhereCbecauseitprovidesanuncertaintyestimateforeachprediction without requiring an ensemble of predictions. Note that although an uncertainty esti- I matecanassistthTeassessmentofthetrustworthinessofaprediction,itisnoteffective in selecting the most accurate reconstruction out of many predictions for a molecule. R This is indicated by the small correlation coefficients between uncertainty estimate A and prediction errors plotted in Supplementary Fig. 6.

Figure 2 also depicts the dependency of prediction errors on the size of the moleculesexamined.Fordiatomics,themeanDEis0.155a.u.Asareference,themean DE’s from calculating the bond length of the same set of diatomic molecules with the classical 1 model and the optimized empirical 0.797 model are 1.27 and 0.49 KER KER a.u.,respectively.(KERstandsforthekineticenergyreleasefromCoulombexplosion. Theempiricalfactor0.797wasdeterminedbyminimizingtheerrorofthe c model KER on the test data.) For larger molecules, the average DE and AE distributions gradu- ally shift upwards, implying that it becomes more difficult for the model to learn the underlying X-ray-induced dynamics in molecules containing more atoms. Although the increased prediction error for molecules with eight or nine atoms is still accept- able,itisexpectedtoincreaseevenmoreforlargermolecules.Correspondingexample predictions for the 1,3-cyclohexadiene molecule containing fourteen atoms are shown inSupplementaryFig.7.Asexpected,theyshowlargediscrepanciesfromtheground- truth structures. In addition to the dependence on the number of atoms, predictions for molecules with larger atomic distances tend to result in larger errors as indicated by Supplementary Fig. 8. Training with a more diverse dataset that includes larger moleculeswithlargeratomicdistanceswouldbeneededtobreaktheselimitationsand further extend the applicability of MOLEXA in the future.

The model has two other major limitations. One originates from the non-unique mapping between ion momentum vectors and real-space molecular geometries as revealed by Ref.41, which can increase the overall prediction error. In future devel- opments, the non-uniqueness issue can be mitigated by a model that makes use of multiple coincidence channels as input data, in contrast to a single channel used by the current model. The other limitation is that the input data must include all ions produced by Coulomb explosion of the molecule. But in a CEI experiment, the com- pletecoincidentiondetectionfromeachexplosionisarelativelylow-probabilityevent, S especially for larger systems. Our current way to circumvent this limitation is by dis- tillingtherequiredfull-coincidencemomentumvectorsfromtShedataaccumulatedover many explosion events21–24. This approach is used to prEepare the input data of the experimentally studied molecules discussed in the next section. Further development R of the neural network is needed to make it capable of molecular structure prediction P from partial coincidence sets, which would be particularly beneficial for single-event reconstruction.

Application of MOLEXA

In this section, we first demonstrateEthe ability of MOLEXA to perform the inversion of experimental data into real spaLce molecular geometries. For this, we used the data acquired during several experiments carried out at the European X-ray Free-Electron C Laser Facility to reconstruct the equilibrium geometry of ground-state molecules I including water, tetrafluoromethane and ethanol. No further molecule-specific input T was provided to the model for retrieving the structures. The experiments were per- R formed in multiple CEI beamtimes using the COLTRIMS (Cold Target Recoil Ion A Momentum SpeActroscopy) Reaction Microscope42 at the Small Quantum Systems (SQS) instrument. The beamtimes and X-ray pulse parameters are summarized in Supplementary Table 4. In all experiments, the molecular samples were delivered to theinteractionregionthroughasupersonicexpansionfollowedbythreeskimmersand an adjustable collimator. The distance from the nozzle to the interaction region is about 54 cm. The pressure in the main chamber was maintained at 1 × 10−11 mbar. The focal spot size of the X-ray beam was about 1.5 –3 µm and the X-ray pulse duration was less than 25 fs based on the 250pC electron bunch charge. The ion frag- ments produced in the interaction region were guided by a homogeneous electric field to a time- and position-sensitive detector. The lab-frame momentum vectors of the ion fragments were then reconstructed from the detector readouts, which were subse- quentlytransformedintothemolecularframewiththealignmentprocedurediscussed in the Training subsection.

In real molecules, the nuclear ground state exhibits a spatial distribution due to the uncertainty principle, such that different nuclear geometries are sampled even in the absence of excitation. Because the mapping between ion momentum and molecu- lar geometry is nonlinear, averaging in momentum space does not strictly correspond to averaging in real space. In principle, the equilibrium molecular structure is best estimated by reconstructing geometries from individual single-shot, full-coincidence Coulomb explosion events and averaging the predicted structures in real space. But single-shotfull-coincidencedetectionremainschallengingfortoday’sCEIexperiments. In the following, the structures are instead predicted from the coincident momen- tum vectors (Supplementary Figs. 9-14) obtained by averaging over the single-shot data. This procedure yields an approximate equilibrium geometry from the ensemble- averagedobservables,ratherthanthefullnuclearprobabilitydistributionoranysingle instantaneous molecular configuration.

As a first example, Figure 3a shows the reconstruction of the molecular structure ofwatermolecules.Theemployeddataset43 usedinthisanalysisisidenticaltotheone used in Ref.16. The 2D heatmap in the left-most part oSf the panel displays the exper- imentally measured molecular-frame momentum diEstribution of two protons detected in coincidence with a singly charged oxygen ion. Next to it in the middle, we show an R illustration of the centroids of the momentum distributions of the three ions, which P serve as the input for MOLEXA. The results obtained from the reconstruction are shown on the right. The reconstructed mol ecular geometry (opaque) is plotted on top N ofthegroundtruth(semi-transparent),withthecorrespondingMAEbeing0.296a.u. Themean(maximum)DEandAEaIre0.674(1.199)a.u.and18.459(27.689)degrees, I respectively. Next, we test the m odel on tetrafluoromethane, a molecule consisting E of five atoms. The corresponding reconstruction is shown in Fig. 3b. The left-most L panel depicts the momentaLof the three of the four fluorine ions in a molecular frame C spanned by the fourth flCuorine ion and one of the three. The reconstructed position- space geometry has a MAE of 0.238 a.u., a mean (maximum) DE of 0.66 (1.117) a.u., I and a mean (maxiTmum) AE of 5.943 (17.173) degrees. The data was recorded during R the commissioning of the SQS reaction microscope44.

Asabenchmarkforanapplicationtomoleculeswithuptonineatoms,weapplied MOLEXA A toaCEIdatasetrecordedforethanolmolecules45.Theresultsareshownin Fig. 3c. The 2D maps show the molecular-frame momentum distributions of protons in the coincidence channel O+/C+/C+/H+, viewed from three different perspectives. We added the corresponding orientation of the real-space molecule to the top of each graph to aid the identification of the six protons in the momentum maps. The input tothemodelisagainobtainedbytakingthecentroidsofthemomentumdistributions of the nine ions. The retrieved molecular geometry plotted together with the ground truth at the right has a MAE of 0.429 a.u. The mean (maximum) DE and AE are 1.024 (2.007) a.u. and 9.011 (32.319) degrees, respectively. More details on the recon- structionfromexperimentaldata,includingthemomentumdistributionsofotherions not displayed in Fig. 3, the momentum centroid data, and the reconstructed atomic coordinates as well as the predicted error estimates, can be found in Supplementary Figs. 9-14 and Supplementary Tables 5-7.

The ultimate aim of CEI is to directly observe molecular dynamics during a chemical reaction in a time-resolved manner. In order to achieve this, coincident momentum-space fragmentation patterns are measured at different instants during the chemical reaction, thus allowing to study the molecular structural changes as the chemicalreactionunfoldsonfemtosecondorlongertimescales.Inthefollowingexam- ple, we exploit MOLEXA to reconstruct the different geometries of cyclobutene as predicted by ab initio simulations46. The electrocyclic reactions of cyclobutene repre- sent a textbook example of pericyclic reactions that are among the important classes of chemical reactions in organic chemistry. Figure 4a shows that MOLEXA is capable ofreconstructingdifferentpossiblegeometricalchanges,includingringopening,twist- ing, and proton migration, after cyclobutene is excited from the ground state (S ) to 0 the S state. In Fig. 4b, MOLEXA is used to reconstruct position-space ”snapshots” 1 ofcyclobuteneasitundergoesaring-openingreaction.Thereconstructionsshowthat themodelcanprovideinsightabouttheoverallstructureandidentifygrossstructural rearrangements such as proton migration and ring opening. Further details on the reconstruction of cyclobutene geometries can be found in Supplementary Tables 8-14.

More examples demonstrating the model’s capability to reSconstruct varying struc- tures of molecules are displayed in Supplementary Fig. 15E. It is worth noting that the reconstructions in Fig. 4 are idealized scenarios where the input momentum vectors correspond to a single well-defined molecular geometry. This can work if the input coincidence data is assembled from a single explosion event. For time-resolved CEI experiments,theaccumulateddataforeachtime delaytypicallyresultsfromamixture of geometries corresponding to multiple quantum states. Classification of the experi- mental data into individual states, basedIon, e.g., ion charge state characteristics and kinetic energies18, is hence required pri or to applying MOLEXA for reconstruction.

Methods
Dataset creation

Fortheabinitiosimulation,weusedatheoreticalmodelthatcombinesMonteCarlo/ MolecularDynamicssimulations(MC/MD)40 withaclassicalover-the-barrier(COB) model47–49 totrackinner-shellphotoionization,Auger-Meitnercascades,valenceelec- tron redistribution, and nuclear dynamics. Photoabsorption and inner-shell cascade processes were modeled using a Monte Carlo method to calculate quantum elec- tron transition probabilities across all participating electronic configurations (ECs), S including ground, core-excited, and valence-excited states. The electronic-structure calculations were based on the relativistic Hartree-FocSk-Slater (HFS) method, which provided bound-state and continuum wavefunctions for computing cross sections of E photoionization,shake-off,electron-impactionization,andelectron-ionrecombination, R as well as Auger-Meitner and fluorescence decay rates. The molecular-dynamics com- ponenttrackedthemotionofatoms,ions,anPddelocalizedionizedelectrons.TheCOB model simulates electron-transfer dynami cs in the valence shell. In this model, an electron fills a vacancy in the valencNe shell of a neighboring atom when its bind- ing energy is higher than the Coulomb barrier. When the atoms are far apart, the I resulting Coulomb barrier suppre sses electron transfer. Electron transfer takes place instantaneously when the electEron orbital energy is higher than the Coulomb barrier.

With this ab initio moLdel, the Coulomb explosion of three hundred different moleculeswithfewerthaCntenatomswasfirstsimulatedintheirequilibriumgeometry. TheX-raypulseshaveaphotonenergyof2keV,pulseenergyof1mJ,pulseduration I of 15 fs, and focal spot size of 1 µm. In order to expand the dataset, the simulations T were additionally performed tens of times on each of these molecules after randomly R varyingtheirstructures.BecauseofthestochasticnatureoftheX-rayinteractionwith molecules,Athe atomic charge-state combination of the resulting ion fragments from a moleculecanvaryfromonesimulationtrajectorytoanother.Withahundredthousand trajectories simulated for each molecule at a fixed structure, the number of trajecto- ries ending at each of the possible charge-state combinations was enumerated. Only combinations with a count greater than three hundred were considered. The momen- tum of each ion fragment was obtained by averaging all trajectories. For every such charge-state combination, the atomic number, charge state, and momentum of all ion fragments,aswellastheinitialatomiccoordinatesofthemolecule,wereincludedinto the dataset as a single entry. With an average of ten structures simulated for each of the three hundred molecules and an average of about ten charge-state combinations producedfromamoleculeataparticularstructure,thedatasetcontains76000entries. It was further split into training (80%), validation (10%), and test (10%) datasets. The training dataset was used in the second step of the two-stage modeling process. ExemplarysamplesfromthetestdatasetareshowninSupplementaryFig.17together with the corresponding predictions.

Becausetheabinitiosimulationwascomputationallyexpensiveandcouldonlybe usedtogenerateasmalldataset,anapproximateforwardCoulombexplosionmodel22 wasusedtocreateadatasetabouttwoordersofmagnitudelarger.Themodeldescribes thecharge-upofeachatominamoleculewithamodifiederrorfunctionthatincreases fromzerotothefinalchargenumberwithinatimewindowcontrolledbytheconstant τ. For the simulation, τ was set to be 45 fs, which was determined by minimizing the discrepancy of the results of the approximate model with respect to those of the ab initio simulations. Using the time-dependent charge states given by the modified errorfunction,theCoulombexplosiondynamicsweresimulatedwiththeRunge-Kutta approachthatpropagatesthetime-dependentpositionsandvelocitiesofionfragments accordingtoclassicalmechanics.The”molecules”usedforthisapproximatesimulation were generated by enumerating all possible combinations of the 9 elements (H, C, N, O, F, Si, P, S, and Cl), with the number of atoms in each combination less than 10. The positions of the atoms in a ”molecule” were sampled from a uniform distribution S ranging from -10 a.u. to 10 a.u. The dataset produced from the simulation with these ”molecules” consists of six million entries and was used forSthe first step of the two- stage modeling process. E

Training details

During training, the reverse diffusion process (S tructure Denoising Module) was run only once for each training step. Instead oNf taking a random structure as input, it starts with a noisified ground-truth structure with the noise level controlled by σ . I i The Structure Denoising Module was trained to denoise this input and generate a geometry Gprediction that is a reconEstruction of the ground truth Ggroundtruth. The i i corresponding loss function is L
(cid:16)C (cid:17)
L =E w ∥Gprediction−Ggroundtruth∥2 , (1)
x I i i i i 2
where the weight wTis set according to Ref.34 and given by
i
R
σ2+σ2
A w i = i σ2σ2 data, (2)
i data
withσ determinedbythestandarddeviationofthemolecularstructuresinthe
data
dataset. In addition to structure reconstruction, MOLEXA was trained to estimate the uncertainty of its predicted structures. As already mentioned in the main text, it first gets the probability si that the uncertainty of the ith predicted coordinate n xprediction falls into the nth bin of the pre-defined uncertainty list [r ,...,r ]. The i 0 200
absoluteerror|xprediction−xgroundtruth|isclassifiedaccordingtothislistasaone-hot
i i
encoded vector qi. The loss function for the uncertainty estimate is then calculated as the averaged cross entropy
L =−E (cid:0) qilog(si) (cid:1) . (3)
u n,i n n

The combined loss function used during training is given by
L=c L +c L , (4)
x x u u
where c and c are the weights of the structure retrieval and uncertainty
x u
estimation loss functions, respectively.

The neural network was trained through two stages. The weights were initialized using the orthogonal Glorot initialization50,51 with a scale of 2 for the linear layers
√ √
and sampled from the uniform distribution with a range from − 3 to 3 for the
embedding layers. In the first stage, it was trained on the large dataset generated by the approximate forward model. The weight c in the loss function was set to 400.
x
And c was set to 0.1 for the first seven epochs and 1 afterwards. For the second
u
stage, the training was performed with the dataset which is about a hundred times smaller and generated by the ab initio forward model. The weights c and c in the
x u
lossfunctionweresetto400and0.01,respectively.Inordertoimprovetheaccuracyof uncertainty predictions, after the two-stage training, the network was further trained onthevalidationdataset.OnlytheUncertaintyEstimationModulewastrainedwhile theothermoduleswerekeptfrozen.Theweightc
x
wassetSto0andc
u
to0.01.During
alltrainingphases,theAdamoptimizer52wasusedforoptimization.Itsparametersβ ,
S 1
β ,andϵwerefixedat0.9,0.99,and10−5,respectively.Thelearningratewaskeptat
2
0.001.Trainingwithlearningratedecaywastested,Ebutdidnotimprovetheprediction
errors.Moredetailsonthetwo-stagetrainingarResummarizedinSupplementaryTable
1.

Coordinate frame transformations for experimentally studied
N
molecules

I
The molecular frames used by the 2D maps in Fig. 3 were defined with a procedure E similar to that described in the main text. The flying direction of a reference ion (O+
for water and ethanol, andLF+ for tetrafluoromethane, as indicated by the arrows in the 2D maps of Fig. 3)Cwas set as x axis. The y axis was then defined such that the momentum vectoIr of a second reference ion (H+ for water, and F+ for tetrafluo- romethane) falls wTithin the positive x-z plane. For ethanol, the second reference ion was chosen to be the C+ ion that has a smaller cosine similarity with respect to R the first reference ion (O+). The y axis was defined such that the momentum vector A of this second reference ion falls within the positive x-y plane. The molecular-frame momentum distributions and their centroids are shown for the four molecules in Sup- plementary Figs. 9-14. The atomic coordinates reconstructed by MOLEXA together with the ground truth are listed in Supplementary Tables 5-7.

Peer review rounds
[{"round": 1, "reviewer_comments": {"reviewer_1": "This work presents a neural network model termed MOLEXA, which utilizes Machine learning technique to retrieve\nmolecular structure from experimental Coulomb explosion imaging (CEI) data. The model is intelligently designed and\nemploys a two-stage training strategy to reconstruct the molecular structures and estimate associated errors. Furthermore,\nwhen combined with time-resolved CEI techniques, MOLEXA will be capable of capturing ‘snapshots’ of a molecule at\ndifferent instant time during a chemical reaction. Although the current version of MOLEXA can provide reconstruction of a\nmolecule structure with fewer than ten atoms with reasonable reliabilities, it represents a significant and important step\ntoward resolving the highly nonlinear inverse problem of inferring molecular configurations from the moment distributions of\nfragment ions produced in rapid Coulomb explosion processes.\nThis work holds considerable importance for the fields of ultrafast imaging in chemistry and physics, particularly in the\ncontext of Coulomb explosion techniques. The machine learning approaches applied are well-founded, and the data,\nanalysis, and figures strongly support the reported conclusions and perspectives.\nThe study demonstrates clear novelty, and the manuscript is well organized with a logical flow. I recommend its publication\nin Nature Communications after minor revisions\nThe following comments and corrections should be addressed.\n1. Under real experimental conditions, the complete detection of all fragment ions in coincidence become increasingly\ndifficult, if not impossible. Reference 33 provides an example in which molecular structure was successfully reconstructed\nfrom only a subset of detected fragment ions – assuming complete ionization and fragmentation of the molecule. As the\nnumber of atoms in a molecule increases, manual reconstruction of molecular structure using such partial data becomes\ninfeasible. Can the machine learning approach presented in this work be extended to reconstruct the molecular structure\nfrom partially detected fragment ions following a complete Coulomb explosion?\n2. In Fig.2, molecular structure is illustrated using ball-and-stick models. In some cases, however, sticks between atoms are\nmissing. please explain the reason for this omission.\n3. Minor points:\nLine 101, the word ‘before’ is duplicated; one should be removed.\nLine 244, “FIG. 4C” should be corrected to “FIG. 3C”", "reviewer_2": "The manuscript “Generative Modeling Enables Molecular Structure Retrieval from Coulomb Explosion Imaging” described a\nmuch-expected effort to use a neutral network model for solving the problem of reconstructing molecular structure from\nCoulomb explosion imaging data.\nThe authors describe their Model and training method of choice, in particular the generation (calculation) of a large lower-\nlevel training data-set and a smaller refined set calculated with a higher-level method. Testing was performed primarily on\nsimulated data, but also applied to experimental Coulomb explosion imaging measurements.\nThis topic is very interesting and can be suitable for the broad audience of the journal. Nevertheless, there are major issues\nthat need to be addressed and revised.\n1. The main issue is the Model performance evaluation criteria:\nThe model is computed in the center of mass frame, so the RMSE and MAE can be misleading – in particular for small\nmolecules. For example, consider the H2O structure described in figure 3: In the center of mass geometry the Oxygen\nposition is automatically correct – giving the impression of successful reconstruction, while it is just a successful choice of\nthe frame of reference. One would benefit much more from considering the error in bond lengths and angles. One would\nexpect Coulomb explosion imaging to have better performance for bond-lengths compared with angles between bonds. And\nthen too, it would be valuable to learn what is the Maximum Error – not the average one which can be small even if one out of 5 atoms is completely in the wrong position…\n2. In the context of Coulomb explosion imaging of angles, it would be valuable to cite and discuss the\nhttps://doi.org/10.1103/PhysRevA.97.033412 work that shows the intrinsic non-unique mapping Coulomb explosion imaging\non classical electrostatic potentials.\n3. One example for the MAE criteria bias is the comparison of H2O and CF4: while inspection of the figure seems to indicate\nthat reconstructed CF4 structure is better – the MAE shows better H2O ( most likely due to the exact Oxygen position that\nreflects only the CM choice ). A good test of this would be to look at CF4 vs CH4 which would appear better when using the\nMAE criteria to evaluate performance.\n4. The problem of the performance criteria is even more pronounced when considering diatomic structures: On pg. 5, the\nauthors claim 100% accuracy when considering diatomic structure. However, MAE of 0.4 a.u. might be not so different from a\nball park guess with most diatomic bonds being somewhere between 2-4 a.u.? How does MOLEXA perform compared to a\nthumbs rule based on the kinetic energy release for the dication: R ~ 1.44/KER ? ( or a more phenomenological\n~1.44/1.2/KER, considering some screening from the remaining electrons ) ? This comparison is important for the readers to\nunderstand the real potential of MOLEXA at least compared with thumb rule approximations….\n5. Following up on point 1, one would also benefit from a discussion of which bonds are better predicted which are less ?\ne.g. C-C single bonds vs. C=C double bonds and C-H bonds in figure S2 etc…?\n6. When considering the experimental Coulomb explosion Imaging: choosing the momentum centroids could be misleading:\nThe ground state wavefunction has an intrinsic spread due to the QM uncertainty principle that samples different initial\ngeometries. Could the authors use MOLEXA to map the experimentally measured distribution and then take the centroid\ngeometry rather than take a centroid position in the momentum correlations that could misrepresent the most probable initial\ngeometry due to the non-linear mapping ( same as the shift between lambda0 and w0 in spectroscopy arising due to their\nnon-linear relation )\n7. Fig 4 shows that while MOLEXA does not seem to provide reliable information about bond lengths and angles - it can\nprovide insight about the overall structure and identify gross structural rearrangements such as ring-opening and proton\nmigration. In the community of Coulomb explosion imaging, there has been a long established effort to identify molecular\nisomerization. An early “classical” example is resolving acetylene and vinylidene , see for example\nhttps://doi.org/10.1103/PhysRevLett.81.3347 , another critical example is identification of chirality by Coulomb explosion\nimaging, see for example https://doi.org/10.1126/science.1246549 that could be cited in addition to the already cited\nPitzer et al. – the MOLEXA ability of reconstructing the isomer state could be much more convincing than reporting on the maybe\nlimited success in ACCURATE structure reconvery.\n8. I think that if the authors focus on the ability of MOLEXA to correctly predict the topological structure ( e.g. identify the\nmolecular bonds in the system )\n9. There are many seemingly technical, but crucial details that are not so well described in the main text. For example: the\nchoice of molecular orientation or the exact nature of the approximations made in the Coulomb explosion simulations. I\nunderstand the desire to limit technical details for the broad audience, but it would be good to read in the main text if the\nsimulated interaction potentials are merely electrostatic approximations – this is important as the readers should grasp the\nquality of the simulated data set as state of the art ab initio molecular dynamics simulations often are required to involve\nhigher level electronic structure Coulomb explosion calculations that help extract time resolved structural information from\nexperimental data, e.g. https://www.nature.com/articles/s41467-024-50759-2 and\nhttps://www.science.org/doi/full/10.1126/sciadv.abq8084\nAdditional minor comments:\n- It is not clear for the readers what exactly was done and achieved in the cyclobutene test on pg.6: I understand that few\nstructures were sampled from ref 38 time-resolved simulations – Coulomb explosion were calculated ( using an unspecified\napproximation ) and the simulated momentum distributions were FED to MOLEXA to reconstruct structure. Without taking\nanything away from the authors achievement, this is not exactly what one would naively understand from “By employing\ntime-resolved CEI datasets, MOLEXA is able to provide ”snapshots” of a molecule at different instants during a chemical\nreaction.” as stated in the conclusions section. I mean that MOLEXA did not provide new information about a chemical\nreaction – maybe one can say that it “demonstrates” the potential of MOLEXA to provide such information.\n- In figure 2, it is not very clear which are the correct positions and which are the simulated ones – specially in the panel\nshowing low accuracy examples. Please revise.", "reviewer_3": "Li et al. shows that it is possible to determine the structure of a molecule from the momentum vectors of the recoil ions\nproduced in X-ray laser-induced Coulomb explosion based on neural networks.\nThe authors describe the two-level training protocol where one level comprises a data set obtained by classical trajectory\ncalculation of the Coulomb explosion process and the second level comprises a data set obtained from an ab initio\ncalculation. In terms of application to experimental data, results are shown for three different molecules (Fig. 2). The\nreconstructed molecular structures are fairly close to the known structures. Finally, to illustrate the use of the retrieval method\nto image molecules during time-dependent processes, it is applied to a simulated data set of different structures of\ncyclobutene as it undergoes ring-opening. Again, the reconstructed structures agree quite well with those from the\nsimulation.\nAs the authors write, timed Coulomb explosion is a promising technique to measure the structure of molecules as they\nundergo various nuclear dynamics. Experimental data consists of time-resolved momentum distributions, and the challenge\nis to convert these to molecular structures, i.e. binding angles and distances. This is a notorious challenge for molecules\nbeyond the diatomic case. The neural network approach, which I have not seen reported previously, introduced here is\ninteresting and could be most useful in future studies.\nThere is certainly room for improvement of the structure accuracy of the method, but I consider the results presented and the\nprospects sufficiently interesting and promising to make a fit for Nature Communications.\nThere is, however, one point the authors should address. The results presented in Fig. 4 are, with due respect, just four\ndifferent static structures of the same molecule that happen to represent the structures of cyclobutene at selected times\nduring the ring-opening process. Since the neural network algorithm was already shown to work for static structure of similar\nsize molecules (Fig. 2, panel 8 and 9) it is hardly surprising that it also works for the four cyclobutene structures. In an\nexperiment, only a small fraction (< 10 %) of all molecules are photoexcited to undergo subsequent dynamics. The majority\nof the molecules are left in the ground state or, depending on the pump pulse intensity, may be singly ionized. Thus, the\ntarget in an experiment is mixed sample. Ideally, the authors should have considered this in their example to more\nconvincingly demonstrate the potential of their novel structure retrieval algorithm to address time-dependent process. Is that\npossible? If not, the authors should at least include a discussion of this point.", "reviewer_4": "The manuscript proposes a new method, MOLEXA, based on diffusion models to reconstruct atomistic positions from\nCoulomb Explosion Imaging data. The approach consists of several components: an embedding module for the ion\nmomenta, a transformer-based architecture to compute subsequent conditioning features, and a denoising model employing\na similar transformer architecture for structure reconstruction. In addition, a confidence model is trained to estimate the\nuncertainties of the reconstructed geometries. For training and validation, the authors develop a new dataset that combines\nboth inexpensive classical simulations and more accurate ab initio calculations, which is also made available online.\nAlthough the model is trained only on molecules containing up to seven atoms, it can also be applied to slightly larger\nsystems (up to nine atoms). Moreover, the method is evaluated on some experimental data.\nThe introduction is exceptionally well written, clearly explaining the task, background, and proposed approach.\nSome architectural and model choices are unclear from the manuscript and require more clarifications:\n- As described in Fig. 2 and Algorithm 12, the diffusion sampler performs only five denoising steps. It is not clear why such a\nsmall number of steps was chosen. Providing ablation studies or a more detailed justification for this design choice would be\nhelpful.\n- The proposed memory transformer module could be motivated more clearly, for example through ablation studies\ncomparing it against simple skip connections or learnable residual gates.\n- The noise schedule is unclear. Is it just a straight application of the EDM schedule?\nThe uncertainty model requires some clarifications as well:\n- Why 200 bins/classes and not a continuous uncertainty/error score?\n- The uncertainty model seems to underestimate the error of reconstructed samples, particularly for samples with high\ndeviations. Especially in those cases, a correct uncertainty value might be important. Could this behavior be a consequence\nof training with a cross-entropy loss where some categories are sparsely populated? Potentially, a weighted cross-entropy\nloss or continuous error could help mitigate this issue.\n- Could the uncertainty model be used to select the best reconstruction of a molecule from multiple samples generated by\nthe model? An ablation study investigating whether this leads to improved reconstruction quality would be interesting. For\nexample, scatter plots similar to those shown in the fourth row of Figure 2, but evaluated for individual molecules.\n- Figure 2, third row: How are the bins selected, given that the uncertainty model is trained to predict 200 bins? Additionally,\nare the bar heights representing the mean values, and the error bars the standard deviation?\nIt is generally unclear how the reconstructed structures are selected, given that the diffusion model can generate multiple\ncandidate structures per molecule. Is only a single reconstruction generated and reported for each molecule? It would be\ninteresting to assess how consistently the model reconstructs structures for a given molecule.\nThe model does not generalize that well to the unseen larger molecules with 8 or 9 atoms. Could a reason for this be that the\nionization process for larger molecules differs from that of smaller ones, e.g. maybe the larger molecules break into smaller\nfragments first? If so, training data in this regime may be crucial for the model to accurately reconstruct such molecules.\nHowever, this also implies that scaling the method to even larger systems would again require corresponding training data,\nwhich might be expensive to obtain.\nAdditional questions and suggestions:\n- A runtime evaluation might help the work, as the inference time of the model is most likely much faster than any existing\nmethod for the reconstruction\n- It may be useful to compare not only absolute positions of the reconstructed structures, but also structural properties such\nas bond lengths and bond angles. Otherwise, a small difference in the alignment could result in large differences.\n- In the tables in the supplementary material, it is unclear how the reported errors are computed.\n- How are the ground truth structures in figure 3 generated?\n- Are there specific types or classes of molecules for which the model performs particularly well or poorly?", "reviewer_5": "I co-reviewed this manuscript with one of the reviewers who provided the listed reports. This is part of the Nature\nCommunications initiative to facilitate training in peer review and to provide appropriate recognition for Early Career\nResearchers who co-review manuscripts."}, "author_response": "We thank the editor and reviewers for their valuable time and insightful comments which are\ncrucial for further improving our paper. We have revised it accordingly. The following is a\npoint-by-point response.\nREVIEWER COMMENTS\nReviewer #1 (Remarks to the Author):\nThis work presents a neural network model termed MOLEXA, which utilizes Machine\nlearning technique to retrieve molecular structure from experimental Coulomb explosion\nimaging (CEI) data. The model is intelligently designed and employs a two-stage training\nstrategy to reconstruct the molecular structures and estimate associated errors.\nFurthermore, when combined with time-resolved CEI techniques, MOLEXA will be capable of\ncapturing ‘snapshots’ of a molecule at different instant time during a chemical reaction.\nAlthough the current version of MOLEXA can provide reconstruction of a molecule structure\nwith fewer than ten atoms with reasonable reliabilities, it represents a significant and\nimportant step toward resolving the highly nonlinear inverse problem of inferring molecular\nconfigurations from the moment distributions of fragment ions produced in rapid Coulomb\nexplosion processes.\nThis work holds considerable importance for the fields of ultrafast imaging in chemistry and\nphysics, particularly in the context of Coulomb explosion techniques. The machine learning\napproaches applied are well-founded, and the data, analysis, and figures strongly support\nthe reported conclusions and perspectives.\nWe appreciate the reviewer’s acknowledgement of the importance of our study.\nThe study demonstrates clear novelty, and the manuscript is well organized with a logical\nflow. I recommend its publication in Nature Communications after minor revisions\nThe following comments and corrections should be addressed.\n1. Under real experimental conditions, the complete detection of all fragment ions in\ncoincidence become increasingly difficult, if not impossible. Reference 33 provides an\nexample in which molecular structure was successfully reconstructed from only a subset of\ndetected fragment ions – assuming complete ionization and fragmentation of the molecule.\nAs the number of atoms in a molecule increases, manual reconstruction of molecular\nstructure using such partial data becomes infeasible. Can the machine learning approach\npresented in this work be extended to reconstruct the molecular structure from partially\ndetected fragment ions following a complete Coulomb explosion?\nThe reviewer pointed out one of the major limitations of MOLEXA. Currently the model is\nincapable of reconstructing molecular structures from incomplete sets of detected\nfragment ions. We agree with the reviewer that prediction from partial coincidence sets is a\nvery promising research direction. We hope to report our progress on this challenge in\nfuture work. To emphasize this limitation, we have added the following to the main text (line\n265-274):\nThe other limitation is that the input data must include all ions produced by Coulomb\nexplosion of the molecule. But in a CEI experiment, the complete coincident ion detection from\neach explosion is a relatively low-probability event, especially for larger systems. Our current\nway to circumvent this limitation is by distilling the required full-coincidence momentum\nvectors from the data accumulated over many explosion events 21-24. This approach is used to\nprepare the input data of the experimentally studied molecules discussed in the next section.\nFurther development of the neural network is needed to make it capable of molecular structure\nprediction from partial coincidence sets, which would be particularly beneficial for single-event\nreconstruction.\n2. In Fig.2, molecular structure is illustrated using ball-and-stick models. In some cases,\nhowever, sticks between atoms are missing. please explain the reason for this omission.\nWe thank the reviewer for notifying us of this missing piece of information. The model only\npredicts the atomic positions but cannot provide bonding information. The bonding is\ninstead automatically decided by the visualization software (PyMOL) based on the atomic\ndistances. The following was added to the caption of Fig. 2:\nThe ball-and-stick models were plotted with PyMOL which considers two atoms bonded if their\ndistance is smaller than a tolerance-expanded sum of their covalent sizes.\n3. Minor points:\nLine 101, the word ‘before’ is duplicated; one should be removed.\nLine 244, “FIG. 4C” should be corrected to “FIG. 3C”\nWe thank the reviewer for pointing them out. They have been corrected.\nReviewer #2 (Remarks to the Author):\nThe manuscript “Generative Modeling Enables Molecular Structure Retrieval from Coulomb\nExplosion Imaging” described a much-expected effort to use a neutral network model for\nsolving the problem of reconstructing molecular structure from Coulomb explosion imaging\ndata.\nThe authors describe their Model and training method of choice, in particular the generation\n(calculation) of a large lower-level training data-set and a smaller refined set calculated with\na higher-level method. Testing was performed primarily on simulated data, but also applied\nto experimental Coulomb explosion imaging measurements.\nThis topic is very interesting and can be suitable for the broad audience of the journal.\nNevertheless, there are major issues that need to be addressed and revised.\nWe appreciate the reviewer’s acknowledgement of our work’s potential broad interest.\n1. The main issue is the Model performance evaluation criteria:\nThe model is computed in the center of mass frame, so the RMSE and MAE can be\nmisleading – in particular for small molecules. For example, consider the H2O structure\ndescribed in figure 3: In the center of mass geometry the Oxygen position is automatically\ncorrect – giving the impression of successful reconstruction, while it is just a successful\nchoice of the frame of reference. One would benefit much more from considering the error\nin bond lengths and angles. One would expect Coulomb explosion imaging to have better\nperformance for bond-lengths compared with angles between bonds. And then too, it would\nbe valuable to learn what is the Maximum Error – not the average one which can be small\neven if one out of 5 atoms is completely in the wrong position…\nWe agree with the reviewer that it would be beneficial to present the prediction errors in\nbond lengths and angles. The model performance has been re-evaluated with the mean and\nmaximum distance errors of all atomic pairs, and the mean and maximum angle errors of\nall directional atomic triplets . The following changes have been made to the manuscript:\n1. The mean and maximum distance error (DE) and angle error (AE) evaluated on the\ntest datasets are now reported as (line 195-198 and 202-204):\nUsing the test dataset of molecules with less than eight atoms, the mean absolute\nerror (MAE) is 0.52 a.u., the mean (maximum) distance error (DE) of all atomic pairs\nin a molecule is 0.98 (2.11) a.u., and the mean (maximum) angle error (AE) of all\ndirectional triplets in a molecule is 13.97 (38.39) degrees.\nFor these molecules larger than those included in the training dataset, the MAE is 0.66\na.u., the mean (maximum) DE is 1.16 (3.43) a.u., and the mean (maximum) AE is 14.12\n(57.91) degrees.\n2. In the top two rows of Fig. 2, the RMSE and MAE values previously displayed for\nexemplary reconstructions were replaced with the mean DE and AE values. And their\nmaximum DE and AE values are now summarized in Supplementary Tables 2-3.\n3. The previous accuracy and MAE plots in the bottom two rows of Fig. 2 have been\nrelocated to Supplementary Fig. 4. They were replaced with the distributions of mean\nDE and AE values as a function of the predicted uncertainties.\n4. The distributions of maximum DE and AE as a function of the predicted uncertainties\nwere added to Supplementary Fig. 5.\n2. In the context of Coulomb explosion imaging of angles, it would be valuable to cite and\ndiscuss the https://doi.org/10.1103/PhysRevA.97.033412 work that shows the intrinsic\nnon-unique mapping Coulomb explosion imaging on classical electrostatic potentials.\nWe thank the reviewer for recommending this work. It is now cited as Ref.41. The following\ndiscussion was added in order to tell the readers that the performance of our current model\ncould be affected by the non-uniquiness problem and suggest a potential route to mitigate this\nissue in the future development (line 260-265):\nThe model has two other major limitations. One originates from the non-unique mapping\nbetween ion momentum vectors and real-space molecular geometries as revealed by Ref.41,\nwhich can increase the overall prediction error. In future developments, the non-uniqueness\nissue can be mitigated by a model that makes use of multiple coincidence channels as input\ndata, in contrast to a single channel used by the current model.\n3. One example for the MAE criteria bias is the comparison of H2O and CF4: while\ninspection of the figure seems to indicate that reconstructed CF4 structure is better – the\nMAE shows better H2O ( most likely due to the exact Oxygen position that reflects only the\nCM choice ). A good test of this would be to look at CF4 vs CH4 which would appear better\nwhen using the MAE criteria to evaluate performance.\nWe agree with the reviewer that MAE is not as good a metric to compare the H O and CF\n2 4\nreconstructions. To address this, the mean and maximum DE and AE values are now\nreported for the reconstructions of all three experimentally studied molecules as the\nfollowing (line 327-329, 333-335, and 345-347):\nH O: The reconstructed molecular geometry (opaque) is plotted on top of the ground truth\n2\n(semi-transparent), with the corresponding MAE being 0.296 a.u. The mean (maximum) DE\nand AE are 0.674 (1.199) a.u. and 18.459 (27.689) degrees, respectively.\nCF : The reconstructed position-space geometry has a MAE of 0.238 a.u., a mean (maximum)\nC H O: The retrieved molecular geometry plotted together with the ground truth at the right\n2 6\nhas a MAE of 0.429 a.u. The mean (maximum) DE and AE are 1.024 (2.007) a.u.and 9.011\n(32.319) degrees, respectively.\n4. The problem of the performance criteria is even more pronounced when considering\ndiatomic structures: On pg. 5, the authors claim 100% accuracy when considering diatomic\nstructure. However, MAE of 0.4 a.u. might be not so different from a ball park guess with\nmost diatomic bonds being somewhere between 2-4 a.u.? How does MOLEXA perform\ncompared to a thumbs rule based on the kinetic energy release for the dication: R ~\n1.44/KER ? ( or a more phenomenological ~1.44/1.2/KER, considering some screening from\nthe remaining electrons ) ? This comparison is important for the readers to understand the\nreal potential of MOLEXA at least compared with thumb rule approximations….\nWe agree with the reviewer that it is important to compare MOLEXA predictions with classical\ncalculations based on the thumb rule and have added the following text (line 241-246):\nFor diatomics, the mean DE is 0.155 a.u. As a reference, the mean DE’s from calculating the\nbond length of the same set of diatomic molecules with the classical 1/KER model and the\noptimized empirical 0.797/KER model are 1.27 and 0.49 a.u., respectively. (KER stands for the\nkinetic energy release from Coulomb explosion. The empirical factor 0.797 was determined by\nminimizing the error of the c/KER model on the test data.)\nWe would like to note that the optimized empirical factor (0.797) is less than 1 because the\nmolecule takes time to be charged up, which results in smaller KER than the instantaneous\ncharge-up model 1/KER. A related experimental example can be found in [Phys.Rev.Research 4,\n013029 (2022)].\nThe mean DE and AE errors from using these two classical models have also been plotted in\nthe bottom left corner of Fig. 2 to facilitate their comparison with MOLEXA prediction errors.\n5. Following up on point 1, one would also benefit from a discussion of which bonds are\nbetter predicted which are less ? e.g. C-C single bonds vs. C=C double bonds and C-H bonds\nin figure S2 etc…?\nWe did not identify the performance dependence on the bond types. As discussed in the paper,\nthe main dependence we found is on the number of atoms in a molecule. In addition, we\ninvestigated the dependence on other factors and found that prediction error also tended to\nincrease with atomic distances. The following text was added along with the new\nSupplementary Fig. 8 showing the dependence of maximum DE and AE values on the\nmaximum atomic distance of a molecule (line 255-257):\nIn addition to the dependence on the number of atoms, predictions for molecules with larger\natomic distances tend to result in larger errors as indicated by Supplementary Fig. 8.\n6. When considering the experimental Coulomb explosion Imaging: choosing the\nmomentum centroids could be misleading: The ground state wavefunction has an intrinsic\nspread due to the QM uncertainty principle that samples different initial\ngeometries. Could the authors use MOLEXA to map the experimentally measured distribution\nand then take the centroid geometry rather than take a centroid position in the momentum\ncorrelations that could misrepresent the most probable initial geometry due to the non-linear\nmapping ( same as the shift between lambda0 and w0 in spectroscopy arising due to their\nnon-linear relation )\nWe thank the reviewer for raising this point and agree that, due to the nonlinear mapping\nbetween ion momentum and real-space molecular geometry, averaging in momentum\nspace does not strictly correspond to averaging in real space. In principle, the equilibrium\nstructure is best estimated by reconstructing molecular geometries from individual\nsingle-shot, full-coincidence Coulomb explosion events and then averaging the predicted\nstructures in real space. At present, however, single-shot full-coincidence detection of all\nionic fragments remains experimentally challenging. In the current work and as the reviewer\nwrote, we obtain the input full-coincidence momentum vectors by averaging the data over\nmany explosion events. They are then used to estimate the equilibrium molecular geometry.\nThe following discussion was added in order to emphasize to the readers that our approach\nis an approximate way to obtaining the equilibrium geometry (line 306-318):\nIn real molecules, the nuclear ground state exhibits a spatial distribution due to the uncertainty\nprinciple, such that different nuclear geometries are sampled even in the absence of\nexcitation. Because the mapping between ion momentum and molecular geometry is\nnonlinear, averaging in momentum space does not strictly correspond to averaging in real\nspace. In principle, the equilibrium molecular structure is best estimated by reconstructing\ngeometries from individual single-shot, full-coincidence Coulomb explosion events and\naveraging the predicted structures in real space. But single-shot full-coincidence detection\nremains challenging for today’s CEI experiments. In the following, the structures are instead\npredicted from the coincident momentum vectors (Supplementary Figs. 9 - 14) obtained by\naveraging over the single-shot data. This procedure yields an approximate equilibrium\ngeometry from the ensemble-averaged observables, rather than the full nuclear probability\ndistribution or any single instantaneous molecular configuration.\n7. Fig 4 shows that while MOLEXA does not seem to provide reliable information about\nbond lengths and angles - it can provide insight about the overall structure and identify\ngross structural rearrangements such as ring-opening and proton migration. In the\ncommunity of Coulomb explosion imaging, there has been a long established effort to\nidentify molecular isomerization. An early “classical” example is resolving acetylene and\nvinylidene , see for example https://doi.org/10.1103/PhysRevLett.81.3347 , another critical\nexample is identification of chirality by Coulomb explosion imaging, see for example\nhttps://doi.org/10.1126/science.1246549 that could be cited in addition to the already cited\nPitzer et al. – the MOLEXA ability of reconstructing the isomer state could be much more\nconvincing than reporting on the maybe limited success in ACCURATE structure reconvery.\nWe thank the reviewer for suggesting these two pioneering studies and have added them as\nRefs.6-7. We also agree with the reviewer that the science cases in them are among the best\nexamples for demonstrating the capability of MOLEXA and hope to make use of them in our\nfuture work.\n8. I think that if the authors focus on the ability of MOLEXA to correctly predict the\ntopological structure ( e.g. identify the molecular bonds in the system )\nWe share this view with the reviewer and have added the texts below (line 366-368 and\n381-384):\nThe reconstructions show that the model can provide insight about the overall structure and\nidentify gross structural rearrangements such as proton migration and ring opening\nIt allows for inverting momentum-space datasets to position space, providing the structure of a\nmolecule right before its explosion by an X-ray pulse and showing particular effectiveness in\nreconstructing the overall structure of molecules.\n9. There are many seemingly technical, but crucial details that are not so well described in\nthe main text. For example: the choice of molecular orientation or the exact nature of the\napproximations made in the Coulomb explosion simulations. I understand the desire to limit\ntechnical details for the broad audience, but it would be good to read in the main text if the\nsimulated interaction potentials are merely electrostatic approximations – this is important\nas the readers should grasp the quality of the simulated data set as state of the art ab initio\nmolecular dynamics simulations often are required to involve higher level electronic\nstructure Coulomb explosion calculations that help extract time resolved structural\ninformation from experimental data, e.g.\nhttps://www.nature.com/articles/s41467-024-50759-2 and\nhttps://www.science.org/doi/full/10.1126/sciadv.abq8084\nWe thank the reviewer for pointing this out and have added more details to the main text. They\ninclude\n1. Moving the following molecular orientation discussion from Methods to the main text (line\n172-181)\nSpecifically, for each molecule, the emission direction of the heaviest ion fragment is\ndefined as the x axis. Then the cosine similarities of the momenta of the other ion\nfragments relative to the x direction are calculated. The ion momentum with the smallest\ncosine similarity is chosen to define the y axis through Gram-Schmidt orthonormalization\nso that this momentum vector would fall within the first quadrant of the x-y plane. Both\nthe momenta of all ion fragments and the initial molecular structure, after being centered\nat the origin, are transformed into this coordinate frame. For diatomics, the coordinate\nsystem is fully determined once the emission direction of the heavier ion fragment is\ndefined as the x axis.\n2. Providing a brief explanation of the nature of the ab initio simulation (line 154-156)\nOne level involves the ab initio calculation of the XFEL-induced Coulomb explosion of\nmolecules, which tracks the quantum transition probabilities across all participating\nelectronic configurations while treating the nuclei as moving in a classical force field 40.\n3. Moving the experimental details from Methods to the main text (line 292-305)\nThe experiments were performed in multiple CEI beamtimes using the COLTRIMS (Cold\nTarget Recoil Ion Momentum Spectroscopy) Reaction Microscope42 at the Small\nQuantum Systems (SQS) instrument. The beamtimes and x-ray pulse parameters are\nsummarized in Table 4. In all experiments, the molecular samples were delivered to the\ninteraction region through a supersonic expansion followed by three skimmers and an\nadjustable collimator. The distance from the nozzle to the interaction region is about 54\ncm. The pressure in the main chamber was maintained at 1 × 10−11 mbar. The focal\nspot size of the X-ray beam was about 1.5 –3 μm and the X-ray pulse duration was less\nthan 25 fs based on the 250pC electron bunch charge. The ion fragments produced in\nthe interaction region were guided by a homogeneous electric field to a time- and\nposition-sensitive detector. The lab-frame momentum vectors of the ion fragments were\nthen reconstructed from the detector readouts, which were subsequently transformed\ninto the molecular frame with the alignment procedure discussed in the Training\nsubsection and Supplementary Methods.\n4. Several details related to the neural network, including model architecture, inference\ntime and uncertainty estimation, are now added to the main text. (line 121-123, 133-134,\n143-146, 204-210, 216-219, and 224-229)\nAdditional minor comments:\n- It is not clear for the readers what exactly was done and achieved in the cyclobutene test\non pg.6: I understand that few structures were sampled from ref 38 time-resolved\nsimulations – Coulomb explosion were calculated ( using an unspecified approximation )\nand the simulated momentum distributions were FED to MOLEXA to reconstruct structure.\nWithout taking anything away from the authors achievement, this is not exactly what one\nwould naively understand from “By employing time-resolved CEI datasets, MOLEXA is able\nto provide ”snapshots” of a molecule at different instants during a chemical reaction.” as\nstated in the conclusions section. I mean that MOLEXA did not provide new information\nabout a chemical reaction – maybe one can say that it “demonstrates” the potential of\nMOLEXA to provide such information.\nWe agree with the reviewer that our previous way of presenting the cyclobutene example\ncan give readers a false impression of MOLEXA’s capability. As written in the response to\npoint 8, the following text is added to the cyclobutene discussion to clarify our message\n(line 366-368):\nThe reconstructions show that the model can provide insight about the overall structure and\nidentify gross structural rearrangements such as proton migration and ring opening.\nWe have also taken the reviewer’s advice and revised the related text in the Conclusions\nsection as the following (line 387-390):\nMOLEXA has the potential to provide ”snapshots” of a molecule at different instants during a\nchemical reaction, which can enable the use of the CEI technique for direct reconstruction of\nmolecular dynamics in position space as they unfold on their natural time scales.\n- In figure 2, it is not very clear which are the correct positions and which are the simulated\nones – specially in the panel showing low accuracy examples. Please revise.\nWe thank the reviewer for bringing out this issue. To make the predicted and ground-truth\nstructures more differentiable, we have increased the transparency of the latter in Fig. 2. This\nchange was also applied to the ball-and-stick models in Fig. 3 and Supplementary Fig. 15.\nReviewer #3 (Remarks to the Author):\nLi et al. shows that it is possible to determine the structure of a molecule from the\nmomentum vectors of the recoil ions produced in X-ray laser-induced Coulomb explosion\nbased on neural networks.\nThe authors describe the two-level training protocol where one level comprises a data set\nobtained by classical trajectory calculation of the Coulomb explosion process and the\nsecond level comprises a data set obtained from an ab initio calculation. In terms of\napplication to experimental data, results are shown for three different molecules (Fig. 2).\nThe reconstructed molecular structures are fairly close to the known structures. Finally, to\nillustrate the use of the retrieval method to image molecules during time-dependent\nprocesses, it is applied to a simulated data set of different structures of cyclobutene as it\nundergoes ring-opening. Again, the reconstructed structures agree quite well with those\nfrom the simulation.\nAs the authors write, timed Coulomb explosion is a promising technique to measure the\nstructure of molecules as they undergo various nuclear dynamics. Experimental data\nconsists of time-resolved momentum distributions, and the challenge is to convert these to\nmolecular structures, i.e. binding angles and distances. This is a notorious challenge for\nmolecules beyond the diatomic case. The neural network approach, which I have not seen\nreported previously, introduced here is interesting and could be most useful in future\nstudies.\nThere is certainly room for improvement of the structure accuracy of the method, but I\nconsider the results presented and the prospects sufficiently interesting and promising to\nmake a fit for Nature Communications.\nWe are thankful for the reviewer’s recognition of our work’s potential.\nThere is, however, one point the authors should address. The results presented in Fig. 4 are,\nwith due respect, just four different static structures of the same molecule that happen to\nrepresent the structures of cyclobutene at selected times during the ring-opening process.\nSince the neural network algorithm was already shown to work for static structure of similar\nsize molecules (Fig. 2, panel 8 and 9) it is hardly surprising that it also works for the four\ncyclobutene structures. In an experiment, only a small fraction (< 10 %) of all molecules are\nphotoexcited to undergo subsequent dynamics. The majority of the molecules are left in the\nground state or, depending on the pump pulse intensity, may be singly ionized. Thus, the\ntarget in an experiment is mixed sample. Ideally, the authors should have considered this in\ntheir example to more convincingly demonstrate the potential of their novel structure\nretrieval algorithm to address time-dependent process. Is that possible? If not, the authors\nshould at least include a discussion of this point.\nWe thank the reviewer for raising this point. It made us realize our previous way of\npresenting Fig. 4 could give readers a false impression that the neural network could be\ndirectly applied to a time-resolved CEI experiment. Currently, the model can only make\nreconstructions if the input momentum vectors originate from a single initial molecular\ngeometry and hence cannot deal with mixed structures. To address this point, we added the\nfollowing to the main text (line 371-378):\nIt is worth noting that the reconstructions in Fig. 4 are idealized scenarios where the input\nmomentum vectors correspond to a single well-defined molecular geometry. This can work if\nthe input coincidence data is assembled from a single explosion event. For time-resolved CEI\nexperiments, the accumulated data for each time delay typically results from a mixture of\ngeometries corresponding to multiple quantum states. Classification of the experimental data\ninto individual states, based on, e.g., ion charge state characteristics and kinetic energies18, is\nhence required prior to applying MOLEXA for reconstruction.\nReviewer #4 (Remarks to the Author):\nThe manuscript proposes a new method, MOLEXA, based on diffusion models to\nreconstruct atomistic positions from Coulomb Explosion Imaging data. The approach\nconsists of several components: an embedding module for the ion momenta, a\ntransformer-based architecture to compute subsequent conditioning features, and a\ndenoising model employing a similar transformer architecture for structure reconstruction.\nIn addition, a confidence model is trained to estimate the uncertainties of the reconstructed\ngeometries. For training and validation, the authors develop a new dataset that combines\nboth inexpensive classical simulations and more accurate ab initio calculations, which is\nalso made available online. Although the model is trained only on molecules containing up\nto seven atoms, it can also be applied to slightly larger systems (up to nine atoms).\nMoreover, the method is evaluated on some experimental data.\nThe introduction is exceptionally well written, clearly explaining the task, background, and\nproposed approach.\nWe thank the reviewer for providing this encouraging feedback for the introduction.\nSome architectural and model choices are unclear from the manuscript and require more\nclarifications:\n- As described in Fig. 2 and Algorithm 12, the diffusion sampler performs only five denoising\nsteps. It is not clear why such a small number of steps was chosen. Providing ablation\nstudies or a more detailed justification for this design choice would be helpful.\nWe agree with the reviewer and have added Supplementary Fig. 18 which shows the\ndependence of prediction errors on the number of denoising (sampling) steps. It can be\nseen that with 5 denoising steps, the prediction errors are already getting close to the\nminimum values that are achievable with less than 100 denoising steps. The following\njustification is added (page 51, line 737-742):\nThe prediction errors evaluated as a function of the number of sampling steps (NSS) are\ndisplayed in Supplementary Fig. 18. The overall minimum is reached at the NSS of 33. Since\nfive sampling steps can already get to an error within 1% of the minimum value (achievable\nwith less than 100 sampling steps), it was chosen as a good compromise between prediction\naccuracy and inference time.\n- The proposed memory transformer module could be motivated more clearly, for example\nthrough ablation studies comparing it against simple skip connections or learnable residual\ngates.\nWe agree that a clearer justification of the memory transformer module is important. We\ncarried out an ablation study in which the memory mechanism was replaced by skip\nconnections while keeping all other components unchanged. The results show that the\noriginal model consistently outperforms the “skip connection” variant, reducing the mean\n(maximum) atomic distance errors by 3.6% (4.0%) and the mean (maximum) angle errors by\n1.3% (4.4%). We added the following to the main text (line 121-123):\nIn comparison to using skip connections, the memory mechanism is found to suppress the\nmean (maximum) atomic distance and angle errors by 3.6% (4%) and 1.3% (4.4%), respectively.\n- The noise schedule is unclear. Is it just a straight application of the EDM schedule?\nYes, it is mostly an application of the EDM schedule from Ref.34, with the detailed\nparameters provided in Algorithm 12. We added the following to the main text (line\n133-134):\nDuring inference, this structure is iteratively refined by a diffusion sampler with a noise\nschedule adapted from Ref.34.\nThe uncertainty model requires some clarifications as well:\n- Why 200 bins/classes and not a continuous uncertainty/error score?\nIndeed, the model first predicts the probabilities of the prediction error falling into each of\nthe 200 bins/classes. Then the uncertainty is calculated by taking a probability-weighted\nsum of the 200 bin values. This uncertainty produced this way is continuous. To clarify this,\nwe added the following to the main text (line 143-146):\nIt pre-defines the uncertainty bins r = [0, 0.05, 0.1, . . . , 9.95] and estimates the probability that\nthe prediction error falls within each of these bins. For each coordinate of the predicted atomic\npositions, the uncertainty is then calculated as the probability-weighted sum of the bin values.\n- The uncertainty model seems to underestimate the error of reconstructed samples,\nparticularly for samples with high deviations. Especially in those cases, a correct\nuncertainty value might be important. Could this behavior be a consequence of training with\na cross-entropy loss where some categories are sparsely populated? Potentially, a weighted\ncross-entropy loss or continuous error could help mitigate this issue.\nWe thank the reviewer for raising this issue. It led us to spot a mistake in the accuracy and\nMAE plots in the previous Fig. 2. The predicted uncertainty was designed to match the MAE.\nPrevious plots accidentally used the square root of the uncertainties and RMSE (typically\nlarger than MAE), which made the prediction error look larger than the predicted uncertainty.\nThis mistake has been corrected and the accuracy and MAE plots in the previous Fig.2 have\nbeen relocated to Supplementary Fig. 4. It shows a more reasonable correspondence\nbetween the MAE’s and the predicted uncertainties.\nWe agree with the reviewer, even with the corrected plot, the model can underestimate the\nprediction errors, especially for samples with large prediction errors. As the reviewer wrote,\nit could be caused by some uncertainty bins being sparsely populated and could potentially\nbe mitigated with a weighted cross-entropy loss. To emphasize this limitation, we added the\nfollowing to the main text (line 224-229):\nIt can be observed from these heatmaps that the predicted uncertainties can underestimate\nthe reconstruction errors, especially for larger molecules with larger prediction errors. We\nattribute this to a training bias caused by the Uncertainty Estimation Module being trained on\nreconstructions most of which have relatively low errors. Such bias can potentially be\nmitigated with a weighted cross-entropy loss in future developments.\n- Could the uncertainty model be used to select the best reconstruction of a molecule from\nmultiple samples generated by the model? An ablation study investigating whether this\nleads to improved reconstruction quality would be interesting. For example, scatter plots\nsimilar to those shown in the fourth row of Figure 2, but evaluated for individual molecules.\nWe thank the reviewer for making this suggestion. We randomly selected 1000 molecules,\nperformed 1000 independent reconstructions for each molecule. The distributions of the\ncorrelation coefficient between reconstruction errors and predicted uncertainties are plotted\nin Supplementary Fig. 6. They show no clear positive correlation, although their means are\nslightly positive. It makes us conclude that the benefit of using uncertainty values to select\nreconstructions is marginal. The following was added to the main text to discuss this idea\n(line 235-239):\nIt needs to be noted that although an uncertainty estimate can assist the assessment of the\ntrustworthiness of a prediction, it is not effective in selecting the most accurate reconstruction\nout of many predictions for a molecule. This is indicated by the small correlation coefficients\nbetween uncertainty estimate and prediction errors plotted in Supplementary Fig. 6.\n- Figure 2, third row: How are the bins selected, given that the uncertainty model is trained to\npredict 200 bins? Additionally, are the bar heights representing the mean values, and the\nerror bars the standard deviation?\nThe predicted uncertainty is a continuous value given by the probability-weighted sum of\nthe 200 bin values. Each of the bins in the third row of Fig. 2 has a size of 0.1 a.u. The bar\nheight represents the percentage of the model predictions with a MAE less than 0.6 a.u. out\nof all predictions which have an uncertainty falling into that bin. The error bar is derived\nfrom the standard deviations. We added the text below to the caption of Supplementary Fig.\n4:\nAt each uncertainty bin of size 0.1 a.u., the bar height is the percentage of predictions with a\nMAE below 0.6~a.u out of all predictions which have an uncertainty belonging to that bin. The\nerror bar is calculated from standard deviations, assuming a Poisson distribution for the\nnumber of counts falling into each bin.\nIt is generally unclear how the reconstructed structures are selected, given that the diffusion\nmodel can generate multiple candidate structures per molecule. Is only a single\nreconstruction generated and reported for each molecule? It would be interesting to assess\nhow consistently the model reconstructs structures for a given molecule.\nWe thank the reviewer for this suggestion. Each reconstruction is indeed from a single model\nprediction. We found little discrepancies between independent predictions. Quantitatively, this is\nreflected by the distribution of standard deviations of model predictions in Supplementary Fig. 2,\nwhich reveals that the distribution is mostly below 0.075 a.u. We added the text below to discuss\nthis point (line 206-210):\nEach of the reconstructed molecular structures is from a single model prediction. Based on the\nstandard deviations of the model predictions plotted in Supplementary Fig. 2, the expected\ndiscrepancy between the corresponding atomic coordinates of two independent reconstructions\nfor a single molecule is about 0.07 a.u.\nThe model does not generalize that well to the unseen larger molecules with 8 or 9 atoms.\nCould a reason for this be that the ionization process for larger molecules differs from that\nof smaller ones, e.g. maybe the larger molecules break into smaller fragments first? If so,\ntraining data in this regime may be crucial for the model to accurately reconstruct such\nmolecules. However, this also implies that scaling the method to even larger systems would\nagain require corresponding training data, which might be expensive to obtain.\nAs the reviewer wrote, the charge-up and fragmentation processes for larger molecules\ncould be different from small ones. And we agree that training with larger molecules is\ncrucial to improve its performance. We are working on ways to more efficiently produce\ntraining datasets for larger molecules and hope to report our progress in this respect in a\nfuture work.\nAdditional questions and suggestions:\n- A runtime evaluation might help the work, as the inference time of the model is most likely\nmuch faster than any existing method for the reconstruction\nThe runtime evaluation has been added to Supplementary Fig. 1, with the corresponding\ntext added as the following (line 204-206):\nThe inference time distribution evaluated on all test molecules is plotted in Supplementary Fig.\n1 and has a mean of 59.8 ms.\n- It may be useful to compare not only absolute positions of the reconstructed structures,\nbut also structural properties such as bond lengths and bond angles. Otherwise, a small\ndifference in the alignment could result in large differences.\nWe agree with the reviewer and have added evaluation metrics related to bond lengths and\nangles. The model was re-evaluated based on the distance errors of all atomic pairs and\nangle errors of all directional triplets in a molecule. Related changes in the manuscript\ninclude\n1. The mean and maximum distance error (DE) and angle error (AE) evaluated on the\ntest datasets are now reported as (line 195-198 and 202-204):\nUsing the test dataset of molecules with less than eight atoms, the mean absolute\nerror (MAE) is 0.52 a.u., the mean (maximum) distance error (DE) of all atomic pairs\nin a molecule is 0.98 (2.11) a.u., and the mean (maximum) angle error (AE) of all\ndirectional triplets in a molecule is 13.97 (38.39) degrees.\nFor these molecules larger than those included in the training dataset, the MAE is 0.66\na.u., the mean (maximum) DE is 1.16 (3.43) a.u., and the mean (maximum) AE is 14.12\n(57.91) degrees.\n2. In the top two rows of Fig. 2, the RMSE and MAE values previously displayed for\nexemplary reconstructions were replaced with the mean DE and AE values. And their\nmaximum DE and AE values are now summarized in Supplementary Tables 2-3.\n3. The previous accuracy and MAE plots in the bottom two rows of Fig. 2 have been\nrelocated to Supplementary Fig. 4. They were replaced with the distributions of mean\nDE and AE as a function of the predicted uncertainties.\n4. The distributions of maximum DE and AE as a function of the predicted uncertainties\nwere added to Supplementary Fig. 5.\n- In the tables in the supplementary material, it is unclear how the reported errors are\ncomputed.\nThe errors were the uncertainties directly predicted by the model after taking the\nprobability-weighted sum of the 200 bin values. To clarify it, we added the following to the\nmain text (line 216-219):\nThe Uncertainty Estimation Module produces uncertainties for all coordinates of all atoms in a\nmolecule. The overall uncertainty displayed here (Fig. 2) was obtained by taking an average of\nthe uncertainties of all atomic positions of a molecule.\nThe following text is added to the Uncertainty Estimation Module subsection in\nSupplementary Methods (page 55, line 754-758):\nFor each predicted coordinate, the uncertainty is then calculated as the probability-weighted\nsum of the bin values. Such values are directly used in Supplementary Tables 5 - 7. The\nuncertainty used in Fig. 2 and Supplementary Figs. 4- 5 was calculated by taking an average of\nthe uncertainty values predicted for all atomic coordinates in a molecule.\n- How are the ground truth structures in figure 3 generated?\nThey are from the NIST Computational Chemistry Comparison and Benchmark Database.\nThe text below has been added to the caption of Fig. 3:\nThe ground truth structures are from the NIST Computational Chemistry Comparison and\nBenchmark Database 47.\n- Are there specific types or classes of molecules for which the model performs particularly\nwell or poorly?\nThe main dependence we found is on the number of atoms as discussed in the main text. We\nalso investigated other potentially influencing factors and found that the model also tended to\nhave larger prediction errors for molecules with larger atomic distances. The following text was\nadded along with the new Supplementary Fig. 8 showing the dependence of maximum DE and\nAE on the maximum atomic distance of a molecule (line 255-257):\nIn addition to the dependence on the number of atoms, predictions for molecules with larger\natomic distances tend to result in larger errors as indicated by Supplementary Fig. 8.\nReviewer #4 (Remarks on code availability):\nThe code enables both training the model and performing inference using the trained model\nor pre-trained weights. In addition, the training, validation, and test sets are provided.\nHowever, the packages used and their corresponding versions should be documented\nsomewhere.\nWe thank the reviewer for this suggestion and have listed the following dependencies in the\nenvironment.yml file of the MOELXA repository on GitHub:\ndependencies:\n- python=3.12.2\n- pytorch=2.3.0\n- pytorch-cuda=12.1\n- cudnn=8.9.2\n- fastai=2.7.15\n- fastcore=1.5.33\n- fastprogress=1.0.3\n- numpy=1.26.4\n- pandas=2.2.1\n- matplotlib=3.8.4\n- gdown=5.2.0\n- torch-geometric==2.5.3\n- torcheval==0.0.7\n- einops==0.8.0\nReviewer #5 (Remarks to the Author):\nI co-reviewed this manuscript with one of the reviewers who provided the listed reports.\nThis is part of the Nature Communications initiative to facilitate training in peer review and\nto provide appropriate recognition for Early Career Researchers who co-review manuscripts.\nWe thank the reviewer for helping review our manuscript."}, {"round": 2, "reviewer_comments": {"reviewer_1": "The authors have satisfactorily addressed all of my concerns and have made substantial improvements to the manuscript.\nHaving also reviewed the full rebuttal to the other reviewers, I believe that the authors have thoughtfully considered all\ncomments and critiques. the authors have corrected errors, refined the content of main text ans supplementary materials, and\nenhanced the clarity of the figures. with these revisions, the manuscript is now well-structured and clearly presented for\nreaders. Therefore, I recommend its acceptance for publication in NC in its current form.", "reviewer_2": "The authors made substantial revisions that improved their manuscript and addressed the concerns and comments of the\nreviewers. Specifically, MOLEXA performance for diatomic systems shown to exceed the typically applied thumb rules for\ninterpreting KER in terms of interatomic distances emphasizes the relevance of the presented work to a broad audience.", "reviewer_4": "I thank the authors for their additional experiments and clarifications. All my concerns have been addressed, and I think the\npaper improved significantly.", "reviewer_5": "I co-reviewed this manuscript with one of the reviewers who provided the listed reports. This is part of the Nature\nCommunications initiative to facilitate training in peer review and to provide appropriate recognition for Early Career\nResearchers who co-review manuscripts."}, "author_response": ""}]

Decision letter