Progressive Alignment with VLM-LLM Feature to Augment
Defect Classification for the ASE Dataset

Defect Recognition has witnessed remarkable enhancements with the integration of deep and machine learning algorithms. In overall, defect recognition can be toward divided into two groups: (1) defect detection and (2) defect classification. The major issue of defect recognition is that when the dataset is small or the quality of image is disappointing, the performance of deep model will be limited.

Defect Detection. Generally speaking, defect detection aims to point out where the defect exist. Defect detection tasks are more common than classification. There have been various works and open-sources specifically designed to handle different types of defect detection, such as [56, 3, 70, 54, 19]. Recently, in order to deal with the challenges in the above mentioned, CS-Flow [42] leverage cross-scale feature flow to integrate the high and low-level feature, enhancing learned representation. XDNet [21] is based on self-supervised learning strategy to explore the information within given dataset.

Defect Classification. Different from detection, the goal of defect classification is to figure out which class does sample belong to. It can be seen as image classification task. The simple way is to train a CNN [53, 25, 57] with suitable augmentation pipeline [52] or data synthesis [22], and ensembling different CNN backbone to enhance the few-shot capabilities [31]. Some researchers have proposed the incorporation of low-level features such as SIFT [41] or depth information [46, 55, 10] to bolster the model performance.

While traditional deep learning typically relies on single-modality data, CLIP [49] have pioneered the integration of vision-language. Benefit from the large-scale of high quality pre-training on vision-language data, multi-modality-based model [33, 34, 35, 11, 67] have shown astonishing zero-shot capabilities. The VLM paradigm introduced by CLIP, along with the versatility of LLM like GPT-3 [16] and LLaMA [27, 20], have a significant impact on the deep learning community. VPT [30] further enhance downstream task performance using fine-tuning-free prompt techniques.

VLM Prompting in Vision With the famous CoCoOp and MaPLe [72, 71, 32], which proposed prompting learning with VLM in vision, there has been a proliferation of works utilizing VLM to enhance computer vision task. Hu and Awal et al. utilized VLM-feature and accurate prompting to align the depth and visual semantic representation for few-shot depth estimation [64, 14]. Liang et al. proposed CLIP-LIT [36], which deployed iterative prompting strategy with CLIP prior for open-world image enhancement. Subramanyam et al., presented CREPE [50] and explored the potential of VLM to improve the performance of visual-object-relationship prediction.

LLM Prompting in Vision The concept of language instruction has also been embraced by the computer vision community for defining image-to-text tasks. Flamingo [17] stands out as a groundbreaking work in this regard, utilizing both vision and language inputs as prompts and achieving impressive few-shot results across various vision-language tasks like image captioning and Visual Question Answering (VQA) [48, 11, 34, 35]. Prophet [73, 68] is designed to guide LLM using answer heuristics for knowledge-based VQA without incorporating external knowledge. VisionLLM [61], presented remarkable results on several benchmarks.

Multi-modality alignment and fusion are both the most essential topics in multi-modal learning, which mainly focus on how to incorporate modality-wise features into a joint representation for downstream application. Different modality feature which embedded on high-dimensional space may contains complementarity, its information may benefits each other [7]. Originated from the limitation of single-modality approaches, which suffer from the variation of dataset and their poor generalizablilty, CLIP [49] and ALBEF [33], which aim to align image-text pairs representation embedded in high-dimensional latent space, has been introduced, lighting up an alternative avenue for tackling zero/few-shot and low-quality dataset scenarios [47, 37]. If multi-modal feature can be processed effectively, rich feature representation can be obtained [47, 8, 45, 38, 15, 26], outperforming the single-modality-based methods.

Multi-modal fusion strategy can be typically classified into feature-level fusion, decision-level fusion and hybrid fusion [44, 23, 2, 4]. Feature-level feature aims to conjuncte the different modality feature to archive high-level joint representation within single model [59, 60]. Decision-level fusion like the classical Mixture-of Experts (MoE) [43] and model-ensemble strategy, each expert specializes in a subset of all given modalities. Hybrid fusion is the most complicated design [63, 13], this type of strategy achieves the best results but has the highest computational burden.

The dataset is provided by ASE corporation and consists of two parts, the details is shown in Figure 1. This dataset contains five classes, with detailed insights provided in Figure 2. (1) Image, which records the drilling machine’s drilling within a fixed range. The goal is to ensure that the drilled holes are positioned as close to the center of the circle as possible, with no shape deviation (without defect). (2) The information about the drilled hole positions, including frequency statistics within specific intervals, as well as corresponding mean $\mu$ and standard deviation $\sigma$. The second part is recorded by meticulous machine, ensuring high quality.

Limitations and Challenges. The ASE dataset itself exhibits a monotonous texture/pattern compared to conventional images. Due to its inductive biases, such as locality and spatial invariance, CNNs model can not learn sufficiently strong low-level feature within shallow layers. As for Vision Transformer (ViT) [58] approaches, while effectively focusing on the global dependencies in the image thanks to its attention-based design, is constrained by the limited number of samples [74], leading to awful performance. To sum up, conventional vision-only model is ineffective to address ASE dataset. Due to mainstream defective classification methods are CNN-based, it is necessary to develop the multi-modal-based method to deal with the difficulty. The GradCAM visualization is in Figure 4, as the figure shows, CNN focuses on the middle region but is not sensitive to capture the shape or distribution of pink dot. ViT-based model fails to concentrate on the central area,

LLMs possess the capability to retain long-term memory of data and accommodate extensive textual input through their large token capacity. Additionally, they can engage in high-level decision-making through iterative question-answering processes. These attributes form the primary motivation for leveraging LLMs to enhance defect classification. In contrast, VLMs are limited to receiving input in the form of individual image-text pairs. While they can perform rudimentary visual reasoning based on the textual content, their ability is restricted to providing basic descriptions of images.

Industrial defect classification, medical diagnosis, product identification, and similar datasets often contain textual or numerical records in addition to images, especially datasets like ASE dataset, as shown in Figure 1. Taking advantage of this, we extract numerical information using OCR and combine with our prior knowledge to serve as input for LLM to accomplish subsequent tasks. Recognizing the superior capability of VLM in simple visual reasoning and image description, we utilize straightforward prompting such as ”Please comprehensively describe the distribution and shape of the image” to obtain basic textual information. Subsequently, combining the output of LLM with the data extracted through OCR and $g(\cdot)$, where $g(\cdot)$ means data preparation. We employ more complex prompting, as depicted in Figure 3, to perform advanced reasoning. The formula can be written as:

	VLM’s answer	$\displaystyle=\text{VLM}(\text{prompt}_{\text{VLM}},\text{Image})$		(1)
	LLM’s answer	$\displaystyle=\text{LLM}(\text{prompt}_{\text{LLM}},g(\text{OCR}(\text{Image})))$

Thanks to the robust zero-shot recognition abilities of VLMs and LLMs, remarkable results are attainable without additional fine-tuning. Additionally, to mitigate the risk of catastrophic forgetting and maintain zero-shot capabilities, we train two adapters tailored to the ASE dataset.

After the VLM-LLM answering, we use a pre-trained image encoder $f_{\text{image}}(\cdot)$ (Resnet-50 [25]) and text encoder $f_{\text{text}}(\cdot)$ (BERT-base [28]) to get a representation for modality fusion. The encoded representation $\mathbf{Z}$ is defined as the Formula 2. we fine-tune on the ASE dataset with 30 epochs for warming up those three encoders before alignment and fusion (refer to Section 3-3, 3-4).

	$\displaystyle\mathbf{Z}_{\text{encoded}}^{\text{Image}}$	$\displaystyle=f_{\text{image}}(\text{Image})$		(2)
	$\displaystyle\mathbf{Z}_{\text{encoded}}^{\text{VLM}}$	$\displaystyle=f_{\text{text}}^{\text{VLM}}(\text{VLM's answer})$
	$\displaystyle\mathbf{Z}_{\text{encoded}}^{\text{LLM}}$	$\displaystyle=f_{\text{text}}^{\text{LLM}}(\text{LLM's answer})$

At recent years, progressive training strategy (PTS) [18] have shown the potential to improve model performance, enhancing the representation understanding without sample-size burden. PTS begins by selecting a subset of positive and negative samples from the training dataset ${D}_{train}$, based on the given initial sampling rate. In our work, we define the positive and negative sample by ranking self-similarity within image-text pairs among training set, negative (low self-similarity) sample will be given priority in the early stages (refer to Figure 5). PTS divides the entire training set into sub-blocks, trains only a small portion of the whole dataset, and gradually increases the training data with each stage. The remaining samples are used for validation, denoted as ${D}_{pend}$. Benefited by PTS, the network’ parameters can be greatly initial and easy to converge. Inspired by [33, 37, 51], which aims to learn a powerful unimodal representation before fusion, we also use the contrastive learning manner for our model, as shown in Figure 6.

Directly aligning features is not effective when there are insufficient samples [62]. The philosophy behind the design of Progressive Feature Alignment (PFA) is that leveraging the PTS’s advantage to deal with the inefficacy of multi-modality learning, resulted from the insufficient-large dataset. In proposed PFA block, we use PTS to gradually align different modal representation encoded by VLM-LLM and image encoder, represented as $\mathbf{Z}_{\text{encoded}}^{\text{VLM}}$, $\mathbf{Z}_{\text{encoded}}^{\text{LLM}}$ and $\mathbf{Z}_{\text{encoded}}^{\text{Image}}$ respectively, where all of representation vectors are encoded to lower-dimensional (256-d) representations. For the every single-batch $B$ at the sub-training set at PTS framework, we can formulate the two data-pairs, $(\mathbf{I},T_{\text{LLM}})$ and $(\mathbf{I},T_{\text{VLM}})$ as the input of FPA block. We first align the feature representation ($\mathbf{Z}_{\text{encoded}}^{\text{Image}}$,$\mathbf{Z}_{\text{encoded}}^{\text{LLM}}$) at divided training set, then aligh ($\mathbf{Z}_{\text{encoded}}^{\text{Image}}$,$\mathbf{Z}_{\text{encoded}}^{\text{VLM}}$). Subsequently, we train for 15 epochs and add more training data progressively, until all training data are used.

For each image and text, we calculate the softmax-normalized image-to-text and text-to-image similarity as:

	$\displaystyle p^{i2t}_{b}(I)=\frac{\exp(s(I,T_{b})/\tau)}{\sum_{b=1}^{B}\exp(s(I,T_{b})/\tau)}$		(3)
	$\displaystyle p^{t2i}_{b}(T)=\frac{\exp(s(T,I_{b})/\tau)}{\sum_{b=1}^{B}\exp(s(T,I_{b})/\tau)}$

where $\tau$ is a learnable parameter and $s(\cdot)$ denote as cosine similarity. Finally, the image-text contrastive loss of proposed PFA is defined as the cross-entropy function $H(\cdot)$ between the similarity $p$ and the ground-truth $y$:

	$\displaystyle L_{itc}^{\text{llm}}=\frac{1}{2}\mathbb{E}\left[H(y^{i2t}(I),p^{i2t}(I))+H(y^{t2i}(T_{\text{llm}}),p^{t2i}(T_{\text{llm}}))\right]$		(4)
	$\displaystyle L_{itc}^{\text{vlm}}=\frac{1}{2}\mathbb{E}\left[H(y^{i2t}(I),p^{i2t}(I))+H(y^{t2i}(T_{\text{vlm}}),p^{t2i}(T_{\text{vlm}}))\right]$

$\displaystyle L_{itc}^{\text{total}}=L_{itc}^{\text{llm}}+L_{itc}^{\text{vlm}}$

(5)

Directly concatenating the representation from different branches may loss some information, leading to reduced performance. Therefore, we proposed cross-modality attention fusion (CMAF) module to incorporate these altogether, avoiding the information missing. Suppose that the aligned features of $f_{\text{text}}(\text{VLM's answer})$, $f_{\text{text}}(\text{LLM's answer})$, and $f_{\text{image}}(\text{Image})$ denoted by $\mathbf{Q}$, $\mathbf{K}$, $\mathbf{V}$, the projection is defined as follows:

	$\displaystyle\mathbf{Q}_{i}$	$\displaystyle=\text{Proj.}(\mathbf{Q})$	$\displaystyle=\mathbf{Z}_{\text{aligned}}^{\text{VLM}}$		(6)
	$\displaystyle\mathbf{K}_{i}$	$\displaystyle=\text{Proj.}(\mathbf{K})$	$\displaystyle=\mathbf{Z}_{\text{aligned}}^{\text{LLM}}$
	$\displaystyle\mathbf{V}_{i}$	$\displaystyle=\text{Proj.}(\mathbf{V})$	$\displaystyle=\mathbf{Z}_{\text{aligned}}^{\text{Image}},$

where $\text{Proj.}(\cdot)$ consists of multiple stages to project the input feature into lower dimensional space. First, the $1\times 1$ convolution is used to project the $\mathbf{X}\in\mathbb{R}^{b\times c\times h\times w}$ into $\mathbf{X^{\prime}}\in\mathbb{R}^{b\times r\times h\times w}$, where $r$ is the reduced number of dimension. We perform the cross-attention [9] by

	$\displaystyle\mathbf{A}_{i}$	$\displaystyle=\text{softmax}\left(\frac{\mathbf{Q}_{i}\cdot\mathbf{K}_{i}^{T}}{\sqrt{r}}\right),$		(7)
	$\displaystyle\mathbf{C}$	$\displaystyle=\text{Concat.}(\mathbf{A}_{0}\cdot\mathbf{V}_{0},\mathbf{A}_{1}\cdot\mathbf{V}_{1},...,\mathbf{A}_{head}\cdot\mathbf{V}_{head}),$
	$\displaystyle\mathbf{W}$	$\displaystyle=\text{Sigmoid}(\text{Proj.}(\mathbf{C}))$

where $\text{Proj.}_{C}$ projects the concatenated cross-attentions to get adaptive weights $W$, i.e., $b\times 3\times h\times w$. In this way, we could judiciously fuse the different modalities, the proposed CMAF still remains strong due to its adaptivity, as follows:

$$\mathbf{Z}_{\text{fused}}=W_{1}\cdot\mathbf{Q}+W_{2}\cdot\mathbf{K}+W_{3}\cdot\mathbf{V}$$

(8)

where $\mathbf{W}_{i}$ indicates $i$-th channel of $\mathbf{W}$. Eventually, the predicted class $\mathbf{Y}^{*}$ is obtained via a multi-layer perceptron by $\mathbf{Y}^{*}$ = MLP ($\mathbf{Z}_{\text{fused}}$).

Data augmentation is crucial to improve model performance, especially when there is a domain gap or limited sample size. It aims to improve the recognition ability of the model by increasing the diversity of data domain. However, the text and numeric information is paired with image in ASE dataset, common data augmentation strategies, such as geometry and HSV transformation, are incompatible. We design a simple but effective Task-specific data augmentation (TDA), following offline synthesis manner [52, 22], to address the data insufficiency.

To synthesis the ASE data, we sample data points $(x_{i},y_{i})$ N times from bivariate Gaussian distribution $\mathcal{N}(\mu_{cls},\Sigma_{cls})$ for the different classes. For each sampled data points, set the pixel at coordinates of sampled point in the image to pink dot. The number of samples N obtained depends on the dataset’s configuration. For instance, we can utilize OCR-ed data to compute the total points across all radius (refer to Figure 1). After sampling process, we could obtain the augmented image-text pair data $\mathbf{X}_{aug,cls}^{\text{image}}$ and $\mathbf{X}_{aug,cls}^{\text{text}}$ for ASE dataset. The TDA pipeline can be simply defined as:

$$\mu_{cls}=[\mu_{x,cls},\mu_{x,cls}],\Sigma_{cls}=\begin{pmatrix}\sigma_{x,cls}^{2}&\pm\sigma_{xy,cls}\\ \pm\sigma_{xy,cls}&\sigma_{y,cls}^{2}\end{pmatrix}$$

(9)

$$\mathbf{X}_{aug,cls}^{\text{image}}=\{(x_{i},y_{i})\mid(x_{i-1},y_{i-1})\sim\mathcal{N}(\mu_{cls},\Sigma_{cls}),\,i=1,2,\ldots,N\}$$

(10)

General settings. The details of ASE dataset have been described in Section 3-1. It contains 455 samples, 325 of which are selected as the training set and 130 as the testing set in our experiments. We extracted external modal information utilizing CnOCR [1, 69], an efficient OCR package. Instruct-BLIP2 [11], a powerful VLM renowned for its adeptness in generating long-context and exhibiting promising performance across several VQA benchmarks, is employed. LLaMa2-7B [27], is chosen for generating high-low-level context about decision making. We use the proposed TDA to synthesis the training sample, from 325 to 650. Both VLM and LLM employed fixed prompting, as detailed in Section 3-2. For the vision encoder, ResNet50 [25] is employed. On the other hand, the text encoder utilized BERT-base [28] following inference from VLM and LLM. Our experiments is on NVIDIA GeForce RTX 3090.

Loss function and Hyperparameter-settings. As for alignment phase, the loss function was described in Section 3-3. After PFA, the loss function is changed to Cross-Entropy for training remains network. The AdamW optimizer [40, 39] is employed, where ($\beta_{1}$, $\beta_{2}$) are (0.9, 0.999), and the epsilon $\epsilon$ is $1\textrm{e}\!-\!8$. The initial learning rate is $1\textrm{e}\!-\!2$ with the step-wise learning rate decaying every 15 epochs with scaling factor $3/4$, totally trains 60 epochs in fusion phase, and the batch size is $10$.

Evaluation metric. In our experiments, we evaluate model performance using the f1-score for each class and the macro-f1-score. The f1-score is defined as:

$$\text{f1-score}=2\times\frac{\text{precision}\times\text{recall}}{\text{precision}+\text{recall}}$$

(11)

where precision and recall are computed for each individual class. The macro f1-score is the average of the f1-scores for all classes:

$$\text{macro f1-score}=\frac{1}{N}\sum_{i=1}^{N}\text{f1-score}_{i}$$

(12)

where $N$ is the number of classes, and $\text{f1-score}_{i}$ is the f1-score for the $i$-th class. These metrics provide a balanced evaluation of the model’s ability to classify each class accurately and its overall performance across all classes.

To demonstrate the efficacy of our method in the ASE dataset, we selected three prominent image classification models as baselines. These include CrossViT-18 [9], DeiT-B [58], and EfficientNet-b6a [57]. Additionally, we evaluated our approach against three specialized few-shot-based defect classification frameworks: those proposed by Karmakar et al. [31], which mainly based on transfer learning and ensemble method, Xie et al. [66], which is rooted on exploring extra-feature within given image.

The experiment results has shown in the Table 2. Our method demonstrated a outstanding results. Despite being designed to enhance performance through the fusion of multi-scale features and to increase data utilization efficiency, respectively, both CrossViT [9] and DeiT [58] under-perform compared to the simple EfficientNet classifier [57] on the ASE dataset, whether in binary or multi-class classification tasks.

In this section, we ablate important design elements in the proposed method. The results is in the Table 3. In overall, the PFA block is the most beneficial for the model’s effectiveness. It aligns the features of AOI images with those of the VLM/LLM image-text pairs, enabling the model to address the minority samples in multi-class issues, especially for the tail classes such as type 2, 3, and 4. Furthermore, our proposed TDA designed for the ASE data also effectively enhances the model’s performance.

Stride settings in PFA block. The stride is crucial as it pertains to the proportion of data added to the progressive training with each iteration. A smaller stride can more precisely align the embedded vectors between image-text pairs in negative samples, yielding more desirable performance. However, it may lead to additional training duration and computational load. Conversely, a too large stride may result in limited effectiveness, as shown in Figure 8. An appropriate stride can guide the model to effectively align with challenging samples without compromising performance. Ablation studies measuring performance across different strides are presented in Table 4.

Feature Fusion Strategy. The method of feature fusion may affect whether certain semantic information can interact well during the training process. Direct concatenation without the use of additional fusion networks, such as self-attention mechanisms, may lead to information loss. In contrast, our proposed CMAF contains the sigmoid function, allowing features from different modalities to adaptively update each other’s weights without overly relying on any single modality. The experimental results demonstrating this, as indicated in Table 5, affirm the efficacy of our approach.