Gated Relational Alignment via Confidence-based Distillation for Efficient VLMs

AI summary

Overview

• Research area: Efficient vision-language model (VLM) compression, combining quantization-aware training (QAT) with knowledge distillation (KD) for low-bit inference.
• Technical level: Advanced. The paper assumes familiarity with knowledge distillation, KL divergence, quantization (PTQ/QAT, LSQ), and the Information Bottleneck principle, and it presents formal theorems and propositions.
• Scope: The paper introduces GRACE, a unified QAT-plus-distillation framework that lets INT4-quantized VLMs match or exceed their BF16 counterparts while delivering faster throughput and lower memory use.

What This Paper Is About

Vision-language models give strong multimodal performance but are expensive to deploy, and pushing them to INT4 precision usually degrades accuracy because visual, projection, and language components react differently to low-bit perturbation. The authors argue that QAT for VLMs is underexplored, and that naively bolting KD onto QAT fails because some teacher predictions are unreliable and a quantized student has limited capacity to absorb everything the teacher knows. GRACE reframes the problem as a capacity-allocation question under the Information Bottleneck principle: quantization fixes the bit budget, and distillation decides which teacher knowledge gets preserved.

Key Contributions

1. An Information Bottleneck framing of QAT-based VLM compression. The paper formalizes low-bit training as maximizing the student's mutual information with teacher-provided task-relevant information, subject to a hard capacity limit imposed by quantization, and connects capacity allocation to teacher-guided knowledge preservation.
2. The GRACE framework with three complementary components. Confidence-gated decoupled knowledge distillation (GDKD), relational centered kernel alignment (RCKA) over visual tokens, and an adaptive Information Bottleneck controller that tunes distillation strength online.
3. Group-wise learned step-size quantization for the student. Quantization is applied only to the LLM decoder linear layers, with learnable log-space scales initialized by MSE-based grid search and optimized jointly with the weights.
4. Empirical evidence that INT4 students can beat BF16 baselines. Evaluations on LLaVA-1.5 and Qwen2-VL show accuracy gains over BF16 baselines plus 3x throughput and 54% memory reduction with real INT4 kernels.

Main Findings

• Teacher entropy predicts teacher error. On ScienceQA with LLaVA-1.5 13B as teacher, teacher output entropy correlates positively with sample-level error (Pearson r = 0.484); regressing binned error rates on entropy deciles yields R² = 0.901. This motivates using entropy as a confidence signal.
• Logit distillation misses visual attention structure. Attention visualizations of LLaVA-1.5 13B versus 7B on "What object is being used as the telephone receiver?" show the 13B teacher progressively localizing the banana while the 7B model's attention stays scattered — a capability the authors say logit-level distillation alone cannot transfer.
• GRACE leads KD baselines on LLaVA-1.5. Averaged over VQA v2, GQA, TextVQA, VizWiz, POPE, SQA, MME, and MMB, GRACE reaches 69.0% versus 66.5% for the LLaVA-1.5-7B baseline (a 2.5% gain as labeled in the table), approaching the 13B teacher's 69.1%. MoVE-KD-v1.0 scores 68.0, MoVE-KD-v1.1 and HAWAII both 68.5. GRACE also leads on TextVQA (61.5 vs. 58.2 baseline and 61.3 teacher) and SQA (71.7 vs. 66.8 baseline, 71.6 teacher).
• INT4 beats BF16 and PTQ baselines on LLaVA-1.5. At 4-bit, RTN drops 1.3% and AWQ drops 0.9% relative to the BF16 baseline, while vanilla QAT only reaches 65.9% (−0.6%). GRACE reaches 67.2%, exceeding the BF16 baseline by 0.7% and beating AWQ and QAT by 1.6% and 1.3% absolute. At 8-bit GRACE reaches 68.3%, surpassing the full-precision baseline by 1.8%.
• All three components are complementary (Qwen2-VL 7B → 2B). From a vanilla KD average of 74.8%, GDKD alone reaches 76.6%, RCKA alone 76.8%, and the adaptive IB controller alone 75.5%. GDKD + RCKA reaches 77.5%, and all three together reach 78.2% (+5.5% over the 2B baseline of 72.7%), against a 7B teacher at 80.3%.
• QAT alone loses accuracy; GRACE recovers it. Qwen2-VL 2B with QAT only scores 72.2 at 8-bit (−0.5%) and 71.3 at 4-bit (−1.4%) relative to the BF16 baseline of 72.7%. Adding naive KD gives 74.4 and 73.4. Adding GRACE gives 77.9 at 8-bit (+5.2%) and 77.2 at 4-bit (+4.5%).
• Soft exponential gating is the best gating design. Under standard KL distillation on Qwen2-VL (7B → 2B, BF16), no gating gives 74.8%, hard thresholding falls to 72.3%, linear gating improves over no gating, and exponential gating peaks at λ = 1.0 with 76.2%. Performance follows an inverted-U: λ = 0.5 under-suppresses noise and λ = 3.0 over-suppresses supervision.
• Qwen2-VL headline gains in the abstract. The abstract reports INT4 Qwen2-VL-2B at 76.9 on MMBench versus 72.6 for the baseline (the Qwen2-VL BF16 baseline appears as 71.6 on MMBench in Tables 3 and 4), and the introduction reports 79.1 vs. 73.7 on ScienceQA, with full Qwen2-VL results deferred to Appendix A.1, Table 5, which is not included in the provided content.

Methodology in Plain English

The setup is a standard teacher-student arrangement: a frozen, full-precision teacher and a trainable quantized student process the same input, and the student is trained with a combination of three losses — cross-entropy on the task, a gated distillation loss, and a relational alignment loss.

Filtering unreliable teacher signals. Following decoupled knowledge distillation, the distillation objective is split into a target-class term (how confident the teacher is in the correct answer) and a non-target-class term (the "dark knowledge" spread over incorrect classes), with the non-target term weighted more heavily. Each token's contribution is then down-weighted by the teacher's entropy: the weight is exp(−h̃), where h̃ is the token's entropy normalized by log of the vocabulary size. Confident tokens dominate the loss; uncertain ones are suppressed. The authors prove that this gating shifts the loss by a covariance term, which is negative when entropy and loss move together, and they further show via a variational bound that minimizing the gated loss maximizes the mutual information between the student's representation and the teacher's knowledge.

Transferring visual structure, not just answers. Instead of matching hidden features point by point, the student is asked to reproduce the teacher's relational geometry among visual tokens. The penultimate-layer representations of visual tokens (text tokens are excluded) are row-normalized, turned into pairwise cosine-similarity Gram matrices, centered, and compared with Centered Kernel Alignment. The loss is simply 1 − CKA. Because CKA operates on n × n matrices, it works even when teacher and student have different hidden dimensions and needs no projection layers — an important property when the student is quantized.

Balancing fidelity against capacity. The Information Bottleneck view treats the distillation loss as a constraint: keep CE loss low subject to the distillation loss staying below a budget τ. Rather than fixing the trade-off weight β by hand, the authors relax this into a Lagrangian and update β online with projected dual ascent on an EMA-smoothed distillation loss, clamped between minimum and maximum values. If the student is failing to retain teacher knowledge, β increases; if the constraint is satisfied, β decreases so task performance dominates.

Quantization details. Weights in the LLM decoder's linear layers are split into contiguous groups (group size 128 by default), each with its own learnable scale stored in log space. Scales are initialized by sampling 20 candidate values in the range [0.3, 1.2] · max(|W_i|)/Q_p and picking the one minimizing fake-quantization error, then optimized jointly with the weights at a learning rate 10x that of the weights, with no weight decay and straight-through gradients through the rounding operation. The vision encoder stays frozen in BF16.

Training setup. Training uses the ShareGPT4V dataset (1.3M high-quality image-text pairs generated by GPT-4V and an extended caption model) on 8 NVIDIA H100 GPUs — 12 hours for LLaVA-1.5 and 8 hours for Qwen2-VL. LLaVA-1.5 uses 13B as teacher and 7B as student; Qwen2-VL uses 7B as teacher and 2B as student. All inference experiments run on a single NVIDIA A100 GPU. The code and data are released at ForeverBlue816/GRACE.

Why This Matters

The paper argues that good low-bit VLM deployment is not just a numerical precision problem but a question of what information is worth keeping under a hard bit budget. By showing that INT4 students can surpass BF16 baselines — LLaVA-1.5-7B at 70.1 vs. 66.8 on SQA, Qwen2-VL-2B at 79.1 vs. 73.7 on ScienceQA — it challenges the assumption that quantization must cost accuracy, and it offers a concrete recipe (confidence gating plus relational alignment) that is architecture-agnostic and requires no complex multi-teacher or multi-encoder designs.

Real-world applications:

• On-device multimodal assistants on phones and laptops, where 54% memory reduction and 3x throughput make billion-parameter VLMs practical.
• Robotics and vision-language-action systems, which the paper cites as a target domain and which are constrained by onboard compute.
• Visual question answering and scene understanding in bandwidth- or power-limited settings such as edge cameras and industrial inspection.
• Text-heavy visual tasks like scene-text reading, where GRACE shows large gains (TextVQA 61.5, above the 13B teacher's 61.3 on the LLaVA-1.5 comparison).

Industry relevance: One co-author is affiliated with Qualcomm AI Research, and the framing throughout is deployment-oriented — throughput, memory footprint, and real INT4 kernels. The results suggest that compression pipelines built only on post-training quantization or plain QAT leave substantial accuracy on the table, which matters directly for serving costs and for any product that must run VLMs without data-center GPUs.

Future Directions

• Pushing below INT4. All reported low-bit results are INT8 and INT4; the paper does not test 2-bit or mixed-precision schemes, leaving open whether the confidence gating and relational alignment still hold when the capacity constraint becomes much tighter.
• Scaling to more model families and sizes. Evaluations cover LLaVA-1.5 (13B → 7B) and Qwen2-VL (7B → 2B) only. The provided content does not report results on other VLM families, video-capable models, or larger teachers.
• Understanding the interaction between the three components more deeply. The adaptive IB controller is described as mainly regulating distillation strength rather than adding new supervision, and its standalone gain (75.5% vs. 74.8%) is the smallest of the three components; a finer analysis of when β adaptation actually helps is not reported.
• Broader deployment validation. Results are measured on a single NVIDIA A100 GPU; the provided content does not report latency, power, or accuracy measurements on actual edge accelerators. The paper's conclusion is also truncated in the provided text, so an explicit limitations section and author-stated future work are not available here.

Target Audience

Researchers and engineers working on model compression, quantization, and multimodal model deployment will get the most from this paper, particularly those already comfortable with knowledge distillation and the Information Bottleneck. Practitioners building on-device or low-cost VLM serving pipelines will find the practical QAT recipe and the INT4-versus-BF16 comparisons directly useful. Readers without a background in quantization-aware training will need to work through the formal sections (Theorem 3.1, Proposition 3.2, and the appendices), but the motivation and ablation sections stand on their own.

Authors’ abstract

Vision-Language Models (VLMs) achieve strong multimodal performance but are costly to deploy, and post-training quantization often causes significant accuracy loss. Despite its potential, quantization-aware training for VLMs remains underexplored. We propose GRACE, a framework unifying knowledge distillation and QAT under the Information Bottleneck principle: quantization constrains information capacity while distillation guides what to preserve within this budget. Treating the teacher as a proxy for task-relevant information, we introduce confidence-gated decoupled distillation to filter unreliable supervision, relational centered kernel alignment to transfer visual token structures, and an adaptive controller via Lagrangian relaxation to balance fidelity against capacity constraints. Across extensive benchmarks on LLaVA and Qwen families, our INT4 models consistently outperform FP16 baselines (e.g., LLaVA-1.5-7B: 70.1 vs. 66.8 on SQA; Qwen2-VL-2B: 76.9 vs. 72.6 on MMBench), nearly matching teacher performance. Using real INT4 kernel, we achieve 3$\times$ throughput with 54% memory reduction. This principled framework significantly outperforms existing quantization methods, making GRACE a compelling solution for resource-constrained deployment. Code and data are available at: https://github.com/ForeverBlue816/GRACE.