T3C: Test-Time Tensor Compression with Consistency Guarantees

TL;DR

The paper addresses adapting deep models to variable runtime budgets without retraining by introducing T3C, a train-once, test-time budget-conditioned compression framework. It blends elastic tensor factorization with rank-tied mixed-precision quantization and a lightweight budget controller to produce per-layer $k_\ell$ and $q_\ell$ that map to hardware-aligned profiles. A fast consistency certificate upper-bounds logit drift, regularizing training and enabling monotone, certified accuracy–latency–size trade-offs across devices. Empirically, a single T3C checkpoint achieves predictable frontiers across CNNs, ViTs, and NLP models and consistently outperforms strong PTQ/QAT baselines on ImageNet-1k and GLUE benchmarks.

Abstract

We present T3C, a train-once, test-time budget-conditioned compression framework that exposes rank and precision as a controllable deployment knob. T3C combines elastic tensor factorization (maintained up to a maximal rank) with rank-tied mixed-precision quantization and a lightweight controller that maps a latency/energy/size budget token to per-layer rank/bit assignments; the policy snaps to hardware-aligned profiles and is monotone in the budget. A fast, layerwise consistency certificate, computed from spectral proxies and activation statistics, upper-bounds logit drift and regularizes training, yielding a practical reliability signal with negligible overhead. On ImageNet-1k, T3C shifts the vision Pareto frontier: for ResNet-50 at matched accuracy (\leq 0.5% drop), p50 latency is 1.18ms with a 38MB model, outperforming PTQ-8b (1.44ms, 88MB); for ViT-B/16, T3C reaches 2.30ms p50 with 59MB, improving over strong PTQ/QAT baselines. A single T3C checkpoint therefore provides predictable, certificate-backed accuracy-latency-size trade-offs on demand across devices.

Figures (4)

• Figure 1: Detailed T3C pipeline (train once, control at test time). Given an input $x$, the full (teacher) model evaluated at high rank $k_{\max}$ produces teacher logits/distribution$p_{\text{full}}(x)$ and an optional input summary$s(x)$. A budget token$b$ (e.g., latency/energy/size target) together with $s(x)$ is consumed by the budget controller$\pi_\phi(b,s(x))$, which outputs per-layer rank and precision assignments$\{k_\ell, q_\ell\}$. In parallel, each layer’s weight tensor $W$ is stored in an elastic factorization (SVD/Tucker/CP) maintained up to $k_{\max}$. A differentiable Gumbel-Top-$k$ soft mask activates the first $k_\ell$ spectral/tensor components (rank control), and a rank-tied mixed-precision quantizer$Q_{q_\ell}$ with STE rounding applies $q_\ell$-bit quantization (bit-width control) to form compressed weights $\tilde{W}_{q}(k)$. The recomposed operator yields compressed predictions $p_k(x)$, trained using (i) the task loss (CE) and (ii) consistency/self-distillation via $\mathrm{KL}(p_{\text{full}}\|p_k)$ (optionally also under light augmentations). A certificate module estimates the logit-drift$\hat{\Delta}(k)$ and contributes a drift-penalty term (e.g., $\max(0,\hat{\Delta}(k)-\epsilon)$) to the total objective $\mathcal{L}$. For deployment, continuous $(k_\ell,q_\ell)$ choices are snapped to a discrete profile set$\{(k,q)\}_j$ matched to hardware-efficient kernels, enabling predictable runtime behavior. Solid arrows denote forward/data flow; dashed arrows denote control/meta or gradient pathways.
• Figure 2: Certificate bound $\hat{\Delta}(k)$ versus observed logit drift $\|\delta z\|_2$ across discrete profiles $(k,q)$. The shaded region marks a tolerance band around $y{=}x$; dashed lines show its boundaries. Arrows and labels are positioned in axis coordinates and point to representative profiles.
• Figure 3: ImageNet-1k Pareto (A100, p50). T3C produces controllable frontiers (Tiny$\rightarrow$Max) that dominate PTQ/QAT and LR+FT across CNNs and ViTs.
• Figure 4: Ablations at a fixed latency target (lower is better). Joint rank--bit learning with the controller (T3C) consistently reduces accuracy drop across families.

Theorems & Definitions (1)

• Proposition 4.1: Layerwise truncation certificate