Multi-Agent System for Cosmological Parameter Analysis

TL;DR

The paper addresses the challenge of automating cosmological data analysis pipelines by introducing cmbagent, a multi-agent system (MAS) that uses retrieval-augmented generation (RAG) and local code execution to orchestrate tasks. It demonstrates the approach on ACT DR6 CMB lensing parameter inference, with agents handling information retrieval, code development, planning, and execution within a controlled, largely human-in-the-loop workflow. Key contributions include a detailed MAS architecture, a deterministic workflow to minimize stochasticity in LLM outputs, and a successful reproduction of ACT DR6 constraints with minimal manual coding. The work suggests MAS can accelerate quantitative cosmology and related fields, while highlighting practical constraints such as cost, robustness of RAG for complex literature, and the need for benchmarking and potential future automation enhancements.

Abstract

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.

Figures (4)

• Figure 1: Overview of system hierarchy and transitions. Manager agents are shown in salmon, coder agents are shown in blue, experiment RAG agents in yellow, community software RAG agents in light green, and research software RAG agents in magenta. Arrows indicate current allowed transitions in https://github.com/CMBAgents/cmbagent.
• Figure 2: Reproducing of pipelines for the ACT DR6 CMB lensing cosmological parameter constraints. Contours obtained by our system, cmbagent are in red and the original chains downloaded from ACT repository are in blue. These results are obtained in the examples presented in the online documentation.
• Figure 3: Number-count likelihood for a Simons-Observatory-like simulated catalogue as a function of the mass bias, computed by cosmocnc via cmbagent. The true input value of the mass bias is shown as the dashed line; it is consistent with the likelihood constraint obtained by cmbagent. Our GitHub repository contains the cosmocnc agent instructions (see its yaml file).
• Figure 4: CMB power spectra for ten values of the cosmological parameter $f_\mathrm{EDE}$ computed by classy_sz via cmbagent. The notebook provided as part of our online documentation contains the full output of the session where Figure \ref{['fig:newcode']} was made. Our GitHub repository also contains the classy_sz agent instructions (see yaml file).