BiteAI: Attention-Guided Distillation and Weight-Only Quantization for Compact Insect-Bite Classification

Abstract

Insect bites are a common cause of skin irritation and can contribute to disease transmission through vector-borne pathogens. Early identification of the likely biting organism can assist preliminary guidance (e.g., monitoring for warning signs, considering exposure history) and may reduce complications through timely follow-up. This paper studies a compact attention-guided learning framework for multiclass insect-bite image classification under strict storage constraints. A teacher network (BiteAI-T) based on MobileNetV3-Small is trained with spatial attention pooling to emphasize lesion-relevant regions while maintaining an efficient backbone. A lightweight depthwise-separable student (BiteAI-S) is trained using multi-level knowledge distillation that combines softened-logit matching with intermediate supervision through attention-map alignment and pooled-feature matching. Model storage is further reduced through weight-only quantization-aware training using an LSQ-inspired learnable scaling factor; BatchNorm running statistics are frozen during quantization fine-tuning to improve stability. Experiments on an eight-class dataset (ants, bed bugs, chiggers, fleas, mosquitos, no bites, spiders, ticks) show that BiteAI-T reaches 93.75% test accuracy. For deployment, we export (i) a TorchScript Lite teacher artifact (BiteAI-T_Lite, 2.35 MB) and (ii) a weight-only int8 student artifact (BiteAI-S_int8, 0.992 MB). Comparative results are also reported for an SVD-compressed + fine-tuned FP16 variant (92.66% test accuracy, 2.84 MB), illustrating accuracy–size trade-offs across compression strategies.

1. Introduction

Insect bites and stings are a frequent cause of dermatological complaints and may be associated with allergic reactions, secondary bacterial infection, or exposure to vector-borne diseases. In many settings, the first step of bite assessment relies on visual inspection of the lesion pattern (e.g., clustered papules, erythema, wheals) together with contextual information such as geographic region, seasonality, and exposure history. However, accurate recognition is difficult in practice because bite appearance varies with skin tone, time since the bite, scratching, and imaging conditions (illumination, camera pipeline, background, distance to skin). These factors lead to substantial variability within the same class and visual overlap between different insects.

Automated insect-bite classification from images can support preliminary screening in lightweight applications and assist non-experts in organizing observations (e.g., identifying likely classes, deciding whether to seek medical advice, and monitoring changes over time). Deep learning has achieved strong performance in related dermatological recognition tasks, but many high-performing approaches employ large backbones or ensembles. Such designs are often impractical when the target constraint is strict model storage, low-latency inference, or deployment inside software environments where memory and packaging limits are significant.

Recent work on insect-bite recognition has started exploring deep-learning solutions, but most designs are not explicitly optimized for strict storage budgets. For example, DeepBiteNet relies on an ensemble of heterogeneous backbones to improve robustness, which increases footprint and deployment complexity [1]. In parallel, the broader skin-lesion literature (e.g., ISIC/HAM10000) has produced strong recognition pipelines, yet many high-performing solutions are based on heavier backbones and ensembles that assume GPU-class resources [2,3,4]. These trends highlight a practical gap: achieving reliable multiclass bite recognition with attention/interpretability signals under sub-megabyte to few-megabyte storage constraints.

Compression-aware learning is therefore essential for practical insect-bite recognition. Knowledge distillation is a key technique for transferring discriminative behavior from a higher-capacity teacher to a smaller student that might otherwise underfit [5,6]. In parallel, attention mechanisms can strengthen performance by emphasizing lesion-relevant spatial regions and can provide coarse interpretability through attention maps [7,8,9]. The combination of attention-guided representation learning, multi-level distillation, and stable quantization remains challenging when the student must satisfy extreme storage constraints (sub-megabyte footprints) while retaining meaningful accuracy.

1.1. Scope and Generalization

This study is evaluated on a single public eight-class dataset with predefined splits, and therefore the reported performance primarily reflects in-distribution. generalization under the dataset’s acquisition conditions. In real-world usage, insect-bite appearance may vary with skin tone, lesion stage (time since bite), scratching, comorbid skin conditions, and capture pipeline (illumination, device sensor, compression). Domain shift and dataset bias are widely recognized challenges in dermatology-related vision systems [10,11,12]. Accordingly, this paper does not claim cross-dataset or clinical generalization; instead, the objective is to demonstrate a storage-efficient attention-guided teacher–student framework and quantify accuracy–size trade-offs within the available dataset, together with transparent reporting of dataset structure, preprocessing, and limitations.

This paper introduces BiteAI, a compact attention-guided framework for insect-bite recognition under strict model storage limits. A MobileNetV3-Small teacher (BiteAI-T) is trained with spatial attention pooling to capture discriminative local cues while remaining computationally efficient. Knowledge is transferred into a compact depthwise-separable student (BiteAI-S) through multi-level distillation that aligns (i) class logits, (ii) attention maps, and (iii) pooled feature representations. To reduce storage without modifying the network topology, the distilled student is then optimized using weight-only quantization-aware training with learnable scaling factors, enabling a weight-only int8 export under 1 MB [13,14,15].

1.2. Main Contributions

• An attention-guided teacher–student framework for insect-bite image classification based on compact convolutional architectures.
• A multi-level distillation objective combining softened-logit matching with attention-map and feature-level alignment to improve learning under severe capacity constraints.
• An LSQ-inspired weight quantization-aware training procedure producing a sub-megabyte weight-only int8 representation while preserving a substantial fraction of teacher accuracy.
• A deployment-oriented analysis that includes a reject option (abstention by confidence thresholding) and a transparent discussion of scope and limitations under domain shift.
• A comparative evaluation against an alternative compression pathway (SVD factorization + fine-tuning) to contextualize accuracy–size trade-offs.

2. Related Work

2.1. Insect-Bite Image Classification

Compared with general dermatological image analysis, insect-bite recognition remains relatively underexplored and presents distinct challenges. Intra-class variability is driven by skin tone diversity, changes in lesion appearance over time, scratching and secondary irritation, and imaging variation (lighting, camera quality, compression artifacts). Inter-class similarity is also common because different bites can produce clinically similar erythematous papules or clustered patterns, especially among insects that tend to bite multiple times.

DeepBiteNet is a representative deep-learning approach for multiclass insect-bite recognition that uses a stacked ensemble of heterogeneous CNN backbones (including DenseNet121, EfficientNet-B0, and MobileNetV3) and a meta-classifier to improve robustness. Although ensembling improves accuracy, it increases model footprint and inference complexity, which is undesirable under strict packaging and memory constraints [1]. Related work also explores bug-bite classification as an end-to-end system using transfer learning and classifier fusion/ensembling, supporting the feasibility of bite categorization while typically not targeting sub-megabyte deployment artifacts [16].

Table 1 summarizes representative insect-bite/skin-lesion recognition and efficiency-oriented references used to contextualize our approach.

Table 1. Representative insect-bite/skin-lesion recognition and efficiency-oriented references (NR: not reported).

Study	Dataset	Backbone	Efficiency	Key Result(s)	Diff. vs. Ours
Khasanov et al. (2025) [1]	bug-bite (NR)	Ensemble (heterogeneous CNNs)	Ensemble design	Robust multiclass bite recognition (NR)	Higher footprint; no sub-MB target
Tschandl et al. (2018) [2]	HAM10000	Various CNNs	Benchmarking	Large-scale lesion benchmark	Different task; not bite-specific
Codella et al. (2019) [3]	ISIC 2018	Various CNNs	Challenge setup	Standardized evaluation protocols	Different task; often heavy backbones
Hinton et al. (2015) [5]	general (NR)	KD framework	Distillation	Logit distillation	BiteAI adds attention/feature alignment
Howard et al. (2019) [13]	multi-task (NR)	MobileNetV3	Mobile NAS	Strong accuracy–latency trade-off	Used here as compact teacher
Ilijoski et al. (2023) [16]	bug-bite (NR)	Transfer learning + fusion	Mobile-oriented system	End-to-end bite classification system (NR)	Larger footprint; no attention-aligned KD

2.2. Lightweight CNNs and Compact Inference

Efficient CNN design has been extensively studied for mobile and embedded vision [17,18]. MobileNet introduced depthwise-separable convolutions to reduce parameters and multiply–add operations. ShuffleNet improved efficiency further using pointwise group convolution and channel shuffling to preserve information flow under constrained budgets. EfficientNet proposed compound scaling to jointly scale depth, width, and resolution under a principled budget, producing model families optimized for different deployment constraints. MobileNetV3 combined hardware-aware neural architecture search with architectural refinements to improve accuracy–latency trade-offs on mobile CPUs [13,19,20,21].

2.3. Attention Mechanisms and Interpretability

Attention mechanisms enhance discriminative learning by emphasizing informative spatial regions and feature channels. This is particularly relevant for skin imagery, where class evidence is often localized around a lesion. CBAM is a lightweight module that sequentially applies channel and spatial attention and can be inserted into CNN backbones with minimal overhead, often improving performance. Attention visualizations can also provide coarse interpretability by indicating which regions contributed most strongly to a prediction [7].

Beyond lightweight attention blocks, localization-oriented pooling mechanisms are widely used to extract class-relevant spatial evidence with minimal architectural overhead. Class activation mapping (CAM) and its gradient-based variants (Grad-CAM) are representative approaches that highlight discriminative regions and are frequently used for qualitative inspection in medical and dermatological imaging pipelines. While CAM/Grad-CAM are typically used as post-hoc explanations, attention pooling can be integrated directly into the forward path to produce an explicit spatial weighting that is optimized end-to-end for recognition. This design is well aligned with strict footprint constraints because it adds only a small $1\times 1$ convolution and a pooling operation, and it also produces attention maps that can be used as an intermediate target during distillation [8,9].

2.4. Model Compression, Distillation, and Quantization

Knowledge distillation transfers information from a high-capacity teacher to a smaller student by encouraging the student to match softened teacher outputs. This strategy is especially useful for compact models in medical imaging, where small networks may underfit fine-grained patterns [5].

Quantization-aware training complements distillation by reducing storage and potentially improving inference efficiency through low-precision weights. LSQ introduced learnable step sizes and training procedures that preserve accuracy at low bit-widths. In applied pipelines, LSQ-inspired fake quantization with learnable scaling is frequently used to obtain stable training and compact representations [15,22].

2.4.1. Distillation Beyond Logits

While classical distillation matches softened logits, many works show that compact students benefit from intermediate supervision, especially when the student capacity is severely constrained. FitNets introduced hint-based training using intermediate feature regression to guide thin networks. Attention Transfer (AT) aligns spatial attention summaries between teacher and student and is particularly relevant when the task evidence is localized, as in skin imagery. More generally, feature-level and relational distillation variants aim to reduce representation mismatch and stabilize optimization for small students. These findings motivate the multi-level objective used in BiteAI, where logits, attention maps, and pooled descriptors are aligned jointly [6,23].

2.4.2. Quantization-Aware Training

Integer quantization is a common strategy for shrinking model storage and enabling efficient deployment. Practical QAT methods simulate quantization during training to reduce the accuracy loss observed in post-training quantization, and they are often combined with per-channel or per-layer scaling factors to better match weight distributions. Integer-arithmetic-only inference has been demonstrated for CNNs using quantization with learned or calibrated scales, and QAT has become a standard approach for maintaining accuracy under low precision. In addition, recent work highlights that rounding and scale selection are key contributors to the final accuracy after quantization, motivating careful training-time handling of scale parameters. These results support the LSQ-inspired learnable scaling used here, as well as the choice to keep activations in full precision in order to simplify training and reduce accuracy degradation under a strict size budget [14,24,25].

2.4.3. Dataset Diversity and Domain Shift in Skin Imagery

Recent studies in dermatology AI emphasize that performance can degrade under shifts in skin type distribution and acquisition settings, and that dataset reporting practices often lack sufficient demographic and capture metadata [10,11,12]. These findings motivate transparent reporting of dataset structure, split protocols, and limitations when evaluation is performed on a single dataset.

3. Proposed Method

3.1. Overview

The proposed framework, termed BiteAI, follows a teacher–student paradigm designed to satisfy strict model size constraints. The pipeline contains three stages: (1) training an attention-guided teacher network (BiteAI-T), (2) training a compact student (BiteAI-S) using multi-level knowledge distillation, and (3) compressing the student through weight-only quantization-aware fine-tuning to obtain the deployment artifact (BiteAI-S_int8).

The teacher network uses MobileNetV3-Small as an efficient backbone and adds spatial attention pooling to strengthen localization of discriminative lesion regions. The student network is a smaller depthwise-separable CNN that uses the same attention pooling mechanism to enable architectural alignment during distillation. After distillation, weight-only quantization-aware training is applied to the student to produce a compact weight-only int8 export without changing the model topology.

3.1.1. Process and Data Flow (End-to-End)

For an input RGB image $x\in {\mathbb{R}}^{3\times 256\times 256}$, both teacher and student compute an intermediate feature map $F\in {\mathbb{R}}^{C\times H\times W}$. A spatial attention map $A\in {\mathbb{R}}^{H\times W}$ is produced by a $1\times 1$ convolution followed by a spatial softmax. Two pooled descriptors are then formed: global average pooling $\mathbf{g}\in {\mathbb{R}}^C$ and attention-weighted pooling $\mathbf{l}\in {\mathbb{R}}^C$. The concatenated descriptor $\mathbf{f}=[\mathbf{g};\mathbf{l}]\in {\mathbb{R}}^{2C}$ is passed to the classifier head to obtain logits. During training, BiteAI-Tis optimized with supervised learning and attention regularization; BiteAI-S is trained with supervised loss plus distillation losses aligning logits, attention maps, and pooled descriptors. Finally, weight-only QAT is applied to BiteAI-S to produce BiteAI-S_int8 by storing int8 weights and learned scales.

3.1.2. Design Rationale Under Strict Footprint Constraints

The methodology choices in BiteAI are driven by a size-first deployment target. MobileNetV3-Small is selected as the teacher backbone because it offers a strong accuracy–efficiency trade-off from mobile-oriented architectural search while remaining substantially smaller than common dermatology backbones [13]. For the student, depthwise-separable blocks preserve the MobileNet design philosophy and provide a controllable parameter budget [19]. Attention pooling is used instead of inserting heavier attention blocks because it provides an explicit spatial weighting with minimal parameter overhead and produces attention maps that are naturally compatible with attention-based distillation [23]. Finally, weight-only QAT is favored over full integer activation quantization because it delivers most of the storage benefit (int8 weights) while avoiding additional accuracy loss and engineering complexity from activation quantization in small models [14].

3.1.3. Novelty Relative to Prior Work

The novelty of BiteAI is not the use of distillation or quantization in isolation, but their integration under a strict storage objective through (i) attention pooling as a lightweight, explicit spatial interface; (ii) multi-level distillation that aligns logits and intermediate attention/pooled descriptors to stabilize learning for an extremely compact student; and (iii) a weight-only, LSQ-inspired QAT recipe with frozen BatchNorm statistics that enables a sub-megabyte export while retaining practical accuracy.

3.2. Attention-Guided Feature Aggregation

Let $F\in {\mathbb{R}}^{C\times H\times W}$ denote an intermediate feature map produced by a convolutional backbone. Two complementary descriptors are extracted

$$\mathbf{g}=\mathrm{GAP}\left(F\right)\in {\mathbb{R}}^C,$$

which summarizes global context via global average pooling, and a spatial attention distribution

$$A=\mathrm{softmax}\left({\mathrm{Conv}}_{1\times 1}\left(F\right)\right)\in {\mathbb{R}}^{H\times W}.$$

The attention distribution assigns higher weight to spatial locations that are more informative for classification. A localized descriptor is computed by attention-weighted pooling:

$$\mathbf{l}=\sum _{i,j}{A}_{i,j}·F_{:,i,j}.$$

The final representation concatenates global and local descriptors,

$$\mathbf{f}=[\mathbf{g};\mathbf{l}]\in {\mathbb{R}}^{2C},$$

and is fed to a lightweight classifier head. An entropy regularization term on A is included to discourage degenerate attention (e.g., overly diffuse distributions) and to stabilize training by promoting consistent spatial focus.

Why Attention Pooling (Rather than Adding Attention Blocks)

A common alternative to spatial attention pooling is to inject lightweight attention modules such as SE or CBAM inside the backbone [7,26]. While effective, these modules introduce additional parameters and intermediate activations across multiple stages. In contrast, the proposed attention pooling is concentrated at the aggregation stage: it encourages the network to form a spatial evidence distribution that directly controls pooling, yielding a compact representation $[\mathbf{g};\mathbf{l}]$ with a clear interpretation. This choice is particularly suitable for distillation because it exposes an explicit attention map A that can be aligned between teacher and student, which is consistent with prior findings that attention alignment can improve student learning when capacity is limited [23].

3.3. Training Objective and Attention Entropy Regularization

The model is trained end-to-end using cross-entropy loss with label smoothing, together with an attention entropy penalty that encourages spatially focused attention maps. Let $\mathbf{a}\in {\mathbb{R}}^{HW}$ denote the spatial attention distribution for an input image. The attention entropy term is

$${\mathcal{L}}_{attn}=-\sum _{u=1}^{HW}{a}_{u}log(a_u+ϵ),$$

and the overall objective is

$$\mathcal{L}={\mathcal{L}}_{CE}+\lambda {\mathcal{L}}_{attn},$$

where $\lambda$ controls the strength of the attention regularization and $ϵ$ is a small constant for numerical stability.

3.4. Student Network Architecture

The student model is a compact depthwise-separable convolutional network designed for parameter efficiency. It includes a convolutional stem followed by multiple depthwise-separable blocks with increasing channel widths. Depthwise separable convolutions reduce computation by decomposing a standard convolution into a depthwise spatial convolution followed by a pointwise ($1\times 1$) projection. The attention pooling module is shared with the teacher to maintain a comparable representation interface, which simplifies intermediate alignment during distillation.

3.5. Multi-Level Knowledge Distillation

Let ${\mathbf{z}}_s$ and ${\mathbf{z}}_t$ denote student and teacher logits, respectively. The standard logit distillation term is defined as

$${\mathcal{L}}_{KD}=T^2·\mathrm{KL}\left(\mathrm{softmax}({\mathbf{z}}_t/T)\parallel \mathrm{softmax}({\mathbf{z}}_s/T)\right),$$

where T is a temperature parameter controlling the softness of the probability distribution. This form corresponds to the standard teacher→student distillation objective, where the student is optimized to match the teacher soft targets.

In addition to logit matching, intermediate behavior is aligned through two mean-squared-error terms: (i) attention-map alignment between teacher and student attention distributions, and (ii) pooled-feature alignment between teacher and student descriptors after attention pooling. This multi-level supervision constrains the student not only to predict similar outputs, but also to attend to similar regions and to build comparable pooled representations, which is important when capacity is limited.

Ablating the Roles of the Three Distillation Signals

The three components of the distillation objective serve complementary purposes. Logit matching primarily transfers inter-class similarity structure (“dark knowledge”), which is valuable when the dataset is small or labels are visually ambiguous [5]. Attention-map alignment encourages the student to focus on lesion-relevant regions rather than spurious background cues, which is important in uncontrolled capture conditions. Pooled-feature alignment constrains the student representation after attention pooling, reducing the risk that the student matches logits using a qualitatively different and less transferable internal representation. This division of roles follows the broader distillation literature where intermediate alignment (hints/attention) improves thin-network optimization beyond what logits alone can provide [6,23].

3.6. Complete Objective and Variable Definitions

For an input image x with class label $y\in \{1,\dots ,C\}$, the teacher and student produce logits ${\mathbf{z}}_t\left(x\right),{\mathbf{z}}_s\left(x\right)\in {\mathbb{R}}^C$. The attention pooling module produces a normalized spatial attention map $A\left(x\right)\in {\mathbb{R}}^{H\times W}$ with ${\sum }_{i,j}{A}_{i,j}=1$, and a pooled descriptor $\mathbf{f}\left(x\right)=[\mathbf{g}\left(x\right);\mathbf{l}\left(x\right)]\in {\mathbb{R}}^{2C^{\prime }}$, where $C^{\prime }$ is the channel dimension at the pooling layer. The student is optimized with a weighted sum:

$${\mathcal{L}}_{student}=\alpha {\mathcal{L}}_{CE}(y,{\mathbf{z}}_s)+\beta {\mathcal{L}}_{KD}({\mathbf{z}}_s,{\mathbf{z}}_t;T)+\gamma {\mathcal{L}}_A(A_s,A_t)+\delta {\mathcal{L}}_F({\mathbf{f}}_s,{\mathbf{f}}_t),$$

where ${\mathcal{L}}_{CE}$ is cross-entropy, ${\mathcal{L}}_{KD}$ is KL divergence between temperature-softened outputs, ${\mathcal{L}}_A$ is an MSE alignment between teacher and student attention maps, and ${\mathcal{L}}_F$ is an MSE alignment between pooled descriptors. The scalar weights $\alpha ,\beta ,\gamma ,\delta \ge 0$ balance the supervised and distillation components.

3.7. Quantization-Aware Training

To reduce storage, the distilled student is fine-tuned using weight-only quantization-aware training (QAT). Convolutional and linear layers are replaced by quantization-simulating counterparts that apply symmetric fake quantization to weights with learnable scaling factors, while activations remain in full precision. For a weight tensor W and a learnable per-layer step size (scale) $s>0$, we use signed b-bit symmetric quantization with the integer range

$$Q=2^{b-1}-1(\mathrm{we}\mathrm{use}b=8).$$

Fake quantization is implemented by a quantize–dequantize step in the forward pass

$$\begin{array}{cc}\hfill \widehat{W}& =\mathrm{clip}\left(⌊{\displaystyle \frac{W}{s}}⌉,-Q,Q\right),\hfill \end{array}$$

(1)

$$\begin{array}{cc}\hfill \tilde{W}& =s·\widehat{W},\hfill \end{array}$$

(2)

where $\lfloor ·\rceil$ denotes rounding to the nearest integer and $\mathrm{clip}(·)$ clamps to the valid integer range. The layer output is then computed using $\tilde{W}$ (with full-precision activations), e.g.,

$$y=\tilde{W}x+b.$$

Quantization is simulated during training using a straight-through estimator (STE) so that gradients can pass through the rounding operation, i.e., we approximate the derivative of rounding by identity inside the clipping range:

$${\displaystyle \frac{\partial \tilde{W}}{\partial W}}\approx \mathbf{1}\left(\left|{\displaystyle \frac{W}{s}}\right|\le Q\right).$$

The procedure is LSQ-inspired because it learns per-layer scaling parameters s, but it does not implement the full LSQ gradient-scaling rule for the step size; the method is therefore described as LSQ-like. Batch normalization layers are kept in evaluation mode during quantization fine-tuning, freezing running statistics to reduce training instability caused by distribution shifts introduced by quantized weights. After convergence, the model is exported by storing int8 weights $\widehat{W}$ along with the corresponding scaling factors s and biases, yielding a compact weight-only representation (0.992 MB in the reported experiments).

3.7.1. Evaluation Note (QAT vs. Exported Artifact)

In this work, the reported accuracy for BiteAI-S with QAT refers to evaluation with quantization simulation enabled (fake-quantized weights with full-precision activations) using the same learned scales that are stored in the weight-only int8 export. The exported BiteAI-S_int8 artifact stores the resulting int8 weights and scales for deployment; full integer-kernel evaluation on target devices is outside the scope of this study and is left for future work.

Why Freeze BatchNorm Statistics During QAT

When weights are discretized, the induced distribution shift can destabilize BatchNorm running statistics if they continue to update, especially under small batch sizes or imbalanced datasets. Freezing BatchNorm (keeping BN layers in evaluation mode) is a pragmatic stabilization strategy: it prevents drift in running mean/variance and isolates the effect of quantized weights during fine-tuning. This is consistent with common deployment practice where BatchNorm is typically folded into preceding convolutions for inference, and with QAT pipelines that treat normalization statistics carefully to reduce quantization-induced mismatch [14,24].

4. Experimental Results

4.1. Data Preprocessing and Split Construction

The dataset used in this work is obtained from Kaggle (“Bug Bite Images”) [27]. Table 2 summarizes the dataset structure, preprocessing, and evaluation protocol.

Table 2. Dataset description and evaluation protocol for the Kaggle Bug Bite Images dataset [27].

Item	Details
Dataset name	Bug Bite Images (Kaggle)
Kaggle identifier	moonfallidk/bug-bite-images
Task	Multi-class insect-bite image classification
Classes ($C=8$)	ants; bed_bugs; chiggers; fleas; mosquitos; no_bites; spiders; ticks
Data organization	Folder-structured splits: training/, testing/, skin_of_color_testing/
Validation protocol	20% stratified split from training/ (VAL_SPLIT = 0.2, SEED = 1337)
In-distribution test	testing/ copied to dataset/test/
OOD/domain shift	skin_of_color_testing/ copied to dataset/test_skin/
Preprocessing	Resize to $256\times 256$; ImageNet normalization
Augmentation	Training-only offline augmentation for class balancing (saved images); in addition, the teacher training uses online mixup. No augmentation is applied to val/test/test_skin.
Metrics reported	Top-1 accuracy (primary); confusion matrix and qualitative examples for error interpretation

4.1.1. Dataset Source, Structure, and Preprocessing Details

The raw dataset is organized under data/ with three predefined subsets: training/, testing/, and skin_of_color_testing/. Each subset follows a class-folder structure, and the canonical class list is defined from the folder names in data/training/. Only common image formats are retained (e.g., JPG/PNG/TIFF/WebP). The dataset release does not provide explicit metadata describing skin type, lesion stage, anatomical site, or acquisition device. Therefore, the present study uses image content only, and conclusions are restricted to in-distribution performance under this dataset’s conditions. This is consistent with broader dermatology AI findings that missing demographic/capture metadata limits claims about fairness and cross-domain generalization [10,11,12].

4.1.2. Reproducible Stratified Split

A validation set is created from data/training/ to avoid contaminating the held-out test subsets. The partition uses a stratified per-class split with VAL_SPLIT = 0.2 and SEED = 1337. For each class c with image set ${\mathcal{D}}_c$, indices are shuffled with the fixed seed and

$$k_c=⌊0.2·|{\mathcal{D}}_c|⌋$$

samples are assigned to dataset/val/, with the remainder assigned to dataset/train/. The predefined subsets data/testing/ and data/skin_of_color_testing/ are copied, unchanged, to dataset/test/ and dataset/test_skin/, respectively. In this study, test_skin is retained as an out-of-distribution (OOD) proxy split provided by the dataset; due to its small size and missing metadata, it is used for descriptive reporting only and is not treated as sufficient evidence for demographic fairness or clinical generalization. A manifest file (dataset/manifest.json) records the class list, seed, split ratio, and split provenance to support reproducibility.

4.1.3. Training-Only Class Balancing

Class imbalance is mitigated using offline augmentation applied only to dataset/train/. For any class c with $n_c<N_{\mathrm{target}}$ (TARGET_COUNT = 500), additional samples are generated by repeatedly sampling existing training images from the same class and applying one randomly selected transformation: horizontal flip, rotation $\{90°,180°,270°\}$, color jitter, contrast jitter, additive Gaussian noise, or zoom-in via center crop and resize. Augmented images are saved in the same class directory with an aug_ prefix to preserve traceability.

4.2. Dataset and Experimental Setup

Experiments are conducted on an eight-class insect-bite dataset comprising ants, bed_bugs, chiggers, fleas, mosquitos, no_bites, spiders, and ticks. The Kaggle dataset provides predefined folders training/, testing/, and skin_of_color_testing/. We construct dataset/train/ and dataset/val/ by applying a stratified split (VAL_SPLIT = 0.2, SEED = 1337) to training/. Figure 1 and Figure 2 report the class distribution of dataset/train/ (i.e., the subset actually optimized by SGD) before and after offline augmentation for balancing. The in-distribution test set is taken from testing/ (copied unchanged to dataset/test/), and an additional small OOD subset test_skin is taken from skin_of_color_testing/ for future domain-shift analysis.

Figure 1. Dataset distribution in dataset/train/ before offline augmentation.

Figure 2. Dataset distribution in dataset/train/ after offline augmentation (balanced to 500 images/class).

All models are implemented in PyTorch. Images are resized to $256\times 256$ and normalized with ImageNet mean and standard deviation. Optimization uses AdamW with an initial learning rate of $3\times {10}^{-4}$, batch size 32, and cosine annealing. The BiteAI-T training includes label smoothing and online mixup to improve generalization on the curated distribution. The BiteAI-S training uses standard cross-entropy without label smoothing. Since offline augmentation balances dataset/train/ to 500 images per class, uniform class weights are used in practice for BiteAI-S training. During quantization-aware fine-tuning, weight-only LSQ-inspired fake quantization is enabled and BatchNorm layers are kept in evaluation mode to freeze running statistics [14,15].

4.3. Overall Performance, Baselines, and Model Size

Table 3 summarizes the main accuracy–size trade-offs for models evaluated in FP32, alongside serialized TorchScript FP16 footprints when applicable. The best BiteAI-T checkpoint reaches 94.53% validation accuracy and 93.75% test accuracy. Low-rank SVD factorization of $1\times 1$ and linear layers followed by fine-tuning reaches 94.84% validation accuracy and 92.66% test accuracy with a TorchScript FP16 size of 2.84 MB.

Table 3. Accuracy and storage footprint for evaluated models (input $256\times 256$).

Model	Val Acc (%)	Test Acc (%)	Size/Format
BiteAI-T (MobileNetV3-Small + attention pooling)	94.53	93.75	∼4.90 MB (FP32)
SVD + Fine-tune (FP32 eval)	94.84	92.66	2.84 MB (TS FP16)
BiteAI-S (Tiny depthwise, FP32)	68.91	67.50	∼3.71 MB (FP32)
BiteAI-S (KD + weight-only QAT simulation)	86.56	83.75	∼3.71 MB (FP32 ckpt)

Direct supervised training of the compact student underperforms substantially (68.91% validation accuracy), indicating that the restricted capacity limits direct optimization. Multi-level distillation combined with weight-only quantization-aware training improves the student to 86.56% validation accuracy and 83.75% test accuracy under QAT simulation (see Section 3.7). Deployment artifacts produced through export and packaging are summarized separately in Table 4.

Table 4. Deployment artifacts produced by compression/export.

Artifact	Size/Format	Notes
BiteAI-TLite	2.35 MB (.ptl)	TorchScript Lite package
BiteAI-Sint8 (weight-only)	0.992 MB	int8 weights + scaling factors

Interpreting the int8 Student Accuracy Drop

The weight-only int8 student achieves the smallest footprint but exhibits lower accuracy than the FP32 teacher and FP16-compressed variants. This reduction is expected because discretizing weights perturbs learned filters, and weight-only export prioritizes storage over maximal accuracy retention. A practical mitigation is to operate conservatively using the reject option (Section 4.6), abstaining on low-confidence predictions instead of forcing a potentially misleading label.

To enable on-device inference under strict storage constraints, we export two compressed deployment artifacts. First, the Lite artifact (BiteAI-T_Lite, 2.35 MB, .ptl) is obtained by converting the trained teacher into a TorchScript Lite package for mobile-friendly execution. Second, the distilled student and learned scales are exported using weight-only int8 quantization, producing BiteAI-S_int8 (0.992 MB) by storing quantized weights together with the corresponding scaling factors. Together, these artifacts represent (i) a compact model suitable for direct deployment and (ii) an ultra-compact student model optimized for minimal footprint.

4.4. Ablation Study and Component-Wise Analysis

Because the target deployment constraint is primarily storage, the ablation analysis is organized around the sequential pipeline used to obtain the final artifact: (i) direct training of a compact student, (ii) progressively richer distillation, and (iii) quantization-aware fine-tuning for weight-only export. Rather than treating these steps as independent options, we analyze them as interdependent choices that trade accuracy for size and training stability.

4.4.1. Supervision vs. Distillation

The supervised tiny student (trained without teacher guidance) underperforms substantially (68.91% validation accuracy in Table 3). This gap indicates that the student architecture is capacity-limited for direct empirical risk minimization on this dataset. Introducing distillation changes the optimization target from “fit labels only” to “fit labels plus teacher structure”, which typically improves sample efficiency and decision boundary smoothness. In BiteAI, the best student results are achieved only when distillation is used, supporting the conclusion that knowledge transfer is not a marginal improvement but a requirement under the chosen footprint constraint [5,6].

4.4.2. Logit-Only vs. Multi-Level Distillation

Logit matching transfers class similarity structure but does not explicitly constrain where the model should look. The additional attention-map alignment addresses this limitation by encouraging spatial evidence consistency between teacher and student, which is particularly relevant in bite imagery where lesion evidence is localized and backgrounds vary. Pooled-feature alignment further constrains the representation after attention pooling, reducing representational drift and making training more stable for very small students. Overall, the multi-level objective is designed to reduce two common failure modes in tiny models: (i) learning spurious background shortcuts and (ii) matching outputs with poorly transferable internal features [9,23].

4.4.3. Distillation-Only vs. Distillation + Weight-Only QAT

After distillation, quantization-aware fine-tuning is applied to reach the final storage target (0.992 MB int8 weight-only export). Weight discretization inevitably perturbs learned filters; QAT mitigates this by exposing the network to fake-quantized weights during training. The reported results show that the final quantized student retains meaningful accuracy (83.75% on the test split under QAT simulation) while meeting the sub-megabyte constraint, demonstrating that the proposed LSQ-inspired scaling strategy is effective in this low-capacity regime [14,15,25].

4.4.4. Training Stability Choices in QAT

Two practical choices are emphasized for stability: learnable per-layer scaling and frozen BatchNorm statistics. Learnable scales help adapt quantization ranges to layer-wise weight distributions, while freezing BN running statistics prevents quantization-induced distribution shift from corrupting normalization estimates during fine-tuning. These measures are especially important in small-data and class-imbalanced settings, where BN statistics are already noisy and can become a major source of instability once weights are discretized [14,24].

4.4.5. Connection to Alternative Compression Routes (SVD)

SVD factorization modifies the compute graph by replacing certain layers with low-rank decompositions, which can reduce FLOPs but does not necessarily reduce serialized size under FP16 TorchScript in the tested configuration. The comparison in Table 3 indicates that SVD decreases test accuracy relative to BiteAI-T (92.66% vs. 93.75%), highlighting that storage-driven compression often benefits more from precision reduction (int8 weights) than from low-rank factorization unless ranks are aggressively tuned and serialization overheads are carefully controlled.

4.5. Confusion Matrix Analysis (Deployment Artifact)

Figure 3 presents the confusion matrix of the deployed teacher artifact BiteAI-T_Lite evaluated on the test split. This artifact is obtained via TorchScript Lite export for mobile-friendly execution; as a deployment-oriented representation, its evaluation is reported separately from the FP32 teacher checkpoint. From the confusion matrix, BiteAI-T_Lite achieves 590/640 correct predictions on the test split (92.19% top-1 accuracy), indicating that most errors occur among visually similar bite patterns.

Figure 3. Confusion matrix of the deployed teacher artifact BiteAI-T_Lite evaluated on the test set.

Most classification errors occur between visually similar bite patterns. In particular, ants and bed_bugs exhibit occasional mutual confusion, while chiggers and bed_bugs are sometimes misclassified, reflecting similarities in clustered bite appearances. Minor confusion is also observed between fleas and neighboring classes, which is consistent with overlapping visual characteristics such as lesion size and distribution.

4.6. Selective Prediction (Reject Option)

A reject option is evaluated by discarding predictions whose maximum class probability falls below a confidence threshold ${\tau }_{\mathrm{conf}}$. Table 5 reports coverage and accuracy computed on the non-rejected subset. Higher thresholds reduce coverage but increase accuracy on accepted predictions, reaching 98.8% at ${\tau }_{\mathrm{conf}}=0.9$.

Table 5. Reject-option analysis on the validation set.

${\mathit{\tau }}_{\mathbf{conf}}$	Coverage	Acc. (Accepted Only)
0.50	0.895	0.972
0.60	0.891	0.974
0.70	0.855	0.984
0.80	0.784	0.986
0.90	0.520	0.988

4.7. Qualitative Evaluation

To qualitatively evaluate the proposed models, inference was performed on randomly selected images from the held-out test split. For each sample, the ground-truth label was displayed alongside the class predicted by the deployed system. This procedure provides an intuitive assessment of model behavior beyond aggregate metrics and highlights its ability to identify insect-bite categories in individual cases. Figure 4 illustrates several example predictions obtained using the web-interface deployment setup.

Figure 4. Qualitative prediction examples from the deployed student model BiteAI-S_int8 (web interface).

4.8. Impact of Distillation, Quantization, and SVD Compression

The student trained purely with supervised learning exhibits a large accuracy gap relative to BiteAI-T, indicating underfitting under severe capacity constraints. Multi-level distillation reduces this gap by imposing output-level and intermediate-level alignment, encouraging the student to reproduce both the teacher’s decision boundaries and its spatial focus [5,23]. Weight-only quantization-aware training enables a sub-megabyte export, but introduces additional accuracy degradation due to weight discretization [14,15]. In practice, this motivates conservative operation (e.g., abstention under low confidence) when deploying extremely compact students.

SVD factorization offers a different compression route that modifies the compute structure of $1\times 1$ convolutions and linear layers. In the reported configuration, SVD + fine-tuning reduces test accuracy relative to BiteAI-T (92.66% vs. 93.75%) while producing a TorchScript FP16 size of 2.84 MB, motivating careful layer selection and rank tuning for storage-driven compression goals.

5. Web-Based Evaluation of BiteAI

To assess the practical behavior of the proposed lightweight classifier, a simple evaluation interface was implemented using the Flask web framework. The application integrates the quantized student model (BiteAI-S_int8), since the student provides the smallest footprint and fast interactive response. The interface allows a user to upload an image of a skin lesion or suspected insect bite. Upon submission, the image is decoded, resized, and normalized using the same preprocessing steps as in offline evaluation, then forwarded to the deployed network for inference. The server returns the predicted class label corresponding to the highest posterior probability, and can optionally apply a conservative confidence threshold to abstain on low-confidence inputs.

The goal of this Flask-based prototype is not production deployment, but functional validation and end-to-end latency characterization under realistic usage conditions. Unlike offline benchmarking, this pipeline captures the complete response time, including input handling, preprocessing, model execution, and output formatting. This provides a more representative measure of interactive performance and verifies that the compressed student can support real-time or near-real-time feedback in resource-constrained settings while preserving the strict storage constraints targeted by BiteAI. Figure 5 shows the home page of the implemented web interface.

Figure 5. Home page of the BiteAI web application.

In our development environment, end-to-end inference for a single image (batch size 1 with $256\times 256$ preprocessing) takes approximately 1–2 s. Actual latency will vary with hardware, runtime settings, and network conditions. Figure 6 presents example interactions with the BiteAI web interface and the corresponding predicted outputs.

Figure 6. Examples of use of the BiteAI web interface and the corresponding predicted outputs.

6. Discussion

The results demonstrate that attention-guided compact CNNs can achieve strong in-distribution performance for multiclass insect-bite recognition, while highlighting important limitations related to domain shift and dataset bias. The BiteAI-T model reaches 94.53% validation accuracy and 93.75% test accuracy on the in-distribution split. Errors are concentrated among visually similar categories, which is consistent with the overlap in clinical appearance across clustered bite reactions.

A large performance gap is observed between direct supervised training of the compact student and training with knowledge transfer. Multi-level distillation (logits, attention maps, and pooled features) recovers a substantial fraction of BiteAI-T performance, supporting distillation as a practical requirement when the student capacity is extremely limited [5,6,23].

Compression results highlight distinct trade-offs. FP16 TorchScript models preserve accuracy at a few megabytes, whereas weight-only int8 export achieves sub-megabyte storage with a larger accuracy reduction. This degradation is expected under aggressive discretization and may not be uniform across classes; consequently, the reject option provides a practical abstention mechanism to reduce the rate of confident errors when deployed under strict size constraints. In the tested configuration, SVD compression decreases accuracy without a clear footprint advantage, indicating that SVD must be carefully configured to meet storage-driven objectives.

The evaluation is performed on a single public dataset, and the dataset release does not provide comprehensive metadata describing skin type, lesion stage, or acquisition device. Therefore, the reported results should not be interpreted as evidence of demographic fairness or cross-dataset clinical generalization. Prior work in dermatology AI has shown that performance can vary across skin types and that domain shift can materially impact accuracy [10,11,12]. Future work should prioritize broader external validation and richer metadata collection, while maintaining the strict efficiency constraints that motivate BiteAI.

7. Conclusions

This paper introduced BiteAI, a storage-efficient insect-bite image classification framework that integrates (i) an attention-pooled MobileNetV3-Small teacher (BiteAI-T), (ii) a compact depthwise-separable student (BiteAI-S), (iii) multi-level knowledge distillation aligning logits and intermediate representations, and (iv) LSQ-inspired weight-only quantization-aware fine-tuning for ultra-compact deployment (BiteAI-S_int8). Across the evaluated splits, BiteAI shows that attention-guided distillation combined with weight-only int8 export can substantially reduce model footprint while preserving practical predictive performance on the in-distribution test set.

Future work will focus on (1) expanding and balancing data collection across demographic and acquisition conditions, (2) directly evaluating and improving robustness under domain shift using external datasets and better metadata, and (3) strengthening the current reject-option setting by adding calibrated uncertainty estimation to support more reliable abstention behavior. These directions aim to improve reliability and safety for unconstrained, real-world use cases while maintaining strict efficiency requirements.