#  cosmosage: A Natural-Language Assistant for Cosmologists

Tijmen de Haan

###### Abstract

cosmosage is a natural-language assistant intended for a wide audience, from
laypersons interested in cosmology to students, teachers, and professional
cosmologists. cosmosage provides a novel way to access knowledge and reason
about cosmology. Leveraging the power of advanced large language models
(LLMs), cosmosage has learned from a vast corpus of open-access source texts,
including textbooks and papers. cosmosage is found to be state-of-the-art on
the narrow task of answering questions about cosmology, outperforming all
general-purpose models. The model parameters and code are publicly available.

##  1 Introduction

Large language models (LLMs) are a recent and exciting development in
artificial intelligence that have shown remarkable capabilities in
understanding and generating human-like text. Powered by a decoder-only
version of the transformer architecture [1, 2], OpenAI’s GPT-3 caught the
public’s imagination with the launch of ChatGPT in November 2022. By training
on a significant fraction of the text data available on the internet, LLMs
have learned to understand and produce language with human-like fluency and
coherence. Surprisingly, LLMs appear to have the capability to reason about
never-before-seen problems [3]. This opens up new possibilities for natural
language interfaces that can assist humans in a wide range of tasks, from
answering questions to aiding in idea generation, coding, and writing tasks.

In this work, I show that a relatively small model that is narrowly tailored
to a particular application outperforms state-of-the-art models with general
capabilities. The model is named cosmosage, an LLM that is fine-tuned from
publicly available base models such as Mistral-7B-v0.1 [4], a relatively
modest-sized LLM base model with 7.24 billion learned parameters. After
continued pretraining in the form of text completion on the data corpus, I
performed assistant-tuning to create a natural-language assistant specifically
tailored for answering questions about the field of cosmology. By training on
a large body of open-access cosmology textbooks, papers, and other resources,
cosmosage has acquired a deep understanding of cosmological concepts,
theories, and research.

##  2 Development of cosmosage

The development of cosmosage involved several steps, summarized in Figure 1.

Figure 1: Overview of the training procedure for cosmosage. The initial
continued pretraining adds knowledge to the base model. The subsequent tuning
on question-answer pairs creates a useful assistant model.

###  2.1 Data Collection and Preprocessing

The source documents for cosmosage consist of textbooks, lecture notes,
papers, and other open-access datasets. The textbooks were found manually by
searching the internet. The papers were selected using a mixture of search
queries and preference of the cosmic microwave background community as
determined by upvotes on my website that hosts a modified version of arxiv-
sanity-lite [5]. All papers were downloaded with the arXiv API. Finally, all
non-PDF documents, such as textbooks in non-PDF eBook formats were converted
to PDF using command line scripts.

The PDFs were then converted to human-readable lightweight markdown format
using Nougat OCR [6]. This markdown text required some cleaning by removing
repeated words and repeated sequences via a recursive algorithm. Furthermore,
malformatted outputs were spotted by eye and for each of these occurrences an
appropriate regular expression was created for search-and-delete. This process
was continued until no further malformatted markdown was spotted, requiring a
total of 10 regular expressions.

###  2.2 Synthetic Question-Answer Pair Generation

The second stage of training – assistant tuning – requires triplets of a
system prompt, user prompt, and assistant completion. The assistant completion
is the target over which the loss of each token is minimized given the
previous tokens. Of paramount importance are quality, quantity, and diversity.
The following procedure was employed.

First, the input data was separated into overlapping segments. The segment
length was set to 1800 characters, with 600 characters of overlap. The 8-bit
GPTQ version of the ``bagel-34b-v0.2 [7] model was used with the following
prompt. The system prompt was given as

> Assume the role of a cosmology lecturer with a deep understanding of all
> areas of modern cosmology, astrophysics. You must precisely follow the
> instructions given by the user.

and the user prompt:

> As a cosmology expert, your task is to create precise and self-contained
> question-answer pairs from a specified PASSAGE of {shorthand_title}. The
> paper is titled {title} and has the arXiv ID {arxiv_id}. Here is a short
> summary: {summary} Ensure that each question incorporates all necessary
> context, allowing it to be fully understood on its own. Answers should be
> clear, specific, and provide comprehensive information based on the PASSAGE.
> The goal is for each question and answer pair to be understandable
> independently, ensuring they are complete and contextually clear without
> external references.
>
> PASSAGE: {input text segment}

In order to encourage variety, an additional instruction was inserted into the
user prompt prior to the “PASSAGE”, randomly selected from 20 manually crafted
additional instructions. These additional instructions include phrases such as
“Consider whether a mathematical representation would be helpful in
summarizing the PASSAGE.” and “Create questions based on hypothetical
scenarios or ’what-if’ questions inspired by the PASSAGE.”

To enhance the quality of the question-answer (QA) pairs, we implemented a
refinement procedure. This process was designed to mimic an expert teacher
evaluating and improving a student’s work, with the ultimate goal of creating
high-quality synthetic QA pairs. Given the well-documented tendency of Large
Language Models (LLMs) to overestimate the quality of their own outputs [8],
we employed the Mixtral-8x7B-v0.1-instruct model [9] as an independent
evaluator.

The refinement procedure was structured as follows:

  1. 1.

Each QA pair was presented to the Mixtral model along with the original
context used to generate the pair.

  2. 2.

The model was instructed to assign a numerical grade (0-100%) to the answer
based on its accuracy, completeness, and relevance to the question.

  3. 3.

If the original answer received a grade of 90% or higher, it was retained in
the dataset.

  4. 4.

If the original answer scored below 90%, the Mixtral model was prompted to
generate an improved “solution key” answer.

  5. 5.

This solution key was then evaluated using the same grading criteria.

  6. 6.

If the solution key received a grade of 90% or higher, it replaced the
original answer in the QA pair.

  7. 7.

If both the original answer and the solution key scored below 90%, the entire
QA pair was discarded.

This iterative refinement process resulted in a curated dataset of over 1.3
million high-quality QA pairs. The combination of the initial generation and
refinement used approximately 500 A100 GPU-hours.

###  2.3 Model Training

The base model for cosmosage v2 is Mistral-7B-v0.1 [4]. It was selected
because it is relatively small, has a permissive Apache 2.0 license, and
scores well on various benchmarks, implying good reasoning abilities for its
size. The initial step in training was to continue pretraining this base model
on our corpus of cosmology texts. While early versions of cosmosage used a
custom pytorch training loop, the axolotl [10] training framework was found to
be preferable given its large number of implemented algorithms and ease of
switching between these. The AdamW optimizer [11, 12] was used with the
deepspeed_zero_stage_2 configuration provided by the DeepSpeed library [13].
The learning rate was set to $1\times 10^{-5}$ with a micro batch size of 1,
over 8 GPUs. The training was performed for three epochs, taking approximately
180 A100 GPU hours.

Following the continued pretraining phase, we implemented a fine-tuning
process to enhance cosmosage’s capabilities in natural language interaction
and cosmology-specific question answering. This phase used supervised learning
on a diverse dataset of question-answer pairs and dialogue examples, using the
checkpoint obtained from the continued pretraining as the starting point.

The fine-tuning dataset comprised the refined cosmology-specific QA pairs
generated in Section 2.2, as well as procedurally generated QA pairs about
metadata, and QA pairs based on synthetically generated summaries of the
cosmology papers. To enhance the model’s ability to engage in multi-turn
conversations and follow general instructions, we also incorporated the
OpenHermes 2.5 dataset [14].

The fine-tuning process employed the same axolotl framework and DeepSpeed
configuration as the continued pretraining phase, ensuring consistency in the
training pipeline. However, we adjusted the hyperparameters to suit the fine-
tuning task by decreasing the initial learning rate to $2\times 10^{-6}$,
adding a linear decay with warmup over 100 steps. The fine-tuning process was
performed over two epochs, using approximately 300 A100 GPU hours.

##  3 Evaluation

The capabilities of cosmosage were primarily assessed through practical use.
Figures 2, 3, and 4 illustrate instances where cosmosage demonstrates superior
performance compared to state-of-the-art generalist models when answering
cosmology-related questions. However, it is important to note that as a
relatively small language model, cosmosage’s reasoning capabilities are not as
strong as those of very large models such as Claude or ChatGPT. For questions
that demand extended logical reasoning or precise adherence to complex
instructions, larger models remain preferable.

Question: How many parameters are in the wCDM model? cosmosage answer: The
wCDM model is a cosmological model that includes seven parameters:
$\Omega_{\mathrm{b}}h^{2}$, $\Omega_{\mathrm{c}}h^{2}$,
$100\theta_{\mathrm{MC}}$, $\tau_{\mathrm{reio}}$, $n_{\mathrm{s}}$,
$\ln(10^{10}A_{\mathrm{s}})$, and $w$. These parameters describe the
composition and evolution of the universe. Let’s break down each parameter: 1.
$\Omega_{\mathrm{b}}h^{2}$: This parameter represents the density of baryonic
matter (normal matter) in the universe, multiplied by the square of the Hubble
constant ($H_{\mathrm{0}}$). It is expressed as a fraction of the critical
density. 2. $\Omega_{\mathrm{c}}h^{2}$: This parameter represents the density
of cold dark matter in the universe, multiplied by the square of the Hubble
constant ($H_{\mathrm{0}}$). It is also expressed as a fraction of the
critical density. 3. $100\theta_{\mathrm{MC}}$: This parameter is related to
the angular size of the sound horizon at recombination. The sound horizon is
the distance that sound waves could travel before the universe became
transparent to radiation. This value is used to determine the geometry of the
universe. 4. $\tau_{\mathrm{reio}}$: This parameter represents the optical
depth due to reionization, which is the process where neutral hydrogen atoms
in the early universe were ionized by photons from hot stars. It quantifies
the amount of scattering of photons during this period. 5. $n_{\mathrm{s}}$:
This parameter represents the spectral index of primordial scalar
perturbations. It describes the distribution of density fluctuations in the
early universe, which eventually led to the formation of galaxies and other
large-scale structures. 6. $\ln(10^{10}A_{\mathrm{s}})$: This parameter
represents the amplitude of the initial power spectrum of density
fluctuations. It is a measure of the strength of these fluctuations, which
also contributed to the formation of galaxies and other large-scale
structures. 7. $w$: This parameter represents the equation of state for dark
energy. Dark energy is a mysterious component that drives the accelerated
expansion of the universe. By allowing $w$ to vary, the wCDM model can account
for different models of dark energy, such as quintessence or phantom energy.
In summary, the wCDM model has seven parameters that describe the composition,
geometry, and evolution of the universe, including the properties of matter,
dark energy, and the initial conditions of the Big Bang. Claude answer: The
wCDM model, which is an extension of the standard $\Lambda$CDM (Lambda Cold
Dark Matter) model, typically has 6 main parameters: 1. $\Omega_{\mathrm{m}}$
($\Omega_{\mathrm{m}}$): The density parameter for matter (including both
baryonic and dark matter). 2. $\Omega_{\mathrm{b}}$ ($\Omega_{\mathrm{b}}$):
The density parameter for baryonic matter. 3. $h$: The Hubble parameter
($H_{\mathrm{0}}=100h$ km/s/Mpc). 4. $n_{\mathrm{s}}$: The spectral index of
primordial density fluctuations. 5. $\sigma_{8}$ ($\sigma_{8}$): The amplitude
of matter fluctuations on $8\,\mathrm{Mpc}/h$ scales. 6. $w$: The dark energy
equation of state parameter. The key difference between $\Lambda$CDM and wCDM
is the inclusion of $w$ as a free parameter. In $\Lambda$CDM, $w$ is fixed at
-1, corresponding to a cosmological constant. In wCDM, $w$ is allowed to vary,
potentially differing from -1. It’s worth noting that depending on the
specific analysis or dataset being used, additional parameters might be
included, such as: • $\tau$ ($\tau$): The optical depth to reionization. •
$\Omega_{\mathrm{k}}$ ($\Omega_{\mathrm{k}}$): The curvature parameter. •
$N_{\mathrm{eff}}$: The effective number of relativistic species. However, the
core wCDM model is typically described by the 6 parameters listed above. The
exact number of parameters can vary slightly depending on the specific
formulation or application of the model.

Figure 2: Comparison of cosmosage and Claude Sonnet 3.5 (2024-06-20 version)
on a cosmology question testing model knowledge. cosmosage consistently
outperforms large generalist models on questions of this nature.

A common concern when presenting a new model is that the examples shown are
carefully selected, or “cherry-picked,” to highlight the model’s strengths and
downplay its weaknesses. To avoid such concerns, here is the explicit way in
which the examples showcased in this work were selected.

The generalist models used for comparison (ChatGPT and Claude) were chosen due
to their community perception as the strongest generalist models at the time
of writing. Each model is shown once and this selection was based on the
particular output. The cosmosage outputs presented in Figures 3 and 2 were
generated only once. The calculation example in Figure 4 was generated twice
due to an incorrect initial response.

Lacking a cosmology-specific benchmark, no quantitative evaluation was
performed in this work. However, I did submit cosmosage to the publicly
available benchmark suite “Open LLM Leaderboard”. The results are shown in
Table 1, showing that the specialization in cosmology only slightly reduced
the general capabilities of the model, and in fact has enhanced the capability
of the model on the grade-school mathematics dataset GSM-8k.

Evaluation | cosmosage | OpenHermes2.5-Mistral7B
---|---|---
Avg. | 60.66 | 61.52
AI2 Reasoning Challenge (25-Shot) | 59.73 | 64.93
HellaSwag (10-Shot) | 80.90 | 84.18
MMLU (5-Shot) | 59.57 | 63.64
TruthfulQA (0-shot) | 50.98 | 52.24
Winogrande (5-shot) | 75.93 | 78.06
GSM8k (5-shot) | 36.85 | 26.08
Table 1: Performance comparison on Open LLM Leaderboard benchmarks for non-
cosmology skills. cosmosage and OpenHermes2.5-Mistral7B share a base model and
instruction-tuning dataset, differing only in cosmosage’s additional
cosmology-specific training data, as well as potential differences in training
hyperparameters. Results suggest that cosmology specialization did not
substantially impact performance on general tasks.

Question: What is ”stacking”? cosmosage answer: Stacking is a technique used
in astronomy to combine multiple images or data sets of the same object into a
single image or dataset. This can be done by aligning and co-adding the
individual images or data points, which helps to improve the signal-to-noise
ratio and reveal faint features that may not be visible in a single image or
dataset. Stacking is often used to study the properties of galaxies, galaxy
clusters, and other astronomical objects. ChatGPT answer: Stacking is an
ensemble learning technique used in machine learning to improve the
performance of predictive models. It involves combining multiple different
models (often referred to as base models or level-0 models) to create a more
robust and accurate predictive model. Here’s how stacking generally works: •
Train Base Models: Multiple base models are trained on the same dataset. These
models can be of different types (e.g., decision trees, logistic regression,
neural networks) or variations of the same type with different
hyperparameters. • Generate Meta-Features: Each base model makes predictions
on the training data. These predictions, along with the original features,
[response continues]

Figure 3: Comparison of responses to an ambiguous query between cosmosage and
ChatGPT (GPT-4o, 2024-07-03). cosmosage, fine-tuned for cosmology, correctly
interprets the context and provides a relevant explanation of the term
“stacking.”

Question: Calculate the angular diameter distance to the Bullet cluster. Show
your work. cosmosage answer: To calculate the angular diameter distance (dA)
to the Bullet cluster, we need to know its redshift (z). The redshift of the
Bullet cluster is approximately 0.296. The angular diameter distance is
related to the comoving distance by the formula: dA(z) = dC(z)/(1 + z) where
dA(z) is the angular diameter distance, dC(z) is the comoving distance, and z
is the redshift. First, let’s find the comoving distance using the Hubble
constant (H0), which is approximately 70 km/s/Mpc. We can use the Hubble-
Lemaître law for the comoving distance: dC(z) $\approx$ c/H0 * z where c is
the speed of light (approximately 300,000 km/s). Now, let’s convert the
redshift from dimensionless to a proper length: z $\approx$ 0.296 c * z
$\approx$ 89,200 km/s Now, we can find the comoving distance: dC(z) $\approx$
89,200 km/s / 70 km/s/Mpc $\approx$ 1273 Mpc Finally, we can find the angular
diameter distance: dA(z) = dC(z)/(1 + z) dA(z) = 1273 Mpc / (1 + 0.296) dA(z)
$\approx$ 1035 Mpc So, the angular diameter distance to the Bullet cluster is
approximately 1035 Mpc.

Figure 4: Here, cosmosage is asked to do something it is not capable of doing:
a non-trivial calculation that requires numerical integration. Despite this,
it makes a valiant attempt by using the Hubble-Lemaître law (which is
approximately correct for low-redshift objects) and gets close to the correct
answer of 918.4 Mpc, despite inaccurate arithmetic division. The completions
for Claude Sonnet 3.5, GPT-4o, LLAMA-70B and Phi-3 are not shown for brevity,
but all perform erroneous calculations, often attempting the numerical
integral and hallucinating the result.

##  4 Availability of Products

The cosmosage project aims to provide a useful natural language assistant to
the cosmology community and anyone interested in cosmology. Therefore, to
promote openness, collaboration, and visibility, various components have been
made publicly available under the extremely permissive MIT license. The source
code used for training and inference, along with documentation and usage
examples, is publicly available on GitHub at:
https://github.com/tijmen/cosmosage

For researchers and developers interested in experimenting with the model
directly, the trained cosmosage model parameters (“weights”) are hosted on
Hugging Face at
https://huggingface.co/Tijmen2/

Furthermore, despite the resource-intensive nature of running large language
models, cosmosage has been deployed for academic use on the Helmholtz AI
Blablador inference server hosted by Jülich Supercomputing Centre (JSC):
https://helmholtz-blablador.fz-juelich.de/

##  5 Potential Enhancements and Extensions

While cosmosage shows significant capabilities in its current state, there are
numerous avenues for improvement. A clear next step is to rerun the training
procedure on an larger, stronger, and newer base model, such as gemma-2-27b.

Another direction that I am in the process of exploring is to give the model
access to the source material at inference time via retrieval augmented
generation (RAG). This technique allows the LLM to access and utilize the
original material, reducing the frequency of hallucinations.

Another exciting direction is to explore the potential of multilingual
training and deployment. While the current version was trained exclusively on
English data, the starting base model is multilingual. Surprisingly, the
multilingual capabilities of the base model persist into the final cosmosage
model. This was found by testing cosmosage in Japanese, where it answers
coherently. By expanding the training corpus to include cosmological resources
in other languages, one could further enhance cosmosage’s multilingual
capabilities and further democratize access to cosmological information.

Finally, as the field of LLMs continues to evolve at a rapid pace, it is
crucial to stay up to date with the latest advancements and incorporate them
into cosmosage’s development. By providing the source code and weights openly,
the hope is that this work can be advanced by anyone in the community.

##  6 Conclusion

cosmosage is a natural-language assistant specifically tailored for cosmology.
Through continued pretraining on a corpus of cosmology literature and
subsequent fine-tuning on a large dataset of synthetically generated question-
answer pairs, cosmosage demonstrates an understanding of cosmological concepts
and surpasses general-purpose LLMs in this domain. While quantitative
benchmarking remains a challenge due to the lack of standardized tests for
specialized LLMs, qualitative evaluation through practical use cases
highlights cosmosage’s superior performance in answering cosmology-related
questions. The model’s development and deployment, including public
availability of the source code, trained weights, and an instance hosted on
the Helmholtz AI Blablador inference server, have created a useful assistant
that can benefit researchers, educators, students, and anyone interested in a
better understanding of the cosmos.

## Acknowledgements

This work was made possible by the Center for Computational Astrophysics,
National Astronomical Observatory, Japan (CfCA, NAOJ). Tijmen de Haan was
supported by World Premier International Research Center Initiative (WPI),
MEXT, Japan. Many thanks to Alex Strube for making cosmosage available on the
Blablador inference service.

## References

  * [1] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A.N. Gomez et al., _Attention Is All You Need_ , 2017.  10.48550/arXiv.1706.03762.
  * [2] T.B. Brown, B. Mann, N. Ryder, M. Subbiah, J. Kaplan, P. Dhariwal et al., _Language Models are Few-Shot Learners_ , July, 2020.  10.48550/arXiv.2005.14165.
  * [3] F. Chollet, _On the Measure of Intelligence_ , Nov., 2019.  10.48550/arXiv.1911.01547.
  * [4] A.Q. Jiang, A. Sablayrolles, A. Mensch, C. Bamford, D.S. Chaplot, D.d.l. Casas et al., _Mistral 7B_ , Oct., 2023.  10.48550/arXiv.2310.06825.
  * [5] A. Karpathy, _arxiv-sanity-lite_ , July, 2024.
  * [6] L. Blecher, G. Cucurull, T. Scialom and R. Stojnic, _Nougat: Neural Optical Understanding for Academic Documents_ , Aug., 2023.  10.48550/arXiv.2308.13418.
  * [7] J. Durbin, _bagel-34b-v0.2_ , 2023.
  * [8] S. Lin, J. Hilton and O. Evans, _TruthfulQA: Measuring How Models Mimic Human Falsehoods_ , May, 2022.  10.48550/arXiv.2109.07958.
  * [9] A.Q. Jiang, A. Sablayrolles, A. Roux, A. Mensch, B. Savary, C. Bamford et al., _Mixtral of Experts_ , Jan., 2024.  10.48550/arXiv.2401.04088.
  * [10] OpenAccess-AI Collective, _OpenAccess-AI-Collective/axolotl_ , July, 2024.
  * [11] D.P. Kingma and J. Ba, _Adam: A Method for Stochastic Optimization_ , Jan., 2017.  10.48550/arXiv.1412.6980.
  * [12] I. Loshchilov and F. Hutter, _Decoupled Weight Decay Regularization_ , Jan., 2019.  10.48550/arXiv.1711.05101.
  * [13] S. Rajbhandari, J. Rasley, O. Ruwase and Y. He, _ZeRO: Memory Optimizations Toward Training Trillion Parameter Models_ , May, 2020.  10.48550/arXiv.1910.02054.
  * [14] Teknium, _OpenHermes 2.5: An Open Dataset of Synthetic Data for Generalist LLM Assistants_ , 2023.