pbcong/llava-1.5-7b-hal-verify-probe

Verify Probe for LLaVA-1.5-7B (hallucination detection)

A verification-pass hallucination probe for LLaVA-1.5-7B: a small per-layer MLP head that reads the host model's own hidden states while it answers a visibility question, and predicts whether an object mention is hallucinated. On a COCO CHAIR-80 human-GT benchmark it beats a 122B VLM judge that reads the model's outputs (AP .876 vs .826), despite adding only ~22 MB per seed on top of the frozen host.

What it does

For a caption produced by LLaVA-1.5-7B, each mentioned object noun becomes an (image, noun) pair. The probe runs one extra forward pass of the host model:

USER: <image>\nIs at least one {noun} visible in this image? Answer yes or no.\nASSISTANT:

and captures the hidden states at the answer position at layers [8, 12, 16, 20, 24]. The feature per layer is the contrast vector: answer-position hidden state with the image minus the same state without the image (d = 4096 per layer). Each layer has its own MLP head (4096 → 256 → 256 → 1, GELU + LayerNorm); the pair's hallucination score is the mean sigmoid over the 5 layer heads, averaged over the 3 seed checkpoints (ensemble).

Results (COCO CHAIR-80 human GT, 7,548-pair holdout)

detector (same eval universe)	AP	F1	within-word AUROC
LLaVA-7B self logit (same forward pass)	.718	—	—
generation-time probe	.781	—	—
122B VLM judge (logit readout)	.826	.777	.930
verify probe 7B (this repo, 3-seed ens)	.876 [.862, .889]	.814	.937

Paired ΔAP vs the 122B judge: +.049 [+.030, +.069], p < 1e-6. The probe and the self-logit rung are measured on the byte-identical forward pass, so the +.16 AP gap is purely internals vs outputs.

Training

• Data: 202k (image, noun) pairs from 50,405 COCO train2014 images; captions generated by the host model (greedy), object spans via CHAIR-80, labels from COCO human ground truth. Train/holdout images disjoint.
• Fit: BCE with positive re-weighting, label smoothing 0.98/0.01, AdamW lr 3e-4, weight decay 0.05, batch 256, 12 epochs; model selection on val within-word AUROC (10% of images held out by image id).
• Seeds 0/1/2 (files probe_verify_contrast_n50k_s{0,1,2}.pt); single-seed AP .872–.876, ensemble .876.

Files

• probe_verify_contrast_n50k_s0.pt, ..._s1.pt, ..._s2.pt — PyTorch state dicts of the per-layer MLP heads (one PairMLP each).
• probe_config.json — layers, dims, feature mode, metrics.

Usage

import torch, torch.nn as nn

LAYERS = [8, 12, 16, 20, 24]

class PairMLP(nn.Module):
    def __init__(self, d_in=4096, hidden=256):
        super().__init__()
        self.mlp = nn.ModuleList([
            nn.Sequential(nn.Linear(d_in, hidden), nn.GELU(), nn.LayerNorm(hidden),
                          nn.Linear(hidden, hidden), nn.GELU(), nn.LayerNorm(hidden),
                          nn.Linear(hidden, 1))
            for _ in LAYERS])

    def forward(self, X):  # X: {layer: (N, d_in) contrast features}
        return [m(X[l]).squeeze(-1) for m, l in zip(self.mlp, LAYERS)]

models = []
for s in (0, 1, 2):
    m = PairMLP()
    m.load_state_dict(torch.load(f"probe_verify_contrast_n50k_s{s}.pt",
                                 map_location="cpu"))
    m.eval()
    models.append(m)

# p_halluc = mean over seeds of (mean over layers of sigmoid(head(x)))
with torch.no_grad():
    p = torch.stack([
        torch.stack([torch.sigmoid(z) for z in m(X)]).mean(0)
        for m in models]).mean(0)

Feature extraction (the verify forward pass + contrast features) is general_hallucination/scripts/cocogt/verify_extract.py in the training repo; fitting/eval is verify_fit.py in the same directory.

Companion model

A 13B-hosted version of the same probe (AP .883) is at pbcong/llava-1.5-13b-hal-verify-probe.