BitMedViT: Ternary-Quantized ViT for Edge Imaging

• BitMedViT is a ternary-quantized Vision Transformer designed for resource-constrained edge deployment in medical imaging, dramatically reducing model size and energy usage.
• It utilizes multi-query attention, abs-mean ternary quantization, and task-aware distillation to deliver high diagnostic performance with low-precision operations.
• Optimized for platforms like NVIDIA Jetson Orin Nano, BitMedViT achieves up to 43× smaller model size and 21.7× faster inference compared to full-precision models.

BitMedViT is a ternary-quantized Vision Transformer (ViT) architecture specifically designed for resource-constrained edge deployment in medical imaging tasks. It combines an efficient quantization scheme, multi-query attention, task-aware distillation, and hardware-optimized inference to deliver high diagnostic accuracy while dramatically reducing model size, memory bandwidth requirements, and energy consumption, particularly on platforms such as the NVIDIA Jetson Orin Nano (Walczak et al., 15 Oct 2025).

1. Architecture and Quantization Strategy

BitMedViT adopts a compact Vision Transformer backbone consisting of 3 layers with an embedding dimension ($E_n$) of 512, 8 attention heads ($H=8$), and a feed-forward expansion factor of 4 (FFN dimension 2048). The standard Multi-Head Self-Attention (MHSA) is replaced by Multi-Query Attention (MQA), wherein the key and value projections are shared across all heads—reducing the number of key and value parameters by $1/H$ with negligible impact on accuracy when trained via distillation from a high-capacity teacher.

All transformer linear layers, including projection matrices in attention and FFN blocks, employ “abs-mean” ternary quantization. Given $W \in \mathbb{R}^{k \times n}$, each entry $w_{ij}$ is quantized as:

$$q(w_{ij}) = \text{RoundClip}\left( \frac{w_{ij}}{\beta}, -1, +1 \right), \quad \beta = \frac{1}{k n} \sum_{i=1}^k \sum_{j=1}^n |w_{ij}|$$

where $\text{RoundClip}(x, -1, 1)$ restricts the rounded value to $[-1, 1]$. These ternary weights ($W_2 \in \{-1, 0, +1\}^{k \times n}$) are stored using 2 bits per parameter.

Input activations are quantized using a per-row “abs-max” scheme: for each input row $p$, a scale $H=8$0 is computed, and activations are quantized to 8-bit signed integers. Matrix multiplications are performed as $H=8$1 operations, with dequantization via per-layer scaling ($H=8$2) applied at each layer output.

In the MQA mechanism, queries remain head-specific and are ternarized identically. All subsequent matrix multiplications in self-attention and FFN blocks are conducted in this 2-bit (weights), 8-bit (activations) regime.

2. Training Pipeline and Task-Aware Distillation

BitMedViT employs a composite multi-term loss for knowledge distillation from a full-precision MedViTv2-L teacher. The loss optimized for each student model is:

$H=8$3

where:

• $H=8$4: Cross-entropy loss on ground-truth labels $H=8$5.
• $H=8$6: Temperature-scaled KL divergence (with $H=8$7) between student and teacher logits.
• $H=8$8: Feature alignment loss, using a projection to align student and teacher intermediate features.

Regularization includes weight decay ($H=8$9), cosine annealing learning rates, and per-layer scale tracking in FP16 to retain stability during back-propagation. Gradient clipping ($1/H$0) is employed to avoid explosion due to the quantizer.

Ablation studies identify the optimal distillation weights as $1/H$1; removing feature-level distillation results in up to 1.8% accuracy degradation.

3. Edge Inference: CUDA Implementation and Deployment

On the Jetson Orin Nano, BitMedViT leverages a custom CUDA kernel pipeline designed for efficient memory and computation. The architecture comprises 4 GPCs and 8 SMs, each with 4 Ampere Tensor Cores supporting int8 WMMA. Memory hierarchies include 192 KB L1 cache per SM, 4 MB shared L2, and 8 GB DRAM.

Each linear transformation is tiled into $1/H$2 blocks for compatibility with the Tensor Core's TMMA (Tensor Matrix Multiply-Accumulate) operations. Sixteen 2-bit ternary weight values are packed into each 32-bit column-major word; activations are similarly packed for memory-efficient transfer.

The custom CUDA kernel implements the Blocked_Ternary_GEMM pattern:

• Loads packed ternary weights from DRAM and unpacks to int8.
• Asynchronously copies 8-bit activations, maximizing pipeline concurrency.
• Executes WMMA (TF32 accumulation) on the Tensor Cores for each tile.
• Applies dequantization using per-tile $1/H$3 scales stored in shared memory.
• Writes output coalesced in BF16.

Integration with TensorRT is accomplished via a custom plugin, enabling mixed W2A8 I/O and dequantization-fused layer fusion. Two kernel variants enable BF16 and legacy FP16 fallback on the Orin Nano.

4. Performance Benchmarks and Ablation

On the MedMNIST collection (12 medical image tasks), BitMedViT attains an average validation accuracy of 86%, compared to 89% for the full-precision MedViTv2-L teacher (3% differential). Test set performance follows a similar trend (82% versus 85%).

Per-dataset results demonstrate high performance on PathMNIST (99.0% validation ACC), OCTMNIST (95.0%), but a larger gap on tasks such as PneumoniaMNIST (86.0% vs. 97.3%) and RetinaMNIST (53.0% vs. 57.8%).

Summary of key compression and efficiency metrics on the Jetson Orin Nano:

Metric	BitMedViT	Teacher (MedViTv2-L)	Fold Change
Model Size	10.5 MB	447.7 MB	43× smaller
Memory Traffic	-	-	39× lower
Inference Latency	16.8 ms	366.6 ms	21.7× faster
Throughput (GOP/s)	683	35	19.5×
Energy Efficiency (GOP/J)	183.6	8.3	22×
Peak Power Draw	3.72 W	4.25 W	1.14× less

Ablation studies indicate that W2A8 (2-bit weights, 8-bit activations) is the optimal trade-off, with W1A8 causing a 5% accuracy drop and W4A8 yielding only minimal accuracy gains at double the memory cost. A kernel tile size of $1/H$4 maximizes DRAM bandwidth utilization by 12% over larger tiles.

5. Technical and Clinical Implications

The BitMedViT design demonstrates the feasibility of deploying ViT-based medical diagnosis models on resource-limited edge devices for real-time operation. The combined use of ternary quantization, low-precision multi-query attention, and targeted distillation enables a $1/H$5 reduction in model size and $1/H$6 decrease in memory bandwidth demand, with only a modest 3% performance decrease relative to state-of-the-art floating-point baselines.

Energy consumption is markedly reduced, with BitMedViT achieving $1/H$7 GOP/J compared to $1/H$8 GOP/J for the teacher, and inference latency reduced below 17 ms. These results underscore the suitability of such architectures for integration into clinical workflows on portable, low-power platforms.

6. Limitations and Future Directions

BitMedViT incurs a measurable accuracy gap (up to 3%) relative to full-precision ViT baselines, with the impact more pronounced on some datasets (e.g., PneumoniaMNIST, RetinaMNIST). While the extreme model and memory compression are beneficial for edge inference, tasks with higher data complexity or noise sensitivity may experience a disproportionate loss in diagnostic accuracy.

A plausible implication is that future research may explore hybrid quantization, dynamic bit-width selection, or improved feature alignment strategies to further mitigate these deficits without sacrificing deployment efficiency. The co-optimization of hardware and algorithm remains crucial in translating academic advances into robust, clinically deployable edge AI systems.

(Walczak et al., 15 Oct 2025)