SQuaT: Self-Supervised Knowledge Distillation via Student-Aware Quantized Teacher Features

SQuaT (Student-Aware Quantized Teacher Features) is a label-free Quantization-Aware Training (QAT) framework with Knowledge Distillation (KD) that addresses a fundamental limitation in prior work combining QAT with KD: during distillation, the range mismatch between the teacher and the quantized student model induces an unattainable residual, resulting in an irreducible lower bound on the distillation loss. The paper observes that the intermediate feature maps of a FP teacher model and a low-bit quantized student model have markedly different value ranges, and that using FP features as distillation targets for the student induces an unattainable residual between the teacher and the student, resulting in a theoretically irreducible lower bound on the distillation loss.

The core idea of SQuaT is to "quantize the intermediate features of the teacher model using the quantization parameters of the student model during distillation, thereby projecting the teacher’s features directly into the student’s quantized representation space. This student-aware projection aligns the teacher and student features on the same quantization lattice and enables feature-level knowledge distillation using only targets that are attainable by the student. The paper formally proves in Theorem 1 that supervising a quantized student with the FP target fT,l cannot drive the expected squared error below the quantization mismatch E ΠϕS,l (fT,l)− fT,l 2, whereas using the student-quantized target f˜T,l = ΠϕS,l (fT,l) removes this lower bound."

The framework performs label-free joint alignment at both the feature and logit levels, where the former preserves spatial/channel patterns, while the latter maintains class-level semantics. The joint objective is LSQuaT = λfeat Lfeat + λlogit Llogit, where feature-level alignment uses L2 distance between quantized student features and student-aware projected teacher features, and logit-level alignment uses KL divergence with temperature smoothing.

The paper's main contributions are: (1) proposing SQuaT, which introduces a paradigm shift from teacher-aware to student-aware projection by leveraging the student model’s quantization parameters; (2) showing through theoretical and empirical analyses, that SQuaT effectively removes the lower bound caused by unattainable residuals in prior teacher-aware projection-based methods; and (3) conducting extensive experiments to validate the proposed approach, demonstrating that SQuaT is agnostic to model design choices and broadly applicable, with particularly pronounced gains in extreme low-bit settings.

Experimental results show consistent improvements across diverse settings. On CIFAR-10 with ResNet-20, under 1-bit quantization, SQuaT improves over SQAKD by +0.45 pp. On CIFAR-100 with ResNet-32, under 1-bit quantization, SQuaT improves over SQAKD by +1.74 pp. On ImageNet-1K with ResNet-18, under 1-bit quantization, it improves over SQAKD by +0.40 pp. The paper notes that the performance gap widens in extreme low-bit settings, highlighting the robustness of the proposed projection strategy under stringent quantization constraints.

The approach demonstrates generality across model architectures and domains. For DeiT-Tiny (Vision Transformer) on CIFAR-100, SQuaT achieves improvements of +1.27, +3.27, +3.42 for W2A2, W3A3, and W4A4 respectively. For BERTBASE on the GLUE benchmark, improvements include +3.07 on RTE, +0.69 on SST-2, and +1.26 on QNLI for W3A3. The paper states that "the effectiveness of SQuaT does not rely on a particular architectural assumption or quantization implementation, but rather stems from the general principle of aligning teacher features to the student’s quantization lattice."

The paper also compares projection strategies empirically, showing that the Student-aware scheme achieves higher accuracy than both Non-aware and Teacher-aware approaches across all bit-widths and model configurations, and that Teacher-aware projection often performs worse than the Non-aware variant. The training curves show that the Student-aware approach converges to the lowest feature distillation loss throughout training, indicating stable optimization.

Finally, on-device evaluation on NVIDIA Jetson Nano shows that compared with FP32 inference, INT8 inference increases throughput from 294.32 FPS to 2046.03 FPS and reduces latency from 3.44 ms to 0.53 ms, corresponding to a 6.49× speedup.

Based on the paper, here are the specific improvements I can make to AI systems:

1. Quantization-Aware Training with Student-Aware Feature Alignment: I will implement a training pipeline where the teacher model’s intermediate features are quantized using the student’s quantization parameters (e.g., scale and zero-point) before computing the distillation loss. This removes the irreducible lower bound on the distillation loss caused by range mismatch, enabling the student to achieve lower expected squared error and higher accuracy, especially at extreme low-bit widths (e.g., 1-bit weights and activations).

2. Joint Feature-Logit Distillation with Attainable Targets: I will combine feature-level L2 loss (on student-quantized teacher features) with logit-level KL divergence (temperature-smoothed) in a label-free manner. This ensures the student learns both spatial/channel patterns and class-level semantics without relying on ground-truth labels, improving generalization in semi-supervised or privacy-sensitive settings.

3. Architecture-Agnostic Quantization Framework: I will apply SQuaT to diverse model families—CNNs (ResNet), Vision Transformers (DeiT), and Transformers (BERT)—without modifying the underlying architecture or quantization scheme. This allows me to deploy the same training recipe across vision and NLP tasks, yielding consistent gains (e.g., +3.42 pp for DeiT-Tiny on CIFAR-100 at W4A4, +3.07 pp for BERT on RTE at W3A3).

4. Theoretical Guarantee for Loss Minimization: I will use the paper’s Theorem 1 to certify that my training objective has no unattainable residual lower bound. This enables me to set training stopping criteria or learning rate schedules based on the actual achievable loss, improving optimization stability and convergence speed.

5. Extreme Low-Bit Deployment with Minimal Accuracy Drop: I will use SQuaT to train models at 1-bit or 2-bit precision, achieving higher accuracy than prior KD-based QAT methods (e.g., +1.74 pp on CIFAR-100 with ResNet-32 at 1-bit). This allows me to deploy models on edge devices with severe memory and compute constraints while maintaining task performance.

6. On-Device Inference Acceleration: By training with SQuaT and then quantizing to INT8, I can achieve a 6.49× throughput increase (from 294 FPS to 2046 FPS) and 6.49× latency reduction (from 3.44 ms to 0.53 ms) on hardware like NVIDIA Jetson Nano, enabling real-time inference for embedded AI applications.

7. Stable Training Dynamics: I will adopt the student-aware projection to ensure the feature distillation loss converges to the lowest value throughout training (as shown in the paper’s training curves), reducing variance and improving reproducibility across random seeds and data splits.

8. Label-Free Distillation for Unlabeled Datasets: Since SQuaT does not require labels for feature alignment, I can use it to distill knowledge from a large pre-trained teacher into a small quantized student using only unlabeled data, enabling model compression in domains where labeled data is scarce or expensive to obtain.

What the improved AI system can do:

• Achieve state-of-the-art accuracy for quantized models at 1–4 bit precision across vision (CIFAR, ImageNet) and NLP (GLUE) benchmarks.

• Run efficiently on low-power edge devices (e.g., Jetson Nano) with >6× speedup and minimal accuracy loss.

• Be trained without labels, making it suitable for semi-supervised or unsupervised compression.

• Generalize across CNNs, ViTs, and Transformers without architectural changes.

• Provide theoretically grounded training that avoids loss plateaus and unstable optimization.