# Assembly Quality Detection Based on Class-Imbalanced Semi-Supervised Learning
## 2 Class-Imbalanced Semi-Supervised Learning
### Model Framework for Assembly Quality Detection  
The method adopted in this paper is based on the mean teacher algorithm, which includes a teacher model and a student model with the same network structure. The overall algorithm is shown in Figure 2. The models use the DenseNet121 [29] structure, famous for achieving high performance while reducing the scale of parameters. As shown in Figure 2, the DenseNet121 includes four dense blocks, each followed by a transition layer. The transition layer consists of one \(1\times 1\) convolutional layer and one \(2\times 2\) pooling layer. The dense block is composed of multiple dense layers, and the output of each dense layer is connected to other layers in a feedforward manner to prevent the disappearance of features. The dense block and dense layer are illustrated in Figure 3, and each dense layer is a sequence of BN-ReLU-\(1\times 1\) Conv-BN-ReLU-\(3\times 3\) Conv. Figure 3: Dense block composed of five Dense layers. Figure 2: Class-imbalanced semi-supervised learning framework. Inspired by knowledge distillation, the mean teacher (MT) method [26] uses the teacher-student structure. The weights of the teacher are the exponential moving average of the weights of the student. The mean teacher algorithm introduces perturbations in the model weights and input data and encouraging the predictions to remain the same. Defining \(\theta^{s}\) as the weights of the student model, then the corresponding weight \(\theta^{t}\) of the teacher model is  
\[\theta^{t}\leftarrow\mu\theta^{t}+(1-\mu)\theta^{s}, \tag{1}\]  
\[\mu=\min(1-\frac{1}{iter},\mu_{0}), \tag{2}\]  
where \(\mu\) is the smoothing coefficient hyper-parameter, \(iter\) is the global iteration step and \(\mu_{0}\) is the maximum value of \(\mu\). At the early training period, \(\mu\) is small, therefore the teacher is rapidly updated by the new student weights. In the later training period, when \(\mu\) reaches \(\mu_{0}\), the teacher will have a longer memory since the improvement of the student is slow down. As the perturbations to the input, random noise enhancement \((\eta,\eta^{\prime})\) is applied to the original sample \(x_{i}\) before input to the models, so the predictions of the student and the teacher are \(\hat{y}_{i}^{s}=f(x_{i},\theta^{s},\eta)\) and \(\hat{y}_{i}^{t}=f(x_{i},\theta^{t},\eta^{\prime})\). The usual mean teacher approach uses the mean square error (MSE) as the consistency regularization loss (CRL) to minimize the Euclidean distance between the teacher prediction and the student prediction. \[L_{CRL}(\hat{y}_{i}^{s},\hat{y}_{i}^{t})=\big{\|}\hat{y}_{i}^{s}-\hat{y}_{i}^{ t}\big{\|}^{2}, \tag{3}\]  
Not all teacher predictions are reliable, and consistent constraints on the unreliable predictions will damage model performance. To make the student model dynamically select reliable predictions from the teacher, this paper adopts the certainty driven mechanism, which is explained in detail in Section 2.2. At each iteration, \(m\) samples with reliable teacher predictions are chosen from the mini-batch to form the subset \(M\), and the consistent constraint is computed on \(M\). Then, the cost function is the sum of the consistency regularization loss of each sample. \[I_{consistency}=\sum_{x_{i}\in M}L_{CRL}(\hat{y}_{i}^{s},\hat{y}_{i}^{t}), \tag{4}\]  
To improve the robustness of MT to class imbalance, LDAM loss is applied to \(b_{l}\) labeled samples in the mini-batch to enhance the accuracy of the student model in the supervised training. Moreover, the compression consistency loss (CCL) is introduced to weaken the decision boundary smoothing effect of samples predicted to be the majority class. Details of the class-imbalanced loss functions are descripted in Section 2.3. Finally, the semi-supervised cost function of the mini-batch is  
\[ \begin{array}{l}J=J_{classification}+J_{consistency}\\ =\sum_{i=1}^{b_{l}}L_{\text{LDAM}}(\hat{y}_{i}^{s},y_{i})+\sum_{x_{i}\in M}L_{ CCL}(\hat{y}_{i}^{s},\hat{y}_{i}^{t})\end{array}