PaperlayerView exact revision on arXiv

Source-linked AI summary

AM-RADIO: Agglomerative Vision Foundation Model -- Reduce All Domains Into One

Mike Ranzinger, Greg Heinrich, Jan Kautz, Pavlo Molchanov

arXiv:2312.06709v5cs.CV

TL;DR

Visual foundation models provide complementary capabilities but are trained separately, motivating a unified model that can cover their distinct strengths. AM-RADIO uses multi-teacher feature distillation to train such a model, and E-RADIO supplies a more efficient architecture. The resulting students combine teacher capabilities, outperform teachers on many benchmarks, and support downstream applications.

  • Problem

    CLIP, DINOv2, and SAM offer complementary strengths across vision-language, dense understanding, and segmentation tasks, but their capabilities remain distributed across separate models.

  • Method

    AM-RADIO trains a vision foundation model from scratch by matching heterogeneous teachers independently through teacher-specific projection heads.

  • Results

    AM-RADIO outperforms its teachers on 6 of 9 tasks, while E-RADIO improves efficiency and is at least 6x faster than teacher models at matched resolution.

  • Takeaways & Limitations

    A single distilled student can combine zero-shot, dense, and segmentation capabilities while serving as a drop-in teacher replacement or visual encoder for downstream applications.

  • Takeaways & Limitations

    The approach assumes its selected teacher models represent a broad range of internet images, including images from datasets such as ImageNet, LAION-400M, and DataComp-1B.

Abstract

from arXiv · show

A handful of visual foundation models (VFMs) have recently emerged as the backbones for numerous downstream tasks. VFMs like CLIP, DINOv2, SAM are trained with distinct objectives, exhibiting unique characteristics for various downstream tasks. We find that despite their conceptual differences, these models can be effectively merged into a unified model through multi-teacher distillation. We name this approach AM-RADIO (Agglomerative Model -- Reduce All Domains Into One). This integrative approach not only surpasses the performance of individual teacher models but also amalgamates their distinctive features, such as zero-shot vision-language comprehension, detailed pixel-level understanding, and open vocabulary segmentation capabilities. In pursuit of the most hardware-efficient backbone, we evaluated numerous architectures in our multi-teacher distillation pipeline using the same training recipe. This led to the development of a novel architecture (E-RADIO) that exceeds the performance of its predecessors and is at least 7x faster than the teacher models. Our comprehensive benchmarking process covers downstream tasks including ImageNet classification, ADE20k semantic segmentation, COCO object detection and LLaVa-1.5 framework. Code: https://github.com/NVlabs/RADIO

NVIDIA

AM-RADIO distills multiple pretrained vision foundation models into one RADIO model that combines their distinctive capabilities. Figure 1 presents the framework, feature visualizations, and downstream benchmarks.

  • AM-RADIO distills CLIP, DINOv2, and SAM into a single model called RADIO.
  • RADIO combines zero-shot classification and open set instance segmentation capabilities in one forward pass.
  • Figure 1 shows model feature visualizations, the distillation framework, and benchmarks across classification, segmentation, and vision-language modeling.

1. Introduction

AM-RADIO distills multiple visual foundation models into one model that combines their complementary capabilities and can surpass individual teachers. The framework also evaluates efficient student architectures, leading to E-RADIO for faster deployment.

  • The benchmark suite includes ImageNet classification, semantic segmentation, LLaVA vision-language tasks, and COCO instance segmentation.
  • The framework combines capabilities associated with CLIP, DINOv2, and SAM, including language grounding, dense representations, and segmentation.
  • AM-RADIO distills multiple foundation models into a single student that agglomerates their distinct attributes.
  • AM-RADIO students can outperform their teachers on representative benchmarks and support downstream applications such as CLIP zero-shot tasks and Segment Anything.
  • More than 10 architectures were evaluated under the same training recipe, revealing that CNN-like models are faster but struggle to distill ViT-based foundation models.
  • E-RADIO is a hybrid architecture developed to improve the efficiency-quality trade-off, achieving at least 6x faster inference than teacher models at matched resolution.

2. Related Work

AM-RADIO builds on knowledge distillation and multi-teacher distillation, adapting feature matching to heterogeneous vision foundation models. Its method maps a unified student representation to each teacher independently.

  • Knowledge distillation trains a student using outputs or feature representations from an already-trained teacher model.
  • Multi-teacher distillation jointly transfers knowledge from multiple teachers whose domains or capabilities differ.
  • Because the teachers cover heterogeneous domains, AM-RADIO matches each teacher independently through teacher-specific projection heads rather than a shared label.

3. Knowledge Agglomeration

AM-RADIO trains one vision foundation model from scratch by matching heterogeneous teachers through teacher-specific adaptor heads. The framework combines teacher-specific summary and spatial features, with spatial supervision improving downstream dense-task quality.

  • Knowledge Agglomeration: AM-RADIO distills multiple foundation models simultaneously so one student can agglomerate their distinct attributes.The framework is designed for heterogeneous teachers with different domains and resolutions.
  • Adaptor Heads: Teacher-specific 2-layer MLP heads map the unified student representation to each teacher’s summary and spatial feature dimensions.Each head uses LayerNorm and GELU between the input and output projections.
  • Feature Matching: The student matches teacher vision-encoder features because per-image ground-truth labels are unavailable for every teacher.The loss distinguishes summary vectors from spatial feature vectors.
  • Loss Formulation: Adding spatial-feature supervision is always beneficial, and matching both CLIP and DINOv2 spatial features gives the best ablation results.The ablation uses a ViT-L/14 student trained on LAION-400M and retains SAM for its spatial features.
  • Loss Formulation: Spatial supervision improves downstream dense tasks and the model’s holistic quality, while cosine distance outperforms L1, MSE, and Smooth-L1 for feature matching.The spatial loss combines cosine similarity with smooth L1 to preserve both direction and vector magnitude.
  • Loss Formulation: AdaLoss performs best on ImageNet tasks but worst on ADE20K, so the paper uses naive manual loss balancing thereafter.The authors report little advantage for more elaborate balancing schemes.

4. Implementation Details

The implementation addresses heterogeneous teacher resolutions and hardware costs through staged multi-resolution training and efficient student-architecture design. CPE enables resolution generalization while reducing position-embedding artifacts, and the study compares feature summarization choices.

  • Training Setup: Training uses AdamW, batch size 1024, cosine annealing, learning rate 0.001, and 600k steps totaling 614M examples.The recipe is applied across student models for controlled architecture comparisons.
  • Multi-scale Teachers: Two-stage training uses CLIP and DINOv2 for 300k steps at 256px, then adds higher-resolution CLIP and SAM supervision for 300k steps.The second stage uses CLIP+DINOv2 at 432px and SAM at 1024px.
  • Multi-scale Teachers: The student uses 432² inputs for patch-14 CLIP and DINOv2 teachers, while DINOv2 outputs are upsampled from 224px without degrading results.SAM features are matched using the expected 1024² resolution.
  • Position Embeddings: CPE lets a ViT-L/14 student trained at 224² generalize to 518², improving over non-CPE and surpassing DINOv2-g.The comparison uses otherwise identical student models.
  • Position Embeddings: CPE regularizes position-embedding space and eliminates high-frequency artifacts associated with regular position-embedding patterns.Figure 3 visualizes the position embeddings with PCA.
  • Efficient Architecture: Efficient training reformulates high-resolution processing as sampled flexible window sizes, reducing SAM-teacher training cost and enabling inference throughput scaling.The approach is also evaluated as a replacement for SAM’s encoder.
  • Feature Summarization: Student feature summarization compares a separate CLS token with average pooling over patch tokens.Both options are evaluated in Table 6.

5. Results

AM-RADIO evaluates distilled student backbones for efficiency, teacher comparison, dense vision, multimodal reasoning, and 3D awareness. E-RADIO is introduced as a hybrid architecture after existing efficient backbones expose speed–accuracy and spatial-resolution trade-offs.

  • Backbone efficiency: The study evaluates backbone efficiency with unchanged distillation hyperparameters across multiple architectures and compares them with CLIP, DINOv2, and SAM teachers.The evaluated architecture set includes EfficientNetV2, ResNetv2, RegNetY, FasterViT, EfficientViT, ConvNext, NFNet, SwinV2, MaxViT, PoolformerV2, and MViTV2.
  • Backbone efficiency: CNN-like models are significantly faster than ViTs, while ViTs are more accurate; existing efficient backbones also lag on dense ADE20k segmentation.Their final feature maps use a spatial reduction factor of 32, limiting fine-grained spatial information.
  • Efficiency results: E-RADIO delivers results 10 times faster and with a 20% improvement over teacher models under the reported Figure 4 evaluation.The study uses 224px resolution and a window size of 7 for E-RADIO in this comparison.
  • E-RADIO: E-RADIO combines convolutional and transformer stages with local and global windowed attention to preserve fine-grained spatial information efficiently.Its convolutional stages use YOLOv8 C2f blocks, while transformer stages use windowed attention and multi-resolution attention.
  • Comparison with teachers: The full ViT-H/16 model is as fast as the teachers and outperforms them on 6 of 9 tasks, while ERADIO-L is significantly faster than all ViT models.ERADIO-L also strongly outperforms MetaCLIP on most metrics at matched throughput and enables zero-shot capability absent in DINOv2 and SAM.
  • Downstream evaluation: The models are evaluated through linear probing on ADE20k and PascalVOC2012, LLaVA 1.5 tasks, SAM encoder replacement, and NAVI 3D probing.The NAVI probing covers monocular depth, surface normals, and multi-view correspondence.

6. Conclusion and Key Insights

AM-RADIO unifies complementary capabilities from CLIP, DINOv2, and SAM through distillation while addressing efficiency and resolution behavior. The framework is presented as a route to models that can combine broad downstream functionality with improved efficiency.

  • Key insights: CLIP, DINOv2, and SAM provide complementary language grounding, dense correspondence, and detailed segmentation capabilities, but each has capability gaps.AM-RADIO targets a single model that combines these properties.
  • Key insights: Distillation can unite the teachers’ properties in one model that often outperforms individual teachers on downstream tasks.The paper reports an 18% relative ADE20K improvement from full feature distillation.
  • Resolution behavior: RADIO’s DINOv2 head exhibits a sudden MSE error jump at 720px, corresponding to a complete change in image color space.The comparison uses DINOv2 features at 518px as the reference after bilinear interpolation.
  • Efficiency: E-RADIO uses a hybrid CNN-Transformer design to improve the efficiency–quality Pareto frontier of vision foundation model backbones.The conclusion frames the architecture as addressing the poor scaling of existing approaches to VFM settings.

Supplementary Material

The supplementary material presents the high-level E-RADIO architecture and its multi-resolution attention component.

  • Architecture: The E-RADIO network contains a stem, two YOLOv8 convolutional blocks, and two transformer blocks with multi-resolution windowed self-attention.These components define the architecture’s main stages.
  • Attention: Multi-resolution attention is identified as a dedicated E-RADIO component.The supplied supplementary caption names the component without reporting a comparison or outcome.

A. E-RADIO architecture details

E-RADIO is a hybrid CNN-Transformer backbone that combines convolutional early stages with transformer-based later stages and multi-resolution attention.

  • Architecture: E-RADIO begins with two YOLOv8 C2f convolutional stages and uses two later transformer stages with windowed and multi-resolution attention.Every stage except the last is followed by a strided-convolution downsampling block with batch normalization.

A.1. Multi-Resolution Attention

E-RADIO uses multi-resolution and windowed attention to reduce high-resolution transformer costs while retaining global communication. Its architecture combines convolutional and transformer stages and supports flexible inference window sizes.

  • Attention design: Multi-resolution attention addresses high-resolution cost by avoiding SWIN’s window-shifting overhead while retaining communication across resolutions.The design is inspired by EdgeViT’s Local-Global-Local attention.
  • Architecture: The architecture uses two convolutional and two transformer stages, with multi-resolution attention followed by windowed attention.The third stage uses multi-resolution attention, while the fourth uses windowed attention.
  • Input flexibility: DINOv2 and RADIO support arbitrary resolutions and aspect ratios, unlike most compared models that handle fixed square inputs.The distinction is illustrated through PCA visualizations of spatial features.
  • Attention implementation: Alternating windowed and global attention is implemented by reordering patches so contiguous chunks correspond to windows.The final block always uses global attention.
  • Flexible windows: Window sizes are sampled during training, reducing SAM-teacher training cost and providing an inference throughput scaling mechanism.The same reformulation replaces a fixed architecture choice with training augmentation.

E.1. Uncertainty

The uncertainty formulation learns how to balance multiple loss terms during training. It uses predicted uncertainty parameters, stabilizing transformations, and an alternative AdaLoss weighting scheme.

  • Parameterization: The student predicts b_k = log σ_k^2 instead of σ_k to avoid division by zero and permit unconstrained scalar regression.The σ_k values are predicted by the student.
  • Uncertainty weighting: Learned uncertainty weights replace manual λ scalars and add a penalty for large uncertainties to stabilize training.The resulting objective sums weighted losses together with an uncertainty-dependent regularization term.
  • Teacher-specific balancing: The formulation predicts balance parameters separately for each teacher and for summary-weight or feature-vector-weight losses.Teacher-specific terms are indexed by i and loss categories by m ∈ {s,v}.
  • Stabilization: The uncertainty weighting function uses ψ(x), the softplus activation, to keep loss weights positive and improve stability for large weights.For negative inputs corresponding to weights above 1, the derivative tends toward zero.
  • AdaLoss: AdaLoss was also explored as an alternative weighting formulation based on the inverse expected loss.Its objective retains separate weights for teacher and loss categories.

F. Visual Question Answering Samples

The LLaVA experiments provide qualitative visual question-answering examples using several image encoders on GQA and TextVQA. Figures indicate correctness by coloring answers green for matches and pink for mismatches.

  • VQA examples: Figures 9–13 show sample questions and answers from GQA and TextVQA in LLaVA setups using the paper’s vision encoders.The examples compare outputs obtained with various image encoders.
  • Attention visualization: Attention maps trace how LLaVA attends to RADIO visual tokens across language-model layers for a GQA helmet-color question.Attention scores are averaged over heads and overlaid on the input image.
  • Figure encoding: Answer color encodes correctness against ground truth: green indicates a match, while pink indicates an incorrect answer.This encoding is stated in the captions for the GQA and TextVQA figures.
  • Text reading: TextVQA examples include reading text such as “Dealing in Desire,” “Kimberly Kay Hoang,” and “Silicon Graphics” from images.The sample answers show variants of these textual readings.
Loading 2312.06709v5…