Align Once to Explain: Feature Alignment for Scalable B-cosification of Foundational Vision Transformers

AI summary

Overview

Research area: Computer vision interpretability — specifically, converting large pre-trained Vision Transformer (ViT) foundation models into inherently interpretable "B-cos" networks.

Technical level: Intermediate. Readers should know roughly what a Vision Transformer, a linear probe, and an attribution map are; the paper's core idea is explained without requiring deep math.

Scope: The paper introduces ALOE, a one-time, label-free feature-alignment recipe that turns frozen foundation ViT encoders (supervised, self-supervised, and vision–language) into B-cos models that keep most of the teacher's accuracy while producing faithful explanations by construction.

What This Paper Is About

Foundation vision models such as DINOv3 and SigLIP2 are accurate but opaque, and post-hoc explanation methods can be noisy or unfaithful to the model they describe. An alternative is to use inherently interpretable architectures like B-cos Networks, which produce faithful explanations by design — but retraining such models at foundation scale is prohibitively expensive. The paper's goal is to make the existing "B-cosification" conversion work for modern ViTs, which it previously did not: prior work gave only modest gains for ViTs and sometimes failed to match training from scratch.

Key Contributions

1. ALOE (ALign Once to Explain): a universal, one-time, label-free feature-alignment approach that converts a foundational ViT backbone into a B-cos variant, which can then be reused as a drop-in encoder across downstream tasks so that the cost of interpretability is amortized.
2. A design study of alignment choices: the authors systematically compare alignment targets (layer-wise versus pooled; single versus multi-layer), loss families (MSE, cosine, SigLIP, InfoNCE), model scales, and alignment dataset sizes.
3. A recipe that works across pre-training paradigms: ALOE is applied to a fully supervised ViT-B/16, to DINOv3 self-supervised encoders (ViT-S/16, ViT-B/16, ViT-L/16), and to SigLIP2 vision–language encoders (ViT-B/16, ViT-L/16, so400m/16, so400m/14@384), with special tokens such as [CLS] and registers preserved so alignment can be one-to-one.
4. Evidence that interpretability need not cost much accuracy or data: the aligned models retain linear-probing, k-NN, zero-shot, and dense-prediction performance competitive with their teachers, while yielding well-localized explanations.

Main Findings

• Large gains over vanilla B-cosification. For ViT-B/16, average linear-probe accuracy across ten datasets improves by +6.05 p.p. for the supervised model (74.56% → 80.61%), +8.82 p.p. for SigLIP2 (79.96% → 88.78%), and +19.28 p.p. for DINOv3 (70.71% → 89.99%). On ImageNet-1k specifically, the gains are +3.65, +8.82, and +10.50 p.p. respectively.
• Close to the teachers. ALOE models nearly match their teachers on the same ten-dataset average: SigLIP2 at 88.78% versus a teacher at 89.69%; DINOv3 at 89.99% versus 90.20%. On ImageNet-1k linear probing, the SigLIP2-aligned B-ViT-B/16 reaches 83.87% versus the teacher's 84.21%.
• Larger aligned variants scale up. SigLIP2-aligned B-ViT-L/16 and B-ViT-so400m/16 reach 87.08% and 87.83% ImageNet-1k linear-probe accuracy.
• Strong data efficiency. With only unlabeled YFCC15M, accuracy stays nearly flat when alignment data shrinks from 100% to 0.1% (within 0.78 p.p.). At 1% of YFCC15M (150k images) the model reaches 83.79%, within 0.42 p.p. of the SigLIP2 teacher (84.21%) and 8.74 p.p. above the B-cosification reference; the full 15M-image setting reaches 83.87%, a gap of 0.34 p.p. For context, the paper notes SigLIP2 training from scratch required roughly 10B images and 12B alt-texts.
• k-NN and zero-shot transfer survive. With k=20, ALOE gets 81.56% on ImageNet-1k and 83.38% averaged across datasets, versus B-cosification's 71.03% and 50.78%, and the DINOv3 teacher's 82.28% and 83.71% — a +10.53 p.p. ImageNet-1k gain and +32.60 p.p. average gain over B-cosification. Zero-shot ImageNet top-1 with the SigLIP2 ViT-B/16 text encoder is 77.17% for ALOE versus 58.66% for B-cosification and 78.07% for the SigLIP2 teacher.
• Spatially structured features are preserved. On monocular depth with a linear probe (ViT-B/16), ALOE improves relative δ₁ from 0.83 to 0.93 and reduces RMSE from 0.46 to 0.33, approaching the DINOv3 teacher (0.95 / 0.28).
• Better localization. On GridPG, the SigLIP2 ALOE model reaches 84.4% versus 55.5% for the teacher explained with AttnLRP, and ALOE yields higher localization scores than AttnLRP across supervised, SigLIP2, and DINOv3 teachers, with object-centric, class-specific maps.
• Three alignment depths beat one or all. For SigLIP2 ViT-B/16, average linear-probe accuracy goes 80.86 (global only) → 83.06 (+L) → 88.66 (+{2/3, L}) → 88.79 (+{1/3, 2/3, L}) → 88.41 (all layers), so three evenly spaced depths are used.
• Cosine is the most reliable loss. Among MSE, cosine, SigLIP, and InfoNCE tested under identical settings with only the global term, cosine and SigLIP were most consistent across models, while MSE and InfoNCE showed inconsistencies; cosine was adopted for its simplicity.
• Extensions beyond classification. ALOE-aligned SigLIP2 encoders produce zero-shot, model-inherent explanations when combined with the original text encoder, and in preliminary experiments inside a LLaVA-style Gemma-2-9B multimodal model they enable token-level visual grounding (with AttnLRP still used to propagate relevance through the language model).

Methodology in Plain English

The starting point is a frozen "teacher" model — any pre-trained ViT, whether supervised, self-supervised, or vision–language. The authors build a "student" by applying the existing B-cosification conversion: linear and MLP layers become bias-free B-cos transforms (with B=2), all biases are removed including in normalization, and image inputs are expanded from 3 to 6 channels as (r, g, b, 1−r, 1−g, 1−b), with a matching transform on the first-layer weights so explanations can be shown in color. Self-attention and positional embeddings are left untouched, and the [CLS] and register token pathways are kept identical to the teacher's so tokens can be matched one-to-one. The patch embedding is replaced with a B-cos convolutional stem of the same stride so the token grid still lines up.

The student is then trained on unlabeled images (CC3M, CC12M, or YFCC15M; main results use YFCC15M) to reproduce the teacher's internal features. The objective has two parts: a global term that makes the final image embedding of student and teacher point in the same direction (cosine similarity), and a layer-wise term that does the same for individual tokens at three evenly spaced depths (⌊L/3⌋, ⌊2L/3⌋, and L). Both weights are set to 1. Which tokens to align depends on the teacher: [CLS] for the supervised model, [CLS] plus registers plus image tokens for DINOv3, and attention-pooled output versus image tokens for SigLIP2. No projection heads are needed because the student mirrors the teacher's width.

Training uses AdamW with a cosine learning-rate schedule, mixed precision, gradient-norm clipping at 1.0, no weight decay, batch size 1024, and a sweep over learning rates of 3e-3, 1e-3, and 5e-4, with early stopping on a held-out 30k-image split. Random resized crops and horizontal flips are used as augmentation. After this single alignment, the backbone is frozen and evaluated with a linear head (no augmentation during probing), a weighted k-NN classifier at k=20, zero-shot classification using the original SigLIP2 text encoder with the 80 ImageNet prompt templates, and a linear depth head on NYUv2. Explanations come directly from the model's dynamic linear weights W(x) and are scored with GridPG and pixel-deletion against post-hoc methods including AttnLRP, Integrated Gradients, LeGrad, CheferCAM, and LIME.

Why This Matters

The work suggests that trustworthy vision backbones do not require billion-scale retraining: a single alignment pass over a modest unlabeled image set is enough to recover most of a foundation model's generalization while gaining explanations that are faithful by construction rather than post-hoc guesses. This directly addresses a bottleneck in deploying opaque foundation models in settings where decisions must be justified.

• Safety-critical perception: medical imaging, autonomous driving, and industrial inspection, where a practitioner needs to see which image regions drove a prediction.
• Regulated or audited decisions: domains where explanation is a compliance requirement, and where post-hoc attributions are considered insufficiently faithful.
• Debugging and dataset auditing: model-inherent, class-specific maps can reveal spurious correlations or label errors in training data.
• Explainable multimodal assistants: the paper's preliminary LLaVA-style experiment points to vision encoders inside multimodal LLMs that can localize the visual evidence behind individual generated tokens.

Industry relevance: the recipe is deliberately cheap and general — it works across supervised, self-supervised, and vision–language pre-training, scales from ViT-S/16 to ViT-L/16 and so400m variants, and requires no labels. That makes it practical for organizations that already own a strong foundation checkpoint and want interpretability as a drop-in upgrade rather than a retraining project.

Future Directions

• End-to-end interpretable multimodal LLMs. The current MLLM demonstration still relies on AttnLRP to propagate relevance through the language model; the authors explicitly defer a fuller treatment to future work.
• Closing the remaining teacher gap. Small gaps persist (for example 83.87% versus 84.21% on ImageNet-1k for SigLIP2 ViT-B/16, and 88.78% versus 89.69% on the ten-dataset average), leaving room for better alignment objectives or schedules.
• Extending beyond the tested settings. The alignment study covers cosine, MSE, SigLIP, and InfoNCE objectives, three depth schedules, and three unlabeled datasets; other teachers, architectures, and tasks (the paper reports a preliminary MLLM result and dense-prediction probes but notes more dense results in the appendix) remain open.
• Making the conversion fully faithful at extreme resolutions. The paper documents a concrete wrinkle: for SigLIP2-so400m/14@384, replacing the patch embedding with a stride-16 convolution changes patch size from 14 to 16, so the aligned model is evaluated at 432×432 and denoted B-ViT-so400m/16@432 — a reminder that architecture-preserving conversion needs care at non-standard configurations.

Target Audience

Researchers and engineers working on explainable AI for vision, interpretability of foundation models, and efficient model conversion or distillation. It is also useful for practitioners who already have a strong ViT checkpoint and want faithful explanations without retraining from scratch, and for students interested in inherently interpretable architectures as an alternative to post-hoc attribution. Readers without background in Vision Transformers or attribution methods will need to consult the background section and cited B-cos work first.

Authors’ abstract

Foundational vision models have become the de facto standard for many vision tasks due to their strong performance. However, they are notoriously opaque and remain hard to interpret. We present ALOE (ALign Once to Explain), a one-time, label-free feature alignment approach that efficiently converts foundational vision models into inherently interpretable B-cos variants. Once aligned, the B-cos backbone is used as a drop-in replacement across several downstream tasks-amortizing the cost of interpretability. ALOE is robust across pre-training paradigms (supervised, self-supervised, vision-language) and can preserve teacher-level generalization with orders of magnitude less data than training from scratch. On classification, it strongly outperforms vanilla B-cosification (e.g., +10.5 p.p. top-1 on ImageNet for DINOv3 ViT-B/16), retains strong linear probing, k-NN, and zero-shot transfer competitive with foundational backbones (DINOv3, SigLIP2) across diverse downstream datasets. It also preserves spatially structured features useful for dense prediction, while yielding well-localized and highly human-interpretable explanations by design.