genome-assembly-general-purpose-ft-bitnet-2b

GAQA v1 — Fine-Tuned BitNet b1.58 2B for Genome Assembly Quality Assessment

This is a fine-tuned version of microsoft/bitnet-b1.58-2B-4T-bf16 for genome assembly quality assessment (GAQA): automated diagnosis of assembly problems (low coverage, contamination, misassembly), remediation advice, and assembler/tool selection. It is the model from the paper:

GAQA v1: Fine-Tuning Ternary BitNet b1.58 for Genome Assembly Quality Assessment — A 2B Model Outperforming General-Purpose 7B LLMs on External Heldout Data (Hmade & Guelzim).

Fine-tuning uses LoRA (rank r=32, scaling α=64, dropout 0.05) over the frozen ternary base weights. The adapter adds ~43M trainable parameters (about 372M trainable in total, including the token embeddings — 15.15% of the model). For deployment the ternary base is packed into ~1.58 bits/weight, so the full model — packed ternary base + LoRA adapter — runs in ~1.22 GB of VRAM and outperforms general-purpose 7B instruct models on organism-heldout assembly QC.

Results (from the paper)

Per-family heldout: tool selection 100%, qc_diagnosis AUC 0.959, remediation 39.5%. Peak VRAM ≈ 1.22 GB.

Training

• Base model: microsoft/bitnet-b1.58-2B-4T-bf16 (ternary weights in {−1, 0, +1}, ~1.58 bits/weight).
• Method: LoRA over all attention (q/k/v/o) and MLP (gate/up/down) projections — 210 layers total.
• Loss: completion-only (answer tokens only), straight-through estimator (STE).
• Schedule: 3 epochs (~4,143 steps), sequence length 1,024, AdamW 8-bit, bfloat16.
• Hardware: a single A100 80 GB GPU, ~35 minutes.
• Dataset: GAQA v1 — 11,043 training examples (4,740 qc_diagnosis + 4,740 remediation + 159 tool_selection + 1,404 knowledge).

Usage

This repository contains the packed ternary base weights (packed_weights.pt) and the LoRA adapter (lora_weights.pt), plus a self-contained loader. Load the full model with:

from load_model import load_gaqa_v1_model

model, tokenizer = load_gaqa_v1_model()
# model is the packed ternary BitNet + GAQA v1 LoRA adapter, ready for inference

For the paper's logit-based forced-choice MCQ scoring (used to reproduce the results above), see the reproduction code at Zenodo 10.5281/zenodo.22760609.

Files

• packed_weights.pt — packed ternary base weights (~1.58 bits/weight).
• lora_weights.pt — final LoRA adapter (used in the paper's evaluations).
• lora_weights_best.pt — best-validation LoRA checkpoint (alternative).
• load_model.py — self-contained loader (BitLinear packing + LoRA application).
• adapter_config.json — LoRA hyperparameters (r=32, α=64, dropout=0.05) and base model reference.
• config.json, tokenizer.json, tokenizer_config.json, chat_template.jinja — base model configuration and tokenizer, for reference.

Limitations

Trained on short-read (Illumina) bacterial assemblies produced by SPAdes and MEGAHIT; see the paper for coverage of long-read/hybrid assemblers and other limitations.

License

MIT (inherited from the base model). Dataset: CC BY 4.0.

Citation

Please cite the paper (add the published DOI here) and the dataset:

@article{hmade_gaqa_v1,
  title = {GAQA v1: Fine-Tuning Ternary BitNet b1.58 for Genome Assembly Quality Assessment},
  author = {Hmade, Abdellah and Guelzim, Ibrahim},
  note = {dataset: https://doi.org/10.5281/zenodo.21762025}
}