PaperlayerView exact revision on arXiv

Source-linked AI summary

The AI Cosmologist I: An Agentic System for Automated Data Analysis

Adam Moss

arXiv:2504.03424v1astro-ph.IMastro-ph.COastro-ph.GAcs.AIphysics.data-an

TL;DR

Large cosmological datasets and the expertise and iterative experimentation required for practical ML create a need for more automated research workflows. AI Cosmologist coordinates agents to generate, implement, evaluate, and refine approaches. Across the demonstrated tasks, later iterations improved over initial implementations, while the system remained limited in conceptual novelty and independent question formulation.

  • Problem

    Practical ML deployment still requires domain-specific expertise and iterative experimentation, while existing AutoML struggles with novel tasks requiring creative problem-solving or refinement.

    Source: §1 · p. 2

  • Method

    AI Cosmologist combines specialized agents for planning, coding, execution, analysis, and synthesis to automate cosmological data-analysis and ML research workflows.

    Source: §1 · p. 2, §3.1 · p. 5

  • Results

    Across the demonstrated tasks, later iterations consistently improved over initial implementations; the system explored 50 Galaxy Zoo variations in approximately 72 hours.

    Source: §5 · p. 16

  • Takeaways & Limitations

    The system demonstrates that AI can assist in or potentially automate significant portions of the cosmological data-analysis and machine-learning research process.

    Source: §6 · p. 18

    The system adapts established patterns more readily than it makes conceptual innovations and requires well-specified problems rather than formulating its own research questions.

    Source: §5.1 · p. 17

Abstract

from arXiv · show

We present the AI Cosmologist, an agentic system designed to automate cosmological/astronomical data analysis and machine learning research workflows. This implements a complete pipeline from idea generation to experimental evaluation and research dissemination, mimicking the scientific process typically performed by human researchers. The system employs specialized agents for planning, coding, execution, analysis, and synthesis that work together to develop novel approaches. Unlike traditional auto machine-learning systems, the AI Cosmologist generates diverse implementation strategies, writes complete code, handles execution errors, analyzes results, and synthesizes new approaches based on experimental outcomes. We demonstrate the AI Cosmologist capabilities across several machine learning tasks, showing how it can successfully explore solution spaces, iterate based on experimental results, and combine successful elements from different approaches. Our results indicate that agentic systems can automate portions of the research process, potentially accelerating scientific discovery. The code and experimental data used in this paper are available on GitHub at https://github.com/adammoss/aicosmologist. Example papers included in the appendix demonstrate the system's capability to autonomously produce complete scientific publications, starting from only the dataset and task description

1 Introduction

Large, high-dimensional astronomy datasets make traditional manual analysis impractical, while existing AutoML systems struggle with novel tasks requiring creative problem-solving. The paper introduces an agentic system for automating cosmological data analysis and machine-learning research workflows.

1 Introduction

  • Upcoming surveys and simulations will produce petabyte-to-exabyte-scale datasets, making traditional manual or semi-manual analysis workflows impractical.The introduction cites the Vera C. Rubin Observatory, Euclid, SKA, DESI, and cosmological simulations as examples.Source: §1 · p. 2
  • Existing AutoML can streamline predefined workflows but struggles with novel tasks requiring creative problem-solving or iterative refinement.The paper identifies domain-specific expertise and iterative experimentation as bottlenecks in practical ML deployment.Source: §1 · p. 2
  • The AI Cosmologist integrates AutoML, LLM-driven code generation, and autonomous reasoning to formulate hypotheses, run experiments, evaluate results, and refine methods.The stated workflow operates without manual intervention.Source: §1 · p. 2

1.1 Contributions

  • The paper's contributions include an autonomous research pipeline and demonstrations on galaxy morphology classification and cosmological parameter inference.The listed contributions also include publication-ready results and visualizations, with example scientific papers in the appendix.Source: §1.1 · p. 3, §1 · p. 3

2 Related Work

Prior work automates parts of machine learning, programming, and scientific research, but the paper emphasizes the remaining need for systems that support creative, iterative scientific workflows. AI Cosmologist is positioned as an effort to automate the full scientific ML lifecycle in cosmology.

2.1 Automated Machine Learning (AutoML)

  • AutoML automates pipeline design, model selection, and tuning, but scientific systems often explore predefined spaces and still require substantial human guidance.Prior approaches include Bayesian optimization, meta-learning, evolutionary optimization, neural architecture search, and meta-learning for rapid task adaptation.Source: §2.1 · p. 3

2.2 AI-Assisted Programming

  • LLM-based programming advances natural-language-to-code generation and may support scientific prototyping, standardized data pipelines, and error detection.The cited systems include Codex, CodeGen, AlphaCode, InCoder, and StarCoder.Source: §2.2 · pp. 3–4

2.3 Agentic Systems and Autonomous Scientists

  • Agentic systems combine planning, tool use, and iterative refinement; AI Cosmologist focuses on automating the full scientific ML lifecycle in cosmology.The paper situates its work alongside domain-specific agents and early astrophysics and cosmology explorations.Source: §2.3 · p. 4

3 An Agentic System for Automated Cosmological Data Analysis

AI Cosmologist coordinates specialized agents through a research cycle that generates, implements, evaluates, and refines machine-learning approaches, followed by a dissemination workflow for scientific communications. Cross-experiment synthesis supports both refinement of successful ideas and exploration of unexamined approaches.

3.1 Agent Architecture

  • Specialized planning, coding, execution, analysis, synthesis, and literature agents coordinate through a directed graph with conditional execution and shared global context.Agents can use external tools, including code execution and visualization environments.Source: §3.1 · p. 5

3.2 Research Phase

  • The research phase begins with dataset information and a task specification, then generates distinct implementation ideas for development and evaluation.The workflow cycles through idea generation, implementation, evaluation, and collaborative synthesis.Source: §3.2 · p. 5, §3.2.1 · pp. 6, 8, §3.2.2 · p. 8

3.2.3 Plan Development

  • For each idea, the system develops and may reflect on a plan covering data handling, model architecture, training, and evaluation, then turns it into end-to-end executable code.Generated implementations include preprocessing, training and evaluation procedures, experiment tracking, checkpoint handling, and result visualizations.Source: §3.2.3 · p. 8, §3.2.4 · p. 8

3.2.5 Execution and Evaluation

  • The execution agent runs each implementation, diagnoses runtime errors, and applies diff-based fixes until success or the retry limit is reached.An unresolved approach is marked unsuccessful after the maximum number of attempts.Source: §3.2.5 · p. 9

3.2.6 Synthesis

  • Synthesis ranks experimental results and analyzes both patterns in top-performing ideas and unexplored regions to generate improved and diverse ideas.The two analyses combine successful approaches with exploration of previously unexamined parts of the solution space.Source: §3.2.6 · pp. 9–10

3.2.7 Collaborative Rounds

  • Collaborative rounds feed synthesized ideas back into planning, sending them through the full development and evaluation pipeline using accumulated experimental evidence.The research workflow diagram also depicts parallel analysis of successful patterns and unexplored approaches.Source: Figure 2 · p. 7, §3.2.7 · p. 10

3.3 Research Dissemination Phase

  • After promising results, dissemination agents plan a scientific narrative, review literature, and draft conventional paper sections with technical details and result visualizations.The current system produces drafts that benefit from human review and refinement before formal submission.Source: §3.3 · pp. 10–11

3.4 Implementation Details

  • Gemini 2.5 Pro supports most agent functions, while o3mini-high handles diff-based code editing because frequent requests can make the rate-limited Gemini model restrictive.The paper reports that a complete end-to-end experimental process typically costs several dollars.Source: §3.4 · p. 11

4 Experimental Results

Across two cosmological machine-learning tasks, AI Cosmologist iteratively developed approaches, with the strongest reported outcomes in Galaxy Zoo morphology regression and Quijote parameter inference. The experiments used 50 implementation attempts per dataset, starting from basic dataset and task information.

4.1 Experimental Setup

  • The evaluation covered Galaxy Zoo 2 morphology regression on 304,122 galaxies and Quijote inference from 2,000 simulations across five cosmological parameters.With only basic dataset and task information, the system generated 20 initial ideas and ran five collaborative rounds of six ideas each, for 50 attempts per dataset.Source: §4.1 · p. 12

4.2 Galaxy Zoo Results

  • For Galaxy Zoo, analysis found pre-trained CNNs—particularly EfficientNet and ConvNeXt—performed consistently well, while overfitting and decision-tree dependencies remained challenges.The system's refinements included hierarchical losses, domain-specific augmentation, self-supervised learning, multi-resolution inputs, and ensembles.Source: §4.2 · p. 13
  • 0.07235 RMSE on the Kaggle public leaderboard was achieved by the final Galaxy Zoo implementation, which also exceeded the original competition winner on the private leaderboard.The final approach combined test-time augmentation with multi-scale resizing; Figure 4 tracks validation RMSE and public Kaggle score across implementation cycles.Source: §4.2 · p. 13, Figure 4 · p. 15

4.3 Quijote Results

  • Quijote experiments revealed that Ωb and h were harder to predict than Ωm and σ8, while standard pooling could lose information and 3D training faced computational constraints.These observations guided exploration of physics-informed architectures, uncertainty quantification, and multi-scale feature extraction.Source: §4.3 · p. 16
  • The physics-augmented 3D CNN achieved state-of-the-art Quijote performance, improving constraints on Ωb and h while maintaining excellent accuracy for Ωm and σ8.It combined deep feature extraction with explicitly computed power-spectrum and density-probability-distribution features.Source: §4.3 · p. 16

5 Discussion

The system rapidly explores implementations, improves them through empirical iteration, and can outperform conventional approaches on some tasks. Its autonomy and scope remain limited, motivating future work on theory integration, transfer, and evaluation.

5 Discussion

  • The system matched human-programmer baselines and sometimes exceeded conventional approaches, including the original Kaggle winner on Galaxy Zoo.It explored 50 Galaxy Zoo implementations in approximately 72 hours.Source: §5 · p. 16
  • Later iterations consistently improved on initial implementations, indicating transfer of knowledge across experimental runs.The authors interpret this as accumulating insights and refining methods from empirical evidence.Source: §5 · pp. 16–17

5.1 Limitations

  • The system recombines and refines established approaches, but struggles with conceptual innovation and cannot formulate its own research questions.It requires well-specified problems.Source: §5.1 · p. 17
  • Generated-code efficiency may limit practical problem scale, and the system lacks deep theoretical understanding to guide principled approaches.It learns primarily from empirical results rather than physical theory.Source: §5.1 · p. 17
  • The demonstrated datasets are particularly well-suited to machine learning, so broader evaluation requires more challenging and diverse cosmological datasets.The authors call for more exhaustive studies across cosmological research problems.Source: §5.1 · p. 17

5.2 Future Directions

  • Future directions include theory knowledge bases, improved meta-learning, human collaboration, distributed exploration, multimodal data, and better evaluation of scientific significance.Proposed evaluation should address robustness, generalizability, and physical plausibility.Source: §5.2 · p. 17

6 Conclusion

The AI Cosmologist automates a research workflow from idea generation through experimental evaluation, demonstrating AI participation in methodological exploration and iteration. Its speed and breadth of experimentation could accelerate scientific discovery, although it has not yet achieved human-like creative leaps.

6 Conclusion

  • The AI Cosmologist implements a complete workflow from idea generation to experimental evaluation, automating significant portions of cosmological data analysis and ML research.The paper describes it as a first step toward automating this research process.Source: §6 · p. 18
  • Methodical exploration, implementation, and iteration can be performed effectively by AI systems, though the system cannot yet make groundbreaking human-like creative leaps.The authors suggest AI may increasingly participate in its own development.Source: §6 · p. 18
  • Running dozens of experiments in parallel and completing in days what might take researchers weeks or months could accelerate scientific discovery.The authors suggest scaled solution-space exploration may uncover approaches missed under traditional research timelines.Source: §6 · p. 18

Declarations

The appendix presents a Galaxy Zoo 2 classification study using ConvNeXt, multi-resolution training, and test-time augmentation to predict detailed morphological probabilities. It reports low RMSE against human consensus and a top competition result, while the main paper makes code and data available.

Declarations

  • The paper reports no competing interests and makes its system code and experimental data available on GitHub.It also identifies the Galaxy Zoo 2 and Quijote data sources.Source: Declarations · p. 18

ABSTRACT

  • On the private leaderboard, the model achieved an RMSE of 0.07243, placed first in the Galaxy Zoo Kaggle competition, and test-time augmentation improved performance by 4%.The appendix abstract describes the approach as closely approximating human-level classification accuracy.Source: ABSTRACT · p. 20

1 INTRODUCTION

  • Automated classification must predict detailed, probabilistic features such as bars, spiral arms, and bulge prominence, beyond broad galaxy types.Manual classification at the scale of modern surveys is increasingly challenging.Source: §1 · pp. 20–21
  • The study aims to predict the full set of 37 Galaxy Zoo 2 morphological probabilities, enabling efficient processing of current and future large surveys.Source: §1 · p. 21
  • The approach uses ConvNeXt with multi-resolution training for scale-invariant features and test-time augmentation to combine transformed-image predictions for rotational and flip invariance.The authors identify varying apparent galaxy sizes and orientation effects as classification challenges.Source: §1 · p. 21
  • The model achieved an RMSE of 0.07235 against human consensus classifications using over 60,000 training images and 37 Galaxy Zoo 2 probabilities.The introduction describes the result as a significant improvement over previous benchmarks and near the theoretical limit of agreement between human classifier groups.Source: §1 · p. 21

2 RELATED WORK

The review traces galaxy morphology classification from visual and citizen-science catalogs to automated methods, then identifies unresolved challenges in detailed features, scale, and orientation. The paper presents ConvNeXt, multi-resolution training, and test-time augmentation as responses to those challenges.

2.1 Galaxy Morphology Classification: Evolution and Significance

  • Galaxy morphology supports studies of galaxy formation and evolution, while expert visual inspection has become difficult to scale with growing survey data.The historical progression from the Hubble sequence to more nuanced visual classifications is described.Source: §2.1 · p. 21

2.2 Citizen Science and the Galaxy Zoo Project

  • Galaxy Zoo and GZ2 used crowdsourced classifications, with GZ2's 11-question decision tree covering 37 features and producing probabilistic classifications.The project enabled studies of morphology-environment relationships, rare objects, and galaxy quenching pathways.Source: §2.2 · p. 21

2.3.1 Traditional Machine Learning Methods

  • Earlier automated methods included neural networks for basic morphological types and CAS structural parameters, but did not capture the full detail of GZ2 classifications.Banerji et al. (2010) achieved over 90% accuracy for primary types, while Conselice (2006) introduced concentration, asymmetry, and clumpiness parameters.Source: §2.3.1 · p. 21

2.3.2 Environmental and Physical Correlations

  • Studies connect morphology with environment, color, stellar mass, and star-formation pathways, including distinct evolutionary routes through the green valley.Bamford et al. (2009) found most of the morphology-density relation was driven by environmental morphological fractions at fixed stellar mass; Skibba et al. (2009) attributed much to color-density relations.Source: §2.3.2 · p. 22

2.3.3 Specific Morphological Features

  • Research on bars, dust, and mergers shows that specific features and observational effects complicate morphological interpretation and automated classification.Masters et al. (2011) found over half of red, bulge-dominated disk galaxies possess a bar; Masters et al. (2010a) reported inclination-related reddening, and Darg et al. (2010a) found an approximately twofold higher spiral-to-elliptical ratio in mergers.Source: §2.3.3 · p. 22

2.4 Research Gap and Our Contribution

  • Prior approaches leave detailed features, scale invariance, and orientation sensitivity unresolved; this work combines ConvNeXt, multi-resolution training, and test-time augmentation to address them.The stated goal is prediction across all 37 Galaxy Zoo 2 morphological classes.Source: §2.4 · p. 22

3 DATASET

The dataset section describes Galaxy Zoo 2 images and volunteer classifications as probabilistic labels across a branching 37-class morphology scheme. It also explains how those labels are weighted and validated, while noting known dataset biases.

3.1 Data Sources

  • GZ2 supplies detailed morphological classifications for 304,122 SDSS galaxies, represented here as color JPEG images.The images combine SDSS g, r, and i bands; the JPEG artifacts are consistent with those in the original classification process.Source: §3.1 · p. 22

3.2 Classification Methodology

  • An 11-question decision tree branches from broad galaxy categories to 37 morphological classes, using responses from typically 40–50 volunteers per galaxy.Volunteer responses form probability distributions, and later questions depend on earlier answers.Source: §3.2 · pp. 22–23

3.3 Data Processing

  • Later-question probabilities are weighted by the probability of the preceding classification path, so downstream labels depend on earlier branches.Known dataset biases include harder-to-detect fine features in distant or small galaxies and dust-related changes to spiral-galaxy appearance (Bamford et al. 2009; Masters et al. 2010a).Source: §3.3 · p. 23

3.4 Dataset Characteristics

  • The dataset contains 61,578 training and 79,975 testing galaxies, with roughly 35% smooth, 55% featured/disk, and 10% stars or artifacts.Featured galaxies include varied bars, edge-on systems, spiral arms, bulges, and merger signatures.Source: §3.4 · p. 23

3.5 Data Validation

  • Prior expert agreement and over 90% accuracy for primary morphology frame the harder task of predicting probabilities across all 37 classes.Galaxy Zoo classifications also surfaced unusual objects such as Green Peas and Hanny’s Voorwerp, supporting their use as human ground truth for automated methods.Source: §3.5 · p. 23

4 METHODS

Galaxy morphology is treated as a 37-output regression task, using ConvNeXt, multi-resolution training, and test-time averaging across transformed views. Full test-time augmentation lowers validation RMSE from 0.07571 to 0.07273.

4.1 Morphological Classification Framework

  • Galaxy morphology is framed as predicting 37 class probabilities from each image, with the approach combining convolutional feature extraction, multi-resolution training, and symmetry-aware test-time augmentation.The classification target reflects the distribution of human classifications, including consensus and uncertainty.Source: §4.1 · p. 23

4.2 Galaxy Zoo 2 Dataset

  • GZ2 volunteer classifications of SDSS images follow an 11-question decision tree, producing 37 probabilities weighted by the probability of reaching each question.Figure 1 shows representative training images from different morphological categories.Source: §4.2 · pp. 23–24, Figure 1 · p. 24

4.3 Deep Convolutional Architecture

  • The ConvNeXt backbone extracts image features that a prediction head maps to the 37-class probability vector, while hierarchical blocks capture both large-scale structure and fine details.The architecture builds on ConvNeXt (Liu et al. 2022).Source: §4.3 · p. 24

4.4 Multi-Resolution Training

  • Training images are randomly resized from 256 to 456 pixels and cropped to 224×224, exposing the model to varying apparent scales.The authors state that this strategy encourages scale-invariant features.Source: §4.4 · p. 24

4.5 Loss Function and Optimization

  • Mean Squared Error fits predicted probabilities to the GZ2 distributions, while AdamW optimization uses a learning rate of 10^-4 and weight decay of 10^-5.The learning rate drops by a factor of 0.1 when validation RMSE plateaus, and early stopping is based on validation performance.Source: §4.5 · pp. 24–25

4.6 Test-Time Augmentation

  • At inference, the model averages predictions across transformed galaxy views to account for scale, rotation, and reflection invariances.The tested transformations combine resolutions of 256, 384, and 456 pixels; four rotations; and three flip settings.Source: §4.6 · p. 25

4.7 Evaluation Metrics and Error Analysis

  • RMSE measures error across predicted and ground-truth probability vectors, while class-specific RMSE examines how accurately individual morphological features are predicted.The analysis gives particular attention to smooth and features/disk classes.Source: §4.7 · p. 25

4.8 Implementation Details

  • 0.07235 RMSE on the public test set accompanied 0.07273 validation RMSE with TTA, versus 0.07571 without TTA (4% error reduction).Training used 61,578 images; ImageNet-statistic normalization and mixed precision were used, and the lowest-validation-RMSE model was selected.Source: §4.8 · p. 25, §4.7 · p. 25

5 RESULTS

The galaxy morphology model combines multi-resolution training with test-time augmentation and performs strongly across 37 classes, with predictions closely matching human classifications. Results include a leaderboard-leading reported score, high primary-class accuracy, and improved validation RMSE with augmentation.

5.2 Overall Model Performance

  • Validation RMSE improved from 0.075710 with standard evaluation to 0.072727 with test-time augmentation; the public leaderboard RMSE was 0.07235.The authors report that validation and leaderboard performance were closely matched.Source: §5.2 · p. 26

5.3 Performance Across Morphological Categories

  • Across 37 morphological classes, RMSE ranged from 0.0198 to 0.1291, and rare or subtle features generally had higher errors.Reported strong examples include Class8.7 at RMSE 0.0198 and Class11.5 at RMSE 0.0268.Source: §5.3 · p. 26

5.4 Prediction Accuracy Analysis

  • Predicted probabilities correlate strongly with human-consensus probabilities, while bar-feature predictions vary more at intermediate probabilities of 0.3–0.7.The passage attributes this uncertainty to human classifiers also disagreeing about bars in these cases.Source: §5.4 · p. 26, Figure 3 · p. 27

5.5 Error Distribution

  • 78% of predictions had absolute error below 0.1, and 94% were below 0.2.The error distribution is shown as a histogram with a cumulative proportion overlay.Source: §5.5 · pp. 26–27, Figure 4 · p. 27

5.6 Impact of Test-Time Augmentation

  • Averaging predictions across 36 augmented image views reduced validation RMSE by 4.0%.The views address orientation and scale variation; the TTA-enhanced validation score closely matched the public leaderboard score.Source: Figure 5 · p. 27, §5.6 · p. 27

5.7 Example Predictions

  • The model reproduced hierarchical human-classification patterns, including downstream spiral-arm and tightness probabilities after a confident spiral classification.The examples compare predictions with human classifications on representative validation galaxies.Source: §5.7 · p. 27

5.8 Benchmark Comparison

  • Primary-class accuracy was 95.7%, above Banerji et al. (2010)'s 90% benchmark; average correlation across all 37 classes was 0.8103.The reported primary-class correlations were 0.9118, 0.9140, and 0.7179 for Classes 1.1–1.3.Source: §5.8 · p. 27
  • The public test RMSE of 0.07235 (private: 0.07243) would have ranked first, ahead of the next-best private score of 0.07491.The authors report strong performance across diverse features, including edge-on orientation, spiral patterns, and bar presence.Source: §5.8 · p. 27

6 CONCLUSIONS

The conclusions report that multi-resolution training and test-time augmentation produce accurate Galaxy Zoo 2 morphology predictions across 37 features, with a public leaderboard RMSE of 0.07235. The same supplied material also describes an example Quijote study that combines 3D density fields and physics-based features for cosmological parameter inference.

6 CONCLUSIONS

  • Multi-resolution training and comprehensive test-time augmentation achieved a public leaderboard RMSE of 0.07235 for Galaxy Zoo 2 morphology classification.The authors report accurate probabilities across the decision tree and a 4% error reduction from test-time augmentation.Source: §6 · pp. 27–28
  • Automated morphology classification can support large-scale statistical studies of links between galaxy morphology, color, and environment.The conclusions frame automation as useful for image datasets impractical to classify manually, even with crowdsourcing.Source: §6 · p. 28
  • The model depends on Galaxy Zoo 2 training-data quality and biases; rare features may be underrepresented, and optimal TTA requires 36 forward passes per image.The authors also note that the model does not provide uncertainty estimates for its probability predictions.Source: §6 · p. 28

ABSTRACT

  • The Appendix B Quijote example combines a 3D ResNet with power-spectrum and density-PDF features to estimate five cosmological parameters.Its reported R² values are 0.939 for Ωm, 0.992 for σ8, 0.468 for Ωb, 0.480 for h, and 0.587 for ns.Source: ABSTRACT · p. 31

1 INTRODUCTION

  • The introduction motivates extracting cosmological information from large-scale structure and describes the two-pathway model evaluated on Quijote density fields.The reported results include R² = 0.939 for Ωm and R² = 0.992 for σ8, with improved constraints also reported for Ωb, h, and ns.Source: §1 · pp. 31–32

2 RELATED WORK

Prior work applies deep learning to cosmological inference across weak-lensing maps, 3D density fields, and other probes, while hybrid methods combine learned features with physical statistics. The supplied material identifies limited systematic combination across the full parameter space as a gap, then outlines the Quijote data and processing used for the proposed approach.

2.1.1 Weak Lensing Applications

  • Weak-lensing studies report tighter constraints than traditional statistics and show that CNNs can use non-Gaussian information; later work examined peak gradients and real observational data.Reported results include up to 50% tighter constraints in Fluri et al. (2018), about five times tighter constraints in Gupta et al. (2018), and a 16% S8 improvement in Fluri et al. (2022).Source: §2.1.1 · p. 32

2.1.2 From 2D to 3D: Full Field Analysis

  • Research extended CNN inference from 2D projections to 3D dark-matter fields and other probes, including 21-cm maps and combinations of weak lensing with galaxy clustering.Pan et al. (2020) reported uncertainties δΩm = 0.0015 and δσ8 = 0.0029 for 3D dark-matter distributions.Source: §2.1.2 · pp. 32–33

2.2 Hybrid and Physics-Informed Approaches

  • Hybrid inference combines neural networks with physical statistics; prior examples report improved constraints over either component alone or while accounting for baryonic effects.The cited work includes CNNs combined with power-spectrum-based networks and weak-lensing models trained to account for baryonic effects.Source: §2.2 · p. 33

2.3 Model Interpretation and Feature Analysis

  • Interpretability studies have linked network predictions to extreme convergence values and investigated which initial-density features contain information about halo mass.Other cited analyses identify factors from the linear matter power spectrum associated with halo-mass-function behavior.Source: §2.3 · p. 33

2.4 Technical Innovations and Optimizations; 2.5 Applications in Related Areas

  • Technical work spans denoising, architecture search, simulation resources, and related astrophysical applications such as cluster-mass estimation.CosNAS is reported to decrease estimation errors by 85.5%; denoising work reports 30–40% better constraints, and cluster-mass studies report lower scatter or avoidance of observational biases.Source: §2.4 · p. 33, §2.5 · p. 33

2.6 Research Gap and Our Approach

  • The identified gap is that few studies systematically combine CNNs with power-spectrum and density-PDF statistics to constrain the full ΛCDM parameter set from 3D fields.The proposed framework targets all five parameters, emphasizes Ωb and h, and incorporates attention mechanisms.Source: §2.6 · pp. 33–34

3.1 Simulation Data; 3.2 Selection Criteria; 3.3 Data Processing; 3.4 Dataset Characteristics

  • The study uses Quijote’s 2,000 Latin Hypercube simulations, extracts z = 0 CDM density fields on 64³ grids, and splits them 70%/15%/15% into training, validation, and test sets.The five varied parameters are Ωm, Ωb, h, ns, and σ8, with ranges Ωm [0.1, 0.5], Ωb [0.03, 0.07], h [0.5, 0.9], ns [0.8, 1.2], and σ8 [0.6, 1.0].Source: §3.1 · p. 34, §3.2 · p. 34, §3.3 · p. 34

3.3 Data Processing; 3.5 Data Validation

  • Processing applies log transformation, normalization, and rotational/reflection augmentation; the feature set includes power-spectrum and density-PDF statistics.The paper reports that 64³ resolution captures relevant large-scale features but limits sensitivity to small-scale nonlinear structure.Source: §3.3 · p. 34, §3.5 · p. 35

4 METHODS

The method combines an attentive 3D CNN with physics-derived statistics to infer cosmological parameters from density fields. It transforms the fields, extracts complementary power-spectrum and density-distribution features, and fuses them with learned features for prediction.

4.1 Theoretical Framework

  • The hybrid framework pairs 3D CNN feature learning with physics-motivated statistics to address parameters that standard CNNs struggle to constrain.The statistics target information that may be missed by summary statistics alone, including non-Gaussian density-field features.Source: §4.1 · p. 35, §4.7 · p. 38

4.2.2 Density Field Transformation

  • Log transformation compresses the density field’s dynamic range and yields a more symmetric distribution; standardization then supports training stability and convergence.The transformed field is standardized to zero mean and unit variance using training-set statistics.Source: §4.2.2 · p. 35

4.3.1 Power Spectrum Features

  • Power-spectrum features encode BAO-related scale dependence, spectral shape, and amplitude, targeting Ωm, Ωb, h, ns, and σ8 information.The spectra are estimated in logarithmically spaced bins from k=0.01 h/Mpc to k=1.0 h/Mpc; Figure 2 illustrates parameter-dependent amplitude and shape variation.Source: §4.3.1 · p. 36, Figure 2 · p. 36

4.3.2 Density Probability Distribution Function

  • Density PDFs complement the power spectrum by capturing non-Gaussian structure through distribution moments, percentiles, and void-volume fractions.Figure 3 shows parameter-dependent PDF shape variation, particularly in the width and tail associated with σ8.Source: Figure 3 · p. 36, §4.3.2 · p. 37

4.4.1 Physics-Augmented Attentive 3D ResNet

  • The architecture fuses pooled 3D ResNet features with normalized physics features, while Squeeze-and-Excitation blocks recalibrate channel responses.The regression head predicts five cosmological parameters from the combined feature vector.Source: §4.4.1 · p. 37, §4.4.1.2 · p. 37, §4.4.1.3 · p. 37

4.5.1 Objective Function

  • Training minimizes normalized MSE, uses AdamW with plateau-based learning-rate reduction and early stopping, and augments fields with 90-degree rotations and reflections.The data split is 70% training, 15% validation, and 15% test; augmentation transformations preserve relevant statistical properties.Source: §4.5.1 · p. 38, §4.5.2 · p. 38, §4.5.3 · p. 38

4.6 Performance Evaluation

  • Test evaluation uses MSE, MAE, R2, and range-normalized errors per parameter, with joint constraints used to inspect parameter degeneracies.The joint-constraint analysis particularly considers Ωm and σ8.Source: §4.6 · p. 38

4.7 Uncertainty Estimation

  • Prediction uncertainty is estimated from empirical test-error distributions, with percentile-based 68% intervals accommodating potentially non-Gaussian errors.Kernel density estimation converts discrete predictions into continuous distributions for visualization and comparison with other cosmological probes.Source: §4.7 · p. 38

5 RESULTS

The physics-augmented attentive 3D ResNet performs strongly on some cosmological parameters and improves on CNN-only inference for more challenging ones. Its results also show parameter-dependent accuracy and remain bounded by idealized simulation data and resolution.

5.1 Overall Model Performance

  • The average R2 across the five parameters is 0.693, with the reported test-set metrics summarized in Table 1.The training and validation losses stabilize after approximately 50 epochs without significant overfitting.Source: §5.1 · p. 38, Table 1 · p. 39, Figure 4 · p. 39

5.2 Parameter-Specific Performance

  • R2 is 0.939 for Ωm and 0.992 for σ8, compared with 0.587 for ns, 0.480 for h, and 0.468 for Ωb.The moderate results for Ωb and h are attributed to their subtler effects on matter distribution, which are harder to extract from finite-resolution fields.Source: §5.2 · p. 39, Figure 5 · p. 39
  • Test errors are approximately Gaussian and centered near zero; their distributions are narrowest for σ8 and Ωm and broader for Ωb, h, and ns.The broader distributions reflect greater difficulty constraining those parameters from the density field alone.Source: §5.2 · p. 39, Figure 6 · p. 40

5.3 Comparison with Other Methods

  • On Ωm, the hybrid reaches R2 = 0.939 versus CNN-only R2 = 0.78 and power-spectrum R2 = 0.91; on σ8 and ns, it trails power-spectrum methods.Across all five parameters, it significantly outperforms the CNN-only method, particularly for Ωb, h, and ns.Source: §5.3 · pp. 39–40, Table 2 · p. 40

6 CONCLUSIONS

  • The authors interpret the hybrid’s gains as evidence of complementary information in CNN-learned spatial patterns and explicit scale-dependent features.They suggest that physics-motivated features targeting effects such as BAO likely contribute to improved Ωb and h constraints.Source: §6 · p. 40
  • The study uses gravity-only Quijote simulations at 643-grid resolution; real-data application would require addressing survey geometry, masks, and instrumental systematics.The authors also note that the resolution may miss fine-scale information relevant to ns and Ωb.Source: §6 · p. 40

Related papers

Loading 2504.03424v1…