Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Preserving Angles Improves Feature Distillation

Evelyn J. Mannix, Liam Hodgkinson, Howard Bondell

TL;DR

Cosine-similarity Preserving Compression (CosPress) is introduced, a feature distillation technique that learns a mapping to compress the latent space of the teacher model into the smaller latent space of the student, by preserving the cosine similarities between image embeddings.

Abstract

Knowledge distillation methods compress models by training a student network using the classification outputs of a high quality teacher model, but can fail to effectively transfer the properties of computer vision foundation models from the teacher to the student. While it has been recently shown that feature distillation$\unicode{x2013}$where a teacher model's output features are replicated instead$\unicode{x2013}$can reproduce performance for foundation models across numerous downstream tasks, they fall short in matching critical properties such as robustness and out-of-distribution (OOD) detection performance. This paper overcomes this shortcoming by introducing Cosine-similarity Preserving Compression (CosPress), a feature distillation technique that learns a mapping to compress the latent space of the teacher model into the smaller latent space of the student, by preserving the cosine similarities between image embeddings. This enables direct optimisation of the student network and produces a more faithful reproduction of the teacher's properties. It is shown that distillation with CosPress on a variety of datasets, including ImageNet, produces more accurate models with greater performance on generalisability, robustness and OOD detection benchmarks, and that this technique provides a competitive pathway for training highly performant lightweight models on small datasets. Code is available at github.com/emannix/cospress.

Preserving Angles Improves Feature Distillation

TL;DR

Cosine-similarity Preserving Compression (CosPress) is introduced, a feature distillation technique that learns a mapping to compress the latent space of the teacher model into the smaller latent space of the student, by preserving the cosine similarities between image embeddings.

Abstract

Knowledge distillation methods compress models by training a student network using the classification outputs of a high quality teacher model, but can fail to effectively transfer the properties of computer vision foundation models from the teacher to the student. While it has been recently shown that feature distillationwhere a teacher model's output features are replicated insteadcan reproduce performance for foundation models across numerous downstream tasks, they fall short in matching critical properties such as robustness and out-of-distribution (OOD) detection performance. This paper overcomes this shortcoming by introducing Cosine-similarity Preserving Compression (CosPress), a feature distillation technique that learns a mapping to compress the latent space of the teacher model into the smaller latent space of the student, by preserving the cosine similarities between image embeddings. This enables direct optimisation of the student network and produces a more faithful reproduction of the teacher's properties. It is shown that distillation with CosPress on a variety of datasets, including ImageNet, produces more accurate models with greater performance on generalisability, robustness and OOD detection benchmarks, and that this technique provides a competitive pathway for training highly performant lightweight models on small datasets. Code is available at github.com/emannix/cospress.

Paper Structure

This paper contains 44 sections, 6 theorems, 31 equations, 3 figures, 21 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Foundation models
  4. Knowledge distillation.
  5. Feature distillation.
  6. Dimensionality reduction.
  7. Methods
  8. Notation.
  9. Motivation.
  10. Proteus.
  11. CosPress.
  12. Dimensionality reduction objective.
  13. Teacher head architecture.
  14. Student objective.
  15. Experiments: Feature distillation
  16. ...and 29 more sections

Key Result

Lemma 1

Let $\mathbf{M} \in \mathbb{R}^{m \times d}$ with $m < d$ and $\mathrm{rank}(\mathbf{M}) = m$. Then Moreover, the converse inequality does not generally hold.

Figures (3)

  • Figure 1: Patch features. PCA visualisation of patch features for the DINOv2 ViT-S/14 model, and the distilled ViT-Ti/14 model produced using the CosPress feature distillation approach.
  • Figure 2: Feature distillation frameworks. In Proteus, student heads $g$ are used to map the outputs of the student network $S_\theta$ into the latent space of the teacher $T$, so that a MSE loss can be applied. In CosPress, a teacher head $h$ is trained to compress the teacher $T$ outputs into the student latent space, preserving the cosine similarity of image embeddings and allowing direct optimisation. The Proteus student head does not preserve cosine similarity, even when the projection matrices are forced to be right-orthogonal.
  • Figure 3: Visualising orthogonality. Kernel density estimate plots of the diagonal and non-diagonal elements for the scaled Gram matrices of the linear maps $\mathbf{W}$ in the teacher and student heads, drawn from CosPress and Proteus respectively. The dashed coloured lines represent $\mathbf{W}^\top \mathbf{W}/\alpha$ and the solid lines represent $\mathbf{W} \mathbf{W}^\top/\beta$, where $\alpha, \beta$ are defined as in \ref{['tbl:orthgonal_measurement']}. A perfectly orthogonal matrix $\mathbf{W}$ will have a Gram matrix with density on the black dashed vertical lines.

Theorems & Definitions (9)

  • Definition 1: Approximately Orthogonal Matrix
  • Lemma 1
  • Lemma 2: Johnson–Lindenstrauss; freksen2021introduction
  • Lemma 3: Norm preservation; freksen2021introduction
  • Lemma 4: Angles; magen2007dimensionality
  • Lemma 5: Approximate left and right orthogonality for JL transforms
  • proof
  • Lemma 1
  • proof