Publications
Discovering symbolic models from deep learning with inductive biases
We develop a general approach to distill symbolic representations of a learned deep model by introducing strong inductive biases. We focus on Graph Neural Networks (GNNs). The technique works as follows: we first encourage sparse latent representations when we train a GNN in a supervised setting, then we apply symbolic regression to components of the learned model to extract explicit physical relations. We find the correct known equations, including force laws and Hamiltonians, can be extracted from the neural network. We then apply our method to a non-trivial cosmology example-a detailed dark matter simulation-and discover a new analytic formula which can predict the concentration of dark matter from the mass distribution of nearby cosmic structures. The symbolic expressions extracted from the GNN using our technique also generalized to out-of-distribution data better than the GNN itself. Our approach offers alternative directions for interpreting neural networks and discovering novel physical principles from the representations they learn.
Lagrangian neural networks
Accurate models of the world are built upon notions of its underlying symmetries. In physics, these symmetries correspond to conservation laws, such as for energy and momentum. Yet even though neural network models see increasing use in the physical sciences, they struggle to learn these symmetries. In this paper, we propose Lagrangian Neural Networks (LNNs), which can parameterize arbitrary Lagrangians using neural networks. In contrast to models that learn Hamiltonians, LNNs do not require canonical coordinates, and thus perform well in situations where canonical momenta are unknown or difficult to compute. Unlike previous approaches, our method does not restrict the functional form of learned energies and will produce energy-conserving models for a variety of tasks. We test our approach on a double pendulum and a relativistic particle, demonstrating energy conservation where a baseline approach incurs dissipation and modeling relativity without canonical coordinates where a Hamiltonian approach fails. Finally, we show how this model can be applied to graphs and continuous systems using a Lagrangian Graph Network, and demonstrate it on the 1D wave equation.
The CHIME fast radio burst project: system overview
The Canadian Hydrogen Intensity Mapping Experiment (CHIME) is a novel transit radio telescope operating across the 400-800-MHz band. CHIME is comprised of four 20-m x 100-m semi-cylindrical paraboloid reflectors, each of which has 256 dual-polarization feeds suspended along its axis, giving it a >200 square degree field-of-view. This, combined with wide bandwidth, high sensitivity, and a powerful correlator makes CHIME an excellent instrument for the detection of Fast Radio Bursts (FRBs). The CHIME Fast Radio Burst Project (CHIME/FRB) will search beam-formed, high time-and frequency-resolution data in real time for FRBs in the CHIME field-of-view. Here we describe the CHIME/FRB backend, including the real-time FRB search and detection software pipeline as well as the planned offline analyses. We estimate a CHIME/FRB detection rate of 2-42 FRBs/sky/day normalizing to the rate estimated at 1.4-GHz by Vander Wiel et al. (2016). Likely science outcomes of CHIME/FRB are also discussed. CHIME/FRB is currently operational in a commissioning phase, with science operations expected to commence in the latter half of 2018.
Interpretable machine learning for science with PySR and SymbolicRegression.jl
PySR is an open-source library for practical symbolic regression, a type of machine learning which aims to discover human-interpretable symbolic models. PySR was developed to democratize and popularize symbolic regression for the sciences, and is built on a high-performance distributed back-end, a flexible search algorithm, and interfaces with several deep learning packages. PySR's internal search algorithm is a multi-population evolutionary algorithm, which consists of a unique evolve-simplify-optimize loop, designed for optimization of unknown scalar constants in newly-discovered empirical expressions. PySR's backend is the extremely optimized Julia library SymbolicRegression.jl, which can be used directly from Julia. It is capable of fusing user-defined operators into SIMD kernels at runtime, performing automatic differentiation, and distributing populations of expressions to thousands of cores across a cluster. In describing this software, we also introduce a new benchmark, "EmpiricalBench," to quantify the applicability of symbolic regression algorithms in science. This benchmark measures recovery of historical empirical equations from original and synthetic datasets.
Learned coarse models for efficient turbulence simulation
Turbulence simulation with classical numerical solvers requires high-resolution grids to accurately resolve dynamics. Here we train learned simulators at low spatial and temporal resolutions to capture turbulent dynamics generated at high resolution. We show that our proposed model can simulate turbulent dynamics more accurately than classical numerical solvers at the comparably low resolutions across various scientifically relevant metrics. Our model is trained end-to-end from data and is capable of learning a range of challenging chaotic and turbulent dynamics at low resolution, including trajectories generated by the state-of-the-art Athena++ engine. We show that our simpler, general-purpose architecture outperforms various more specialized, turbulence-specific architectures from the learned turbulence simulation literature. In general, we see that learned simulators yield unstable trajectories; however, we show that tuning training noise and temporal downsampling solves this problem. We also find that while generalization beyond the training distribution is a challenge for learned models, training noise, added loss constraints, and dataset augmentation can help. Broadly, we conclude that our learned simulator outperforms traditional solvers run on coarser grids, and emphasize that simple design choices can offer stability and robust generalization.
Free-space quantum key distribution to a moving receiver
Technological realities limit terrestrial quantum key distribution (QKD) to single-link distances of a few hundred kilometers. One promising avenue for global-scale quantum communication networks is to use low-Earth-orbit satellites. Here we report the first demonstration of QKD from a stationary transmitter to a receiver platform traveling at an angular speed equivalent to a 600 km altitude satellite, located on a moving truck. We overcome the challenges of actively correcting beam pointing, photon polarization and time-of-flight. Our system generates an asymptotic secure key at 40 bits/s.
Rediscovering orbital mechanics with machine learning
We present an approach for using machine learning to automatically discover the governing equations and hidden properties of real physical systems from observations. We train a "graph neural network" to simulate the dynamics of our solar system's Sun, planets, and large moons from 30 years of trajectory data. We then use symbolic regression to discover an analytical expression for the force law implicitly learned by the neural network, which our results showed is equivalent to Newton's law of gravitation. The key assumptions that were required were translational and rotational equivariance, and Newton's second and third laws of motion. Our approach correctly discovered the form of the symbolic force law. Furthermore, our approach did not require any assumptions about the masses of planets and moons or physical constants. They, too, were accurately inferred through our methods. Though, of course, the classical law of gravitation has been known since Isaac Newton, our result serves as a validation that our method can discover unknown laws and hidden properties from observed data. More broadly this work represents a key step toward realizing the potential of machine learning for accelerating scientific discovery.
Predicting the long-term stability of compact multiplanet systems
We combine analytical understanding of resonant dynamics in two-planet systems with machine learning techniques to train a model capable of robustly classifying stability in compact multi-planet systems over long timescales of $10^9$ orbits. Our Stability of Planetary Orbital Configurations Klassifier (SPOCK) predicts stability using physically motivated summary statistics measured in integrations of the first $10^4$ orbits, thus achieving speed-ups of up to $10^5$ over full simulations. This computationally opens up the stability constrained characterization of multi-planet systems. Our model, trained on $\approx 100,000$ three-planet systems sampled at discrete resonances, generalizes both to a sample spanning a continuous period-ratio range, as well as to a large five-planet sample with qualitatively different configurations to our training dataset. Our approach significantly outperforms previous methods based on systems' angular momentum deficit, chaos indicators, and parametrized fits to numerical integrations. We use SPOCK to constrain the free eccentricities between the inner and outer pairs of planets in the Kepler-431 system of three approximately Earth-sized planets to both be below 0.05. Our stability analysis provides significantly stronger eccentricity constraints than currently achievable through either radial velocity or transit duration measurements for small planets, and within a factor of a few of systems that exhibit transit timing variations (TTVs). Given that current exoplanet detection strategies now rarely allow for strong TTV constraints (Hadden et al., 2019), SPOCK enables a powerful complementary method for precisely characterizing compact multi-planet systems. We publicly release SPOCK for community use.
Learning symbolic physics with graph networks
We introduce an approach for imposing physically motivated inductive biases on graph networks to learn interpretable representations and improved zero-shot generalization. Our experiments show that our graph network models, which implement this inductive bias, can learn message representations equivalent to the true force vector when trained on n-body gravitational and spring-like simulations. We use symbolic regression to fit explicit algebraic equations to our trained model's message function and recover the symbolic form of Newton's law of gravitation without prior knowledge. We also show that our model generalizes better at inference time to systems with more bodies than had been experienced during training. Our approach is extensible, in principle, to any unknown interaction law learned by a graph network, and offers a valuable technique for interpreting and inferring explicit causal theories about the world from implicit knowledge captured by deep learning.
Multiple physics pretraining for physical surrogate models
We introduce multiple physics pretraining (MPP), an autoregressive task-agnostic pretraining approach for physical surrogate modeling of spatiotemporal systems with transformers. In MPP, rather than training one model on a specific physical system, we train a backbone model to predict the dynamics of multiple heterogeneous physical systems simultaneously in order to learn features that are broadly useful across systems and facilitate transfer. In order to learn effectively in this setting, we introduce a shared embedding and normalization strategy that projects the fields of multiple systems into a shared embedding space. We validate the efficacy of our approach on both pretraining and downstream tasks over a broad fluid mechanics-oriented benchmark. We show that a single MPP-pretrained transformer is able to match or outperform task-specific baselines on all pretraining sub-tasks without the need for finetuning. For downstream tasks, we demonstrate that finetuning MPP-trained models results in more accurate predictions across multiple time-steps on systems with previously unseen physical components or higher dimensional systems compared to training from scratch or finetuning pretrained video foundation models. We open-source our code and model weights trained at multiple scales for reproducibility.
xval: A continuous number encoding for large language models
Large Language Models have not yet been broadly adapted for the analysis of scientific datasets due in part to the unique difficulties of tokenizing numbers. We propose xVal, a numerical encoding scheme that represents any real number using just a single token. xVal represents a given real number by scaling a dedicated embedding vector by the number value. Combined with a modified number-inference approach, this strategy renders the model end-to-end continuous when considered as a map from the numbers of the input string to those of the output string. This leads to an inductive bias that is generally more suitable for applications in scientific domains. We empirically evaluate our proposal on a number of synthetic and real-world datasets. Compared with existing number encoding schemes, we find that xVal is more token-efficient and demonstrates improved generalization.
A deep-learning approach for live anomaly detection of extragalactic transients
There is a shortage of multiwavelength and spectroscopic follow-up capabilities given the number of transient and variable astrophysical events discovered through wide-field optical surveys such as the upcoming Vera C. Rubin Observatory and its associated Legacy Survey of Space and Time. From the haystack of potential science targets, astronomers must allocate scarce resources to study a selection of needles in real time. Here we present a variational recurrent autoencoder neural network to encode simulated Rubin Observatory extragalactic transient events using 1% of the PLAsTiCC data set to train the autoencoder. Our unsupervised method uniquely works with unlabeled, real-time, multivariate, and aperiodic data. We rank 1,129,184 events based on an anomaly score estimated using an isolation forest. We find that our pipeline successfully ranks rarer classes of transients as more anomalous. Using …
Bifrost: A Python/C Framework for High-Throughput Stream Processing in Astronomy
Radio astronomy observatories with high throughput back end instruments require real-time data processing. While computing hardware continues to advance rapidly, development of real-time processing pipelines remains difficult and time-consuming, which can limit scientific productivity. Motivated by this, we have developed Bifrost: an open-source software framework for rapid pipeline development. Bifrost combines a high-level Python interface with highly efficient reconfigurable data transport and a library of computing blocks for CPU and GPU processing. The framework is generalizable, but initially it emphasizes the needs of high-throughput radio astronomy pipelines, such as the ability to process data buffers as if they were continuous streams, the capacity to partition processing into distinct data sequences (e.g. separate observations), and the ability to extract specific intervals from buffered data. Computing …
A Bayesian neural network predicts the dissolution of compact planetary systems
Despite over three hundred years of effort, no solutions exist for predicting when a general planetary configuration will become unstable. We introduce a deep learning architecture to push forward this problem for compact systems. While current machine learning algorithms in this area rely on scientist-derived instability metrics, our new technique learns its own metrics from scratch, enabled by a novel internal structure inspired from dynamics theory. Our Bayesian neural network model can accurately predict not only if, but also when a compact planetary system with three or more planets will go unstable. Our model, trained directly from short N-body time series of raw orbital elements, is more than two orders of magnitude more accurate at predicting instability times than analytical estimators, while also reducing the bias of existing machine learning algorithms by nearly a factor of three. Despite being trained on compact resonant and near-resonant three-planet configurations, the model demonstrates robust generalization to both non-resonant and higher multiplicity configurations, in the latter case outperforming models fit to that specific set of integrations. The model computes instability estimates up to five orders of magnitude faster than a numerical integrator, and unlike previous efforts provides confidence intervals on its predictions. Our inference model is publicly available in the SPOCK package, with training code open-sourced.
Mitigating radiation damage of single photon detectors for space applications
Single-photon detectors in space must retain useful performance characteristics despite being bombarded with sub-atomic particles. Mitigating the effects of this space radiation is vital to enabling new space applications which require high-fidelity single-photon detection. To this end, we conducted proton radiation tests of various models of avalanche photodiodes (APDs) and one model of photomultiplier tube potentially suitable for satellite-based quantum communications. The samples were irradiated with 106 MeV protons at doses approximately equivalent to lifetimes of 0.6 , 6, 12 and 24 months in a low-Earth polar orbit. Although most detection properties were preserved, including efficiency, timing jitter and afterpulsing probability, all APD samples demonstrated significant increases in dark count rate (DCR) due to radiation-induced damage, many orders of magnitude higher than the 200 counts per second (cps) required for ground-to-satellite quantum communications. We then successfully demonstrated the mitigation of this DCR degradation through the use of deep cooling, to as low as -86 degrees C. This achieved DCR below the required 200 cps over the 24 months orbit duration. DCR was further reduced by thermal annealing at temperatures of +50 to +100 degrees C.
Robust simulation-based inference in cosmology with Bayesian neural networks
Simulation-based inference (SBI) is rapidly establishing itself as a standard machine learning technique for analyzing data in cosmological surveys. Despite continual improvements to the quality of density estimation by learned models, applications of such techniques to real data are entirely reliant on the generalization power of neural networks far outside the training distribution, which is mostly unconstrained. Due to the imperfections in scientist-created simulations, and the large computational expense of generating all possible parameter combinations, SBI methods in cosmology are vulnerable to such generalization issues. Here, we discuss the effects of both issues, and show how using a Bayesian neural network framework for training SBI can mitigate biases, and result in more reliable inference outside the training set. We introduce cosmoSWAG, the first application of Stochastic Weight Averaging to cosmology, and apply it to SBI trained for inference on the cosmic microwave background.
Mangrove: Learning Galaxy Properties from Merger Trees
Efficiently mapping baryonic properties onto dark matter is a major challenge in astrophysics. Although semianalytic models (SAMs) and hydrodynamical simulations have made impressive advances in reproducing galaxy observables across cosmologically significant volumes, these methods still require significant computation times, representing a barrier to many applications. Graph neural networks have recently proven to be the natural choice for learning physical relations. Among the most inherently graph-like structures found in astrophysics are the dark matter merger trees that encode the evolution of dark matter halos. In this paper, we introduce a new, graph-based emulator framework, Mangrove, and show that it emulates the galactic stellar mass, cold gas mass and metallicity, instantaneous and time-averaged star formation rate, and black hole mass—as predicted by an SAM—with an rms error up to 2 …
HIFlow: Generating diverse HI maps and inferring cosmology while marginalizing over astrophysics using normalizing flows
A wealth of cosmological and astrophysical information is expected from many ongoing and upcoming large-scale surveys. It is crucial to prepare for these surveys now and develop tools that can efficiently extract most information. We present HIFlow: a fast generative model of the neutral hydrogen (HI) maps that is conditioned only on cosmology ($\Omega_{m}$ and $\sigma_{8}$) and designed using a class of normalizing flow models, the Masked Autoregressive Flow (MAF). HIFlow is trained on the state-of-the-art simulations from the Cosmology and Astrophysics with MachinE Learning Simulations (CAMELS) project. HIFlow has the ability to generate realistic diverse maps without explicitly incorporating the expected 2D maps structure into the flow as an inductive bias. We find that HIFlow is able to reproduce the CAMELS average and standard deviation HI power spectrum (Pk) within a factor of $\lesssim$ 2, scoring a very high $R^{2} > 90\%$. By inverting the flow, HIFlow provides a tractable high-dimensional likelihood for efficient parameter inference. We show that the conditional HIFlow on cosmology is successfully able to marginalize over astrophysics at the field level, regardless of the stellar and AGN feedback strengths. This new tool represents a first step toward a more powerful parameter inference, maximizing the scientific return of future HI surveys, and opening a new avenue to minimize the loss of complex information due to data compression down to summary statistics.
The SZ flux-mass (Y–M) relation at low-halo masses: improvements with symbolic regression and strong constraints on baryonic feedback
Feedback from active galactic nuclei (AGNs) and supernovae can affect measurements of integrated Sunyaev–Zeldovich (SZ) flux of haloes (YSZ) from cosmic microwave background (CMB) surveys, and cause its relation with the halo mass (YSZ–M) to deviate from the self-similar power-law prediction of the virial theorem. We perform a comprehensive study of such deviations using CAMELS, a suite of hydrodynamic simulations with extensive variations in feedback prescriptions. We use a combination of two machine learning tools (random forest and symbolic regression) to search for analogues of the Y–M relation which are more robust to feedback processes for low masses (); we find that simply replacing Y → Y(1 + M*/Mgas) in the relation makes it remarkably self-similar. This could serve as a robust multiwavelength mass proxy for low-mass clusters and galaxy groups. Our methodology can …
Augmenting astrophysical scaling relations with machine learning: Application to reducing the Sunyaev–Zeldovich flux–mass scatter
Complex astrophysical systems often exhibit low-scatter relations between observable properties (e.g., luminosity, velocity dispersion, oscillation period). These scaling relations illuminate the underlying physics, and can provide observational tools for estimating masses and distances. Machine learning can provide a fast and systematic way to search for new scaling relations (or for simple extensions to existing relations) in abstract high-dimensional parameter spaces. We use a machine learning tool called symbolic regression (SR), which models patterns in a dataset in the form of analytic equations. We focus on the Sunyaev-Zeldovich flux−cluster mass relation (YSZ − M), the scatter in which affects inference of cosmological parameters from cluster abundance data. Using SR on the data from the IllustrisTNG hydrodynamical simulation, we find a new proxy for cluster mass which combines YSZ and concentration …
AstroCLIP: a cross-modal foundation model for galaxies
We present AstroCLIP, a single, versatile model that can embed both galaxy images and spectra into a shared, physically meaningful latent space. These embeddings can then be used - without any model fine-tuning - for a variety of downstream tasks including (1) accurate in-modality and cross-modality semantic similarity search, (2) photometric redshift estimation, (3) galaxy property estimation from both images and spectra, and (4) morphology classification. Our approach to implementing AstroCLIP consists of two parts. First, we embed galaxy images and spectra separately by pretraining separate transformer-based image and spectrum encoders in self-supervised settings. We then align the encoders using a contrastive loss. We apply our method to spectra from the Dark Energy Spectroscopic Instrument and images from its corresponding Legacy Imaging Survey. Overall, we find remarkable performance on all downstream tasks, even relative to supervised baselines. For example, for a task like photometric redshift prediction, we find similar performance to a specifically-trained ResNet18, and for additional tasks like physical property estimation (stellar mass, age, metallicity, and sSFR), we beat this supervised baseline by 19\% in terms of $R^2$. We also compare our results to a state-of-the-art self-supervised single-modal model for galaxy images, and find that our approach outperforms this benchmark by roughly a factor of two on photometric redshift estimation and physical property prediction in terms of $R^2$, while remaining roughly in-line in terms of morphology classification. Ultimately, our approach represents the first cross-modal self-supervised model for galaxies, and the first self-supervised transformer-based architectures for galaxy images and spectra.
Astroclip: Cross-modal pre-training for astronomical foundation models
We present AstroCLiP, a strategy to facilitate the construction of astronomical foundation models that bridge the gap between diverse astronomical observational modalities. We demonstrate that a cross-modal contrastive learning approach between images and spectra of galaxies yields highly informative embeddings of both modalities. In particular, we apply our method on multi-band images and spectrograms from the Dark Energy Spectroscopic Instrument (DESI), and show that: (1) these embeddings are well-aligned between modalities and can be used for accurate cross-modal searches, and (2) these embeddings encode valuable physical information about the galaxies - in particular redshift and stellar mass - that can be used to achieve competitive zero- and few- shot predictions without further finetuning. Additionally, in the process of developing our approach, we also construct a novel, transformer-based model and pretraining approach for galaxy spectra.
Unsupervised Resource Allocation with Graph Neural Networks
We present an approach for maximizing a global utility function by learning how to allocate resources in an unsupervised way. We expect interactions between allocation targets to be important and therefore propose to learn the reward structure for near-optimal allocation policies with a GNN. By relaxing the resource constraint, we can employ gradient-based optimization in contrast to more standard evolutionary algorithms. Our algorithm is motivated by a problem in modern astronomy, where one needs to select-based on limited initial information-among $10^9$ galaxies those whose detailed measurement will lead to optimal inference of the composition of the universe. Our technique presents a way of flexibly learning an allocation strategy by only requiring forward simulators for the physics of interest and the measurement process. We anticipate that our technique will also find applications in a range of resource allocation problems.
GaMPEN: a machine-learning framework for estimating Bayesian posteriors of galaxy morphological parameters
We introduce a novel machine learning framework for estimating the Bayesian posteriors of morphological parameters for arbitrarily large numbers of galaxies. The Galaxy Morphology Posterior Estimation Network (GaMPEN) estimates values and uncertainties for a galaxy's bulge-to-total light ratio ($L_B/L_T$), effective radius ($R_e$), and flux ($F$). To estimate posteriors, GaMPEN uses the Monte Carlo Dropout technique and incorporates the full covariance matrix between the output parameters in its loss function. GaMPEN also uses a Spatial Transformer Network (STN) to automatically crop input galaxy frames to an optimal size before determining their morphology. This will allow it to be applied to new data without prior knowledge of galaxy size. Training and testing GaMPEN on galaxies simulated to match $z < 0.25$ galaxies in Hyper Suprime-Cam Wide $g$-band images, we demonstrate that GaMPEN achieves typical errors of $0.1$ in $L_B/L_T$, $0.17$ arcsec ($\sim 7\%$) in $R_e$, and $6.3\times10^4$ nJy ($\sim 1\%$) in $F$. GaMPEN's predicted uncertainties are well-calibrated and accurate ($<5\%$ deviation) -- for regions of the parameter space with high residuals, GaMPEN correctly predicts correspondingly large uncertainties. We also demonstrate that we can apply categorical labels (i.e., classifications such as "highly bulge-dominated") to predictions in regions with high residuals and verify that those labels are $\gtrsim 97\%$ accurate. To the best of our knowledge, GaMPEN is the first machine learning framework for determining joint posterior distributions of multiple morphological parameters and is also the first application of an STN to optical imaging in astronomy.
Charting Galactic Accelerations with Stellar Streams and Machine Learning
We present a data-driven method for reconstructing the galactic acceleration field from phase-space measurements of stellar streams. Our approach is based on a flexible and differentiable fit to the stream in phase-space, enabling a direct estimate of the acceleration vector along the stream. Reconstruction of the local acceleration field can be applied independently to each of several streams, allowing us to sample the acceleration field due to the underlying galactic potential across a range of scales. Our approach is methodologically different from previous works, since a model for the gravitational potential does not need to be adopted beforehand. Instead, our flexible neural-network-based model treats the stream as a collection of orbits with a locally similar mixture of energies, rather than assuming that the stream delineates a single stellar orbit. Accordingly, our approach allows for distinct regions of the stream to have different mean energies, as is the case for real stellar streams. Once the acceleration vector is sampled along the stream, standard analytic models for the galactic potential can then be rapidly constrained. We find our method recovers the correct parameters for a ground-truth triaxial logarithmic halo potential when applied to simulated stellar streams. Alternatively, we demonstrate that a flexible potential can be constrained with a neural network, though standard multipole expansions can also be constrained. Our approach is applicable to simple and complicated gravitational potentials alike, and enables potential reconstruction from a fully data-driven standpoint using measurements of slowly phase-mixing tidal debris.
Interpretable symbolic regression for data science: Analysis of the 2022 competition
Symbolic regression searches for analytic expressions that accurately describe studied phenomena. The main attraction of this approach is that it returns an interpretable model that can be insightful to users. Historically, the majority of algorithms for symbolic regression have been based on evolutionary algorithms. However, there has been a recent surge of new proposals that instead utilize approaches such as enumeration algorithms, mixed linear integer programming, neural networks, and Bayesian optimization. In order to assess how well these new approaches behave on a set of common challenges often faced in real-world data, we hosted a competition at the 2022 Genetic and Evolutionary Computation Conference consisting of different synthetic and real-world datasets which were blind to entrants. For the real-world track, we assessed interpretability in a realistic way by using a domain expert to judge the trustworthiness of candidate models.We present an in-depth analysis of the results obtained in this competition, discuss current challenges of symbolic regression algorithms and highlight possible improvements for future competitions.
Stability Constrained Characterization of the 23 Myr Old V1298 Tau System: Do Young Planets Form in Mean Motion Resonance Chains?
A leading theoretical expectation for the final stages of planet formation is that disk migration should naturally drive orbits into chains of mean motion resonances (MMRs). In order to explain the dearth of MMR chains observed at Gyr ages ($<1\%$), this picture requires such configurations to destabilize and scramble period ratios following disk dispersal. Strikingly, the only two known stars with three or more planets younger than $\lesssim 100$ Myrs, HR 8799 and V1298 Tau, have been suggested to be in such MMR chains, given the orbits' near-integer period ratios. We incorporate recent transit and radial velocity observations of the V1298 Tau system, and investigate constraints on the system's orbital architecture imposed by requiring dynamical stability on timescales much shorter than the system's age. We show that the recent radial-velocity mass measurement of V1298 Tau $b$ places it within a factor of two of the instability limit, and that this allows us to set significantly lower limits on the eccentricity ($e_b \leq 0.17$ at $99.7\%$ confidence). Additionally, we rule out a resonant chain configuration for V1298 Tau at $\gtrsim 99\%$ confidence. Thus, if the $\sim 23$ Myr-old V1298 Tau system did form as a resonant chain, it must have undergone instability and rearrangement shortly after disk dispersal. We expect that similar stability constrained characterization of future young multi-planet systems will be valuable in informing planet formation models.
Normalizing flows for hierarchical bayesian analysis: A gravitational wave population study
We propose parameterizing the population distribution of the gravitational wave population modeling framework (Hierarchical Bayesian Analysis) with a normalizing flow. We first demonstrate the merit of this method on illustrative experiments and then analyze four parameters of the latest LIGO/Virgo data release: primary mass, secondary mass, redshift, and effective spin. Our results show that despite the small and notoriously noisy dataset, the posterior predictive distributions (assuming a prior over the parameters of the flow) of the observed gravitational wave population recover structure that agrees with robust previous phenomenological modeling results while being less susceptible to biases introduced by less flexible models. Therefore, the method forms a promising flexible, reliable replacement for population inference distributions, even when data is highly noisy.
Symbolic regression with a learned concept library
We present a novel method for symbolic regression (SR), the task of searching for compact programmatic hypotheses that best explain a dataset. The problem is commonly solved using genetic algorithms; we show that we can enhance such methods by inducing a library of abstract textual concepts. Our algorithm, called LaSR, uses zero-shot queries to a large language model (LLM) to discover and evolve concepts occurring in known high-performing hypotheses. We discover new hypotheses using a mix of standard evolutionary steps and LLM-guided steps (obtained through zero-shot LLM queries) conditioned on discovered concepts. Once discovered, hypotheses are used in a new round of concept abstraction and evolution. We validate LaSR on the Feynman equations, a popular SR benchmark, as well as a set of synthetic tasks. On these benchmarks, LaSR substantially outperforms a variety of state-of-the-art SR approaches based on deep learning and evolutionary algorithms. Moreover, we show that LaSR can be used to discover a novel and powerful scaling law for LLMs.
Symbolic Regression on FPGAs for Fast Machine Learning Inference
The high-energy physics community is investigating the potential of deploying machine-learning-based solutions on Field-Programmable Gate Arrays (FPGAs) to enhance physics sensitivity while still meeting data processing time constraints. In this contribution, we introduce a novel end-to-end procedure that utilizes a machine learning technique called symbolic regression (SR). It searches the equation space to discover algebraic relations approximating a dataset. We use PySR (a software to uncover these expressions based on an evolutionary algorithm) and extend the functionality of hls4ml (a package for machine learning inference in FPGAs) to support PySR-generated expressions for resource-constrained production environments. Deep learning models often optimize the top metric by pinning the network size because the vast hyperparameter space prevents an extensive search for neural architecture. Conversely, SR selects a set of models on the Pareto front, which allows for optimizing the performance-resource trade-off directly. By embedding symbolic forms, our implementation can dramatically reduce the computational resources needed to perform critical tasks. We validate our method on a physics benchmark: the multiclass classification of jets produced in simulated proton-proton collisions at the CERN Large Hadron Collider. We show that our approach can approximate a 3-layer neural network using an inference model that achieves up to a 13-fold decrease in execution time, down to 5 ns, while still preserving more than 90% approximation accuracy.
TNT: Vision transformer for turbulence simulations
Turbulence is notoriously difficult to model due to its multi-scale nature and sensitivity to small perturbations. Classical solvers of turbulence simulation generally operate on finer grids and are computationally inefficient. In this paper, we propose the Turbulence Neural Transformer (TNT), which is a learned simulator based on the transformer architecture, to predict turbulent dynamics on coarsened grids. TNT extends the positional embeddings of vanilla transformers to a spatiotemporal setting to learn the representation in the 3D time-series domain, and applies Temporal Mutual Self-Attention (TMSA), which captures adjacent dependencies, to extract deep and dynamic features. TNT is capable of generating comparatively long-range predictions stably and accurately, and we show that TNT outperforms the state-of-the-art U-net simulator on several metrics. We also test the model performance with different components removed and evaluate robustness to different initial conditions. Although more experiments are needed, we conclude that TNT has great potential to outperform existing solvers and generalize to additional simulation datasets.
Modeling the gaia color-magnitude diagram with bayesian neural flows to constrain distance estimates
We demonstrate an algorithm for learning a flexible color-magnitude diagram from noisy parallax and photometry measurements using a normalizing flow, a deep neural network capable of learning an arbitrary multi-dimensional probability distribution. We present a catalog of 640M photometric distance posteriors to nearby stars derived from this data-driven model using Gaia DR2 photometry and parallaxes. Dust estimation and dereddening is done iteratively inside the model and without prior distance information, using the Bayestar map. The signal-to-noise (precision) of distance measurements improves on average by more than 48% over the raw Gaia data, and we also demonstrate how the accuracy of distances have improved over other models, especially in the noisy-parallax regime. Applications are discussed, including significantly improved Milky Way disk separation and substructure detection. We conclude with a discussion of future work, which exploits the normalizing flow architecture to allow us to exactly marginalize over missing photometry, enabling the inclusion of many surveys without losing coverage.
Disentangled sparsity networks for explainable AI
We develop “Disentangled Sparsity Networks,” a model for fine-grained control over the dependencies between input features and output features in a multi-layer perceptron. Latent variables in a deep neural network are difficult to interpret because there are often many redundant such variables, and each relies on the interactions of many input features. Existing sparsity strategies target the parameters of a neural network, but do not guarantee a sparse dependency mapping from inputs to outputs, and as one can show empirically, often result in latent features which are dependent on the complex interactions of many input features. Our architecture and regularization scheme allows one to minimize the dependency of each latent feature in a deep network, making them much easier to interpret. This network can be used as a drop-in replacement for any multi-layer perceptron, and allows the user to fine-tune the dependencies of each output feature on each input feature. We show that this can greatly increase the applicability of symbolic regression-based Explainable AI approaches: for the first time, we show that we can recover the majority of the compressible Euler fluid equations from a trained graph neural network, purely by studying its internal latent features as trained with our model.
SRBench++: Principled benchmarking of symbolic regression with domain-expert interpretation
Symbolic regression searches for analytic expressions that accurately describe studied phenomena. The main promise of this approach is that it may return an interpretable model that can be insightful to users, while maintaining high accuracy. The current standard for benchmarking these algorithms is SRBench, which evaluates methods on hundreds of datasets that are a mix of real-world and simulated processes spanning multiple domains. At present, the ability of SRBench to evaluate interpretability is limited to measuring the size of expressions on real-world data, and the exactness of model forms on synthetic data. In practice, model size is only one of many factors used by subject experts to determine how interpretable a model truly is. Furthermore, SRBench does not characterize algorithm performance on specific, challenging sub-tasks of regression such as feature selection and evasion of local minima. In …
Hierarchical inference of the lensing convergence from photometric catalogs with Bayesian graph neural networks
We present a Bayesian graph neural network (BGNN) that can estimate the weak lensing convergence ($\kappa$) from photometric measurements of galaxies along a given line of sight. The method is of particular interest in strong gravitational time delay cosmography (TDC), where characterizing the "external convergence" ($\kappa_{\rm ext}$) from the lens environment and line of sight is necessary for precise inference of the Hubble constant ($H_0$). Starting from a large-scale simulation with a $\kappa$ resolution of $\sim$1$'$, we introduce fluctuations on galaxy-galaxy lensing scales of $\sim$1$''$ and extract random sightlines to train our BGNN. We then evaluate the model on test sets with varying degrees of overlap with the training distribution. For each test set of 1,000 sightlines, the BGNN infers the individual $\kappa$ posteriors, which we combine in a hierarchical Bayesian model to yield constraints on the hyperparameters governing the population. For a test field well sampled by the training set, the BGNN recovers the population mean of $\kappa$ precisely and without bias, resulting in a contribution to the $H_0$ error budget well under 1\%. In the tails of the training set with sparse samples, the BGNN, which can ingest all available information about each sightline, extracts more $\kappa$ signal compared to a simplified version of the traditional method based on matching galaxy number counts, which is limited by sample variance. Our hierarchical inference pipeline using BGNNs promises to improve the $\kappa_{\rm ext}$ characterization for precision TDC. The implementation of our pipeline is available as a public Python package, Node to Joy.
Predicting the thermal Sunyaev–Zel’dovich field using modular and equivariant set-based neural networks
Theoretical uncertainty limits our ability to extract cosmological information from baryonic fields such as the thermal Sunyaev–Zel'dovich (tSZ) effect. Being sourced by the electron pressure field, the tSZ effect depends on baryonic physics that is usually modeled by expensive hydrodynamic simulations. We train neural networks on the IllustrisTNG-300 cosmological simulation to predict the continuous electron pressure field in galaxy clusters from gravity-only simulations. Modeling clusters is challenging for neural networks as most of the gas pressure is concentrated in a handful of voxels and even the largest hydrodynamical simulations contain only a few hundred clusters that can be used for training. Instead of conventional convolutional neural net (CNN) architectures, we choose to employ a rotationally equivariant DeepSets architecture to operate directly on the set of dark matter particles. We argue that set-based …
Histogram Pooling Operators: An Interpretable Alternative for DeepSets
In this paper we describe the use of a differentiable histogram as a pooling operator in a Deep Set. For some set summarization problems, we argue this is the more interpretable choice compared to traditional aggregations, since it has functional resemblance to manual set summarization techniques. By staying closer to standard techniques, one can make more explicit connections with the learned functional form used by the model. We motivate and test this proposal with a large-scale summarization problem for cosmological simulations: one wishes to predict global properties of the universe via a set of observed structures distributed throughout space. We demonstrate comparable performance to a sum-pool and mean-pool operation over a range of hyperparameters for this problem. In doing so, we also increase the accuracy of traditional forecasting methods from 20% error using our dataset down to 13%, strengthening the argument for using such methods in Cosmology. We conclude by using our operator to symbolically discover an optimal cosmological feature for cosmic voids (which was not possible with traditional pooling operators) and discuss potential physical connections.
Learning general-purpose cnn-based simulators for astrophysical turbulence
Datasets and domain generality Model Stability Temporal coarsening Spatial coarsening Running time Constraint preservation Gener Page 1 Datasets and domain generality Learning general-purpose CNN-based simulators for astrophysical turbulence Alvaro Sanchez-Gonzalez*,1 Kimberly Stachenfeld*,1 Drummond Fielding,2 Dmitrii Kochkov,3 Miles Cranmer,4 Tobias Pfaff,1 Jonathan Godwin,1 Can Cui,2 Shirley Ho,2 Peter Battaglia1 1 DeepMind 2 Flatiron Institute 3 Google Research 4 Princeton University Model dCNN block: stack of 7 dilated CNNs to gradually increase and decrease perceptual range. Our model uses N=4 of these stacks. Perceptual range at each convolution in the stack dCNN Stability Temporal coarsening Spatial coarsening Running time Constraint preservation Generalization to different states Generalization to larger boxes One model → Accurate performance in 4 different domains …
Adversarial noise injection for learned turbulence simulations
Machine learning is a powerful way to learn effective dynamics of physical simulations, and has seen great interest from the community in recent years. Recent work has shown that deep neural networks trained in an end-to-end manner seem capable of learning to predict turbulent dynamics on coarse grids more accurately than classical solvers. All these works point out that adding Gaussian noise to the input during training is indispensable to improve the stability and roll-out performance of learned simulators, as an alternative to training through multiple steps. In this work we bring in insights from robust machine learning and propose to inject adversarial noise to move machine learning systems a step further towards improving generalization in ML-assisted physical simulations. We advocate that training our models on these worst case perturbation instead of model-agnostic Gaussian noise might lead to better rollout and hope that adversarial noise injection becomes a standard tool for ML-based simulations. We show experimentally in the 2D-setting that for certain classes of turbulence adversarial noise can help stabilize model rollouts, maintain a lower loss and preserve other physical properties such as energy. In addition, we identify a potentially more challenging task, driven 2D-turbulence and show that while none of the noise-based attempts significantly improve rollout, adversarial noise helps.
Automated discovery of interpretable gravitational-wave population models
We present an automatic approach to discover analytic population models for gravitational-wave (GW) events from data. As more gravitational-wave (GW) events are detected, flexible models such as Gaussian Mixture Models have become more important in fitting the distribution of GW properties due to their expressivity. However, flexible models come with many parameters that lack physical motivation, making interpreting the implication of these models challenging. In this work, we demonstrate symbolic regression can complement flexible models by distilling the posterior predictive distribution of such flexible models into interpretable analytic expressions. We recover common GW population models such as a power-law-plus-Gaussian, and find a new empirical population model which combines accuracy and simplicity. This demonstrates a strategy to automatically discover interpretable population models in the ever-growing GW catalog, which can potentially be applied to other astrophysical phenomena.
Anomaly detection for multivariate time series of exotic supernovae
Supernovae mark the explosive deaths of stars and enrich the cosmos with heavy elements. Future telescopes will discover thousands of new supernovae nightly, creating a need to flag astrophysically interesting events rapidly for followup study. Ideally, such an anomaly detection pipeline would be independent of our current knowledge and be sensitive to unexpected phenomena. Here we present an unsupervised method to search for anomalous time series in real time for transient, multivariate, and aperiodic signals. We use a RNN-based variational autoencoder to encode supernova time series and an isolation forest to search for anomalous events in the learned encoded space. We apply this method to a simulated dataset of 12,159 supernovae, successfully discovering anomalous supernovae and objects with catastrophically incorrect redshift measurements. This work is the first anomaly detection pipeline for supernovae which works with online datastreams.
The well: a large-scale collection of diverse physics simulations for machine learning
Machine learning based surrogate models offer researchers powerful tools for accelerating simulation-based workflows. However, as standard datasets in this space often cover small classes of physical behavior, it can be difficult to evaluate the efficacy of new approaches. To address this gap, we introduce the Well: a large-scale collection of datasets containing numerical simulations of a wide variety of spatiotemporal physical systems. The Well draws from domain experts and numerical software developers to provide 15TB of data across 16 datasets covering diverse domains such as biological systems, fluid dynamics, acoustic scattering, as well as magneto-hydrodynamic simulations of extra-galactic fluids or supernova explosions. These datasets can be used individually or as part of a broader benchmark suite. To facilitate usage of the Well, we provide a unified PyTorch interface for training and evaluating models. We demonstrate the function of this library by introducing example baselines that highlight the new challenges posed by the complex dynamics of the Well. The code and data is available at https://github.com/PolymathicAI/the_well.
Machine Learning with Physics Knowledge for Prediction: A Survey
This survey examines the broad suite of methods and models for combining machine learning with physics knowledge for prediction and forecast, with a focus on partial differential equations. These methods have attracted significant interest due to their potential impact on advancing scientific research and industrial practices by improving predictive models with small- or large-scale datasets and expressive predictive models with useful inductive biases. The survey has two parts. The first considers incorporating physics knowledge on an architectural level through objective functions, structured predictive models, and data augmentation. The second considers data as physics knowledge, which motivates looking at multi-task, meta, and contextual learning as an alternative approach to incorporating physics knowledge in a data-driven fashion. Finally, we also provide an industrial perspective on the application of these methods and a survey of the open-source ecosystem for physics-informed machine learning.
Workshop Summary: Exoplanet Orbits and Dynamics
Exoplanetary systems show a wide variety of architectures, which can be explained by different formation and dynamical evolution processes. Precise orbital monitoring is mandatory to accurately constrain their orbital and dynamical parameters. Although major observational and theoretical advances have been made in understanding the architecture and dynamical properties of exoplanetary systems, many outstanding questions remain. This paper aims to give a brief review of a few current challenges in orbital and dynamical studies of exoplanetary systems and a few future prospects for improving our knowledge. Joint data analyses from several techniques are providing precise measurements of orbits and masses for a growing sample of exoplanetary systems, both with close-in orbits and with wide orbits, as well as different evolutionary stages. The sample of young planets detected around stars with …
Learning integrable dynamics with action-angle networks
Machine learning has become increasingly popular for efficiently modelling the dynamics of complex physical systems, demonstrating a capability to learn effective models for dynamics which ignore redundant degrees of freedom. Learned simulators typically predict the evolution of the system in a step-by-step manner with numerical integration techniques. However, such models often suffer from instability over long roll-outs due to the accumulation of both estimation and integration error at each prediction step. Here, we propose an alternative construction for learned physical simulators that are inspired by the concept of action-angle coordinates from classical mechanics for describing integrable systems. We propose Action-Angle Networks, which learn a nonlinear transformation from input coordinates to the action-angle space, where evolution of the system is linear. Unlike traditional learned simulators, Action-Angle Networks do not employ any higher-order numerical integration methods, making them extremely efficient at modelling the dynamics of integrable physical systems.
Equivariant and Modular DeepSets with Applications in Cluster Cosmology
We design modular and rotationally equivariant DeepSets for predicting a continuous background quantity from a set of known foreground particles. Using this architecture, we address a crucial problem in Cosmology: modelling the continuous electron pressure field inside massive structures known as “clusters.” Given a simulation of pressureless, dark matter particles, our networks can directly and accurately predict the background electron pressure field. The modular design of our architecture makes it possible to physically interpret the individual components. Our most powerful deterministic model improves by 70% on the benchmark. A conditional-VAE extension yields further improvement by 7%, being limited by our small training set however. We envision use cases beyond theoretical cosmology, for example in soft condensed matter physics, or meteorology and climate science. 2
Five parameters are all you need (in CDM)
The standard cosmological model, with its six independent parameters, successfully describes our observable Universe. One of these parameters, the optical depth to reionization , represents the scatterings that Cosmic Microwave Background (CMB) photons will experience after decoupling from the primordial plasma as the intergalactic medium transitions from neutral to ionized. depends on the neutral hydrogen fraction , which, in turn, should theoretically depend on cosmology. We present a novel method to establish the missing link between cosmology and reionization timeline using symbolic regression. We discover the timeline has a universal shape well described by the Gompertz mortality law, applicable to any cosmology within our simulated data. Unlike the conventional tanh prescription, our model is asymmetric in time and a good fit to astrophysical constraints on . By combining CMB with astrophysical data and marginalizing over astrophysics, we treat as a derived parameter, tightening its constraint to . This approach reduces the error on the amplitude of the primordial fluctuations by a factor of 2.3 compared to Planck's PR3 constraint and provides a commanding constraint on the ionization efficiency . We expect further improvements in the near term as reionization constraints increase and our understanding of reionization advances.
Interpretable Machine Learning for the Physical Sciences
Would Kepler have discovered his laws if machine learning had been around in 1609? Or would he have been satisfied with the accuracy of some black box regression model, leaving Newton without the inspiration to find the law of gravitation? In this thesis I will present a review of machine learning and its use cases in the physical sciences. I will emphasize a major issue facing their use in science: a lack of interpretability. Overparameterized black box models are susceptible to memorizing spurious correlations in training data. Not only does this threaten reported research advances made with machine learning, but it also deprives scientists of our most powerful toolbox: symbolic manipulation and logical reasoning. With this in mind, I will demonstrate a framework for interpretable machine learning, using physically-motivated inductive biases and a new technique “symbolic distillation”. The combination of these …
A Neural Network Subgrid Model of the Early Stages of Planet Formation
Planet formation is a multi-scale process in which the coagulation of $\mathrm{\mu m}$-sized dust grains in protoplanetary disks is strongly influenced by the hydrodynamic processes on scales of astronomical units ($\approx 1.5\times 10^8 \,\mathrm{km}$). Studies are therefore dependent on subgrid models to emulate the micro physics of dust coagulation on top of a large scale hydrodynamic simulation. Numerical simulations which include the relevant physical effects are complex and computationally expensive. Here, we present a fast and accurate learned effective model for dust coagulation, trained on data from high resolution numerical coagulation simulations. Our model captures details of the dust coagulation process that were so far not tractable with other dust coagulation prescriptions with similar computational efficiency.
The ones that got away: chemical tagging of globular cluster-origin stars with Gaia BP/RP spectra
Globular clusters (GCs) are sites of extremely efficient star formation, and recent studies suggest they significantly contributed to the early Milky Way's stellar mass build-up. Although their role has since diminished, GCs' impact on the Galaxy's initial evolution can be traced today by identifying their most chemically unique stars--those with anomalous nitrogen and aluminum overabundances and oxygen depletion. While they are a perfect tracer of clusters, be it intact or fully dissolved, these high-[N/O], high-[Al/Fe] GC-origin stars are extremely rare within the current Galaxy. To address the scarcity of these unusual, precious former GC members, we train a neural network (NN) to identify high-[N/O], high-[Al/Fe] stars using low-resolution Gaia BP/RP spectra. Our NN achieves a classification accuracy of approximately $\approx99\%$ and a false positive rate of around $\approx7\%$, identifying 878 new candidates in the Galactic field. We validate our results with several physically-motivated sanity checks, showing, for example, that the incidence of selected stars in Galactic GCs is significantly higher than in the field. Moreover, we find that most of our GC-origin candidates reside in the inner Galaxy, having likely formed in the proto-Milky Way, consistent with previous research. The fraction of GC candidates in the field drops at a metallicity of [Fe/H]$\approx-1$, approximately coinciding with the completion of spin-up, i.e. the formation of the Galactic stellar disk.
The Multimodal Universe: Enabling Large-Scale Machine Learning with 100TB of Astronomical Scientific Data
We present the MULTIMODAL UNIVERSE, a large-scale multimodal dataset of scientific astronomical data, compiled specifically to facilitate machine learning research. Overall, the MULTIMODAL UNIVERSE contains hundreds of millions of astronomical observations, constituting 100\,TB of multi-channel and hyper-spectral images, spectra, multivariate time series, as well as a wide variety of associated scientific measurements and "metadata". In addition, we include a range of benchmark tasks representative of standard practices for machine learning methods in astrophysics. This massive dataset will enable the development of large multi-modal models specifically targeted towards scientific applications. All codes used to compile the MULTIMODAL UNIVERSE and a description of how to access the data is available at https://github.com/MultimodalUniverse/MultimodalUniverse
Reusability report: Prostate cancer stratification with diverse biologically-informed neural architectures
In Elmarakeby et al., "Biologically informed deep neural network for prostate cancer discovery", a feedforward neural network with biologically informed, sparse connections (P-NET) was presented to model the state of prostate cancer. We verified the reproducibility of the study conducted by Elmarakeby et al., using both their original codebase, and our own re-implementation using more up-to-date libraries. We quantified the contribution of network sparsification by Reactome biological pathways, and confirmed its importance to P-NET's superior performance. Furthermore, we explored alternative neural architectures and approaches to incorporating biological information into the networks. We experimented with three types of graph neural networks on the same training data, and investigated the clinical prediction agreement between different models. Our analyses demonstrated that deep neural networks with distinct architectures make incorrect predictions for individual patient that are persistent across different initializations of a specific neural architecture. This suggests that different neural architectures are sensitive to different aspects of the data, an important yet under-explored challenge for clinical prediction tasks.
Astronomical Image Coaddition with Bundle-Adjusting Radiance Fields
Image coaddition is of critical importance to observational astronomy. This family of methods consisting of several processing steps such as image registration, resampling, deconvolution, and artifact removal is used to combine images into a single higher-quality image. An alternative to these methods that are built upon vectorized operations is the representation of an image function as a neural network, which has had considerable success in machine learning image processing applications. We propose a deep learning method employing gradient-based planar alignment with Bundle-Adjusting Radiance Fields (BARF) to combine, de-noise, and remove obstructions from observations of cosmological objects at different resolutions, seeing, and noise levels–tasks not currently possible within a single process in astronomy. We test our algorithm on artificial images of star clusters, demonstrating powerful artifact removal and de-noising.
InceptionSR: Recursive Symbolic Regression for Equation Synthesis
Symbolic regression (SR) algorithms generate equations that fit a dataset. We propose a symbolic regression variant which we call InceptionSR. Motivated by gradient boosting, our algorithm iteratively runs a base SR algorithm, freezes partial equations, and seeds the next round of SR using those frozen equations as features. Our algorithm can also be viewed as a simple form of “library learning”, a technique common in the related field of program synthesis. We evaluate our algorithm on toy problems as well as a real world scientific discovery problem of generating equations to predict planetary instability. Results show that our method consistently improves SR performance, and may also generate more interpretable equations. Keywords: symbolic regression; equation discovery; library learning, recursion; boosting
Multi-Agent System for Cosmological Parameter Analysis
Multi-agent systems (MAS) utilizing multiple Large Language Model agents with Retrieval Augmented Generation and that can execute code locally may become beneficial in cosmological data analysis. Here, we illustrate a first small step towards AI-assisted analyses and a glimpse of the potential of MAS to automate and optimize scientific workflows in Cosmology. The system architecture of our example package, that builds upon the autogen/ag2 framework, can be applied to MAS in any area of quantitative scientific research. The particular task we apply our methods to is the cosmological parameter analysis of the Atacama Cosmology Telescope lensing power spectrum likelihood using Monte Carlo Markov Chains. Our work-in-progress code is open source and available at https://github.com/CMBAgents/cmbagent.
SymbolFit: Automatic Parametric Modeling with Symbolic Regression
We introduce SymbolFit, a framework that automates parametric modeling by using symbolic regression to perform a machine-search for functions that fit the data while simultaneously providing uncertainty estimates in a single run. Traditionally, constructing a parametric model to accurately describe binned data has been a manual and iterative process, requiring an adequate functional form to be determined before the fit can be performed. The main challenge arises when the appropriate functional forms cannot be derived from first principles, especially when there is no underlying true closed-form function for the distribution. In this work, we develop a framework that automates and streamlines the process by utilizing symbolic regression, a machine learning technique that explores a vast space of candidate functions without requiring a predefined functional form because the functional form itself is treated as a trainable parameter, making the process far more efficient and effortless than traditional regression methods. We demonstrate the framework in high-energy physics experiments at the CERN Large Hadron Collider (LHC) using five real proton-proton collision datasets from new physics searches, including background modeling in resonance searches for high-mass dijet, trijet, paired-dijet, diphoton, and dimuon events. We show that our framework can flexibly and efficiently generate a wide range of candidate functions that fit a nontrivial distribution well using a simple fit configuration that varies only by random seed, and that the same fit configuration, which defines a vast function space, can also be applied to distributions of different shapes, whereas achieving a comparable result with traditional methods would have required extensive manual effort.
Accelerating Giant-impact Simulations with Machine Learning
Constraining planet formation models based on the observed exoplanet population requires generating large samples of synthetic planetary systems, which can be computationally prohibitive. A significant bottleneck is simulating the giant impact phase, during which planetary embryos evolve gravitationally and combine to form planets, which may themselves experience later collisions. To accelerate giant impact simulations, we present a machine learning (ML) approach to predicting collisional outcomes in multiplanet systems. Trained on more than 500,000 $N$-body simulations of three-planet systems, we develop an ML model that can accurately predict which two planets will experience a collision, along with the state of the post-collision planets, from a short integration of the system's initial conditions. Our model greatly improves on non-ML baselines that rely on metrics from dynamics theory, which struggle to accurately predict which pair of planets will experience a collision. By combining with a model for predicting long-term stability, we create an ML-based giant impact emulator, which can predict the outcomes of giant impact simulations with reasonable accuracy and a speedup of up to four orders of magnitude. We expect our model to enable analyses that would not otherwise be computationally feasible. As such, we release our training code, along with an easy-to-use API for our collision outcome model and giant impact emulator.
Contextual Counting: A Mechanistic Study of Transformers on a Quantitative Task
Transformers have revolutionized machine learning across diverse domains, yet understanding their behavior remains crucial, particularly in high-stakes applications. This paper introduces the contextual counting task, a novel toy problem aimed at enhancing our understanding of Transformers in quantitative and scientific contexts. This task requires precise localization and computation within datasets, akin to object detection or region-based scientific analysis. We present theoretical and empirical analysis using both causal and non-causal Transformer architectures, investigating the influence of various positional encodings on performance and interpretability. In particular, we find that causal attention is much better suited for the task, and that no positional embeddings lead to the best accuracy, though rotary embeddings are competitive and easier to train. We also show that out of distribution performance is tightly linked to which tokens it uses as a bias term.
Improving astrophysical scaling relations with machine learning
Finding low-scatter relationships in properties of complex systems (eg, stars, supernovae, galaxies) is important to gain physical insights into them and/or to estimate their distances/masses. As the size of simulation/observational datasets grow, finding low-scatter relationships in the data becomes extremely arduous using manual data analysis methods. I will show how machine learning techniques can be used to expeditiously search for such relations in abstract high-dimensional data-spaces. Focusing on clusters of galaxies, I will present new scaling relations between their properties obtained using machine learning tools. Our relations can enable more accurate inference of cosmology and baryonic feedback from upcoming surveys of galaxy clusters such as ACT, SO, eROSITA and CMB-S4.
Mangrove: Infer galaxy properties using dark matter merger trees
Mangrove uses Graph Neural Networks to regress baryonic properties directly from full dark matter merger trees to infer galaxy properties. The package includes code for preprocessing the merger tree, and training the model can be done either as single experiments or as a sweep. Mangrove provides loss functions, learning rate schedulers, models, and a script for doing the training on a GPU.
Learning Galaxy Properties from Merger Trees
Efficiently mapping baryonic properties onto dark matter is a major challenge in astrophysics. Although semi-analytic models (SAMs) and hydrodynamical simulations have made impressive advances in reproducing galaxy observables across cosmologically significant volumes, these methods still require significant computation times, representing a barrier to many applications. Graph Neural Networks (GNNs) have recently proven to be the natural choice for learning physical relations. Among the most inherently graph-like structures found in astrophysics are the dark matter merger trees that encode the evolution of dark matter halos. In this paper we introduce a new, graph-based emulator framework, $\texttt{Mangrove}$, and show that it emulates the galactic stellar mass, cold gas mass and metallicity, instantaneous and time-averaged star formation rate, and black hole mass -- as predicted by a SAM -- with root mean squared error up to two times lower than other methods across a $(75 Mpc/h)^3$ simulation box in 40 seconds, 4 orders of magnitude faster than the SAM. We show that $\texttt{Mangrove}$ allows for quantification of the dependence of galaxy properties on merger history. We compare our results to the current state of the art in the field and show significant improvements for all target properties. $\texttt{Mangrove}$ is publicly available.
Interpretable Deep Learning for Computational Fluid Dynamics
Can deep learning or symbolic regression supplement traditional simulators in fluid dynamics? How well do such models generalize outside of the dataset they learn from, and how well do they preserve statistical properties of the simulated fluid? In this talk, I will present some key observations from our recent research in this area, which aims to answer these questions. I will highlight our new method:''Disentangled Sparsity Networks,''which allow one to interpret the internals of a neural network trained on fluids simulation. Not only does this give us a way of interrogating how the deep learning model is making predictions, but it also allows one to replace the learned model with a symbolic expression and embed that model inside a traditional solver. We show that this technique can improve the applicability of symbolic regression to high-dimensional datasets, such as those in fluid dynamics, without imposing priors …
SIRFING: Sparse Image Reconstruction For INterferometry using GPUs
We present a deconvolution code for radio interferometric imaging based on the compressed sensing algorithms in Garsden et al.(2015). Being computationally intensive, compressed sensing is ripe for parallelization over GPUs. Our compressed sensing implementation generates images using wavelets, and we have ported the underlying wavelet library to CUDA, targeting the spline filter reconstruction part of the algorithm. The speedup achieved is almost an order of magnitude. The code is modular but is also being integrated into the calibration and imaging pipeline in use by the LEDA project at the Long Wavelength Array (LWA) as well as by the Murchinson Widefield Array (MWA).
Compute-Adaptive Surrogate Modeling of Partial Differential Equations
Modeling dynamical systems governed by partial differential equations (PDEs) presents significant challenges for machine learning-based surrogate models. While vision transformers have shown potential in capturing complex spatial dynamics, their reliance on fixed-size patches limits flexibility and scalability. In this work, we introduce two convolutional architectural blocks—Convolutional Kernel Modulator (CKM) and Convolutional Stride Modulator (CSM)—designed for patch processing in autoregressive prediction tasks. These blocks unlock dynamic patching and striding strategies to balance accuracy and computational efficiency during inference. Furthermore, we propose a rollout strategy that adaptively adjusts patching and striding configurations throughout temporally sequential predictions, mitigating patch artifacts and long-term error accumulation while improving the capture of finer-scale structures of physics-based PDEs. We show that the use of these blocks improves predictive accuracy against fixed-patch baselines, while also enabling inference time scaling.
Spectral Shaping for Neural PDE Surrogates
Neural surrogates for PDE solvers suffer from an inability to model the spectrum of solutions adequately, especially in the medium to high frequency bands. This impacts not only correct spectral shapes, but also stability and long-term rollout accuracy. We identify three convergent factors that exacerbate this phenomenon, namely: distribution shift over unrolls, spectral bias of the MSE loss, and spurious high frequency noise, or _spectral junk_, introduced by the use of pointwise nonlinearities. We find that _spectral shaping_, filtering the spectrum of activations after every layer of pointwise nonlinearities, is enough to reduce spectral junk and improve long-term rollout accuracy. We show spectral shaping not only fixes the learned spectrum (down to machine precision in some cases), but also leads to very stable neural surrogates. We validate these findings on a suite of challenging fluid dynamics problems in the field of neural PDE surrogacy, promoting a clear need for more careful attention to surrogate architecture design and adding a new and simple trick to the practitioner toolbox.