codewithdark/EVLMs-v2-MIMIC-CXR

EVLMs-v2 — Medical VLM (mHC + CSA + Muon) baseline on MIMIC-CXR-VReason

Explainable vision-language model for chest X-ray finding classification and structured radiology-style rationale generation, trained on 16,000 MIMIC-CXR-VReason interleaved multi-image chat samples.

Architecture

• Vision: Swin-B patch4 window7 224 (timm), primary image keeps the full 49-token spatial map; tool-returned ROI images are mean-pooled to 1 token.
• Language: Gemma-3-1B (LoRA fine-tuned), interleaved slots spliced into the token stream at their exact positions.
• DeepSeek-V4-inspired components: mHC (multi-head hyper-connections, doubly-stochastic routing), CSA (compressed sparse attention), Muon optimizer (Newton-Schulz orthogonalized momentum) + AdamW for 1-D params.
• Loss: adaptive (Kendall uncertainty) weighting of classification + language losses; per-class positive weighting for ChestX-ray14 labels.

Training summary

• 16,000 samples, batch 16, 6 epochs (10,000 optimizer steps incl. an earlier segment), cosine schedule with 10% warmup, lr 1e-4, bf16, SDPA.
• Final segment metrics (this repo's weights = best_model):
  • epoch 1: train_loss=0.8919 (cls=0.8082, lang=0.2424) | val_loss=0.9098 mean_auc=0.6783
  • epoch 2: train_loss=0.8133 (cls=0.7954, lang=0.2256) | val_loss=0.9022 mean_auc=0.6758
  • epoch 3: train_loss=0.7690 (cls=0.7894, lang=0.2142) | val_loss=0.9048 mean_auc=0.6786
  • epoch 4: train_loss=0.7455 (cls=0.7822, lang=0.2081) | val_loss=0.8900 mean_auc=0.6810

Files

• final_model/pytorch_model.bin — raw model.state_dict() (EVLMsV2 custom nn.Module). Load: rebuild EVLMsV2(EVLMsV2Config(...)) from the repo code (github.com/codewithdark-git/EVLMs) then load_state_dict(torch.load(...)).
• checkpoints/ — ALL resumable checkpoints (model + optimizer + scheduler + step counter) incl. best_model.pt.
• artifacts/epoch_*/ — GradCAM explainability overlays, sample generations, per-epoch metrics.
• config.json — intentionally ONLY the HuggingFace token and Weights & Biases fields (per project requirement). Values are redacted because this repo is public; authenticate via HF_TOKEN / WANDB_API_KEY env vars.

W&B

Training internals (mHC doubly-stochastic matrices, CSA gates, adaptive loss weights, per-component losses) are logged on Weights & Biases: entity codewithdark-islamia-college-peshawar, project evlms-v2, run mhc-csa-muon-16k-resume-fixed.