Source-linked AI summary

Cross-Architecture Knowledge Distillation from a Vision Foundation Model to a Lightweight Visual State Space Model for Tea Leaf Disease Classification

Zibo Zhou, Zongsen Qiu, Rui Chen, Yujie Yao, Yue Zhou, Jianjun Wang

arXiv:2608.26771v1cs.CV

TL;DR

Accurate tea disease classification is difficult to deploy on edge devices because foundation models are large and small-data lightweight models may underfit. The paper distills a fine-tuned DINOv2 teacher into a compact LVSSM student and finds that logit-only KD improves the student while feature alignment hurts. The distilled student retains most of the teacher’s accuracy with substantially fewer parameters, within the study’s single-dataset and simplified-SSM scope.

  • Problem

    Cross-architecture KD from ViT teachers to SSM students is less understood because their token-mixing mechanisms and intermediate representations differ, while edge deployment needs compact accurate models.

  • Method

    The study fixes two from-scratch LVSSM training failures and distills a fine-tuned DINOv2 teacher using temperature-scaled logit matching, with feature-alignment ablations.

  • Results

    +3.09 accuracy points on average raise the student from 92.32% to 95.41% across three seeds; the distilled student is 5.0× smaller and retains 98.3% of the teacher’s accuracy.

  • Takeaways & Limitations

    Logit-only KD is the strongest reported ViT→SSM configuration and is most useful when a from-scratch student starts below its teacher.

  • Takeaways & Limitations

    The evaluation uses one tea dataset, a simplified non-official SSM implementation, and server GPU/CPU latency rather than embedded-device measurements.

Abstract

from arXiv · show

Automated tea leaf disease classification supports precision agriculture, yet deploying accurate models on edge devices remains challenging under tight compute budgets. Self-supervised vision foundation models such as DINOv2 provide strong features but are too large for field deployment, while lightweight models trained from scratch on small agricultural datasets often underfit. We study cross-architecture knowledge distillation (KD) from a fine-tuned DINOv2 teacher (Vision Transformer) to a compact bidirectional Visual State Space Model (LVSSM) student, an underexplored direction because the architectures use fundamentally different token-mixing mechanisms. We identify and fix two training-stability problems that prevent the from-scratch SSM student from learning on limited data: a single large patch-embedding convolution and a fusion layer that severs the residual path. With a progressive convolutional stem and gated bidirectional selective-scan block, the 4.45M-parameter student trains stably. Across three seeds, temperature-scaled logit distillation raises test accuracy from 92.32+/-2.14% to 95.41+/-1.17% (best single run: 96.20%; macro-F1: 94.45%), a +3.09 percentage-point mean gain. The student uses 5.0 times fewer parameters than the 22M-parameter teacher while retaining 98.3% of its accuracy. Ablations show that intermediate feature-alignment losses reduce accuracy, making simple logit-level KD the strongest configuration. A fair from-scratch comparison shows the gain is specific to students that start below the teacher. We report per-class metrics, confusion matrices, bootstrap confidence intervals, and FLOPs/latency measurements, and discuss limitations including the single-dataset scope and simplified non-official SSM implementation.

I. INTRODUCTION

Tea disease classification needs accurate automation for labor-intensive inspection, but edge deployment remains difficult because strong foundation models are expensive and small-data models often underfit. This work studies cross-architecture distillation from a ViT teacher to a lightweight visual SSM student.

  • Edge deployment remains challenging because high-accuracy agricultural models exceed resource constraints.
  • 22M-parameter DINOv2 variants are expensive for embedded inference, while compact models trained from scratch often generalize poorly on small datasets.
  • Vision SSMs offer linear-time sequence modeling, but bidirectional scanning and weak patch embeddings make small-data training difficult.
  • Cross-architecture KD between ViT teachers and SSM students is less understood because their token-mixing mechanisms and intermediate representations differ.
  • The study fixes two SSM training failures, evaluates logit and feature distillation, and reports accuracy, class metrics, uncertainty intervals, and efficiency measurements.
  • Plant-disease recognition commonly relies on ImageNet-pretrained CNNs and large curated datasets, whereas tea datasets are smaller and class-imbalanced.

B. Vision State Space Models

Vision SSMs adapt selective state-space scanning to images by serializing patches, offering efficient sequence modeling but requiring careful training choices. The paper positions LVSSM as a simplified self-contained SSM student for cross-architecture distillation.

  • Mamba provides input-dependent selective state parameters for linear-time sequence modeling with a hardware-aware scan.
  • Vision Mamba and VMamba serialize image patches and scan them bidirectionally or along multiple spatial orders.
  • Cross-architecture transfer is less studied because spatial attention maps and scan-order state trajectories use structurally different coordinate systems.
  • The training objective makes LVSSM approximate the fine-tuned DINOv2 teacher’s predictive distribution using fewer parameters.

B. Teacher Model: Fine-tuned DINOv2

The teacher is a fine-tuned DINOv2-Small model, while the student is a compact LVSSM designed to address instability in from-scratch bidirectional SSM training.

  • DINOv2-Small has 22.06M parameters and reaches 97.86% test accuracy with 96.76% macro-F1 on the tea disease dataset.
  • The LVSSM student contains 4,454,982 parameters and targets two observed training-stability failure modes.
  • A single 16×16 patch convolution and nonidentity fusion layer produced vanishing gradients and approximately chance accuracy near 1/6.
  • A four-layer progressive stride-2 convolutional stem replaces the single patch embedding to provide local feature extraction before sequential processing.

3) Gated-Residual Bidirectional Selective Scan:

The gated bidirectional block combines forward and backward scan outputs while retaining an explicit identity path. This design lets the gate control contextual injection and restores stable gradient flow.

  • Each block scans the sequence forward and backward, then combines both directions through a learned gate with an identity residual path.
  • The residual output is o = x + g ⊙Wp c, where the gate controls how much bidirectional context enters the block.
  • The additive residual guarantees an identity route for gradients and restored stable training in all reported runs.
  • The student uses six gated bidirectional SSM blocks with embedding dimension d = 192, state dimension 16, and 4,454,982 total parameters.
  • Logit-level distillation uses temperature-scaled soft-target matching between student and teacher logits, with τ = 2.0.

2) Combined Objective:

The training objective combines groundtruth cross-entropy with knowledge-distillation loss, while optional feature-alignment terms are evaluated separately. The main configuration weights cross-entropy and KD equally and uses specified optimization, augmentation, batch-size, and training-duration settings.

  • Combined Objective: The total loss combines groundtruth cross-entropy and KD.The main configuration sets αCE = αKD = 0.5.
  • Feature Alignment: Optional feature losses align teacher attention with student scan states or projected intermediate features.The two terms are scan-attention alignment and progressive feature distillation.
  • Optimization: The models train for 30 epochs using AdamW, warmup followed by cosine annealing, gradient clipping, and image augmentation.The distilled students use batch size 12, while the standalone student uses batch size 24.

IV. EXPERIMENTS

Experiments evaluate the models on a leakage-checked tea leaf disease dataset under a shared measurement pipeline. Distillation improves the compact student while preserving a substantial efficiency advantage, although pretrained CNNs remain stronger.

  • Dataset: 4,259 training and 421 validation images are reported after converting detection annotations to classification labels and discarding 79 empty-annotation images.The dataset contains six tea leaf disease classes and uses source-photograph group-aware splits with zero cross-split source-photo overlap.
  • Evaluation: The shared evaluation pipeline reports accuracy, F1 scores, per-class metrics, confusion matrices, bootstrap 95% confidence intervals, FLOPs, and latency.Models are evaluated under one pipeline for comparability.
  • Main Results: 3.09 accuracy points separate the distilled and standalone students on the three-seed mean comparison.The best-run comparison in the cited results raises accuracy from 92.16% to 96.20%.
  • Training Dynamics: Distillation converges faster and reaches a higher validation-accuracy plateau than standalone training.This comparison is shown in Fig. 1.
  • Main Results: 5.0× fewer parameters and 98.3% retained teacher accuracy characterize the distilled student’s efficiency tradeoff.The cited best-run accuracies are 96.20% for the student and 97.86% for the teacher.
  • Main Results: 98.1–98.6% accuracy is achieved by ImageNet-pretrained CNNs, exceeding the from-scratch SSM student.The distilled SSM comes within approximately 2.4 points of EfficientNet-B0 while using only the tea images for training.

D. Per-Class Analysis and Error Structure

Per-class analysis shows that distillation most benefits classes that the standalone student handled poorly, while residual errors concentrate between visually similar diseases. Controlled ablations show that feature alignment reduces accuracy despite retaining the full KD signal.

  • Per-Class Analysis: Helopeltis gains +12.7 F1 and Red Leaf Spot gains +8.6 F1 with distillation.These are the classes the standalone student handled worst.
  • Error Structure: Brown Blight and Gray Blight are mutually confused, with 5 misclassifications in each direction.The teacher also confuses these two visually similar necrotic-lesion diseases.
  • Ablation: The accuracy falls from 96.20% to 94.77–95.49% when feature-alignment losses are added with the KD signal held at full strength.The controlled ablation adds scan-attention or progressive feature-distillation terms without reducing the primary KD weight.
  • Ablation: Feature alignment can interfere with logit-level KD because ViT attention geometry and SSM scan-state or depth-wise representations are structurally different.The paper identifies representational mismatch as the explanation and selects logit-only KD as the strongest configuration.

F. Is Cross-Architecture KD Special to the SSM Student?

The fair comparison shows that cross-architecture distillation substantially benefits the underperforming SSM student but not compact CNNs that already match or exceed the teacher. Controlled ablations and residual-error analysis support logit-only KD as the strongest configuration.

  • Residual errors: Brown Blight and Gray Blight remained the dominant residual error, with five misclassifications in each direction.The teacher also confuses these visually similar diseases, suggesting the error is intrinsic rather than a distillation artifact.
  • Fair cross-architecture comparison: +3.09 points was the SSM student's mean distillation gain, compared with −1.19 and +0.24 for the CNN students.All students used the same teacher and logit-KD recipe from random initialization.
  • Fair cross-architecture comparison: 96.4–97.2% standalone CNN accuracy left little useful teacher knowledge to transfer, whereas the weaker SSM had substantial headroom.The teacher's own accuracy was 97.86%.
  • Efficiency: The distilled LVSSM is 5.0× smaller than the teacher, but its unfused selective scan measured 10.05 ms on GPU and 1410 ms on CPU.The simplified pure-PyTorch implementation makes parameter count and FLOPs understate real inference cost.
  • Controlled ablation: Adding feature alignment reduced accuracy when CE and KD weights remained fixed at 0.5, making logit-only KD strongest.The controlled ablation added feature terms rather than trading them against the primary KD objective.

H. Multi-Seed Stability

Across three seeds, logit distillation consistently improved LVSSM accuracy and macro-F1 while reducing run-to-run variance. The reported mean, rather than the best single run, is used for cross-model comparisons, but deployment efficiency remains limited by the unfused scan implementation.

  • Robustness across seeds: +3.09 points in accuracy and +3.44 points in macro-F1 were observed across three seeds with logit KD.Accuracy increased from 92.32% to 95.41%, while macro-F1 increased from 90.19% to 93.63%.
  • Robustness across seeds: Accuracy standard deviation fell from 2.14 to 1.17 with distillation.The authors relate this reduction to soft targets providing a more stable optimization signal than one-hot labels alone.
  • Reporting convention: The best single logit-KD run reached 96.20%, while the three-seed mean was 95.41%.The difference was less than one standard deviation, so the mean is used where variance matters.
  • Interpretation: Logit-level KD leaves the SSM free to learn scan-suited internal representations, unlike feature alignment to the teacher's patch-attention geometry.This architecture-agnostic output-space rationale is consistent with the ablation favoring logit-only KD.
  • Limitations: The study's efficiency evidence is limited by a simplified unfused SSM implementation and server-side GPU/CPU measurements rather than embedded-device testing.The paper identifies fused kernels and on-device measurement as prerequisites for assessing edge deployment suitability.

VI. CONCLUSION

The study shows that logit-level cross-architecture distillation can improve a lightweight LVSSM for tea leaf disease classification while using far fewer parameters than the teacher. Feature-alignment losses reduce accuracy, and broader validation plus optimized kernels remain important for deployment.

  • Conclusion: +3.09 accuracy points on average over three seeds lift the 4.45M-parameter student to 95.41% test accuracy.The best run reached 96.20% accuracy and 94.45% macro-F1.
  • Conclusion: 5.0× fewer parameters let the student retain 98.3% of the teacher’s accuracy.
  • Conclusion: Feature-alignment losses reduce accuracy, making simple logit-level KD the strongest ViT→SSM configuration.
  • Conclusion: The distillation gain is specific to students that start below the teacher: it helps the from-scratch SSM but not already-matching CNNs.
  • Conclusion: Fused selective-scan kernels and multi-dataset validation are identified as key steps toward practical edge deployment.
Loading 2608.26771v1…