--- title: Token Relation Distillation (TRD) url: https://www.emergentmind.com/topics/token-relation-distillation-trd type: topic --- # Token Relation Distillation (TRD) Token Relation Distillation (TRD) introduces a methodology for enhancing knowledge distillation by explicitly transferring the token-level relational structure learned by a powerful teacher network to a typically smaller student model. Departing from conventional distillation approaches that emphasize either logits or instance-level relationships, TRD leverages a Token-level Relationship Graph (TRG) to encapsulate both intra-instance semantic context and cross-instance token-wise similarities. This graph-centric strategy enables the student to emulate fine-grained, higher-order semantic dependencies from the teacher, with demonstrated advantages on balanced and imbalanced classification tasks across CNN and ViT architectures [2306.12442]. ## 1. Motivation and Theoretical Foundations Traditional knowledge distillation (KD) as introduced by Hinton et al. [arXiv:1503.02531] focuses on transferring class probability distributions through softened logits: $$ L_{KD} = \frac{1}{N} \sum_{i=1}^N \mathrm{KL}(p_i^S(\tau) \| p_i^T(\tau)), $$ yielding the standard loss formulation: $$ L_\text{logit} = L_\text{CE}(p^S, y) + \lambda L_{KD}. $$ Extensions incorporating feature-based alignment or instance-level relational graphs (e.g., RKD [Park et al., 2019], IRG [Liu et al., 2019]) improved transfer, but failed to explicitly model intra-image structure and higher-order semantic patterns. The central hypothesis motivating TRD is that transferring the rich, token-level relational information—particularly relevant for patch-based architectures (e.g., ViT [Dosovitskiy et al., 2021]) or feature-aggregating CNNs—enables more complete knowledge transfer. This is especially beneficial in long-tailed settings, where rare classes may share semantic micro-patterns best captured at the token level. Token-level relationships encode: - **Inner-instance semantic context:** How patches or regions within an image are related. - **Cross-instance patch-to-patch similarity:** How a token in one image relates to tokens in others. By explicitly distilling this graph-structured data, TRD aims to bridge the capability gap between teacher and student, surpassing instance- or feature-level approaches [2306.12442]. ## 2. Construction of the Token-level Relationship Graph (TRG) ### Token Representation - **ViT-like networks:** Images $x \in \mathbb{R}^{C \times H \times W}$ are partitioned into $M = HW/P^2$ patches $x_p \in \mathbb{R}^{M \times D}$ with $D = P^2 C$. Teacher tokens $T^T \in \mathbb{R}^{B \cdot M \times d_T}$, student tokens $T^S \in \mathbb{R}^{B \cdot M \times d_S}$. - **CNN-like networks:** The penultimate feature map $F \in \mathbb{R}^{C_\ell \times H_\ell \times W_\ell}$ is split into $M$ patches $\rightarrow$ tokens $T \in \mathbb{R}^{M \times d}$. Every token, whether from teacher or student, is a $d$-dimensional vector $h_i$ used as a node attribute in the graph. ### Random Token Sampling A full batch contains $B \times M$ tokens, often computationally unwieldy. TRD samples $K$ tokens per image using a shared random mask for both models, yielding $N = B \cdot K$ tokens for both teacher and student: $$ T^T = \{ h^T_i \}_{i=1}^N, \quad T^S = \{ h^S_i \}_{i=1}^N. $$ ### Graph Construction Two attributed graphs, $G^T = (V, A^T)$ and $G^S = (V, A^S)$, share vertex set $V = \{1, ..., N\}$, with adjacency matrices $A^T, A^S \in \mathbb{R}^{N \times N}$. For $i \neq j$: $$ A^\{\cdot\}_{ij} = \begin{cases} \exp \left(- \frac{\|h_i - h_j\|^2}{2\sigma^2} \right) & \text{if } i \in k\text{NN}(j) \text{ or } j \in k\text{NN}(i) \\ 0 & \text{otherwise} \end{cases} $$ A dense token-wise relational similarity matrix $R$ can also be formed: $$ R_{ij} = f(h_i, h_j) = \cos(h_i, h_j) = \frac{h_i \cdot h_j}{\|h_i\| \|h_j\|}, $$ optionally normalized: $$ P_{ij} = \mathrm{Softmax}_j(R_{ij}/\tau_R). $$ ## 3. TRG-based Distillation Objectives The total loss combines logit-based KD with several graph and token-level objectives: $$ L_\text{total} = L_\text{logit} + \alpha L_\text{inner} + \beta L_\text{local} + \gamma L_\text{global} $$ with hyperparameters $\lambda, \alpha, \beta, \gamma, \tau, \tau_g^{\mathrm{init}}, W_U$. ### 3.1 Local Preserving Loss This term matches local neighborhood structure between student and teacher graphs: $$ L_\text{local} = \sum_{i=1}^N \mathrm{KL}(\mathrm{Softmax}_j(A^S_{ij}) \,\|\, \mathrm{Softmax}_j(A^T_{ij})). $$ ### 3.2 Global Contrastive Loss Employing an InfoNCE objective, corresponding tokens across teacher and student are aligned, with negatives pushed apart. Student tokens are projected to teacher dimension if $d_S \neq d_T$: $$ \mathrm{SIM}(h_i^S, h_j^T) = \frac{\mathrm{Proj}(h_i^S) \cdot h_j^T}{\|\mathrm{Proj}(h_i^S)\| \|h_j^T\|} $$ $$ L_\text{global} = - \sum_{i=1}^N \log \frac{\exp(\mathrm{SIM}(h_i^S, h_i^T)/\tau_g)}{\sum_{j=1}^N \exp(\mathrm{SIM}(h_i^S, h_j^T)/\tau_g)} $$ A **dynamic temperature** $\tau_g(e)$, where $e$ is the epoch, is adopted: $$ \tau_g(e) = \begin{cases} \tau_g^{\mathrm{init}}, & \text{if } e \leq W_U \\ \tau_g^{\mathrm{init}} / \log_{W_U}(e), & \text{if } e > W_U \end{cases} $$ ### 3.3 Token-wise Contextual Loss To transfer **inner-instance** semantic context, the method matches self-similarity matrices of patch tokens within each image. For penultimate feature map $F$: $$ \mathrm{CS} = \mathrm{Softmax}\left(\frac{F F^\top}{\sqrt{d}} \right), \quad \mathrm{CS}^T, \mathrm{CS}^S \in \mathbb{R}^{M \times M} $$ $$ L_\text{inner} = \|\mathrm{CS}^T - \mathrm{CS}^S\|^2_F $$ This constrains the student to preserve internal patch arrangements consistent with the teacher. ## 4. Empirical Evaluation ### Datasets and Architectures - **Datasets:** CIFAR-100, CIFAR-100-LT (imbalance ratios 10, 50, 100), ImageNet-1K, ImageNet-LT (imbalance 10). - **CNN-based models:** ResNet-32×4 → ResNet-8×4, ResNet56 → ResNet20, VGG13 → VGG8, WRN-40-2 → WRN-40-1, ShuffleNet → MobileNet. - **ViT-based models:** DeiT-Tiny/Small students, ResNet-101 or CeiT-Base teachers. ### Training Configurations - **CIFAR-100:** SGD with Nesterov ($0.9$ momentum, $5\times 10^{-4}$ weight decay), 240 epochs, reduced LR at 150, 180, 210. $\tau=4$, $\tau_g^{\mathrm{init}}=0.1$, warm-up $W_U=15$. - **ImageNet & ViTs:** CNNs – 100 epochs, cosine LR schedule, batch 128×4. ViTs – AdamW, LR=$5\times10^{-4}$, 200 epochs, 10-epoch warm-up. All experiments on 4×RTX3090 GPUs. ### Performance Summary TRD achieves superior accuracy and robustness, notably: | Setting | TRD Top-1 (%) | Strongest Baseline (%) | |----------------------------|---------------|--------------------------| | ResNet-32×4 → 8×4 (CIFAR) | 76.42 | HKD: 76.21, KD: 74.12 | | ShuffleV1←ResNet32×4 | 76.42 | DKD: 76.42, HKD: 75.99 | | ResNet34→ResNet18 (ImageNet)| 71.31 | KD +1.56% | | DeiT-Tiny←ResNet101 | 75.5 | KD: 74.8, HKD: 75.2 | | DeiT-Small←CeiT-Base | 81.8 | HKD: 81.3 | On long-tailed variants: - **CIFAR-100-LT:** TRD degrades less as imbalance increases and can surpass teacher accuracy. - **ImageNet-LT:** TRD Top-1 50.32% ($-20.99\%$ drop) vs. KD 46.70% ($-24.00\%$ drop). Ablation studies confirm each loss ($L_\text{inner}$, $L_\text{local}$, $L_\text{global}$) contributes $0.3\%$–$1.0\%$ accuracy, with token-level graphs outperforming instance-level graphs. Dynamic $\tau_g$ provides smoother optimization and lower embedding divergence. ## 5. Analysis: Contextualization and Visualizations Several key insights arise from empirical analysis: - Larger batch sizes (e.g., 512 tokens) strengthen the representational capacity and graph structure, yielding marginal accuracy improvements. - t-SNE projections show TRD features are more class-separable than those from KD, IRG, or HKD, indicating successful transfer of fine-grained relational information. - The addition of each loss component leads to measurable, compositional gains, confirming efficacy of the multi-term objective [2306.12442]. ## 6. Limitations and Prospective Directions Identified challenges and future research include: - **Computational demands:** Full construction of k-NN graphs for $O(BK)$ tokens is resource-intensive. Practical deployments may require efficient or approximate graph building techniques such as locality-sensitive hashing. - **Hyperparameter sensitivity:** Optimal settings for $\alpha$, $\beta$, $\gamma$, $k$, $\sigma$, $\tau_g^{\mathrm{init}}$, $W_U$ must be tuned for datasets and architectures. - **Beyond classification:** Extending TRD to object detection, semantic segmentation, or temporal token graphs for video. - **Graph topology learning:** Instead of fixed $k$-NN, learning the adjacency matrix during training. - **Self-distillation and multiscale transfer:** Applying token-level relations for intra-network feature alignment. - **Cross-modal distillation:** Adapting the framework to transfer relational information across modalities (e.g., image-text pairs in vision-language models). ## 7. Related Work and Positioning TRD advances the progression from basic logit-based KD [Hinton et al., 2015] to feature and relation-based approaches, such as RKD [Park et al., 2019], IRG [Liu et al., 2019], and graph-based distillation [Zhou et al., 2021]. Distinct from prior work that either matches holistic features or instance relationships, TRD’s explicit modeling of token-level graphs enables new forms of fine-grained and spatially-aware knowledge transfer, especially applicable to architectures with patch- or region-based representations (e.g., ViT [Dosovitskiy et al., 2021]). Use of InfoNCE and contrastive paradigms is aligned with the research trajectory outlined in [Tian et al., 2019; Wang & Isola, 2020]. A plausible implication is that the token-centric relational framework may generalize to settings with structured semantic dependencies beyond image classification, suggesting future exploration into relational and multi-modal knowledge transfer regimes [2306.12442]. Source: https://www.emergentmind.com/topics/token-relation-distillation-trd