Source-linked AI summary
Knowledge Transfer with Jacobian Matching
Suraj Srinivas, Francois Fleuret
TL;DR
The paper studies how to choose and justify Jacobian-matching losses for knowledge transfer, including distillation and transfer learning across different architectures. It equates Jacobian matching with distillation under input noise, connects LwF to distillation, and reports improvements in distillation, noisy-input robustness, and transfer learning. The authors also identify optimization difficulties for directly reducing Jacobian loss in practical deep networks and a remaining gap to oracle transfer learning.
Problem
Existing Jacobian-matching methods leave the appropriate loss function unclear, while the relationship between a recent transfer-learning procedure and distillation is also unresolved.
Method
The paper derives Jacobian-matching losses from an equivalence with distillation under input noise and applies the analysis by viewing LwF as distillation.
Results
Jacobian-based penalties improve regular distillation, robustness to noisy inputs, and transfer learning in experiments on standard image datasets.
Takeaways & Limitations
Jacobian matching provides a noise-augmentation interpretation for transfer across arbitrary architectures and supports its use in distillation and transfer learning.
Takeaways & Limitations
Direct Jacobian-loss optimization failed to reduce the loss in practical deep networks, and a substantial gap remains between distillation methods and oracle pre-trained-network transfer.
Abstract
from arXiv · showhide
Classical distillation methods transfer representations from a "teacher" neural network to a "student" network by matching their output activations. Recent methods also match the Jacobians, or the gradient of output activations with the input. However, this involves making some ad hoc decisions, in particular, the choice of the loss function. In this paper, we first establish an equivalence between Jacobian matching and distillation with input noise, from which we derive appropriate loss functions for Jacobian matching. We then rely on this analysis to apply Jacobian matching to transfer learning by establishing equivalence of a recent transfer learning procedure to distillation. We then show experimentally on standard image datasets that Jacobian-based penalties improve distillation, robustness to noisy inputs, and transfer learning.
1. Introduction
The paper addresses knowledge transfer between neural networks with different architectures, focusing on unresolved choices for Jacobian matching and its relation to distillation and transfer learning. It proposes Jacobian-based connections and reports benefits for distillation, noisy-input robustness, and transfer learning.
- Motivation: Knowledge transfer is challenging when source and target networks have different architectures, because direct weight copying is unavailable.Distillation uses the same dataset with different architectures, while transfer learning uses related datasets and may also require architectural flexibility.
- Open problems: Prior Jacobian-matching methods did not establish which loss function to use or how they relate to classical distillation.The paper also identifies uncertainty about applying distillation-like transfer learning methods when both architectures are arbitrary.
- Contributions: The paper shows that matching Jacobians is a special case of classical distillation with noise added to the inputs.This equivalence provides a basis for selecting Jacobian-matching losses.
- Contributions: A recent transfer learning method, LwF, can be viewed as distillation, enabling Jacobian matching in that setting.The paper also develops methods for matching Jacobians when both networks have arbitrary architectures.
- Empirical findings: Experiments provide evidence that Jacobian matching helps regular distillation and transfer learning, while Jacobian-norm penalties learn models robust to noise.These results are reported as experimental validation of the paper’s theoretical connections.
2. Related Work
Related work established several uses for derivative-based regularization, distillation, and feature matching, but this paper clarifies how Jacobian matching connects to activation-based distillation and loss design.
- Jacobian-based regularization: Sobolev training showed that combining higher-order derivatives with targets can help train models with less data.The paper presents this as closely related to its own use of derivative information.
- Paper’s connection: The paper distinguishes its contribution by relating Jacobian matching to regular distillation and showing that the activation-matching loss determines the corresponding Jacobian loss.This addresses a design question left unclear in earlier Jacobian-matching work.
- Jacobian-based regularization: Earlier work penalized Jacobian norms to make models more robust to small input changes.The paper states that this intuition conforms to its analysis of Jacobian-based regularization.
- Distillation: Prior distillation research used temperature-scaled softmax, squared logit error, intermediate-feature matching, and noise added to logits.These approaches differ in their targets or matching mechanisms.
3. Jacobians of Neural Networks
The paper motivates Jacobian matching because Jacobians are architecture-independent and can therefore compare networks with different structures, while local linearization explains how they characterize nearby behavior.
- Local linearization: A first-order Taylor expansion approximates f(x + ∆x) using the function value, its input gradient, and higher-order error terms.The expansion is considered in a small neighborhood around x.
- Local linearization: Neural-network linearization requires handling elementwise nonlinearities and pooling operators, which are the sources of network nonlinearity.ReLU and sigmoid are given as examples of elementwise nonlinear activations.
- Piecewise-linear networks: For ReLU, the local Taylor approximation is exact except at zero because its derivative is 0 or 1; max-pooling has a similar property.The same local behavior supports piecewise-linear analysis.
- Piecewise-linear networks: For piecewise-linear networks, sufficiently small perturbations on the same linear surface make the higher-order terms exactly zero.Under this condition, f(x + ∆x) equals its first-order linearization.
- Architecture independence: For k output classes and input dimension D, the Jacobian has dimension D × k, independent of network architecture.This dimensional property allows Jacobians from different architectures to be compared.
- Architecture independence: Different weight configurations can produce the same Jacobian because of hidden-neuron permutation symmetry, model redundancy, and non-convexity.These properties mean Jacobians can represent shared function behavior without uniquely identifying weights.
- Transfer motivation: The paper proposes using Jacobians together with output activations for knowledge transfer, while leaving the practical matching procedure as the key design question.Both quantities share architecture-independent dimensions, according to the passage.
4. Distillation
The paper shows that Jacobian matching is equivalent to distillation with noisy inputs, yielding a principled Jacobian penalty tied to the activation-matching loss. This connection also motivates noise-robust regularization while exposing computational costs of full-Jacobian computation.
- Noise-augmented distillation: Jacobian matching is equivalent to matching teacher and student outputs on inputs augmented with noise.The equivalence follows by expanding the noisy-input distillation loss around each data point.
- Noise-augmented distillation: The noisy-input loss decomposes into ordinary distillation on observed samples plus a Jacobian-matching regularizer.For piecewise-linear networks, the higher-order error term is exactly zero.
- Choosing the loss: The activation-matching loss determines the corresponding Jacobian-matching loss, so noise assumptions can guide loss selection.The paper discusses both squared-error and cross-entropy formulations, finding squared error more stable experimentally for distillation.
- Interpretation: Matching outputs at nearby noise-perturbed inputs explains why Jacobian matching can make the student mimic the teacher beyond the observed data points.The paper interprets Jacobian matching as matching outputs over infinitely many nearby points.
- Computational considerations: Full Jacobian computation is computationally expensive, motivating approximations that retain only selected output-variable terms.The paper proposes label-based or largest-magnitude output heuristics for selecting terms.
5. Transfer Learning
The paper extends transfer learning to networks with arbitrary architectures by interpreting LwF as distillation and matching outputs, attention maps, and their Jacobians. Its analysis connects target-domain losses to source-domain distillation through dataset distance, while practical Jacobian matching uses approximations because full Jacobians are costly and direct optimization can fail.
- Transfer Learning: Fine-tuning restricts the pre-trained and target networks to the same architecture, motivating transfer methods that allow arbitrary architectures.The target dataset is typically smaller and differs statistically from the source dataset.
- 5.2. Jacobian Matching: The method matches output activations, aggregated attention maps, and Jacobians of attention maps with respect to inputs between teacher and student networks.Attention maps aggregate squared channel activations; full Jacobians are approximated by differentiating the largest teacher attention-map pixel at the same index in the student.
- Transfer Learning: LwF trains a student with target labels while matching teacher activations on the target dataset, without explicitly using source data.The student can use separate output branches for target and source label spaces.
- 5.1. LwF as Distillation: LwF approximates distillation on part of the source dataset, with a source distillation loss bounded by a target-domain max-loss plus terms depending on source–target distance.When the Hausdorff distance is small, reducing the max-loss also reduces the distillation loss.
- 5.1. LwF as Distillation: Adding noise to target samples is expected to tighten the transfer bound because it decreases the minimum distance between source and target points.The result is stated for any superset of the target dataset that includes the augmented samples.
- 5.2. Jacobian Matching: Direct Jacobian-loss optimization failed for practical deep networks, possibly because second-order gradients vanish or cause network degeneracy, so an alternative matching method was required.The paper also notes that computing full Jacobians is computationally costly.
6. Experiments
Experiments evaluate Jacobian matching for limited-data distillation, noisy-input robustness, and transfer learning across CIFAR100 and MIT Scenes. Jacobian matching improves distillation and transfer methods, while feature depth and pooling affect its effectiveness.
- Experimental setup: The experiments evaluate distillation on CIFAR100, noise robustness on CIFAR100, and transfer learning from ImageNet to MIT Scenes.The studies use teacher–student networks with different data regimes and architectures.
- Distillation: 52.43% accuracy with 100 data points per class approaches 54.28% from full-data raw training, using one-fifth of the data.This result uses activations, Jacobians, and regular cross-entropy training.
- Distillation: Matching activations and Jacobians with cross-entropy performs best for limited-data distillation settings.The VGG-4 student is distilled from a VGG-9 teacher; the teacher reaches 64.78% accuracy.
- Noise robustness: Higher Jacobian-norm regularization improves robustness to Gaussian noise, whereas ℓ2 regularization and dropout are less robust.The comparison evaluates VGG-9 models on noisy CIFAR100 images.
- Transfer learning: Jacobian matching improves activation- and attention-based transfer methods, but none reaches the oracle performance of a pretrained target model.The performance gap is especially large at 10 and 25 images per class.
7. Conclusion
The paper interprets Jacobian matching as noise-augmented distillation and connects LwF transfer learning to distillation. It reports benefits for knowledge transfer but acknowledges a remaining gap to pretrained-model transfer.
- Conclusion: Jacobian matching can be viewed as data augmentation with Gaussian noise, which motivates its use and informs the matching loss.The analysis treats Jacobian matching as matching outputs at nearby noise-perturbed inputs.
- Conclusion: Connecting LwF to distillation enables Jacobian matching for transfer learning.This extends the Jacobian-matching analysis beyond regular distillation.
- Conclusion: Distillation-based transfer learning still has a large performance gap relative to using pretrained networks directly.The paper suggests more structured data augmentation than simple noise as future work.
8. Proof for Proposition 1
The proof establishes locally that squared-error matching under small Gaussian input perturbations decomposes into ordinary output matching plus a Jacobian-based penalty. It relies on a first-order Taylor expansion and zero-mean, diagonal-covariance noise.
- Proposition setup: The proposition considers squared error between k-dimensional teacher and student outputs at an input x, with perturbation ξ = σz.The perturbation is a scaled unit normal vector in the input space.
- Proof: For sufficiently small perturbations, a first-order Taylor expansion supports the local equivalence.The proof explicitly assumes ξ is small enough for the expansion to hold.
- Proof: Zero-mean noise removes the first-order perturbation term in expectation.The proof invokes the fact that the mean of ξ is zero.
- Proof: A diagonal covariance assumption completes the reduction to coordinatewise Jacobian contributions.The stated proof uses this covariance structure after applying the zero-mean condition.
- Extensions: The same proof strategy extends to cross-entropy using a second-order Taylor expansion of log.The appendix states that other statements have analogous proofs.
9. Proof for Proposition 3
The proof bounds target-set loss using nearest source points, a Lipschitz condition, and the directed Hausdorff distance. Enlarging the source set cannot increase this distance and therefore tightens the bound.
- Assumptions: The proof assumes the loss discrepancy ρ is Lipschitz with constant K under an input-space distance ψx.This assumption relates loss changes to distances between corresponding data points.
- Nearest-point construction: For each target point, the proof selects a closest point from the source dataset, allowing repeated selections and a subset of source points.This nearest-point construction is used to compare target and source losses.
- Bounding step: The maximum source loss bounds the average loss over the selected nearest-point subset.The argument uses the fact that a maximum is at least as large as any convex combination.
- Bounding step: The nearest-point distance is bounded by the directed Hausdorff distance between the target and source datasets.This converts pointwise proximity into a dataset-level bound.
- Corollary: For any enlarged source dataset D′s ⊇ Ds, the directed Hausdorff distance satisfies Ha(Dl, D′s) ≤ Ha(Dl, Ds).Adding source points can only make the nearest source point no farther from each target point.
10. Experimental details
The experiments use VGG-style convolutional architectures, specified distillation losses, and dataset-specific optimization settings. Jacobian matching is implemented with defined activation/Jacobian terms and feature-level choices, with smaller batches for memory-intensive computation.
- Architectures: The convolutional blocks use 3 × 3 stride-1 kernels, while max-pooling uses stride 2 and a 2 × 2 window.
- Architectures: VGG-9 and VGG-4 use sequential convolutional blocks with max-pooling, followed by global average pooling and a fully connected classifier.VGG-9 uses channel widths 64, 128, 256, 256, 512, 512, 512, and 512; VGG-4 uses 64, 128, and 512.
- Distillation setup: Distillation compares cross-entropy, squared error between pre-softmax teacher and student activations, and squared error between their Jacobians.The loss function is introduced for the distillation experiments, with activation matching written as (y_s−y_t)^2.
- Optimization: CIFAR100 training runs for 500 epochs with Adam, batch size 128, and learning-rate annealing from 1e−3 to 1e−4 at epoch 400.
- Optimization: MIT Scene training uses SGD with momentum 0.9 for 75 epochs, batch size 8, and learning-rate reductions after epochs 40 and 60 because Jacobian computation is memory intensive.For MIT Scene, regularization constants are set to either 10 or 0 depending on the method; Jacobians are pooled and matched at feature level 1.
- Distillation setup: Regularization constants are set to 1 or 0 in the reported experiments, with all three constants equal to 1 when enabled.