chenyeikai/VisProbe-Molmo2-4B-LoRA

VisProbe · Molmo2-4B LoRA (probabilistic visual classifier)

A LoRA adapter that turns allenai/Molmo2-4B into a probabilistic closed-set visual classifier. At inference the model does not generate free-form text: the answer-position logits are restricted to the candidate codes (A/B/C/D, or A=false/B=true) and renormalised with a softmax, producing an auditable probability distribution.

Released as part of VisProbe — code, data pipeline, and the full experiment log: https://github.com/Hanzoghos/visprobe

Results — frozen baseline vs this adapter

Same test sets, same scoring protocol, single RTX 3090, bf16. Accuracy (and absolute gain):

Test set	Baseline	After LoRA	Gain
SNLI-VE (visual entailment, 3-way, n=21,482)	0.8720	0.9096	+3.76 pp
VQAv2 yes/no (seeded 1/3 subset, n=26,846)	0.8671	0.8867	+1.96 pp
Visual7W (4-way multiple choice, n=6,695)	0.8774	0.9159	+3.85 pp
Macro average (equal weight)	0.8722	0.9041	+3.2 pp

Full accuracy/ECE/Brier tables and loss curves: see results/SUMMARY.md and EXPERIMENTS.md in the VisProbe repository. Molmo2-4B had the strongest frozen baseline of the released models; after fine-tuning it ties for the highest macro accuracy, with the best post-finetune ECE/Brier calibration of the set.

Training

Item	Value
Data	99,999 mixed examples (33,333 per dataset, stratified by label, seed=42)
Epochs / steps	1 / 25,000 (batch size 4; per-example micro-batches with gradient accumulation)
LoRA	r=8, alpha=16, dropout=0.05, bias=none
Target modules	att_proj, attn_out, ff_proj, ff_out (language decoder only; vision tower frozen)
Optimisation	lr 1e-4, 3% linear warmup + cosine decay, grad clip 1.0, AdamW
Precision	bfloat16, gradient checkpointing (use_reentrant=False)
Checkpoint	final-from-best: lowest recent-window training loss
Hardware / time	1× NVIDIA RTX 3090 24 GB / roughly 22–30 h (3.2–5.2 s per step)
Software	transformers 4.57.1 (pinned), peft, trust_remote_code=True

Adapter weights: 57.3 MB in safetensors. About 14.3M trainable parameters (0.29%).

Compatibility note: Molmo2 ships custom remote modelling code whose embedding build uses an in-place write that crashes under LoRA + gradient checkpointing. VisProbe applies an automatic out-of-place patch (vlm/training/molmo_patch.py) during training and always loads the processor from the base model. Use the VisProbe loading path below for a tested setup; see docs/molmo2-training-fix.md in the repository for details.

How to use

With VisProbe (recommended — gives probabilities, not text)

from vlm.scoring.vlm_scorer import VLMScorer

# trust_remote_code=True is applied automatically for the molmo2 family
scorer = VLMScorer("molmo4b", adapter="chenyeikai/VisProbe-Molmo2-4B-LoRA")
result = scorer.score(image, fields, context_text=context)
# result["fields"][name]["scores"] -> {"A": p, "B": p, ...} candidate probabilities

Generic PEFT loading

The processor/tokenizer must be loaded from the base repo; the adapter upload intentionally contains no processor files:

import torch
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoProcessor

base = AutoModelForCausalLM.from_pretrained(
    "allenai/Molmo2-4B", torch_dtype=torch.bfloat16, device_map="auto",
    trust_remote_code=True)
model = PeftModel.from_pretrained(base, "chenyeikai/VisProbe-Molmo2-4B-LoRA")
processor = AutoProcessor.from_pretrained("allenai/Molmo2-4B", trust_remote_code=True)

Scope & limitations

• The adapter is specialised for closed-set choice / boolean questions rendered with the VisProbe prompt schema (single-token answer codes). It is not a general-purpose chat upgrade.
• Training data are three academic visual-reasoning benchmarks; in-domain behaviour is measured by the tables above, out-of-domain behaviour is not.
• VQAv2 results use a seeded random 1/3 test subset (standard error ≈ ±0.5 pp).

License

Adapter weights: Apache-2.0, following the base model allenai/Molmo2-4B. Datasets remain under their upstream licenses (Flickr30k / COCO / Visual Genome based).