A relational knowledge distillation method based on category boundary uncertainty

Abstract

This invention discloses a relational knowledge distillation method based on category boundary uncertainty. 1) An image is simultaneously processed by a feature encoder of two network branches to obtain an intermediate feature map, which is then processed by a decoder to obtain the final segmentation prediction map; 2) Based on the intermediate feature map obtained in step 1), the feature-level consistency distillation loss is calculated; 3) Based on the segmentation prediction map obtained in step 1), the pixel-level consistency distillation loss is calculated; 4) Based on the two-branch segmentation prediction map obtained in step 1) and its corresponding true class label A, the target loss for branches 1 and 2 is calculated respectively; 5) The total optimized loss function is calculated; 6) The model is trained based on the loss function of step 5). This invention introduces the BURKD algorithm, which can better model the relationships between samples and between categories, and uses a method based on category boundary uncertainty to adaptively and dynamically weight each category, thereby improving the performance of semantic segmentation.

Description

Technical Field

[0001] This invention belongs to the field of relational knowledge distillation, specifically constructing a relational knowledge distillation method based on category boundary uncertainty to assist in enhancing semantic segmentation methods. Background Technology

[0002] Knowledge distillation (KD) aims to transfer valuable knowledge from a large teacher model to a small student model to improve the accuracy of the student model's reasoning. Generally, distillation methods can be categorized into three types: response-based distillation (matching soft scores between teachers and students), representation-based distillation (matching feature maps), and relation-based distillation (matching sample relations). Key factors for the success of KD include: 1) the type of knowledge distillation (e.g., feature or score levels), 2) the distillation strategy employed (e.g., offline distillation, online distillation, and self-distillation), and 3) the teacher-student architecture (e.g., multi-teacher KD, attention-based KD, and graph-based KD).

[0003] While knowledge distillation originated from resource-efficient deep learning, it has found wider application in fields such as incomplete multimodal learning. For incomplete multimodal medical image segmentation tasks, models are highly sensitive to feature representations; coarse, generalized KD may be ineffective and can be considered aimless and redundant, resulting in a significant performance gap between the obtained and expected results. Furthermore, because the inputs to teacher and student networks differ, knowledge transfer via representation-based methods can lead to overfitting. Recent methods have focused on transferring modal-specific information through relation-based approaches. However, current methods typically consider different instances equally, ignoring their specificity, which leads to suboptimal performance.

[0004] Given the above limitations, we propose a relational knowledge distillation method based on category boundary uncertainty. We introduce this method into a semantic segmentation network and verify its effectiveness in an incomplete multimodal brain tumor segmentation scenario. Summary of the Invention

[0005] To address the problems in the background art, this invention provides a relational knowledge distillation method based on category boundary uncertainty. The purpose of this invention is to use this method to determine the importance of instance categories based on entropy, thereby assigning different weights to different instances. This allows for the specific consideration of different semantic segmentation instance categories, ultimately enhancing the accuracy of semantic segmentation.

[0006] The technical solution adopted in this invention is as follows, specifically including the following steps:

[0007] 1) Input an image into two network branches simultaneously, first passing it through the feature encoder. The intermediate feature maps F1 and F2 are obtained; then they are decoded respectively. The final segmentation prediction map is obtained.

[0008] 2) Based on the intermediate feature maps F1 and F2 obtained in step 1), calculate the uniform distillation loss L at the feature level. fc ;

[0009] 3) Based on the segmentation prediction map obtained in step 1), Calculate pixel-level uniform distillation loss

[0010] 4) Based on the bi-branch segmentation prediction map obtained in step 1), Calculate the target loss for the first network branch and the second network branch, respectively, based on the corresponding true class label A.

[0011] 5) The L obtained from step 2) fc And obtained in step 3) and the result obtained in step 5) Calculate the total optimization loss L i ;

[0012] 6) Train the model based on the loss function from step 5);

[0013] 7) Use the model trained in step 6) to obtain the segmentation or classification results of the image to be processed.

[0014] In step 1):

[0015] The first network branch includes a feature encoder. decoder

[0016] The second network branch includes a feature encoder. decoder

[0017] Step 2) specifically involves: based on the intermediate feature maps F1 and F2 obtained in step 1), sequentially calculating the class-specific prototype, the contextual relationships between modeled samples and between classes, and the reweighting of class boundary uncertainty, to obtain the feature-level consistency loss L. fc The specific steps are as follows:

[0018] Step 2.1) Category-Specific Prototype Generation: Given an intermediate feature map and its corresponding real label A∈RH×W×Z Calculate the class-specific prototypes for the first and second network branches respectively:

[0019]

[0020] in, H i W i Z i d represents the length, width, and height of the feature. f The number of channels represents the feature; H, W, Z represent the length, width, and height of the true label A; N represents the number of channels. c This represents the number of pixels belonging to category c. It is an indicator function that represents a value of 1 when the internal condition is True, and 0 otherwise.

[0021] Step 2.2) Contextual Modeling of Relationships Between Samples and Categories: Based on the category-specific prototype obtained in Step 2.1), the relationship graphs R1 and R2 between samples and categories are obtained using the following formulas:

[0022]

[0023] in, Let R1 and R2 represent the prototypes of the c1-th class of the b1-th sample in the first and second network branches, respectively, where R1 and R2 ∈ R. (B*C)×(B*C) B represents the number of samples in a batch, and C represents the number of semantic categories;<a,b> R represents the operation of multiplying the dot product of elements a and b. i (b1,c1; b2,c2) represents the feature similarity between the c1th class of the b1th sample and the c2th class of the b2th sample in the i-th network branch; then the distance difference between samples and between classes between the two branches is calculated:

[0024]

[0025] in, This represents the overall relationship difference between the b1-th sample and the c1-th category between the first and second network branches;

[0026] Step 2.3) Reweighting based on boundary uncertainty: To prioritize consistency learning on more challenging categories, we introduce a reweighting method based on category boundary specificity uncertainty; based on the final segmentation prediction map obtained in Step 1). The prediction results are normalized using the softmax method. Then calculate the category-specific weight.

[0027]

[0028] in, C represents the number of categories. The c-th category of branch i is represented; σ() represents the softmax function;

[0029] The final uncertainty weights are obtained by averaging the values ​​of the two network branches:

[0030]

[0031] Step 2.4) The final feature-level relation consistency loss is:

[0032]

[0033] Among them, W b1,c1 W c A subset of W c The b1 sample.

[0034] Step 3) specifically involves: based on the bi-branch predicted segmentation map obtained in step 1). Calculate the pixel-level consistency loss for branches 1 and 2 using KL divergence:

[0035]

[0036]

[0037] in, This indicates that the second network branch constrains the first network branch. This indicates that the first network branch is used to constrain the second network branch, and τ represents the temperature factor hyperparameter.

[0038] Step 4) specifically involves: based on the segmentation prediction map obtained in step 1). And its corresponding real label A, to obtain the target task loss of the network branch.

[0039]

[0040] Where i represents the i-th network branch, and CE(.) and Dice(.) represent the cross-entropy and Dice loss functions, respectively.

[0041] Step 5) specifically involves: based on the feature-level consistency loss L from step 2). fc Pixel-level consistency loss in step 3) and the task loss in step 4). Calculate the overall objective loss L for branch i. i :

[0042]

[0043] The beneficial effects of this invention are:

[0044] The Boundary Uncertainty based Relation Knowledgeable Distillation (BURKD) method of this invention mainly comprises three steps: class-specific prototype generation, contextual modeling of relationships between samples and classes, and uncertainty-based reweighting. We integrate this method into a mutual learning-based bi-branch semantic segmentation network and explore its effectiveness in an incomplete multimodal brain tumor segmentation scenario. In this invention, we validate the effectiveness of this method in an incomplete multimodal brain tumor segmentation model scenario, surpassing the current state-of-the-art (SOTA) performance and achieving the best segmentation results. The relation knowledge distillation method based on uncertain classification boundaries designed in this invention fully considers the specificity of different instances in semantic segmentation. This enables the model to achieve SOTA performance on the multimodal brain tumor semantic segmentation dataset (BraTs2018), surpassing other existing methods. Attached Figure Description

[0045] Figure 1 This is a flowchart of the method of the present invention;

[0046] Figure 2 This is a schematic diagram of the BURKD algorithm flow of the present invention. Detailed Implementation

[0047] The present invention will be further described in detail below with reference to the accompanying drawings and specific embodiments.

[0048] like Figure 1 As shown, a relational knowledge distillation method based on category boundary uncertainty according to the present invention includes the following steps:

[0049] Step 1): An image is simultaneously passed through the feature encoders of two network branches. The intermediate feature maps F1 and F2 are obtained;

[0050] Step 2): Then, based on their respective decoders The final segmentation prediction map is obtained.

[0051] Step 3): Based on the intermediate feature maps F1 and F2 obtained in Step 1), calculate the feature-level uniformity distillation loss L. fc ;

[0052] Step 3) specifically involves: based on the intermediate feature maps F1 and F2 obtained in step 1), sequentially calculating the class-specific prototype, the contextual relationships between modeled samples and between classes, and the reweighting of class boundary uncertainty, to obtain the feature-level consistency loss L. fc The specific operations are as follows: steps 3.1) to 3.3).

[0053] 3.1) Category-Specific Prototype Generation: Given an intermediate feature map and its corresponding real label A∈R H×W×Z Calculate the class-specific prototypes for branches 1 and 2 respectively:

[0054]

[0055] in H i W i Z i d represents the length, width, and height of the feature. f The number of channels represents the feature; H, W, Z represent the length, width, and height of the true label A; N represents the number of channels. c This represents the number of pixels belonging to category c. It is an indicator function that represents a value of 1 when the internal condition is True, and 0 otherwise.

[0056] 3.2) Contextual Modeling of Relationships Between Samples and Categories: Based on the category-specific prototype obtained in step 3.1), the relationship graphs R1 and R2 between samples and categories are obtained using the following formulas:

[0057]

[0058] in, Let R1 and R2 represent the prototypes of the c1th class of the b1th sample in branches 1 and 2, respectively. (B*C)×(B*C) B represents the number of samples in a batch, and C represents the number of semantic categories.<a,b> R represents the operation of multiplying the dot product of elements a and b. i (b1,c1; b2,c2) represents the feature similarity between the c1th class of the b1th sample and the c2th class of the b2th sample in network branch i. Then, the distance differences between samples and between classes are calculated between the two branches:

[0059]

[0060] in This represents the overall relationship difference between the b1-th sample and the c1-th category in network branches 1 and 2.

[0061] 3.3) Reweighting Based on Boundary Uncertainty: To prioritize consistency learning on more challenging categories, we introduce a reweighting method based on category boundary specificity uncertainty. This method is applied to the final segmentation prediction map obtained in step 2). The prediction results are normalized using the softmax method. Then calculate the category-specific weight.

[0062]

[0063] in C represents the number of categories. Let represent the c-th category of branch i. The final uncertainty weight is the average of the two network branches:

[0064]

[0065] The final feature-level relational consistency loss;

[0066]

[0067] Step 4): Based on the segmentation prediction map obtained in Step 2). Calculate pixel-level uniform distillation loss Specifically as follows:

[0068] Based on the bi-branch predicted segmentation map obtained in step 2), Calculate the pixel-level consistency loss for branches 1 and 2 using KL divergence:

[0069]

[0070] in This means that branch 1 is constrained by branch 2 of the network. This indicates that branch 2 is constrained by network branch 1, and τ represents the temperature factor hyperparameter.

[0071] Step 5): Based on the bi-branch segmentation prediction map obtained in Step 2), Calculate the target loss for branches 1 and 2, respectively, based on the corresponding true class label A. Specifically as follows:

[0072] Based on the segmentation prediction map obtained in step 2) And its corresponding real label A, to obtain the target task loss of the network branch.

[0073]

[0074] Where i represents network branches 1 and 2, and CE(.) and Dice(.) represent cross-entropy and Dice loss functions;

[0075] Step 6): Combine the L obtained from Step 3) fc And obtained in step 4) and the result obtained in step 5) Calculate the total optimization loss L i The details are as follows:

[0076] Based on the feature-level consistency loss L in step 3) fC Pixel-level consistency loss in step 4) and the task loss in step 5). Calculate the overall objective loss L for branch i. i :

[0077]

[0078] Step 7): Train the network model based on the BURKD algorithm, select the comprehensive loss as the final loss function, and use error backpropagation to adjust the weights and biases of the network model. Select parameters based on the Dice Similarity Coefficient (DSC), and finally determine a set of parameters that maximizes the accuracy of brain tumor segmentation (DSC score).

[0079] The BURKD algorithm of this invention can be used to assist the knowledge distillation framework in knowledge transfer learning. This invention applies it to brain tumor segmentation in a multimodal missing scenario to verify the effectiveness of the algorithm. The model ultimately determines whether each pixel in the image belongs to its corresponding tumor category. According to the evaluation criteria, the method has a significant effect on multi-class semantic segmentation.

[0080] The BURKD algorithm and some state-of-the-art (SOTA) algorithms perform as follows on the current BraTs2018 brain tumor segmentation dataset with arbitrary modalities missing:

[0081]

[0082] Note:

[0083] (1)MCTSeg, M3AE, U-Net-MFI, RFNet, mmFormer and SMU-Net are models that previously achieved state-of-the-art performance on the BraTs2018 dataset, and BURKD is the model of this invention.

[0084] (2) WT, TC, and ET represent Whole Tumor, Tumor Core, and Enhancing Tumor, respectively.

[0085] (3) The values ​​in the table are Dice similarity coefficients (DSC). The larger the value, the better the segmentation effect.

[0086] (4) Avg: represents the average DSC score of the three tumor categories. The higher the value, the better the segmentation effect.

[0087] (5) The values ​​in bold black in the table indicate the best segmentation performance;

[0088] (6) Compared with the state-of-the-art models MCTSeg, M3AE, U-Net-MFI, RFNet, mmFormer and SMU-Net, the three class segmentation metrics DSC of the present invention on the same dataset are improved by 0.5%, 1.1% and 1.6% respectively, with an average improvement of 1.1%.

Claims

1. A relational knowledge distillation method based on category boundary uncertainty, comprising the following steps: 1) Input an image into two network branches simultaneously, first passing it through the feature encoder. The intermediate feature maps F1 and F2 are obtained; then they are decoded respectively. The final segmentation prediction map is obtained. 2) Based on the intermediate feature maps F1 and F2 obtained in step 1), calculate the uniform distillation loss L at the feature level. fc ; 3) Based on the segmentation prediction map obtained in step 1), Calculate pixel-level uniform distillation loss 4) Based on the bi-branch segmentation prediction map obtained in step 1), Calculate the target loss for the first network branch and the second network branch, respectively, based on the corresponding true class label A. 5) The L obtained from step 2) fc And obtained in step 3) and the result obtained in step 4) Calculate the total optimization loss L i ; 6) Train the model based on the loss function from step 5); 7) Use the model trained in step 6) to obtain the segmentation or classification results of the image to be processed; Step 2) specifically involves: based on the intermediate feature maps F1 and F2 obtained in step 1), sequentially calculating the class-specific prototype, the contextual relationships between modeled samples and between classes, and the reweighting of class boundary uncertainty, to obtain the feature-level consistency loss L. fc The specific steps are as follows: Step 2.1) Category-Specific Prototype Generation: Given an intermediate feature map and its corresponding real label A∈R H×W×Z Calculate the class-specific prototypes for the first and second network branches respectively: in, H i W i Z i d represents the length, width, and height of the feature. f The number of channels represents the feature; H, W, Z represent the length, width, and height of the true label A; N represents the number of channels. c This represents the number of pixels belonging to category c. It is an indicator function that represents a value of 1 when the internal condition is True, and 0 otherwise. Step 2.2) Contextual Modeling of Relationships Between Samples and Categories: Based on the category-specific prototype obtained in Step 2.1), the relationship graphs R1 and R2 between samples and categories are obtained using the following formulas: in, Let R1 and R2 represent the prototypes of the c1-th class of the b1-th sample in the first and second network branches, respectively, where R1 and R2 ∈ R. (B*C)×(B*C) B represents the number of samples in a batch, and C represents the number of semantic categories;<a,b> R represents the operation of multiplying the dot product of elements a and b. i (b1,c1; b2,c2) represents the feature similarity between the c1th class of the b1th sample and the c2th class of the b2th sample in the i-th network branch; then the distance difference between samples and between classes between the two branches is calculated: in, This represents the overall relationship difference between the b1-th sample and the c1-th category between the first and second network branches; Step 2.3) Reweighting based on boundary uncertainty: Based on the final segmentation prediction map obtained in Step 1). The prediction results are normalized using the softmax method. Then calculate the category-specific weight. Among them, W i ={W i c }∈R C C represents the number of categories. The c-th category of branch i is represented; σ() represents the softmax function; The final uncertainty weights are obtained by averaging the values ​​of the two network branches: Step 2.4) The final feature-level relation consistency loss is: Among them, W b1,c1 W c A subset of W c The b1 sample.

2. The relational knowledge distillation method based on category boundary uncertainty according to claim 1, characterized in that, In step 1): The first network branch includes a feature encoder. decoder The second network branch includes a feature encoder. decoder 3. The relational knowledge distillation method based on category boundary uncertainty according to claim 1, characterized in that, Step 3) specifically involves: based on the bi-branch predicted segmentation map obtained in step 1). Calculate the pixel-level consistency loss for branches 1 and 2 using KL divergence: in, This indicates that the second network branch constrains the first network branch. This indicates that the first network branch is used to constrain the second network branch, and τ represents the temperature factor hyperparameter.

4. The relational knowledge distillation method based on category boundary uncertainty according to claim 1, characterized in that, Step 4) specifically involves: based on the segmentation prediction map obtained in step 1). And its corresponding real label A, to obtain the target task loss of the network branch. Where i represents the i-th network branch, and CE(.) and Dice(.) represent the cross-entropy and Dice loss functions, respectively.

5. The relational knowledge distillation method based on category boundary uncertainty according to claim 1, characterized in that, Step 5) specifically involves: based on the feature-level consistency loss L from step 2). fc Pixel-level consistency loss in step 3) and the task loss in step 4). Calculate the overall objective loss L for branch i. i :

Images