Technical Summary: Colinearity Decay for Quantization-Friendly Vision Transformers

1. Problem Statement

Low-bit quantization is a critical strategy for deploying Vision Transformers (ViTs) in resource-constrained environments. However, fully quantized deployment is hindered by activation outliers—a small number of channels or tokens exhibiting significantly larger values than the rest. These outliers force coarse quantization scales, degrading the effective precision for the majority of activations.

Existing approaches generally fall into two categories:

1. Post-training methods: Calibration, scale reparameterization, or reconstruction techniques applied after the full-precision model is trained.
2. Training-time suppression: Methods that aggressively suppress large activations during training (e.g., TWEO [23]) to stabilize low-precision optimization.

The authors argue that the second approach has a critical flaw: aggressively minimizing outliers can harm the representation learned by the full-precision model, leading to a suboptimal trade-off between full-precision (FP) accuracy and quantized accuracy. Furthermore, direct activation regularization often incurs significant training overhead. The paper posits that the training objective should not simply minimize outlier magnitude but should instead control the structural amplification that makes these outliers harmful to quantization.

2. Methodology: Colinearity-Decay (CD)

The proposed method, Colinearity-Decay (CD), is a structural regularizer designed to mitigate the amplification of activations without altering the model architecture or task loss.

2.1 Structural Source of Outliers

The authors identify that large activations often arise not from individual matrices but from the alignment between ordered matrix pairs $(W_1, W_2)$ in the data flow.

• Let $W_1 = USV^\top$ be the SVD of an upstream matrix.
• If a row of the downstream matrix $W_2$ aligns with a dominant left singular vector of $W_1$ (associated with a large singular value), the output is amplified.
• Standard weight decay regularizes matrices independently and fails to penalize this cross-matrix alignment.

2.2 The Regularizer

CD penalizes the pairwise energy $\|W_2 W_1\|_F^2$. This term is large when rows of $W_2$ place mass on high-singular-value directions of $W_1$.

• Target Pairs: CD applies to specific pairs within Transformer blocks:
  • Composable Pairs: Normalization scales followed by projections (e.g., LayerNorm scales and $W_Q, W_K, W_V$ or $W_{F1}$).
  • Functional Pairs: Matrices separated by non-linearities but forming a functional route (e.g., $W_V \to W_O$ in attention, and $W_{F1} \to W_{F2}$ in FFN).
• Decoupled Update: Instead of adding an auxiliary loss to the training graph (which increases memory and time complexity), CD is implemented as a decoupled weight update. After computing task gradients, the downstream matrix $W_2$ is updated as:
  $$W_2 \leftarrow W_2 - \eta \lambda_{cd} W_2 W_1 W_1^\top$$
  This update shrinks $W_2$ rows along the high-singular-value directions of $W_1$. The authors prove that this gradient direction forms a non-obtuse angle with standard weight decay, ensuring they are cooperative.

2.3 Implementation Details

• CD is applied only to the downstream matrix in each pair to avoid redundant decay.
• The regularization budget is managed by reducing the standard weight decay coefficient ($\lambda_{wd}$) to accommodate $\lambda_{cd}$, keeping the total regularization strength comparable to the baseline.
• The method is non-invasive, preserving the inference graph and architecture.

3. Key Contributions

1. Empirical Insight: The authors demonstrate that keeping outliers in a "suitable range" is beneficial. Aggressive suppression (reducing max activations to very low values) harms FP accuracy and, consequently, quantized performance.
2. Structural Regularization: They identify excessive cross-matrix alignment as a structural source of outliers and introduce CD, a non-invasive regularizer with negligible overhead.
3. Orthogonality: CD is orthogonal to existing quantization pipelines and architectural outlier-mitigation methods (e.g., attention sinks), allowing for combined use.
4. Comprehensive Validation: CD is validated across ImageNet-1K pre-training, COCO detection, and nine downstream fine-tuning tasks.

4. Experimental Results

4.1 ImageNet-1K Pre-training

• Quantized Accuracy: CD consistently improves W4A4 quantized accuracy (RepQ-ViT and APHQ) compared to baselines and TWEO. For example, on ViT-B, CD achieves 77.78% (RepQ-ViT) vs. 76.47% (Baseline) and 73.86% (TWEO).
• Full-Precision Accuracy: CD preserves or slightly improves FP Top-1 accuracy (81.63% vs. 81.44% Baseline).
• Outlier Control: CD reduces maximum activations significantly (e.g., from 250.60 to 162.44 in blocks) but not as aggressively as TWEO (18.56), supporting the hypothesis that moderate control is optimal.

4.2 COCO Detection

• CD improves quantized detector performance (mAP) substantially while maintaining or slightly improving FP performance.
• In contrast, TWEO and Noise baselines often incur FP accuracy drops.
• CD shows consistent gains across different quantization pipelines (Percentile, RepQ-ViT, APHQ, GPLQ) and model sizes (Swin-T, Swin-S).

4.3 Downstream Fine-tuning

• Across nine datasets (including Flowers, Pets, Cars, etc.), CD yields the best average FP accuracy, best average quantized accuracy, and the smallest average accuracy drop compared to Baseline and TWEO.
• This confirms that controlling activations to a reasonable range is more effective than minimizing them as much as possible.

4.4 Efficiency

• Overhead: As a decoupled update, CD adds minimal time and memory overhead. In contrast, TWEO's activation regularization loss scales linearly with batch size and hidden dimension, incurring significantly higher costs.

4.5 Ablation Studies

• Strength: A moderate CD strength ($\lambda_{cd} = 0.005$) yields the best results; too strong or too weak is suboptimal.
• Form: The decoupled update form outperforms a direct loss-term implementation, which was found to be unstable.
• Target: Applying CD to the second matrix in both composable and functional pairs yields the best results. Applying it to both matrices in a functional pair degrades performance.

5. Significance and Claims

The paper claims that structural regularization of Transformer matrix pairs is a practical and effective way to prepare vision models for low-bit deployment.

• Zero Inference Overhead: Since CD only modifies training dynamics, inference architecture and latency remain unchanged.
• Better Trade-off: By avoiding the aggressive suppression of outliers that harms full-precision representations, CD achieves a superior balance between FP and quantized accuracy.
• Generalizability: While focused on ViTs, the method is defined on ordered matrix pairs and shows initial promise on GPT-2 language models (Appendix A.6), suggesting applicability beyond vision.

The authors conclude that rather than treating outliers as a problem to be eliminated, they should be viewed as a diagnostic quantity where the goal is to control the structural mechanisms (cross-matrix alignment) that amplify them to harmful levels.