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:

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

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

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:

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.