AstroSage-Llama-3.1-8B: Expert Astronomy LLM

• AstroSage-Llama-3.1-8B is a domain-specialized large language model focused on astronomy, astrophysics, and instrumentation, built on Meta’s Llama-3.1-8B architecture.
• It achieves state-of-the-art performance with 80.9% accuracy on the AstroMLab-1 benchmark by leveraging continued pretraining on a curated astronomy corpus and 8.8 million synthetic Q&A pairs.
• The model features mechanistic interpretability using Sparse Autoencoders and open-access integration, making it a cost-effective tool for research and educational applications in astronomy.

AstroSage-Llama-3.1-8B is a domain-specialized LLM designed to deliver expert-level performance in astronomy, astrophysics, cosmology, and astronomical instrumentation. Built upon Meta’s Llama-3.1-8B Transformer architecture, AstroSage combines high-quality continued pretraining on a curated astronomy corpus, large-scale synthetic question-answering supervision, and targeted model merging. It sets a new standard in domain-specialized LLMs by achieving 80.9% accuracy on the stringent AstroMLab-1 benchmark—exceeding all other 8-billion parameter models and matching the performance of proprietary systems 40 times its size, such as GPT-4o. The model and its training resources are openly released, enabling broad integration into both research and educational workflows (Haan et al., 2024).

1. Model Architecture and Design

AstroSage-Llama-3.1-8B inherits all core architectural hyperparameters from its Llama-3.1-8B base:

• Model size: 8 billion parameters, decoder-only transformer.
• Layers: 32 transformer blocks.
• Hidden dimension: $d_\text{model} = 4096$.
• Feed-forward dimension: $d_\text{ff} = 11,008$ (approximately $2.68 \times d_\text{model}$).
• Attention heads: 32, each with $d_k = 128$ head dimension.
• Positional encoding: Rotary positional embeddings (RoPE).
• Activation: Gated GeLU in MLP blocks.
• Normalization: LayerNorm after each block.
• Tokenizer: UTF-flexible tiktoken, with no added astronomy-specific tokens.
• Specialization: No use of bespoke modules such as adapters or LoRA; all specialization is achieved via data-centric methods.

The attention mechanism applies

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{Q K^\top}{\sqrt{d_k}}\right)V$$

with softmax applied row-wise. No deviations from the core Llama-3.1 block structure are present, ensuring compatibility with standard open-source tools and deployment platforms (Haan et al., 2024).

2. Corpus Construction and Supervised Fine-Tuning

Continued Pretraining (CPT) Corpus

• Primary data: ≈250,000 astronomy/arXiv preprints (astro-ph., gr-qc.), 2007–2024.
• Augmentation: Depth-2 Wikipedia crawl (astronomy/astrophysics), selected textbooks (OCR and markdown conversion).
• Token volume: 3.3 billion tokens (19.9 GB).
• Cleaning: Paragraph-wise perplexity filtering under the base model; top 2% ppx paragraphs removed.
• Corpus finalization: Cleaned documents reassembled, preserving ≈98% of initial volume.

Synthetic Question-Answer (Q$$) Data for SFT</h3> <ul> <li><strong>Generation:</strong> 11 million Q$$ pairs from the CPT corpus using the de Haan et al. method.

Filtering: Four criteria (correctness, stand-alone clarity, pertinence, quality) assessed by LLMs; final set includes 8.8 million high-quality Q$ pairs.</li> <li><strong>Enrichment:</strong> Variations in question phrasing, referencing style, randomized preambles.</li> <li><strong>Other tasks:</strong> Filtered entries from Infinity-Instruct-7M, synthetic &quot;paper summary&quot; prompts, metadata Q$s, and five hand-curated web corpora.

Total SFT tokens: ≈2.0 billion.

Training Protocols

Stage	Hardware	Epochs	LR	Optimizer (β₂)	Weight Decay	Precision
CPT	1,472 AMD MI250X GPUs	2	$1.5 \times 10^{-4}$ (const, quad. warmup)	0.95	$1 \times 10^{-3}$	BF16
SFT	1,472 AMD MI250X GPUs	6	$1 \times 10^{-4}$ (quad. warmup, cosine decay)	0.95	0	BF16

• Sequence length: 8,192. Micro-batch per GPU: 3.
• Gradient clipping: Max norm=3.0.
• Sharding: FSDP, auto-wrap.
• Wall time: 10 hours CPT; 9.5 hours SFT.

Loss curves during both phases indicated stable convergence and no overfitting. No proprietary modules were used in the specialization pipeline (Haan et al., 2024).

3. Evaluation, Benchmarking, and Cost Analysis

AstroMLab-1 Benchmark

• Test set: 4,425 human-verified multiple-choice astronomy questions, sourced from Annual Review of Astronomy & Astrophysics; all source ARAA papers excluded from training to enforce generalization.
• AstroSage-Llama-3.1-8B accuracy: 80.9%.
• 8B class comparison: Llama-3.1-8B base: 72.9%; all other open and proprietary 8B models: <72%.
• GPT-4o accuracy: 80.4%; Claude-3.5-Sonnet: 85.0%.
• Human experts baseline: ~68%.

Statistical Significance

• 95% Wilson interval for AstroSage (n = 4,425): $0.809 \pm 0.008$
• Margin vs. base model: 7.2 percentage point absolute gain; parity with GPT-4o is statistically robust.

Cost Efficiency

• Inference cost (per $10^5$ tokens):
  • AstroSage-8B ≈ 1/100× cost of open 90B models
  • AstroSage-8B ≈ 1/1,000× cost of GPT-4o
• This suggests that domain specialization can yield high-fidelity expert performance at orders-of-magnitude lower compute budgets (Haan et al., 2024).

4. Mechanistic Interpretability with Sparse Autoencoders

AstroSage-Llama-3.1-8B is equipped for mechanistic interpretability via integration with Llama Scope’s suite of 256 Sparse Autoencoders (SAEs) (He et al., 2024):

• SAE Coverage: 32 transformer layers × [post-MLP residual, attention output, MLP output, transcoder input→MLP output] × [8×, 32× expansion] = 256 SAEs.
• Feature counts: 32K (8×) and 128K (32×) per SAE.
• Sparsity: Top-K selection (typically $d_\text{ff} = 11,008$0 active features per token); “JumpReLU” thresholded and annealed during training.
• Reconstruction loss: MSE; explained variance (EV) and ΔLM-loss recorded.
• Generalization: SAEs trained on the base model generalize directly to AstroSage, maintaining performance even on out-of-distribution contexts, longer sequences (up to 8,192 tokens), and in instruction-tuned models.

Applications

• Activation probing: Extraction and interpretation of monosemantic features (e.g., “stellar nucleosynthesis,” “Kepler transit,” “dark-matter halo”).
• Feature splitting: Wider (128K) SAEs uncover rare, domain-specific representations (e.g., a distinct “Brexit” or “dark-matter halo” feature).
• Scientific probing: Lightweight linear probes on SAE latents achieve robust classification of astrophysical sub-domains.
• Causal tracing: Zeroing single features in the residual stream enables measurement of their causal impact on scientific keyword prediction.

Integration Workflow

• Plug-and-play: Llama Scope SAEs can be downloaded and directly hooked into AstroSage’s forward pass with normalization constants and mask selection to extract or ablate features during inference.
• Recommended configuration: Post-MLP residual SAEs at 32K width for best sparsity-fidelity tradeoff; 128K for rare features (He et al., 2024).

5. Domain Specialization Versus Scale: Comparative Analysis

AstroSage-Llama-3.1-8B demonstrates that targeted continued pretraining and task-specific supervised fine-tuning can outperform naively scaled models within a restricted domain:

• Peer comparison: AstroSage-8B surpasses all other models in its parameter class by >8 points on AstroMLab-1.
• Large model parity: Achieves GPT-4o-level performance (despite a 40× parameter gap).
• Generalization: Maintains core language modeling and mathematical capabilities; minimal trade-off against generalist instruction following (IF-EVAL) due to strategic model merging (25% of Meta’s instruct model weights).
• Failure cases: Multi-step derivations and granular metadata recall remain difficult.

A plausible implication is that domain-specialized models, if trained on well-curated, high-fidelity data and supervised with high-quality synthetic Q&A relevant to real scientific tasks, offer a more resource-efficient path to expert-level AI assistance than universal scaling alone (Haan et al., 2024).

6. Applications, Accessibility, and Future Directions

AstroSage-Llama-3.1-8B is open access under the Llama 3.1 Community License via https://huggingface.co/AstroMLab/AstroSage-8B, with weights in PyTorch and safetensors formats. Integration into research and educational stacks is facilitated by standard APIs and compatibility with Hugging Face Inference, Jupyter, and Virtual Observatory services.

Principal Use Cases

• Research Automation: Literature review, summarization of new arXiv papers, data-driven hypothesis generation, and code generation for data reduction pipelines.
• Education: Adaptive “astronomy tutor,” problem set and worked solution creation across subfields.
• Instrumentation: Documentation, troubleshooting, and workflow support for astronomy software and hardware.
• Mechanistic interpretability: Extraction, annotation, and causal testing of medium/high-level scientific representations using pre-trained SAEs.

Ongoing and Future Work

• Scaling: Extension of the current paradigm to 70B-parameter models, aiming to support more nuanced scientific reasoning and retrieval-augmented tasks.
• Real-time augmentation: Addition of retrieval mechanisms and real-time literature updating to ensure models remain current.
• Community collaboration: Open release of code, training scripts, and interpretability tools, with ongoing support for extension into other scientific domains (Haan et al., 2024, He et al., 2024).