Using Logit Space of VLMs for Attention to Detail

Give a Vision-Language Model one RGB image and ask it to sort it into one of n categories. Sounds standard. The catch is that the correct category here doesn't hinge on the big, obvious object in the frame — it hinges on a small detail buried inside it. Classifying a casualty and localizing a wound isn't a question of "is there a person in this photo," it's a question of "where exactly on that person is the blood." That distinction — big-object recognition versus small-detail localization — turned out to be the whole problem, and this post walks through the methods we tried and where we landed.

This work was built as the trauma and hemorrhage detection stage of a robotics perception pipeline for the DARPA Triage Challenge (Stage 2). We didn't use a public benchmark — that was out of scope for where we were in the project — and instead worked directly against DARPA's private dataset of simulated casualties.

Starting point: do existing vision backbones even look in the right place?

Before building anything, it's worth checking whether the problem is even hard for current models, or whether it's just a matter of reading out what they already know. So we took an image from the dataset and ran it through several state-of-the-art vision transformer and vision-language backbones, purely to inspect their attention maps.

SigLIP's attention spreads broadly across the person's torso and the ground around them — it's attending to "there is a human-shaped object here," not to any specific region of interest. That's a reasonable thing for a contrastive image-text model to learn, since its pretraining objective rewards matching a whole image to a whole caption. It's just not what we need.

DINOv3 without register tokens shows the same pattern — attention pooled loosely over the whole body, no real concentration on the limbs or the hand pressed against the wound. Adding register tokens changes this meaningfully: DINOv3-with-registers is the only backbone in this comparison whose attention starts to sharpen around the more localized, informative regions rather than the coarse silhouette. That's consistent with what register tokens are supposed to do — soak up the diffuse, low-information attention that would otherwise spread across background patches, freeing the rest of the attention budget to concentrate on content. Even so, it's a partial win, not a solved problem. Across the board, the models are telling us: I found a person, not I found the injury. That's the signal that pushed us toward vision-language models next — models that could, at least in principle, be asked directly about the detail rather than left to infer it from pixel salience alone.

Shifting to VLMs, and immediately hitting a wall

The DARPA data we had was sparse — nowhere near enough to fine-tune anything, whether full SFT or a lightweight LoRA adapter. That ruled out training a classifier from scratch and pushed us toward a zero-shot approach: prompt a pretrained VLM directly and use its output. The problem was that classification accuracy this way wasn't good enough, and our working hypothesis was that the failure wasn't in the model's language reasoning but in its localization — it wasn't reliably finding the small detail the classification depended on. That hypothesis is what the next set of experiments was built to test.

Attention-based cropping: does the VLM actually find the wound?

The first test was direct: take a VLM's attention weights over the image patches (its hidden_states), find the patch or patches with the highest attention, and crop around them. If the model's internal attention reliably lands on the injury, that crop should contain it.

LLaVA did this well — the high-attention patches clustered tightly around the blood, and cropping to that region reliably isolated it. That's an encouraging result, but a single crop throws away a lot of context: a wound doesn't exist in isolation from the limb it's on. So instead of cropping down to a single box, we shifted to keeping every high-attention patch plus its immediate one-block neighborhood in every direction — a small buffer zone around each hotspot, preserving the surrounding context the model was implicitly using before it answered.

This buffered version held up across different images and poses — the retained patches consistently tracked visible wounds and the limbs holding or near them, rather than drifting onto background clutter.

The gap between where a model looks and what it concludes

A natural question falls out of this immediately: if we ask for a confident patch indicating "blood" and none show up, does that mean there's no blood in the image? The answer is no, and the reason matters. Attention only tells you where the model directed its focus before generating a response — it's a map of where the model looked, not a map of what it decided was there. A patch can receive low attention and still contain the detail that matters; a patch can receive high attention and still be misread. Localization and recognition are separate questions, and attention only answers the first one. That gap is what pushed us toward inspecting the model's actual per-patch judgments, not just where it pointed its focus — which is where the logit lens comes in.

The logit lens: reading out what each layer already "believes"

The logit lens is an interpretability technique that takes intermediate transformer hidden states and projects them straight into the model's output vocabulary space, using the model's own (frozen) output head. Instead of only ever seeing the final-layer prediction, you get a prediction at every layer along the way — a running readout of what token the model would output if you stopped it right there. Applied across layers, it lets you watch a concept like "blood" or "injury" solidify as information moves through the network, rather than only appearing as a black-box answer at the end.

Why reach for this at all? Because the core problem seemed to sit deeper than localization. Even when a VLM's attention was landing in the right place, the model would still frequently answer incorrectly — which points to fragility in how the vision and language components actually communicate, not just where the vision side is looking. That fragility isn't surprising once you consider how these models are typically built: a vision tower (CLIP) and a language tower (LLaMA-family) trained separately, then stitched together through a single MLP projection layer. Leaning on autoregressive generation immediately after that thin a fusion point feels less like a principled design and more like a working patch. The logit lens gives a way to look inside that seam and see what's actually happening layer by layer, rather than trusting the final output alone.

We implemented this for LLaVA — the project is open-sourced as LLaVA-Lens.

Running the lens across image patches surfaces a clean split in the logit space. Patches with high confidence — above roughly 1% probability mass on a real token — carry actual semantic information: they're the patches doing real interpretive work, along with the context immediately around them. Patches below that threshold mostly aren't saying anything meaningful at all; they function more like connective tissue, steering the model's internal representation from one word to the next without encoding real content of their own. We started calling these low-confidence patches "soft prompts," since their role looks more structural than semantic. And the informative patches aren't spread evenly through the network — they concentrate heavily in the last 10 layers, where soft-prompt patches thin out and real, decodable information takes over.

Building a classifier directly on filtered logit-space embeddings

This split gives a natural way to build a lightweight, targeted classifier — one that doesn't require training on raw pixels at all. Instead of asking "what does the whole image show," the classifier is built to ask a narrower question: "which internal representations strongly indicate trauma?" The pipeline is:

• ▸Extract intermediate transformer layer embeddings from a pretrained VLM, precomputed rather than recomputed per training step.
• ▸Apply the logit lens (the frozen LM head) to identify which patches strongly activate meaningful, trauma-relevant tokens.
• ▸Keep only the high-confidence patches — the ones likely to encode real trauma cues rather than soft-prompt filler.
• ▸Add a small number of random background patches back in, so the classifier retains some sense of surrounding context rather than seeing only isolated hotspots.
• ▸Train a lightweight classifier on these filtered representations, not on the full patch sequence.

Concretely: precomputed hidden states (shape L × S × 4096, layers by sequence by hidden dim) go through logit-lens thresholding to select a reduced patch set S′, a CLS token is prepended, the result is projected down through an MLP (4096 → 512), passed through several multi-head self-attention blocks, layer-normed, and the CLS token is read out through a final linear layer (512 → 2) to produce a trauma / no-trauma prediction.

The dataset loader reflects this design directly: rather than loading raw images, it pulls precomputed hidden states from an HDF5 file for a selected band of transformer layers (roughly layers 21–32, where the earlier analysis showed the informative patches concentrating). Sinusoidal positional encodings are added across the layer dimension so the model retains a sense of which layer each embedding came from. Only the patch indices flagged by attention and logit-lens filtering are kept; the frozen LM head is applied to compute per-patch token probabilities, and patches above the confidence threshold are retained as the semantically meaningful set that feeds the classifier. Sequences vary in length after filtering, so batches are padded dynamically with attention masks.

Augmentation happens in embedding space rather than on pixels, which fits naturally with a pipeline that never touches raw images at training time:

• ▸Balanced resampling: class-0 samples are randomly downsampled each epoch to match class-1 count, so the classifier isn't trained against a skewed label distribution.
• ▸Context injection: 16–20 random non-salient patches are sampled from the discarded regions and appended back in, preventing the model from overfitting to only the highest-confidence trauma patches and losing surrounding structure.
• ▸Noise regularization: Gaussian noise with randomly varying variance (0.1–0.9) is added to embeddings during training.
• ▸Random horizontal flip, applied at the image preprocessing stage before hidden states are extracted.

Inference: select, filter, classify, explain

At inference time the same three-stage logic runs end to end. First, a new RGB image is passed through the pretrained VLM to extract intermediate hidden states, and attention maps identify the patches where the model concentrates most — the regions likely to correspond to wounds, blood, or the limbs around them, rather than processing the full image uniformly. Second, those selected patch embeddings go through the frozen LM head for logit-based filtering: per-patch vocabulary logits and softmax probabilities are computed, and only the patches that strongly activate trauma-relevant concepts are retained and aggregated. Third, the filtered embeddings are fed into the lightweight transformer classifier, which produces a trauma / no-trauma prediction from the CLS token representation.

The interpretability benefit falls out of the design rather than needing to be bolted on afterward. Because the final decision is built from a small, explicitly selected set of patches and layers, the intermediate logits and attention weights that fed into it double as an explanation — you can point to exactly which regions and which emerging semantic concepts drove the classifier's output, instead of treating the prediction as a single opaque number.