grKnight/astraq-vl-stage2

AstraQ-VL Stage-2 (connector + LoRA instruction tuning)

AstraQ-VL Stage-2 is the public name for this connector-plus-LoRA checkpoint.

A LLaVA-style vision–language model that lets Qwen2.5-1.5B-Instruct answer questions about astronomy images encoded by CLIP ViT-L/14. This is the AstraQ-VL Stage-2 model: it warm-starts the AstraQ-VL Stage-1 connector and continues training it jointly with LoRA adapters on the Qwen LLM, on the caption + GPT-4 QA records of UniverseTBD/AstroLLaVA_convos. The CLIP vision tower stays frozen. Trained on a disjoint held-out test split so it can be evaluated on unseen images.

Stage 1 aligned the connector with the LLM frozen — it grounds coarse visual structure but hallucinates fine specifics. Stage 2 opens up the LLM (via LoRA) so the model learns to use the visual evidence when committing to answers — the recipe's instruction-tuning step.

⚠️ This repository ships the connector + LoRA adapter only (not full LLM weights). It is not a standalone transformers model — it needs the custom VLM code from the astraq-vl repo, the two base models (auto-downloaded from the Hub), and peft to run.

Download

The repository contains checkpoints saved every 200 steps and the final checkpoint at step 2526. For inference, download the final checkpoint directory:

Artifact	Contents
checkpoints/checkpoint-2526/	Final connector, LoRA adapter, metadata, and training state.
astraq-vl-stage2-metrics.zip	Stage-2 aggregate and per-sample metrics.
eval_loss_curve.zip	Held-out loss curve in CSV, JSON, and PNG formats.

checkpoint-2526/ contains the continued-trained connector (connector.safetensors), the trained LoRA adapter (lora/adapter_model.safetensors + adapter_config.json), optimizer/scheduler state (training_state.pt), and meta.json (step + final loss). Both the connector and the LoRA are required at inference.

Evaluation artifacts

Artifact	Scope	Contents
astraq-vl-stage2-full-heldout-eval-v1.zip	Full held-out: captions + QA	Predictions and aggregate/per-sample metrics for all 3,271 held-out records: 586 caption records and 2,685 QA records, plus comparisons, config, test split, and reproduction notes.
phase0_stage2_results.zip	Phase 0 (captions only)	Caption predictions for 591 held-out images, with NLI and SBERT aggregate/per-sample scores; 586 images have reference captions used for scoring.

The Phase 0 archive does not include the held-out QA records. Use the full-heldout artifact for the combined caption + QA evaluation.

Architecture

image ─► CLIP ViT-L/14 (FROZEN) ─► MLP connector (TRAINED, init from Stage-1) ─► Qwen2.5-1.5B + LoRA (base FROZEN, LoRA TRAINED) ─► text
                                    1024 → 1536 → 1536

• Vision: openai/clip-vit-large-patch14, penultimate-layer patch features (frozen)
• Connector: 2-layer MLP with GELU, 1024→1536→1536; warm-started from Stage-1 checkpoint-3789 and kept trainable
• LLM: Qwen/Qwen2.5-1.5B-Instruct, base frozen + LoRA adapters (r=16, α=32, dropout 0.05) on q/k/v/o/gate/up/down_proj across all 28 layers
• Trainable / total: 22,400,000 / 1,868,879,360 (1.20%) — connector 3,935,232 + LoRA 18,464,768

Training

Data	UniverseTBD/AstroLLaVA_convos, same per-image held-out split as Stage-1: train 161,653 recs / 29,151 imgs, test 591 imgs / 3,271 recs
Initialization	connector ← Stage-1 checkpoint-3789 (epoch 3); LoRA ← fresh (no-op init)
Objective	next-token cross-entropy on answer tokens only (connector + LoRA trainable)
Epochs / steps	1 epoch, 2,526 update steps
Effective batch	64 (per-device 4 × grad-accum 16)
LR / schedule	2e-4, cosine with 3% warmup (75 steps)
Max length	512 (+256 image tokens)
Precision	bf16 (autocast) + gradient checkpointing
Hardware	1× RTX 6000 Ada (48 GB), 15 samples/s (3 h)
Held-out loss	1.60 (step 200) → 1.452 (step 2526), decreasing monotonically — see Training curve below

The full-LLM backward pass (absent in Stage-1) is the memory driver, hence per-device batch 4 + gradient checkpointing to fit ~48 GB. One epoch is the LLaVA instruction-tuning convention — the model only needs to learn to use the already-aligned visual features, not to align them from scratch.

Training curve

Held-out validation loss, recomputed per checkpoint on a fixed 512-sample subset of the unseen test.json and averaged over its answer tokens. (The per-step training log wasn't retained, so this was reconstructed from the saved checkpoints with scripts/eval_loss_curve.py — which makes it a true held-out curve rather than a noisy train-loss trace.) It falls monotonically and flattens by the end of the single epoch, consistent with the 1-epoch choice:

step	200	400	600	800	1000	1200	1400	1600	1800	2000	2200	2400	2526
held-out loss	1.605	1.571	1.548	1.526	1.508	1.494	1.479	1.471	1.462	1.456	1.454	1.452	1.452

Regenerate with python scripts/eval_loss_curve.py --config configs/finetune_astraq_vl_stage2.yaml --checkpoint-dir checkpoints/astraq-vl-stage2 --records-json datasets/astrollava_llava/test.json --image-dir datasets/astrollava_llava/images --num-samples 512 --plot (full series in eval_loss_curve.csv).

Usage

# 1. get the code
git clone https://github.com/crimsonKn1ght/astraq-vl && cd astraq-vl
pip install -r requirements.txt        # includes peft

# 2. download the final checkpoint directory
hf download grKnight/astraq-vl-stage2 --include "checkpoints/checkpoint-2526/**" --local-dir astraq-vl-stage2

# 3. answer a question about an image (CLIP + Qwen auto-download; peft loads the LoRA)
python inference.py \
  --config configs/finetune_astraq_vl_stage2.yaml \
  --checkpoint astraq-vl-stage2/checkpoints/checkpoint-2526 \
  --image your_astro_image.jpg \
  --prompt "What type of object is this and what is notable about it?" \
  --temperature 0

Pass the Stage-2 config so the LoRA modules are built before the adapter weights load; the loader then restores both the connector and the LoRA automatically. Caption-only predictions are in the Phase 0 archive; combined caption + QA predictions are in the full-heldout archive.

Capabilities & limitations

Stage 2 fine-tunes the LLM (LoRA) jointly with the connector, so — unlike Stage-1 — the language model itself learns from the QA pairs rather than improvising specifics from its frozen prior. The intended effect is fewer hallucinated fine details (catalog numbers, instruments, dates) on question-answering prompts, on top of Stage-1's coarse visual grounding. Compare Stage 2's predictions_full_heldout.jsonl with the corresponding Stage-1 held-out predictions to inspect the difference on the same held-out split.

Limitations carried over from the design: CLIP's 224×224 input discards fine astronomical detail; the base LLM is small (1.5B); and LoRA is a low-rank adaptation, not a full fine-tune. Evaluation is a held-out generation set, not a full quantitative benchmark — read results qualitatively.

Reproduction

The full-heldout evaluation archive contains REPRODUCE_FULL_HELDOUT.md, the Stage-2 config, and the exact test.json split used for that evaluation. The split is seeded, so the build command below reproduces the train/test partition.

prereq: Stage-1 connector checkpoint-3789 (grKnight/astraq-vl-stage1 ep3 bundle)
build:  python scripts/build_astrollava_trainset.py --include-qa --max-image-size 384 --test-fraction 0.02 --seed 42
train:  python train.py --config configs/finetune_astraq_vl_stage2.yaml
eval:   python scripts/batch_inference.py --config configs/finetune_astraq_vl_stage2.yaml --records-json datasets/astrollava_llava/test.json --num-samples 0 ...

Citation

If you use AstraQ-VL, this checkpoint, or its evaluation artifacts, please cite:

Roy, G. (2026). AstraQ-VL: Parameter-Efficient Astronomy Vision-Language Modeling with Connector Alignment and LoRA Tuning (Version v1). Zenodo. https://doi.org/10.5281/zenodo.21284851

@misc{roy2026astraqvl,
  author    = {Roy, Gourab},
  title     = {AstraQ-VL: Parameter-Efficient Astronomy Vision-Language Modeling with Connector Alignment and LoRA Tuning},
  year      = {2026},
  publisher = {Zenodo},
  version   = {v1},
  doi       = {10.5281/zenodo.21284851},
  url       = {https://doi.org/10.5281/zenodo.21284851}
}

License & attribution

• Weights: cc-by-sa-4.0, inherited from the training data.
• Training data: UniverseTBD/AstroLLaVA_convos (CC-BY-SA-4.0); imagery from NASA APOD, ESO, and NASA/ESA Hubble.
• Base models: Qwen2.5-1.5B-Instruct (Apache-2.0), CLIP ViT-L/14 (OpenAI, MIT).
• Builds on: AstraQ-VL Stage-1 and the AstroLLaVA work.