Relational Representation Distillation

In plain words

A small student network learns from a large teacher network. The usual way is to make the student's final answers match the teacher's. That copies the answers but not how the teacher arranges its internal picture of the data: which examples it treats as alike and which as different.

Contrastive methods copy that arrangement by pushing every example away from all the others. That can hurt when two examples belong to the same class, because they get pushed apart too.

This paper keeps the relationships instead of forcing separation. For each example, the teacher and the student each rate how similar it is to a memory of recent teacher examples, and the student is trained to match the teacher's ratings. The teacher's ratings are made sharp, so its strongest matches stand out. The student's ratings are softer, so it keeps its weaker similarities too. The authors report that this beats the contrastive baseline on most classification tests and gives gains on detection, all from their own runs.

TL;DR

• The student matches the teacher's distribution over similarities to a FIFO memory of teacher features, using cross-entropy §3.2 eq. 8.
• The teacher uses a sharp temperature (τ_t = 0.02), the student a softer one (τ_s = 0.1) §4.1. The authors call this an information bottleneck §3.3.
• With τ_t → 0 the loss becomes InfoNCE §3.2.1 eq. 9.
• Author-reported CIFAR-100 results: RRD 75.85% vs CRD 75.48% vs KD 74.92% for WRN-40-2 → WRN-16-2 tab. 1.
• ImageNet and COCO results come from a single run each tab. 2, tab. 3.
• The official code is MIT-licensed. It differs from the paper in several places (see Paper vs code).

Problem

• Logit-only distillation with KL divergence aligns class probabilities but ignores the structure among the teacher's internal representations §1.
• Feature-matching methods extend distillation to intermediate layers. Contrastive distillation (CRD) enforces instance-level discrimination, which the authors say causes a "class collision": semantically similar samples get pushed apart under uniform negative sampling §1 §2.
• The authors' experiments use five benchmarks (CIFAR-100, ImageNet ILSVRC-2012, STL-10, Tiny ImageNet, MS-COCO) and thirteen teacher–student pairs §4.1.

Method

The method keeps a FIFO memory of teacher embeddings. Both networks embed their features into a shared space, and each batch's similarities to the memory give a distribution per network. The student is trained to match the teacher's distribution. The paper states the loss is the cross-entropy of the two distributions, with the teacher's distribution fixed §3.2.

1. Pass each image through the frozen teacher and the student. Use the features before the classifier §3.2 train_student.py.
2. Project both feature vectors to d = 128 dimensions and L2-normalise them. The paper uses a two-layer MLP head; the code uses one linear layer (see Paper vs code) §3.2 §4.1.
3. Push the current batch of teacher embeddings into a first-in-first-out memory of K = 16,384 entries, dropping the oldest entries §3.2 §4.4.2.
4. For each embedding, take the inner products with all memory entries. Divide the teacher's by τ_t = 0.02 and the student's by τ_s = 0.1, then apply a softmax over the memory §3.2 eq. 2–eq. 5.
5. Let the teacher's softmax be the target and compute the cross-entropy with the student's log-softmax. Average over the batch §3.2 eq. 8.
6. Add the supervised cross-entropy and the standard KL distillation term on logits, with τ = 4, weighted by λ. Add the relational term with weight β §3.2.3 eq. 10 eq. 11.

Teacher distribution over memory entries, with the paper's notation:

p_j^T = \frac{\exp(\mathbf{z}_i^T \cdot \mathbf{m}_j / \tau_t)}{\sum_{k=1}^{K+1} \exp(\mathbf{z}_i^T \cdot \mathbf{m}_k / \tau_t)} §3.2 eq. 3

Relational loss:

\mathcal{L}_{RRD} = - \sum_{j=1}^{K+1} p_j^T \log p_j^S §3.2 eq. 8

Full objective, with λ on the logit KL term and β on the relational term:

\hat{\theta}_S = \arg\min_{\theta_S} \sum_i \left( \mathcal{L}_{sup} + \lambda \, \mathcal{L}_{KL} + \beta \, \mathcal{L}_{RRD} \right) §3.2.3 eq. 11

Training step in our own notation. The paper defines the memory as K+1 entries (see Pitfalls); the code uses K:

# One RRD training step with a FIFO memory  [§3.2] [eq. 8]
# Input:  images (B, 3, H, W); labels (B,)
# State:  queue (K, d), K = 16384, d = 128  [§4.1] [§4.4.2]
# Output: total loss (scalar)

def rrd_loss(images, labels, queue):
    # Frozen teacher gives pooled features and logits  [§3.2]
    teacher_feats, teacher_logits = teacher(images)
    # Student gives pooled features and logits  [§3.2]
    student_feats, logits = student(images)
    # Project to d dimensions and L2-normalise  [§3.2] [§4.1]
    z_teacher = normalize(teacher_head(teacher_feats))  # (B, d)
    z_student = normalize(student_head(student_feats))  # (B, d)
    # Write the current teacher batch into the FIFO memory  [§3.2]
    queue = push_fifo(queue, z_teacher)                 # (K, d)
    # Similarities of each embedding to every memory entry  [§3.2]
    sim_teacher = z_teacher @ queue.T                   # (B, K)
    sim_student = z_student @ queue.T                   # (B, K)
    # Sharp teacher distribution, fixed as the target  [§3.2] [eq. 3]
    target = softmax(sim_teacher.detach() / 0.02)
    # Soft student log-distribution  [§3.2] [eq. 5]
    log_pred = log_softmax(sim_student / 0.1)
    # Cross-entropy between target and prediction  [§3.2] [eq. 8]
    loss_rrd = -(target * log_pred).sum(dim=1).mean()
    # Supervised cross-entropy, weight 1  [eq. 11]
    loss_sup = cross_entropy(logits, labels)
    # Logit KL distillation at temperature 4, weight lambda  [§3.2.3] [eq. 10]
    loss_kl = logit_kl(logits, teacher_logits, temperature=4)
    # Combined objective  [§3.2.3] [eq. 11]
    return loss_sup + lam * loss_kl + beta * loss_rrd

Properties the paper derives:

• As τ_t → 0 the teacher distribution becomes one-hot on the teacher's own entry. The loss then equals InfoNCE with the memory as negatives §3.2.1 eq. 9.
• Writing the memory as K+1 entries includes the current teacher embedding explicitly, so the teacher's own entry has the largest teacher score §3.2.
• The paper bounds the mutual information between teacher and student by the entropy gap between the student and teacher distributions, which it calls the information bottleneck §3.3 eq. 12. This bound is an interpretation in the paper; it is not used to set any value.

Hyperparameters

Results

Author-reported. Setups: CIFAR-100 with thirteen teacher–student pairs, RRD averaged over five runs tab. 1; ImageNet, single run tab. 2; COCO with Faster R-CNN, single run, baseline numbers taken from the DKD paper tab. 3. Baseline numbers in tables 1 and 2 come from the CRD paper. The bold and underline marks in the source are omitted here.

CIFAR-100 top-1 accuracy (%) tab. 1. Pairs 1–7 are same-architecture, 8–13 different-architecture. The source has a different column order; this table is transposed.

Abbreviations in the source: W = WideResNet, rn = ResNet (CIFAR), RN = ResNet (ImageNet-style), MN = MobileNet, SN = ShuffleNet, VGG as named tab. 1.

ImageNet ILSVRC-2012 student top-1 (%) tab. 2.

COCO object detection, Faster R-CNN, AP / AP50 / AP75 on val2017 tab. 3. Baseline numbers from the DKD paper.

Cost on CIFAR-100, ResNet-32x4 → ResNet-8x4, MDistiller, one RTX 6000 tab. 4.

The parameter count does not match the head described in the text (see Pitfalls).

Transfer: WRN-16-2 student from WRN-40-2 teacher, frozen features tab. 5.

Correlation matrix alignment, CIFAR-100, WRN-40-2 teacher, WRN-40-1 student tab. 6.

Semantic structure, CIFAR-100 student embeddings tab. 7.

Logit correlation difference, WRN-40-1 vs WRN-40-2 fig. 2. Mean and max values in the figure caption, lower is better: Vanilla 0.24 / 1.66; KD 0.09 / 0.49; CRD 0.23 / 1.56; CRD+KD 0.10 / 0.57; RRD 0.18 / 0.99; RRD+KD 0.07 / 0.55.

Ablations, CIFAR-100, WRN-40-2 teacher, WRN-16-2 student, three runs each §4.4. The figure values are not given as numbers in the text. Best temperatures are τ_s = 0.1 and τ_t = 0.02; accuracy falls as τ_t grows and falls sharply once τ_t > τ_s; τ_s = 1.0 underperforms §4.4.1 fig. 3. Accuracy grows with memory size up to about K = 16,384 §4.4.2 fig. 3. β between 0.5 and 1.5 works best §4.4.3 fig. 3.

Paper vs code

Official repository: https://github.com/giakoumoglou/rrd, linked on the title page of the paper. The owner is the first author. Commit 83cefdab08b4cb126c359abe41b79e47794d59b4, MIT licence ("Copyright (c) 2024 Nikolaos Giakoumoglou"). The repository has CIFAR-100 and ImageNet code only. There is no detection code, so the COCO numbers could not be checked.

• Projection heads. Paper: two-layer MLP, 512 hidden, 128 output §4.1. Code: a single nn.Linear per branch to 128 dimensions RRD.py RRD.py.
• Teacher head is trained. Paper: teacher frozen §3.2. Code: the teacher's projection head is added to the optimiser's parameter list train_student.py train_student.py. Only the teacher backbone is frozen.
• Memory count. Paper: K+1 entries, with the current teacher embedding appended as m_{K+1} §3.2. Code: the current teacher batch is written into the FIFO queue of length nce_k before the similarities are computed, so the batch replaces the oldest entries and the softmax runs over K entries RRD.py RRD.py.
• Queue initialisation. Code normalises the initial random queue along dimension 0, so its rows are not unit vectors at start RRD.py. The paper does not mention this. Real teacher embeddings overwrite the initial values after K / B batches, which is 256 batches for K = 16,384 and B = 64.
• Temperatures. Paper: τ_t = 0.02, τ_s = 0.1 §4.1. Code defaults match train_student.py. The class docstring in RRD_FifoMemory_Loss lists them the other way round RRD.py; the argument parser is correct.
• Loss weights on CIFAR-100. Paper: λ = 0.9, β = 1.5 §4.1. The CIFAR README commands use -a 0 -b 1 for RRD and -a 1 -b 1 for RRD+KD README.md README.md. The plot scripts' result folder names use a_0.9_b_1.0 plot_corr.py. The README, the plot scripts and the paper disagree on β, and the README's RRD+KD uses λ = 1.
• Supervised term in the RRD-only run. The paper's footnote defines RRD as supervised learning plus the relational loss §4. The ImageNet README RRD command sets -c 0 -d 0, which removes the classification and KD terms README.md. Code default --cls is 1 train_student.py.
• Target gradient. Paper text says the student distribution is "used as a target" and its gradient is clipped §3.2. The loss formula and the code treat the teacher distribution as the detached target RRD.py. The text and the formula disagree; the code follows the formula.
• Feature layer. Paper: "intermediate representations" §4.1. Code: the pooled last-layer features, feat[-1], which is the input to the classifier train_student.py resnet.py.
• Total objective. Paper Eq. 11 has supervised weight 1 eq. 11. Code: loss = gamma * loss_cls + alpha * loss_div + beta * loss_kd, with gamma default 1 train_student.py train_student.py.
• ImageNet distiller does not run as written. imagenet/distillers/RRD.py calls nn.functional.normalize.normalize, nn.functional.normalize.softmax and nn.functional.normalize.log_softmax RRD.py RRD.py. By reading, these raise AttributeError, since normalize is a function. We did not run it. The CIFAR distiller uses the correct torch.nn.functional calls RRD.py RRD.py.
• Momentum memory. The code has a second, momentum-based memory variant selected by --memory_type momentum train_student.py RRD.py. The paper describes only the FIFO memory §3.2. This card covers the FIFO variant.
• Optimiser values. The paper defers to the CRD protocol §4.1. The values in the Hyperparameters table are the code's defaults train_student.py.

No discrepancies were checked for the detection or transfer experiments, because the repository has no code for them.

Pitfalls

• K or K+1 memory entries. The paper's loss sums over K+1 entries, with the current teacher embedding as an extra entry §3.2 eq. 8. The code sums over K entries that already include the current batch RRD.py. Use the code's behaviour if you want to reproduce the reported numbers; the difference is one entry in a 16,384-entry memory when the batch is written in.
• Memory size in the text. The paper says storing 128-d features "requires only ~600MB on ImageNet" §3.2. That matches a bank with one entry per ImageNet training image: 1,281,167 × 128 × 4 bytes ≈ 656 MB (our calculation). A FIFO queue of K = 16,384 entries is about 8 MB at the same precision (our calculation). Size the memory for K, not for the dataset.
• Parameter count in the cost table. The table lists 2,162,944 extra parameters tab. 4. The described heads on the 256-d features of ResNet-32x4 and ResNet-8x4 give 2 × (256 × 128 + 128) = 65,792 parameters with one linear layer each (our calculation from resnet.py). A 512-hidden MLP on both sides would give about 394k. The table's figure cannot be reproduced from the description.
• Claim of consistent CRD beating. The text says RRD "consistently outperforms" CRD on CIFAR-100 §4.2.1. Table 1 shows RRD below CRD on pair 8 (69.61 vs 69.73) and pair 13 (75.98 vs 76.05), and a tie on pair 10 (74.30 vs 74.30) tab. 1.
• Own teacher embedding in the memory. The code writes the current teacher batch into the queue before it computes the similarities, so each query's own teacher embedding is among the memory entries RRD.py. The paper says this "ensures a clear peak" in the teacher distribution §3.2. The card has no measurement of how much this matters.
• Sensitivity to temperature. The paper finds τ_t > τ_s harms accuracy sharply §4.4.1. Keep τ_t < τ_s.
• Weight decay on projection heads. The projection heads are optimised with the backbone under the same weight decay train_student.py train_student.py. The paper does not say whether the heads have weight decay.
• Batch divisibility. The FIFO write asserts that K is divisible by the batch size RRD.py. With K = 16,384, batch sizes that divide 16,384 (for example 64 and 256) pass; other sizes fail the assertion.
• Unstated values. Not stated in the paper: the number of epochs, learning rate, batch size and optimiser for CIFAR-100 and ImageNet beyond "CRD protocol"; the detection training recipe; the memory warm-up; and how the 13 teacher–student pairs were trained.
• Comparisons across protocols. Baseline numbers come from CRD tab. 1 tab. 2 and DKD tab. 3. Our five-run and single-run numbers may not use identical protocols. Compare results within a table only.
• Single runs. ImageNet and COCO numbers come from one run each tab. 2 tab. 3. The paper reports no variance for these tables, so small differences between methods there cannot be judged from the card.

Lineage

• Builds on: logit distillation with KL divergence [hinton2015distilling] §3.1; contrastive distillation CRD [tian2022crd] §1 §4.1; InfoNCE as the limiting case [oord2019cpc] §3.2.1; memory banks of past features, which the paper adapts to a FIFO queue §3.2.
• Compared against: KD, FitNet, AT, SP, CC, RKD, FSP, OFD, CRD, CRD+KD, ReviewKD and DKD tab. 1 tab. 3 tab. 4; ImageNet comparison uses KD, AT, SP, CC, RKD and CRD tab. 2.
• Superseded by: not stated.