Paper Title: Bayesian Low-Rank LeArning (Bella): A Practical Approach to Bayesian Neural Networks.






[
    [
    August 5, 2024
==================




§ ABSTRACT 
Computational complexity of Bayesian learning is impeding its adoption in practical, large-scale tasks. Despite demonstrations of significant merits such as improved robustness and resilience to unseen or out-of-distribution inputs over their non-Bayesian counterparts, their practical use has faded to near insignificance. In this study, we introduce an innovative framework to mitigate the computational burden of Bayesian neural networks (BNNs).  
Our approach follows the principle of Bayesian techniques based on deep ensembles, but significantly reduces their cost via multiple low-rank perturbations of parameters arising from a pre-trained neural network. Both vanilla version of ensembles as well as more sophisticated schemes such as Bayesian learning with Stein Variational Gradient Descent (SVGD), previously deemed impractical for large models, can be seamlessly implemented within the proposed framework, called Bayesian Low-Rank LeArning (Bella).
In a nutshell, i) Bella achieves a dramatic reduction in the number of trainable parameters required to approximate a Bayesian posterior; and ii) it not only maintains, but in some instances, surpasses the performance of conventional Bayesian learning methods and non-Bayesian baselines. Our results with large-scale tasks such as , , ,  with ,  demonstrate the effectiveness and versatility of Bella in building highly scalable and practical Bayesian deep models for real-world applications.






§ INTRODUCTION






Bayesian deep learning <cit.> provides mechanisms for building more robust predictive models—i) robustness to adversarial attacks; ii) unseen or out-of-distribution data—and iii) a theoretical framework for estimating model uncertainty <cit.>. In particular, quantifying uncertainty enables more reliable decision-making and facilitates identifying potential model vulnerabilities in uncertain regions of the input space.  Consequently, embracing Bayesian deep neural networks (BNNs) represents a significant stride towards building more reliable and trustworthy AI systems for various real-world applications (e.g., autonomous driving, medical image analysis, etc.). Unfortunately, their practical use is encumbered by their computational complexity. 

Unlike traditional alternatives with point estimates—a single set of model parameters mapping inputs to outputs— BNNs learn the distribution of model parameters to offer a distribution over possible predictions. 
Consider a neural network f(,) with input 𝐱 and weights  and a prior distribution on network weights p(). The likelihood function p(|) determines how well the network with weights  fits the data  <cit.>. Bayesian inference integrates the likelihood and the prior using Bayes' theorem to derive a posterior distribution over the space of weights, given by
p(|) =  p(|)p() /p()
and seek to compute the predictive distribution



    p(y|,𝒟)=∫ p(y|,) p(|𝒟) d.





Despite evidence to support the adoption of Bayesian learning, their widespread use  is hindered by several challenges; a primary issue is the intractability of the posterior distribution p(|) in practical learning tasks.

The exact solution for the posterior, even for networks of moderate size, is impractical. Because of network complexity and the high-dimensional integral of the resulting denominator. 
This necessitates the use of approximations. 









In this paper, we consider approximations of the posterior as ensemble or Stein Variational Gradient Descent (SVGD) but provide a practical method to learn BNNs for large-scale, learning tasks. 

Effectively, SVGD employs a mechanism to separate the parameters of the ensembles (i.e. particles), ensuring the representation of alternative modes in the posterior distribution. But computational demands of the approximation, especially the memory demands for storing all parameters, hinder adoption in large-scale models such as  <cit.>. 
However, a recent conjecture <cit.> suggests that the permutation of optimal model parameters can lead to models that are not functionally different. Further, the path between modes may be linear. Therefore, we consider spawning parameter particles (models) by adding a linear interpolation to a pre-trained model's parameters. But, it would still require updating a large number of model parameters and we are still confronted by computation complexity. The recent conjecture in <cit.> suggests modes may yet exist in the constrained region in the parameter space. Therefore, we propose only adding a linear interpolation to a portion of the model parameters and, subsequently, only learning them as illustrated in Figure <ref> for the  task—a benchmark task for out-of-distribution evaluations. We summarise our key contributions as: 



    
  * We propose a new Bayesian learning framework, , for SVGD approximation of the posterior—spawning particles or models, by exploiting the availability of pre-trained models, with linear interpolation of a constrained set of model parameters for learning with SVGD to approximate the posterior. 
    
  * Our approach more efficiently captures the complexity and multi-modality of the solution space 
    compared to current SVGD but 
    at a fraction of the cost—we observe on-par performance in terms of uncertainty estimation, performance improvement, and robustness, updating only 0.3% or less parameters. 
    
    
    
  * We demonstrate to perform on par or better than baselines—ensembles or current SVGD implementations. consistently outperforms the non-BNN counterparts on 8 datasets; 
    image classification (, ), Out-of-distribution (, [Notably, our approach sets a new state-of-the-art for the  (see <ref> for more details) dataset challenge according to the leaderboard <cit.>. 
    ] <cit.>, ), and  as well as adversarial robustness. 
    Notably, our approach, on a multi-modal architecture such as  <cit.>, leads to improved performance and uncertainty estimation highly correlated with human confidence. 





§ RELATED WORK




Increasing network complexity has led to advances in parameter-efficient fine tuning and recent advances in exploring these methods for Bayesian learning, predominantly for language models. We discuss these in the context of our work.


Parameter-Efficient Fine Tuning. In contrast to fine-tuning all parameters, recent research proposed inserting adapters in between existing neural layers to reduce the number of trainable parameters and, subsequently, the compute (GPU consumption) <cit.>.
Hu et al. <cit.> use a bottleneck structure to impose a low-rank constraint on the weight updates, named LoRA. The key functional difference is that LoRA can be merged with the main weights during inference, thus avoiding the introduction of any latency whilst significantly reducing the number of parameters. 

Fine-Tuning Approaches and Bayesian Deep Learning. 
Previous research investigating the application of fine-tuning approaches for BNNs have predominantly focused on large language models (LLMs), e.g.  <cit.>. Notably, marking a departure from the conventional methods relying primarily on tuning the network's parameters, these studies chose to define priors and approximate posterior over low rank attention weights. Concurrently, Yang et al. <cit.> introduced the concept of Laplace LoRA to incorporate Bayesian concepts to enhance the calibration of fine-tuned LLMs. However, Laplace's method relies on a Gaussian approximation of the posterior distribution. Whilst this can be effective for unimodal and symmetric distributions, the approach does not fully encapsulate the intricacies of more complex posteriors, particularly in neural networks where the posterior has multimodality and asymmetry <cit.>. Deep ensembles <cit.> typically perform better in practice compared with variational and Laplace methods, due to their ability to capture multiple modes. When employing fine-tuning, a direct application of ensembling for LoRAs was considered in <cit.>. Some interpretations of deep ensembles suggest that they approximate gradient flows in function spaces and that building desirable properties into an ensemble (such as repulsive behavior), is possible <cit.>. SVGD <cit.>, can be viewed in a similar vein. However, while these more sophisticated, repulsive, ensembling approaches are highly impractical in the pre-training phase, we argue that their expressivity can be brought to bear precisely in tandem with low-rank fine-tuning, which is the viewpoint we adopt in this contribution.

Building upon these foundations, our research represents a pioneering effort to apply the principles of repulsive ensemble-based low-rank fine-tuning to computer vision. In particular, we bridge pre-training and fine-tuning phases with recent conjectures on mode connectivity <cit.>. Our methodology not only capitalizes on the efficiency of fine-tuning techniques, e.g. <cit.> but also innovatively addresses the scalability challenges inherent to BNNs. 
Overall, our work sets a new precedent in applying Bayesian approaches to computer vision tasks by offering a scalable and efficient framework for enhancing model performance.








§ BACKGROUND ON STEIN VARIATIONAL GRADIENT DESCENT (SVGD) APPROXIMATIONS FOR A BAYESIAN POSTERIOR


Here, we provide: i) a brief overview of methods to address the intractability of the posterior distribution p(|) in Bayesian learning tasks; and ii) we elaborate on our motivation for seeking a practical method to learn BNNs for large-scale tasks with SVGD.



Bayesian inference techniques have been integral to the development of neural networks, with a rich history underscored by previous works <cit.>. These methodologies provide a structured framework for generating reliable and interpretable estimates of uncertainty. However, the focus of prior research on Bayesian neural networks has predominantly been on the initial pretraining stages rather than on subsequent fine-tuning processes <cit.>. This emphasis on pretraining presents significant challenges in terms of scalability and resource requirements, limiting the application of Bayesian neural networks to smaller datasets and network architectures with a minimal number of parameter particles. Consequently, the scalability issues hinder the extension of Bayesian neural networks to larger datasets, such as ImageNet <cit.>, and more complex network architectures, like  <cit.>, which involve millions of parameters.



Variational Inference (VI) <cit.> and Markov Chain Monte Carlo (MCMC) <cit.> are two primary approximate Bayesian inference frameworks. The former substitutes the true posterior with a tractable alternative while the latter involves sampling. 
However, accurately computing the posterior with either MCMC or VI becomes computationally infeasible when dealing with large-scale networks containing millions of parameters. Although approximations can be obtained more efficiently with VI <cit.>, VI is also demonstrably too restrictive to resemble the multi-modality of the true posterior and suffers from mode collapse <cit.>.

Stein Variational Gradient Descent (SVGD) <cit.> is an alternative approximate Bayesian technique which combines the strengths of MCMC and VI by transporting a set of parameter particles to fit the true posterior distribution, while encouraging diversity among the particles, by incorporating a repulsive term in the parameter updates. 
This diversity prevents the mode collapse and enables learning multiple models to represent various patterns in the data. 
Using n samples from the posterior (parameter particles), SVGD modifies the gradient descent as: 

    _i  =   _i    - ϵ_i ^*(_i)  
       with  ^*() = ∑_j=1^n[  k(_j, )  ∇__jlog p(_j|) -γ/n∇__j k(_j, )] .

 Here, _i is the ith particle, k(·, ·) is a kernel function that measures the similarity between particles and γ is a hyper-parameter. 
Notably, the kernel function encourages the particles to be dissimilar in order to capture more diverse samples from the posterior and γ controls the trade-off between the diversity of the samples versus the minimization of the loss.








§ BAYESIAN LOW-RANK LEARNING () FOR SVGD







The problem with the current SVGD approach in large deep neural networks is its huge computational cost. This renders it infeasible to train efficiently and to scale to a sufficient number of parameter particles for accurately approximating the posterior distribution, which currently remains coarse.
In this work, we propose to capitalize on the low-rank representations of fine-tuning in order to construct a practical and scalable variant of SVGD. Consider any dense layer, for which there is a fixed pre-trained weight matrix _0∈ℝ^d_1× d_2 with d_1, d_2 the corresponding numbers of hidden units. We consider n low-rank perturbations of _0 as 

    _i = _0 + Δ_i = _0 + _i_i ,   i=1,…,n.

where _i∈ℝ^d_1 ×, _i∈ℝ^× d_2 are the low-dimensional update parameters,  and r≪ d_1,d_2 is the rank of the update.
Now Bella proceeds as the joint SVGD on (_i,_i), with updates

    _i  =   _i - ϵ_i ∑_j=1^n_j^*(_i),   _i  =   _i - ϵ_i ∑_j=1^n_j^*(_i)   
       with
    _j^*(_i)    =  k(_j_j, _i_i) ∇__iℓ(f(;_0 + _i_i),y) -γ/n∇__i k(_j_j, _i_i),
    _j^*(_i)    =  k(_j_j, _i_i) ∇__iℓ(f(;_0 + _i_i),y) -γ/n∇__i k(_j_j, _i_i) .

Here, we have placed the (improper) uniform prior on (_i,_i), but other choices are possible. Uniform prior implies that we only require the gradients of the log-likelihood computed on a mini-batch of data, as well as the kernel function to encourage particle diversity.   
Note that the kernel function on (_i,_i) is given by k(_0 + _j_j, _0 + _i_i), which ensures that the similarity is computed on the original parameter space, and in the commonly used case of shift-invariant kernels, this simplifies to k(_j_j, _i_i). Further simplifications are obtained for specific kernel functions – in particular, in the case of Gaussian RBF, while the naive implementation would require the cost of O(rd_1d_2) for a single kernel evaluation, we can bring it down to O(r^2(d_1+d_2)) using standard trace manipulation, as described in the appendix. This procedure can be repeated across all dense layers. 

Bella introduces a significant improvement in the efficiency of model training and execution. By utilizing the same pre-trained weights _0 across all parameter particles but allowing for individual low-rank adaptations Δ_i, we achieve a balance between parameter sharing and the diversity necessary for effective learning. 
Bella significantly reduces the parameter space from the full matrix's d_1d_2 to just (d_1 + d_2), thereby enhancing both efficiency and scalability.
This setup not only reduces the computational burden during training but also streamlines the process at inference time. The heavy lifting is done once by loading the large base model _0, and the lightweight low-rank adapters Δ_i can be dynamically applied with minimal overhead in order to approximate the posterior predictive distribution   as  p ( y^* |𝐱^*,  ) ≈1/n∑_i=1^n p(y^* |𝐱, θ_0 + Δ_i).
 This approach is particularly advantageous in large-scale models, where the weight matrices _0 are of substantial dimensions.





§ EMPIRICAL EXPERIMENTS AND RESULTS
 











In this section, we provide an in-depth overview of our experimental setup, detailing the methodology, equipment, and procedures employed in the implementation of our . We meticulously outline the configuration settings, dataset, and the criteria used for evaluating outcomes, ensuring a transparent and reproducible framework. Following this, we present the experimental results, offering a thorough analysis of the data obtained, including statistical evaluations and interpretations of the findings in the context of our research objectives. This comprehensive approach ensures a clear understanding of the experiment's execution and its contributions to the field.






 §.§ Experimental Set-up

Datasets. In this research, we have employed a variety of datasets, each selected for its relevance and contribution to the field of image recognition and computer vision. These datasets include ,  <cit.>,  <cit.>,  <cit.>,  <cit.>,  <cit.>, and  <cit.>. 
We also employ VQA v2 dataset containing billions of questions answers given hundred thousands of images <cit.>, which is utilized in Visual Question Answering (VQA) task.
Detailed information about datasets is presented in <ref> and <ref>.



Networks. In our experiments, we employed the CLIP ViT-B/32 model <cit.>, a pre-trained variant utilizing contrastive supervision from image-text pairs, as initially introduced in the seminal CLIP research <cit.>. We conducted end-to-end fine-tuning, adjusting all model parameters, a strategy typically yielding higher accuracy compared to training only the final linear layer. 
In our methodology, we selectively fine-tune only the image encoder of the CLIP model while maintaining the text encoder frozen as it is common <cit.>. To elaborate, both during training and inference phases, the final output is derived through an inner product between the output produced by the image encoder and the embedding generated by the frozen text encoder. For both our ensemble-based methods which include both traditional ensemble <cit.> and the Stein Variational Gradient Descent (SVGD) approach <cit.>, we employed the technique of utilizing the logits (unnormalized outputs) of the models, as described in <cit.>. This strategy has been consistently applied in our ensemble configurations, encompassing traditional ensembles as well as SVGD.
We also employ Ensemble networks as it was shown to be a good and efficient way to approximate the posterior distribution to be baselines <cit.>.
For VQA task, we employ the SoTA LLaVA-1.5 7B with billions of parameters to showcase the effectiveness of our on large-scale network architecture.
Details regarding hyper-parameters are in <ref>.





 §.§ Benchmark Task Performance and Cost Efficiency





The performance comparison of our Bella models with their respective base models, as well as with Vision Bayesian Lora (VBL)—a derived Laplacian sampling method from <cit.> for vision tasks—is delineated in <ref>. The respective costs comparison on training (in terms of trainable parameters) is shown in <ref>—notably, the grey columns are the ones used for generating the results in <ref>. 
































Impressively, across a spectrum of benchmark computer vision datasets such as , , , and , the models demonstrate superior performance. This is achieved with only a fraction of cost (trainable parameters) as shown in <ref>, underscoring the models' proficiency in parameter efficiency without compromising on accuracy. Further, the models surpass the performance of Single models, while employing a comparable level of computational resources—see the Memory Consumption of models used, in <ref> reporting approximately 4.5 GB for compared to 5.05 GB for Single. The results across the benchmarks attest to the efficacy of our methodology.

Notably, in <ref> with SVGD, with 1.6% of the trainable parameters in comparison to the Single baseline, leads to approximately 2% and 10% increase in performance on the OOD tasks of  and , respectively; whilst achieving comparable performance with the current SVGD implementation (SVGD baseline model). 


Along with <ref>, <ref> demonstrates the primary benefit of our approach—the significant reduction in memory and storage needs. For example, <ref>,


we obtain approximately a 5× reduction in model size—that is from 2,222 MB for a 5 particle SVGD baseline model to just 440 MB for a SVGD model with 5 particles. This efficiency not only reduces the demands on GPUs but also minimizes potential I/O bottlenecks.

Moreover, the reduced GPU demand, as shown in <ref>, facilitates larger mini-batch sizes during training to speedup the training process. More crucially, it enables the enhancement of the Bayesian posterior's parameter particles to over 100, a feat previously unattainable with current SVGD implementations. 
Significantly, constructing a 100 parameter Bayesian approximation consumes only 5.19 GB memory compared to current SVGD implementations for Bayesian models (SVGD base) needing over 6 GB for even a 3 particle approximation.





Interestingly, our on par results of models with current SVGD and ensemble approximations of the posterior provide empirical evidence that with constrained model parameters, it is still possible to reach the diverse modes of the posterior. Further, the results support recent conjectures on mode connectivity <cit.>.










 §.§ Performance on Out-Of-Disrtibution (OOD) Datasets













Assessing the robustness of machine learning systems to unseen conditions is crucial, especially their ability to generalize to out-of-distribution (OOD) data. We evaluate robutness to ODD using multiple OOD benchmarks. 

First, we use the  dataset, one of the most diverse domain adaptation datasets, covering a wide range of visual styles from real images to abstract art. This variety provides a challenging test bed for algorithms aiming to bridge different visual domains. Our study involves training the CLIP network on 's `Real' subset and testing its generalization across various domains (see <ref>). Second, , a corrupted version of  helps assess model generalization and robustness (see <ref>). Third, the  dataset, with significant label overlap with , serves as a relevant OOD test case (see <ref>).

models, with significantly better efficiency, demonstrate competitive performance against more resource-intensive implementations with Ensemble and SVGD baselines across both  and  benchmarks. All of the Bayesian approximations, including our scalable and efficient method, outperform the Single model baseline.


















 §.§ Comparing Uncertainty Estimations


Bayesian models capable of providing a theoretical basis for measuring  model uncertainty. Also known as epistemic uncertainty, refers to uncertainty stemming from limitations in our knowledge or understanding of the underlying data generating process or the model itself. One of the ways to quantify model uncertainty is through mutual information estimates, following  <cit.>.



Mutual Information (MI). 
This is the mutual information between the output prediction and the posterior over model parameters , and can be used as a measure of epistemic (model) uncertainty.

It can be expressed as:
MI(, y |𝒟, ) = H[p(y|𝒟, )] - 𝔼_p(|𝒟) H[p(y |,)]

If the parameters at input are well defined ( e.g., data seen during training), then we would gain little information from the obtaining label, or the MI measured will be low. 

We employ MI to measure uncertainty to investigate whether the approximations of the posterior leads to uncertainty estimates commensurate with those obtained from SVGD baselines. This provides empirical evidence of a functional equivalence of the approximations of the posterior to that obtained from the current computationally intensive implementation of SVGD.



Datasets. We utilize the  task, featuring corrupted images, to examine the uncertainty of model predictions trained on the standard  dataset. Additionally, we assess the uncertainty measures on the  dataset, which is characterized by inherent dataset shifts within itself.


Results.  <ref> demonstrates the effectiveness of our approach to estimate uncertainty. Our perform similarly to the SVGD base model, with a slightly better uncertainty on misclassified images of  and corrupted  datasets (under brightness corruption with the maximum intensity), see details and other corruption types in <ref>.
































 §.§ Robustness against Adversarial Examples






In this section, we examine the resilience of our proposed against adversarial attacks, specifically employing the L_∞ Fast Gradient Sign Method (FGSM) across various attack budgets as detailed in <ref>. This analysis aims to benchmark the robustness of our method in comparison to traditional models under adversarial conditions. We employ the robustness benchmark <cit.> to deploy the attack on  test set and report results in <ref>.




The findings presented in <ref> reveal that conventional models such as SVGD and Ensemble exhibit just slightly greater resistance to adversarial attacks. We attribute this enhanced robustness to the broader diversity in model parameters, which stems from their capacity to adjust the entire network's parameters, unlike the models.

Significantly, despite operating within the same computational constraints as a singular network model, our demonstrates enhanced efficacy in mitigating adversarial attacks, thereby bolstering its robustness.








 §.§ Ablation Studies




This section undertakes a series of ablation studies to examine the effects of various components within . 
Our analysis includes an exploration of the training costs associated with different ranks and their consequent influence on model performance. Given that incorporates multiple parameter particles, we also delve into how varying the number of these particles affects 's efficacy. 
Additionally, we explore the application of low-rank adapters across different layers and assess their impact. Further details on other studies are in <ref>. We show in <ref> that we achieve state-of-the-art performance on .



Ablations on rank . As outlined in <ref>, we substitute the network's extensive full matrix with low-rank matrices, which are defined by the `rank' parameter (). This section aims to assess how this parameter influences 's performance.

To comprehensively demonstrate the rank's impact, we perform our analysis on the demanding large-scale  dataset. The findings, illustrated in <ref>, shed light on how the number of ranks affects performance. Utilizing a small rank significantly reduces the parameter space, but this constriction limits 's ability to learn effectively and achieve optimal performance. Conversely, enhancing the rank to 4 markedly boosts 's efficiency. It's observed that performance plateaus at a rank of 16, indicating a saturation point.





































Ablations on the number of particles n. In this section, we delve into the influence of the quantity of parameter particles on the performance of . The findings, depicted in <ref>, reveal an improvement in 's performance with an increase in the number of parameter particles. 

This outcome is both intuitive and insightful, as a larger ensemble of parameter particles enhances the approximation of the Bayesian posterior more effectively.
















Ablations on the number of trainable parameters. 
We will compare the number of trainable parameters between and full SVGD in model performance. Critically, higher number of trainable parameters means higher cost to train. To show generalization, we also employ another large-scale challenging dataset Imagenet in this experiment. 








The starting point of the plot in <ref> associated with the  dataset is for with =2 (consisting of 5 particles), and the end point is associated with full SVGD (SVGD base) with the same number of particles. As  increases, the accuracy of Bellas improves and after a while becomes plateau around 95.08, comparable to the accuracy of full SVGD (95.21). This underscores the efficiency and advantages of our proposed . Despite utilizing a substantially smaller pool of trainable parameters (around 0.3% for =4) compared to full SVGD models, demonstrates performance on par with these more parameter-intensive alternatives.










































 §.§ Generalization to a Visual Question Answer (VQA) Task


In this section, we extend the application of our to another challenging vision task, Visual Question Answering (VQA), as detailed in <cit.>. We leverage the state-of-the-art, pre-trained, large multi-modal model LlaVA  <cit.> for this purpose. Utilizing LlaVA transforms VQA into a process where an image and a natural-language question are inputs, and the model generates a free-form, open-ended text answer. Answering questions in VQA requires various intelligent capabilities, including not only image recognition but also complex reasoning. For this task, we employ
 <cit.> dataset containing 204,721 images, more than 1 Billion (1B) questions and 10B ground-truth answers in total. 

There are three main types of answers: Yes/No, a Number, and Other. 



Model. In our experiments, we employed our proposed on top of LlaVA-1.5-7B <cit.>. Further details about the dataset, model and metrics are deferred to Appendix <ref>.



Results (Performance). 
The primary outcomes for Yes/No and Number answer queries are detailed in <ref>, respectively. We chose these question types to ensure a fair comparison and avoid semantic mismatches in open-ended answers in Others.
A distinctive feature of our is its reduced uncertainty estimates to  accurate predictions, coupled with increased uncertainty estimates to  incorrect ones. Moreover, it surpasses the Single base model in terms of Accuracy and Exact Match metrics.


Notably, our method is efficient, particularly in contrast to the resource-intesive baseline Bayesian models tailored for this task (<cit.>). The computational requirements for using these baselines render the application of full SVGD or ensemble alternatives impractical, thereby highlighting the impact on practical applications by harnessing the effectiveness of Bayesian models with . 







Results (Model Uncertainty and Human Confidence). Further, we investigate the relationship between model uncertainty attained from and human confidence, facilitated by multiple annotators in the VQA dataset.  For this, we measure the correlation between the model disagreement by measuring the predictive entropy vs the human annotations of the answers. 







In Fig. <ref> we shows the correlation between the entropy of the model outputs and human confidence levels. By setting a specific threshold for the model's entropy, we effectively bifurcate our predictions into two distinct categories: those with lower entropy fall into the `Low Entropy' group, signalling reduced uncertainty within the model's assessments, while predictions surpassing this entropy threshold are allocated to the `High Entropy' segment, indicative of greater uncertainty. Intriguingly, our observations reveal a consistent negative correlation between the model's entropy levels and human confidence across different types of queries, such as Yes/No and Number questions. This pattern suggests that the model's entropy can serve as a reliable gauge of certainty, mirroring human judgment in its response to varying levels of uncertainty.






§ CONCLUSION

In this paper, we introduce an innovative approach for creating an efficient Bayesian Neural Network (BNN) approximation using only a base pre-trained model. Our approach, called Bella, demonstrates remarkable compatibility with full-rank BNN approximations like SVGD or Ensemble methods, outperforming single-network solutions across a range of tasks, from pure classification to out-of-distribution (OOD) generalization and uncertainty quantification. Our research paves the way for effective BNN implementation, facilitating the development of reliable and robust machine-learning models capable of scaling to large networks and datasets. Bella reinforces the effectiveness of employing such a simple, efficient and effective training method with diverse representations over the conventional single fine-tuning techniques prevalent in the field.





10

neal2012bayesian
Radford M Neal.
Bayesian learning for neural networks, volume 118.
Springer Science & Business Media, 2012.

advbnn
Xuanqing Liu, Yao Li, Wu Chongruo, and Hsieh Cho-Jui.
ADV-BNN: Improved Adversarial Defense Through Robust Bayesian Neural Network.
In International Conference on Learning Representations (ICLR), 2019.

doan2022bayesian
Bao Gia Doan, Ehsan M Abbasnejad, Javen Qinfeng Shi, and Damith C Ranasinghe.
Bayesian learning with information gain provably bounds risk for a robust adversarial defense.
In International Conference on Machine Learning (ICML), 2022.

BAL
Nanyang Ye and Zhanxing Zhu.
Bayesian adversarial learning.
In Advances in Neural Information Processing Systems (NeurIPS), 2018.

wilson_evaluating_2022
Andrew Gordon Wilson, Pavel Izmailov, Matthew D. Hoffman, Yarin Gal, Yingzhen Li, Melanie F. Pradier, Sharad Vikram, Andrew Foong, Sanae Lotfi, and Sebastian Farquhar.
Evaluating approximate inference in bayesian deep learning.
In Proceedings of the NeurIPS 2021 Competitions and Demonstrations Track.

abbasnejad2020gold
Ehsan Abbasnejad, Iman Abbasnejad, Qi Wu, Javen Shi, and Anton van den Hengel.
Gold seeker: Information gain from policy distributions for goal-oriented vision-and-langauge reasoning.
In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2020.

Bishop2006
Christopher M. Bishop.
Pattern Recognition and Machine Learning.
Springer, 2006.

Luo2023SegCLIP
Huaishao Luo, Junwei Bao, Youzheng Wu, Xiaodong He, and Tianrui Li.
SegCLIP: Patch aggregation with learnable centers for open-vocabulary semantic segmentation.
In Proceedings of the 40th International Conference on Machine Learning (ICML), 2023.

git_rebasin
Samuel Ainsworth, Jonathan Hayase, and Siddhartha Srinivasa.
Git re-basin: Merging models modulo permutation symmetries.
In The Eleventh International Conference on Learning Representations, 2023.

gueta2023knowledge
Almog Gueta, Elad Venezian, Colin Raffel, Noam Slonim, Yoav Katz, and Leshem Choshen.
Knowledge is a region in weight space for fine-tuned language models.
arXiv preprint arXiv:2302.04863, 2023.

leaderboard
Grand Challenge.
Camelyon17 leaderboard.
<https://camelyon17.grand-challenge.org/evaluation/challenge/leaderboard/>, 2024.
Accessed: 08-Mar-2024.

wilds2021
Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang, Akshay Balsubramani, Weihua Hu, Michihiro Yasunaga, Richard Lanas Phillips, Irena Gao, Tony Lee, Etienne David, Ian Stavness, Wei Guo, Berton A. Earnshaw, Imran S. Haque, Sara Beery, Jure Leskovec, Anshul Kundaje, Emma Pierson, Sergey Levine, Chelsea Finn, and Percy Liang.
WILDS: A benchmark of in-the-wild distribution shifts.
In International Conference on Machine Learning (ICML), 2021.

liu2023llava
Haotian Liu, Chunyuan Li, Qingyang Wu, and Yong Jae Lee.
Visual instruction tuning.
In Advanced in Neural Information Processing Systems (NeurIPS), 2023.

houlsby_parameter-efficient_2019
Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin de Laroussilhe, Andrea Gesmundo, Mona Attariyan, and Sylvain Gelly.
Parameter-Efficient Transfer Learning for NLP.
arXiv:1902.00751, 2019.

rebuffi_learning_2017
Sylvestre-Alvise Rebuffi, Hakan Bilen, and Andrea Vedaldi.
Learning multiple visual domains with residual adapters.
arXiv:1705.08045, 2017.

lin-etal-2020-exploring
Zhaojiang Lin, Andrea Madotto, and Pascale Fung.
Exploring versatile generative language model via parameter-efficient transfer learning.
In Findings of the Association for Computational Linguistics (EMNLP), 2020.

hu2022lora
Edward J Hu, yelong shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen.
LoRA: Low-rank adaptation of large language models.
In International Conference on Learning Representations (ICLR), 2022.

fan2020bayesian
Xinjie Fan, Shujian Zhang, Bo Chen, and Mingyuan Zhou.
Bayesian attention modules.
In Advances in Neural Information Processing Systems (NeurIPS), 2020.

zhang2021bayesian
Shujian Zhang, Xinjie Fan, Bo Chen, and Mingyuan Zhou.
Bayesian attention belief networks.
In International Conference on Machine Learning (ICML), 2021.

yang2023bayesian
Adam X Yang, Maxime Robeyns, Xi Wang, and Laurence Aitchison.
Bayesian low-rank adaptation for large language models.
In International Conference on Learning Representations (ICLR), 2024.

izmailov2021bayesian
Pavel Izmailov, Sharad Vikram, Matthew D Hoffman, and Andrew Gordon Wilson.
What are bayesian neural network posteriors really like?
In International Conference on Machine Learning (ICML), 2021.

lakshminarayanan2017simple
Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell.
Simple and scalable predictive uncertainty estimation using deep ensembles.
In Advances in neural information processing systems (NeurIPS), 2017.

wang2024lora
Xi Wang, Laurence Aitchison, and Maja Rudolph.
LoRA ensembles for large language model fine-tuning, 2024.

WilGhaSejKno2023
Veit David Wild, Sahra Ghalebikesabi, Dino Sejdinovic, and Jeremias Knoblauch.
A Rigorous Link between Deep Ensembles and (Variational) Bayesian Methods.
In Advances in Neural Information Processing Systems (NeurIPS), 2023.

liu2016stein
Qiang Liu and Dilin Wang.
Stein variational gradient descent: A general purpose bayesian inference algorithm.
Advances in Neural Information Processing Systems (NeurIPS), 2016.

opendelta
Ning Ding, Yujia Qin, Guang Yang, Fuchao Wei, Zonghan Yang, Yusheng Su, Shengding Hu, Yulin Chen, Chi-Min Chan, Weize Chen, Jing Yi, Weilin Zhao, Xiaozhi Wang, Zhiyuan Liu, Hai-Tao Zheng, Jianfei Chen, Yang Liu, Jie Tang, Juanzi Li, and Maosong Sun.
Parameter-efficient fine-tuning of large-scale pre-trained language models.
Nature Machine Intelligence, 2023.

ding2022delta
Ning Ding, Yujia Qin, Guang Yang, Fuchao Wei, Zonghan Yang, Yusheng Su, Shengding Hu, Yulin Chen, Chi-Min Chan, Weize Chen, et al.
Delta tuning: A comprehensive study of parameter efficient methods for pre-trained language models.
arXiv preprint arXiv:2203.06904, 2022.

dettmers2023qlora
Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, and Luke Zettlemoyer.
Qlora: Efficient finetuning of quantized llms.
arXiv preprint arXiv:2305.14314, 2023.

neal2011mcmc
Radford M Neal et al.
Mcmc using hamiltonian dynamics.
Handbook of markov chain monte carlo, 2(11), 2011.

MacKay1991
David MacKay.
Bayesian model comparison and backprop nets.
In Advances in Neural Information Processing Systems (NeurIPS), 1991.

welling2011bayesian
Max Welling and Yee W Teh.
Bayesian learning via stochastic gradient langevin dynamics.
In Proceedings of the 28th international conference on machine learning (ICML-11), pages 681–688. Citeseer, 2011.

blei2017variational
David M Blei, Alp Kucukelbir, and Jon D McAuliffe.
Variational inference: A review for statisticians.
Journal of the American statistical Association, 2017.

wilson2020bayesian
Andrew G Wilson and Pavel Izmailov.
Bayesian deep learning and a probabilistic perspective of generalization.
In Advances in neural information processing systems (NeurIPS), 2020.

wilson_deep_post
Andrew Gordon Wilson and Pavel Izmailov.
Deep ensembles as approximate bayesian inference.

russakovsky2015imagenet
Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, et al.
Imagenet large scale visual recognition challenge.
International journal of computer vision, 2015.

blundell2015weight
Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra.
Weight uncertainty in neural network.
In International Conference on Machine Learning (ICML), 2015.

cifar
Alex Krizhevsky, Geoffrey Hinton, et al.
Learning multiple layers of features from tiny images.
2009.

hendrycks2019benchmarking
Dan Hendrycks and Thomas Dietterich.
Benchmarking neural network robustness to common corruptions and perturbations.
arXiv preprint arXiv:1903.12261, 2019.

stl10
Adam Coates, Andrew Ng, and Honglak Lee.
An Analysis of Single Layer Networks in Unsupervised Feature Learning.
In International Conference on Artificial Intelligence and Statistics (AISTATS), 2011.

bandi2018detection
Peter Bandi, Oscar Geessink, Quirine Manson, Marcory Van Dijk, Maschenka Balkenhol, Meyke Hermsen, Babak Ehteshami Bejnordi, Byungjae Lee, Kyunghyun Paeng, Aoxiao Zhong, et al.
From detection of individual metastases to classification of lymph node status at the patient level: the camelyon17 challenge.
IEEE Transactions on Medical Imaging, 2018.

peng2019moment
Xingchao Peng, Qinxun Bai, Xide Xia, Zijun Huang, Kate Saenko, and Bo Wang.
Moment matching for multi-source domain adaptation.
In Proceedings of the IEEE/CVF international conference on computer vision (CVPR), 2019.

vqa_v2
Yash Goyal, Tejas Khot, Douglas Summers-Stay, Dhruv Batra, and Devi Parikh.
Making the V in VQA matter: Elevating the role of image understanding in Visual Question Answering.
In Conference on Computer Vision and Pattern Recognition (CVPR), 2017.

ilharco_gabriel_2021_5143773
Gabriel Ilharco, Mitchell Wortsman, Ross Wightman, Cade Gordon, Nicholas Carlini, Rohan Taori, Achal Dave, Vaishaal Shankar, Hongseok Namkoong, John Miller, Hannaneh Hajishirzi, Ali Farhadi, and Ludwig Schmidt.
Openclip, 2021.

Radford2021LearningTV
Alec Radford, Jong Wook Kim, Chris Hallacy, A. Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Ilya Sutskever.
Learning transferable visual models from natural language supervision.
In International Conference on Machine Learning (ICML), 2021.

ortiz2023task
Guillermo Ortiz-Jimenez, Alessandro Favero, and Pascal Frossard.
Task arithmetic in the tangent space: Improved editing of pre-trained models.
arXiv preprint arXiv:2305.12827, 2023.

gontijo2021no
Raphael Gontijo-Lopes, Yann Dauphin, and Ekin D Cubuk.
No one representation to rule them all: Overlapping features of training methods.
arXiv preprint arXiv:2110.12899, 2021.

gal2016uncertainty
Yarin Gal et al.
Uncertainty in deep learning.
2016.

cleverhans
Nicolas Papernot, Fartash Faghri, Nicholas Carlini, Ian Goodfellow, Reuben Feinman, Alexey Kurakin, Cihang Xie, Yash Sharma, Tom Brown, Aurko Roy, Alexander Matyasko, Vahid Behzadan, Karen Hambardzumyan, Zhishuai Zhang, Yi-Lin Juang, Zhi Li, Ryan Sheatsley, Abhibhav Garg, Jonathan Uesato, Willi Gierke, Yinpeng Dong, David Berthelot, Paul Hendricks, Jonas Rauber, and Rujun Long.
Technical report on the cleverhans v2.1.0 adversarial examples library.
arXiv preprint arXiv:1610.00768, 2018.

vqa_v1
Stanislaw Antol, Aishwarya Agrawal, Jiasen Lu, Margaret Mitchell, Dhruv Batra, C. Lawrence Zitnick, and Devi Parikh.
VQA: Visual Question Answering.
In International Conference on Computer Vision (ICCV), 2015.

liu2023improvedllava
Haotian Liu, Chunyuan Li, Yuheng Li, and Yong Jae Lee.
Improved baselines with visual instruction tuning.
arXiv preprint arXiv:2310.03744, 2023.

gade
M. Ehsan Abbasnejad, Qinfeng Shi, Anton van den Hengel, and Lingqiao Liu.
A generative adversarial density estimator.
In IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2019, Long Beach, CA, USA, June 16-20, 2019, pages 10782–10791. Computer Vision Foundation / IEEE, 2019.

gade2
M. Ehsan Abbasnejad, Javen Shi, Anton van den Hengel, and Lingqiao Liu.
GADE: A generative adversarial approach to density estimation and its applications.
Int. J. Comput. Vis., 128(10):2731–2743, 2020.

matena2022merging
Michael S Matena and Colin A Raffel.
Merging models with fisher-weighted averaging.
Advances in Neural Information Processing Systems (NeurIPS), 2022.

wortsman2022model
Mitchell Wortsman, Gabriel Ilharco, Samir Ya Gadre, Rebecca Roelofs, Raphael Gontijo-Lopes, Ari S Morcos, Hongseok Namkoong, Ali Farhadi, Yair Carmon, Simon Kornblith, et al.
Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time.
In International Conference on Machine Learning (ICML), 2022.

wortsman2022robust
Mitchell Wortsman, Gabriel Ilharco, Jong Wook Kim, Mike Li, Simon Kornblith, Rebecca Roelofs, Raphael Gontijo Lopes, Hannaneh Hajishirzi, Ali Farhadi, Hongseok Namkoong, et al.
Robust fine-tuning of zero-shot models.
In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 7959–7971, 2022.

rame2022diverse
Alexandre Rame, Matthieu Kirchmeyer, Thibaud Rahier, Alain Rakotomamonjy, Patrick Gallinari, and Matthieu Cord.
Diverse weight averaging for out-of-distribution generalization.
Advances in Neural Information Processing Systems (NeurIPS), 2022.












Bayesian Low-Rank LeArning (Bella):
A Practical Approach to Bayesian Neural Networks







§ OVERVIEW OF MATERIALS IN THE APPENDICES






We provide a brief overview of the extensive set of additional experimental results and findings in the Appendices that follows. 





    
    
  * Diversity measures between SVGD and its SVGD base model (<ref>)
    
  * The generalization of CIFAR-10 pre-trained models on OOD dataset (STL-10) (<ref>)
    
  * Detailed uncertainty measured on corrupted CIFAR-10-C dataset (<ref>)
    
  * Additional study on weight averaging (<ref>)
    
  * Impact of the pushing parameter on the robustness of models (<ref>)
    
  * Detailed information about VQA task (<ref>)
    
  * Detailed information regarding RBF kernel computation in SVGD (<ref>)
    
  * Detailed information regarding datasets and hyper-parameters utilized in the paper (<ref>)








§ DIVERSITY MEASURES



This section assesses the diversity of the trained , comparing it to the foundational SVGD model. The goal is to determine if can maintain the diversity level of the SVGD base model while utilizing significantly fewer parameters.

Given the lack of a conventional metric for evaluating the diversity among parameter particles, we suggest employing the Kullback–Leibler (KL) Divergence. This measure compares the expected parameters of a Bayesian model to the softmax output for each parameter particle, serving as an indicator of model diversity. We calculate this divergence across 10,000 test samples from the CIFAR-10 dataset. In particular,

    Diversity=1/N∑_i=1^NKL[𝔼_[p(y|_i,)],p(y|_i,)]

where KL is the Kullback–Leibler divergence, N is the number of samples. 



Results. <ref> presents intriguing findings, where the SVGD exhibits comparable and even greater diversity compared to the base SVGD model. We speculate that this increased diversity may stem from the 's use of significantly fewer trainable parameters, reducing the likelihood of overfitting and facilitating the identification of parameter particles with enhanced diversity, all while preserving model performance. These results further underscore the efficiency and effectiveness of our method in implementing a Bayesian neural network with substantial diversity.






§ GENERALIZATION OF CIFAR-10 MODELS ON OOD DATASET (STL-10)



In this section, we evaluate the generalization of networks trained on  on a similar datasets, which is . 

STL-10 dataset. The  dataset is a benchmark for image recognition algorithms, containing 5,000 labeled training images, 8,000 labeled test images across 10 classes, and 100,000 unlabeled images for unsupervised learning. Images are 96x96 pixels, larger than those in similar datasets like , facilitating more detailed models and promoting unsupervised learning techniques.

Experimental Set-up. We choose eight labels in STL-10 test set, same labels observed in  dataset to set a OOD dataset. This includes 6,400 images acting as the OOD dataset for  pre-trained models.

Results. The data presented in Figure <ref> indicate that our approach slightly outperforms the SVGD and Ensemble baseline models, and significantly surpasses the capabilities of a singular network. We propose that the incorporation of a low-rank adaptation mechanism contributes to reducing overfitting, thereby enhancing the model's generalization ability on out-of-distribution (OOD) datasets. This underscores the efficacy of our method in fostering robust generalization across OOD datasets.












§ MEASURE UNCERTAINTY ON CORRUPTED  DATASET





In this section, we apply Mutual Information (MI) as a metric for assessing uncertainty. This approach enables us to compare the performance of our method with models based on Stein Variational Gradient Descent (SVGD) in handling Out-of-Distribution (OOD) tasks, focusing on the CIFAR-10-C dataset. 
We adhere to the corruption types listed in <ref> for this analysis. 



Results. The outcomes, depicted in <ref>, affirm the efficacy of . It demonstrates comparable, and in some instances, superior performance to SVGD-based models, especially notable in most cases such as involving contrast adjustments, Gaussian blur, and pixelation. This indicates a robust ability to detect increased uncertainty in OOD datasets, which is a valuable attribute for enhancing the reliability of machine learning models in unpredictable environments.















§ ADDITIONAL STUDIES




In this section, we conduct further studies on the benefits of our approach utilizing recently developed weight averaging concept.

Recent research has increasingly explored the concept of weight averaging, a technique where parameters from multiple models are combined to create a unified model that may offer superior predictive performance <cit.>. This approach is grounded in the theory that models starting from the same pre-trained state tend to have a linear loss landscape, making averaging a viable strategy. For a detailed mathematical discussion, see <cit.>. Corroborating this, <ref> features density plots and <ref> that illustrate the predictive entropy for samples that are correctly and incorrectly classified by the  and  datasets using the `soup' models. These results show that the performance of models, post-averaging, is comparable to that of more computationally intensive baseline models. This suggests that models, despite having fewer trainable parameters, are capable of achieving the same level of performance as their more complex counterparts.



r0.3
  

0.26!


2lLayer name     2*Trainable
Params     2*Accuracy 
 (r)1-2
FC     Proj.                    
 0-11      0-11        1.85M     94.89 
 0-11      –          925K     95.11 
 –        0-11        925K     95.83  
 0-5       0-5         925K     92.7 
 6-11      6-11        925K     91.39 
 
Choosing different layers for fine-tuning with Bella, with =4, n=5 on  achieves state-of-the-art performance <cit.>.
 
  


Ablations on fine-tuned layers. 
In our paper, we propose the concept of substituting the network's complete weight matrices with low-rank adapters to approximate BNNs. This section is dedicated to assessing the effectiveness of these low-rank adapters across various layer configurations within CLIP, utilizing .  It should be noted that CLIP consists of 12 ResidualAttentionBlocks, each comprising Fully Connected (FC) and Projection (Proj.) linear layers, which we aim to modify for our . The data, presented in <ref>, reveals interesting performance of our as it surpasses the SOTA in one of the experiments. This highlights the adaptability of our proposed approach, which does not necessitate extensive fine-tuning to select specific layers for optimal performance.


















Comparison of SVGD Bella and VBL. 

In this section, we compare the performance of SVGD Bella and VBL in distinguishing between correctly and misclassified samples. Specifically, we assess their ability to assign high entropy to misclassified samples and low entropy to correctly classified ones. Figure <ref> shows that SVGD Bella outperforms VBL in this regard.






§ IMPACT OF Γ VALUES ON ROBUSTNESS. 





This section assesses how the repulsive force parameter, denoted as γ, influences model robustness. Experiments were carried out on the  dataset using SVGD with a configuration of r=16, n=5, while varying γ values. As illustrated in <ref>, adjustments in the repulsive force significantly affect robustness. Increasing γ enhances robustness by promoting greater diversity among parameter particles. However, prioritizing diversity to an excessive degree (i.e., γ values above 0.05) does not further enhance robustness and may even counteract improvements, although performance remains superior to scenarios with minimal repulsive force (γ=0.01).









§ DETAILED INFORMATION REGARDING VQA TASK


In this section, we provide more details about the experiments we conducted on the VQA task in <ref>.


 §.§ Task Definition and Experiment Setup

Task. 
Visual Question Answering (VQA) <cit.> is a free-form and open-ended task, taking as input an image and a natural-language question about the image and producing a natural-language answer as the output.
Questions in VQA require various intelligence capabilities to answer, including image recognition and object detection, as well as reasoning like commonsense reasoning.
Below is an overview of the VQA dataset utilized in our paper:

    
  * VQA v2 <cit.>: This dataset contains 204,721 images, more than 1 Billion (1B) questions, and 10B ground truth answers in total.
    There are three main types of answers: Yes/No, Number, and Other. The evaluation set contains 80,541, 28,134 and 105,679 questions for Yes/No, Number, and Other respectively.

As open-ended questions may result in a diverse set of possible answers, VQA gather 10 human annotations for each question as the ground-truth answers.
These answers can be different from each other, and even incorrect.



Models. 
In our experiments, we applied on top of LLaVA <cit.> to address the VQA task.
Specifically, we utilized the LLaVA-1.5 7B model, an improved version of the original LLaVA with superficial modifications.
We followed the public training method using a deepspeed codebase. 

However, manipulating parameter gradients is non-trivial, and there is no public way of doing this using deepspeed.
Therefore, we simulated the low-rank approximation from the same initialization by conducting several end-to-end fine-tuning settings with different random seeds, learning rates, and gradient accumulation steps.
Together with the official public model, based on our , there are four generated variants we call -0, -1, -2, -3 respectively.



Metrics.
Below are a detailed metrics used during our evaluations.

    
  * Accuracy: In VQA dataset, there are 10 human annotations for each question. The model prediction accuracy is calculated by
    
    Accuracy(ans) = min{# humans that said ans/3, 1}.

    In order to be consistent with human accuracy, this metric is averaged over all 10 chosen 9 sets of human annotations.
    
    
  * Exact Match: We defined the metric as follows:
    
    EM =
                    1,    if Accuracy=1.
    
                    0,    otherwise.

    When Accuracy equals 1, it means the predicted answer is 100%, same with the ground-truth annotation, i.e., Exact Match with one another.

To measure the model uncertainty along with the human confidence in <ref>, we define Entropy and Human Confidence as follows.

    
  * Entropy: Given a single question, we can calculate the entropy of the model prediction after applying softmax over the logits of LLaVA:
    
    Entropy = 1/N∑_i^N∑_j^V-p_ij×ln(p_ij)

    Note that N is the output sequence length, V is the vocabulary size, and p_ij is the output of softmax function.
    We expect the entropy to be lower for correct predictions, as it stands for lower uncertainty.
    We expect the entropy to be higher for incorrect predictions, i.e., more uncertain on the prediction.
    
    
  * Human Confidence (HC): We calculate the confidence of 10 human annotators as follows:
    
    HC = 1/10max{# humans that said ans}

    More annotators agree with the same answer, higher HC will be.




 §.§ Experimental Results

Correlation. 
Given a question, our model generates a prediction accompanied by logits. We calculate the Entropy using the logits and assess Human Confidence by leveraging 10 ground-truth responses for this prediction. After compiling all predictions, we introduce a specific threshold, segregating them into two distinct categories: those not meeting the threshold are classified under `Low Entropy', while those exceeding it are allocated to the `High Entropy' segment. In <ref>, the threshold values are set within the range {0.1, 0.2, ..., 0.8}. For each threshold level, we compute the average Human Confidence for both `Low Entropy' and `High Entropy' segments, and these averages are plotted as distinct curves.



Similarity.  
Initially, we computed the pairwise distances among the trained models. We employed cosine similarity, which varies between 0 and 1, as our metric. As shown in <ref>, the choice of training strategies influences the degree of similarity among model parameters. Notably, -0 exhibits the greatest dissimilarity when compared to the rest, whereas -2 and -3 display the highest similarity.





Results.
The performance of different single models and ensembles are presented in <ref>, <ref>.
We ensembled two models by simply averaging two softmax output scores.
Given the model similarity matrix in <ref>, some observations and conclusions can be made from these tables.

    
  * Across all models, the averaged entropy of correct predictions is lower than the incorrect predictions. It shows models usually have higher certainty in correct predictions but lack confidence in wrong predictions.
    
  * To some extent, entropy can be used as an indicator of the model's accuracy. For example, models with high accuracy often have lower entropy for correct predictions and higher entropy for incorrect predictions.
    
  * We find that an ensemble of two dissimilar models, e.g., SVGD (-0, -1) (parameters are away from each other), usually ends with higher accuracy and exact match scores, as well as lower entropy over correct predictions. While two similar models ensemble ends with lower accuracy and exact match scores, e.g., Ensemble (-2, -3).
    
  * For other questions in VQA, the EM score of the SVGD is slightly lower than the Single base. As models would generate longer textual answers for those questions, it is non-trivial to investigate a better ensemble strategy. In <ref>, we present some incorrect predictions yet reasonable generation by SVGD , where human annotators show high uncertainty as well.




In addition, we show the evaluation of uncertainty using Mutual Information in <ref> and provide case studies in  <ref> and <ref> to show predicated answers from different models.
The <ref> show the efficacy of . 
It enables large-scale networks like LLaVA to predict outcomes with uncertainty for both VQA `Yes/No' and `Number' questions, a capability that a single base model lacks.
 













§ RBF KERNEL COMPUTATION IN SVGD



We wish to evaluate RBF kernel k(_i_i,_j_j) where
_i∈ℝ^d_1× r,
_i∈ℝ^r× d_2:



    k(_i_i,_j_j)   =exp( -1/2σ^2‖_i_i-_j_j‖_F^2 )
       =exp(-1/2σ^2Tr[(_i_i-_j_j)^⊤ (_i_i-_j_j)]).


 The trace term can be written as

    Tr[_i_i^⊤_i^⊤_i]+Tr[_i_i^⊤_i^⊤_i]-Tr[_j_i^⊤_i^⊤_j]-Tr[_i_j^⊤_j^⊤_i].

To compute _j_i^⊤ and _j^⊤_i the cost is
O(r^2d_2) and O(r^2d_1) respectively. And each trace term is of
the form
Tr(𝐔^⊤𝐕)=⟨𝐔,𝐕⟩_F=vec(𝐔)^⊤vec(𝐕)
which is an O(r^2) operation and requires no further matrix
multiplication.




§ DETAILED PARAMETERS AND DATASETS



In this section, we mention in detail the description for each of the dataset utilized in our experiment as below:


    
  * CIFAR-10: This dataset comprises 60,000 32×32 color images divided into 10 distinct classes, with each class containing 6,000 images. It is partitioned into 50,000 training images and 10,000 test images. The classes cover a range of subjects from animals to vehicles, providing a fundamental challenge in image classification.

    
  * CIFAR-10-C: This dataset is an extension of the  dataset, designed to evaluate the robustness of machine learning models against common image corruptions. It contains the same 60,000 images as CIFAR-10, however, the images in  have been systematically altered using a range of corruption techniques, including noise, blur, weather, and digital effects, resulting in 19 different corruption types each at 5 severity levels. This dataset is used to test the performance of models in recognizing objects under various real-world conditions, making it a valuable tool for improving the reliability and robustness of image recognition systems.

    
  * STL-10: This dataset is a benchmark for evaluating image recognition algorithms, featuring 13,000 color images. This dataset is divided into 5,000 training images and 8,000 test images, distributed across 10 different classes that include a variety of objects such as animals and vehicles. Each image in the dataset is 96×96 pixels, offering higher resolution than many similar datasets such as . The  dataset is tailored for supervised learning tasks in image recognition, providing a structured framework for developing and testing algorithms' ability to classify images into predefined categories.
    
    
  * CIFAR-100: Similar to , the  dataset is composed of 100 classes, each with 600 images, offering a more detailed classification challenge compared to CIFAR-10.
    
    
  * CAMELYON17: The  dataset is utilized in a domain generalization context, where the domains are represented by different hospitals. The primary objective is to develop models capable of generalizing to data from hospitals not included in the training set. Focusing on binary classification, the dataset comprises 96×96 histopathological images as input, with the task to identify the presence of tumor tissue in the central 32×32 region, indicated by a binary label. 
    
    
  * ImageNet (ILSVRC2012): which is a subset of the ImageNet dataset specifically used for 
    the ImageNet Large Scale Visual Recognition Challenge in 2012, contains over 1.2 million images
    distributed across 1,000 different classes.
    
    
  * DomainNet:  is one of the largest and most diverse datasets available for domain adaptation studies. It contains approximately 600,000 images across 345 categories, spanning six distinct visual domains (Real, Clip-Art, Infograph, Paint, Sketch, Quick). This diversity in domains and categories enables the dataset to simulate real-world scenarios where models must adapt to different visual representations and styles.





Hyper-Parameters.


Detailed information about the hyper-parameters used can be found in <ref>.




















 Paper Title: Learned Compression of Point Cloud Geometry and Attributes in a Single Model through Multimodal Rate-Control.


Ground-to-UAV and RIS-assisted UAV-to-Ground Communication Under Channel Aging: Statistical Characterization and Outage Performance
     
	Thanh Luan Nguyen, Georges Kaddoum, Tri Nhu Do, Zygmunt J. Haas,




    
    
	T.-L. Nguyen and G. Kaddoum are with the Department of Electrical Engineering, the École de Technologie Supérieure (ÉTS), Université du Québec, Montréal, QC H3C 1K3, Canada. 
    G. Kaddoum is also with Artificial Intelligence & Cyber Systems Research Center, Lebanese American University (emails: thanh-luan.nguyen.1@ens.etsmtl.ca, georges.kaddoum@etsmtl.ca).

	T. N. Do is with Electrical Engineering Department, Polytechnique Montréal, Montréal, QC, Canada (email: tri-nhu.do@polymtl.ca).
	
	Z. J. Haas is with Department of Computer Science, University of Texas at Dallas, TX 75080, USA, and also with the School of Electrical and Computer Engineering, Cornell University, Ithaca, NY 14853, USA (e-mail: haas@cornell.edu).

    Received August 5, 2024
==============================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================



§ ABSTRACT 
Point cloud compression is essential to experience volumetric multimedia as it drastically reduces the required streaming data rates.
Point attributes, specifically colors, extend the challenge of lossy compression beyond geometric representation to achieving joint reconstruction of texture and geometry.
State-of-the-art methods separate geometry and attributes to compress them individually. 
This comes at a computational cost, requiring an encoder and a decoder for each modality.
Additionally, as attribute compression methods require the same geometry for encoding and decoding, the encoder emulates the decoder-side geometry reconstruction as an input step to project and compress the attributes.


In this work, we propose to learn joint compression of geometry and attributes using a single, adaptive autoencoder model, embedding both modalities into a unified latent space which is then entropy encoded.
Key to the technique is to replace the search for trade-offs between rate, attribute quality and geometry quality, through conditioning the model on the desired qualities of both modalities, bypassing the need for training model  ensembles.
To differentiate important point cloud regions during encoding or to allow view-dependent compression for user-centered streaming, conditioning is pointwise, which allows for local quality and rate variation.
Our evaluation shows comparable performance to state-of-the-art compression methods for geometry and attributes, while reducing complexity compared to related compression methods.









§ INTRODUCTION






Point clouds are a popular representation for immersive multimedia content <cit.>, benefiting from low processing overhead compared to other formats for three-dimensional data such as meshes.
The idea of using point clouds as a data structure for streaming dynamical volumetric multimedia experiences over the Internet leads to significant bandwidth requirements, which establishes the need for efficient compression approaches.

Standardization proposes solutions for point cloud compression through G-PCC <cit.>, which uses octree data structures for the geometry and tailored approaches for the attributes like Region Adaptive Hierarchical Transform (RAHT) <cit.>. When considering point cloud sequences, V-PCC <cit.> utilizes video codecs to compress projections of the point cloud, thus allowing to exploit inter-frame correlations in 2D.



Inspired by and going beyond classical compression methods, Deep Learning has shown promising results for the compression of images <cit.>, videos <cit.> and point clouds <cit.>. 
Similar to the traditional approach, learned point cloud compression mostly treats geometry and attribute compression separately.
A geometry codec handles the compression of the points, while an attribute compression model handles the compression of colors given the knowledge of the undistorted geometry.


When combining the individual geometry and attribute codecs, whether learned or traditional, while relying on lossy compression, the reconstructed geometry at the decoder diverges from the geometry necessary as a prior for attribute decoding. 
This impacts the quality of the attribute reconstruction.
To address this issue, the geometry is decoded and the attributes are projected on the reconstruction <cit.> before attribute encoding, ensuring they are compressed using the same geometric prior which would be available at the receiver side.
We observe, however, that this introduces additional distortions and adds to the computational complexity of the encoding, requiring not only two individual encoding steps but additional geometry decoding and attribute projection, mostly done through costly nearest-neighbor searches, at the sender.

As depicted in Fig. <ref>, we propose a joint variable-rate compression method for point clouds in a single, adaptive autoencoder that embeds both modalities into a unified latent space. Figure <ref> shows that using this joint model, we can individually gauge the quality of the geometry and the attributes.
Using our single model dispenses with training an ensemble of models to fit a fixed number of individual tradeoffs between the bitrate and the qualities of the geometry and attributes.

The contributions of this paper are:

    
  * We formulate the rate-distortion loss for multimodal compression of geometry and attributes to incorporate adaptive weighting of conditional loss components. This allows to train a single adaptive model covering a wide range of geometry and attribute qualities and thus rates.
    
  * We propose a joint compression model, handling the geometry and attributes in one architecture, avoiding the need for geometry decoding and reprojection on the sender side.
    
  * Utilizing conditional encoding, attribute and geometry quality can be modulated independently, allowing variable rate-encoding with a single model, as well as, region-of-interest encoding of point clouds.






§ RELATED WORK


Attribute compression is essential for point cloud multimedia applications in contrast to solely compressing geometry, e.g., in LiDAR applications <cit.>.
Merkuria et al. <cit.> developed a codec utilizing octree partitions and intra-frame prediction while projecting attributes to 2D grids and compressing them with legacy image codecs.
MPEG standards <cit.> distinguish between static point cloud compression using Geometry-based PCC (G-PCC) <cit.> and dynamic point cloud compression with Video-based PCC (V-PCC) <cit.>.
While G-PCC relies on octree partitions for geometry compression and methods like Region Adaptive Hierarchical Transform (RAHT) <cit.> for attributes, V-PCC projects geometry and attribute information into video frames and leverages video codecs for compression. 
Driven by promising results in learned image compression <cit.>, learning based solutions for point cloud geometry or attribute compression showed competitive results in both domains.
However, most of these studies <cit.> do not consider the case where both, geometry and attributes, are subjected to distortions, as it would be the case for a realistic transmission use case. 

We note that early work explored learned compression of attributes and geometry in a single model <cit.>, however, learned compression methods using separate models for both modalities <cit.> were shown to outperform the single model approach.
The latter further boosts performance through wrapping RAHT <cit.> implemented in G-PCC <cit.> in the attribute branch and applying adaptive filters <cit.>, which need to be optimized during each encoding, adding to the already large computational cost.



 §.§ Learned Geometry Compression

Learning-based methods demonstrated competitive performance for point cloud compression, where early work <cit.> proposed to use 3D convolutions on voxelized point clouds and leverage learned entropy models <cit.> to optimize a rate-distortion loss.
The reconstruction is formulated as a binary classification problem using the (weighted) Binary Cross Entropy or the Focal Loss <cit.> to handle class imbalance due to point clouds sparsity in 3D space.
The authors of <cit.> employed entropy models from <cit.>, allowing to model spatial correlations combined with auto-regressive prediction based on already decoded coefficients.
Utilizing the above reviewed contributions, the authors of <cit.> employ a Variational Autoencoder (VAE) architecture based on Voxception-ResNet <cit.>, allowing to show on-par performance with V-PCC <cit.> in terms of geometry compression. 
We note, however, as all previously mentioned methods utilize 3D convolutions on sparse data represented in dense data structures, processing the point cloud in cubes is necessary to handle the memory requirements at the cost of limiting throughput and performance.
To overcome this, sparse convolutions <cit.> have been employed in <cit.> to process large-scale point clouds in one pass and to exploit global context during hyper-encoding, while a multi-scale classification in each stage of the decoder prunes non-occupied voxels.
Instead of using an encoder-decoder approach, <cit.> propose to encode the upsampling for each stage individually while relying on group-based decoding to leverage correlations between upsampled voxels. The model in <cit.> is then extended for inter-coding of dynamic point cloud sequences <cit.>.



 §.§ Learned Attribute Compression

Attribute compression is still dominated by traditional approaches, relying on Graph Transform <cit.> or RAHT <cit.>, with the latter being part of G-PCC.
In the field of learning-based approaches, models using sparse convolutions achieved comparable performance <cit.>. 
Combining traditional methods with learning-based approaches, the authors of <cit.> estimate the probabilities of transform coefficients from RAHT with deep entropy models, further improving compression efficiency.
The attribute compression branch in YOGA <cit.> relies on a thumbnail point cloud compressed at low resolution through RAHT to then predict residuals of the attributes through a conditioned autoencoder similar to <cit.>.  
To further increase the quality of the reconstructed attributes, YOGA <cit.> employs adaptive filters <cit.>, which are optimized during encoding and the parameters are added as side information to the bitstream to be used at the receiver.




 §.§ Variable-Rate Control for Learned Compression

Numerous compression models for images <cit.> and point clouds <cit.> based on the non-linear transform coding framework <cit.> achieve remarkable rate-distortion performance.
A major drawback of these methods is that they are conditioned on a fixed tradeoff between compression and reconstruction at training time.
As a result, an ensemble of models needs to be trained when multiple qualities are required, as it is commonly done in adaptive streaming <cit.>. 
Additionally, this further reduces flexibility as encoding configurations need to be fixed through hyper-parameters at training time. 

Variable-rate control and progressive coding schemes are promising techniques to freely adapt the rate-distortion tradeoff using a single model for encoding and decoding.
For example, for learned image compression, progressive coding is realized by iteratively refining the quantization grids <cit.>. 
Trit-planes <cit.> or dead-zone quantizers <cit.> extend this idea to better fit the distribution of quantization errors.
Following a different approach, quantization or reconstruction results are compressed in cascading <cit.> or recurrent compression <cit.> models, allowing for progressive coding. 
In contrast to these architectural decisions, training with uniformly sampled tradeoff parameters λ while dropping a corresponding number of feature channels in the entropy bottlenecks allows progressive coding through sorting feature channels by importance <cit.> and selecting a number of coefficients to decode.

Compared to progressive methods, adaptive rate-control targets conditioning a single model on a rate variable to allow controlling the rate-distortion of the compressed representation.
Early work such as in <cit.> introduces a gain control layer, multiplying and shifting channels of the feature representation to control the quantization before encoding the transformed coefficients in the entropy bottleneck. 
By incorporating multiple gain control layers in the encoder and decoder, features are scaled hierarchically <cit.>. 
Additionally, invertible layers allow increasing the reconstruction quality <cit.>.
YOGA <cit.> adopts these ideas for variable-rate compression with the layers proposed in <cit.> using sinusoidal embeddings, but requires an ensemble of entropy models as well as selecting appropriate rate parameters for the wrapped color compression using G-PCC.

Finally, we note that works such as <cit.> emphasizes regions of interest in images by conditioning compression models on 2D quality maps instead of a single scalar quality parameter.
As the quality map is not available on the decoder side, it is inferred from the compressed representation for decoding.
Extending this idea from image compression to point cloud compression is promising as it allows view-dependent encoding of three-dimensional content as done with traditional codecs <cit.>.



§ CHALLENGE AND APPROACH





 §.§ Problem Statement



We consider the problem of compressing a point set 𝒫 = (𝐠, 𝐚) containing a set of point locations 𝐠, denoted geometry, and corresponding attributes 𝐚. 
Specifically, for the geometry, quantized coordinates 𝐠={ (x_i,y_i,z_i) : x_i,y_i,z_i ∈ℕ_0,x_i,y_i,z_i < R } with maximum value R∈ℕ are assumed[This allows using sparse data structures and thus memory and computational efficient convolutional neural networks.]. Note that for real-valued coordinates, a voxelized representation could be easily derived by rounding and partitioning the full 3-dimensional space into blocks with size R^3.
We refer to attributes (and features) of a sparse tensor at the geometrical location u ∈𝐠 as 𝐚_u.


We observe that state-of-the-art compression methods <cit.> handle both modalities, i.e., geometry and attributes, separately.
Specifically, methods such as <cit.> or  <cit.> solely compress geometry or attributes, respectively. 
Methods that compress both modalities such as <cit.> still resort to separating the individual compression procedures for geometry and attributes. 
In essence, they encode the geometry first, which is then decoded to project the attributes to the possibly distorted geometry after reconstruction.
Finally, they compress the projected attributes on the distorted geometry. 
On the decoder side, the geometry is decoded and then used as side-information to allow attribute decoding. 
While this process allows compression of point cloud attributes under geometric distortion, it has a number of disadvantages: i) Encoding and decoding requires a sequence of non-parallelizable steps, resulting in high latency and ii) reprojecting the attributes on distorted geometries at encoding introduces additional distortions before compression.

Now, we aim at training a model for joint attribute and geometry compression.
In a Lagrangian fashion, we seek to minimize the rate ℛ under the quality constraints on geometry 𝒟'_G and attributes 𝒟'_A given the  loss function 

    ℒ = ℛ + λ_A𝒟'_A + λ_G𝒟'_G  ,

with corresponding deterministic, positive scalars λ_A,λ_G.
This can also be understood as finding a tradeoff between geometry and attribute quality at a given rate.
Compared to compression models that handle only a single modality in form of ℒ = ℛ + λ𝒟, with fixed hyperparameter λ, the loss function above adds another dimension to the number of models required to train, tune and deploy, making ensembles of joint models inadequate for most use cases. 
Hence, a key to solving this challenge is providing a single, adaptive model. This stands in contrast to the image compression literature <cit.> that coined the rate-distortion tradeoff in a bijection sense between the rate R and the quality D. In this work, this bijection does not exist as the rate results from defining a tradeoff between geometry and attributes.

A seemingly promising approach is variable rate compression <cit.>, in the sense

    ℒ = ℛ + 𝒟_A + 𝒟_G  ,

where we condition a single model for compression at multiple rates, thus only requiring a single training run and reducing the memory requirements at deployment.
In contrast to (<ref>), 𝒟_A,𝒟_G now integrate locality weighting functions  λ_A, λ_G instead of scalar quality parameters λ_A,λ_G allowing to define a location-specific tradeoff between the geometry, attributes and the rate. 
This is especially suited for multimedia point clouds, as it allows to differentiate important regions such as faces, or to offer view-dependent encoding. 
Additionally, instead of searching for tradeoffs between both modalities and the rate in (<ref>) at training time, e.g. through grid search, attribute and geometric quality can be balanced post-training. 



 §.§ Intermezzo: From Non-Linear Transform Coding to Conditional Transform Coding


Next, we briefly review the non-linear transform coding framework from <cit.> in addition to its extension towards conditional variable rate coding <cit.>.
The following review, also depicted as building blocks in Fig. <ref>, considers an arbitrary data vector 𝐱. 
In a later section, we show our building blocks where some are based on this framework. 

Figure <ref> shows the operational diagram of the mean-scale hyperprior <cit.> where an input data vector 𝐱 is non-linearly transformed through an encoder g_a. 
The transform output, i.e. the coefficient vector 𝐲, is quantized as ŷ = Q(𝐲) and fed to the decoder g_s producing the reconstruction 𝐱̂ = g_s(ŷ).
Note that the quantized coefficients are usually entropy coded for transmission while the encoder and decoder non-linear transforms g_a, and g_s are learned. 

Entropy coding the quantized coefficients ŷ requires an estimate of their distribution. 
Given the observation that the coefficients vary locally in mean and scale, a hyperprior model allows to parameterize a local estimate in form of a Gaussian distribution <cit.>. 
A hyperprior 𝐳=h_s(𝐲) is inferred from the coefficients 𝐲 as depicted in  Figure <ref>.
The quantized hyperprior ẑ is then entropy coded using a factorized entropy model <cit.> and used to locally reconstruct the parameters of the Gaussian distributions necessary for encoding/decoding ŷ using the hyperdecoder h_s.
Since backpropagation through a quantization operation results in gradients becoming zero almost everywhere, quantization is replaced through additive uniform noise <cit.> for training, i.e. ỹ =𝒬̃(𝐲) = y + Δ with Δ being an independently and identically distributed uniform random variable, Δ∼𝒰_[-0.5, 0.5]. In the following, we refer to all variables affected by this quantization proxy through (̃·̃)̃.

Extending this approach to allow variable rate coding, the encoder g_a in Fig. <ref> is now conditioned on a variable 𝐜 which is inferred through the conditional encoder p_a from a quality variable 𝐪 <cit.>.
This quality variable allows to locally dictate the quality of corresponding elements in the data vector 𝐱.
As the quality variable is not available for decoding, a proxy quality variable 𝐪̂ is inferred from the hyperprior ẑ.
This allows to derive the condition variable ĉ required in the synthesis transform g_s.





 §.§ Multimodal Conditioning for Geometry and Attributes


Next, we show how to achieve variable rate compression of geometry and attributes in a single adaptive model through integrating the multimodal rate-distortion loss in (<ref>) into the conditioned transform coding (cf. Sec. <ref>). 
To this end, we provide a per-point quality map 𝐪=[𝐪^(G), 𝐪^(A) ] to dictate the target quality individually for both modalities, i.e. geometry and attributes, to the model.
The quality map instantiates the quality variable 𝐪 in Sec. <ref> residing on the object geometry 𝐠. 
During training, each element of the quality map, which is set to 𝐪^(A)_u, 𝐪^(G)_u ∈ [0, 1], with u ∈𝐠, is mapped to a corresponding weighting map element λ^(A)_u, λ^(G)_u through a monotonically increasing function <cit.>.
This allows to individually define the range of the weights in λ^(A),λ^(G) and thus the target distortions for the respective modality. 
We use a quadratic weighting function[We note that exponential functions can be used as in <cit.>, however, experiments show that a root function is more suitable.], i.e. T_A: [0, 1] ↦[λ_min, λ_max] with T_A(x) = (λ^A_max - λ^A_min) x^2 + λ^A_min, such that λ^(A)_u = T_A(𝐪^(A)_u), with u ∈𝐠.
The same mapping holds for λ^(G)_u.
The range [λ^A/G_min, λ^A/G_max] defines the minimum and maximum values in the form of distinct hyperparameters.
At training time, quality maps 𝐪 are either uniformly drawn or supplied as a gradient in a random direction, allowing the model to vary the compression ratio both globally and locally.



  §.§.§ Geometry Loss


For the geometry reconstruction, we use the Focal Loss <cit.>
at every level of the decoder as proposed by <cit.> to allow iterative upsampling of the features and geometry in K steps[corresponding to K resolutions also denoted as scales throughout this work.].
In a nutshell, iterative upsampling in K steps is required for constructing the point cloud geometry from the downsampled representation of the encoder. 
More specifically, in each upsampling step k ∈{0,…,K-1}, a set of possible point locations 𝐝̃^(k), is generated from the lower scale geometry 𝐠̃^(k+1) s.t. 𝐠̃^(k+1)⊂𝐝̃^(k) <cit.>[Later, we omit the scale superscript for k=0, i.e. the original resolution 𝐠̃^(0):=𝐠̃, 𝐠^(0):=𝐠.].
Finally, an occupancy probability 𝐩̃_u^(k) at location u and scale k is inferred for the possible point locations 𝐝̃^(k).
Incorporating the geometry loss map λ_u^(G, k) in each scale k, the distortion loss 𝒟_G in (<ref>) is

    𝒟_G = - ∑_k=0^K-1∑_u ∈𝐝̃^(k)λ_u^(G, k) (1 -𝐩̇_u^(k) ))^γlog(𝐩̇_u^(k)) 
        with 𝐩̇_u^(k) = 𝐩̃_u^(k)   if  u ∈𝐠^(k)
    
            1 - 𝐩̃_u^(k)   else  ,

where 𝐠^(k) is the ground truth geometry  at scale k and the hyperparameter γ known from the focal loss <cit.> reduces the gradients for well-classified samples, allowing the model to emphasize harder samples.
As the loss map λ_u^(G, k) is only given on ground truth voxel geometry, i.e. u ∈𝐠, substitute values λ_u^(G, k) are derived through downsampling and upsampling using average pooling and transposed average pooling, respectively, resulting on the same base geometry as in 𝐝̃^(k). 





  §.§.§ Attribute Loss


The attribute reconstructions is typically cast as an element-wise regression problem <cit.>.
Using the L_2-norm the attribute component 𝒟_A in  (<ref>) becomes

    𝒟_A = ∑_u ∈𝐠∩𝐠̃λ_u^(A)‖𝐚_u - ã_u ‖_2^2   ,

where  𝐚_u denotes the attributes at the geometrical location u ∈𝐠 and ã_u, 𝐠̃ are, as discussed in Sect. <ref>,  affected by the quantization proxy. Recall that the attribute loss map λ_u^(A) contains individual, i.e. local, weighting to gauge the attributes to target distortion. 

Note that the loss can only be computed on locations u ∈𝐠∩𝐠̃, i.e. the intersection of ground truth geometry 𝐠 and the reconstruction 𝐠̃, where we compare the ground truth attributes 𝐚_u with the reconstruction 𝐚̃_u.






  §.§.§ Rate Loss

The rate loss ℛ is computed according to <cit.>, aiming to minimize the cross-entropy as in

    ℛ =    - ∑_u ∈𝐠^(K)∑_i log_2(p_ỹ |  μ̃, σ̃(ỹ_u,i |  μ̃_u,i, σ̃_u,i))
       - ∑_u ∈𝐠^(H)∑_ilog_2(p_𝐳̃ | ϕ(𝐳̃_u,i  | ϕ^(i))),

where Φ={ϕ^(i)}_i are the learnable parameters of the factorized entropy model <cit.>. 
We use ỹ_u,i and 𝐳̃_u,i to denote the transformed coefficient of ỹ and 𝐳̃ at the i^th channel in the feature dimension at the geometry position u, respectively.

In addition, μ̃_u,i, σ̃_u,i denote the local estimates of the means and scales that are obtained from the hyperprior 𝐳̃ as in (μ̃, σ̃) = h_s(𝐳̃).
Hereby, the hyperprior 𝐳̃ capturing the estimates resides on lower scale geometry 𝐠^(H), i.e. H > K.




 §.§ Model Architecture



We propose a sparse neural network architecture as depicted in Fig. <ref>, following the framework proposed by <cit.> (cf. <ref>).
Recall the notation conventions from Sec. <ref>, with (̂·̂)̂ denoting variables affected by quantization (i.e. inference time), and their training time counterparts (quantization proxy) given as (̃·̃)̃. 

The model receives a point cloud 𝒫 = (𝐠, 𝐚).
The encoder g_a downsamples the point cloud 𝒫 to the (resolution) scale k ∈{0,…,K-1}, defined in Sect. <ref>, in combination with the quality map 𝐪 residing on the same geometry 𝐠, to obtain a transformed representation (ĝ^(K), ŷ).

The key to rate adaptive feature transformation at each scale k is the Conditional Feature Extraction (CFE) block (see Fig. <ref>) in the encoder that takes the quality map 𝐪 as input.
On the decoder side, the surrogate[we call estimates at the decoder side surrogate.] quality map 𝐪̂^(𝐊) required for adaptive decoding at the lowest scale[in this work we set the number of resolution scales to K=3.] is inferred from the hyperprior ẑ and the transformed coefficients ŷ. 
The second ingredient to our method is the Upsampling and Prunining (UP) block during decoding (see Fig. <ref>) which prunes the geometry at each scale k to maintain the sparsity of the input geometry, thus overcoming memory restrictions <cit.>.





  §.§.§ Conditional Feature Extraction on Geometry and Attribute Importance
 

We employ local conditioning of feature representations through a scale and shift operation, following common procedure to allow variable-rate control in image <cit.> and point cloud compression <cit.>.
Hereby, features 𝐭^(k) residing on the geometry 𝐠^(k) at a given scale k are transformed through the local scale and shift operation into 𝐭'^(k) as

    𝐭_u'^(k) = α_u^(k)𝐭_u^(k) + β_u^(k)  ∀ u ∈𝐠^(k).

using the scales α_u^(k) and shifts β_u^(k). 
In contrast to the related work <cit.>, which only predicts element-wise scales α_u^(k) and shifts β_u^(k) from a single quality map value 𝐪^(k)_u per geometry location u, joint adaptation of the geometry and attribute quality requires learning a mapping from geometry and attribute quality to the local transformation variables. 

In essence, the feature extraction layer learns the importance of features for optimizing a given task by varying its scale and shift.
This results in broader distributions at the bottleneck when aiming for higher rates, thus reducing quantization errors.
Vice versa, dictating low qualities through 𝐪^(k) results in more narrow distributions or in disabling certain feature channels completely through trivial distributions.
Additionally, the additive shifting operation allows to encode the quality map 𝐪 into the transformed coefficients 𝐲, embedding information required to infer the surrogate quality map from ŷ.



  §.§.§ Upsampling Sparse Features with Geometry Pruning
 

During encoding, lower scale geometry 𝐠^(k+1) is derived from 𝐠^(k) through rounding in a dyadic scheme. 
As a result, when upsampling a representation ĝ^(k+1) in the decoder, the correct voxels to be occupied in 𝐠̂^(k) are unknown.
Hence, to prevent information loss, voxels are generated at each possible location through a Generative Transposed Convolution <cit.>, resulting in a dense geometry 𝐝̂^(k), which is a superset of 𝐠^(k).
While the continuation of decoding with 𝐝̂^(k) is conceptually possible, it leads to significant increase in memory and computational requirements.
Geometry compression methods overcome this challenge by performing multiscale classification <cit.>, pruning the upsampled geometry 𝐝̂^(k) at each scale k by inferring per-point occupancy probabilities 𝐩̂_𝐮^(k) for all u ∈𝐝̂^(k).
Note that attribute decoders, such as <cit.>, do not face this challenge by assuming the availability of undistorted geometry 𝐠^(k) from a preceding step, which is in practice not true at the decoder. 

Now, when combining both modalities, i.e. geometry and attributes, we need to predict the occupancy and prune the feature representation for the next scale.
The challenge hereby is maintaining an expressive feature representation after pruning under potential geometry miss-predictions for attribute reconstruction.
To increase robustness against these miss-predictions, we directly transform the dense, upsampled features, arriving at a dense feature representation, as shown in Fig. <ref>.  
Then, an occupancy prediction branch infers occupancy probabilities 𝐩̂^k_u from the transformed, dense feature representation 𝐟^(k).
Finally, top-k selection is performed, keeping the k most probable voxel locations in a pruning step on the transformed, dense features. 










§ EVALUATION RESULTS






 §.§ Evaluation Setup



  §.§.§ Implementation


We implement our approach in PyTorch, using MinkowskiEngine <cit.> to leverage sparse convolutions. 
For training, we collect a variety of human point clouds from the UVG dataset <cit.>, select a number of frames for each point cloud sequence and crop them into cubes with edge length 128, resulting in 18816 distinct samples.
Opposed to projecting image textures on the samples during training <cit.>, we solely subject the attributes to color jitter and randomly rotate each sample. 
For each sample, a random quality map (𝐠, 𝐪) is derived according to Sec. <ref>. 
We train for 200 epochs using a batch size of 16.
Adam <cit.> is used for optimization with an initial learning rate of 10^-3, which we decrease with factor 0.1 every 50 epochs. 
To stabilize training, gradient norm clipping with threshold 1.0 is integrated.




  §.§.§ Baselines


We compare our method against YOGA <cit.>, which we identify as the current state-of-the-art method for learned compression of geometry and attributes. 
Additionally, we use the tmc2v24 test model for V-PCC <cit.>, using all-intra test conditions and the tmc13v23 implementation of G-PCC <cit.> employing octree compression for geometry and RAHT <cit.> for attributes. For both methods, we select the first 4 rate configurations specified in the respective test models, as we find that they result in a similar range of rates as our model.



  §.§.§ Evaluation Metrics


We measure the bitrate in bits per point (bpp), using the cardinality |𝐠| as divisor.
Geometric quality is assessed using the D1-PSNR <cit.> metric, computing the mean distances between points of the tested geometry 𝐠 and their nearest neighbor in the reference 𝐠̂. 
Utilizing the derived neighbor association, the Y-PSNR and YUV-PSNR metrics <cit.> are computed to assess the attribute quality, with the former only using the luminance channel and the later referring to the weighted average over all channels using weights (6/8, 1/8, 1/8), following evaluation protocols in related literature <cit.>. 
We consistently present symmetric metrics, i.e. computation is performed for swapped source and reference while reporting the stricter result.
Finally, PCQM <cit.> is used as a joint metric, incorporating geometric and attribute features. 
Note that for all presented results, the reported bitrates correspond to the compression of geometry and attributes together.

We select 4 frames each from the 8iVFBv2 <cit.> and MVUB <cit.> dataset for testing, which are specified in the supplementary material. 
If not specified otherwise, uniform quality maps 𝐪^(A) and 𝐪^(G) are used in our model to allow fair comparison.







 §.§ Perceptual Quality Optimized Pareto Front


Following the evaluation protocol specified in <cit.>, we traverse the cross product of quality configurations for our approach with step size 0.05, resulting in 400 configurations for the point clouds in 8iVFBv2 <cit.>. 
Then, the Pareto front for the tradeoff between PCQM <cit.> and rate is determined and depicted in Fig. <ref>. 
Two examples of the chosen configuration pairs can be found in Fig. <ref>, which shows the pairs of attribute quality q^(A) and geometric quality q^(G) contributing to the Pareto front as a path over all possible combinations. 
The underlying contour in each plot shows the bitrate, attribute quality and geometric quality, resulting from all possible combinations.
Similarly, Pareto fronts for V-PCC <cit.> and G-PCC <cit.> are determined by traversing the full range of attribute QPs for each geometry setting.
We also compare against the results reported in the original paper of YOGA <cit.> for the 8iVFBv2 <cit.> dataset.

From Fig. <ref> we find that our model shows improved PCQM performance compared to all baselines for medium to high rates, indicating comparable or better perceived quality of the reconstructions at this compression ratio. 
However, when aiming for lower rates, the quality deteriorates faster compared to related methods.
We attribute this to the fact that our approach is the only method solely relying on learned compression for the point clouds, while, e.g. YOGA <cit.> not only inherits low rate compression capabilities by compressing a thumbnail attribute representation using RAHT <cit.> but also employs adaptive filters <cit.> which are optimized during encoding on the specific content to counteract quality degradation.

We note that the shown results present optimized configurations for the respective point clouds.
This is also apparent when comparing the selected configuration pairs for the point clouds redandblack and loot contributing to the Pareto fronts in Fig. <ref>, which are visualized as paths in Fig. <ref>.
Depending on the individual attribute and geometric properties of the content, we notice that the paths and thus the chosen set of configurations to arrive at the Pareto fronts, differ significantly.
Recall that this evaluation required traversing a large number of possible configurations to select the best tradeoff points to arrive at the pareto fronts and is thus not practical.
Nevertheless, this evaluation provides insights into the capabilities of the tested compression methods.





 §.§ Fixed-Configuration Encoding Performance



Now, we set the conditioning of our model to a fixed set of configurations, i.e. (𝐪^(G)_u, 𝐪^(A)_u) ∈{(0.05, 0.1), (0.1, 0.2),  (0.2, 0.4), (0.4, 0.8)} ∀ u ∈𝐠, independent of the input point cloud.
This allows to compare against G-PCC <cit.> and V-PCC <cit.> using their respective standard configurations. 
We exclude YOGA <cit.>, as such a set of configurations is not publicly available.

Selected rate-distortion curves are presented in Fig. <ref> and corresponding Bjønte­gaard-Deltas to G-PCC <cit.> are shown in Tab. <ref>. 
We provide further rate-distortion curves in the supplementary material.
Note that our approach allows for significantly higher geometric quality according to D1-PSNR, thus small overlap of the resulting rate-distortion curves in the y-axis. This renders the Bjønte­gaard-rate Delta Δ_r for the D1-PSNR metric less expressive on a subset of the tested point clouds. 


Over all tested point clouds, our approach shows significantly higher geometric quality compared to G-PCC <cit.> and V-PCC <cit.>. 
The corresponding attribute rate-distortion curves according to the Y-PSNR metric show notable gains over G-PCC, but underperform when compared against V-PCC. 
Considering attribute quality according to the YUV-PSNR metric, Tab. <ref> shows similar results, with the exception of sequences with complex textures, i.e longdress, redandblack and andrew9.
For the Bjønte­gaard-rate metrics Δ_r and Δ_D of the joint metric PCQM <cit.>, we see high potential for rate-savings in the fixed configuration range compared to G-PCC <cit.>.








 §.§ Latency Analysis


Using a single, joint model in our approach allows to significantly simplify the encoding process, only requiring a single pass through the encoder. 
This stands in contrast to the standard procedure, which requires projection of the attributes to the reconstructed geometry during attribute encoding, as it is done in G-PCC <cit.>, V-PCC <cit.> and YOGA <cit.>.
Consequently, the encoding latency of our model (0.75s) in Table <ref> is more then one order of magnitude lower then the one for the implementations of traditional codecs <cit.>. 
When comparing to the learned approach in YOGA <cit.>, we also notice a lower encoding latency by one order of magnitude compared to the results reported in the paper.
We note that the original paper utilizes different hardware, however, the encoding speed advantage of our approach, due to the single pass, is apparent.

Considering the decoding latency, we find that our method is competitive with respect to the related approaches, although all methods may require further optimization to allow application in practical use cases.
Note, however, that the latency for all approaches significantly depends on the content (i.e. number of points) and the quality configuration. 
A more detailed, per-content and per-configuration analysis is given in the appendix.





 §.§ Showcase: View-dependent Compression


Recall that conditioning our model on a quality map (𝐠, 𝐪) instead of a scalar quality parameter allows to perform region-of-interest rate control.
To showcase this capability, we derive a view-dependent compression of a point cloud given some arbitrary user viewing angle.
This is achieved by instantiating quality maps using a gradient along the viewing direction, assigning higher qualities to the front of the point cloud which decreases to the back.
For comparison, we consistently apply the same procedure using a uniform quality map in Fig. <ref> and a quality map according to the step function in Fig. <ref>.
For all compressed point clouds in Fig. <ref>, the geometry quality map 𝐪^(G) is instantiated following an identical trend. 
The reconstructions are rendered from the given viewing-angle and standard image metrics, i.e. SSIM and the PSNR of the MSE in the YUV color space, are computed.
Considering the image metrics of the renders, we find some potential for reducing the rate by locally dictating the quality 𝐪 in both image metrics, i.e., PSNR and SSIM. 













§ CONCLUSION AND OUTLOOK

We presented a method for training a single, joint compression model for point cloud geometry and attributes.
Key to our approach is extending variable-rate compression techniques to multimodal compression, allowing to gauge the model on a tradeoff between the triplet of the quality of both modalities and the rate.
Equipped with this, we proposed a joint compression architecture operating on a sparse tensor representation. Our technique achieves comparable compression performance to the state of the art in learned point cloud compression <cit.> which, however, has significantly higher encoding latency as it requires geometry decoding at the sender to encode the attributes, while outperforming G-PCC <cit.> and further closing the gap to V-PCC <cit.>. 
Our adaptive, variable-rate model illustrates the local degrees of freedom and flexibility in the joint compression of geometry and attributes. 

However, compressing point clouds according to content specific configurations  allows to fully utilize our model capabilities, but requires computationally expensive grid search in the space of possible configuration pairs.
In some use cases, this may result in resorting to fixed configuration sets as evaluated.
Tackling a similar issue in video compression, per-title encoding <cit.> offers more efficient methods for optimizing a configuration set on specific content.

Deriving meaningful quality maps to locally condition the point cloud geometry and attribute quality using our method is a promising extension of this work. 
Our evaluation gives a perspective  on the capabilities of view-dependent and region-of-interest based compression.
Possible directions are incorporating per-point importance, e.g.  through harnessing fixation maps of the content <cit.> or incorporating task-specific optimization <cit.>. 





unsrtnat
 








§ ARCHITECTURE

We extend the visualization in Fig. <ref>, presenting a more detailed version in Fig. <ref> with implementation channel sizes for reproducability. 
All convolution layers, denoted Conv are three-dimensional sparse convolutions, for which we use Minkowski­Engine <cit.>.
For layers with the extension ↑, Generative Transposed Convolutions <cit.> are used, allowing to generate geometry locations and corresponding features during upsampling.
Note that the hyperdecoder h_s is prone to non-deterministic behaviour <cit.>, which arises when using Minkowski­Engine <cit.>.
To this end, we reduce the kernel size for generative convolutions in the hyperdecoder to 2, as we found that this allows to counteract non-determinism in combination with sorting point locations before and after the convolution.

The geometry representation 𝐠^(3) at the entropy bottleneck is encoded using G-PCC in lossless geometry mode, which is common practice <cit.>. 
The base geometry 𝐠^(5) required to decode the hyperpriors 𝐳 is derived through dyadic downsampling of 𝐠^(3).

In addition, the number of voxels to keep at each pruning stage in the UP block (cf. <ref>) is appended as side-information in the header of the bitstream.






§ ADDITIONS TO THE EVALUATION



 §.§ Test Data

The point cloud frames select for testing are presented in Tab. <ref>, consisting of the 8iVFBv2 <cit.> and MVUB dataset <cit.>.
While point clouds in the 8iVFBv2 dataset <cit.> consist of watertight surfaces and high quality geometry and textures at voxel resolution 1024, point clouds in MVUB dataset <cit.> are frontal scans of human upper bodies containing visible (geometry) noise at voxel resolution 512. 
Incorporating both datasets in the evaluation allowed us to assess the performance of our model over broad input conditions.





 §.§ Extended Results for the Latency

Complementary to the results presented in Table <ref> in Sec. <ref>, we showcase per content results of the encoding latency in Table <ref> and decoding latency in Table <ref>.
Since YOGA <cit.> is not open-sourced at the time of our evaluation, we can only report the averaged latencies from their results, which were measured on different hardware.
For all other approaches, we use the configurations for the generation of the pareto-fronts in Sec. <ref> and measure the times for each encoding and decoding.
We notice that the encoding and decoding latency of G-PCC <cit.> heavily relies on the quality configuration for the geometry and thus split the results for each tested geometry setting. 
It becomes apparent that the latency highly depends on the content, more specifically on the number of points. 

Overall, we observe similar behaviour as presented in Sec. <ref>: Our method allows significantly faster encoding times as other methods, which can be mainly attributed to the fact that the geometry does not need to be reconstructed for attribute re-projection during the encoding process. 
For decoding, we find our method performing comparable to related work. 

 
 

  



 §.§ Complementary Rate Distortion Results for Fixed-Configuration Encoding

We showcase the rate-distortion results on all point clouds according to the evaluation procedure described in Sec. <ref>, as only selected curves were shown in the main document. 
Fig. <ref> shows results for the 8iVFBv2 dataset <cit.> and Fig. <ref> shows results for the MVUB dataset <cit.>.
Overall, we find increased performance for the geometry metric D1-PSNR over V-PCC <cit.> and G-PCC <cit.>. 
For attribute quality our approach achieves a better rate-distortion tradeoff then G-PCC but falls short of the attribute qualities achieved by V-PCC. 
Especially the point cloud andrew9, which shows a checkerboard texture pattern on the shirt, is hard to compress using our learned approach. 
Finally, for the joint metric PCQM <cit.>, we find that our approach consistently achieves higher quality at the given rates compared to G-PCC <cit.>. 
For higher rate configurations, our method can deliver comparable and sometimes better quality as V-PCC <cit.>, but falls of when increasing the compression ratio.

Additionally, we notice that the range of compressed bitrates for our model is small compared to the rate-distortion rates achievable by both traditional methods. 












 Paper Title: DDU-Net: A Domain Decomposition-based CNN for High-Resolution Image Segmentation on Multiple GPUs.



SPECtrophotometer for TRansmission spectroscopy of exoplanets (SPECTR)
    Byeong-Cheol Lee
    August 5, 2024
======================================================================




§ ABSTRACT 
The segmentation of ultra-high resolution images poses challenges such as loss of spatial information or computational inefficiency. In this work, a novel approach that combines encoder-decoder architectures with domain decomposition strategies to address these challenges is proposed. Specifically, a domain decomposition-based U-Net (DDU-Net) architecture is introduced, which partitions input images into non-overlapping patches that can be processed independently on separate devices. A communication network is added to facilitate inter-patch information exchange to enhance the understanding of spatial context. Experimental validation is performed on a synthetic dataset that is designed to measure the effectiveness of the communication network. Then, the performance is tested on the DeepGlobe land cover classification dataset as a real-world benchmark data set. The results demonstrate that the approach, which includes inter-patch communication for images divided into 16×16 non-overlapping subimages, achieves a 2-3 % higher intersection over union (IoU) score compared to the same network without inter-patch communication. The performance of the network which includes communication is equivalent to that of a baseline U-Net trained on the full image, showing that our model provides an effective solution for segmenting ultra-high-resolution images while preserving spatial context. The code is available at https://github.com/corne00/HiRes-Seg-CNN.




Convolutional neural networks, deep learning, ultra-high-resolution images, image processing, parallel processing, semantic segmentation, U-Net, spatial context.









§ INTRODUCTION


The vast majority of deep learning models in computer vision focuses on low-resolution 2D and 3D images, typically 256 × 256 pixels or smaller. However, the increased utilization of high-resolution image datasets introduces new challenges due to the memory constraints of a single GPU, especially for memory-intensive tasks such as semantic segmentation of images <cit.>. Semantic segmentation is the computer vision task of classifying the pixels in the input into distinct, non-overlapping semantic categories. Ultra-high-resolution image segmentation holds significance in diverse fields such as object segmentation in satellite images <cit.>, metallic surface defect detection <cit.>, and computer-aided medical diagnosis <cit.>. Whereas deep convolutional neural networks (CNNs) have achieved remarkable success in image segmentation <cit.>, most of these models are unsuitable for model training and inference for high-resolution images due to their high memory requirements.

To illustrate this, consider computed tomography (CT) scans with sub-millimeter resolution, resulting in voxel image data with a typical resolution of 512×512×512 voxels. Even with half-precision floating-point numbers and a modest batch size of 8, processing such images with a 1-layer convolutional neural network with 64 filters demands over 137GB of GPU/TPU memory, as highlighted in <cit.>. Dealing with such high-resolution inputs using conventional strategies, like down-sampling or patch cropping, generally leads to the loss of detailed information or spatial context, resulting in a lower segmentation accuracy <cit.>. 

In this paper, we propose a novel network architecture based on the well-known U-Net <cit.> developed for biomedical image segmentation tasks. The U-Net is a convolutional neural network (CNN) with a bottleneck structure and skip connections between the encoder and decoder paths. This foundational architecture and variations of it <cit.> have demonstrated remarkable, state-of-the-art accuracy in semantic segmentation and other image-to-image tasks <cit.>. However, its substantial memory requirements are prohibitive for application to high-resolution images on computational devices with limited memory <cit.>. The high memory demands for training stem from the required storage of intermediate feature maps obtained during forward propagation to be used during the subsequent backward propagation pass. This is especially acute when dealing with high-resolution inputs, as it leads to high-dimensional feature maps in the first and last few layers of the network. 

Our proposed model integrates the established U-Net architecture with a divide-and-conquer strategy inspired by domain decomposition methods (DDMs) <cit.> to deal with computational device memory limitations. Previous parallelization strategies employ a decomposition of the image into subimages (subdomains) <cit.>, with communicated margins before each convolution or redundant computations providing global contextual information between subimages.
Our approach similarly decomposes the image. However, communication is explicitly limited to only the bottleneck within the U-Net architecture. This minimizes the communication overhead while preserving essential contextual information.

We summarize our main contributions as follows:

    
  * We propose a novel approach combining U-Net architectures with domain decomposition strategies to segment ultra-high-resolution images efficiently while preserving spatial context. 
    
  * We show that the communication network, which is an important component of our approach, can be employed to exchange information between different subimages. It enhances understanding spatial context without significant computational overhead and with minimal extra communication and memory cost. 
    
  * By evaluating our architecture on a synthetic and a realistic image dataset, we demonstrate competitive segmentation performance compared to baseline U-Net models. Our approach remains scalable, even when training is confined to the largest image portion that the available devices can handle.


The remainder of this paper is organized as follows: In <ref>, we discuss the U-Net architecture in more detail as well as existing memory parallelization strategies. Furthermore, we highlight approaches using ideas from the field of domain decomposition to speed up and parallelize CNNs.  <Ref> introduces our novel parallel architecture and training methodology as well as the considered test datasets. In <ref>, we discuss the influence of this new architecture on memory requirements and the receptive field. Then, we evaluate our proposed network model and the related training strategy using these datasets in <ref>. We provide experimental results for the different datasets in terms of the segmentation accuracy of our approach compared to global semantic segmentation U-Net models. Finally, in <ref>, we conclude and present possible future research directions. 



§ RELATED WORK


In this section, we discuss the U-Net, which serves as the basis for our proposed approach, and the receptive field, a key concept for understanding CNNs. Additionally, we provide an overview of past efforts to tackle the memory challenges associated with the U-Net and other CNN-based semantic segmentation models; we discuss both strengths and drawbacks of these methods. Finally, we introduce the idea of domain decomposition as a natural way to handle memory constraints, highlighting some previous works that adopt this approach for both classification and  segmentation.



 §.§ The U-Net
 
The new segmentation approach introduced in the this paper is based on the U-Net architecture <cit.>. Since its publication, many variants and extensions of the U-Net architecture have been developed, for instance, 3D U-Net <cit.>, UNet++ <cit.>, Attention U-Net <cit.>, and ResUNet-a <cit.>. We assume that the reader is familiar with the fundamental concepts underlying CNNs <cit.> and the building blocks of such models – convolutional and pooling layers. For a detailed explanation of these concepts, we refer to the rich literature on this topic, for instance, <cit.> or <cit.>. 



The U-Net architecture <cit.>, depicted in <ref>, consists of two pathways: the contraction path or encoder path, and the expansive path or decoder path. The contraction path involves repeated blocks, each consisting of two successive 3 × 3 convolutions, followed by a ReLU activation and a max-pooling layer. Conversely, the expansion path employs blocks that up-sample the feature map using 2 × 2 transposed convolutions. Subsequently, the feature maps from the corresponding layer in the contracting path are cropped (if the up-sampling path discards the boundary pixels) and concatenated to the up-sampled feature maps. This concatenation is followed by two successive 3 × 3 convolutions and a ReLU activation. 

In the final stage, the number of feature maps is reduced to the desired number of classes by a 1 × 1 convolution, producing the segmented image. The skip connections between the contraction path, which captures (global) context and features, and the expansive path, which enables precise, fine-grained localization, distinguishes the U-Net from other related network architectures, such as the Fully Convolutional Network (FCN) presented in <cit.>. Note that the originally proposed U-Net returns a cropped segmentation mask (due to the loss of border pixels in every convolution). However, more recent implementations of the U-Net often use zero padding at the borders to obtain output dimensions equal to the input dimensions, such that the input dimensions are preserved and cropping is not required; cf. <cit.>.



 §.§ The Receptive Field

An important concept for understanding CNNs is their receptive field, or field of view. The receptive field of a CNN is the area of the input image that influences the output of a pixel in the CNN output. Unlike fully-connected neural networks, where each input neuron is connected with each output neuron, CNNs have outputs that depend on specific regions of the input layer. In each layer of the CNN, the receptive field expands as we move deeper into the network, allowing higher-level features to be influenced by larger portions of the input image.

The theoretical size of the receptive field can be analyzed both experimentally and theoretically, as has been discussed extensively in <cit.>. The size of the receptive field significantly impacts the predictive performance of the CNN. If the receptive field is too small, the network may not capture all relevant information, leading to suboptimal predictions. For tasks like segmentation, the receptive field should be large enough to include all pixels in the input image that are relevant for predicting the correct pixel class. Note that the optimal size of the receptive field depends on both the network architecture as well as the data on which the network is trained.

In practice, the receptive field size of a CNN is influenced by several factors, including the number of convolutional layers, the size of the convolutional kernels, and the presence of pooling layers; cf. <cit.>. Increasing the depth of the network and the size of the kernels can expand the receptive field, allowing the network to capture more context from the input image. However, this also increases computational complexity and the risk of overfitting.



 §.§ Overcoming U-Net Memory Limits


One issue shared by most U-Net variants is the large memory requirements due to the storage of intermediate feature maps during training, making the model unsuitable for high-resolution applications with limited-memory computing devices <cit.>. A natural approach to overcome these constraints 
is to partition the memory load of the feature maps over different computational devices by using data and/or model parallelism. Several parallelized forms of the U-Net employ parallelism to reduce the memory per device. Both in <cit.> and in <cit.>, the authors implement a spatial partitioning technique that decomposes the input and output of the convolutional layers into smaller, non-overlapping subimages. Before each convolution operation, devices exchange patch margins of the feature maps of half the convolution kernel size with each other. While this approach introduces a communication overhead, it effectively distributes memory across devices without altering the fundamental architecture of the U-Net. Our approach, on the other hand, involves architectural modifications to the U-Net to further optimize performance and limit the communication to the bottleneck of the architecture.

The authors of <cit.> partition the image into overlapping subimages, with the overlap size determined by the receptive field size. Due to redundant computations (𝒪(q4ϵ/N) for an N × N pixel image partitioned into q × q subimages and a U-Net with receptive field size ϵ×ϵ), this approach allows for a fully parallelized execution of both forward and backward passes. <cit.> adopts a similar strategy, but with an application to image classification rather than image segmentation using the ResNet architecture <cit.>.

These approaches have been successful in partitioning the U-Net in such a way that the model can be trained and evaluated in parallel, even for ultra-high-resolution image datasets. However, they either involve communicating margins before each convolutional operation (this is the case for <cit.>), which leads to communication overhead through the many point-to-point messages, or they entail many redundant computations (this is the case for <cit.>). 



 §.§ Memory Optimization for CNNs Using Domain Decomposition Approaches
 
For improving memory efficiency, DDMs are inspiring due to their inherent parallelization and scalability properties, resulting from localization of the computations. DDMs are effective iterative solvers for (discretized) partial differential equations (PDEs), exhibiting scalability through a divide-and-conquer strategy that partitions the computational domain into overlapping or non-overlapping subdomains; see, for instance, <cit.>). The PDE problem is partitioned into subproblems defined on these subdomains. This enables parallel execution of computations within the subdomains. Whereas global convergence is ensured by well-balanced neighbor communication at potentially overlapping subdomain boundaries, and a limited amount of global communication. Highly scalable state-of-the-art DDMs are, for instance, variants of overlapping Schwarz methods <cit.> or balancing domain decomposition by constraints (BDDC) <cit.> and finite element tearing and interconnecting - dual primal (FETI-DP) <cit.>.

Recently, there has been an increasing interest in integrating DDMs with machine learning (ML) algorithms, combining the strengths of both fields <cit.>. Up to now, the majority of works focuses on combining DDMs and ML for solving PDEs. The present paper, however, pursues the combination of CNN with DDM techniques for semantic image segmentation. Notably, to the best of the authors' knowledge, no prior studies have tackled image segmentation tasks using ML explicitly based on domain decomposition strategies, although some methods have been proposed that have similarities with domain decomposition strategies. Existing research mainly focuses on image classification, which is closely related to semantic segmentation; however, due to the different model output, network architectures are somewhat different. Here, we provide a concise overview of methods combining DDM strategies with ML for image classification and semantic image segmentation tasks.



In <cit.>, the authors propose a image partitioning approach for image segmentation, reducing network complexity and enhancing parallelizability. They partition input images into non-overlapping subimages and subsequently train smaller local CNNs on these subimages instead of a large global CNN. This enhances model specialization to specific image regions and reduces overall complexity without a significant impact on test errors compared to a global CNN. Importantly, the authors process the subimages completely independently, with no coupling between them.
In <cit.>, DDM strategies are applied to train a CNN-DNN (convolutional neural network - deep neural network) architecture for image classification. In particular, the authors partition the input image into subimages, each used to train an individual CNN for predicting the class locally. A DNN then aggregates these local predictions into one global prediction. The authors interpret the DNN as a coarse problem solver that combines the finer-grained information from the local networks. In <cit.>, the authors decompose a global CNN for image classification into a finite number of smaller, independent local subnetworks along its width. The weights obtained from training these local subnetworks serve as an initialization for the subsequent training of the global network, employing a transfer learning strategy. In <cit.>, the authors propose a physically-motivated neural network topology for estimating extensive parameters such as energy or entropy. They decompose the domain into subdomains with focus and context regions, overlapping based on the locality of the extensive parameter estimated. The same subnetwork weights are used on the different subdomains (we employ this strategy as well in our approach), leading to a reduction of the number of parameters, enhanced parallelizability, and faster inference. Finally, in <cit.>, a DDM-inspired segmentation algorithm is proposed that divides the input image into several overlapping subimages and trains one segmentation network on these subimages, allowing for parallel inference. The authors conclude that this approach leads to better segmentation accuracy of small objects.




§ METHODOLOGY


In this section, we introduce our domain decomposition-based U-Net (DDU-Net) architecture, along with our training approach. We also describe the datasets used for testing the model and outline the model training procedure.



 §.§ Network architecture
 


  §.§.§ Decomposition of the images and masks

<Ref> shows a schematic representation of our proposed network architecture. Starting with a dataset of (high-resolution) images and corresponding segmentation masks, the model processes a 2D pixel image with dimensions H × W (left) and outputs a probability distribution for K ∈ℕ classes for each pixel of the image (right). 

Following DDM strategies, we decompose the input data into N × M non-overlapping subimages, with N, M ∈ℕ. In particular, each image is divided into N × M subimages with heights H_i and widths W_j, i=1,…, N and j=1,…,M, such that ∑_i=1^N H_i = H and ∑_j=1^M W_i = W. The corresponding segmentation masks are partitioned similarly. The subimages and sub-masks are distributed across the computational devices (e.g., GPUs/TPUs) to balance the workload evenly.

After partitioning, the subimages are processed independently, in parallel, by the encoders on the separate computational devices. For more details on the architectures of encoders and decoders, we refer to <ref>. A chosen number of feature maps from the last layer of the encoder is sent to the communication network (see <ref>), allowing for the exchange of relevant context between subimages. The communication network concatenates the encoded feature maps from all subimages corresponding to the arrangement of the subimages in the original full resolution image and processes them. Each output feature map of the communication network is again partitioned into N × M sub-feature maps, that are then sent back to the devices. The decoders then produce sub-predictions for each subimage, which are concatenated to form a global predicted mask for the entire high-resolution image.



  §.§.§ Encoder-decoder networks
 
<Ref> shows the encoder-decoder architecture used in this work, which closely aligns with the classical U-Net architecture <cit.>; cf. <ref>. The primary difference is a modification in the deepest layer of the encoder path to facilitate communication between the local encoder-decoder clones, which will be discussed below.



As discussed in <ref>, each computational device contains a separate encoder-decoder network, allowing for parallel processing of the subimages. However, these networks share their weights to ensure consistent segmentation across subimages, making them local clones of a global encoder-decoder network. This implies that during training, the weights and gradients need to be synchronized after each backward propagation step.

Each local encoder network processes its corresponding input subimage with several blocks of 3 × 3 convolutions, each followed by a batch normalization layer and a ReLU activation. During training, a dropout layer is also applied to prevent over-fitting. After two convolution blocks, max-pooling is performed to reduce the spatial dimension of the feature maps. The deepest encoder layer produces 256 spatially coarse feature maps (this number can vary with changes in the network architecture), some of which are employed for communication between the encoder-decoder networks. The output of each encoder network are 256 feature maps, with small spatial dimensions. The last F of these 256 feature maps are directed to the communication network, which processes them and returns modified feature maps with the same dimensions that replace the original input feature maps; see <ref>. After communicating, the feature maps then progress through the expansive path, resulting in a final segmented sub-mask for each subimage. These sub-masks are concatenated to form a full mask covering the entire full-resolution image.



  §.§.§ Communication network

The communication network, depicted in <ref>, transfers contextual features, captured by the encoders operating on subimages, between the encoded subimage feature maps. 



For this paper, the communication network consists of three layers of 5 × 5 convolutions. The DDU-Net architecture is not restricted to this specific network configuration. Adjustments such as altering the number of layers, dilation, and kernel size may enhance the model's receptive field size for specific applications. We choose our communication network to be fully convolutional to ensure adaptability to arbitrary input sizes as well as scaling the number of subimages.

The communication network receives as input F feature maps from the deepest encoder layers for each of the N × M subimages; cf. <ref>. These feature maps are concatenated along the height and width dimensions, aligning with the subimages' positions in the full-resolution input image; cf. <Ref>. The communication network processes this concatenated input and produces output feature maps with the same spatial dimensions as the inputs. These feature maps are partitioned back along height and width into N × M sub-feature maps, which are then sent back to the corresponding encoder-decoder network for further processing; they replace the original feature maps generated by the encoder and used as inputs to the communication network.



 §.§ Notation


In the remainder of this paper, we compare various model configurations using a naming convention for our domain decomposition-based U-Net that encapsulates key parameters. Each model is denoted as DFC, where:


    
  * D (∈ℕ_>0) indicates the number of up- and down-sampling blocks in the encoder-decoder network.
    
  * F (∈ℕ_≥0) specifies the number of feature maps processed by the communication network.
    
  * C (Y or N) denotes whether communication between communication feature maps is enabled (C=Y) or disabled (C=N)).

Note that, if communication is disabled (C = N) but the number F of feature maps processed by the communication network is non-zero, this means that F feature maps are sent through the communication network independently, without concatenation. This effectively results in a baseline U-Net architecture with extra convolutions in the bottleneck layer.

For instance, 464Y refers to a model where the encoder-decoder network has a depth of 4, 64 feature maps are sent to the communication network, and communication is enabled. The case D0N operating on 1 × 1 subimages is equivalent to the baseline U-Net architecture. In this case, DFN, with F>0 refers to a baseline U-Net with extra convolutions in the bottleneck layer. Additionally, for the case where F = 0, it does not matter whether communication is enabled (Y) or not (N), as the communication network uses 0 feature maps.



 §.§ Model training

Loss Function. For the network training, we employ the dice loss function, as this function addresses the issue of class imbalance for semantic segmentation of images; cf. <cit.>. The dice loss is defined as

    DL = 1 - 2 ∑_p=1^P ∑_k=1^K y_k,pŷ_k,p + ϵ/∑_p=1^P∑_k=1^K y_k + ∑_p=1^P∑_k=1^K ŷ_k + ϵ.

Here, K represents the number of classes, P is the total number of pixels in a batch, and ŷ_k,p is the predicted probability of pixel p for class k, obtained by applying a softmax function to the model's output so that all the output logits are in the range [0,1]. Furthermore, y_k,p denotes the true probability of pixel p for class k, with values restricted to {0,1} as the true mask is known, and ϵ serves as a small numerical stability constant (to avoid division by zero), set to ϵ=10^-7 in this paper. 

Parameter optimization. The weights of the network are initialized using the method proposed by He et al. <cit.>. We employ the Adam optimizer <cit.>, with momentum parameters β_1 = 0.9, β_2 = 0.999, and a plateau learning rate decay strategy, where the learning rate decays with a factor of 2 when the validation loss does not decrease, maintaining a patience of a chosen number of epochs. The initial learning rate, batch size, and early stopping criterion were chosen depending on the task and available memory and will be provided in the results section.

Training procedure. For training the DDU-Net architecture, we first initialize the encoder-decoder network with shared weights and distribute clones of this network  across all M × N devices. The coarse network for communication is initialized on a selected device and, due to the small size of the communicated feature maps, may optionally reside on the same device as one of the encoder-decoder networks. During forward propagation, each device independently computes the encoded feature maps. After processing a selected number F of those feature maps using the communication network, the decoding of the feature maps again takes place fully in parallel. Backward propagation through the encoders and decoders can, again, be done in parallel without dependencies between the local encoder-decoder networks. However, after backward propagation through the decoders, communication is necessary for backward propagating through the communication network. After backward propagation, gradients are accumulated centrally on a main device and used to update weights, ensuring uniform updates across all devices. The updated weights are broadcasted to the other devices for synchronization.

Implementation. The implementation was done using PyTorch <cit.> (version ), an open-source machine learning library.  We carried out the training and testing on the DelftBlue supercomputer <cit.> at the Delft University of Technology, and we employed NVIDIA Tesla V100S GPUs, with a memory of 32 GB, for the training.



 §.§ Datasets

For testing our model, we use two different image datasets: 1) a synthetically generated dataset designed to test the capabilities of the communication network and 2) a realistic image semantic segmentation dataset for multi-class land cover segmentation to assess the effectiveness of the proposed model both in terms of segmentation quality and memory efficiency. 



  §.§.§ Synthetic dataset


In our approach, only deep feature maps, with low spatial resolution, are exchanged between encoder-decoder networks. This differs significantly from previous U-Net parallelization methods that involve the exchange of a number of feature maps at each U-Net layer. To assess the level of spatial context that these low-resolution feature maps can capture, we designed a synthetic dataset. This dataset comprises one-channel gray-scale images with dimensions (k · 32) × 32 pixels, where k ∈{2, 3, 4, 6, 8, 16}. For each k, we generate 4 000 training, 1 000 testing and 1 000 validation images, resulting in a total synthetic dataset size of 36 000 images. This design allows the decomposition of images into k subimages of 32×32 pixels; see <Ref> for examples with k=4 subimages.



Each image in this gray-scale dataset displays a black background with two randomly placed white circles, each with a 4-pixel radius, placed completely within a subimage such that the subimage boundaries do not intersect the circles. The corresponding mask resembles the image, except for the addition of a third class of pixels: a line segment drawn between the centers of the two circles. Two examples of images and corresponding masks are shown in <Ref>.

When processing an image in this dataset with the DDU-Net, the segmentation of the line segment connecting the two circles relies entirely on the communication network. As the two subimages are processed independently and only linked by the communication network, the effectiveness of the segmentation is a direct reflection of the communication network's capability to transfer global information accurately. 



  §.§.§ DeepGlobe land cover classification dataset
 
The DeepGlobe land cover classification dataset <cit.> is a semantic segmentation dataset for land cover types. The dataset contains 803 high-resolution (2 448 × 2 448 pixels) annotated satellite images with 7 classes: urban, agriculture, rangeland, forest, water, barren, and unknown. The images have 50 cm/pixel resolution and span a total area of 1 716.9 km^2. In addition to the high resolution of the images, segmenting this dataset is challenging due to the large class imbalance; see <ref>. In <ref>, we show two example images and masks from the dataset.



Other challenges include the limited number of images, inexact ground truth, and the presence of multi-scale relevant contexts; cf. <cit.>. To illustrate this, note that trees can exist in various land types such as urban, rangeland, or forest areas. As a result, the network needs to assimilate contextual information from a wider region around the tree to predict the correct class for the tree pixels accurately.





 §.§ Evaluation Metrics

After training the model, we analyze the results. Because of the large class imbalance, we used the (mean) class-wise intersection over union (IoU) score as a metric, as was suggested in the paper introducing the DeepGlobe dataset <cit.>. Given a dataset with n images, the IoU score for class j is defined as: 

    IoU_j = ∑_i=1^n TP_ij/∑_i=1^n TP_ij + ∑_i=1^n FP_ij + ∑_i=1^n FN_ij,

where:

    
  * TP_ij is the number of pixels in the i-th image correctly predicted as class j,
    
  * FP_ij is the number of pixels in the i-th image incorrectly predicted as class j, and
    
  * FN_ij is the number of pixels in the i-th image that belong to class j but were predicted as another class.
 
The mean IoU score (mIoU) provides a single metric by averaging the IoU scores across all classes, and is then defined as:

    mIoU = 1/K∑_j=1^K IoU_j,

where K is the number of classes.




§ ARCHITECTURE DISCUSSION
 

In this section, we analyze some architectural properties of the DDU-Net.  First, we investigate the memory requirements of our approach both experimentally and theoretically and compare the results against a standard U-Net model. Then, we analyze the size of the receptive field of DDU-Net models with different architectures, that is, with varying depth of the subnetworks and the communication network.



 §.§ Memory requirements
 
We compare the memory requirements for the baseline U-Net depicted in <ref> and the proposed DDU-Net model architecture, both with a depth of 4 up-sampling and down-sampling blocks, operating on 2 subimages; cf. <ref>. We specifically present results for a configuration with F=64 communicated feature maps. The channel distribution for the encoder-decoder networks in both models follows the same scheme depicted in <ref>.

<Ref> presents a detailed analysis of the memory requirements for the encoder and decoder of the U-Net architecture, which are identical to the encoder and decoder used in the DDU-Net, during training on a 1 024 × 1 024 image. The table also includes the memory requirements for the proposed communication network with F=64. The values in this table were derived theoretically (see, for instance, <cit.>) and validated experimentally using the  library <cit.>. 

From the analysis, it is evident that storing the feature maps demands significantly more memory than the storing of model weights. The shallow layers, including the input block, the first decoder block, and the last decoder block, collectively contribute to nearly half of the total memory allocation for feature maps. Conversely, the number of weights increases for the deeper layers of the U-Net due to the higher number of kernels and larger kernel sizes in these blocks. The coarse feature maps used in the communicated layer result in relatively low memory requirements for the communication network, demonstrating the efficiency of the proposed DDU-Net in terms of memory utilization.



Figure <ref> shows the peak memory usage during inference for the baseline U-Net architecture and for a GPU containing both the U-Net encoder-decoder and the proposed communication network, for various image resolutions. For smaller image resolutions, the memory allocation can be largely attributed to model weights. However, as image size surpasses a moderate size of 2^7 × 2^7 (128 × 128) pixels, the memory required to store the feature maps begins to dominate the total memory cost. Beyond this point, a noticeable increase in memory requirements is observed, as illustrated in <ref>. 
It is evident from <ref> that the memory requirements for the DDU-Net and U-Net become relatively significantly closer at higher resolutions: for a resolution of 32 pixels, the relative difference is 25.9%, whereas for 4 096 pixels, this difference reduces to only 0.65 %. Additionally, it is important to note that peak memory scales quadratically with resolution; therefore, doubling the resolution results in a fourfold increase in memory consumption.
The results show that the proposed DDU-Net, with its communication network, requires only a small memory overhead compared to the baseline U-Net, despite the added cost of the communication network. This demonstrates the DDU-Net's efficiency in memory utilization: it facilitates communication between sub-images while increasing the memory requirements only slightly.



 §.§ Receptive field analysis
 
The proposed communication network operates on the coarse bottleneck layer of the encoder-decoder networks. Outputs of those layers typically have large receptive fields, as they are the result of numerous convolutional and downsampling operations. Therefore, the communication network is very effective in increasing the receptive field size of the model architecture. In <ref>, we compare the (theoretical) receptive field size of a standard U-Net with that of our proposed model at different depths of the encoder-decoder network, following the theoretical approach as presented in <cit.>. For the DDU-Net, we calculate the theoretical receptive field size for one infinitely large subdomain, providing an upper boundary for the true receptive field size when the network operates on multiple subimages. The comparison shows that the communication network significantly enlarges the receptive field size of the encoder-decoder network. 





§ EXPERIMENTAL RESULTS


In this section, we compare the DDU-Net approach against the corresponding baseline U-Net model based on the segmentation quality on the datasets introduced in <ref>. Moreover, we conduct an ablation study to examine the effectiveness of the communication between encoder-decoder networks processing subimages.





 §.§ Synthetic Dataset Results

In this section, we present the results obtained using our synthetic dataset, which is described in <ref>. The dataset comprises gray-scale images of varying dimensions, specifically (k · 32) × 32 pixels, where k ∈{2, 3, 4, 6, 8, 16}. For each k, we trained separate baseline U-Nets and DDU-Nets.

For our experiments, we trained baseline U-Nets using entire k · 32 × 32 images as inputs. In contrast, the DDU-Net architecture received k subimages of size 32 × 32 pixels, processed by k separate encoder-decoder clones. Both the U-Net and DDU-Net models were designed for 1-channel grayscale input images, generating outputs with three channels corresponding to the logits for the pixel classes: background, line segment, and circle.

We evaluate  

	
  * the impact of varying the number of down- and up-sampling blocks (depth D) in the encoder-decoder networks, 
	
  * the influence of the depth of the communication network, which is equal to the number of feature maps sent to the communication network (denoted as F), and 
	
  * the effect of inter-feature map communication (C, with values Y for enabled and N for disabled).

As mentioned in <ref>, we use the notation DFC to represent a DDU-Net with parameters D ∈ℕ_≥ 0, F ∈ℕ_≥ 0, and C ∈{Y, N}. For a baseline U-Net with depth D we use the notation D. The training hyperparameters used to train the networks are listed in <ref>. The number of weights in the encoder, decoder, and communication network for the different experiments can be found in the appendix, specifically in <ref>.



  §.§.§ Qualitative results

To show the effectiveness of the coarse network with communication, we compare the qualitative segmentation masks for three architectures: the 3, the 332Y and the 332N for the 32 × 64 pixel images (k=2) in <ref>. We observe that the 332N is unable to predict the position of the line segment correctly, as it is unaware of the circle position in the other subimage due to the lack of communication. In contrast, the 332Y predicts the location of the line segment correctly.



<Ref> shows that the proposed encoder-decoder network can capture the relevant features in the F=32 communicated feature maps, and that the communication network is able to transfer this information between the encoder-decoder networks. While we currently communicate all 32 feature maps in the bottleneck of the 332Y for this example, we will examine the impact of varying the number of feature maps in <ref>.  This is particularly noteworthy, as the communication only happens on a very coarse level (with 4×4 pixel feature maps) and the line segments are drawn on the finer grained 64× 32 resolution. Apparently, communication on this coarse level is sufficient for the decoder to produce a good segmentation result on the fine level.



  §.§.§ Varying the number of subimages and communicated feature maps


In <ref>, the IoU score (<ref>) for the line segment class is depicted for the baseline 3 as well as the 3FY, with F∈{0, 1, 2, 4, 8, 16, 32} feature maps communicated, for different numbers of subimages k. Recall that the case F=0 corresponds to a DDU-Net where the communication network has a depth of 0, meaning it is not used, also implying that there is no communication between encoder-decoder clones.

We observe two phenomena. First, there exists a clear positive correlation between the number of communicated feature maps F and the quality of the results. Secondly, the DDU-Net performs even better than the baseline U-Net for larger images and numbers and subimages, respectively. This is related to the sizes of the receptive fields as reported in <ref>: the DDU-Net has a larger receptive field (188 × 188 pixels versus 92 × 92 pixels), yielding better results when the circles are further apart.



  §.§.§ Impact of the depth of the encoder-decoder network

In the DDU-Net, the encoder-decoder network has depth D, that is, the number of up- and down-sampling blocks. The ideal depth is a trade-off: while shallow depths cannot capture all relevant features due to a limited receptive field size (cf. <ref>), overly deep encoder-decoders contain more parameters and may produce too coarse-grained and less informative feature maps for communication.

To see the effect of the depth of the encoder-decoder on the predictions, we compare the segmentation masks generated by a D16Y to the segmentation masks generated by a D, with D ∈{2,3,4}. The resulting predictions of those predictions are shown in <ref>. For encoder-decoder networks with D=2, we observe that the line segment is only predicted correctly for the pixels that are in the center region between the two circles. This can also be explained by the receptive field: shallow networks have a limited receptive field, such that pixels far away from one of the two circles are not affected by this circle; this leads to the incorrect segmentation result.

It can also be seen that the DDU-Net with communication enabled gives better results than the U-Net for D=2 and D=3, as the communication network increases the receptive field size; cf. <ref>. 



  §.§.§ Generalization to different numbers of subimages


The baseline U-Net and the DDU-Net are both size-agnostic, as both are fully convolutional neural networks. This allows them to process images with sizes different from the training images. In case of the DDU-Net, this means that we can vary both the size of subimages and the number of subimages. In order to test the generalization of the DDU-Net with respect to the number of subimages for different amounts of communication, we train for each k∈{2,3,4,6,8,16} a 3FY on input images with spatial dimensions of (k · 32) × 32 pixels, as described in section <ref>, with F ∈{0,1, 2,4,8,16,32}. Each of these models is then evaluated on a dataset with 6 subimages. The results are shown in <ref>.


We can make the following observations: 


  * Models trained on 3, 4, and 6 subimages perform the best on the 6 subimage test dataset. This is likely because the images in these datasets are small enough to be (almost) fully covered by the 188× 188 pixels receptive field of the DDU-Net with D=3; cf. <ref>.

  * The more inaccurate results for the 8 and 16 subimages can be explained by the limited receptive field size, 188× 188 pixels, of the DDU-Net with D = 3; cf. <ref>. 

  * The IoU score deteriorates when the model is trained on only two subimages and then evaluated on a larger number of subimages. This is likely due to the fact that every subimage in the training data set contains one circle and a part of the line segment; however, there are no samples with subimages that do not contain circles or line segments.


These findings suggest that the DDU-Net can be trained on a different number of subimages than on which it is evaluated. 



  §.§.§ Summary of results on the synthetic dataset

We summarize our findings for the synthetic dataset as follows:

    
  * The communication network is able to transfer contextual information across subimages, which becomes clear both from qualitative comparison in <ref> and quantitative results in <ref>.
    
  * The segmentation quality increases when the number of feature maps involved in communication increases.
    
  * The larger receptive field size of the DDU-Net with communication leads to better results compared to the baseline U-Net; cf. <ref>.
    
  * The DDU-Net can be trained successfully on a fixed number of subimages and evaluated on different number of subimages. 




 §.§ DeepGlobe Dataset

Now, we assess the effectiveness of the DDU-Net for a real-world dataset with high-resolution images: the DeepGlobe dataset; cf. <ref>.



Given the small number of 803 images in the DeepGlobe dataset relative to the segmentation task's complexity, we made several adjustments to the encoder-decoder architecture, including batch normalization, random dropout layers, and data augmentation (horizontal and vertical flipping, random rotation of R radians, with R ∈{0, π/2, π, 3π/4}), and we incorporated a pre-trained image encoder model, the ResNet-18 <cit.>. This model, trained on over 1 000 classes on the ImageNet data set <cit.>, offers a rich feature representation for diverse images. The ResNet-18 model consists of residual blocks with skip connections allowing an effective gradient flow. We used the first four residual blocks of the ResNet-18 to initialize our encoder. This strategy leverages the pre-existing knowledge within the pre-trained model to enhance the network's ability to generalize patterns from the limited dataset. A visualization of the employed model architecture is given in <ref>. For a detailed overview of the distribution of model parameters within each component of the network, we refer to <ref>.
 
During training, we kept the ResNet-18 model's weights fixed, only adjusting the weights of the decoder and the communication network. Additionally, we inserted two extra 3×3 convolutional layers in the bottleneck layer. These layers allow the network to restructure and refine the feature maps produced by the ResNet encoder generate relevant information for the communication network.

To ensure a fair comparison between the segmentation quality of the baseline U-Net and the DDU-Net, we trained both models on equally large images. Due to the baseline U-Net architecture's inherent inability to parallelize across different GPUs and the limited training device memory (32 GB), we cropped 1 024 × 1 024 non-overlapping patches from the DeepGlobe dataset, resulting in a training dataset consisting of 2 412 images. It is crucial to distinguish between these large “global patches” and the subimages in the DDU-Net architecture. During training of the DDU-Net, the global 1 024 × 1 024 patches are further partitioned into smaller subimages, which then are distributed across the encoder-decoder networks. Using mixed precision training <cit.>, this approach allowed us to train with a mini-batch size of 12 images on a single GPU. This and other hyperparameters used for training are shown in <ref>. Furthermore, to take into account random initialization of the network parameters, we trained every network repeatedly for three times with the same settings and training dataset. For each configuration, the best performing model (in terms of IoU score on the test dataset) was selected. 


We want to investigate if the DDU-Net can perform equivalently to the baseline U-Net. However, when we include communication in the DDU-Net architecture the additional trainable parameters imply a potential ability to learn more complex patterns. To isolate the effects of (1) the extra parameters in the communication network and (2) the communication itself, we vary the number of feature maps F for both 4FY and 4FN. The 4FN scenario represents an encoder-decoder network with an extra coarse network at the bottleneck layer (with coarse network chosen the same as the communication network), but without information exchange between sub-domains. In this case, the coarse network operates solely on the local bottleneck feature maps, rather than on the concatenated feature maps of all subimages. Notably, F=0 corresponds to the baseline 4. The difference between the cases with communication enabled (Y) and disabled (N) shows the impact of communication across subimages.



  §.§.§ Quantitative results

<Ref> shows the mean IoU scores for different DDU-Net configurations and image partitions, trained on 1 024× 1 024 images and evaluated on the 2 048 × 2 048 test dataset with fixed subimage size, using the approach discussed in <ref>. 

From this figure, two main observations emerge. First, the segmentation quality improves with increased number of feature maps F in the coarse network. This is expected since a higher number of feature maps leads to an increase in the number of parameters and a larger receptive field, enabling the model to capture more complex and distant patterns.

Another trend in <ref> is that the quality of the predictions decreases as the number P× P of subimages increases, or, equivalently, as the subimage size decreases. However, for the 4FY, this decrease in quality is much less pronounced compared to the 4FN. This indicates that the communication combined with the coarse network effectively transfers contextual information between subimages, also for this realistic dataset. 





  §.§.§ Qualitative results

In <ref>, example predictions are shown for different training and model configurations. We observe that the DDU-Net trained on 1× 1 and the DDU-Net trained on 8 × 8 subimages with communication enabled both produce good segmentation results, although there are differences compared to the true mask. Conversely, the DDU-Net trained on 8× 8 subimages without communication shows poorer predictions.To better illustrate prediction errors, we include <ref> in the Appendix, where black indicates correct predictions and white indicates incorrect predictions. 
The key distinction lies in the consistency across neighboring subimages: when communication is enabled, predictions for neighboring subimages exhibit smoother boundaries instead of patchy patterns. This highlights the effectiveness of communication between subimages in the DDU-Net architecture.



§ CONCLUSION


This paper develops a new domain decomposition-based U-Net (DDU-Net) architecture for semantic segmentation tasks. Our results show that by including communication between subimages, the DDU-Net can handle high-resolution image segmentation efficiently without sacrificing accuracy or memory efficiency. Our approach improves segmentation accuracy by leveraging inter-subimage communication. Future research will focus on refining communication strategies, applying DDU-Net to more complex datasets, and evaluating the benefits of parallelization on computing times. Another direction is extending the parallelization strategy developed in this paper to other encoder-decoder architectures.



§ ACKNOWLEDGMENT


The authors acknowledge the use of computational resources of the DelftBlue supercomputer, provided by Delft High Performance Computing Centre (https://www.tudelft.nl/dhpc) <cit.>.
This work performed in part at Sandia National Laboratories was supported by the U.S. Department of Energy, Office of Science, Office of Advanced Scientific Computing Research through the SEA-CROGS project in the MMICCs program.
Sandia National Laboratories is a multimission laboratory managed and operated by National Technology & Engineering Solutions of Sandia, LLC, a wholly owned subsidiary of Honeywell International Inc., for the U.S. Department of Energy’s National Nuclear Security Administration under contract DE-NA0003525. This paper describes objective technical results and analysis. Any subjective views or opinions that might be expressed in the paper do not necessarily represent the views of the U.S. Department of Energy or the United States Government. Sandia report number SAND2024-09671O.



§ NETWORKS USED FOR SYNTHETIC DATASET









The properties of the encoder-decoder and coarse networks used for the synthetic dataset experiments are detailed in <ref>. The properties are shown for different configurations in terms of encoder-decoder depth D. The numbers of weights for the encoder-decoder and coarse network architectures used for the DeepGlobe land cover classification dataset are detailed in <ref>. The encoder-decoder network consists of multiple ResNet18 blocks followed by an inter-convolutional and up-sampling path. We show the number of weights in the coarse network for various choices of its architecture, corresponding to the different experiments as shown, for instance, in <ref>. <Ref> summarizes the hyperparameters employed during training for synthetic data experiments. <Ref> visualizes the difference between predicted and true masks for different configurations of the DDU-Net.



siamplain




 Paper Title: 3D U-KAN Implementation for Multi-modal MRI Brain Tumor Segmentation.












Tang et al.





Department of Biostatistics, School of Global Public Health, New York University
 Equal contributions
* Correspondence: hs120@nyu.edu
 








3D U-KAN Implementation for 
Multi-modal MRI Brain Tumor Segmentation
    Tianze Tang*0009-0007-6244-5813 Yanbing Chen*0009-0009-1218-0129 Hai Shu**0000-0002-6968-4063
    
=================================================================================================




§ ABSTRACT 
We explore the application of U-KAN, a U-Net based network enhanced with Kolmogorov-Arnold Network (KAN) layers, for 3D brain tumor segmentation using multi-modal MRI data. 
We adapt the original 2D U-KAN model to the 3D task,  and introduce a variant called UKAN-SE, which incorporates  Squeeze-and-Excitation modules for global attention. We compare the performance of U-KAN and UKAN-SE against existing methods such as U-Net, Attention U-Net, and Swin UNETR, using the BraTS 2024 dataset. Our results show that U-KAN and UKAN-SE, with approximately 10.6 million parameters, achieve exceptional efficiency, requiring  only about 1/4 of the training time of U-Net and Attention U-Net, and 1/6 that of Swin UNETR, while surpassing these models across most evaluation metrics.
Notably,
UKAN-SE slightly outperforms  U-KAN.







§ INTRODUCTION


Gliomas are a common type of malignant brain tumors and  a leading cause of cancer-related deaths among adults. Diagnosing gliomas is particularly challenging due to their invasive nature and their ability to occur in any part of the brain <cit.>. Accurate detection of these tumors requires multi-protocol magnetic resonance imaging (MRI), which remains the gold standard for post-treatment imaging of gliomas. MRI provides crucial information on tumor size, location, and morphological changes. The most commonly-used MRI modalities include T1-weighted (T1), post-contrast T1-weighted (T1Gd), T2-weighted (T2), and T2 Fluid Attenuated Inversion Recovery (FLAIR) <cit.>.


The Brain Tumor Segmentation (BraTS) Challenge, an annual event hosted by the International Conference on Medical Image Computing and Computer Assisted Intervention (MICCAI) <cit.>, has been held since 2012. 
In the BraTS 2024 Challenge, we participated in Task 1: Segmentation–Adult Glioma Post Treatment <cit.>,  which focuses on evaluating advanced techniques for segmenting brain diffuse gliomas. 
Participants are provided with a  multi-modal MRI dataset to develop automated algorithms to accurately segment the  tumors. This initiative is valuable for creating tools to objectively assess residual tumor volumes, which is crucial for treatment planning and predicting patient outcomes.



 U-Net <cit.> has become 
 the benchmark network
for medical image segmentation, and U-Net based models have successfully demonstrated their efficacy in past BraTS challenges <cit.>. Originating from a straightforward design that features convolutional layers and optionally normalization layers, U-Net has been significantly enhanced through various modifications. For example, Attention U-Net <cit.> incorporates attention blocks to improve focus on relevant features,  while Swin UNETR <cit.> integrates a Swin Transformer as the encoder and connects it to a CNN-based decoder at various resolutions through skip connections.

While previous U-Net based models have made progress in medical image segmentation, they face significant challenges with conventional kernels. These kernels often struggle to capture complex nonlinear patterns and rely on empirical designs that lack interpretability, limiting their clinical trustworthiness. 
Recently, Kolmogorov-Arnold Networks (KANs) <cit.> have emerged as a promising alternative, using learnable activation functions on edges instead of fixed ones on nodes. This approach enhances the network's ability to model complex functions more accurately, improves interpretability, and more effectively links the network's structure to its performance.

Given the potential of KANs, U-KAN <cit.> has been developed by integrating KAN layers into the U-Net architecture. This network design includes a tokenized KAN block that captures complex patterns more effectively. Empirical evaluations of U-KAN have shown its superior performance in medical segmentation tasks and its potential as an advanced alternative in various visual applications.



In this paper, we compare five segmentation models
on the BraTS 2024 Challenge's Task 1,
including 
U-Net <cit.>, Attention U-Net <cit.>, Swin UNETR <cit.>, U-KAN <cit.>, and 
our  UKAN-SE,
a new U-KAN variant with
  Squeeze-and-Excitation
modules <cit.> for global attention. We evaluate these models based on total parameter size, training time, 
and segmentation accuracy. 




§ METHODS





 §.§ Dataset

The dataset for Task 1 of the 2024 BraTS Challenge <cit.> comprises 1,350 labeled training samples and 188 unlabeled validation samples. Each sample includes four modalities (T1, T1Gd, T2, and FLAIR), as illustrated in Figure <ref>. The data is derived from multi-institutional, routine post-treatment clinically-acquired multi-parametric MRI (mpMRI) scans of glioma patients. The MRI scans undergo preprocessing to align them to a standard anatomical template, normalize the resolution to 1  mm^3, and remove the skull. The final dimensions of the processed scans are 218 × 182 × 182.





The ground-truth segmentation labels for the validation set are hidden from the Challenge participating teams, who can only access  the Dice scores and the 95% Hausdorff distance of their segmentation results via the online participation platform. The labels for the MRI scans include non-enhancing tumor core (NETC), surrounding non-enhancing FLAIR hyperintensity (SNFH), enhancing tissue (ET), and resection cavity (RC), as shown in Figure <ref>.





 §.§ Data Augmentation


Data Processing. Our data augmentation pipeline for MRI images incorporates several transformations to enhance the diversity of the training data. We combine the four modalities into a single 4D image (C × H × W × D, where C = 4) and save it, along with the segmentation labels, in the  file format. Each image is stored in 32-bit floating-point format, while the labels are stored as 8-bit unsigned integers. The grayscale values of each image are normalized, ensuring a background value of zero. This process helps reduce the file size.



Data Enhancing. 
The augmentation process begins with a fixed crop of 192 × 160 × 160 to ensure a consistent input size, followed by random flips along different axes to introduce spatial variability. Gaussian noise is added to simulate variations in image quality, and random rotations are applied to account for potential misalignments. Additionally, the contrast of the images is adjusted to vary the intensity distribution. Finally, the images and labels are converted into tensor format, preparing them for further processing in the neural network. This augmentation strategy enhances the model's robustness.





 §.§ Network Architectures




Loss Function.
The loss function 
is given in equation (<ref>),
which combines the soft Dice score and Cross Entropy (CE) loss, with uniform weighting applied across all labels in the CE loss. We experimented with different weights for the CE loss, but no improvement in training performance was observed. The loss is calculated on a batch-by-batch basis. Specifically,

    ℒ_total = (1 - α) ·ℒ_CE + α· (1 - Soft Dice),

where α is a weight factor, and we use 0.5. The softmax output excluding the background class is used as the input for the soft Dice loss, and the one-hot encoded target excluding the background class is used as the target.



Previous U-Net based models.
We consider
three U-Net based models, including the classic U-Net <cit.>, Attention U-Net (Att-Unet) <cit.>, and Swin UNETR <cit.>, 
which have demonstrated strong performance in previous challenges. The U-Net architecture captures features at multiple resolutions through its downscaling and upscaling layers (encoding and decoding), making it particularly suitable for tasks like brain tumor segmentation. Att-Unet  incorporates attention gates, enhancing segmentation accuracy. Swin UNETR combines a U-shaped network design with a Swin Transformer encoder, utilizing multi-scale attention mechanisms to capture both global and local contextual information, making it effective for complex pattern recognition in medical images.



U-KAN.
The U-KAN model <cit.> integrates KAN layers <cit.> into the traditional U-Net structure. In our experiment, we use the default U-KAN configuration, with channel sizes D_1 = 128, D_2 = 160, and D_3 = 256 (see Figure <ref>). The architecture employs a two-phase design: a Convolution Phase for initial feature extraction, followed by a Tokenized KAN (Tok-KAN) Phase where the KAN layers refine the feature representations. This design enables U-KAN to effectively capture complex non-linear patterns, making it particularly suitable for medical image segmentation tasks. The KAN layers are mathematically described as

    KAN(𝐙) = Φ_K-1∘Φ_K-2∘⋯∘Φ_0(𝐙),

where Φ_i 
is the i-th KAN layer
consisting of learnable activation functions, enhancing the model's ability to capture intricate patterns in medical images.

Given the strong performance of U-KAN  in 2D image segmentation, we hypothesize that it can also excel in 3D tasks like brain tumor segmentation. Thus, we adapt it to a 3D version by replacing the 2D convolutional and batch normalization blocks with their 3D counterparts, as shown in Figure <ref>. Notably, the KAN itself does not impose any dimensional constraints, as all input data are transformed into 1D channels, facilitating a straightforward transition to 3D.






UKAN-SE.
We introduce a new U-KAN variant,
called
UKAN-SE,
which incorporates
 a  Squeeze-and-Excitation (SE) module <cit.>  after each convolutional block to provide global attention and enhance feature representations.
Figure <ref> illustrates the network structure.

















 §.§ Evaluation  Metrics


Besides the 
Dice score  
and 95% Hausdorff distance (HD95)
for the whole image segmentation, 
the lesion-wise Dice and lesion-wise HD95 are used in the BraTS 2024 Challenge <cit.> to evaluate the performance of segmentation models. The
two lesion-wise metrics are defined as


    Lesion-wise Dice = ∑_i=1^LDice(l_i)/TP + FN + FP,



    Lesion-wise HD95 = ∑_i=1^LHD95(l_i)/TP + FN + FP,


where L is the number of ground-truth lesions, l_i denotes the i-th lesion, TP is the number of true positives, FN is the number of false negatives, and FP is the number of false positives.

The lesion-wise Dice measures the overlap between predicted and ground-truth segmentations, excluding true negative voxels. The lesion-wise HD95  evaluates the distance between the closest points of the predicted and ground-truth segmentations, emphasizing the most significant segmentation errors. The lesion-wise metrics provide a detailed understanding of a model's ability to accurately segment complex and multi-focal lesions, and help prevent bias toward models that excel only in detecting larger lesions.

The Dice and HD95 metrics are computed 
for each of the six tumor sub-regions:
non-enhancing tumor core (NETC), surrounding non-enhancing FLAIR hyperintensity (SNFH), enhancing tissue (ET), resection cavity (RC),
 tumor core (TC = ET + NETC),
 and  whole tumor (WT = TC +
SNFH).
 
 To compare the efficiency of the five segmentation models, we report the total number of  parameters and calculate the average training time per epoch. 
All models are implemented on 
 a NVIDIA RTX8000 GPU (48GB memory).




 §.§ Training Details


To find the optimal hyperparameters for U-KAN, we aim to achieve the highest lesion-wise Dice and smallest lesion-wise HD95 scores to ensure the best performance. The best-performing set of hyperparameters includes a combination of batch size, epochs, optimizer, weight decay and learning rate.  Due to the large-scale BraTS images, we set the batch size to 2. Model training is conducted for 100 epochs, beginning with a 10-epoch warm-up period. The AdamW optimizer <cit.> is utilized, with a weight decay of 0.0001. The learning rate starts from 0.005, reaches a maximum value of 0.01 after 10 warm-up epochs, and then gradually decreases to 0.005 through cosine annealing <cit.>. For consistency, this set of hyperparameters is also applied to the other four models in our study.












§ RESULTS






The lesion-wise Dice and whole-image Dice results in Table <ref> indicate that UKAN-SE  achieves the best performance across various lesion types, particularly excelling in ET and RC. It also shows consistent high scores in WT segmentation.  Att-Unet  performs well in some areas, notably in SNFH and WT, but does not surpass  UKAN-SE in overall performance.  U-KAN  also demonstrates strong performance, with competitive scores across multiple categories. 


In Table <ref>, which  presents lesion-wise HD95 and whole-image HD95 metrics, the UKAN-SE method continues to lead, showcasing lower values in most lesion types, indicative of better boundary delineation and contour accuracy. This method is particularly strong in ET, RC, and WT segmentation. The U-KAN method also shows commendable performance, particularly in RC and SNFH. 


Table <ref> shows
the parameter size and training time
for each
of the five models.
Notably,
while U-KAN and U-Net have comparable parameter sizes, with 10.612 million and 10.743 million parameters, respectively, U-KAN achieves significantly faster training time, averaging only 803 seconds per epoch  compared to U-Net's 3322 seconds per epoch. Although UKAN-SE has slightly more parameters and requires slightly more training time than U-KAN, it generally performs better in segmentation accuracy 
due to its global attention mechanism. An example of segmentation results is presented in Figure <ref>.


























§ DISCUSSION


This study evaluated the performance of the latest KAN-enhanced U-Net model, U-KAN, on  3D multi-modal MRI brain tumor segmentation using the BraTS 2024 dataset. The results indicate that the U-KAN structure achieves strong performance and exceptional efficiency in the brain tumor segmentation. Compared to traditional models like U-Net and Swin UNETR, the 3D U-KAN model delivers comparable performance  with approximately only 1/4 to 1/6 of the computational cost, highlighting the significant potential of the KAN structure.

In this study, we implemented only simple data augmentation techniques and did not utilize ensemble learning, unlike the more complex approaches described by the past challenge's winner <cit.>. A medium-sized KAN layer configuration ([D_1,D_2,D_3] = [128, 160, 256]) was used. Future research may explore varying these dimensions to assess their impact. Besides, the KAN structure's inherent advantages <cit.> in interpretability suggest promising avenues for further investigation in image segmentation tasks, potentially enhancing the understanding of model predictions.















§ ACKNOWLEDGEMENTS

The data used in this publication were obtained as part of the Challenge project through Synapse ID (syn51156910). This work was partially supported by the NYU IT High Performance Computing resources, services, and staff expertise.

















































































































































 Paper Title: MicroMIL: Graph-based Contextual Multiple Instance Learning for Patient Diagnosis Using Microscopy Images.









MicroMIL: Graph-based Contextual Multiple Instance Learning for Patient Diagnosis Using Microscopy Images
    
JongWoo Kim1Equal contribution Bryan Wong1[1] YoungSin Ko2 MunYong Yi1Corresponding author 1KAIST and 2Seegene Medical Foundation  

1{gsds4885, bryan.wong, munyi}@kaist.ac.kr, 2noteasy@mf.seegene.com  
    
===============================================================================================================================================================================================================




§ ABSTRACT 
Current histopathology research has primarily focused on using whole-slide images (WSIs) produced by scanners with weakly-supervised multiple instance learning (MIL). However, WSIs are costly, memory-intensive, and require extensive analysis time. As an alternative, microscopy-based analysis offers cost and memory efficiency, though microscopy images face issues with unknown absolute positions and redundant images due to multiple captures from the subjective perspectives of pathologists. To this end, we introduce MicroMIL, a weakly-supervised MIL framework specifically built to address these challenges by dynamically clustering images using deep cluster embedding (DCE) and Gumbel Softmax for representative image extraction. Graph edges are then constructed from the upper triangular similarity matrix, with nodes connected to their most similar neighbors, and a graph neural network (GNN) is utilized to capture local and diverse areas of contextual information. Unlike existing graph-based MIL methods designed for WSIs that require absolute positions, MicroMIL efficiently handles the graph edges without this need. Extensive evaluations on real-world colon cancer (Seegene) and public BreakHis datasets demonstrate that MicroMIL outperforms state-of-the-art (SOTA) methods, offering a robust and efficient solution for patient diagnosis using microscopy images. The code is available at <https://anonymous.4open.science/r/MicroMIL-6C7C>




§ INTRODUCTION




Cancer remains a significant global health challenge, prompting extensive research in various fields including medical informatics <cit.>. Timely detection is crucial for improving survival rates, yet conventional cancer diagnosis conducted by pathologists suffers from subjectivity, bias, and high manual labor costs <cit.>. To address these challenges, advancements in deep learning and tissue digital scanners have facilitated the development and analysis of whole-slide images (WSIs), enhancing diagnostic accuracy and efficiency.

While WSIs provide detailed views of human tissue structures and disease-related anomalies, their typical composition of billions of pixels makes direct training of deep neural networks impractical. Consequently, the prevailing methods follow a patch-based paradigm <cit.>. To handle the large number of patches and alleviate the workload on pathologists from labeling each patch, weakly-supervised multiple instance learning (MIL) has emerged as a viable solution, requiring only WSI-level labels without fine-grained (patch-level) annotations <cit.>.

In histopathology, understanding tissue context through patterns is crucial <cit.>, yet existing MIL models <cit.> struggle to analyze this effectively using only individual patches. Graph-based MIL approaches address this by modeling the relationships between patches, enabling a more comprehensive understanding of tissue anomalies and their spatial context. Models such as WSI-HGNN <cit.> and H2-MIL <cit.> enhance the exchange of contextual information across regions by considering relationships among multi-resolution patches. Additionally, GDS-MIL <cit.> introduces a novel approach by integrating graph attention layers (GAT) to infuse spatial contextual information into MIL models, further improving the accuracy and robustness of cancer detection.

Recently, microscopy-based analysis has gained attention for its cost and memory efficiency compared to WSIs produced by scanners <cit.>. Despite these advances, capturing contextual information in microscope images remains a critical challenge. Existing graph-based MIL models designed for WSIs <cit.> require absolute positions to construct graph edges, which are not available in microscope datasets. As shown in Figure <ref> (Left), tiling patches from WSIs using a scanner follows a window-sliding approach, allowing access to absolute positions. In contrast, Figure <ref> (Right) illustrates that microscopy images are acquired from the subjective perspective of pathologists, resulting in unknown absolute positions and redundant patches due to multiple captures. This lack of absolute positions hinders the use of existing graph-based MIL models, preventing effective analysis of relationships between microscope images within a patient. Moreover, redundancy in images is problematic because creating edges based on similarity significantly increases connections between duplicate images, reducing connectivity with images from diverse areas. This could slow down overall information exchange and ultimately hampers contextual understanding.

Given the unique characteristics of microscopy images, such as unknown absolute positions that hinder graph edge construction and the presence of redundant images, we propose a weakly-supervised MIL framework to alleviate the workload of pathologists. Our approach extracts representative images and incorporates their relative positions to capture contextual information, thereby improving the prediction of patient diagnoses. To this end, we introduce Graph-based Multiple Instance Learning for Patient
Diagnosis Using Microscopy Images(MicroMIL). Specifically, we dynamically cluster microscopy images using deep cluster embedding (DCE) <cit.> and employ Gumbel Softmax <cit.> to extract a representative image from each cluster. Subsequently, we utilize graph neural networks (GNNs) to connect representative images based on their similarity and model the relationships between images. By integrating local contextual information within clusters and combining it with contextual information between clusters, MicroMIL effectively captures and leverages the spatial and contextual relationships present in microscopy images.

We extensively evaluate our method on two microscopy datasets: a real-world colon cancer dataset (Seegene), which contains a larger number of microscopy images and exhibits a more severe redundancy issue compared to the public dataset, and the public BreakHis dataset <cit.> for breast cancer. Our method outperforms various state-of-the-art (SOTA) methods used for WSI scanners, thanks to the DCE that effectively addresses the redundancy issue and the GNNs that extract contextual information based on representative image similarity. Our results demonstrate that MicroMIL provides a robust and efficient solution for patient diagnosis using microscopy images, overcoming key limitations of existing methods designed for scanner-derived WSIs.



§ RELATED WORK




 §.§ MIL for WSIs (Scanner)


Weakly-supervised MIL has significantly impacted the histopathology field by utilizing WSI labels without requiring individual patch annotations. Initial methods used Mean or Max Pooling for feature aggregation, but these have evolved into more sophisticated techniques. ABMIL <cit.> calculates attention scores for each patch to enhance the aggregation process. Building on this, MSDAMIL <cit.> combines multiple-instance, domain adversarial, and multi-scale learning frameworks to improve robustness. DSMIL <cit.> further refines MIL models with self-supervised contrastive learning and proposes considering the distance between each patch and the critical patch. Transformer-based approaches like TransMIL <cit.> and attention modules in MIMIL <cit.> have effectively captured patch relationships. DTFD-MIL <cit.> leverages Grad-CAM <cit.> to derive patch probabilities, offering more precise predictions and better interpretability. A recent method, HAGMIL <cit.>, incorporates multiple resolutions of the WSIs to further improve analysis.

Graph-based MIL methods have recently been proposed to extract contextual information between patches. WSI-HGNN <cit.> and H2-MIL <cit.> actively model the exchange and integration of contextual information across regions by considering the relationships among multi-resolution patches. GDS-MIL <cit.> proposes a novel approach by integrating graph attention layers (GAT) to infuse spatial contextual information into MIL models. While these methods effectively address the unique characteristics of WSIs, they do not account for the specific characteristics of microscopy images: unknown absolute positions (making graph-based MIL models designed for WSIs inapplicable) and redundant images.








 §.§ Patient Diagnosis Using Microscopy Images


Previous research often relies on simple statistical techniques for patient-level diagnosis. Nguyen <cit.> employ ensemble methods for colorectal tissue image classification, utilizing soft voting on VGG16 and CapsNet outcomes. Gandomkar <cit.> propose an ensemble approach using meta-decision trees (MDT) for patient classification while Han <cit.> leverage GoogleNet and majority voting for breast tissue histopathology image classification.

MIL has also been applied to microscopy images. Pal <cit.> utilize sparse attention for instance aggregation in a deep MIL framework to analyze cervical histopathological images more cost-effectively using microscopy images. Recently, Yang <cit.> introduce the hierarchical aggregation network for MIL, evaluating its performance on protein subcellular localization using immunofluorescence images and gene annotation with spatial gene expression images.

Despite these advancements, existing studies often rely on simple machine learning and statistical methods, which exhibit limitations in both accuracy and efficiency. Furthermore, current MIL models are not specifically designed to address the unique characteristics of H&E microscopy images. Our study emphasizes the importance of eliminating redundant images and considering the contextual information between representative images to achieve accurate patient-level predictions.



 §.§ Microscope Datasets


In the field of microscopy image analysis, several benchmark datasets are commonly used for research. The ICIAR2018 dataset[https://iciar2018-challenge.grand-challenge.org/Dataset/] contains 1376 breast cancer samples at magnifications from x20 to x200. Additionally, datasets from the Netherlands Cancer Institute (NKI) and Vancouver General Hospital (VGH) offer 248 and 328 breast cancer samples at x200 magnification <cit.>. These datasets are complemented by others like BreakHist <cit.> for breast cancer and MHIST <cit.> for various cancers.

While these research-focused datasets are invaluable, they offer differ significantly from real-world clinical data, which typically includes more redundant images per patient (Table <ref>). Benchmark public datasets are curated to remove unnecessary data, which is not reflective of the complexity, encountered in clinical settings. To enhance the practicality and robustness of our approach, we utilize a real-world microscopy dataset from the Seegene Medical Foundation (details available in the Supplementary Information) and the public BreakHis dataset <cit.> for breast cancer. This study aims to predict patient diagnoses using microscope images, relying solely on patient labels without image-level annotations.



§ PROPOSED METHOD


We present MicroMIL, a weakly-supervised MIL framework adapted to the unique characteristics of microscopy images (Figure <ref>). It consists of three main components: 1) extracting features using a pre-trained feature extractor model, 2) removing redundancy through deep cluster embedding (DCE) and then extracting representative images with Gumbel Softmax, and 3) constructing a graph neural network (GNN) based on representative patch similarity to capture both local (within-cluster) and regional (between-cluster) contextual information. The details of the process can be found in Algorithm <ref>.






 §.§ Feature Extractor


MIL can be categorized into two types: instance-based and embedding-based MIL <cit.>. Embedding-based MIL, which provides richer representation, has become the preferred research focus due to its effectiveness in classifying WSIs <cit.>. Following the common practice in embedding-based MIL, we extract microscopy images I_s using a pre-trained feature extractor E (Equation <ref>) where S_p represents the number of images in patient p.


    f_s^(p) = E(I_s^(p)),    s ∈{1, …, S_p}   (1)




 §.§ Deep Cluster Embedding


We utilize deep cluster embedding (DCE) <cit.> as a key component of our methodology to remove redundant images, which occur frequently in microscope datasets, as shown in Table <ref>.

Let F_p = {f_1^(p), f_2^(p), …, f_S_p^(p)} denote the set of features for patient p obtained from the feature extraction process. Additionally, let C_p = {c_1^(p), c_2^(p), …, c_K^(p)} be the set of image cluster assignments for patient p, where K denotes the number of clusters.

The deep cluster function can be formulated as follows:


    DCE(F_p) = min_C∑_i=1^S_p∑_k=1^Kf_i^(p) - μ_k^(p)^2 ·α (c_i^(p) = k)


Here, μ_k^(p) represents the centroid of cluster k and α is the indicator function that equals 1 if the condition inside the parentheses is true and 0 otherwise.

The deep cluster function aims to minimize the distances between data points and their respective cluster centroids while simultaneously updating the cluster assignments to ensure compact and well-separated clusters in the feature space. This process iteratively refines the embeddings and cluster assignments until convergence, resulting in meaningful representations of the data that are suitable for subsequent classification tasks.



 §.§ Representative Images


Through the aforementioned method, we perform image clustering to understand the relationships between images within the same area. To eliminate redundancy, we then select the most significant image within each area. This process reduces interactions within the same area (cluster) while facilitating information exchange between different areas (between clusters). Initially, we use Equation <ref> to calculate the weights between images within the clusters. Subsequently, we employ the Gumbel Softmax <cit.>, which enables backpropagation to select the most influential representative image.

The weight t_k^(p) for cluster k of patient p is calculated as follows:


    t_k^(p) = W_1 · (DCE(F^(p))_k · F^(p)),    k ∈{1,2,…,K}


where W_1 denotes the weight of self-attention, where W_1 ∈ℝ^d × 1., DCE(F^(p))_k represents the features of cluster k, and F^(p) is the set of features for patient p.

The Gumbel Softmax function is then used to calculate the attention weight αk, i^(p) for each image i in cluster k of patient p:


    α_k,i^(p) = exp(t_k,i^(p) + g_i/τ)/∑_j=1^N_k^(p)exp(t_k,j^(p) + g_j/τ)


where g_i is the noise sampled from a Gumbel distribution and τ is the temperature parameter.

The representative feature q_k^(p) for cluster k of patient p is then computed as follows:


    q_k^(p) = ∑_i=1^N_k^(p)α_k,i^(p)· (DCE(F^(p))_k · F^(p))


Let Q^(p) = {q_1^(p), q_2^(p), …, q_K^(p)} represent the set of representative features of patient p in cluster k, where K denotes the total number of clusters. N_k^(w) is the number of patches included in the k-th cluster of patient p



 §.§ Graph-based Aggregation


Patches derived from WSI via scanners have defined absolute positions, facilitating the identification of neighboring patches. However, microscopy images lack these absolute coordinates, complicating the task of neighbor identification. To address this challenge, we determine the neighbors of representative images through relative positioning and maximize interactions using graph-based aggregation.



  §.§.§ Constructing Graph


The graph is constructed by defining similarities between representative images, represented as nodes. The similarity in Equation <ref> computes a similarity matrix A_i,j^(p) based on image features. To form the graph, we use the upper triangular part of the similarity matrix, ensuring that each node connects to its most similar neighbors using the Gumbel Softmax function.


    A_i,j^(p) = similarity(q_i^(p), q_j^(p))


Here, A_i,j^(p) denotes the similarity between image features q_i^(w) and q_j^(w) for patient p. To construct the graph 𝒢^(p), we create an edge matrix ℰ^(p) using the upper triangular part of A^(p) and the Gumbel Softmax function, as defined in Equation <ref>.


    ℰ_i,j^(p) = 
    1,    if  A_i,j^(p) = max_k ∈𝒩(i)(σ(A_k,j^(p))), 
    
    0,    otherwise


ℰ_i,j^(p) represents the presence (1) or absence (0) of an edge between node i and node j in the graph for patient p, with i < j ensuring that only the upper triangular part is considered. The function σ is the Gumbel Softmax function.

The graph 𝒢^(p) is then formed by connecting nodes based on the edge matrix ℰ^(p):


    𝒢^(p) = (𝒱^(p), ℰ^(p))


where 𝒱^(p) is the set of nodes representing the representative images (𝒱^(p) = Q^(p)), and ℰ^(p) is the set of edges as defined by the edge matrix ℰ^(p).




  §.§.§ Graph Neural Network


In the graph 𝒢^(p), nodes represent representative images and edges represent their relationships. The graph attention network (GAT) effectively models interactions between nodes by assigning different weights to adjacent nodes. This can be expressed as follows:


    h_i^(l+1)(p) = σ(∑_j ∈𝒩(i)α_i,j^(l)(p) W^(l) h_j^(l)(p))


h_i^(l)(p) is the feature representation of node i for patient p at layer l. The set 𝒩(i) represents the adjacent nodes to node i in graph 𝒢^(p) while α_i,j^(l)(p) indicates the attention coefficient between nodes i and j at layer l for patient p. The matrix W^(l) is the weight matrix at layer l and σ denotes the activation function, such as ReLU <cit.>.

The attention coefficients α_i,j^(l)(p) are computed as follows:


    α_i,j^(l)(p) = exp(σ(a^(l)^T[W^(l) h_i^(l)(p)∥ W^(l) h_j^(l)(p)]))/∑_k ∈𝒩(i)exp(σ(a^(l)^T[W^(l) h_i^(l)(p)∥ W^(l) h_k^(l)(p)]))


Here, a^(l) is a learnable weight vector at layer l, [·∥·] denotes concatenation, and σ represents the Leaky ReLU activation function.

The representations of nodes are aggregated to obtain a WSI representation h̅^(p) by averaging all node representations:


    h̅^(p) = 1/|𝒱^(p)|∑_i ∈𝒱^(p) h_i^(L)(p)


This aggregated representation is then fed into a multi-layer perceptron (MLP) to predict the target variable y^(p):


    ŷ^(p) = MLP(h̅^(p))


Where ŷ^(p) is the predicted value of the target variable for patient p.




 §.§ Training and Optimization


The entire framework, including the feature extractor, clustering, representative image extraction, and GNN is trained end-to-end. The loss function used for optimization is the cross-entropy loss, defined as follows:


    ℒ = -1/P∑_p=1^P[ y^(p)log(ŷ^(p)) + (1 - y^(p)) log(1 - ŷ^(p)) ]


P is the total number of patients in the training set, y^(p) is the ground-truth label for patient p, and ŷ^(p) is the predicted probability for patient p.







§ EXPERIMENTS




 §.§ Datasets


Prior studies have primarily used datasets collected in controlled environments as benchmark datasets. However, real-world microscopic datasets contain significantly more images (Table <ref>) and exhibit a higher number of redundant images (Table <ref>). To demonstrate the practicality and robustness of our approach, we utilize a real-world microscope dataset for colon cancer from the Seegene Medical Foundation (details available in the Supplementary Information) and the public BreakHis dataset <cit.> for breast cancer. The goal of this study is to predict patient diagnosis based on microscope images for each patient. It is important to note that only patient labels are used for both datasets, without image-level labels.






 §.§ Baselines


We compared MicroMIL against SOTA MIL models used for WSI from scanners <cit.> in the weakly-supervised setting (only patient labels are used). Additional details regarding the baselines can be found in the Supplementary Information. The reported result of DTFD-MIL is based on the AFS variant. Importantly, direct comparison with existing graph-based MIL models <cit.> is not feasible because these models require absolute positional information to construct graph edges, which is not available in the microscope datasets.



 §.§ Implementation Details


To ensure fair comparison, we used ResNet18 pre-trained on ImageNet as the feature extractor for both MicroMIL and the baselines. We set the hidden dimension across all MIL models to 128 to maintain fairness. For other hyperparameters, we followed the default settings from the original papers to achieve their best performance across all datasets. For MicroMIL, we experimented with a dropout rate of 0.5 and learning rates within the range of {1, 3, 5}× 10^-3. We utilized the Adam optimizer <cit.> for model optimization and applied an early stopping mechanism with a patience of 5 epochs. The implementations were made using the PyTorch framework and trained on an NVIDIA GeForce RTX 3080 GPU.



 §.§ Metrics


We evaluated the performance using three metrics: Macro-F1, Micro-F1, and area under the curve (AUC). The Macro-F1 score balances precision and recall, making it suitable for imbalanced datasets. Micro-F1 measures overall classification correctness, though it may be affected by class imbalances. AUC assesses the model's ability to distinguish between positive and negative cases, with higher values indicating better discrimination.



§ RESULTS


The primary goal of these experiments is to address the following research questions:



  * RQ1: Does the proposed MIL method outperform existing SOTA MIL models used for WSI from scanners?

  * RQ2: Do the extracted representative images in diverse areas enhance contextual information?

  * RQ3: Does connecting each node to its most similar negihbors during edge generation in GNN enhance contextual information?









 §.§ Baselines Comparison (RQ1)


Table <ref> shows that MicroMIL consistently surpasses baseline MIL models across all evaluated metrics on both real-world and public datasets. This performance advantage stems from MicroMIL's specific design to address two key characteristics of microscope datasets: image redundancy and missing absolute positions. In contrast, existing MIL models are designed for WSIs obtained from scanners, which do not account for these characteristics. By effectively tackling the unique characteristics present, MicroMIL proves to be exceptionally well-suited for patient diagnosis using microscopy images.







 §.§ Representative Images in Diverse Areas (RQ2)


To validate the hypothesis that extracting representative images from diverse areas enhances contextual information, we employ K-means clustering to extract all images in one patient based on the number of clusters. Figure <ref> illustrates that when K=1 (all data points are assigned to a single cluster), the connections (red line) based on the maximum similarity are between images in diverse areas. However, as K increases, the connectivity becomes overly localized. The trend is further shown in Table <ref> which the highest performance of the evaluated metrics is achieved when K=1. Following this observation, it is crucial to extract representative images from diverse areas to obtain more effective contextual information. 








 §.§ Similarity-based Edges in GNN (RQ3)


Contrary to WSIs, where absolute positions are available, microscopy images lack this spatial information, posing a challenge for creating edges. To address this, we introduce a GNN-based approach that connects each node to its most similar neighbors during edge generation using Gumbel Softmax. As shown in Table <ref>, the similarity-based edges outperform random and dissimilarity-based methods across all evaluated metrics in both datasets. This demonstrates that using similarity-based edges effectively enhances contextual information.




§ ABLATION STUDY




 §.§ Without vs. With Representative Images


The performance without using representative images is lower compared to with using representative images in both datasets, as shown in Table <ref>. This is because the model struggles to convey contextual information between diverse areas effectively. By selecting only the representative image for each area, the model can capture better contextual details, leading to improved interaction between diverse areas and thus enhancing overall performance. BreakHis shows a smaller difference between scenarios without and with representative images compared to the Seegene  dataset (0.023 vs. 0.0472) as the number of redundant images is not as severe in BreakHis (Table <ref>). Conversely, in the Seegene  dataset, performance is significantly improved when representative images are used as they capture essential contextual details among a high number of redundant images.





 §.§ Representative Image Methods




This study compares the performance of the Mean, Centroid, and Gumbel Softmax (MicroMIL) methods for selecting representative images. The Mean method averages all the image features in each cluster. The Centroid method selects the centroid of each cluster optimized by the DCE model. Lastly, the Gumbel Softmax method selects the image with the highest attention weight within the cluster. As shown in Table <ref>, performance is lowest with the Mean method, followed by the Centroid, and highest with Gumbel Softmax Method. This indicates that methods which consider the importance of individual images, like Gumbel Softmax, better capture the essential features of each cluster, leading to improved performance.





 §.§ Representative Image Clusters




Table <ref> presents the effect of varying the number of representative image clusters in both datasets. For the Seegene  dataset, with an average of 150 redundant images per patient, the best performance is achieved with 36 clusters. In contrast, BreakHis, which has fewer redundant images, achives optimal peformance with 16 clusters, suggesting a need for more representative images as the number of redundant images increases.





 §.§ Graph Layers Effect


The effects of graph edge generation under different number of graph layers is shown in Table <ref>. The result indicate that there is no difference in performance between Random Shuffling and Similarity (MicroMIL) when no layers are used, likely due to the lack of interaction with other images. Additionally, when many layers are used, the performance difference is negligible as it incorporates not only adjacent representative images but also most of the images. With a small number of layers, consideration is limited to adjacent representative images, leading to relatively high performance especially when using the Similarity method for edge generation.





§ CONCLUSION


We present MicroMIL, a weakly-supervised MIL framework that specifically designed for the unique characteristics of microscopy image analysis, particularly addressing the issues of unknown absolute positions and redundant images. By employing deep cluster embedding (DCE) to dynamically cluster images, using Gumbel Softmax to select representative images, and integrating a graph neural network (GNN) to model relationships between images, MicroMIL effectively captures both local and global contextual information. Validated on both real-world colon cancer data and the public BreakHis microscope dataset, MicroMIL outperforms existing MIL models. Our results demonstrate that MicroMIL significantly improves patient diagnostic accuracy and efficiency using microscopy images, setting a new standard in the field.



§ SUPPLEMENTARY MATERIAL




 §.§ A. Private (Real-World) Dataset Description

We utilized a dataset obtained from the Seegene-Medical Foundation, which was approved by both the foundation's institutional review board (SMF-IRB-2020-007)) and the collaborating partner university's institutional review board (KAIST-IRB-23-214). This dataset comprises colorectal tissue images of 899 patients treated for colorectal cancer between 2021 and 2023. It was split by imaging dates, with post-November 2023 data designated as the test set. The training set consists of 272 malignant and 330 normal cases, while the test set comprises 134 malignant and 123 normal cases. Overall, the dataset includes 82,761 malignant and 52,339 normal images.

The data was collected by Seegene-Medical Foundation, which has one of the largest diagnosis centers in the country. After collection, the tissue samples are fixed in 10% formalin and delivered to the pathology lab within 12 to 24 hours. The tissues are embedded in paraffin blocks, sectioned at approximately 3 micrometers (µm) by clinical pathologists, and mounted on glass slides. Subsequently, the paraffin is removed, and the slides are stained with hematoxylin-eosin (H&E) before being covered with a cover glass for examination by the pathologists primarily at 40x and 100x magnifications.

For this study, an Augmentiqs Optical Module was attached to the microscope used by a pathologist. Cases selected by the pathologist for review were reassessed, and the device was used to automatically capture images by changing the magnification and field of view of the slide on the microscope equipped with this module for image generation purposes.




 §.§ B. Baseline code

Table <ref> shows the urls of the authors' implementations of the compared methods.




ieee_fullname






 Paper Title: Deepfake Media Forensics: State of the Art and Challenges Ahead.

































Deepfake Media Forensics: State of the Art and Challenges Ahead
    Irene Amerini^5, Mauro Barni^8, Sebastiano Battiato^1, Paolo Bestagini^6, Giulia Boato^9, Tania Sari Bonaventura^5,
Vittoria Bruni^5, Roberto Caldelli^2, Francesco De Natale^9, Rocco De Nicola^10, Luca Guarnera^1, Sara Mandelli^6,
Gian Luca Marcialis^4, Marco Micheletto^4, Andrea Montibeller^9, Giulia Orrù^4, Alessandro Ortis^1,
Pericle Perazzo^3, Davide Salvi^6, Stefano Tubaro^6, Claudia Melis Tonti^5, Massimo Villari^7, Domenico Vitulano^5
^1 University of Catania, ^2 CNIT, Florence, and Universitas Mercatorum, ^3 University of Pisa, ^4 University of Cagliari, 

^5 Sapienza University of Rome, ^6 Politecnico di Milano, ^7 University of Messina,
^8 Università di Siena, 

^9 University of Trento, ^10 Scuola IMT Alti Studi Lucca



{sebastiano.battiato, luca.guarnera, alessandro.ortis}@unict.it, roberto.caldelli@unifi.it, pericle.perazzo@unipi.it, 

{marco.micheletto, giulia.orru, marcialis}@unica.it, {irene.amerini, domenico.vitulano, vittoria.bruni}@uniroma1.it, 

{davide.salvi, sara.mandelli, paolo.bestagini, stefano.tubaro}@polimi.it, mvillari@unime.it, mauro.barni@unisi.it,

 {giulia.boato, andrea.montibeller,francesco.denatale}@unitn.it, rocco.denicola@imtlucca.it,  {bonaventura, melistonti}@diag.uniroma1.it



    August 5, 2024
=======================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================



§ ABSTRACT 

AI-generated synthetic media, also called Deepfakes, have significantly influenced so many domains, from entertainment to cybersecurity. Generative Adversarial Networks (GANs) and Diffusion Models (DMs) are the main frameworks used to create Deepfakes, producing highly realistic yet fabricated content. While these technologies open up new creative possibilities, they also bring substantial ethical and security risks due to their potential misuse. The rise of such advanced media has led to the development of a cognitive bias known as Impostor Bias, where individuals doubt the authenticity of multimedia due to the awareness of AI's capabilities. As a result, Deepfake detection has become a vital area of research, focusing on identifying subtle inconsistencies and artifacts with machine learning techniques, especially Convolutional Neural Networks (CNNs). Research in forensic Deepfake technology encompasses five main areas: detection, attribution and recognition, passive authentication, detection in realistic scenarios, and active authentication. Each area tackles specific challenges, from tracing the origins of synthetic media and examining its inherent characteristics for authenticity. 
This paper reviews the primary algorithms that address these challenges, examining their advantages, limitations, and future prospects.





Multimedia Forensics, Deepfakes




§ INTRODUCTION

The advent of Deepfakes, synthetic media generated by Artificial Intelligence (AI) that mimics real images, audio, and video, has significantly impacted various domains including entertainment, politics, and cybersecurity. Deepfakes leverage deep learning techniques, particularly GANs <cit.> and DMs <cit.>, to create highly convincing but falsified representations of individuals. While these technologies offer creative opportunities, they also pose serious ethical and security challenges due to their potential for misuse.
The emergence of such advanced AI-generated media has led to the development of a cognitive bias known as the Impostor Bias <cit.>, which refers to the tendency to doubt the veracity of multimedia elements due to the knowledge that they can be realistically generated by AI models.
Deepfake detection has become an essential field of research, aiming to develop methods to distinguish between real and artificially generated media. Techniques for Deepfake detection often involve analyzing inconsistencies and artifacts that are not easily perceptible to the human eye <cit.> but can be detected using proper detectors based on machine learning algorithms. These detection methods typically focus on both spatial and temporal anomalies in the data, utilizing Convolutional Neural Networks (CNNs) <cit.> for enhanced accuracy. 
Starting with the Deepfake detection task, the scientific community has over the years taken on several other new challenges to study the nature of synthetic data in detail. We can therefore distinguish 5 main areas of research in the Forensic Deepfake domain, namely Deepfake Detection (Section <ref>) Deepfake Attribution And Recognition (Section <ref>), Passive Deepfake Authentication Methods (Section <ref>), Deepfakes Detection Method On Realistic Scenarios (Section <ref>), and Active Authentication (Section <ref>).





In this context, authors of the proposed papers are involved in the FF4ALL initiative (FF4ALL - Detection of Deep Fake Media and Life-Long Media Authentication), 
which aims to develop theoretical and practical tools for detecting and combating media counterfeits or Deepfakes, tracing their origin and limiting their dissemination. 

In the following sections, a brief overview of the main algorithms that aim to address the above-mentioned challenges will be presented. 



§ DEEPFAKE DETECTION































Deepfake technology poses significant challenges due to its potential for misuse, which can severely impact public well-being and trust. While current detection methods, primarily based on convolutional neural networks and deep learning paradigms, have shown promising results, they often struggle to generalize across the varied techniques employed in digital content manipulation. This issue primarily arises from the intricate interplay between textures and artifacts in Deepfake data, which traditional detection methods frequently overlook.
In this context, artifacts are unintentional distortions or irregularities that occur during the Deepfake generation process, including unusual pixel formations or edge anomalies. Conversely, textures refer to the inherent patterns and fine details present in authentic images, such as the natural appearance of skin and hair.
In fact, it is well established that synthetic manipulations typically disrupt the texture consistency of original images <cit.> and often leave detectable traces in the form of artifacts in both spatial <cit.> and frequency domains <cit.>, particularly in specific facial regions <cit.>. Consequently, numerous studies focus their analysis on specific portions of face images to identify these inconsistencies.
One promising approach involves using both No-Reference (NR) and Full-Reference (FR) quality measures to detect subtle manipulations in video frames <cit.>. This method has significantly improved cross-manipulation generalization by focusing on areas susceptible to artifacts, such as the mouth and eyes, and analyzing the image quality degradation caused by Deepfake algorithms.
In addition to artifacts, texture analysis provides another robust basis for distinguishing between real and fake images. In some Deepfake technologies dedicated to face-swapping operations, the inner and outer faces have different identities, making texture inconsistencies particularly evident <cit.>.
However, focusing exclusively on either artifacts or textures in Deepfake detection can be limiting. While these approaches yield high accuracy in specific contexts, they often fail to adapt to new and evolving Deepfake techniques. To address this limitation, a novel framework called the Texture and Artifact Detector (TAD) has been proposed <cit.>. The TAD framework enhances Deepfake detection by leveraging both texture and artifact inconsistencies, thereby improving model generalization across various forgery scenarios through ensemble learning.
Unfortunately, the performance of these methods is often hindered by the challenges posed by highly compressed data. High compression ratios can obscure subtle manipulations, leading to a significant degradation in detection accuracy <cit.>.
A promising solution involves leveraging a learnable adaptive high-frequency enhancement framework to enrich weak high-frequency details in compressed content, thereby enhancing the robustness of Deepfake detection under compression <cit.>. Further details on compression impacts and related detection strategies will be discussed in subsequent sections.










§ DEEPFAKE ATTRIBUTION AND RECOGNITION
































 §.§ Deepfake Fingerprint and Attribution


Deepfake attribution, often referred to as Deepfake Model Recognition <cit.>, encompasses methodologies capable of identifying the specific model used to generate synthetic data. This process includes attempts to estimate the unique model weights <cit.> of the architecture instance responsible for creating the Deepfake. SOTA techniques are highly effective in detecting Deepfake content generated by widely-used GANs <cit.>  and DMs <cit.>. These techniques can even specialize to recognize the specific architectures, and, in more details, the specific model used in the creation procedure. Then, a more advanced challenge in this domain is identifying the exact model instance, characterized by a unique set of weights and parameters, within a given architecture: Guarnera et al. <cit.> demonstrated that using a simple ResNET-18 <cit.> engine combined with a metric learning approach <cit.>, excellent results can be achieved in identifying the specific model used for creating synthetic data from 100 different instances of StyleGAN2-ADA <cit.>. Moreover, the same authors obtained good results in a preliminary generalisation test conducted on architectures other than StyleGAN2-ADA. A robust model recognition solution would enable the attribution of an image to a specific model owner, which is crucial for intellectual property rights <cit.>. To establish the ownership or authenticity of an image generated by a particular model within a specific architecture, new strategies and appropriate metrics are required <cit.>.
In the context of forensic investigations involving Deepfake images, videos, or audio, state-of-the-art Deepfake detectors and architecture classifiers can be likened to the task of identifying camera models in traditional forensic analysis. Just as camera source identification is used to trace the origin of a photograph, Deepfake model recognition aims to trace the origin of a Deepfake to a specific model instance within an architecture. This parallel underscores the necessity of developing advanced techniques for Deepfake model attribution to ensure authenticity in digital media.








§ PASSIVE DEEPFAKE AUTHENTICATION METHODS





















In the modern era, where video calls have become a cornerstone of global communication, the importance of authenticating audio and video streams cannot be overstated. The advent of Deepfake technology poses a significant challenge to the integrity of digital communication.
Traditional Deepfake detection methods may fall short as they often focus on either audio or video data in isolation. However, Deepfakes may involve sophisticated manipulations of both audio and video streams, making them harder to detect with monomodal methods.

This highlights the need for a multimodal approach that simultaneously analyzes both audio and visual data <cit.>. By correlating information from these two channels, we can significantly improve the accuracy of Deepfake detection. This approach takes advantage of the fact that inconsistencies are often more noticeable when multiple modalities of data are considered together. For instance, <cit.> leverages the incongruity between emotional cues portrayed by audio and visual modalities. In <cit.>, the authenticity of a speaker is verified by detecting anomalous correspondences between his/her facial movements and what he/she says. Moreover, the results of <cit.> show that an ensemble of audio and visual baselines outperforms monomodal counterparts.

Given that audio-visual authentication methods may exploit monomodal detectors in a synergistic fashion, another step towards better performance is to enhance monomodal audio or visual detectors separately through the use of advanced techniques. Concerning the visual component, it is possible to leverage fusion of multiple detectors trained on purpose to capture orthogonal traces <cit.>. Concerning audio, it is possible to exploit modern solutions such as transformers <cit.>, as well as investigating the use of semantic traces <cit.>.

Despite the great effort of the multimedia forensics community, a series of challenges remains. Concerning multimodal solutions, the need for audio-video Deepfake datasets is becoming more than an urgent necessity. Indeed, most of the effort has been put towards monomodal datasets creation. Moreover, given the trend of large language models, it could be interesting to try using the same logic for audio visual reasoning. Concerning monomodal solutions, explainability has definitely not been reach yet, which still proves a problem in case of court of laws.















§ DEEPFAKES DETECTION METHOD ON REALISTIC SCENARIOS























 §.§ Deepfake Detection of image-videos in the wild 
 



In recent years, there has been a growing interest in the study of techniques for the detection of Deepfake and AI-generated media <cit.>. Consequently, numerous solutions have been proposed to address the problems posed by the increasing spread of fake multimedia content. However, most of these solutions perform well only in controlled settings, such as laboratory experiments, but fail to provide reliable results in real-life conditions typical of practical applications.
Deep Learning (DL) models can effectively detect Deepfake media and identify their source. However, despite promising results, DL-based methods face several significant challenges, particularly in real-world applications where controlled laboratory conditions are absent. Firstly, DL models require vast amounts of labeled data for training, which is often difficult to obtain in real-life scenarios. Additionally, these models must handle unforeseen situations that were not accounted for during training, a common issue in multimedia forensics applied outside the lab. For DL models to be effective in the ongoing battle between forensic analysts and counterfeiters, it is crucial to address the risk of overfitting to training data, which can lead to failures in new, unexpected situations. While it is possible to train highly accurate detectors, these methods struggle to generalize to new generative techniques due to data drift <cit.>. Detectors perform well on the techniques they were trained on but often fail with content from new generative models. Another major obstacle is the black-box nature of DL techniques, making it difficult to interpret analysis results and understand decision-making processes. This lack of transparency hampers the practical application of DL in scenarios where accountability is essential. To address these issues, researchers address their attention on several strategies, including the use of one-class classifiers trained only on pristine images <cit.>, developing classifiers with rejection options to opt out when encountering unfamiliar inputs not well-represented in the training set <cit.> and adopting methods capable of generalization though features fusion <cit.> together with multimodal approaches that combine audio and video streams  <cit.>. 
Even considering these efforts, the practical application of automatic detectors has been minimal. Deploying these tools in commercial or mass verification systems presents numerous challenges beyond generalizing from a few known benchmarks <cit.>. One significant challenge is the need to continuously train these models on new generative/Deepfake techniques in a continual learning fashion <cit.>.
Continual learning, also known as lifelong or incremental learning, is an ongoing approach to maintaining good model performance on evolving tasks without experiencing catastrophic forgetting <cit.>. This approach is well-suited to recognize content generated by new techniques and continuously adjusting models to account for data drift observed during inference versus training.
A promising direction is the creation of an end-to-end Deepfake detection system that supports continuous integration and continuous delivery/deployment (CI/CD) with the design of a Machine Learning Model Operations (MLOps <cit.>) pipeline, enabling the end-to-end development of continuously trained and monitored intelligent detectors with a minimal set of components.



 §.§ Deepfake and Social Media
 


An extremely challenging “real-world" case scenario involves the detection of Deepfake multimedia shared on social networks <cit.>. In fact, to cope with bandwidth and storage limitations <cit.>, social networks apply severe data compression and resizing. However, such processing, while reducing the overall multimedia size, also reduces the presence of forensic features used for the discrimination of real vs. fake multimedia <cit.>.

A first study on GAN images <cit.> shared on Twitter demonstrates the adversarial effects of social network compression on Deepfake detectors. Specifically, while the visual quality of the shared images was untouched, the presence of forensic traces and patterns was reduced. The effects of social network sharing were then extensively observed and studied in <cit.>.
In <cit.>, the authors explore the challenges and advancements in the field of media forensics as applied to social network. The study addresses the increasing concerns regarding the authenticity and reliability of digital media shared on social platforms, focusing on challenges affecting source attribution algorithms <cit.> as well as multimedia verification <cit.>. For the latter, additional effort was devoted to studying whether the multimedia content is consistent with its descriptive text. The paper <cit.> discusses the emerging challenges in the field, such as the rise of Deepfakes and the use of bots for spreading disinformation. The authors highlight the need for advanced forensic tools to keep up with these sophisticated methods of media manipulation.
In <cit.>, the authors composed a large and diverse dataset counting 80k fake images generated with StyleGAN models and 70k real images collected from several state-of-the-art datasets <cit.>. In addition to this, while revealing insightful details on the entity and severity of social network compression applied by Twitter, Facebook and Telegram, the authors shown how their dataset can be used to finetuning new detectors, preserving their accuracy on social network compressed images while without experiencing “catastrophic forgetting loss" <cit.>.
Interestingly, <cit.> showed that while social networks degrade the presence of forensic artifacts used by real vs. fake detectors, they introduce other traces that can be exploited to reconstruct the life cycle of multimedia and determine on which social networks it has been shared. While the life cycle of multimedia does not provide specific information on its nature (i.e., real or fake), it can be fundamental in recovering the version of that multimedia closer to the original, unshared one. This, in turn, allows for more accurate real vs. fake detection.
Finally, preliminary studies have been conducted also on videos shared on social netowrks <cit.>, showing similar effects to those on images. One of these works <cit.> studies the effects of social network compression on FaceForensics <cit.> videos shared on Facebook and Youtube. The study provides a results in line with what observed on images <cit.> and a new dataset of shared videos to be used to finetune real vs. fake detectors.

Nevertheless, while new works on Deepfake detection on social networks are available, the continuous update of social network compression algorithms makes the arms race increasingly challenging. As a consequence, this require additional effort to develop new architectures and updated datasets of social network shared images.




 §.§ Detection of Deepfake Images and Videos in Adversarial Setting


An additional problem affecting virtually all the Deepfake forensic techniques developed so far is that such techniques are thought to operate in a benign setting, that is, by neglecting the possible efforts made by an adversary to mislead the forensic analysis. Yet, recent researches <cit.> have shown how easy is to generate adversarial contents capable of deceiving image and video processing techniques based on DL, when the adversary is informed about the details of the tools employed by the analyst. Some works <cit.> have also studied the transferability of adversarial examples to networks different than those targeted by the attack, opening the way to the development of powerful attacks even when the attacker is unaware, or only partially aware, of the techniques used by the analyst.
Even worse, often it is not even necessary that the adversary applies sophisticated attacks relying on the full or partial knowledge of the to-be-attacked system. By relying on the lack of robustness and the generalization capabilities of the forensic tools already outlined in Section <ref>, the adversary may simply process the deepkfake content in such way to prevent a correct forensic analysis, or at least degrade its performance up to a point to make it unusable. Some examples of this kind of attacks, often referred to as laundering attacks, include the application of moderate to strong lossy compression, geometric processing of images and videos, noise addition, histogram stretching and many others. 

Understanding and ensuring the security of Deepfake forensic tools is a crucial problem, if such tools have to be used under the intrinsically adversarial conditions typical of multimedia forensics applications.

For this reason, several efforts have been made to defend against adversarial attacks <cit.>, both in the realm of computer vision applications and multimedia forensics. Still, no general effective solutions have been found yet <cit.>. Among the solutions developed so far, adversarial training <cit.> has received some consensus and has proven to at least mitigate the effectiveness of adversarial attacks in computer vision applications. As argued in <cit.>, adversarial training forces DL models to focus on robust, possibly semantic, features, which are inherently more difficult to attack. Whether such a beneficial effect of adversarial training also applies to Deepfake forensic applications is still an open problem. It is no clear, in fact, if in multimedia forensics the equivalent of semantic computer vision features exist or not.

With regard to laundering attacks, the solutions proposed so far are similar to those already discussed in Section <ref>, given the, ultimately, the effectiveness of laundering attacks can be drastically reduced by improving the robustness and generalization capabilities of the forensic tools. A common approach to do so, involves the use of data augmentation techniques that enrich the training set with processed samples, thus improving the robustness of the forensic tools against the processing operators included in the data augmentation procedure. Yet, accounting for all possible kinds of processing during training is clearly unfeasible. Among the solutions proposed so far, we mention the possibility of identifying a kind of worst possible laundering attack and include it in the training procedure. Examples of such an approach are described in <cit.>. Despite all the efforts made, even for laundering attacks, a definitive solution has not been devised yet, thus adding yet another point to the to-do list of multimedia forensic researchers. 









§ ACTIVE AUTHENTICATION
























 §.§ Active deepefake detection



Passive Deepfake detection techniques <cit.> work a posteriori, after that the forged content has been generated, distributed and possibly processed, on the contrary, active methods work in a preemptive way, pre-processing the media in such a way to ease the subsequent analysis. This is the case, for instance, of Deepfake detection methods based on DNN watermarking, whereby the content generated by DNNs is watermarked in such a way to ease the distinction between genuine and fake media, and the attribution of the fake content to the network which generated it. An alternative possibility is to modify the computational imaging chain characterizing modern acquisition devices, to insert within the generated content a unique fingerprint to be used later on for authentication purposes. This means a change of paradigm that needs to be properly explored.  Active authentication techniques represent a valid, and more reliable, alternative (or complement) to passive authentication, whenever the operating conditions allow their use.

Watermarking has recently been proposed as a means to protect the IPR of DNNs <cit.>. By tying a watermark to a DNN model, in fact, it would be possible to prove the ownership of the model or trace its illegal use. On this basis, DNN watermarking can be also used to link AI-generated contents <cit.>, like Deepfakes, to the model which generated them, thus providing an easy and convincing  way to distinguish between synthetic and natural contents. Such a goal is achieved by requiring that all the contents generated by a network contain a predefined watermark (a kind of synthetic fingerprint) that can be used later on to distinguish the synthetic images (or videos) generated by a trained model from real ones. This marks a drastic paradigm change with respect to current solutions based on multimedia forensics, since authentication is now achieved with the active help of the party which trained the media-generation network. Though some solutions have appeared in this direction, putting this idea at work requires that considerable advances are made particularly in terms of watermark robustness against image, video and network manipulations, security against adversarial attacks, payload and also imperceptibility. Succeeding into designing a new class of robust and secure solutions, based on active approaches, for Deepfake detection/attribution surely represents an open challenge and an interesting opportunity for scientific research in the field of multimedia forensics.





 §.§ Efficient Media Origin Authentication



Customary deep-fake detection methods, both passive and active, are subject to false positives and false negatives, whose rate highly depends on the employed method and the goodness of the training data. False positives are due to various factors, such as the complexity of the content, the quality of the training data, or the intrinsic limitations of the detection algorithm itself. False negatives could happen if the deep-fake is very well made, or in general if the detection method fails to recognize certain patterns or features that indicate a deep-fake.

On the other hand, cryptographic signatures are “almost perfect” from this point of view, in the sense that false negatives (i.e., authentic signatures which are not recognized so) are zero and false positives (i.e., fake signatures that are taken as authentic) are considered computationally infeasible to forge. This suggests that cryptographic signatures could be used fruitfully to detect deep-fakes (or better, the absence of deep-fakes) with perfect precision.
In this direction, the work-in-progress standard JPEG Trust <cit.> by the Joint Photographic Experts Group (JPEG) aims to establish trust in digital media by addressing aspects of authenticity, provenance, and integrity. JPEG Trust will provide a framework for establishing trust in media through secure annotation of media assets throughout their life cycle, using cryptography as a key component. Cryptographically signing media allows deep-fakes to be repudiated by the interested person, since the signature on them will be absent or invalid. Such an anti-fake signature should allow “good” manipulation of original file (at least cropping), disallow “bad” manipulation, but also be space efficient, to save bandwidth on web servers once the media file is disseminated. Unfortunately, customary signature schemes like ECDSA do not have these properties. 

To address this challenge, a solution could rely on novel aggregatable signatures, such as the Boneh-Lynn-Shacham (BLS) signature <cit.>, which has been successfully used in blockchain technologies like Ethereum 2.0 to optimize storage[https://github.com/ethereum/consensus-specs/blob/dev/specs/phase0/beacon-chain.md#bls-signatures]. The BLS signature scheme makes use of a novel form of cryptography called pairing-based cryptography, which allows for a plethora of new functionalities like attribute-based encryption<cit.>. Aggregatable signatures could be employed in JPEG (possibly within the JPEG Trust standard itself) in such a way to permit benign alterations of the image like cropping while preventing malicious tampering, without increasing too much the bandwidth occupation on web servers.




 §.§ Trusted Remote Media Processing on Cloud and Edge Computing Systems



In the emerging Smart Cities context, systems based on IoT (Internet of Things) play an important role to allow citizens to interact with the environment and to benefit from advanced services, such as video surveillance, intelligent traffic lightning, and air quality sensing. From a technological point of view, using sensors and actuators to automate services is strategic, but managing, configuring, and optimizing the digital infrastructures to adapt their behavior to the specific needs of the context is a big challenge, both in terms of system design and security. Deepefake media detection in these scenarios represent a challenge due to the nature of possible manipulations of future citizens day life, hence a more holistic approach should be considered where the media production occur, since at the Edge Computing Systems <cit.>.  

In just 20 years, with the objective to increase system response and reduce communication latency, computation moved from mainframes and computing rooms towards Cloud Computing, Fog Computing, and lastly, Edge Computing. A Federated Cloud-Edge infrastructure is considered, where different administrative domains are in place and where Machine Learning software artifacts, in the assertion of Federated Learning even at the Edge, help to distribute intelligence in this scenario.

Multimedia acquisition devices based on IoT generate an unprecedented amount of data, with the need of developing Cloud-based video big data analytics frameworks. A distributed approach in video recording and elaboration systems, such as video surveillance systems based on IP cameras, is highly recommended to overcome the maximum storage or throughput limitation of Network Video Recorders installed on single machines. 
To perform such a variety of tasks, and to be able to modify a device’s behavior on-demand, the Function as a Service (FaaS) computational paradigm is generally adopted. FaaS allows to define several minimal applications and to run one or more instances of these on the same device at the same time.
FaaS framework relies on two configuration approaches: a local configuration file, generally YAML, or a secure remote server. However, both come with limitations: a local file configuration requires direct access to the device, physically or through a secure connection, to modify it.
Alternatively, a remote server can store and send updated configuration files, but it might be vulnerable to well-known cyberattacks suck as Man-in-the-middle (MITM) or Distributed Denial of Service (DDoS), making it unusable and unreliable.
To overcome such limitations, it is possible to benefit from three technologies that have been increasingly recognized to be able to address information access problems and system trustiness in different application domains: Federated Learning <cit.>, Blockchain and IPFS (InterPlanetary File System) starting at the Edge:

    
  * Federated Learning is a decentralized approach to training Machine Learning Models. In traditional Machine Learning, data is centralized in the Cloud, where a single model is trained on the entire dataset. Federated Learning, on the other hand, allows for training Machine Learning Models across multiple decentralized devices or servers that hold local data samples without exchanging them. Moreover, Federated Learning at the Edge refers to the application of Federated Learning techniques on Edge devices, such as IoT devices, or Edge Servers. This approach combines the benefits of Federated Learning, which ensures Data Privacy and reduces communication costs, with the advantages of Edge Computing Systems, which enables data processing and model training to occur closer to where the data is generated, hence, fake Media might not exit from the Edge. 
    
  * The use of Blockchain, supported by the flexibility and robustness of Smart Contracts, allows the combining of the well-known FaaS paradigm with the intrinsic features of data non-repudiation and immutability, replacing the service configuration with a Smart Contract, guaranteeing protection against distributed cyber-attacks <cit.>. 
    
  * IPFS is a distributed system for storing and accessing files. Since the block size of the Blockchain does not allow storing files, these can be uploaded to this special file storage, which produces a unique hash value to be used as a key to access its content <cit.>.













§ CONCLUSION


In this paper, we have conducted a comprehensive review of the state-of-the-art techniques and challenges in Deepfake media forensics. Our exploration covered the core areas of Deepfake detection, attribution and recognition, passive authentication, detection in realistic scenarios, and active authentication. Each of these areas addresses specific facets of the Deepfake phenomenon, from the identification of synthetic media and tracing their origins to ensuring the robustness of detection systems in real-world environments and embedding verifiable information within media for instant authentication.

Our review underscores the significant progress made in the field and we also acknowledge the ongoing challenges, particularly in addressing the ethical and security risks posed by the misuse of Deepfake technology. 

Future work will focus on conducting a more in-depth analysis of practical countermeasures and gaining deeper insights into real-world applications (e.g. highly compressed data). 
Furthermore, a thorough study on the generalization capability of state-of-the-art deepfake detectors will be undertaken. 
In addition, we will explore issues related to scalability, real-time processing capabilities and ease of use of state-of-the-art deepfake detector approaches. 








§ ACKNOWLEDGMENT




This study has been partially supported by SERICS (PE00000014) under the MUR National Recovery and Resilience Plan funded by the European Union - NextGenerationEU 




splncs04.bst




 Paper Title: Mitral Regurgitation Recogniton based on Unsupervised Out-of-Distribution Detection with Residual Diffusion Amplification.




Residual Diffusion Amplification








Liu et al.



1National-Regional Key Technology Engineering Laboratory for Medical Ultrasound, School of Biomedical Engineering, Medical School,
Shenzhen University, China

2Medical Ultrasound Image Computing (MUSIC) Lab, Shenzhen University, China


3Key Laboratory for Bio-Electromagnetic Environment and Advanced Medical Theranostics, School of Biomedical Engineering and Informatics,
 Nanjing Medical University, China

4Shenzhen RayShape Medical Technology Co., Ltd, China

5Center for Cardiovascular Disease, The Affiliated Suzhou Hospital of Nanjing Medical University, China




Mitral Regurgitation Recogniton based on Unsupervised Out-of-Distribution Detection with Residual Diffusion Amplification
    
Zhe Liu1,2Zhe Liu, Xiliang Zhu and Tong Han contribute equally to this work.

Corresponding authors’ email: gzhongshan1986@163.com and  xinyang@szu.edu.cnXiliang Zhu1,2⋆Tong Han1,2⋆Yuhao Huang1,2Jian Wang3

Lian Liu1,2,4Fang Wang5Dong Ni1,2Zhongshan Gou5()

Xin Yang1,2()

    August 5, 2024
=====================================================================================================================================================================================================================================================================================



§ ABSTRACT 


Mitral regurgitation (MR) is a serious heart valve disease. 
Early and accurate diagnosis of MR via ultrasound video is critical for timely clinical decision-making and surgical intervention. 
However, manual MR diagnosis heavily relies on the operator's experience, which may cause misdiagnosis and inter-observer variability. 
Since MR data is limited and has large intra-class variability, we propose an unsupervised out-of-distribution (OOD) detection method to identify MR rather than building a deep classifier.
To our knowledge, we are the first to explore OOD in MR ultrasound videos.

Our method consists of a feature extractor, a feature reconstruction model, and a residual accumulation amplification algorithm. 
The feature extractor obtains features from the video clips and feeds them into the feature reconstruction model to restore the original features.
The residual accumulation amplification algorithm then iteratively performs noise feature reconstruction, amplifying the reconstructed error of OOD features.

This algorithm is straightforward yet efficient and can seamlessly integrate as a plug-and-play component in reconstruction-based OOD detection methods.
We validated the proposed method on a large ultrasound dataset containing 893 non-MR and 267 MR videos. 
Experimental results show that our OOD detection method can effectively identify MR samples.








§ INTRODUCTION

Mitral regurgitation (MR) is a significant heart valve disorder that becomes more common with advancing age.
It may cause heart size changes, decreased cardiac function, and even life-threatening consequences <cit.>.
Early detection and accurate assessment of MR enables timely intervention, holding vital clinical significance.
Color Doppler echocardiography is the primary tool to diagnose MR, enabling visualization of blood flow direction and velocity within the heart that is not visible in standard B-mode ultrasound <cit.>. 
However, manual assessment is subjective and easily affected by experts' experience, potentially leading to diagnostic errors and observer inconsistency.  
Figure <ref> shows the challenges in MR recognition, where similar features appear in both negative and positive cases.









Deep learning techniques can potentially address the above problems, e.g., by training a classifier to recognize MR. 
However, the limited availability of MR data and its significant intra-class variability compared to normal data can constrain the performance of binary classifier <cit.>. For example, the imbalanced data distribution caused by the lack of MR data can severely affect the performance of supervised methods <cit.>. 
Recently, there have been significant advancements in out-of-distribution (OOD) detection, opening up new possibilities for MR recognition.
The goal of OOD detection is to identify samples that are inconsistent with the distribution of in-distribution (ID) data (i.e., training data) <cit.>. 
It is noted that in MR recognition, numerous normal samples with similar features are regarded as ID data, and minimal but diverse MR cases are considered OOD data. 
Although promising, we found that studies applying OOD detection to MR recognition have not yet been reported. 

Applying existing OOD research to MR recognition without modification may face several challenges. 


Unsupervised OOD detection methods for videos showed great potential, and they mainly included two main streams: frame prediction and generation techniques.

Among frame prediction methods, those using optical flow maps <cit.> and long sequences of observations <cit.> were deemed inadequate for MR ultrasound videos due to noise, artifacts, and the cardiac cycle's motion pattern.
Besides, generation-based approaches may encounter information bottlenecks caused by discrepancies between potential input dimensions and the dimensions reconstructed by the model <cit.>.
The diffusion model solely employed the mean and standard deviation for denoising the reconstructed image, thereby mitigating the bottleneck issue mentioned previously to ensure the reconstruction of echocardiographic videos <cit.>.




In this study, we propose a diffusion-based unsupervised OOD detection method for recognizing MR from echocardiography videos. 
Specifically, we first employ pre-trained models as feature extractors to extract features from video clips. 
Second, we train a diffusion model for feature-level reconstruction. 
Last, we design a residual accumulation amplification algorithm during the testing phase.
This algorithm iteratively performs noise feature selection and reconstruction, amplifying the reconstructed error of OOD features.

We validate the proposed method on a large four-chamber cardiac (4CC) ultrasound video dataset containing 893 non-MR and 267 MR videos, and the experimental results show that the proposed method is effective. 
We believe we are the first to explore OOD detection in MR recognition with echocardiography videos.











	





















§ METHODOLOGY

As shown in Figure <ref>, our method consists of a feature extractor, a feature reconstruction model, and a residual accumulation amplification algorithm. 
The feature extractor catches features from video clips and inputs them into the reconstruction model to rebuild the features. 
Then, during testing, under the proposed residual accumulation amplification algorithm, the OOD data can gradually be more distinguishable from the ID data.







 §.§ Representative Feature Extractor


The feature extractor aims to capture essential knowledge from video clips. Let V ∈ R^k × 3 × h × w denote a video clip consisting of k consecutive RGB images with a resolution of h × w pixels, the feature v ∈ R^l is obtained by feeding V into the feature extractor f(·). 
The architecture of the feature extractor is flexible and can be convolutional neural networks, long short-term memory, visual transformers, and their hybrids. 
Here, we tried several models, including 3D-ResNet18 <cit.>, 3D-ResNet101 <cit.>, Video MAE <cit.>, and X-CLIP <cit.>. 
These models were pretrained on Kinetics-400 <cit.>, a large-scale dataset containing more than 160,000 videos with 400 categories. We froze the models' weights without any fine-tuning.




 §.§ Feature Reconstruction Model


The diffusion model gains the capability to generate diverse samples by disrupting training samples with noise and learning the reverse process.
For a given video clip V, features are obtained by v=f(V).
The process of feature disruption can be formulated as a progressive addition of Gaussian noise (standard deviation: σ) to the input data v_T obtained by sampling from a data distribution p_data(v) with standard deviation σ_data. After adding noise, the data distribution could be represented as p(v;σ) and p(v;σ_min)≈ p_data(v) ideally. If σ_max≫σ_data, the p(v;σ) becomes isotropic Gaussian and allows to sample a data v_0∼𝒩(0,σ_maxI). The reversed process of diffusion model can be formulated as v_0 being gradually denoised with noise levels σ_0=σ_max>σ_1>…>σ_T=0 into new samples, until it becomes a data with data distribution of p_data(v). 




















In detail, we follow the diffusion model in <cit.>, treating s_θ(v,σ) as a denoising function that directly estimates denoised samples with target function ℒ:

    ℒ=E_v∼ p_dataE_ϵ∼𝒩(0,σ I)‖ s_θ(v+ϵ;σ)- v ‖_2.





Besides, in <cit.>, researchers untangled the design choices of previous diffusion models and offered a framework where each component (Sampling, Network and preconditioning, Training) can be adjusted independently. Therefore, in our study, the denoising function s_θ(· ) is formulated as follows:

    s_θ(v;σ)=c_skip(σ)v+c_out(σ)*F_θ (c_in(σ)v;c_noise(σ)),

where F_θ(·) is a multi-layer perceptron (MLP) with an encoder-decoder structure. 

c_skip(·) modulates the skip connection.
c_out(·) and c_in(·) scale magnitudes of the noise variance.
c_noise(·) scales θ to a suitable value for the input of F_θ(·).

The reverse process of disrupting samples in the feature reconstruction of video clips can be formulated as:

    v_t=v_t-1+ε/2 s_θ(v_t-1, σ )+√(ε)z_t,


where  ε > 0 is a predefined step size and  z_t∼𝒩(0,I) is  random term. With certain constraints, reversing multiple times at very small step sizes results in the final generated sample obeying the distribution p_data(v).


The diffusion process is governed by multiple hyperparameters, with our main focus being the exploration of the effect of noise distribution on model performance. 
For fair comparisons, we follow <cit.> to determine other settings, including model structure, learning rate scheduler, etc.

The diffusion model is central to the OOD detection model, used to recover the samples after being corrupted with Gaussian noise during the testing phase. 
The OOD information is derived by evaluating the discrepancy between samples before and after restoration.



 §.§ Residual Accumulation Amplification


Inspired by the genetic accumulation of mutated genes with inheritance in genetics <cit.>, we designed a residual accumulation amplification algorithm used in the testing stage. 
The core of the algorithm involves selecting MR recognition-related features and amplifying the reconstruction error via multiple iterations of noise feature reconstruction to enhance the distinction between OOD and ID data.

In each iteration, we first randomly sample s Gaussian noises from the normal distribution 𝒩(μ,σ), and then add them to the anchor feature to randomly destroy features that represent different information. The formula is as follows:

    v_t^i = v_t + m · n_t^i,    i ∈{ 1,2,...s },

where n^i denotes the sampled Gaussian noise and m denotes the weight of noise. v and v^i represent the anchor and noise features, respectively. t is the current iteration number. 
In the first iteration, the anchor feature (v_1) is extracted from the video clips using the feature extractor. 
Those noise features are fed into the reconstruction model, and the reconstructed error is calculated as follows:

    e_t^i = ‖ DM(v_t^i) - v_1 ‖_2,    i ∈{ 1,2,...s },

where DM(·) denotes the reconstruction model and e^i_t denotes the reconstructed error in the t-th iteration. ‖·‖_2 represents the L_2 norm. 

As shown in Figure <ref> (D), OOD data is more sensitive to reconstruction, and the difference between the reconstruction error of ID and OOD data will increase after multiple reconstructions.
Therefore, we continue to add Gaussian noise to the current noise features for the next reconstruction. 
The noise feature corresponding to the maximum error is selected as the anchor feature in the next iteration. We believe that the larger error arises from the destruction of features unrelated to MR recognition. 
As these irrelevant features are eliminated after multiple iterative selections, the remaining features that are more related to MR recognition can enhance the detection efficacy: 


    v_t+1 = max_v_t(e_t).


The reconstruction error continues to be amplified until the stopping condition is reached, that is, t = T. The maximum error in the last iteration is regarded as the difficulty of feature reconstruction: diff = max(e_T^1, e_T^2, ... e_T^s).
The sample is classified according to a data-driven threshold (thre), defined as thre=μ_diff + 0.001 σ_diff, where μ_diff and σ_diff are the mean and standard deviation of the difficulty of feature reconstruction in the validation set.


















§ EXPERIMENTAL RESULTS






 §.§ Dataset and Implementations


The 4CC ultrasound video dataset was collected, comprising 893 non-MR and 267 MR videos.


Regions of interest were extracted from frames originally sized at 1016×708 and resized to 224×224 for further processing.
Multiple non-overlapping video clips were sampled from each video, and each clip consists of 16 consecutive frames, resulting in 3480 clips. 
Those clips cropped from MR videos are designated as positive samples (i.e., OOD samples), while others are taken as negative samples (i.e., ID samples). 

The dataset was divided into a training set with 1506 ID clips, a validation set with 621 clips (372 ID/249 OOD), and a testing set with 1353 clips (801 ID/552 OOD).



We adopted log-normal sampling (mean: -0.05, var: 1.5) for noise generation. Training epoch is set to 100 with batch size=64, using the AdamW optimizer with learning rate=0.0002. We used five metrics for evaluation: F1-score, Recall, Precision, Specificity, and Accuracy. All experiments were conducted on the PyTorch 2.0.1 framework with NVIDIA Tesla A40 GPUs.









 §.§ Performance Comparisons

We compared our proposed method with a basic supervised classification method (3D-ResNet50) and other reconstruction-based unsupervised OOD detection methods under the same settings. 
In 3D-ResNet50, we gradually increased the number of OOD clips N in the training set.

As shown in Table <ref>, the proposed method surpasses previous unsupervised methods and even surpasses 3D-ResNet50 with limited OOD data.




































 §.§ Comparison of Feature Extractor

We compared the MR recognition performance under different feature extractors.

Thus, we used X-CLIP as the backbone in the following experiments for it achieved the best results in Table <ref>.
The X-CLIP extractor with Transformer modules pretrained on a vast video dataset has a stronger capability to capture rich semantic information. 
This may be because text-video model fusion allows a more comprehensive understanding of complex features by capturing diverse aspects of the data. 
The performance gap between Convolution-based models and Transformer-based models can be attributed to the superior ability of Transformer with self-attention modules in learning long-range dependence features compared to convolutional modules.
Particularly in video tasks with temporal information, Transformers exhibit advantages due to their enhanced capability to capture temporal dependencies and patterns.


































 §.§ Diffusion Model Analysis
 
We tend to find the optimal T and 𝒩(μ,σ) to preserve sufficient structural information of the video clip while disrupting potential anomaly information. 
Then, higher reconstruction errors can be obtained to judge the associated video frames as abnormal. 
For Gaussian distributions 𝒩(μ,σ), μ ranges from -0.05 to -4 and σ ranges from 0 to 2.
Besides, the start points T of the reverse process are set to 1,5 and 10. 
For each value of T, 25 different Gaussian noise samples can be produced, as shown in Figure <ref> (A-C). It can be observed that μ=-0.05, σ=1.5 and T=5 obtain the best score. 
Overall, increasing the T value while fixing μ and σ can improve the OOD detection results. 
σ is the primary determinant of the model. 
Specifically, the performance gap for changing the value of σ is 3% in the F1-score, with all other hyperparameters remaining constant.













 §.§ Residual Accumulation Amplification Analysis

Residual accumulation amplification consists of residual accumulation (RA) and noise sampling (NS). 
This section analyzes the impact of both components on MR recognition performance.

For RA, the reconstruction is iterated T times (i.e., 1, 2, 3, 4, 5, and 10). Table <ref> shows that when T is set to 5, the diffusion model achieves optimal detection results on all metrics, and T=10 causes a slight performance drop. 
For NS, we use the control variable method to assess how different parameters affect diffusion model performance.
First, we set the noise sampling number s=3 as a trade-off for computational efficiency. Then we change the weight of noise m to analyze its impact. 

Results in Table <ref> show that consistent optimal results across all metrics can not be obtained with a single set of parameters.
But in general, the number of iterations T=5 can get good results easier. 
Finally, we choose the optimal hyperparameters (m=1.8, s=3, T=5) to construct the OOD detection model. 
The impact of each module on model performance is shown in Table <ref>.
RA improves the model performance by amplifing the difference between ID and OOD data reconstruction errors, which is consistent with Figure <ref> (D). However, its impact is limited by irrelevant features in the extracted information. NS addresses this issue effectively, significantly boosting performance by retaining MR recognition-relevant features through iterative selections.

























































§ CONCLUSION

We propose an innovative model tailored for recognizing MR echocardiographic videos using an unsupervised OOD detection approach. 
Our model exclusively leverages features extracted from ID samples by a fixed-weight feature extractor to train the diffusion model. 
Subsequently, by iteratively introducing disturbance, the model selects MR recognition-related features and amplifies the difference between reconstruction errors of OOD and ID data. 


Comprehensive experimentation conducted on MR datasets substantiates the efficacy of our proposed method. 
Future endeavors may explore expanding our approach to encompass multi-classification within the OOD samples.



  §.§.§ Acknowledgement.

This work was supported by the grant from National Natural Science Foundation of China (12326619, 62101343, 62171290), Science and Technology Planning Project of Guangdong Province (2023A0505020002), Shenzhen-Hong Kong Joint Research Program (SGDX20201103095613036) and Suzhou Gusu Health Talent Program (GSWS 2022071, GSWS 2022072).



  §.§.§ Disclosure of Interests.

The authors have no competing interests to declare that are relevant to the content of this article.










splncs04



 Paper Title: An Explainable Vision Transformer with Transfer Learning Combined with Support Vector Machine Based Efficient Drought Stress Identification.





[1,2]Aswini Kumar Patraaswinipatra@gmail.com / akp@nerist.ac.in

3]Ankit Varshney 

[2]Lingaraj Sahools@iitg.ac.in

[1]Dept. of CSE, NERIST, India

[2]Dept. of BSBE, IIT Guwahati, India

[3]Dept. of Physics, IIT Guwahati, India


Early detection of drought stress is critical for taking timely measures for reducing crop loss before the drought impact becomes irreversible. The subtle phenotypical and physiological changes in response to drought stress are captured by non-invasive imaging techniques and these imaging data serve as valuable resource for machine learning methods to identify drought stress. While convolutional neural networks (CNNs) are in wide use, vision transformers (ViTs) present a promising alternative in capturing long-range dependencies and intricate spatial relationships, thereby enhancing the detection of subtle indicators of drought stress. We propose an explainable deep learning pipeline that leverages the power of ViTs for drought stress detection in potato crops using aerial imagery. We applied two distinct approaches: a synergistic combination of ViT and support vector machine (SVM), where ViT extracts intricate spatial features from aerial images, and SVM classifies the crops as stressed or healthy and an end-to-end approach using a dedicated classification layer within ViT to directly detect drought stress. Our key findings explain the ViT model's decision-making process by visualizing attention maps. These maps highlight the specific spatial features within the aerial images that the ViT model focuses as the drought stress signature.  Our findings demonstrate that the proposed methods not only achieve high accuracy in drought stress identification but also shedding light on the diverse subtle plant features associated  with drought stress. This offers a robust and interpretable solution for drought stress monitoring for farmers to undertake informed decisions for improved crop management.







An Explainable Vision Transformer with Transfer Learning
Combined with Support Vector Machine Based Efficient
Drought Stress Identification
    *
    
===========================================================================================================================================





§ INTRODUCTION

Drought stress adversely affects the growth, development and yield of plants. The rise in global temperature coupled with the growing demand for water severely depletes soil water reserves. Water deficit and temperature stress can cause damage to the reproductive growth phase of crops, resulting in a substantial reduction in yield <cit.>. Furthermore, prolonged exposure to water deficit stress particularly in the reproductive stage can significantly affect grain quality. Ensuring sustainable crop yields in the face of escalating drought conditions has become a great challenge <cit.>. The plant's response to drought stress is influenced by a number of factors such as stress duration, severity, genotype and its developmental stage <cit.>. In plants, drought stress in early phase induces reduced stomatal conductance, transpiration and net carbon assimilation, whereas mid- and late-phases result in visible leaf senescence, stagnant growth and wilting. The early detection of drought induced phenotypic changes in plants, facilitates in preventing the adverse impacts of drought from irreversible damage through timely irrigation and other management practices. The point based, contact type and destructive methods based estimation of drought stress impact on crops are limited by good approximation of representative area as well as sample size and hence, difficult to implement over larger agricultural production systems. On the contrary, image-based, non-invasive and high-throughput analyses of subtle phenotypical changes in plants due to drought and other stresses, provides an efficient means of monitoring  stress levels in crops and good prediction of drought stress impact <cit.>. Various imaging techniques and sensors are utilized to precisely capture stress responses, including simple red-blue-green (RGB) , thermal, multispectral, hyperspectral, or fluorescence imaging methods <cit.>. However, traditional image processing generates significant amount of variation in light intensity, shading, occlusion, etc. leading to variation in quality thereby increases complexity of image processing <cit.>. Machine learning (ML) techniques, on the other hand, solve such complex problems which are nonlinear in nature by identifying key feature to be extracted from the acquired image in real time, based on knowledge and/or domain expertise of the user. Hence, ML aids in better decision-making and informed actions in real-world scenarios with minimal human intervention <cit.>. Deep learning (DL) <cit.> techniques however, help in extracting different information from the raw image data during input training, by means of various convolutions which allows larger learning capabilities, that result in the best classification/regression and higher performance and precision. Convolutional Neural Networks (CNNs) have been developed for providing a much-detailed representation of the image features, for a more discriminative and detailed analysis of plant images, leading to highly accurate classifications <cit.>. Many studies have used deep CNN models to detect drought stress in affected plants, mostly solved either by finely tuned CNNs or by training from scratch. 

 Early work by Zhuang et al. <cit.> utilized a combination of segmentation, color, and texture feature extraction, employing a gradient boosting decision tree (GBDT) method to detect water stress in maize. However, a subsequent study using the same RGB dataset demonstrated that a deep convolutional neural network (DCNN) significantly outperformed GBDT in terms of performance <cit.>.
Sankararao et al. <cit.> developed a comprehensive pipeline for detecting water stress in groundnut canopies using hyperspectral imaging. The pipeline included steps for image quality assessment, denoising, and band selection, followed by classification using Support Vector Machine (SVM), Random Forest (RF), and Extreme Gradient Boosting (XGBoost/XGB). Transfer learning with a pre-trained DenseNet-121 model has been utilized to classify soybean color images into four levels of drought severity <cit.>. Similarly, a controlled dataset is created for chickpea crops to identify water stress at different stages—control, young seedling, and pre-flowering—using a CNN-LSTM hybrid model where CNNs that served as feature extractors and LSTM models for predicting water stress levels <cit.>.
The evaluation of various CNN models, including AlexNet, GoogLeNet, and Inception V3, for distinguishing between stress and non-stress conditions in water-sensitive crops such as maize, okra, and soybean, revealed that GoogLeNet achieved the highest identification accuracy <cit.>.
Chlorophyll fluorescence images of wheat canopies have been studied employing a multi-step process involving segmentation and feature extraction with a correlation-based gray-level co-occurrence matrix (CGLCM) and color features, followed by classification using different ML classifiers. The tree-based methods, particularly random forest (RF) and extra trees classifier, is found giving superior performance <cit.>.
A dataset of aerial images of potato canopies has been studied using a deep learning-based model for drought stress identification, leveraging various imaging modalities and their combinations <cit.>. An explainable deep learning framework was later proposed using the same dataset, based on CNN architectures and transfer learning, achieving improved accuracy <cit.>.
A regression approach is applied to predict the drought tolerance coefficient using hyperspectral images of tea canopies; the comparison of SVM, RF, and Partial Least Squares Regression (PLSR) revealed that SVM was the most effective <cit.>. Machine learning and deep learning methods were compared, including DNN, SVM, and RF, for drought stress detection using full spectra and first-order derivative spectra <cit.>. Goyal et al. <cit.> created a dataset of RGB images of maize crops and proposed a custom-designed CNN model that outperformed existing state-of-the-art CNN architectures in early drought stress detection.  However, Vision Transformers (ViTs) offer notable advantages over CNNs in capturing long-range dependencies and global context due to their self-attention mechanism, whereas CNNs rely on local receptive fields and often struggle to capture relationships across various parts of an image <cit.>. Additionally, ViTs exhibit greater flexibility with input image sizes and handle complex patterns and data variations more effectively <cit.>.

The self-attention mechanism in ViTs enables more accurate classification by considering the entire image context, which enhances accuracy and robustness compared to traditional CNNs <cit.>. For example, Dosovitskiy et al. demonstrated that ViTs could outperform CNNs in image classification tasks, highlighting their potential for diversified applications [<cit.>]. Recent studies have also applied ViTs to plant disease identification <cit.>.  In this study, we applied ViTs to drought stress identification to (a) achieve higher accuracy in detecting drought stress, b) leverage the ViT’s superior ability to handle complex patterns and variations inherent in field images captured in natural settings, ensuring the model’s practicality and relevance to real-world conditions, c) interprete model by utilizing ViTs to generate attention maps, which provide insights into the model’s decision-making process and contribute to the transparency and explainability of the predictions.




§ MATERIALS AND METHODS

In this section, we begin by introducing the experimental dataset. Next, we present the drought stress classification model using two different approaches. In the first approach, we employ the vision transformer with transfer learning. In the second approach, we propose a framework that utilizes a vision transformer as a feature extractor, followed by the integration of an SVM as the classifier. Additionally, we investigate the interpretability of the model by generating and analyzing the attention maps. This comprehensive use of the ViT elucidates how the spatial features of drought stress can be precisely identified. Finally, we discuss the performance metrics for the proposed approaches.


 §.§ Preparing the Data

The potato crop aerial images utilized in this study have been sourced from a publicly accessible dataset that encompasses multiple modalities <cit.>. Collected from a field at the Aberdeen Research and Extension Center, University of Idaho, these images serve as valuable resources for training machine learning models dedicated to crop health assessment in precision agriculture applications. Acquired using a Parrot Sequoia multi-spectral camera mounted on a 3DR Solo drone, the dataset features an RGB sensor with a resolution of 4,608 × 3,456 pixels and four monochrome sensors capturing narrow bands of light wavelengths: green (550nm), red (660nm), red-edge (735nm), and near-infrared (790nm), each with a resolution of 1,280 × 960 pixels. The drone flew over the potato field at a low altitude of 3 meters, with the primary objective of capturing drought stress in Russet Burbank potato plants attributed to premature plant senescence.




The dataset consists of 360 RGB image patches in JPG format, each measuring 750×750 pixels. These patches were derived from high-resolution aerial images through cropping, rotating, and resizing operations. The dataset is divided into a training subset of 300 images and a testing subset of 60 images. Ground-truth annotations, provided in both XML and CSV formats, indicate regions of healthy and stressed plants marked by rectangular bounding boxes. These annotations were manually created using the  LabelImg software. An example RGB image and its annotated regions are shown in Fig. <ref>. The testing subset is independent of the training subset and comes from different aerial images. Additionally, the dataset includes corresponding image patches from spectral sensors with red, green, red-edge, and near-infrared bands, each sized 416×416 pixels. However, we are only utilizing the RGB images due to the limitations of the low-resolution monochromatic images.

We utilized the augmented dataset prepared by Butte et al., expanding the original set from 300 to 1,500 images for training. The test set comprised 60 images. From each augmented and test image, annotated windows (i.e., rectangular bounding boxes) were extracted, creating separate folders for "healthy" and "stressed" images. These folders served as input for our proposed model. The final count of training images for the "stressed" and "healthy" classes is 11,915 and 8,200, respectively. The test set, consisting of 60 images, yielded 401 healthy and 734 stressed images. We evaluated our method on this test set and further assessed the model performance through 5-fold cross-validation.



 §.§ Vision Transformer (ViT)
 



A vision transformer (ViT) is a type of neural network architecture that has revolutionized the field of computer vision <cit.>. Unlike traditional convolutional neural networks (CNNs), which process images pixel by pixel, ViT processes the input images in a sequential manner by dividing them into fixed-size patches, linearly embedding these patches, and then applying self-attention mechanisms for capturing global dependencies <cit.>. The VIT architecture used for our work is inspired by Dosovitskiy et al.<cit.> and depicted in Fig <ref>. It processes images by dividing them into fixed-size embedded patches, linearly transforming these patches, and treating them as sequences. The transformer architecture is then applied, incorporating multi-head self-attention mechanisms <cit.> that enable the model to capture long-range dependencies within the image. Layer normalization is applied before and after the multi-head attention, ensuring stable training, and a Multilayer Perceptron (MLP) head is added to the transformer's global representation for task-specific processing. 



  §.§.§ ViT Architecture
 The proposed ViT architecture is based on the layers in <cit.> and Dosovitskiy et al.<cit.> and comprises five main components: patch embedding, positional encoding, transformer encoder , normalization layer and a classification head. This is illustrated in Fig <ref>.

  
  * Patch Embedding: Input images are divided into fixed-size patches, which are then linearly embedded to create a sequence of embeddings.
  
  
  * Positional Encoding: To capture spatial information, positional encodings are added to the patch embeddings, allowing the model to understand the relative positions of different patches.
  
  
  * Transformer Encoder: The embedded patches are fed into a Transformer Encoder, which is the core component of the ViT architecture. This encoder consists of twelve identical encoder layers stacked together. Each attention layer analyzes the relationships between pairs of patches, allowing the model to understand how different image regions interact and influence each other. Each attention layer consists of the following:
        
            
  * Multi-Head Self-Attention: This mechanism allows the model to weigh the importance of different parts of the image. It captures global dependencies between the patches.
            
  * MLP (Multi-Layer Perceptron) Block: This block introduces non-linearity to the network and further processes the information from the attention layer.
            
  * Normalization Layers: Layer normalization is applied after the multi-head attention and MLP blocks to stabilize training.
            
  * Dropout: The model uses dropout at two levels: within the attention mechanism and after the MLP block. Dropout is used to prevent over-fitting by randomly dropping out neurons during training.           
        
   
  * Normalization Layer: Layer normalization applied to the output of the encoder serves several crucial purposes: Reduces internal co-variate shift, improves gradient flow, acts as regularization, handles varying input distributions and accelerates convergence.     
    
  
  * Classification Head: This head typically consists of a simple MLP layer that maps the feature representation to the desired output, such as class probabilities.






  §.§.§ Information Processing in ViT

At the core of the Vision Transformer (ViT) is the concept of a token, which plays a crucial role in the model's ability to process and understand images. A token is a fixed-size vector that represents a small patch of the input image. These tokens are at the core of the Vision Transformer model, forming the input sequence to the Transformer layers. This token-based approach enables the model to process and understand the image by focusing on the relationships between these patches through self-attention. The detailed explanation is given below.


    
  * Dividing the Image into Patches:
    
        
  * The input image is divided into smaller, fixed-size patches. For example, an image of size 224x224 pixels might be divided into patches of size 16x16 pixels, resulting in (224/16)^2 = 196 patches.
    
    
    
  * Flattening and Embedding:
    
        
  * Each image patch is then flattened into a one-dimensional vector. For instance, a 16x16 patch with 3 color channels (RGB) would be flattened into a vector of length 16 × 16 × 3 = 768.
        
  * These flattened vectors (patch representations) are then linearly embedded into a higher-dimensional space. This is typically done using a learnable linear projection, transforming each vector into a fixed-size embedding, say of dimension 768.
    
    
    
  * Tokens:
    
        
  * After linear embedding, each flattened and embedded patch becomes a "token." In this example, the image is transformed into a sequence of 196 tokens, each representing a 16x16 patch of the original image.
    
    
    
  * Adding Positional Encoding:
    
        
  * Since the transformer model does not inherently understand the order or position of tokens, positional encodings are added to each token to incorporate information about its original position in the image. This helps the model understand spatial relationships between patches.
    
    
    
  * Processing by Transformer Layers:
    
        
  * These tokens are then processed by the Transformer layers, which include self-attention mechanisms. The self-attention mechanism computes relationships between these tokens to understand how different parts of the image relate to one another.
    




 §.§ ViT with Transfer Learning

The proposed framework, as shown in Fig <ref> effectively combines transfer learning, the power of the Vision Transformer and attention-based interpretability to address the challenging task of drought stress identification in potato crop images captured in natural settings. A core component of this approach is the utilization of pre-trained weights. This technique, known as transfer learning, involves leveraging knowledge gained from solving one problem (often a large-scale image classification task) and applying it to a different but related problem. By employing pre-trained weights, the model can benefit from the rich feature representations learned from a massive dataset, accelerating training and potentially improving performance. Using this the ViT is trained over multiple epochs using a combination of binary cross-entropy loss, Adam optimizer and learning rate tuning. Experimenting with different learning rates is essential to find the optimal value for convergence and generalization. Eventually, the classification layer is responsible for making the final prediction. It takes the output of the ViT model and maps it to two classes: "healthy" and "stressed." This layer typically consists of a fully connected neural network with a sigmoid activation function for binary classification.

To enhance model interpretability, the architecture incorporates a mechanism to capture attention weights. Attention weights reveal which parts of the input image the model focuses on when making a decision. By visualizing these weights, researchers can gain insights into the model's decision-making process and identify key image features that contribute to the classification. 



  §.§.§ Attention Maps
 





Attention maps provide insights into how the model focuses on different parts of the image during the self-attention mechanism. This mechanism allows the model to selectively attend to specific areas while processing visual information. The image is first divided into patches. Self-attention then helps the model prioritize relevant patches and their relationships for effective feature extraction.

Attention maps act as a visual representation of the weights assigned by the self-attention mechanism to each patch. These maps can be visualized for each self-attention layer within the ViT model, as each layer progressively learns more intricate relationships between these image features. Higher weights indicate a greater focus on a specific patch and its connection to others. By analyzing these maps, we can essentially see the model's "thought process" during image understanding. We can identify which regions it prioritizes for information extraction. The following section highlights the computation behind the attention mechanism <cit.>. 

At its core, self-attention computes a weighted sum of the values (features) based on the similarities (attention scores) between different positions in the sequence. This is achieved through three learnable matrices: Query (Q), Key (K), and Value (V). 

Query Matrix (Q): The query matrix is responsible for capturing the information about the current token being processed. It learns to encode the features of the token in a format suitable for comparison with other tokens. Each token in the sequence is associated with a query vector, which represents its characteristics in the context of the entire sequence.

Key Matrix (K): The key matrix holds information about the relationship between the current token and other tokens in the sequence. It learns to encode the features that determine how relevant each token is to the current token.

Value Matrix (V):  

The key matrix holds information about the relationship between the current token and other tokens in the sequence. It learns to encode the features that determine how relevant each token is to the current token.

Given a sequence of tokens X = [x_1, x_2, ..., x_n], the attention scores A are computed as:


    A = softmax( QK^T/√(d_k))


where Q = XW_Q, K = XW_K, V = XW_V, and W_Q, W_K, W_V are weight matrices. d_k represents the dimensionality of the key vectors.

Once the attention scores are computed, they are used to compute a weighted sum of the values:


    Attention(Q, K, V) = A · V


where A is the attention matrix.


























To capture different relationships between tokens, ViT (and Transformers in general) often employ multiple attention heads. Each head learns different sets of Q, K, V weight matrices and computes separate attention scores and weighted sums. The results from all heads are concatenated and linearly transformed to maintain a consistent output dimension.
Since self-attention does not inherently consider the order of tokens, positional encoding is typically added to the token embeddings to provide positional information.
 
Algorithm for Capturing Attention Weights:
To enhance model interpretability, the architecture incorporates a mechanism to capture attention weights. The algorithm <ref> outlines the steps during the training of ViT to capture and utilize attention weights for the drought identification task. The key components of this class include the initialization, forward pass, attention weight capture, output size determination, and retrieval of attention weights. By visualizing these weights, researchers can gain insights into the model's decision-making process and identify key image features that contribute to the classification. 

Initialization
The VisionTransformerBinary class is initialized with a pre-trained Vision Transformer model passed as . During initialization, the model assigns the provided  to its own  attribute. Additionally, it sets up a fully connected (linear) layer () with a size appropriate to the output of the Vision Transformer. This layer is responsible for converting the output of the Vision Transformer into a format suitable for binary classification. An empty list  is also initialized to store the attention weights captured during the forward pass.

Forward Pass
The  function is central to the operation of the VisionTransformerBinary class. When an input image  is passed through the model, the function first clears any previously stored attention weights. This ensures that the attention weights list is fresh and only contains data relevant to the current input.

Next, the function registers hooks on the self-attention layers of each block in the Vision Transformer. These hooks are responsible for capturing the attention weights during the forward pass. For each block in the Vision Transformer's encoder layers, the self-attention layer is accessed, and a hook is registered to capture its attention weights using the  helper function. The hooks are stored in a list to facilitate their removal later.

With the hooks in place, the input image is passed through the Vision Transformer and then through the fully connected layer. This produces the model's output for the given input. After the forward pass is complete, the hooks are removed to free up memory, ensuring that they do not persist and interfere with future operations.

Attention Weights Capture
The  helper function is designed to capture the attention weights during the forward pass. It is triggered by the hooks registered on the self-attention layers. This function receives the module, input, and output as arguments. From the input, it extracts the query, key, and value components, which are essential for computing the attention weights. These components, along with the output, are appended to the  list. By capturing these components, the model can later compute and analyze the attention scores, which provide insights into the regions of the input image that the model focuses on during classification.

Output Size Determination
The  helper function determines the output size of the Vision Transformer model. This function performs a forward pass with a zero tensor of appropriate dimensions through the Vision Transformer. By doing so, it captures the shape of the output tensor produced by the Vision Transformer. The size of the last dimension of this output tensor is then returned, which is used to initialize the fully connected layer with the correct input size.

Retrieval of Attention Weights
The  function provides a simple interface to retrieve the captured attention weights. It returns the  list, allowing external components or functions to access the attention weights for further analysis or visualization.

The modular design of the class, with separate functions for initialization, forward pass, attention weight capture, output size determination, and retrieval, provides a clear and maintainable structure. This design facilitates the integration of attention-based insights into the classification process, enhancing the interpretability and performance of the model. 






 §.§ Integrating Vision Transformer (ViT) and Support Vector Machine (SVM)


This framework focuses on combining a vision transformer (ViT)  and a support vector machine (SVM) within a three-step approach for effective stress identification.

  
  * Feature Extraction: For each image in the dataset, extract the final hidden state or pooled representation from the ViT model to obtain a feature vector. Let 𝐗 represent the set of embedded features extracted from a dataset of images, and 𝐲 represent the corresponding class labels.
  
  
  * Training SVM: Train an SVM classifier using the extracted features 𝐗 and their corresponding class labels 𝐲.
  
  
  * Classification: For a new, unseen image, extract features using the pre-trained ViT model and use the trained SVM classifier to predict its class.



  §.§.§ Support Vector Machine (SVM)

Support Vector Machines <cit.> aim to find a hyperplane that best separates a given set of data points into different classes. Given a set of labeled data points (𝐱_1, y_1), (𝐱_2, y_2), …, (𝐱_n, y_n) where 𝐱_i is the feature vector of the i-th data point, and y_i is its corresponding class label (y_i ∈{1, 0} for binary classification), SVM seeks to find a hyperplane defined by the equation:


    𝐰·𝐱 + b = 0


where 𝐰 is the weight vector and b is the bias.

The goal is to maximize the margin, which is the distance between the hyperplane and the nearest data point from each class. The margin is computed as the perpendicular distance from a data point 𝐱_i to the hyperplane:


    margin = 1/𝐰· | 𝐰·𝐱_i + b |


Subject to the constraint that for all data points:


    y_i (𝐰·𝐱_i + b) ≥ 1


This constraint ensures that data points are correctly classified and lie on the correct side of the hyperplane.

The SVM optimization problem can be formulated as:


    Minimize  1/2𝐰^2


Subject to the constraints:


    y_i (𝐰·𝐱_i + b) ≥ 1   for all  i


This is the primal form of the optimization problem. However, the SVM problem is often reformulated in its dual form, which introduces Lagrange multipliers α_i to handle the constraints. The dual formulation is:


    Maximize  ∑_i=1^nα_i - 1/2∑_i=1^n∑_j=1^nα_i α_j y_i y_j 𝐱_i ·𝐱_j


Subject to the constraints:


    0 ≤α_i ≤ C   for all  i



    ∑_i=1^nα_i y_i = 0


The solution to the dual problem provides the values of α_i, and the weight vector 𝐰 and bias b can be obtained from these values.

In cases where the data is not linearly separable, SVM can be extended to handle non-linear decision boundaries using the kernel trick. The feature space is implicitly mapped to a higher-dimensional space, making it possible to find a linear separating hyperplane in that space.

Classification is performed based on the decision function derived from the trained model. Once the SVM model is trained with a set of support vectors, it identifies a hyperplane that best separates different classes in the feature space. The decision function for a new data point 𝐱 is given by:


    f(𝐱) = 𝐰·𝐱 + b


Here, 𝐰 is the weight vector, and b is the bias term. The sign of f(𝐱) determines the predicted class:


    Class 1,    if  f(𝐱) > 0 
    Class 0,    if  f(𝐱) < 0


The magnitude of f(𝐱) provides a measure of how far the data point is from the decision boundary. Larger magnitudes indicate greater confidence in the classification.

The key role of support vectors in this process is that they are the data points lying closest to the decision boundary. Support vectors effectively determine the position and orientation of the hyperplane. The optimization process in SVM aims to maximize the margin between the classes, and support vectors are the data points defining the edges of this margin. In practice, many data points do not significantly contribute to the definition of the decision boundary. Only the support vectors, with non-zero Lagrange multipliers α_i in the dual formulation, are crucial for determining the hyperplane. This property makes SVM memory-efficient and computationally faster, especially in high-dimensional spaces.`



  §.§.§ ViT+SVM Framework

The process begins with input images that are fed into the ViT. The ViT, pre-trained on a vast dataset, is adept at extracting meaningful features from the images. Once the features are extracted by the ViT, they are compiled into a feature matrix, which also includes the corresponding labels indicating whether the plants in the images are healthy or stressed. This matrix forms the input to the SVM, a robust classifier known for its effectiveness in handling high-dimensional data. The SVM is trained to discern between the two classes—healthy and stressed—based on the features provided. Finally, the trained SVM predicts the class of new images, categorizing them as either healthy or stressed. The entire approach is depicted in Fig <ref>.
The efficacy of this framework is evaluated using a designated test set and k-fold cross-validation. 


 §.§ Performance Evaluation Metrics

The model's performance underwent assessment using various evaluation metrics, including accuracy, precision, recall (sensitivity) and the Receiver Operating Characteristic (ROC) curve. These metrics are computed based on the counts of true positives (TP), true negatives (TN), false positives (FP), and false negatives (FN), which collectively form a 2x2 matrix known as the confusion matrix. In this matrix, TP and TN indicate the accurate predictions of water-stressed and healthy potato crops, respectively. FP, termed as type 1 error, denotes predictions where the healthy class is inaccurately identified as water-stressed. FN, referred to as type 2 error, represents instances where water-stressed potato plants are incorrectly predicted as healthy. The classification accuracy is a measure of the ratio between correct predictions for both stressed and healthy images and the total number of images in the test set.


    Accuracy = True Positive + True Negative/Total Population



    Precision = True Positive/True Positive + False Positive



    Recall = True Positive/True Positive + False Negative



    F1-score = 2 ·Precision·Recall/Precision + Recall



Cross-validation is crucial in model development for two key reasons: it helps prevent over-fitting by assessing a model's performance across different subsets of the data, and it ensures the model's generalization ability, providing a reliable estimation of its effectiveness under various conditions. The model was trained and evaluated using k-fold cross-validation, a robust technique for assessing the generalization performance of the model. In each iteration of the k-fold cross-validation process, the dataset was partitioned into k folds, and the model was trained on k-1 folds while being validated on the remaining fold. This process was repeated k times, ensuring that every fold had the opportunity to serve as the validation set. For each fold, the Receiver Operating Characteristic (ROC) curve was plotted, illustrating the trade-off between true positive rate and false positive rate at various thresholds. After completing the k-fold cross-validation, the individual ROC curves were aggregated, and the mean ROC curve was calculated and plotted. AUC (area under the curve) measures the entire two-dimensional area underneath the entire ROC curve. AUC provides an aggregate measure of performance across all possible classification thresholds. One way of interpreting AUC is as the probability that the model ranks a random positive example more highly than a random negative example. AUC ranges in value from 0 to 1. A model whose predictions are 100% wrong has an AUC of 0.0; one whose predictions are 100% correct has an AUC of 1.0.This comprehensive approach provides a more reliable estimation of the model's performance, capturing its consistency across different subsets of the data and enhancing the overall assessment of its predictive capabilities.



§ RESULTS AND DISCUSSION

In this section, we present the experimental results of our proposed model for identifying drought stress in potato crop field images. First, we distinguish between healthy and stressed images. Then, we identify the spatial features responsible for the stress. Our experiments with the proposed Vision Transformer (ViT) framework were conducted in three ways:

    
  * ViT with Transfer Learning: Leveraging pre-trained weights.
    
  * ViT+SVM with Pre-trained Weights.
    
  * ViT+SVM with Optimal Weights. The optimal weights are the ones at which the model performs best while executing the ViT with pre-trained weights in the first case.


The PyTorch library and it's various sub-packages are leveraged for deep learning functionalities. In addition to this, sklearn, pandas, numpy packages are used for developing, training, and evaluating deep learning models, as well as for performing data analysis and pre-processing. The data pre-processing involves resizing images to (224, 224) pixels and transforming them into tensors. 




 §.§ Performance of ViT with Transfer Learning

We relied on PyTorch, a core library for deep learning, which provides the essential tools for tensor computations and neural network construction. Specifically,  and  were used to build and train the neural network. To handle image data, we used Torchvision's  for accessing pre-trained architectures and  for preprocessing images. Additionally, TQDM was employed to create progress bars, enhancing visibility into the training process. For data analysis and manipulation, we leveraged scikit-learn for machine learning utilities and Pandas for managing structured data.

The dataset preparation involved mapping class names to numerical labels using , where 'healthy' images were assigned the label 0 and 'stressed' images were assigned the label 1. We applied image transformations to resize the images to 224x224 pixels and convert them into tensor format, ensuring compatibility with the ViT model. The datasets were loaded using  from Torchvision, which facilitated the organization of images into training and testing sets. Data loaders were then created to handle batching, shuffling, and efficient data retrieval during training and validation, with a batch size of 128.

To adapt the Vision Transformer (ViT) architecture (as depicted in Fig <ref>) for our specific task of binary classification, we began by configuring the model using the models.vit_b_16 variant, setting pretrained=False to start with a clean slate. This approach allowed us to build the model from scratch without any influence from pre-existing learned features. Subsequently, we loaded custom pre-trained weights into the ViT model to realize the approach as shown in Fig <ref>. This step was crucial as it transferred learned representations from a previously trained model to our current architecture, leveraging prior knowledge to enhance performance. A custom class was defined to extend the ViT model by adding a fully connected layer for binary classification. This class also included methods to capture attention weights, which are essential for understanding the model's focus on different image regions.

The training process involved setting up a binary cross-entropy loss function to evaluate the performance of the classification model. We used the Adam optimizer with a learning rate of 0.001 and 0.0009 to update the model's weights. The training loop consisted of multiple epochs, during which the model was trained on the training dataset and validated on the test dataset. Training and validation losses were computed to assess the model's performance and adjust parameters accordingly. Throughout the process, attention weights were captured to provide insights into the model's decision-making process.






The learning curve typically plots both the training loss and validation loss curves on the same graph.
This allows to analyze how the model is performing on both the data it's trained on (training loss) and unseen data (validation loss).
The plot shown in Fig <ref> compares the training losses over the first 10 epochs for VIT trained from scratch versus VIT trained with pre-trained weights. The Vision Transformer (ViT) with pre-trained weights demonstrates a notably faster and more consistent reduction in loss compared to the ViT trained from scratch. This indicates superior initial performance and quicker convergence. As a result, we chose to further experiment with pre-trained weights rather than training the model from scratch.

The training loss is plotted in Fig <ref> for 20 epochs when using a Vision Transformer (ViT) with pre-trained weights and a learning rate of 0.001. The goal of training a machine learning model is to minimize the loss on the validation set. In an ideal scenario, the training loss and validation loss will both decrease as the model trains. However, this is not always the case.
In this graph, the training loss decreases steadily over the course of training. This suggests that the model is learning the training data well. However, the validation loss starts to increase after a certain number of epochs, specifically after epoch 8. This is a sign of over-fitting. Over-fitting is a problem that occurs when a model learns the training data too well, and it is no longer able to generalize to unseen data.

The training loss is plotted in Fig <ref>  for 20 epochs using a Vision Transformer (ViT) with pre-trained weights. The learning rate is 0.001 for epochs 1 to 10 and 0.0009 for epochs 11 to 20. The training loss decreases steadily, indicating the model is learning the training data well. However, the validation loss fluctuates and increases after epoch 6, suggesting over-fitting. The training loss is plotted in Fig <ref> for 20 epochs using a Vision Transformer (ViT) with pre-trained weights. The learning rate is 0.001 for epochs 1 to 8, with no dropout. For epochs 9 to 20, the learning rate remains 0.001, with an attention dropout of 0.2 and an MLP dropout of 0.1. The training loss decreases steadily, but the validation loss fluctuates and increases after epoch 8, indicating over-fitting once dropout is applied. The training loss is plotted in Fig <ref> for 20 epochs using a Vision Transformer (ViT) with pre-trained weights. The learning rate is 0.001 for epochs 1 to 8 and 0.0009 for epochs 9 to 20. An attention dropout of 0.2 and an MLP dropout of 0.1 are applied for epochs 9 to 20. The training loss decreases steadily, but the validation loss fluctuates and increases after epoch 8, suggesting over-fitting as dropout is applied and the learning rate is adjusted. 

In summary, all figures show a decrease in training loss, indicating the model is learning from the training data. Validation loss generally starts to increase or plateau after a certain number of epochs, suggesting overfitting is a common issue in these experiments. The introduction of dropout in Figures <ref> and <ref> seems to have exacerbated the over-fitting problem. However, the optimal result for the specified test set, reaching an accuracy of 0.929, is attained with a learning rate of 0.001 and without applying attention dropout or MLP dropout. The Vision Transformer (ViT) with transfer learning demonstrated impressive performance, achieving a precision of 0.95 and recall of 0.94 for stressed images, as well as a precision of 0.89 and recall of 0.91 for healthy images. The F1-scores for stressed and healthy images were 0.95 and 0.90, respectively, highlighting the model’s strong capability in accurately classifying both categories.



  §.§.§ Analyzing Attention Maps

Visualizing attention weights provides insights into how the ViT focus on different parts of an input during processing. By examining these weights, researchers and practitioners can understand the model's decision-making process, diagnose potential biases, and improve interpretability.

The following pseudocode outlines a systematic approach to calculate and visualize attention weights from a Vision Transformer model. This process involves capturing attention weights during the forward pass, computing attention scores, and generating visual representations of these scores.

Initialization:
The process begins by initializing an empty list to store the attention weights that will be captured during the forward pass of the Vision Transformer model. This list will later be used to compute and visualize the attention maps.

Forward Pass and Capture Attention Weights:
The next step involves iterating through each layer of the Vision Transformer model. For each layer, a hook is registered to capture the attention weights. The input image is then passed through the Vision Transformer to compute the output features. This stage ensures that attention weights are collected during the forward pass for later analysis.

Calculate Attention Score:
After capturing the attention weights, the algorithm processes each weight to compute the attention scores. This involves extracting the query, key, and value tensors from the hook outputs. The attention score is computed as per the principles discussed in section <ref>. The attention scores are then normalized to ensure they are in a range suitable for visualization.

Visualize Attention Maps:
In the visualization phase, each normalized attention map is resized to match the dimensions of the input image. A colormap (e.g., 'hot') is applied to the attention map to highlight areas of high attention. The attention map is then overlaid on the original image to create a visual representation of where the model is focusing.

Display:
Finally, both the original image and the overlaid attention maps are displayed, allowing for an interpretation of how the Vision Transformer model is making its decisions based on different regions of the input image.







The stressed image along with the corresponding attention maps from the 12 encoder blocks of the Vision Transformer (ViT) is shown in Fig <ref>. Several key observations can be made from these attention maps, including spatial relevance, hierarchical processing, interpreting model decisions, visualization of learned features, and using them as a diagnostic tool for model improvement. Each attention map is resized and overlaid on the original image, with colors indicating the attention intensity:

    
  * Red/Hot Colors: Indicate areas of high attention.
    
  * Yellow/Warmer Colors: Show areas of moderate  attention.
    
  * Dark/Cold Colors: Represent areas of low attention.
    




Spatial Relevance: The spatial relevance can give insights into which parts of the image the model finds important for differentiating between classes. The observations from the Fig <ref> are as follow:

    
  * Layer 1: The model's attention is broadly spread with higher intensity in some central regions. This indicates the model initially focuses on larger areas, possibly capturing the general structure of the image.
    
  * Layer 2-3: Attention starts to become more localized, with specific regions of high attention emerging. The model begins to identify and focus on more distinct features.
    
  * Layer 4-7: Attention further refines, highlighting specific areas more sharply. These layers likely focus on intermediate features that are critical for further processing.
    
  * Layer 8-12: Attention becomes very concentrated on specific, smaller regions. The model is likely focusing on fine-grained details and specific objects or patterns.

In the initial layers, the model's attention is broadly spread, capturing the overall structure of the image. As we move to higher layers, the attention becomes more localized, focusing on specific features and objects.

Hierarchical Processing: As we move through the layers of the ViT, attention maps can show how the model progressively refines its understanding of the image. In the lower layers, attention is distributed across large regions, capturing global context. As we ascend through the layers, the attention narrows down to more specific features, highlighting finer details and important objects within the image. This hierarchical processing is crucial for the model to effectively balance global and local information.

Interpreting Model Decisions:
By visualizing attention maps, we can interpret why the model made certain predictions. For instance, if the attention maps highlight specific objects or patterns in an image, it suggests that those elements influenced the classification decision. This interpretability can help validate the model’s decisions and identify potential biases or weaknesses.

Visualization of Learned Features:
The attention maps provide a visual representation of the features learned by the ViT. Unlike abstract feature vectors, these maps directly relate model activations to spatial locations in the input image. This visualization helps in understanding how the model processes visual information and forms its internal representations.

Diagnostic Tool for Model Improvement:
Analyzing attention maps can serve as a diagnostic tool for improving model performance. By examining where the model attends and comparing it with ground truth or human perception, we can identify areas where the model might be lacking or where it might over-emphasize certain features. This feedback loop can guide model refinement and training strategies.

In summary, attention maps in image classification tasks with Vision Transformers provide a transparent view into the inner workings of the model, highlighting which parts of the input image contribute most to its decision-making process. 


 §.§ Performance of ViT+SVM

We used a pre-trained Vision Transformer (ViT) model to extract features from both the training and testing datasets. These extracted features were then utilized by an SVM to identify stress, aiming to distinguish between stressed and healthy images. The implementation was done using the PyTorch framework, leveraging both its core library and the torchvision module. Key libraries such as torch, torchvision, csv, pandas, and scikit-learn were imported to facilitate feature extraction, data transformations, file handling, and classification tasks.

The pre-trained ViT model weights were sourced from a specified path and loaded into an initialized ViT model. We used the ViT-B/16 model, whose architecture was previously discussed, and obtained it from torchvision.models. The model's weights were updated using torch.load and subsequently, the model was set to evaluation mode to facilitate feature extraction. Data transformations were defined to standardize the input images. Each image was resized to 224x224 pixels and converted into a tensor. Data loaders were created for both training and testing datasets , with a batch size of 32. The feature extraction process was encapsulated in a specific function. This function is called twice, once for the training data loader and once for the test data loader. This extracts features from each image in the respective datasets and saves them along with their corresponding labels to separate CSV (comma-separated values) files. Feature extraction was performed in a no-gradient context to optimize performance. The parameters used are mention in Table <ref>.




We extracted features from the ViT in two scenarios: firstly, using the pre-trained weights of ViT-B/16, which were trained on ImageNet; and secondly, by training the ViT from scratch as detailed in Section 3, starting with the ImageNet initial weights up to epoch 8. These extracted features were separately fed as inputs to an SVM to evaluate performance in both scenarios. We conducted 5-fold cross-validation to determine the ROC curve, mean AUC, and mean accuracy. Additionally, we examined both cases using the specified dataset provided in the original data repository.




Table <ref> shows the accuracy and AUC (Area Under the Curve) of approaches: ViT+SVM (Pre-trained Weights) and ViT+SVM (Optimal Weights). The models were evaluated on a specified test set and using 5-fold cross-validation. ViT+SVM (Optimal Weights) outperforms ViT+SVM (Pre-trained Weights) in terms of both mean accuracy (0.95 vs. 0.88) and mean AUC (0.96 vs. 0.95) with 5-fold cross-validation. Though, both the approaches achieved the same accuracy on the specified test set, but ViT+SVM (Optimal Weights) generalizes better to unseen data and suggests potentially less over-fitting as evidenced by its higher mean accuracy and AUC with 5-fold cross-validation. Upon examining the mean RoC curves given in Fig <ref>, it is evident that ViT+SVM (Optimal Weights) is performing consistently as the folds have a small variance in AUC (0.01) and each fold has a high AUC (> 0.95). On the other hand, ViT+SVM (Pre-trained Weights) have a mean AUC of 0.95 and mean variance of 0.02.





















 §.§ Comparison of the Models




Table <ref> presents the performance metrics of various models on a specific dataset. The comparison shows that each model has its strengths and is suitable for different scenarios.

Performance on Specified test set:
    ViT with Transfer Learning consistently performs best across most metrics, making it the most balanced and robust model for this task.
    ViT+SVM (Optimal Weights) excels in precision for the stressed class and recall for the healthy class, making it ideal where high precision in detecting stressed instances and high recall in identifying healthy instances are crucial.
    ViT (Pre-Trained Weights) provides a balanced performance but does not outperform the other models in any specific metric.



K-fold cross validation: 
The ViT with Transfer Learning model achieved an average accuracy of 86.8%, with F1-scores of 85.6% for the "Healthy" class and 88.8% for the "Stressed" class. The model demonstrated consistent performance across the five folds, although there was a slight dip in accuracy in the fifth fold (83%).

The ViT with SVM model achieved a significantly higher average accuracy of 95%, with F1-scores of 93.6% for Class 0.0 and 95.6% for Class 1.0. This model consistently performed well across all folds, with the lowest accuracy being 92.7% and the highest reaching 95.8%.




The comparative results, as shown in Fig <ref> indicate that the ViT with SVM model outperforms the ViT with Transfer Learning model in terms of both accuracy and F1-scores. The average accuracy of the ViT with SVM model is substantially higher (95% compared to 86.8%). Furthermore, the F1-scores for both classes are also higher in the ViT with SVM model, indicating better classification performance for both "Healthy" and "Stressed" classes.

The five-fold accuracies for the ViT with SVM model show less variability and higher consistency compared to the ViT with Transfer Learning model, suggesting that the integration of SVM helps in stabilizing the model's performance across different data splits.

In conclusion, integrating Support Vector Machine with Vision Transformer (ViT) significantly enhances the model's performance compared to using transfer learning alone. The ViT with SVM model demonstrates superior accuracy and F1-scores, making it a more reliable choice for classification tasks. The higher consistency across folds further reinforces the robustness of the ViT with SVM model. This comparative analysis highlights the benefits of combining advanced deep learning architectures with classical machine learning techniques for improved performance in identifying drought stress.


§ CONCLUSION

Drought stress represents a severe threat to crop yield and quality, disrupting normal plant growth and survival rates. Detecting early signs of drought stress is crucial for effective crop management and intervention. Traditional methods, primarily reliant on Convolutional Neural Networks (CNNs), have made significant strides in capturing spatial hierarchies in image data. However, Vision Transformers (ViTs) offer a compelling alternative by leveraging self-attention mechanisms to capture long-range dependencies and complex spatial relationships, thus enhancing the detection of subtle drought stress indicators.

Our study successfully addressed the challenge of stress identification using smaller datasets by harnessing the feature extraction capabilities of Vision Transformers combined with the classification prowess of SVM. Unlike conventional CNN architectures, Vision Transformers utilize self-attention mechanisms to effectively capture relationships across different parts of the image, enabling them to model long-range dependencies, which is particularly advantageous for complex datasets where traditional CNNs may struggle to capture global context. Moreover, Vision Transformers can handle images of varying resolutions without necessitating architectural modifications, enhancing their flexibility in handling diverse datasets. Our framework also emphasized explainability by generating attention maps, which provide insights into the model’s focus areas within the images, thereby offering transparency in its decision-making process.

However, our comparative analysis revealed that integrating SVM with ViT significantly enhances classification performance. The ViT+SVM model achieved a higher average accuracy of 95%, with F1-scores of 93.6% for stressed images and 95.6% for healthy images, outperforming the ViT with transfer learning approach. This integration not only improved accuracy but also demonstrated greater consistency across different data splits, highlighting the robustness of the ViT+SVM model.

In summary, our study illustrates the superior performance of combining Vision Transformers with Support Vector Machines over traditional transfer learning methods. This approach provides a more reliable and accurate solution for drought stress identification, while the attention maps offer valuable insights into the model’s decision-making process. The findings affirm the potential of advanced deep learning and machine learning techniques in enhancing agricultural practices and decision-making, paving the way for improved crop management strategies.
 

unsrt







 Paper Title: A Simple Low-bit Quantization Framework for Video Snapshot Compressive Imaging.





 



A Simple Low-bit Quantization Framework for Video SCI






M. Cao et al.




Zhejiang University, Hangzhou 310058, Zhejiang, China School of Engineering, Westlake University, Hangzhou, 310030, Zhejiang, China

{caomiao,wanglishun,wanghuan,xyuan}@westlake.edu.cn

A Simple Low-bit Quantization Framework for Video Snapshot Compressive Imaging
    Miao Cao1,2 Lishun Wang2Huan Wang2Xin Yuan2()
    August 5, 2024
==============================================================================




§ ABSTRACT 
  Video Snapshot Compressive Imaging (SCI) aims to use a low-speed 2D camera to capture high-speed scene as snapshot compressed measurements, followed by a reconstruction algorithm to reconstruct the high-speed video frames.  
  
  State-of-the-art (SOTA) deep learning-based algorithms have achieved impressive performance, yet with heavy computational workload. 
  Network quantization is a promising way to reduce computational cost. However, a direct low-bit quantization will bring large performance drop. 
  To address this challenge, in this paper, we propose a simple low-bit quantization framework (dubbed Q-SCI) for the end-to-end deep learning-based video SCI reconstruction methods which usually consist of a feature extraction, feature enhancement, and video reconstruction module. 
  Specifically, we first design a high-quality feature extraction module and a precise video reconstruction module to extract and propagate high-quality features in the low-bit quantized model.
  
  In addition, to alleviate the information distortion of the Transformer branch in the quantized feature enhancement module, we introduce a shift operation on the query and key distributions to further bridge the performance gap. Comprehensive experimental results manifest that our Q-SCI framework can achieve superior performance, , 4-bit quantized EfficientSCI-S derived by our Q-SCI framework can theoretically accelerate the real-valued EfficientSCI-S by 7.8× with only 2.3% performance gap on the simulation testing datasets. Code is available at <https://github.com/mcao92/QuantizedSCI>. 
  





§ INTRODUCTION


Recently, video Snapshot Compressive Imaging (SCI) has attracted much attention because it can capture high-speed scenes using a low-speed 2D camera with low bandwidth. There are two main steps in the video SCI system: hardware encoding (shown in Fig. <ref>(a)) and software decoding(shown in Fig. <ref>(b)) <cit.>. In the hardware encoding process, we first modulate the dynamic scene with different masks, and then the modulated scene is compressed into a series of snapshot measurements, which are finally captured by a low-cost 2D camera. In the software decoding stage, the captured snapshot measurements and the corresponding modulation masks are fed into a reconstruction algorithm to recover the desired video frames. 



On the one hand, many successful video SCI hardware encoders <cit.> have been built. On the other hand, state-of-the-art (SOTA) reconstruction algorithms <cit.>, mostly based on deep neural networks, have outperformed the traditional model-based methods <cit.> by a large margin not only in the reconstruction quality but also in the running speed. Thus, it seems that the current video SCI systems can already be used in practical applications. Unfortunately, existing deep learning-based reconstruction algorithms usually need a large number of parameters and floating point operations (FLOPs) to achieve satisfactory accuracy. EfficientSCI <cit.> is a significant exploration to design an efficient video SCI reconstruction algorithm, which has 3.78 MB parameters and requires 563.87 GFLOPs to reconstruct the video frames with 35.51 dB on the simulation testing datasets. Our ultimate goal is to deploy video SCI reconstruction methods on AI chips, thus to forming a real end-to-end system integrating capture and reconstruction. However, it is still challenging to deploy previous reconstruction algorithms on resource-limited devices due to their high computational cost. This motivates us to move one step further to reduce the computational workload of the video SCI reconstruction algorithms while preserving the model performance as much as possible.

Substantial efforts have been made to compress and accelerate neural networks for efficient online inference. These methods can be divided into the following categories: network quantization <cit.>, network pruning <cit.>, knowledge distillation <cit.>, and compact network design <cit.>. Among them, network quantization is suitable for deployment on resource-limited platforms because it can reduce the bit-width of network parameters and
activations for efficient inference. However, directly applying low-bit quantization to video SCI reconstruction algorithms may encounter the following issues: i) Directly quantizing previous reconstruction methods into low-bit will lead to a large performance drop from their full-precision counterpart. ii) There exists a distortion of the query and key distributions in the low-bit quantized vision Transformer.    

Bearing the above concerns in mind, in this paper, we propose a simple low-bit quantization framework, dubbed Q-SCI, for the end-to-end deep learning-based video SCI reconstruction methods which usually consist of a feature extraction module, a feature enhancement module, and a video reconstruction module as shown in Fig. <ref>(c). Specifically, we first conduct extensive empirical analysis and identify that the large performance drop in low-bit quantized model mainly comes from the information loss of low-bit quantized features. Therefore, we design a high-quality feature extraction module and a precise video reconstruction module which can extract and propagate high-quality features through the low-bit quantized network. In addition, to alleviate the information distortion of the Transformer branch in the quantized feature enhancement module, we introduce a shift operation on the query and key distributions to further bridge the performance gap. 
Finally, following previous works <cit.> which are built and evaluated on existing backbones. We choose EfficientSCI-S as the full-precision backbone model for our proposed Q-SCI because EfficientSCI-S can achieve SOTA reconstruction quality with much cheaper computational cost. As shown in Fig. <ref>(d), compared with previous methods, the proposed Q-SCI framework can achieve comparable reconstruction quality with much fewer model parameters and FLOPs.

The main contributions of our paper can be summarized as follows:  

    
  * We propose a simple low-bit quantization framework for video SCI reconstruction. 
    To our best knowledge, this is the first network quantization framework in the video SCI reconstruction task.
    
  * We design a high-quality feature extraction module and a precise video reconstruction module to better extract and propagate high-quality features. Furthermore, the generalization ability of the proposed Q-SCI framework is verified on different end-to-end video SCI reconstruction methods.  
    
  * A shift operation on the query and key distributions is introduced to alleviate the information distortion in the quantized Transformer branch.
    
  * Experimental results on the simulated and real datasets demonstrate that the low-bit quantized models derived by our proposed Q-SCI framework can achieve superior performance with much cheaper computational cost. 




§ RELATED WORK





 §.§ Video SCI Reconstruction Algorithms

The existing video SCI reconstruction methods can be divided into two categories: traditional model-based methods and deep learning-based ones.
The traditional model-based methods formulate the video SCI reconstruction process as an optimization problem with a regularization term such as total variation <cit.> and Gaussian mixture model <cit.>, and solve it via iterative algorithms. The major drawback of these model-based methods is that one needs to perform time-consuming iterative optimization. 
In order to improve the running speed, Yuan et al. develop a plug-and-play (PnP) framework, in which a pre-trained denoising model is plugged into every iteration of the optimization process <cit.>. However, they still take a long time on reconstructing large-scale data. 

Recently, more and more researchers begin to utilize deep neural networks in the video SCI reconstruction task.

For example, BIRNAT <cit.> uses a bidirectional recurrent neural network to exploit the temporal correlation.


To save model training memory, RevSCI <cit.> builds an end-to-end 3D convolutional neural network with reversible structure for video SCI reconstruction. More recently, Wang et al. build the first Transformer-based reconstruction network <cit.> with space-time factorization and local self-attention mechanism. After that, Wang et al. design an efficient reconstruction network <cit.> based on dense connections and space-time factorization.

Combining the idea of model-based and deep learning-based methods, researchers propose the deep unfolding networks <cit.>. 
For example, Zheng et al. first introduce the uncertainty estimation mechanism into the deep unfolding framework <cit.>.  

Although the deep learning-based models have achieved impressive results, it is still challenging to deploy them on the resource-limited devices due to their high computational workload. In this paper, we mainly focus on developing light-weight video SCI reconstruction algorithm empowered by network quantization.  



 §.§ Network Quantization
 



Due to its impressive effectiveness in computational compression, network quantization has
been widely applied in high-level vision <cit.> and low-level vision tasks <cit.>. In the high-level vision tasks, Liu et al. develop a dynamic quantization scheme <cit.> to decide the optimal bit-widths for each image. GPUSQ-ViT <cit.> is a compression scheme which maximally utilize the GPU-friendly 2:4 fine-grained structured sparsity and quantization. Q-ViT <cit.> presents the first quantization-aware training framework for accurate and low-bit vision Transformer. ReActNet <cit.> learns the distribution shape and shift to improve the performance of the CNN-based binary network. Li et al. propose a novel post-training quantization framework specifically tailored towards the unique multi-timestep pipeline and model architecture of the diffusion models <cit.>. Li et al. develop a fully quantized network for object detection <cit.>. Considering the low-level vision tasks, PAMS <cit.> learns the quantization intervals of different layers to adapt to vastly distinct feature distributions in the super-resolution (SR) networks. Hong et al. propose CADyQ <cit.> to assign different bit-width for each patch and layer for image SR. Xia et
al. design a binarized convolution unit BBCU <cit.> for image restoration-tasks, , image denoising and JPEG
compression artifact reduction. FQSR <cit.> jointly optimize efficiency and accuracy on the image SR task. Yet, the potential of network quantization for the video SCI reconstruction task has not been explored.



§ PREREQUISITES OF NETWORK QUANTIZATION


In this section, we briefly introduce the network quantization architecture used in our paper. First, given the activation x and weight w, we introduce a general asymmetric activation quantization and symmetric weight quantization scheme, expressed as follows:

    Q_a(x)    =   ⌊clip{(x-z)/α_x,-Q_n^x,Q_p^x}⌉,
    x̂   =    Q_a(x) ×α_x + z,
    
        Q_ w( w)    =   ⌊clip{ w/α_ w,-Q_n^ w,Q_p^ w}⌉,
    ŵ   =    Q_ w( w) ×α_ w,

where clip{x,m_1,m_2} constrains x into [m_1,m_2] by setting x as m_1 when x < m_1 and m_2 when x > m_2. ⌊ x ⌉ rounds x to the nearest integer. α is a scale factor, which can divide the entire range of the input into uniform partitions. z is a zero-point, which can shift the quantized distribution into the a specific range in the asymmetric quantization. During training, α and z are initialized as 1 and 0, and then been optimized as network parameters.  
Following this, if we quantize the activation x to a bits and weight w to b bits, we have:

    Q_n^x = 2^a-1,Q_p^x = 2^a-1-1,
    
        Q_n^ w = 2^b-1,Q_p^ w = 2^b-1-1.


In this way, the forward and back propagation of network quantization can be formulated as Q-Linear(x)=x̂·ŵ=α_xα_ w((Q_a(x)+z/α_x) ⊗ Q_ w( w) and Eq. (<ref>) respectively.



    ∂𝒥/∂ x =  ∂𝒥/∂x̂∂x̂/∂ x = {[                 ∂𝒥/∂x̂ if x ∈ [-Q_n^x, Q_p^x];                      0              otherwise;                        ]. ,  
       ∂𝒥/∂ w =  ∂𝒥/∂ x∂ x/∂ŵ∂ŵ/∂ w = {[            ∂𝒥/∂ x∂ x/∂ŵ if w∈ [-Q_n^ w, Q_p^ w];                       0               otherwise;                         ]. ,

where 𝒥 denotes the loss function, Q(·) is used in the forward process while the straight-through estimator <cit.> is adopted to retain the derivation of gradient in the backward propagation process, and ⊗ denotes the efficient bit-wise matrix multiplication operation. 

Recent methods <cit.> begin to explore the use of Transformer in the end-to-end video SCI reconstruction methods. Thus, we define the quantization process of the self-attention module in the following way. First, we denote the quantized computation on the query , key  and value  as =Q-Linear_q(x), =Q-Linear_k(x), =Q-Linear_v(x), where Q-Linear_q, Q-Linear_k and Q-Linear_v represent the quantized linear layers for ,  and , respectively. Then, the attention weight can be expressed as:

    =1/√(d)(Q_a()⊗ Q_a()^⊤), 
    
        Q_   =Q_a(softmax()).






§ PROPOSED METHODS



In this section, we propose an accurate and fully quantized framework for the end-to-end video SCI reconstruction algorithms. Firstly, let us give a brief introduction of the previous SOTA EfficientSCI which is used as the full-precision backbone model for our proposed Q-SCI framework. The ResDNet module with N residual style ResDNet blocks is used for feature enhancement in EfficientSCI. In each ResDNet block, we put several CFormer blocks with temporal Transformer branch. Please refer to <cit.> and the supplementary material for more details about EfficientSCI and the video SCI mathematical model. Then, we conduct extensive empirical analysis on the severe performance drop of the low-bit quantized model. Following this, we design a high-quality feature extraction module, a precise video reconstruction module and a shifted Transformer branch to improve the performance of the low-bit quantized model. Afterwards, we customize four different quantized network variants based on our Q-SCI framework. Finally, the loss function is defined for the training process.    






 §.§ Performance Analysis on Low-bit Quantization



In the experiments, compared with the full-precision backbone model EfficientSCI-S, we observe a large performance drop (4.11 dB) in the 4-bit quantized baseline model which directly quantize each network layer of EfficientSCI-S into 4-bit. Thus, we try to detect the source of this severe performance drop in the experiments. First, based on the 8-bit quantized baseline model which directly quantize each network layer into 8-bit, we set the bit-width of one module (the feature extraction module, ResDNet module or video reconstruction module) to 4-bit while maintaining that of the other two modules. As we can see in Tab. <ref>, there exists a 2.22 dB performance drop when we quantize the feature extraction module into 4-bit. However, the performance drop is relatively small with 0.52 dB and 0.49 dB for the 4-bit quantized ResDNet module and video reconstruction module. Therefore, we infer that the severe performance drop mainly comes from the low-bit quantized feature extraction module. Stepping forward, we visualize the output feature map of the feature extraction module to further verify our assumption. As shown in Fig. <ref>, compared with full-precision model (Fig. <ref>(a)), the output feature map of the feature extraction module in 4-bit quantized baseline model (Fig. <ref>(b)) degrades severely. 
Therefore, we can conclude that the large performance drop of the low-bit quantized model mainly comes from the low-bit quantized feature extraction module.  





 §.§ Proposed Quantization Framework




  §.§.§ Quantized Feature Extraction Module



We have noticed from the empirical analysis (detailed in Sec. <ref>) that high-quality features play an important role in the performance improvement of the low-bit quantized model. Therefore, in this section, we present a simple and high-quality feature extraction module to improve the performance of the low-bit quantized  reconstruction methods. First, shortcut connections are widely used to propagate high-quality information. However, there exists a dimension mismatch problem during feature reshaping in the previous feature extraction module (shown in the red box of Fig. <ref>(b)). Thus, we propose to adopt several 1×1×1 convolution layers as shortcut connections. Note that a pixel shuffle operation is performed before the last shortcut connected convolution layer to calibrate the spatial size mismatch. Then, we set the bit-width of the shortcut connected convolution layers to 8-bit to precisely propagate high-quality features. In this way, the proposed high-quality feature extraction module is capable of extracting and propagating high-quality features. Finally, when we compare Fig. <ref>(b) and Fig. <ref>(c), obvious quality improvement of the output feature map can be observed, which verifies the effectiveness of our proposed feature extraction module.





  §.§.§ Quantized Transformer Branch


As shown in Q-ViT <cit.>, there exists a distortion of the query and key distributions in the low-bit quantized vision Transformer branch, which brings performance drop. Therefore, we plot the distribution of the query activation in Fig. <ref> to better measure the distribution distortion of the low-bit quantized Transformer branch. We can see from Fig. <ref>(a) and Fig. <ref>(b) that the mean of the query distribution in the first CFormer block of the first ResDNet block between 8-bit quantized baseline model and its full-precision counterpart is 1.207 (1.056 v.s. -0.151). Unfortunately, we cannot directly adapt Q-ViT into the reconstruction methods due to the absence of bell-shaped distributions as shown in Fig. <ref>(b). Therefore, we adapt a shift operation into the Transformer branch to rectify the query and key distributions. Specifically, given the query  and key , we  define a shift operation as =+β_ and =+β_, where β_ and β_ denote the learnable shift bias of the distributions. Finally, comparing Fig. <ref>(a) with Fig. <ref>(c), we observe that the mean of the query distribution in the first CFormer block of the first ResDNet block of 8-bit quantized baseline model equipped with our proposed shifted Transformer branch stays close to that of its full-precision counterpart (1.056 v.s. 0.984), which verifies the effectiveness of the proposed approach. Due to the space limitation, please refer to the supplementary material for more distribution histograms of different CFormer blocks.   



  §.§.§ Quantized Video Reconstruction Module


In this section, to better propagate high-quality features through the network, we design a precise video reconstruction module to further bridge the performance gap with the following considerations: i) As demonstrated in Sec. <ref>, the proposed high-quality feature extraction module can extract and then propagate high-quality features to the ResDNet module. ii) As illustrated in Fig. <ref>(c), there are already some skip connections in the ResDNet module of EfficientSCI, which are able to propagate the high-quality features to the video reconstruction module. Therefore, we want to further propagate the high-quality features to the end of the network. As the video reconstruction module is also stacked with some convolution layers, we present a similar design with Sec. <ref>. Specifically, as shown in Fig. <ref>(c), several 1×1×1 shortcut connected convolution layers are added above the previous video reconstruction module (shown in the red box of Fig. <ref>(c)). Finally, we also set the bit-width of the shortcut connected convolution layers as 8-bit to ensure precise feature propagation.    



 §.§ Architecture and Variants


In this section, we choose EfficientSCI-S as a full-precision backbone model to customize quantized networks due to its superior efficiency performance. However, this does not mean that our proposed Q-SCI framework is equal to a low-bit quantized EfficientSCI-S, because our Q-SCI framework can also be applied on the other end-to-end video SCI reconstruction methods (See Tab. <ref>). 

After that, we customize four quantized network variants including Q-SCI (8-bit), Q-SCI (4-bit), Q-SCI (3-bit), and Q-SCI (2-bit). The network quantization settings are as follows: i) In the Q-SCI (8-bit) network, we directly set the bit-width of all the network layers to 8-bit and then adopt the proposed shifted Transformer branch. ii) In the Q-SCI (4-bit), Q-SCI (3-bit), and Q-SCI (2-bit) network, we first set the bit-width of all the network layers to 4-bit, 3-bit, and 2-bit, respectively. Then, to ensure high-quality reconstruction results, we use our proposed high-quality feature extraction module, precise video reconstruction module and shifted Transformer branch. Note that previous work <cit.> also sets a part of the network layers as 8-bit to ensure high performance. Additionally, since most layers of the low-bit quantized reconstruction networks are set to be (2, 3, or 4)-bit, it can still be considered as a (2, 3, or 4)-bit quantized model.



 §.§ Loss Function


Our proposed low-bit quantized network takes the measurement () and the corresponding masks ({_t}_t=1^T) as inputs, and then generates the reconstructed video frames ({_t}_t=1^T∈ℝ^n_x×n_y). To train the network, we choose the mean squared error (MSE) as our loss function,



    ℒ_MSE=  1/Tn_xn_y∑_t = 1^T‖_t - _t ‖_2^2,

where {_t}_t=1^T denotes the ground truth.



§ EXPERIMENTAL RESULTS






Video SCI Hardware: The optical setup of the real video SCI system is shown in Fig. <ref>. The hardware encoding process of the video SCI system is as follows: First, the reflected light from the dynamic scene is imaged onto the surface of the digital micromirror device (DMD) via a camera lens (Sigma, 17-50/2.8, EX DC OS HSM) and the first relay lens (Coolens, WWK10-110-111). Following this, the projected high-speed scene on the DMD (TI, 2560 × 1600 pixels with 7.6 μ m pixel pitch) is modulated by the pre-stored random binary masks. Finally, the encoded compressed measurement is projected onto the surface of a low-cost CCD camera (Basler acA1920, 1920 × 1200 pixels with 4.8 μ m pitch pixel) with the second relay lens (Coolens, 
WWK066-110-110), which is captured in a single exposure time. Note that the CCD camera works at 50 FPS when capturing real data and the compression ratio of the video SCI system is Cr, then the equivalent sampling rate of our real video SCI system is 50×Cr FPS.



Datasets: Three types of dataset are used in this paper: training dataset, simulated testing dataset, and real testing dataset. i) First, following BIRNAT <cit.>, we choose the  dataset <cit.> with spatial resolution 480×894 (480p) as the training dataset. ii) We then test our Q-SCI framework on six simulated testing datasets, including , , , , , and  (with spatial resolution 256×256 and compression ratio is 8). For the simulated testing datasets, average Peak Signal-to-Noise-Ratio (PSNR) and average Structured Similarity Index Metrics (SSIM) <cit.> are used as the evaluation metrics of reconstruction quality. Besides, following previous works <cit.>, model size (Params) and number of operations (OPs) are used as the evaluation metrics of algorithm efficiency. iii) Finally, we test the proposed Q-SCI framework on two real testing datasets, including  and  (with spatial resolution 512×512 and compression ratio is 10) captured by our real built video SCI system described in the “Video SCI Hardware” section. Moreover, please refer to the supplementary material for more details about the real data capture process.

Training Process: We build our models with PyTorch, and conduct training on 4 NVIDIA RTX 3090 GPUs. First, regular data augmentation operations such as random cropping, random flipping and random scaling are performed on the training dataset. Then, we initialize our Q-SCI with the full-precision EfficientSCI-S and adopt the Adam <cit.> optimizer to optimize the model with an initial learning rate of 0.0001. After iterating for 100 epochs on the training data with a 128×128 spatial size and 20 epochs on the training data with a 256×256 spatial size, we adjust the learning rate to 0.00001 and continue to iterate for 20 epochs on the training data with a 256×256 spatial size. 







 §.§ Results on Simulated Testing Datasets


In this section, we compare our Q-SCI networks with seven deep learning-based video SCI reconstruction methods (including MetaSCI, BIRNAT, RevSCI, Dense3D-Unfolding, ELP-Unfolding, STFormer-S, and EfficientSCI-S) on six simulated testing datasets. As shown in Tab. <ref>, the proposed Q-SCI networks can achieve comparable reconstruction quality with much fewer parameters and cheaper computational cost. Specifically, i) The proposed Q-SCI (2-bit) model achieves comparable reconstruction quality with MetaSCI, while the OPs is reduced by about 2×. ii) Our proposed Q-SCI (3-bit) model achieves 0.31 dB in PSNR higher than BIRNAT, and the OPs is reduced by about 10.4×. iii) The proposed Q-SCI (4-bit) model achieves 1.38 dB and 0.77 dB in PSNR higher than BIRNAT and RevSCI, and the OPs is reduced by about 5.4× and 10.6×. iv) Our Q-SCI (8-bit) model can achieve 0.31 dB and 0.16 dB in PSNR higher than Dense3D-Unfolding and ELP-Unfolding, and the OPs is reduced by about 28.2× and 32.9×. v) Our proposed Q-SCI (8-bit) model achieves comparable performance with EfficientSCI-S, while the OPs is reduced by about 4×. 

For visualization purposes, we also visualize some reconstructed video frames in Fig. <ref>, where we can see from the zooming areas in each selected video frame that our proposed Q-SCI (4-bit) model can provide much clearer reconstructed images than BIRNAT, RevSCI, and STFormer-S. Further, the proposed Q-SCI (8-bit) model can provide comparable high-quality reconstructed images with EfficientSCI-S. Especially for the , , and  scenes, we can observe sharp edges and more details in the reconstructed frames of our proposed models, which verify the superior performance of the proposed Q-SCI quantization framework. 


Moreover, comparison with previous SOTA quantization method Q-ViT is also conducted.

Note that, we adapt both Q-ViT and Q-SCI into EfficientSCI-S for a fair comparison. The computational complexity and reconstruction quality are given in Tab. <ref>. We can see from Tab. <ref> that Q-SCI (8-bit) outperforms Q-ViT (8-bit) by about 0.4 dB with less computational cost. 


Finally, we adapt the proposed Q-SCI framework into previous SOTA end-to-end reconstruction method STFormer-S.

As shown in Tab. <ref>, compared with 4-bit quantized baseline model, adapting our proposed high-quality feature extraction module will bring a 3.23 dB reconstruction quality improvement for STFormer-S. Then, further adapting the proposed precise video reconstruction module can lead to a 0.25 dB reconstruction quality improvement for STFormer-S. Therefore, the proposed Q-SCI framework can well generalize to other end-to-end video SCI reconstruction methods. 







 §.§ Results on Real Testing Datasets


In this section, we test the proposed Q-SCI (8-bit) and Q-SCI (4-bit) models on the real testing datasets. Due to the uncertain noises of the real video SCI system, it is more difficult to reconstruct real measurements. As shown in Fig. <ref>, our proposed Q-SCI (4-bit) model can reconstruct clearer borders than previous real-valued models including GAP-TV, PnP-FFDNet, and BIRNAT. Moreover, we can clearly recognize the letters on the  on the reconstructed frames of the proposed Q-SCI (4-bit) and Q-SCI (8-bit). Finally, our proposed Q-SCI (8-bit) model can present sharper boarders and easier-recognized letters on the  when compared with the real-valued STFormer-S and EfficientSCI-S.  



 §.§ Ablation Study


Here, we conduct a break-down ablation towards higher performance. We first build a baseline method by directly quantizing each layer of EfficientSCI-S into 4-bit. As shown in Tab. <ref>, the baseline model yields 31.40 dB in PSNR and 0.931 in SSIM. When we adopt the shifted Transformer branch, the model achieves a 0.53 dB improvement (which is a moderately large improvement in the video SCI reconstruction task) with neglect computational burden. After that, we apply the high-quality feature extraction module and the precise video reconstruction module successively. The performance improvement are 2.35 dB and 0.43 dB, respectively. Finally, we obtain the Q-SCI (4-bit) model with 34.71 dB which is a totally 3.31 dB improvement against the 4-bit quantized baseline model, while the computational cost only grows by 3.14%. 
These experimental results verify the effectiveness of the proposed module designs of our Q-SCI framework.



§ CONCLUSION AND FUTURE WORK



In this paper, we propose the first low-bit quantization framework for the end-to-end video SCI reconstruction methods, which is simple and effective.
Specifically, we design a high-quality feature extraction module and a precise video reconstruction module to extract and propagate high-quality features in the low-bit quantized model. Additionally, we introduce a shift operation on the query and key distributions of the low-bit quantized Transformer branch to further bridge the performance gap. Finally, we verify that our proposed low-bit quantization framework can generalize well to different end-to-end reconstruction methods (including EfficientSCI and STFormer). 

While our proposed Q-SCI can largely reduce computational cost with small performance drop, there is still room for further improvement, which we discuss as follows. First of all, deploying Q-SCI on the resource-limited devices (, smart phone and autonomous vehicle)  requires extensive engineering and hardware support, making it a challenging task. Secondly, although Q-SCI is a well-designed network quantization framework for video SCI, a novel mobile-friendly reconstruction network will be more suitable than previous reconstruction methods for the quantized video SCI reconstruction task. Finally, algorithm-hardware codesign should be considered to better fit the resource-limited devices and the video SCI hardware encoding system.  

Regarding the future work, we can further explore other model compression methods including network pruning <cit.> and knowledge distillation <cit.> to further optimize the efficiency performance. Moreover, we expect to extend the proposed Q-SCI framework to various imaging systems such as high-speed imaging <cit.>, hyperspectral imaging <cit.>, depth of field imaging <cit.>, and single-pixel imaging <cit.>.



§ ACKNOWLEDGEMENTS

This work was supported by the National Natural Science Foundation of China (grant number 62271414), Zhejiang Provincial Distinguished Young Scientist Foundation (grant number LR23F010001), Zhejiang “Pioneer” and “Leading Goose” R&D Program  (grant number 2024SDXHDX0006, 2024C03182), the Key
Project of Westlake Institute for Optoelectronics (grant number 2023GD007), and
the 2023 International Sci-tech Cooperation Projects under the purview of the “Innovation Yongjiang 2035” Key R&D Program  (grant number 2024Z126). 






splncs04



 Paper Title: A Simple Background Augmentation Method for Object Detection with Diffusion Model.





 



Background Augmentation with Diffusion Model






Y. Li, X. Dong, C. Chen, W. Zhuang, L. Lyu.




Sony AI

{yuhang.li, xin.dong, chen.chen, weiming.zhuang, lingjuan.lyu}@sony.com



A Simple Background Augmentation Method for Object Detection with Diffusion Model
    Yuhang Li10000-0002-6444-7253Xin Dong1 Chen Chen1 Weiming Zhuang1 Lingjuan Lyu1Corresponding Author. Work done during Yuhang Li’s internship at Sony AI.
    August 5, 2024
============================================================================================================================================================




§ ABSTRACT 
In computer vision, it is well-known that a lack of data diversity will impair model performance.
In this study, we address the challenges of enhancing the dataset diversity problem in order to benefit various downstream tasks such as object detection and instance segmentation. We propose a simple yet effective data augmentation approach by leveraging advancements in generative models, specifically text-to-image synthesis technologies like Stable Diffusion. Our method focuses on generating variations of labeled real images, utilizing generative object and background augmentation via inpainting to augment existing training data without the need for additional annotations. We find that background augmentation, in particular, significantly improves the models' robustness and generalization capabilities. 
We also investigate how to adjust the prompt and mask to ensure the generated content 
comply with the existing annotations. 
The efficacy of our augmentation techniques is validated through comprehensive evaluations of the COCO dataset and several other key object detection benchmarks, demonstrating notable enhancements in model performance across diverse scenarios. This approach offers a promising solution to the challenges of dataset enhancement, contributing to the development of more accurate and robust computer vision models.
  





§ INTRODUCTION


The quest for robust and accurate object detection <cit.> and instance segmentation <cit.> models is a cornerstone for contemporary computer vision research. These models underpin a variety of applications, from autonomous vehicles <cit.> to medical image analysis <cit.>, necessitating high levels of precision and reliability. Central to the development of these models is the availability of large, diverse, and accurately annotated datasets. However, creating such datasets is fraught with challenges, including the high costs of data collection and annotation, privacy concerns, and the potential for dataset bias. For example, collecting images of urban driving scenes requires physical car infrastructure, and labeling a segmentation annotation for a single urban image in Cityscapescan takes up to 60 minutes  <cit.>. 
Therefore, it is necessary to navigate new solutions to address these limitations, data augmentations and low-cost annotations. 
 
Recent advancements in generative models, particularly in the domain of text-to-image synthesis <cit.> (e.g., Stable Diffusion <cit.>), offer a promising avenue for addressing aforementioned challenges. 
While Stable Diffusion and similar models can create visually compelling images, generating synthetic images that are accompanied by accurate and detailed object annotations remains a formidable challenge. This complexity stems from the need for precise alignment between the generated visual content and its corresponding annotations. Although several works have 
generated the synthetic data and instances annotations simultaneously <cit.>, they all require finetuning the existing text-to-image on the current dataset and generate a significant amount of synthetic data for pre-training, which increases training complexity significantly. 


To this end, we explore generating variants of labeled real images for object detection and instance segmentation while reusing their annotations, in a way that augments the existing training data. Specifically, we explore generative object augmentation and background augmentation through inpainting, which allows for the enhancement of objects and backgrounds respectively within the dataset. 
Notably, we discover that background augmentation boosts the model performance, which suggests that modifying the 

background context of images plays a crucial role in improving model robustness and generalization.


We rigorously evaluate our augmentation methods on the MS COCO dataset <cit.> and PASCAL VOC <cit.> for both detection and segmentation tasks across multiple architectures. This comprehensive evaluation demonstrates the effectiveness of our proposed augmentation techniques in improving the performance of object detection and instance segmentation models under various challenging scenarios. We summarize our contributions as follows:

    
  * We introduce a simple data augmentation framework for object detection with the text-to-image model without finetuning. We find that augmenting the background brings more advantages than augmenting objects. 
    
  * We carefully craft the text prompt, mask region, and diffusion steps to ensure the generated content is effective for data augmentation purposes. The ablation studies show that they are essential to achieve superior performance.
    
  * We evaluate the background augmentation policies on various tasks including object detection and instance segmentation, significantly improving their performance. For example, our background augmentation increases mAP up to 5.3% when using only 10% of the COCO training data. 
    






§ RELATED WORK


This section provides literature reviews of object detection, data augmentation, and text-to-image generation for visual tasks.

Object Detection. 
Object detection aims to simultaneously predict the category and corresponding bounding box
for the objects in the images. Generally, object detectors <cit.> are trained on a substantial amount of
training data with bounding box annotations and can only
recognize a predetermined set of categories present in the
training data, MS COCO dataset <cit.>. 
Both backbone architecture and the detection neck/head framework are important for object detection performance. For backbone, Vision Transformer <cit.> has emerged as an effective backbone alternative compared to convolutional neural networks <cit.>. These transformer-based models leverage self-attention mechanisms to capture global dependencies within the image. 
Moreover, the development of efficient and scalable detectors, such as EfficientDet <cit.>, further underscores the importance of optimizing both accuracy and computational efficiency. 


Data Augmentation. 
Data augmentation <cit.> plays an indispensable role in the deep learning model, as it forces the model to learn invariant features, and thus helps generalization.
The data augmentation is applied in many areas of vision tasks including object recognition <cit.>, video perception <cit.>, and semantic segmentation <cit.>. 
Apart from learning invariant features, data augmentation also has other specific applications in deep learning. For example, adversarial training <cit.> leverages data augmentation to create adversarial samples and thereby improves the adversarial robustness of the model.
In terms of data augmentation in object detection, Zoph et.al <cit.> propose reinforcement learning to search for complex augmentation policies including color transformations, geometric transformations, and bounding box transformations. CutMix <cit.> proposes a spatial mixing strategy for augmenting the detection results. RandAugment <cit.> uses a reduced search space for searching both classification and detection augmentation policies. 



Text-to-Image Data for Visual Tasks. 
Traditional synthetic data are acquired through the renderings from the graphics engines <cit.>. However, this type of synthesis cannot guarantee the quality and diversity of the generated data, resulting in a large gap with real-world data. Recent success in generative models has made synthesizing photo-realistic and high-fidelity images possible, which could be used for training the neural networks for image recognition due to their unlimited generation. 
For example, early works explored Generative Adversarial Networks (GAN) <cit.> for image recognition tasks. <cit.> uses a class-conditional GAN to train the classifier head and <cit.> uses StyleGAN to generate the labels for object segmentation. <cit.> adopts the GAN as a generator to synthesize multiple views to conduct contrastive learning. 
Until recently, the text-to-images models <cit.> have been leveraged to synthesize high-quality data for neural network training due to their effectiveness and efficiency. 
<cit.> adopts GLIDE <cit.> to synthesize images for classifier tuning on the CLIP model <cit.>. However, simply tuning the classifier may not fully explore the potential of synthetic data. StableRep <cit.> proposes to use Stable Diffusion for generating pre-trained datasets for contrastive learning and leverages the synthetic data from different random seeds as the positive pairs. 
<cit.> explores synthesizing data under ImageNet label space and 
produces improved performance. Fill-up <cit.> balances the long-tail distribution of the training data using the synthetic data from text-to-image models.
<cit.> propose to leverage text-to-image models for segmentation tasks where they adjust the model to generate both masks and images. However, this type of method requires either finetuning the model on the target dataset or synthesizing a large amount of pre-trained data. 



§ GENERATIVE BACKGROUND AUGMENTATION


In this section, we introduce our generative background augmentation method using the text-to-image diffusion model to increase the diversity of the original dataset. We first discuss the principle of the Inpainting technique and then adapt this method to our background augmentation framework. 



 §.§ Generative Augmentation with Inpainting


Our primary data synthesis model is Stable Diffusion <cit.>. 
Formally, 

the Stable Diffusion generates images from a random Gaussian noise 𝐳∼𝒩(0, 1) by applying the denoising model ϵ(·) repeatedly. We denote the text language prompt embedding as 𝐬. In this work, we apply the widely-used classifier-free guidance <cit.> which linearly combines the conditional estimate and the unconditional estimate of noise, given by

    ϵ̃(𝐬, 𝐳_t) = wϵ(𝐬, 𝐳_t) + (1-w) ϵ(𝐳_t),

where w is the guidance scale and t is the step of the denoising process. After multiple denoising steps, Stable Diffusion employs a decoder 𝒟 to decode the latent to the image. 

Given an input image 𝐱, the inpainting task is to re-generate selective areas in the image indicated by a mask variable 𝐦 while keeping other pixels unchanged. Since Stable Diffusion conducts the generation process in the latent space, we consider inpainting in the latent space as well. First, the encoder projects input image to latent space 𝐱̂ = ℰ(𝐱). Similarly, the mask variable is also projected to the latent space by resizing it 𝐦̂=resize(𝐦) to have the same size of 𝐱̂. 

Now, the inpainting denoising step can be formulated as 

    ϵ̂(𝐬, 𝐱̂_t) = 𝐦̂⊙ϵ̃(𝐬, 𝐱̂_t) + (1-𝐦̂)⊙(𝐱̂_t - 𝐱̂_t-1).

Here, ⊙ denotes the element-wise multiplication operation. 


The inpainting process described above modifies only the masked region, denoted as 𝐦̂, within the original image's latent space. Owing to the localized nature of Stable Diffusion's decoding mechanism, this method can also provide a reasonable assurance that solely the region specified by 𝐦 will be regenerated within the image space.



 §.§ Object Augmentation or Background Augmentation?



Data augmentation serves the role of increasing data amount and diversity to optimize for generalization performance. Traditionally, data augmentation in object detection is conducted through image transformations. For example, Zoph et.al <cit.> define three kinds of data augmentation: (1) color operations, (2) geometric transformations, and (3) bounding box operations, which change the image illumination, position, and annotation respectively. However, these augmentation methods only change pixel values and keep the content of the object or background the same as the original images.


In this work, we explore another kind of augmentation method, using generative models like diffusion models to change the content of training images to increase the training data volume and diversity. 
Specifically, consider a training image 𝐱∈ℝ^3× H× W containing n different objects, where each object corresponds to a unique mask variable 𝐦_i∈{0, 1}^H × W. The background mask can be determined by the following formula, 

    𝐦_b = 1 - ∑_i=1^n 𝐦_i.

With different choices from {𝐦_b, 𝐦_i}, we are able to change the content of different regions of the original image. 

Specifically for dense vision tasks like object detection and instance segmentation, a natural question is whether we should augment objects, background, or both of them.
Surprisingly, we found object augmentation is not suitable in practice.

The reasons are multiple. 

First, generating a realistic object that can precisely fill the mask region is quite challenging for existing Stable Diffusion models. For example, if a generated object cannot adequately cover its corresponding mask region, reusing the mask as its annotation will mislead the training of the model. We visualize some exemplary failure cases of object augmentation in Fig. <ref>(a). Although some efforts have been investigated to improve the alignment between generated objects and their corresponding masks <cit.>, they either introduce prohibitive computational overhead for large datasets or achieve an unsatisfactory mask-object alignment level, 
which is harmful for enhancing task performance for object detection and instance segmentation.

In addition, we have observed that object augmentation can yield low-quality inpainting results for certain object classes, especially those involving human figures <cit.>. This can adversely affect the quality of the data and may even degrade the performance of vision-related tasks (as shown in <ref>(b)). Furthermore, augmenting small objects results in negligible improvements because they contain relatively few pixels (<ref>(c)). Consequently, the detection and segmentation of small objects do not benefit from object augmentation.


Finally, data generation costs of object augmentation is much higher than background augmentation.
While for background augmentation, only one inpainting is enough for the background area. For example, the averaged number of objects in the MS COCO training dataset <cit.> is 7.27, amounting to ∼7.3× generation cost when comparing object augmentation to background augmentation (<ref>(d)).










In the experiments section, we present a thorough ablation study demonstrating that object augmentation underperforms compared to background augmentation and can even lead to a degradation in downstream task performance when compared to a baseline without generative augmentation.




 §.§ Utility-Aware Background Augmentation


In this section, we concentrate on background augmentation techniques with high training utility. An optimal background augmentation approach must avoid introducing extraneous objects, as these can confound the training process due to the lack of corresponding annotations. Additionally, it is crucial for a background augmentation method to preserve the integrity of existing objects, thereby maintaining a high level of object-mask alignment. We have investigated several design strategies to fulfill these requirements.


Prompt Selection of Background Augmentation. 
A critical choice for data generation is the text prompt 𝐬 for guiding the inpainting process. For object augmentation, one can safely adopt the class label as the text prompt while for background augmentation a text description for the whole background is needed. 

Some datasets like MS COCO provide image caption annotations for each image, which could be a potential text prompt for background augmentation. 
However, applying an image caption as a text prompt directly brings two disadvantages in practice. 
First, the image caption usually contains object descriptions as well. Using this caption will generate additional objects in the background. Second, not all detection and segmentation datasets contain image captions. It would require other captioning models to generate the new caption for the image, which involves additional computation costs and data cleaning to get accurate captions. 

To this end, we use a simple text prompt “” for background augmentation. We found this choice is sufficient for Stable Diffusion to modify the background significantly without generating additional objects. 





Background Mask Erosion. 
During the process of background augmentation, we occasionally observed that objects would inadvertently extend into the background region. This phenomenon is illustrated in Fig. <ref>(b), where the giraffe appears larger than its corresponding mask, thereby disrupting the alignment with its initial annotation. To address this issue, we introduce the concept of background mask erosion. This technique employs a minimum filter to the background mask, specifically, 𝐦_b=MinimumFilter(𝐦_b, k). We empirically determined that a kernel size k of 7 is effective in preventing object extension. Fig. <ref>(a) demonstrates the process of eroding the background mask with the minimum filter, and Fig. <ref>(b) shows a much better result after adopting our background mask erosion.

Adaptive Augmentation Freedom. 

Considering the variation in background region sizes across different images, there is an increased likelihood of errors during background augmentation—such as the introduction of undesired objects or the unwarranted extension of objects into the background areas—for images with larger backgrounds, potentially leading to the generation of harmful synthetic data. In the Stable Diffusion inpainting process, the number of diffusion steps can be adjusted to calibrate the extent of modifications applied to the original input image, that is, the distance between the original and the generated image. Drawing inspiration from this aspect of the process, we propose tailoring the degree of background augmentation to the size of the background region.

To this end, we introduce an adaptive control method for the inpainting process. Specifically, we compute the ratio of the background area to the overall image area and utilize this metric to determine the appropriate number of diffusion steps for the original training image 𝐱. The relationship is formalized as follows:


    T̂ = T ×(1 - D ×sum(𝐦_b)/W × H).


In this equation, sum(·) denotes the summation of all elements within the input matrix, and W × H denotes the total area of the training image. 
D∈(0, 1) controls the range of diffusion steps variation.
T is the maximum diffusion step where T̂=T indicates background noise starts from random Guassian noise.  Technically, if the background-to-image area ratio is higher, we reduce the diffusion steps to preserve more information from the original data. 






Fig. <ref> illustrates the process of our overall background augmentation pipeline. The image is first diffused based on the area ratio calculated from the background mask and then denoised using the eroded mask to generate the new background.




 §.§ Training with Background Augmented Data

Unlike traditional data augmentation which modifies the image tensor on the fly <cit.>, our method generates the additional images and expands the dataset size. Hence, our data augmentation can be combined with conventional data augmentation methods to improve the generalization ability in object detection. 
We also utilize our method with semi-supervised learning. Specifically, we incorporate our augmented data in Soft Teacher <cit.> framework as part of the “strong augmentation” to optimize the student network.
The “weak augmentation” is applied to the unlabeled images on teacher networks to generate soft annotations for student learning. 




§ EXPERIMENTAL EVALUATION


In this section, we investigate the impact of the augmented data by using our simple background augmentation framework. We first visualize several examples of data augmentation and then report object detection/instance segmentation performance on MS COCO <cit.> and PASCAL VOC <cit.> datasets. 
We also conduct several ablation studies and investigations into different data regimes. 





 §.§ Visualization


We demonstrate a variety of synthesis results (as well as the original image and the background mask image) using our method on the MS COCO dataset <cit.>, as shown in Fig. <ref>.
Our approach can modify the background of training data to various extents. In certain instances, the new background remains within a similar scene as the original (e.g., images at the second row and first and second columns), yet exhibits different styles and content. Conversely, some augmented backgrounds may represent entirely new scenes. For example, in the fourth row and first column, the original background of a stop sign is transformed into a purple sky, while maintaining the original stop sign objects unchanged. In the fifth/sixth row second column, our method converts a background into a photo hanging on the wall.

These background augmentations enhance the diversity of the training dataset without introducing additional annotation burdens or compromising the quality of existing annotations.




 §.§ Object Detection Evaluation


We start with the evaluation of our method on the MS COCO 
dataset. We train a Faster-RCNN <cit.> with FPN <cit.>, using the ResNet-50 <cit.> backbone. We also provide results of other datasets and model architectures in <ref>. 

Implementation Details. 
Specifically, we sampled 1%, 5%, 10%, and 25% of the original training dataset. By default, background augmentation was applied to the selected training images, generating one augmented copy per image. Both the original and the augmented images were utilized during training. For the noise sampling schedule, we employed Stable Diffusion v2-1 <cit.> and DPMSolver <cit.>. We resized the image resolution to 512×512 for the inpainting process and subsequently restored the images to their original resolution post-inpainting.

For the training of the detection model, we adhered to the implementation and hyperparameters provided by MMDetection <cit.> and Soft Teacher <cit.>. We utilized anchors spanning 5 scales and 3 aspect ratios. During the training and inference phases, 2k and 1k region proposals were generated respectively with a non-maximum suppression threshold set at 0.7. During each training iteration, a subset of 512 proposals were sampled from the 2k available to serve as box candidates for training the RCNN. The model underwent training over 12 epochs on 8 GPUs, each processing 5 images. We opted for SGD as the training algorithm, initiating the learning rate at 0.01 and reducing it by a factor of 10 after 110k and 160k iterations. The weight decay was set at 0.0001, while the momentum was maintained at 0.9.

Comparison to Baselines. 
We compare our method against a baseline Faster-RCNN model that does not utilize generative background augmentation. Both the baseline model and our approach incorporate conventional data augmentation techniques, such as horizontal flipping with a 50% probability and multi-scale jittering, where images are randomly resized within the range of 800 to 1333 pixels. Furthermore, we contrasted our generative background augmentation with the state-of-the-art augmentation method, RandAugment <cit.>. RandAugment conducts an exhaustive search for complex augmentation policies across 14 different augmentation methods using grid search. Despite the simplicity and reduced search effort of our method compared to RandAugment, our approach consistently surpasses RandAugment in performance across various data regimes as shown in <ref>.





Table <ref> shows the results using our augmentation policy on the above procedures. The proposed background augmentation policy achieves systematic gains across different amounts of training data used, ranging from +2.7 mAP to +3.5 mAP. 
In comparison, a previous regularization technique applied to ResNet-50 Faster-RCNN achieves a gain ranging from +0.6mAP to +1.5mAP.


Combined with Semi-Supervised Learning. 
We then experiment with semi-supervised learning in object detection. In particular, we select Soft Teacher <cit.> and add our background augmentation data to the  dataset, which is used to train the teacher network to annotate the  dataset for knowledge transfer, which is the remaining data from MS COCO training set. The model (ResNet-50-based Faster-RCNN) is kept the same as in previous experiments. The other training hyper-parameters are based on the MMDetection implementation <cit.>. 

We compare Soft Teacher with or without our background augmentation as the augmentation strategy. The results are again shown in <ref>, from which we can find that our method also effectively improves the mAP. 




 §.§ Comparisons on Transformer Architecture


We also evaluate our background augmentation on Transformer-based architecture. Specifically, we use Swin Transformer small <cit.> as backbone, Mask-RCNN <cit.> as detection framework to train on the full MS COCO dataset. The training hyper-parameters are kept the same as the original configurations <cit.>, AdamW <cit.> optimizer (initial learning rate of 0.0001, weight decay of 0.05,
and batch size of 16), and 3x schedule (36 epochs with the
learning rate decayed by 10× at epochs 27 and 33). 
The standard augmentation includes horizontal flipping and multi-scale training (resizing the input such that the shorter side is between 480 and 800 while the longer side is at most 1333). 
We report the bounding box AP and instance mask AP metrics following the conventions of transformer evaluation.



<ref> summarizes the results on the Swin transformer. We can find that background augmentation effectively improves the detection and segmentation performance (+1.0 AP^box and +0.8 AP^mask on the Swin-T architecture). 









 §.§ Ablation Studies


In this section, we conduct ablation studies that verify the design choices of our background augmentation pipeline. In order to do this, we generate different sets of COCO training data and combine them with original data to train a ResNet-50 RetinaNet <cit.>. We use a similar training schedule and hyper-parameters with previous experiments. 
In particular, we compare several cases: (1) using object augmentation or background augmentation, (2) using background augmentation but with or without our proposed techniques. 
Specifically, for prompt input we either use the COCO caption as prompt or our proposed simple prompt method, and for background mask erosion as well as the adaptive diffusion steps, we choose to either apply them or not. 

Table <ref> summarizes the results of our ablation experiments, from which we can observe that applying the object augmentation will even decrease the mAP performance (-0.3 mAP). The same situation also applies to background augmentation if the prompt is not chosen correctly. This result indicates that augmenting the data in object detection is non-trivial. With background mask erosion and appropriate prompt choice, we can achieve better detection performance than the baseline (+1.3 mAP on the full COCO dataset). Furthermore, adding the adaptive strength control enhances the final performance to 38.3 mAP. 



 §.§ Augmenting More Data

In previous experiments, we only generate 1x of COCO data. This is still far from the ability of text-to-image models as it can generate a different image by randomly sampling a different 𝐳. In this section, we explore the limits of our background augmentation. Specifically, we experiment with two cases. In the first case, we repeatedly generate α times data for each image and study the effect of increasing α (called uniform sampling).
In the second case, we increase the ratio of augmenting small objects in the training data since detecting small objects often poses more challenges than detecting large objects <cit.>. 
Therefore, we increase the proportion of small object detection based on the background area ratio, which is called non-uniform sampling. 
In total, we also generate α× copy of COCO data but with more augmented data for small objects in this case. 




We evaluate these two scaling cases in the 10% COCO data regime. The α is set to 1, 2, 5, and 10 (Note that α=0 indicates no background augmentation applied). We use ResNet-50-based Faster RCNN <cit.> to evaluate different data regimes. 
The training hyper-parameters are kept the same with <ref>. We report four different metrics: mAP, mAP^small, mAP^medium, and mAP^big, as shown in <ref>.

From the figure, we find that uniformly increasing the number of augmented data does not consistently increase the mAP performance. 
Observing the results from different sizes of the objects, we find that the major cause is the small object performance, which is decreasing when α≥ 2. Instead, non-uniform sampling prioritizes the small objects in augmentation and ensures more diverse images are created in small object detection, leading to better-performed detectors. 






 §.§ Evaluation on PASCAL VOC

To evaluate the effectiveness of our proposed background augmentation on an entirely different dataset and another detection algorithm, we train a Faster R-CNN <cit.> model and a RetinaNet model <cit.> with a ResNet-50 backbone on PASCAL VOC dataset <cit.>.
We combine the training sets of PASCAL VOC 2007 and PASCAL VOC 2012, and test our model on the PASCAL VOC
2007 test set (4952 images). Our evaluation metric is the
bounding box mean average precision (mAP^box).
The training hyper-parameters are kept the same with previous experiments and the backbone model is initialized from the ImageNet pre-trained checkpoint <cit.>. 
We also include the performance of RandAugment in the PASCAL VOC evaluation. 

The results are summarized in <ref>, where we can find our method leads to a significant increase in terms of box mAP. 






§ CONCLUSION

In this work, we propose a simple yet effective background augmentation framework for object detection and instance segmentation. We directly utilize the existing mask annotation to perform Inpainting to create new content in the training data. Through a set of rigorous explorations, we confirm that background augmentation, instead of object augmentation, can achieve higher potential. We also propose several techniques to ensure the background contains no additional object so that the generated data can be directly used for augmentation. Results show that our method can be effectively used to improve the detection performance.
Our method has certain limitations: for images containing extremely large objects, background augmentation can hardly alter the content (as shown in <ref> the improvement of mAP^Large is relatively small). 
In this case, our method could be combined with an effective object augmentation (that may require finetuning of the text-to-image model) to enrich the diversity of this type of data.



§ ACKNOWLEDGEMENT

This work is sponsored by Sony AI.







splncs04



 Paper Title: Revocable Backdoor for Deep Model Trading.










207 















A]Yiran Xu[Equal contribution.]
A]Nan Zhong
A]Zhenxing QianCorresponding Author. Email: zxqian@fudan.edu.cn 
A]Xinpeng ZhangCorresponding Author. Email: zhangxinpeng@fudan.edu.cn 

[A]Fudan University




§ ABSTRACT 
Deep models are being applied in numerous fields and have become a new important digital product. Meanwhile, previous studies have shown that deep models are vulnerable to backdoor attacks, in which compromised models return attacker-desired results when a trigger appears. Backdoor attacks severely break the trustworthiness of deep models. In this paper, we turn this weakness of deep models into a strength, and propose a novel revocable backdoor and deep model trading scenario. Specifically, we aim to compromise deep models without degrading their performance, meanwhile, we can easily detoxify poisoned models without re-training the models. We design specific mask matrices to manage the internal feature maps of the models. These mask matrices can be used to deactivate the backdoors. The revocable backdoor can be adopted in the deep model trading scenario. Sellers train models with revocable backdoors as a trial version. Buyers pay a deposit to sellers and obtain a trial version of the deep model. If buyers are satisfied with the trial version, they pay a final payment to sellers and sellers send mask matrices to buyers to withdraw revocable backdoors. We demonstrate the feasibility and robustness of our revocable backdoor by various datasets and network architectures.








§ INTRODUCTION


Deep learning has achieved impressive performance in a wide range of fields <cit.>. However, the security vulnerability of deep models, such as the notorious backdoor attack, hinders the deployment of deep models in some risk-sensitive domains <cit.>. Most attackers <cit.> implant backdoors into a clean model during the training phase by data poisoning. The compromised models behave normally during the evaluation phase, whereas they return attacker-desired results when the predefined trigger appears. Considering the stealthiness and hazard of backdoor attacks, it has attracted a lot of attention in the machine learning security community <cit.>.

Although backdoor attack is seen as a security threat or vulnerability for deep models in most studies, some researchers adopt backdoor attack to conduct some positive purposes <cit.>. There are various positive usages of the backdoor like dataset or model copyright protection <cit.>, artificial intelligence interpretability <cit.>, adversarial example defence <cit.>, etc. Adi et al. <cit.> adopt the backdoor attack to implement model watermarking for model copyright protection. They use significantly abnormal outputs of models triggered by poisonous inputs as the watermarking signal to verify the ownership of the model. A similar idea <cit.> is also implemented to achieve dataset ownership protection. Shan et al. <cit.> utilize the backdoor as a honeypot for deep models to resist adversarial example attack. When an attacker constructs an adversarial example for the backdoored model, the distribution of the interior feature map of the adversarial example is close to that of poisonous ones. Therefore, defenders identify adversarial example input by inspecting the distribution of the interior feature map.


 


In this paper, we first propose a novel concept of the revocable backdoor and explore its promising application which is deep model trading. The revocable backdoor denotes that attackers implant backdoors into clean deep models while attackers can easily withdraw these backdoors. In other words, attackers control poisoning and detoxification for deep models at the same time. The detoxification should be easy and efficient. Revocable backdoor still needs to satisfy the general properties of a backdoor. When inputs contain the trigger, the model with revocable backdoor produces the attacker's predefined outputs. Otherwise, the model should behave normally and produce correct output. Additionally, it should possess stealth and robustness to prevent users from detecting and removing the backdoor. The revocable backdoor can be applied to deep model trading to protect the rights and interests of both buyers and sellers.

The main contributions of this paper are three-fold: 

    
  * We propose a promising deep model trading scenario, and hold that the revocable backdoor can address the issues haunted by buyers and sellers. 
    
  * We propose a practical revocable backdoor, which does not require extra clean or poison data to withdraw backdoors. 
    
  * Extensive experiments demonstrate that our proposed revocable backdoor is feasible and robust. 






§ MODEL TRADING SCENARIO


Since the wild usage of deep learning models in real world, models become a tradable products. However, how to sale a trained model or buy a trained model still remains a question. The Uniform Commercial Code (UCC) <cit.> stipulates that when the goods delivered by the seller do not conform to the contract, the buyer has the significant right to "revoke acceptance". Correspondingly, when the buyer breaches the contract, such as improperly rejecting the goods, incorrectly revoking the acceptance of the goods, failing to pay the price etc., the seller has the right to undertake a series of remedial measures, including recovering the goods, reselling, and so on.

Therefore, the rights and demands of both the buyer and seller in the model transaction can be summarized as follows:

Sellers seek to profit by selling their own models. As models are considered data resources, a model "return" would carry immeasurable loss of benefits. They attempt to control costs as much as possible to increase profits.

Buyers hope to purchase models that meet their needs. When they find that the purchased model does not meet their needs, they want to be able to return it and get a refund.

Trading platform provides a reliable platform for model transactions, safeguarding the reasonable rights and interests of both buyers and sellers.

However, there are still a few issues:

    
  * Q1. How to allow the buyer to return and refund the purchased model when they find it does not meet their needs, ensuring their right to "revoke acceptance"?
    
    
  * Q2. How to protect the seller's interests without being damaged when the buyer returns the model? That is, how to prevent the buyer from saving some of the data after receiving the model, and then returning and refunding it.


A possible solution to these questions could be to send a trial version at the time of delivery. The full version would then be delivered once the buyer confirms receipt. Then, the questions turn into:


    
  * Q3. How to ensure that both the trial version and the full version of the model meet the buyer's requirements?
    
    
  * Q4. How to ensure the model will not be resold by the buyer?


To solve these two questions, we propose the concept of a revocable backdoor. Fig. <ref> illustrates the details of the deep model trading scenario. A buyer distributes his requirement to a seller. Then, the seller trains a model based on the buyer's requirements. The model has a revocable backdoor and can be sent to the buyer as a trail version. The buyer pays a deposit to obtain the trial model whose performance is almost the same as the clean final model over a normal task. Next, the buyer evaluates the trial model on his/her application and determines whether accepts the model and pays a final payment. If the buyer is satisfied with the trial version, they can pay the seller the final balance. After that, the seller sends the mask matrix to the buyer and the revocable backdoor can be removed by the mask matrix. 

 Buyers can disclose the backdoor hidden in the model publicly and make the trial model unusable. Fig. <ref> illustrates the details of the deep model trading scenario.
The advantages of using the revocable backdoor are as follows:


    
  * It allows the user to try the model while securing the rights of the seller. In the worst case, the buyer maliciously refuses to pay the final payment and deploys the trial model in his/her application. The seller can disclose the backdoor hidden in the model publicly and makes the trail model unusable. Or, if there is a third part (sale platform) involved, the seller can register the mask and sue the buyer for infringement. (Address Q1 and Q2.)
    
    
  * It also ensures the relevance between the trial version and the complete version of the model, preventing the buyer from receiving a complete model that does not match the trial version after paying the balance and confirming receipt. (Address Q3.)
    
    
  * There is no backdoor in the complete version of the model, but the seller can still use the mask matrix to identify whether a model is the one they sold (like a serial number of a software). That is, by subtracting the mask matrix from the model, determining whether the model with the mask matrix removed contains a preset backdoor to authenticate the model's genuineness. (Address Q4.)






§ RELATED WORK




 §.§ Backdoor Attack

Backdoor attack is an emerging topic in the machine learning security community <cit.>. The goal of the backdoor attack is to compromise deep models, which return normal results for clean inputs, but return attacker-desired results when the trigger appears. Gu et al. <cit.> first propose the concept and specific practical backdoor attack approach named BadNets. BadNets mounts an attack through data poison during the training phase. A conspicuous colourful square is adopted as a trigger pattern. BadNets generate some poisonous images by adding the trigger pattern to some clean images and changing their label to the target (attacker-desired) label. Such poisonous images are mixed with other clean training sets. The deep model trained over a poisonous training set contains an insidious backdoor which only activated by the trigger pattern. The implementation of the backdoor attack can be roughly grouped into two categories. 

Poisoning-based backdoor mounts an attack by inserting some well-designed poisonous samples into the training set <cit.>. Previous researchers focus on designing stealthy and efficient trigger patterns. Although BadNets proposes a seminal approach for backdoor attacks, its trigger pattern is conspicuous. In the follow-up studies, researchers explore some invisible trigger patterns. Li et al. <cit.> and Li et al. <cit.> adopt image steganographic noise ( Least Significant Bit and DNNs-based steganographic ) as a trigger pattern. The visual quality of poisonous is greatly boosted by invisible trigger patterns. Furthermore, WaNet <cit.> adopts image affine transformation as a trigger pattern. The poisonous images can be viewed as slightly warped from clean ones. WaNet works effectively not only in the digital domain but also attacks successfully in the physical domain.

Training process controlling-based backdoor can achieve more stealthy attacks than poisoning-based ones <cit.>. As a trade-off, these approaches entail controlling the entire training phase. These attacks focus on both the stealthiness of the trigger pattern and compromised models. Backdoor attacks may leave traces in compromised models. Doan et al. <cit.> focus on the feature map separability of compromised models and employ Wasserstein distance to minimize the feature map distance between poisonous and clean inputs. Some similar ideas are proposed in subsequent studies. Besides consideration of the feature map domain, BppAttack <cit.> employs contrastive supervised learning and adversarial training instead of standard cross-entropy loss in the classification task. This learning approach ensures that compromised models steadily and successfully converge.





 §.§ Backdoor Defence

Considering the serious threat of backdoor attacks to machine learning security, backdoor defence, which is the countermeasure against backdoor attacks, rapidly develops in recent studies <cit.>. There are two main categories in the backdoor defence, i.e., backdoor detection <cit.>, backdoor purification <cit.>

In terms of backdoor detection approaches, they aim to determine whether the deep model contains a backdoor. Neural Cleanse (NC) <cit.> is one of the well-known backdoor detection algorithms. NC observes that the backdoor of the classifier constructs a shortcut between the target label and other clean labels. Based on this observation, NC employs a abnormality detection algorithm to determine whether the model is poisoned by analyzing the size of perturbation leading all clean images to the target label. 

Compared with backdoor detection, backdoor purification aims to eliminate backdoors on the premise that the performance of clean inputs does not significantly degrade. Fine-Pruning <cit.> adopts network pruning to erase backdoors hidden in the models. For existing neural networks, most neurons are dormant when the clean inputs come. However, these dormant neurons may be activated by the trigger pattern and lead models to return attacker-desired results. Therefore, Fine-Pruning identifies dormant neurons based on some clean inputs then deletes them and fine-tunes the model. NAD (Neural Attention Distillation) <cit.> is an alternative effective backdoor purification approach, which utilizes model distillation to erase backdoors. NAD fine-tunes backdoored model to generate a teacher model which is used to conduct model distillation. Many other sophisticated model purification approaches are proposed in subsequent studies.







§ REVOCABLE BACKDOOR



 


 §.§ Preliminaries



In this paper, we focus on the image classification task which is consistent with previous studies. Since the revocable backdoor is first proposed by ours, we describe more details of the revocable backdoor requirement applied in the deep model trading. 

Goal The revocable backdoor aims to ensure that buyers complete the final payment when they are satisfied with the trial model. If buyers refuse to complete the final payment, sellers can disclose the hidden backdoor and corresponding trigger pattern. For example, buyers obtain a face recognition model with a revocable backdoor and deploy this model in their company. If the backdoor is disclosed, anyone with the trigger pattern can break the face recognition model, that is, the model is unusable. Note that our revocable backdoor does not aim to protect the copyright of the model, which is usually accomplished by model watermarking <cit.>. Conventional backdoor-based watermarking aims to verify the ownership of the model, whose goal is completely different from ours. Although backdoor-based watermarking also can make the model return incorrect results when the trigger appears, the backdoor is permanent and irrevocable. Even though buyers complete the final payment, the backdoor still exists in the model and brings security threats to the model deployment.

Effectiveness In previous backdoor attacks, attackers aim to make the compromised classifier return the target label when the trigger appears <cit.>. However, in the context of model tradings, the goal of our revocable backdoor for model sellers is different from previous studies. Sellers can distribute the trigger pattern to make the classifier held by buyers unusable in the event that the buyer does not complete the final payment. Therefore, our revocable backdoor aims to decrease the classification accuracy as low as possible when the trigger appears. In other words, we view the attack success if the compromised classifier returns an arbitrarily wrong label. 

Fidelity The classifier with a revocable backdoor can be used as a trial version of the final classifier for buyers. Therefore, the revocable backdoor aims to keep the performance of the backdoored classifier over clean inputs the same as a clean classifier. For buyers, the performance of the trial version over clean inputs is almost identical to the final version (backdoor-free). In this case, buyers can evaluate whether the classifier meets their requirements.

Revocability The crux of our revocable backdoor is that we can easily withdraw the backdoor without requiring an extra training phase or clean/poisonous data. The revocability denotes that the compromised classifier whose backdoor is withdrawn can correctly predict poisonous inputs.

Robustness In the model trading scenario, the backdoors hidden in the compromised model can and only can be withdrawn (erased) by sellers. Therefore, our attack should ensure that the backdoor is resistant to various backdoor purification defences. 





 §.§ Method


Fig. <ref> illustrates the framework of our approach. 
Let 𝒟 = { (x_i, y_i) }_i=1^N denotes the benign training set, where x_i ∈𝒳= {0,1,…, 255}^C× W × H is the image, y_i ∈𝒴 = {1,…, N} is its label, and N is the number of classes. The trigger pattern is denoted as δ, where δ∈𝒫={-t,-t+1,…,t+1,t}^C× W × H. The elements of the trigger pattern are between -t and t. In the backdoor attack, we select a portion of clean images x, which are stamped with the trigger pattern, as poisonous ones. We denote the clean training set and poisonous training set as 𝒟_c and 𝒟_p, respectively. Given a deep neural network-based classifier f_θ(·), we aim to minimize the following equation,

    L_bd = 1/|D_c|∑_(x,y) ∈ D_cℓ( f_θ(x),y) 
     - 1/|D_p|∑_(x,y) ∈ D_p max( ℓ( f_θ(x+δ),y),c),

where ℓ(·) denotes cross-entropy loss widely adopted in classification tasks. We update f_θ(·) and δ to minimize equation <ref>. The core idea of equation <ref> is to achieve the effectiveness and fidelity of the revocable backdoor attack. The first term of equation <ref> ensures that the performance of backdoored classifiers over clean inputs does not drop. The second term denotes maximizing the cross-entropy loss of poisonous inputs, i.e., making the classifier return wrong results when the trigger appears. However, maximizing the cross-entropy loss of poisonous inputs is much easier than minimizing the cross-entropy loss of clean inputs. In other words, classifier f_θ(·) easily learns how to give wrong results for poisonous inputs but it is hard to learn how to give correct results for clean inputs. As a result, we adopt max(·)  and hyperparameter confidence c to restrict the second term. For a N categories classification task, if the cross-entropy loss of input is larger than -log(1/N), the classifier will return an incorrect result. Therefore, we set c=-1.1× log(1/N) that is slightly larger than the minimum correct decision threshold for a N categories classifier. If we remove the restriction of max(·) or set an excessively large c, the classifier f_θ will only learn the second term of equation <ref> and overlook the first term.

Besides L_bd, our approach also entails considering how to withdraw our backdoor. The modern deep neural network-based classifier f_θ(·) is made up of some cascade layers. The output of f_θ(·) can be expressed as,

    f_θ(·) = f_n(⋯f_2( f_1(⋯))),

where n denotes the number of layers. The output of the interior layer is named as feature map. Based on the cascade structure of f_θ(·), we employ some matrices named masks to manipulate feature maps. Compromised classifiers learn two tasks, i.e., a normal task and an insidious backdoor task. These two tasks conduct different inference links. In other words, the feature maps of these two tasks are different. Therefore, we can adopt some mask matrices to break backdoors without affecting clean inputs. We denote the classifier f_θ(·) with mask matrices M_k as f_θ_mask(·). The formulaic expression is shown as follows,

    f_θ_mask(·) =  f_n(⋯ M_2· f_2( M_1· f_1(x))),

where M and · denote trainable mask matrices and Hadamard product, respectively. The size of M_k is the same as the corresponding feature map. The initial elements of M_k are filled with 1. Moreover, we require a regularization term for M_k. The goal of regularization is used to make the elements of M_k as close to 1 as possible. The regularization for M_k is expressed as follows,


    L_R = sum(|M_k_ij - 1 | ),

where sum(·) denotes the sum of all elements.
The usage of L_R forces M_k to change the most key elements related to poisonous inputs in feature maps. L_R can effectively mitigate the impact of M_k on clean inputs. For f_θ_mask(·), it should return correct results regardless of whether the trigger pattern δ appears. Therefore, the loss function for f_θ_mask(·) is expressed as follows,

    L_rev = [1/|D_c|∑_(x,y) ∈ D_cℓ( f_θ_mask(x),y) 
     + 1/|D_p|∑_(x,y) ∈ D_pℓ( f_θ_mask(x+δ),y)].

We update M_k by minimizing L_rev and make M_k adaptive find and break the poisonous inference link.
The final cost function is expressed as follows,

    L = L_bd+L_rev+α· L_R,

where α is the balance factor to control the weight of the regularization. We simultaneously update parameters of f_θ, trigger pattern δ, and mask matrices M_k to minimize L. The backdoor hidden in the well-trained f_θ can easily be withdrawn by adding M_k in the interior layers.  f_θ(·) and f_θ_mask(·) can be used as the trial model and final model for sellers during the model trading.





 §.§ Trigger Fine-tuning
 
In the previous subsection <ref>, we describe our approach to create a compromised classifier whose backdoors are controlled by mask matrices. The trigger pattern δ, whose elements are between -t and t, is trained with the classifier at the same time. In this subsection, we delve into the design of the trigger pattern to control the balance of the trigger pattern imperceptibility and attack robustness. Hyperparameter t controls the modification magnitude of the trigger pattern. In other words, a small t ensures the imperceptibility of the trigger pattern. In this subsection, we frozen classifier f_θ and mask matrices M, i.e., f_θ_mask is also frozen, and update trigger patterns. The loss function of updating trigger δ is as follows,

    L = L_bd+L_rev+β· | δ|_2,

where | δ|_2 denotes the L_2-norm of the trigger δ, and β controls the balance of | δ|. We control the modification magnitude of the trigger pattern δ by changing the value of the hyperparameter β and t. 



 §.§ Backdoor Erasing for Existing Attacks

To the best of our knowledge, we are the first to propose the concept of the revocable backdoor. For existing conventional backdoor attacks like BadNtes <cit.>, WaNet <cit.>, Blend <cit.>, etc. They can not be withdrawn by attackers. In this subsection, we describe a backdoor erasing method to withdraw existing backdoors. In the model trading scenario, we suppose that buyers own some clean images and their corresponding poisonous ones. Then, they update the parameters of classifier f_θ by minimizing the following loss as,

    L_erasing = [1/|D_c|∑_(x,y) ∈ D_cℓ( f_θ(x),y) 
     + 1/|D_p|∑_(x,y) ∈ D_pℓ( f_θ(x+δ),y)].

Compared with equation <ref>, L_erasing minimizes the cross-entropy loss between clean and poisonous images and their label. The number of poisonous images is critical to the performance of backdoor erasing. We employ this method as our baseline. Compared with our proposed method <ref> and <ref>, this method requires gradient calculation and some poisonous images, which is not practical in the model trading scenario.  





§ EXPERIMENTAL RESULTS










 §.§ Experimental Setup




Datasets and Networks In this paper, we adopt three datasets including CIFAR-10 <cit.>, GTSRB <cit.>, and Sub-ImageNet to evaluate the performance of our approach. These datasets are widely used in previous studies <cit.>. The number of categories of CIFAR-10 and GTSRB is 10, and 43, respectively. Sub-ImageNet consists of 10 categories which are randomly selected from the original ImageNet dataset <cit.>. The image size of CIFAR-10, GTSRB, and Sub-ImageNet is 32× 32,  32× 32, and 224× 224, respectively. In terms of network architectures, we adopt ResNet-18 <cit.> and VGG <cit.> in the experiments.  

Backdoor Baselines Since we are the first to propose the concept of the revocable backdoor attack, there are no suitable revocable backdoor attacks used as a comparison. We suppose that users obtain some clean and corresponding poisonous images to finetune to a backdoored classifier to achieve revocability based on existing conventional attacks. We adopt various trigger patterns to conduct backdoor attacks. The trigger pattern includes BadNets <cit.>, Blend <cit.>, SIG <cit.>, LSB <cit.>, WaNet <cit.>, and BppAttack <cit.>. The details setting of baseline triggers can be found in the supplement. The compromised (backdoored) classifier is generated by updating equation <ref> for baselines. The poisoning rate is fixed as r=0.1. In the backdoor erasing phase, we suppose that users obtain o=5% clean data of the training set and corresponding poisonous ones. Then, we finetune the backdoored baseline classifiers with equation <ref>. 


Implement Details
Our approach consists of two steps, training revocable backdoored classifier and finetuning trigger pattern. The training epoch and batch size are 180, and 256, respectively. We adopt SGD (Stochastic Gradient Descent) optimizer for updating the parameters of the classifier. The initial learning rate of SGD is 0.01. For trainable mask matrices and trigger pattern δ, we adopt Adam optimizer with a 0.001 initial learning rate. The poisoning rate is 0.5 for the first 80 epochs. The poisoning rate times by 0.5 in intervals of 10 epochs when the iteration epoch exceeds 80. Hyperparameter α is fixed as 10, 10, and 100 for CIFAR-10, GTSRB, and Sub-ImageNet, respectively. In the second phase, we fixed the poisoning rate as 0.5. The hyperparameter β is 0.01 for all datasets. The hyperparameter t controlling the range of our trigger pattern is 10 in both phases.

























 §.§ Attack Effectiveness

We adopt accuracy to measure the performance of the classifier. In this part, we evaluate the Effectiveness, Fidelity and Revocability of our approach. The results of our approach and baselines are shown in Table <ref>. We also train a clean classifier whose training setup is the same as compromised ones as a reference.  We aim to significantly degrade the classifier accuracy when the trigger appears while keeping the accuracy for clean inputs unchanged. In terms of revocability, the goal of mask matrices employed in our attack is to withdraw the backdoor. From Table <ref>, our approach achieves satisfactory results in the aspect of attack effectiveness, fidelity, and revocability. The accuracy for clean inputs is almost unchanged compared with the clean classifier in most cases. The accuracy for poisonous inputs is decreased to near even less than random guesses, meanwhile, it recovers to standard performance (like a clean classifier) when we insert mask matrices to withdraw the backdoor. For baselines, their performance is significantly inferior to ours. In previous common backdoor studies <cit.>, their attack setting is that a compromised classifier returns the target label when the trigger appears. This attack setting is much easier than our setting, in which the compromised classifier does not need to understand the semantic information of inputs and only builds a strong relation between the trigger and target label. However, the compromised need to understand the semantic information of inputs and then return wrong results when the trigger appears in our attack setting. Therefore, some invisible triggers (like LSB and BppAttack) whose perturbations are small cannot successfully compromise the classifier. The experiments of VGG network can be found in Table 2 of the supplementary material.




Apart from quantitative results, we also visualize our trigger pattern in Fig. <ref>. Our approach achieves satisfactory visual results and users cannot perceive trigger patterns. Based on Fig. <ref> and Table <ref>, we find that only triggers (like BadNets) with large modification (visible) can attack successfully among baselines. Our approach achieve both trigger invisibility and attack effectiveness at the same time.  



 


 §.§ Attack Robustness




We recall our revocable backdoor scenario, i.e., deep model trading scenario. For sellers, the goal of the revocable backdoor is to make the model unusable by distributing the trigger pattern to the public if buyers do not pay the final payment. Sellers can send the trial model to buyers and let them know the existence of the hidden backdoor. Therefore, we do not consider backdoor detection defence in the robustness evaluation. We focus on evaluating the resistance of our backdoor against backdoor purification defences <cit.>. In other words, we evaluate whether buyers can adopt backdoor purification to erase the backdoor without the permission of sellers. 

The easiest backdoor purification method, which may be adopted by buyers, is fine-tuning. Buyers collect some clean images whose distribution is similar to the training set to fine-tune the trial classifier containing our revocable backdoor. 
To simulate fine-tuning, we adopt SGD optimizer with 0.01 learning rate kept the same as the training phase to fine-tune the compromised classifier with 5% clean images from the training set. We fine-tune the classifiers for 50 epochs. The experimental results have been shown in Table <ref>. Fine-tuning cannot remove our backdoor in all cases.
 
 




An alternative well-known backdoor is Fine-pruning <cit.>, which combines network pruning and fine-tuning. As aforementioned related work part, most neurons are dormant for clean inputs. These dormant neurons almost do not affect the final performance of clean inputs. However, they may be activated by poisonous inputs. Therefore, Fine-pruning prunes dormant neurons with some clean inputs and fine-tunes the pruned classifier at the same time. The specific defence setup is as follows. We pruning the last layer of the classifier which is widely adopted in previous studies. The classifier is fine-tuned in the interval of 10 neurons. The optimizer is SGD with a 0.01 learning rate. We set the maximum tolerable drop of the clean task accuracy as less than 1%. Table <ref> shows the results of Fine-pruning. It also cannot remove our backdoor. The details of the accuracy (both clean and poisonous inputs) with regard to the number of pruned neurons can be found in the supplementary material.





Besides fine-tuning and fine-pruning, we also evaluate the robustness of our approach by NAD <cit.>. NAD first trains a teacher model. Then, NAD adopts the teacher model as a guide to conduct attention distillation. The experimental results of NAD are shown in Table <ref>. NAD is ineffective in our approach. Although NAD can erase the backdoor in some cases (GTSRB dataset), the accuracy of clean inputs severely degrades compared with the clean model. 







 §.§ Ablation Studies

In our approach, we adopt two steps to generate the trigger pattern. The second part, which fine-tunes the trigger, aims to control the perturbation of trigger pattern. Table <ref> gives the quantitative comparison between the trigger with or without fine-tuning. The PSNR (poisonous image quality) and SSIM <cit.> is significantly improved benefit from the trigger fine-tuning.


More ablation studies over confidence c, regularization L_R and the uniqueness of the mask matrices can be found in the supplement. 






§ CONCLUSIONS


In this paper, we first propose the concept of the revocable backdoor attack and a novel application, that is, deep model trading. Revocable backdoor aims to create a trial model whose performance is similar to the clean/final model for buyers to evaluate its performance. Meanwhile, sellers can easily withdraw the backdoor hidden in the trial model when they obtain the final payment. The revocable backdoor attack is achieved by controlling the interior feature maps of models. We train models and mask matrices at the same time. The mask matrices are critical to turning the backdoored (trial) model into a clean (final) model. Extensive experiments demonstrate that our approach is feasible and robust. 








This work was supported by the National Natural Science Foundation of China under Grants U20B2051, 62072114, U20A20178, U22B2047










 Paper Title: Outlier Detection in Large Radiological Datasets using UMAP.

















Outlier Detection in Large Radiological Datasets using UMAP
     Mohammad Tariqul Islam and Jason W. FleischerCorresponding Author 

	Department of Electrical and Computer Engineering

	Princeton University

	Princeton, NJ 08544, USA 

	 

    July 2024
===================================================================================================================================================================================



§ ABSTRACT 
The success of machine learning algorithms heavily relies on the quality of samples and the accuracy of their corresponding labels. However, building and maintaining large, high-quality datasets is an enormous task. This is especially true for biomedical data and for meta-sets that are compiled from smaller ones, as variations in image quality, labeling, reports, and archiving can lead to errors, inconsistencies, and repeated samples. Here, we show that the uniform manifold approximation and projection (UMAP) algorithm can find these anomalies essentially by forming independent clusters that are distinct from the main (“good”) data but similar to other points with the same error type. As a representative example, we apply UMAP to discover outliers in the publicly available ChestX-ray14, CheXpert, and MURA datasets. While the results are archival and retrospective and focus on radiological images, the graph-based methods work for any data type and will prove equally beneficial for curation at the time of dataset creation. 







§ INTRODUCTION


Accepted in MICCAI-2024 Workshop on Topology- and Graph-Informed Imaging Informatics (TGI3)
A prominent reason behind the current success of machine learning-based disease detection is the availability of large medical datasets. However, for the machine learning models to be reliable, quality datasets representative of the target population need to be ensured <cit.>. 
The labels in these datasets are often generated from human annotations using automated extraction or entity detection tools. 
However, these annotations (and their structured archiving) from automated tools can have errors due to faulty perceptions, interpretations, and human errors <cit.>. 

Even if the error rate of the annotator is less than 4%, this can lead to millions of annotation errors per year <cit.>. 
This parallels the 3.3% error rate in large computer vision dataset <cit.>.
Thus, there needs to be a better way to identify such errors before they are included in a dataset.




For images, the search can be performed visually. However, examining individual images is a daunting task that requires many human hours. 
A popular alternative is neighbor embedding <cit.>, which can produce a two-dimensional (2-D) cluster plot that can be analyzed visually quickly. 
(This class of methods is also known as nonlinear dimensionality reduction as the 2-D plot preserves the pairwise similarity, i.e., graph structure, of the original high-dimensional space.)
Widely used neighbor embedding algorithms are t-distributed stochastic neighbor embedding (t-SNE) <cit.> and uniform manifold approximation and projection (UMAP) <cit.>. 
UMAP was introduced relatively recently and has become very popular, as this method draws concepts from rich algebraic and topological structures and is computationally fast.

In this paper, we design a UMAP-based visual analytic method for extracting outlier images from large x-ray datasets. 
We validate our method by analyzing three publicly available and widely used medical image datasets. 
We show that the method can successfully cluster image features and produce interpretable visualization. 
We also discover labeling errors and erroneous images that have slipped through the verification process done prior to dissemination. Codes to reproduce the results are available at <https://github.com/tariqul-islam/Outlier_Detection_UMAP>.




§ RELATED WORKS


In the literature, the term outlier is often used interchangeably with abnormality and anomaly <cit.>.
Here, we define outliers as images that do not have sufficient signal for final decision-making or do not belong in the dataset due to specification. 
Generally, outlier detection methods assume an underlying distribution and often model it as normal distribution <cit.>, i.e., a data point is an outlier if it is far away from the mean of the fitted distribution.
Fritsch et al. <cit.> used the minimum covariance determinant estimator and its extensions to find outliers (due to motion or registration issues) in neuroimaging data by analyzing principal components.
Gang et al. <cit.> used a t-SNE plot to find outliers from binary lung masks in terms of size variation and segmentation error. Fleischer and Islam <cit.> employed UMAP on chest x-rays for phenotyping COVID-19 response.




§ SYSTEM OVERVIEW



Following preprocessing (discussed in <ref>), the major parts of the framework are feature extraction and dimensionality reduction (Fig. <ref>). To extract features from these images, we employed DenseNet-121 <cit.> trained on ImageNet <cit.>, a widely used deep neural network architecture designed to efficiently propagate features from earlier layers of a network to deeper layers. 
Importantly, neural algorithms are usually robust to many variabilities in images by design, and thus can accommodate standard images and outliers on equal terms.

Medical images usually vary in resolution, have different contrast, brightness, and alignment, and often suffer from registration issues.
In our framework, the features have been extracted from the final layer (before the softmax layer) of the network, where the features are generally most discriminating. 
Since we are not using a radiologically pre-trained model, these features generally will not be able to identify individual diseases. Rather, we employ other related labels (e.g., x-ray views, study labels) to examine the datasets. After extracting the features, we employ UMAP <cit.>, to obtain a 2-D approximation of the high-dimensional features.  



§ DATASETS


We evaluate our approach on three publicly available datasets: ChestX-ray14 <cit.>, CheXpert <cit.>, and Musculoskeletal Radiographs (MURA) <cit.>. 
ChestX-ray14 contains 112,120 frontal chest x-ray images from 30,805 unique patients. Images are from posterior-anterior (PA) and anterior-posterior (AP) views.
CheXpert dataset contains 224,316 chest x-rays (PA, AP, and Lateral) from 65,240 patients. We used 223,414 JPEG formatted x-rays from the training set of the dataset. 
MURA dataset contains 40,561 musculoskeletal x-rays from 14,863 studies. 
Like CheXpert, we used 36,808 x-rays from the training set and utilized study labels, e.g., finger, wrist, hand, forearm, elbow, humerus, and shoulder, for analysis.




§ EXPERIMENTS


In this section, we start by analyzing embeddings of chest x-rays (ChestX-ray14 and CheXpert datasets) and find different types of outliers. 
Then, we expand on variations of our method by altering the neural architecture, pre-training dataset, and embedding algorithms. Finally, we discuss outliers in musculoskeletal x-rays of the MURA dataset.





 §.§ Analyzing Chest X-ray Datasets




  §.§.§ Lateral X-rays in ChestX-ray14:








Our 2-D embedding (Fig. <ref> (a)) shows two large clusters of PA and AP views, corresponding to the primary topology of the DenseNet features. 
The embedding also includes a few satellite clusters around these large ones. 
These satellite clusters around the larger ones occur because the nearest neighbor graph creates a loop (or isolated sub-graph) of common features that are distinct from the rest of the data. 
In most cases, each of the satellite clusters of x-rays is from a single patient with a unique signature.
However, if any specific image features (such as similar artifacts in multiple images) are present in x-rays of different patients, these can create satellite clusters as well. Another interesting structure in Fig. <ref> (a) is the protruding region from the AP cluster.


Representative examples of anomalous clusters are shown in Figs. <ref> (b)-(e). 
The most surprising finding is the existence of some lateral x-rays in the dataset (Fig. <ref> (b)), as this dataset is supposed to be composed of frontal chest x-rays only. 
We found 92 lateral x-rays using our method. 

The protruding region from the AP cluster marked c and d 
(in Fig. <ref> (a)), consisting of x-rays with dark borders of PA and AP views, respectively. Finally, cluster (e) shown in Fig. <ref> (e) groups 46 x-rays from patient ID 9845, and a single x-ray from patient ID 12562.



  §.§.§ Corrupted Images in CheXpert:


Figure <ref> (a) shows the 2-D embedding of CheXpert dataset. 
As before, the large PA and AP clusters form the bulk of the mapping. The lateral x-rays also form a separate large cluster. 
A few of the large satellite clusters (b-f) have been marked by red circles in Fig. <ref> (a). Four images from each of the clusters are plotted in Fig. <ref> (b)-(e). 
Figure <ref> (b) depicts images with block artifacts, e.g., from poor JPEG compression or accidental splicing. 
We found 107 such images in this cluster. 
Figure <ref> (c) depicts images that are just noise (19 images). 
Figure <ref> (d) shows images with block artifacts and dynamic range issues (53 images). 
Figure <ref> (e) shows x-rays with vertical artifacts (88 images).
Finally, Fig. <ref> (f) shows rotated images. 
This cluster is placed near the large cluster of lateral (L) x-rays. 
Thus, DenseNet considers rotated x-rays to be more similar to lateral images than upright frontal x-rays.








 §.§ Effect of Different Pre-trained Networks


In the previous section, we used DenseNet-121 trained on ImageNet data for feature extraction (Fig. <ref>). 
However, using ImageNet may introduce domain shift, since the dataset is not specific to medical images. 
To assess, we looked at embeddings from different models trained on ImageNet and chest x-rays (Fig. <ref>). 

ImageNet models (Fig. <ref> (a,b)) create similar clusters of PA, AP, and L x-rays, whereas Chest-Xray14 models (Fig. <ref> (c,d)) merge the PA and AP clusters. 
The x-ray model embeddings also contain fewer satellite clusters, and their distinctness is lost. Thus, outlier detection fails. 
Additionally, we highlight the placement of images with vertical artifacts from the cluster (e) of Fig. <ref> in the alternate models (yellow points in Fig. <ref>). 
While both ImageNet models successfully separate these images (with a few scattered around), all the chest x-ray models fail (where the outliers are inside the large cluster and scattered throughout). 

A major difference between models is that ImageNet models use softmax (n-ary) cross entropy loss (for non-overlapping labels), but chest-x-ray models use binary cross entropy loss (for overlapping labels). 
To resolve this discrepancy, we trained chest x-ray models using non-overlapping labels of the ChestX-ray-14 dataset (for details, see Appendix <ref>). 
The resulting embeddings (Fig. <ref>) show a similar characteristic of the chest x-ray models with overlapping labels (Fig. <ref> (c,d)), in that the individual views are weakly separated while the specific outlier images fail to form the satellite clusters. 
For example, the outliers with vertical artifacts scatter within the 2D mapping. The networks consider them as any other x-ray image and ignore the artifacts. 
This result strengthens the idea that training on ImageNet (or a more general computer vision task) that has broader exposure benefits the discovery of outlier images. 



 §.§ Comparing Various Embedding Algorithms


To assess the effectiveness of UMAP, we compared it with a few other dimensionality reduction techniques, specifically principal component analysis (PCA), t-SNE, TriMap <cit.>, and PaCMAP <cit.>.

The PCA embedding in Fig. <ref> (a) shows the top two directions of largest variances. There are two clusters but the AP and PA views overlap. The nonlinear dimensionality reduction algorithms - t-SNE, UMAP, and variants - discover more features and make the clusters more distinct. 

The default t-SNE is tuned to preserve the neighborhood as best as possible. This often causes the individual clusters to spread out and be less compact (Fig. <ref> (b)). This behavior is apparent in the PA, AP, and lateral x-ray clusters: the separation among them is minimal, and there is little room for the satellite clusters. 
However, t-SNE can be tuned to produce a more UMAP-like output. 
Following the findings of Linderman et al. <cit.> and Bohm et al. <cit.>, we used an exaggeration factor of 4, which means we applied four times more attractive force than the repulsive force throughout the optimization procedure. 
The standard early exaggeration factor of 12 was also applied at the start of the optimization. The resulting plot in Fig. <ref> (c) largely resolves the PA and AP clusters, but there are fewer satellite clusters than the UMAP output.

Both PaCMAP <cit.> and TriMap <cit.> aim to preserve the global structure of the data while keeping the clustering properties of UMAP. PaCMAP modifies the pairwise relation of UMAP by considering different neighborhoods at different scales (near and far), while TriMap achieves this with the triplet constraint. These methods obtain the clustering of PA, AP, and lateral views (Figs. <ref> (d,e)), but, similar to exaggerated t-SNE, the satellite clusters are largely absent.

The yellow dots show the placement of x-rays with vertical artifacts from cluster (e) of Fig. <ref> for the different embedding algorithms. In all the alternate embedding algorithms, these images are scattered within the other x-ray images. The outlier images fails to be identified and labeled separately, which demonstrates a superior performance of UMAP.



 §.§ Extracting Mislabeled X-rays from MURA




Since the MURA dataset consists of x-rays from different parts of the arm and the shoulder, there is a natural ambiguity in the labels, e.g., both wrist and hand x-rays may contain the hand of a person, and shoulder x-rays may contain part of the chest. 
In such cases, finding mislabeled x-rays by embedding all the images may be sub-optimal.
To find outliers more directly, we searched for misclassified images by explicitly using labels of the dataset. The method has two parts: 1) introduce target images with a specific label (preferably from a different dataset than the MURA one); and 2) perform neighbor embedding on the joint dataset.

For example, to look for possible chest x-rays that are falsely classified as finger x-rays, we added 100 chest x-rays from the CheXpert dataset to the 5,106 finger x-rays of MURA. 
We then applied the UMAP to the composite set (Fig. <ref> (a)). As shown in Figs. <ref> (b-d), the seeded chest x-rays acted as an attractor for mislabeled images in MURA, with x-rays labeled `finger' now appearing in the (new) chest cluster. 
Interestingly, both of these x-rays were from patient 04547 (another 3 from this patient were labeled correctly). We describe an additional experiment in Appendix <ref>.



§ CONCLUSION


Neighbor embedding algorithms can be an effective tool for summarizing datasets and identifying outlier images. The principle of the method is that the outliers are different from the main data but can have similarities among themselves. 
Thus, the different outlier types form distinct clusters in the embeddings. 
Our experiments, using a DenseNet-121 feature extractor and UMAP neighbor embedding method on the ChestX-ray14, CheXpert, and MURA datasets, distinguished different radiological views of chest x-rays, classified differences, and identified wrongly labeled or corrupted images.  We further found specific types of outliers by seeding the dataset with target images and performing neighbor embedding.


While this study performed retrospective analysis of large x-ray datasets, outlier curation can be achieved during the initial assembly of the dataset as well. 
For suspected outliers, the method of seeding data with known labels can be applied. To streamline the process, appropriate reference datasets may be created beforehand. Undoubtedly, cleaner input data will result in cleaner output data. For larger datasets, more accurate results and faster embedding may be achieved by dividing them into smaller subsets and applying better alignment techniques <cit.>. Finally, since the methods are graph-based and agnostic to the underlying data type, all of the methods here can be applied to arbitrary datasets, including and especially those that are mixed modality.




  §.§.§ Acknowledgement
 The authors gratefully acknowledge financial support from the Schmidt DataX Fund at Princeton University made possible through a major gift from the Schmidt Futures Foundation.



  §.§.§ Disclosure

The authors have no competing interests to declare that are relevant to the content of this article.










 unsrt
 



§ IMPLEMENTATION DETAILS




 §.§ Image Pre-processing

We apply the following transformations: histogram equalization <cit.>, resizing the images, center cropping, and normalization. Images are resized such that the lowest dimension contains 256 pixels and then center-cropped to a 224×224 dimensional image for feature extraction. 
Then the images are normalized according to the specification of ImageNet: mean (0.485, 0.456, 0.406) and standard deviation (0.229, 0.224, 0.225) of red, green, and blue channels, respectively.



 §.§ Feature Extraction

We use a DenseNet-121 architecture <cit.> pre-trained on the ImageNet dataset <cit.> from the PyTorch deep learning library <cit.>. (DenseNet is a deep neural network with many inter-layer connections designed to reduce the numerical instabilities that originate due to the depth of the network. The usage of neural network outputs as features is an effective baseline in machine learning algorithms <cit.> and is extensively used in medical image analysis <cit.>.)  Here, we remove the classification (softmax) layer from the neural network and use the output as features. 



 §.§ Parameter Settings of Embeddings




In general, we kept the number of nearest neighbors k to be low, as increasing k increases the computational budget exponentially.

ChestX-ray14 - Fig. <ref>: For ChestX-ray14 we used k=50. The minimum distance parameter m_d was set to 0.1. We experimented with smaller m_d values to increase the separation of the AP and PA x-rays, but it had little effect. The embedding was optimized for 200 epochs.

CheXpert - Figs. <ref>, <ref>, and <ref>: The embedding was obtained by using k=10. We used a smaller minimum distance m_d=0.001, as we found that this value provided a better separability of the large clusters. We ran the optimization for 300 epochs.
Figure <ref>: (b) t-SNE: early exaggerated for 250 steps, then 500 regular steps, (c) t-SNE: early exaggerated for 250 steps, then 500 exaggerated steps, (d) TriMap: 1000 steps with k=10, compactness parameter =0.1, and (e) PaCMAP: 450 steps with k=10.

MURA - Figs. <ref>, <ref>, and <ref>: Similar to CheXpert, we used k=10, m_d=0.001, and ran the optimization for 300 epochs.



§ PRE-TRAINING MODELS USING CHESTX-RAY14 DATASET


We trained several chest x-ray models using the ChestX-ray14 dataset. We trained traditional x-ray models using binary cross entropy loss for every disease (one vs all model) by following <cit.> which matches with the settings of the ImageNet pre-training. Additionally, to test whether the loss function and multilabel training have major effects on the resulting embeddings, we constructed a dataset by removing overlapping labels (Table <ref>) and trained two models using binary and softmax cross-entropy loss, respectively. For each mode, we optimized for 400,000 iterations using Adam (learning rate =0.001) with a batch size of 30 (2 images per disease label).





 §.§ Comparison to Publicly Available Chest X-ray Models



We compare our DenseNet-121 model to that from TorchXRayVision library <cit.>. It differs from ours in that the image pixels are normalized to [-1024,1024], images are converted to (or read as) grayscale, and `No Finding' label is untrained. 
This model (Fig. <ref>) agrees with our results and ignores the artifacts of the outlier images, showing embeddings are indifferent to alternate preprocessing scheme.



§ ADDITIONAL DISCUSSION ON MURA






Fig. <ref> shows the embedding of 36,808 musculoskeletal radiographs from the MURA dataset. There is decent separation among the x-rays in terms of the labels, but there is also a considerable overlap. For example, finger, wrist, and hand x-rays overlap, as finger x-rays include parts of the wrist and hand, and vice versa. Similarly, wrist and forearm clusters are often merged, since x-ray of the forearm tend to capture a portion of the wrist, and vice versa. The same happens for humerus and shoulder. The images within each cluster thus share similar acquisition views, aspect ratios, and specific features (e.g., circular window function, stitching of multiple x-rays in one image). 
Unlike ChexPert case, analysis of this mapping did not reveal any satellite clusters with corrupted images.

Based on this, we decided to focus on individual labels and extract images with specific labels. The chest x-rays from the CheXpert dataset used to produce Fig. <ref> are shown in Fig. <ref>.



§ EXTRACTING MISLABELED X-RAYS FROM MURA





In a different experiment on the MURA dataset, we used `finger' and `shoulder' x-rays (Fig. <ref>). 
The broad features of the finger and shoulder are easily separable with a few misclassified points. 
Analyzing `finger' x-rays misclassified in `shoulder' clusters, we can find the two chest x-rays labeled as finger (which we found in previous experiment as well) and two images that are just noise/non-x-ray images (Fig. <ref> (b)). 
The latter belongs to patient ID 04687, which uncovers two more outliers (one is an x-ray of keys). The misclassified `shoulder' labels in the `finger' cluster reveal three leg x-rays (Fig. <ref> (c)).



 Paper Title: Analyzing the impact of semantic LoD3 building models on image-based vehicle localization.




































	1 Professorship of Photogrammetry and Remote Sensing, 
    2 Chair of Engineering Geodesy, 

 TUM School of Engineering and Design, Technical University of Munich, 80333 Munich, Germany 
 - (antonia.bieringer, olaf.wysocki, christoph.holst)@tum.de

 3 3D Mapping Solutions GmbH, 83607 Holzkirchen, Germany - sebastian.tuttas@3d-mapping.de

4 Hochschule München University of Applied Sciences, 80335 Munich, Germany - ludwig.hoegner@hm.edu






















Numerous navigation applications rely on data from global navigation satellite systems (GNSS), even though their accuracy is compromised in urban areas, posing a significant challenge, particularly for precise autonomous car localization.

Extensive research has focused on enhancing localization accuracy by integrating various sensor types to address this issue.
This paper introduces a novel approach for car localization, leveraging image features that correspond with highly detailed semantic 3D building models.

The core concept involves augmenting positioning accuracy by incorporating prior geometric and semantic knowledge into calculations. 
The work assesses outcomes using Level of Detail 2 (LoD2) and Level of Detail 3 (LoD3) models, analyzing whether facade-enriched models yield superior accuracy. 
This comprehensive analysis encompasses diverse methods, including off-the-shelf feature matching and deep learning, facilitating thorough discussion. 
Our experiments corroborate that LoD3 enables detecting up to 69% more features than using LoD2 models.

We believe that this study will contribute to the research of enhancing positioning accuracy in GNSS-denied urban canyons. 
It also shows a practical application of under-explored LoD3 building models on map-based car positioning.




Analyzing the impact of semantic LoD3 building models on image-based vehicle localization
    Antonia Bieringer1, Olaf Wysocki1, Sebastian Tuttas3, Ludwig Hoegner4, Christoph Holst1,2
    
=============================================================================================






 





§ INTRODUCTION

Navigating in complex urban environments can pose challenges for positioning systems, primarily when vehicles depend on GNSS. As vehicles move through signal-obstructed urban canyons, the effectiveness of positioning systems can diminish. 
Consequently, vehicles shall turn to alternative cues to determine their location in such scenarios.

Various methods have been developed to address this issue, combining different sensors to enhance location accuracy. 
While GNSS enables global positioning, supporting it with local cues is an intuitive choice. 
For example, directly enhancing GNSS signals by using 3D city models <cit.>.
More frequently, however, the researchers have focused on utilizing cameras for enhanced positioning owing to their abundance on cars and optical features.





In order to localize the vehicle with optical cameras, features have to be found in the images. There are various techniques for feature finding and matching, e.g., SIFT <cit.>, SURF <cit.>, or ORB <cit.>. 
Over time, researchers compared those methods against each other <cit.> and improved them.
However, to globally position the images, a reference 3D map with the correct scale and position shall be used.

Such 3D map can be extracted directly from governmental and proprietary geoportals <cit.>, where semantic 3D building models are available.
These models are used in different Level of Details (LoD). LoD1 models are limited to geometrical information that includes the vertices and edges of the buildings; they are cuboid polyhedral models with flat roofs. 
LoD2 models are more detailed and include the generalized structure of the roof. 
In comparison, LoD3 is enhanced by facade details such as windows and doors <cit.>. 

A few papers have already investigated the usage of LoD1 and LoD2 models for vehicle localization.



One of the examples of utilizing LoD1 and LoD2 models is the paper of Vogel et al., who introduce a localizing approach in the inner-city outdoor environments using 3D city models <cit.>. 
To this end, they modify an iterated extended Kalman filter to be able to integrate additional information into it, e.g., available maps. 
A more reliable and precise georeferencing is achieved since geometric circumstances are considered <cit.>. 
Their idea is refined in later papers of <cit.>.
However, the usage of LoD3 models facade features for localization has yet to be investigated.


Our contributions are as follows:


 
  * Map-supported vehicle localization by using high fidelity LoD3 models and images.
 
  * Harmonizing two distinct modalities: Visual information encapsulated within optical images and the structure of semantic 3D models' non-textured surfaces.
 
  * Comparing the benefits of LoD3 over LoD2 building models for camera localization.




§ RELATED WORK


We devote our work to analyzing the impact of image-based vehicle positioning using semantics of 3D maps. 
Therefore, in this Section, we present the state-of-the-art in semantic 3D city models and image- and map-based positioning.



 §.§ Semantic 3D city models


At city, regional, and national scales, semantic 3D city models offer comprehensive representations of structures, taxonomies, and aggregations.
Semantic 3D city models are frequently described by the internationally adopted CityGML standard, established by the Open Geospatial Consortium (OGC) <cit.>; 
This data model standard adopts two encodings using either Geography Markup Language (GML) or CityJSON <cit.>.
The CityGML standard enables the modeling of urban objects with 3D geometry, appearance, topology, and semantics across four different Levels of Detail (LoD). 

Building descriptions are pivotal to semantic 3D city models due to the foundational role of urban dwellings in cities <cit.>. 
LoD1 and LoD2 building models, currently prevalent, boast around 220 million models available in countries such as Germany, Japan, the Netherlands, Switzerland, the United States, and Poland[1].
While LoD1 building models are represented by prismatic 3D models of height-extruded building footprints, LoD2 models additionally represent complex roof structures.     
The widespread adoption of these models is attributed to robust 3D reconstruction algorithms and the availability of building footprints combined with aerial observations <cit.>.

However, LoD1 and LoD2 models lack a more detailed description of facade elements.
Addressing this gap are LoD3 building models featuring descriptive facades with objects such as windows, doors, balconies, and underpasses <cit.>.
The current practice shows that these highly-detailed objects are manually modeled <cit.>. 
Yet, a great deal of research has been devoted to enabling automatic LoD3 reconstruction, which has sparked advancements in the LoD3 data adoption
[<https://github.com/OloOcki/awesome-citygml>].


The recently developed LoD3 reconstruction methods promise to increase the availability of such models <cit.>.



 §.§ Image-Based Vehicle Positioning





One way of reconstructing the 3D structure of a scene and estimating the camera poses from a collection of 2D images is Structure from Motion (SfM) <cit.>. 
This technique works by leveraging the relationships between features in different images and the scene's geometry to recover the scene's 3D structure and camera pose.

Another approach of estimating the vehicle's position by using a series of camera images is Visual Odometry (VO). 
It is advantageous over other methods due to its insensitivity to soil mechanics and lower drift rates <cit.>. 
Additionally, it performs well in places where a satellite connection is limited, e.g., indoors or on narrow streets. 
However, the trajectories estimated by visual odometry tend to drift along with increasing distance from the start point. Consequently, researchers investigated to improve the concept by expanding VO ideas to another image-based localization method called SLAM <cit.>.

Simultaneous Localization and Mapping (SLAM) reconstructs reality from observation data, e.g., images and points clouds. Compared to SfM and VO, it computes the trajectory in real time and maps the vehicle's environment <cit.>. The map is updated while driving <cit.>, and consequently, the current point is automatically corrected by previously calculated results. 
A prerequisite for the map-supported SLAM are surroundings that are easy to detect and can be used to scale and correctly map the trajectory <cit.>. 
As Lucks et al. show, the prior 3D building map can enhance the SLAM-based calculated trajectory <cit.>. 



§ METHODOLOGY

In Figure <ref>, the workflow of the developed method is shown. 
First, virtual pictures of the building models are created. The resulting images are then used to find features in the corresponding optical images. 
For this purpose, the ORB detector <cit.> is used. The generated 2D-image coordinates are correlated with a 3D coordinate by benefiting from the LoD model semantics and geometry. In the end, these coordinates are needed to estimate the camera position for every single image with spatial resection. 
This workflow is designed for LoD1, LoD2 and LoD3 building models. 
The developed method's code is published in the GitHub repository[<https://github.com/tum-pf/LoD3ForLocalization>].





 §.§ Creating Virtual Images of the LoD Models


For the first step, the virtual images are created by using ray casting. To work with the ray casting algorithm, a polyhedral building model is tranformed into a mesh model using triangulation since all results are based on the hit of a triangle by virtual rays. 
The algorithm creates those virtual rays and sends them through a virtual camera, preferably in the direction of the model. Every time a ray hits a point, the corresponding triangle and the barycentric coordinates of the exact hit point in the triangle are saved. Therefore, the 3D coordinates in the real world can be calculated by only knowing the 2D coordinates of the virtual image. The calculation is described in Section <ref>.

Moreover, the virtual camera's pose is required (see Figure <ref>). As an approximation solution, the position is taken from the GNSS data of the current corresponding optical image. 
On top, the camera's height above the GNSS antenna must be considered, leading to GNSS values adapted to the camera position. 
The camera is pointing to the adapted GNSS position of the consecutive image. 
Afterward, the orientation of the camera requires to be adjusted: The angles for roll, pitch, and yaw are required to approximate the point close to the GNSS position in consideration of the camera tilt. 
A visual example of pitch can be seen in Figure <ref>. The approximation is calculated as follows:

    dy [m] = r_GNSS· dy  [^∘] ·π/180 ^∘


An overview of the above-described virtual setup can be seen in Figure <ref>.


After applying the ray casting algorithm, the results are available in a tensor. There are five results: the distance, the geometry IDs, the primitive IDs, the primitive normals, and the barycentric coordinates. For an example of LoD2 combined with one LoD3 building, see Figure <ref>:


From these results, the primitive-normals-data is saved as an image. As shown in Figure <ref>d, some walls of the buildings have the same color as the background. To make every plane of the building visible, the values of the primitive normals are adjusted so that only the absolute values are considered since the exact value is unimportant for the virtual images.
This adjustment is necessary so that every side of the building has a unique color. Otherwise, some buildings' planes will not be considered during feature matching. 
Note that the LoD3 building in the virtual image in Figure <ref> is more detailed due to the additional information on the model, especially regarding windows and roof eaves.









 §.§ Feature Matching


An overview of the approach is visualized in Figure <ref>. The feature matching is performed on images with reduced details. Therefore, the images become more abstract and only focus on important information. 
Then, the images are matched against themselves: the virtual and the optical image. 

The next step is done for the image of the real world, as well as for the one of the virtual camera. 
First, another abstract, virtual feature image is created and implemented as a matrix with the same values for every pixel. If features were found in matching the images against themselves, those pixels are colored differently. This is realized by assigning a different value to this pixel in the matrix. The image, therefore, only contains information about the essential points of the image before. 

Afterward, these newly created feature images are matched with an ORB-feature detector again. 





 §.§ Correlating the Features with a 3D-Coordinate
 

During ray casting, every hit of the mesh by a ray is saved. 
In this case, the corresponding triangle IDs are stored. 
Then, the corresponding triangles in the mesh are found, which enables the extraction of which vertices belong to a triangle that was hit. 
From the IDs of those vertices, the 3D coordinates from the vertex positions are extracted in the mesh.
Finally, the exact point needs to be weighted by the vertices of its triangle because only the relative positions from the vertices are known <cit.>. Those relative positions are named u, v, s. 
The results only return two barycentric coordinates for each hit point since the sum of the three coordinates equals one:

    1 = u + v + s


This means that s can be directly derived from u and v.
The weighting to receive the final coordinates works as follows:

    P = u*P_vertex_1 + v*P_vertex_2 + s*P_vertex_3




 §.§ Calculating the Trajectory


The vehicle's trajectory is simultaneously calculated by driving and reading the generated images. To this end, we leverage the spatial resection concept. Spatial resection is a set of statistical formulas to compute the accuracy of observations by finding a model that fits them best. Consequently, the accuracy of the observations and related values can be calculated afterward. In this case, the camera position is estimated while the observations are optimized at the same time. The position calculation is performed with the help of the photogrammetrical co-linearity equations, which can be seen in <ref> and <ref>.


    x = x̂_̂0̂ + z (r̂_̂1̂1̂(X-X̂_̂0̂) + r̂_̂2̂1̂(Y-Ŷ_̂0̂) + r̂_̂3̂1̂(Z-Ẑ_̂0̂)/(r̂_̂1̂3̂(X-X̂_̂0̂) + r̂_̂2̂3̂(Y-Ŷ_̂0̂) + r̂_̂3̂3̂(Z-Ẑ_̂0̂)


    y = ŷ_̂0̂ + z (r̂_̂1̂2̂(X-X̂_̂0̂) + r̂_̂2̂2̂(Y-Ŷ_̂0̂) + r̂_̂3̂2̂(Z-Ẑ_̂0̂)/(r̂_̂1̂3̂(X-X̂_̂0̂) + r̂_̂2̂3̂(Y-Ŷ_̂0̂) + r̂_̂3̂3̂(Z-Ẑ_̂0̂)


For the calculation, six values are to be determined, representing the camera's position (X_0, Y_0, Z_0) and orientation (ω, ϕ, κ).
Consequently, for the solution, at least six corresponding points are required. More corresponding points are preferred to get even better accuracy. 

The spatial resection is performed as follows. A direct linear transformation generates the approximate camera orientation values. Input parameters are at least six corresponding points. The approximate camera position is taken from the corresponding GNSS point. The approximate values are required because the spatial resection needs starting points to optimize. Otherwise, the algorithm might run into a local minimum and keeps iterating in an infinite loop. 

In the optimization loop, the derivations of the A-matrix are computed. This is necessary because the refinements of the observations for the best-fitting model are calculated with the help of this matrix. The A-matrix has the size of [n x m], while n is the number of observations and m is the number of values to be determined.

The values that are to be determined are estimated in every iteration by calculating the following:

    δx̂ = (A^T*P_bb*A)/(A^T*P_bb*w)

Where A is the already described A-matrix. P_bb is a matrix that contains weights. Those weights decide how important an observation is, e.g., if the measurement is probably more precise than others because of the geometrical position. The w-vector contains the current errors between the approximated δx̂ and the actual observations. The optimized observation values are computed by:

    δv̂ = A*δx̂ - w




§ EXPERIMENTS
 


The developed method is evaluated in three different test areas, which are visualized in Figure <ref>. 


Notably, there were no instances in the test area where LoD3 models were present on both sides of the street, which mimics a frequently present real-world situation. 
The missing LoD3 models were supplemented from the LoD2 dataset as a solution. 
Therefore, the same models from the LoD2 dataset were used for experiments involving solely LoD2 models. 
This hybrid approach ensures a balanced distribution of potential reference points and promotes an effective configuration. 
Owing to the experimental setup, we were able to directly evaluate the influence of the LoD3 over to the usage of only LoD2 models.

The test site was located in Munich, Germany, and represented a Central European architecture comprising a mixed residential area and a university campus.
The images were acquired by a mobile mapping vehicle of 3D Mapping Solutions GmbH <cit.>.
The LoD2 models were downloaded through the governmental open-geodata portal of Bavaria, Germany <cit.>, whereas the LoD3 models were developed within the interdisciplinary tun2twin project <cit.>.




 §.§ Different Approaches for Feature Matching


The feature matching was performed on five different pillars. Each approach is described in the following subsections.



  §.§.§ Corresponding Image-Pairs of Different Types


 


This approach aims to find features from image pairs that are not the same type. This means that one virtual image was matched with one image of reality. Both images were taken from the same GNSS point. The results show only a median of 9-18 suitable feature matches for each image pair. 



  §.§.§ Segmented Images with Deep Learning
 

 


Matching real images with virtual images presents a challenge due to disruptive elements not being shared between the two image types. To mitigate this issue, the approach involves segmenting the real images. 
The used algorithm was built on the SegFormer network <cit.>, using the MMSegmentation framework <cit.>. 
This algorithm generated an output image where distinct elements (e.g., buildings) were each assigned a unique color.


To match these semantically segmented images, a method is devised wherein the segmented images act as masks for the real images to isolate the building areas from the real images. 
Consequently, only the building-related parts of the real image are retained and matched with their virtual counterparts. 
This strategic approach effectively minimizes any potential disruptions from extraneous objects that could potentially result in erroneous matches.



  §.§.§ Sobel-Filter and Canny Edge Detector


 


An established approaches for feature matching focuse solely on features along object edges, in this case, including structures such as window edges, e.g., Sobel filters and Canny edge detector <cit.>. 

In instances where it was desirable to exclusively match buildings, applying a bounding box to actual images proved effective. 
These filters were harmonized with the segmented images detailed in Section <ref> to ensure precision. 
This strategic combination ensured that only relevant objects for vehicle localization were considered. 



§ RESULTS AND DISCUSSION


The experiments were evaluated by comparing the number of found features and the standard deviation of the camera position. 
For every table, the corresponding gain was computed as a percentage increase or decrease between LoD2 and LoD3. 



 §.§ Number of found features


The different approaches led to different results regarding the number of found features, which represented those that can be successfully matched between two images. 
The selection pertains to the median of all tested images within each area. The best result for each experiment is written in bold, comparing LoD2 and LoD3. An example of the found features in the images is visualized in Figure <ref>.



A comprehensive comparison for each model, namely LoD2 and LoD3, is presented across different test areas in Tables <ref>, <ref>, and <ref>.







As our experiments corroborate, the LoD3 models, on average, result in the same or more found features between the images. We deem the feature images algorithm as the most reliable feature extraction between semantic 3D building models and optical images, as it scored more features on average in the three test areas (60%, 69%, and 58%, respectively). 

In the test area 1 with prevalent LoD2 models, the algorithm matched features for LoD2 until the point where the whole virtual image was facing the model-absent underpass. 
Then, no feature matching was possible again and the algorithm failed.
This happens due to the frequently absent underpasses in LoD2, where instead a vertical wall is rendered to the ground level, as seen in Figure <ref>.



Especially in test area 2, more features lead to a more reliable estimation of the trajectory point: up to 69% more features (see Table <ref>). 

As described in Section <ref>, test area 3 exhibited the optimal situation for localizing the vehicle with the help of high-detail polyhedral models. 
This can be seen in Table <ref>, where 57% more features were found, and seven of the tested methods achieved an increase of found features.
We also observe that the higher the buildings, the more details can be found for feature matching.
Furthermore, narrow streets are advantageous because the buildings are closer to the camera and can be captured in more detail. 
In contrast, GNSS data is most accurate in urban environments when surrounded by only low buildings; narrow streets lead to positionining deviations. 
Consequently, the localization with images and LoD3 models can be advantageous over using GNSS in urban areas. 
Both methods are like counterparts: where one method fails, the other has advantages.



 §.§ Positioning accuracy


To analyze the advantage can be achieved by LoD3 models, Tables <ref>, <ref>, <ref> show the accuracy of the estimated camera position in comparison.







Regarding test area 1, the accuracy of the feature images algorithm drops to [-37%, -202%, -26%]. This shows, that the algorithm is not yet working perfectly for every situation. The configuration with the underpass leads also in LoD3 to problems in the accuracies. Still, other methods show the advantages of LoD3 also in this case. The tackling of the problems is described in Section <ref>.

The developed algorithm works less accurately in urban areas when buildings are available on only one side of the street, see Tables <ref> and <ref> (test area 2), especially for the Feature images algorithm. 


As seen in Tables <ref>, <ref> and <ref>, LoD3 reaches a higher number of features. More found features provide a better distribution in the image and, therefore, result in a better standard deviation of the camera point, e.g., up to [88%, 70%, 79%] more accurate in test area 3 (see Table <ref>).  
Considering the standard deviation of the camera position for the evaluation, LoD3 is favorable in test area 3.



 §.§ Limitations
 
The developed algorithm shows the advantages of using LoD3 over LoD2 models. Still, the deviations of the estimated point from the corresponding GNSS point are large. 
To use those models for positioning, the feature matching has to be optimized so that matched features refer to exactly one single point in the real world. 
This is not the case for all matched features so far.

In general, the deviation of the calculated camera position performs slightly better in LoD3 compared to LoD2. 
A reason for the deviation from the GNSS point is an imprecise calculation of the 3D coordinates. 
The exact coordinates are weighted by the barycentric coordinates in the triangle. 
Consequently, each triangle vertex belongs to exactly one of the barycentric coordinates u, v, s. 
This depends on the sequence in which the triangle coordinates are saved in the mesh. The order of the vertices is a requirement currently investigated for LoD2 and LoD3 models. 
The topology for such models has to be consistent in every model, e.g., in the right-hand rule. 
This is not yet common for all developed LoD3 models since the automatic generation remains in its infancy and is prone to manual, randomly induced errors. 
These errors lead to an imprecise calculation of the 3D coordinates and, therefore, to a deviation of the calculated camera position from the GNSS point. 
The method developed in this paper works assuming that the models are correct and the vertices are saved in the right-hand rule. 



§ CONCLUSION


In this paper, we propose an approach to compare the performance of LoD2 against LoD3 models for map- and image-based vehicle positioning. 
Using a feature image method, we harmonize semantic 3D city models with optical images. 
Our experiments corroborate that although we had different representations, optical images vs. quasi-optical images of the LoD models, our method performs well in feature matching. 
The presented method performs best when buildings cover both sides of the street.
We also observe that the higher the building models, the more details can be found for feature matching.
Additionally, narrow streets are advantageous because the buildings are closer to the camera and can be captured in more detail. 
We are convinced our image and model-based navigation algorithm can inspire upcoming work. 
In the future, we plan to improve upon measurements, e.g., by using SLAM. With this, we hope to decrease the large deviations in our calculations from the GNSS points. Additionally, we will work on reducing false feature matches to achieve higher accuracy. 
We also plan to do more experiments on more datasets and aim to expand the LoD3 datasets.



§ ACKNOWLEDGMENT


The authors would like to thank 3D Mapping Solutions GmbH for providing the images to develop and test the algorithm.
This work was supported by the Bavarian State Ministry for Economic Affairs, Regional Development and Energy within the framework of the IuK Bayern project MoFa3D - Mobile Erfassung von Fassaden mittels 3D Punktwolken, Grant No. IUK643/001.
Moreover, the work was conducted within the framework of the Leonhard Obermeyer Center at the Technical University of Munich (TUM).
We gratefully acknowledge the tum2twin team at the TUM for the valuable insights and for providing the CityGML datasets.









	1.17
		 
	

















































































 Paper Title: Prompting Medical Large Vision-Language Models to Diagnose Pathologies by Visual Question Answering.



Einstein Probe discovery of a super-soft outburst from CXOU J005245.0-722844:
 a rare BeWD binary in the Small Magellanic Cloud
    [
    
===============================================================================================================================



§ ABSTRACT 
Large Vision-Language Models (LVLMs) have achieved significant success in recent years, and they have been extended to the medical domain. Although demonstrating satisfactory performance on medical Visual Question Answering (VQA) tasks, Medical LVLMs (MLVLMs) suffer from the hallucination problem, which makes them fail to diagnose complex pathologies. Moreover, they readily fail to learn minority pathologies due to imbalanced training data. We propose two prompting strategies for MLVLMs that reduce hallucination and improve VQA performance. In the first strategy, we provide a detailed explanation of the queried pathology. In the second strategy, we fine-tune a cheap, weak learner to achieve high performance on a specific metric, and textually provide its judgment to the MLVLM. Tested on the MIMIC-CXR-JPG and Chexpert datasets, our methods significantly improve the diagnostic F1 score, with the highest increase being 0.27. We also demonstrate that our prompting strategies can be extended to general LVLM domains. Based on POPE metrics, it effectively suppresses the false negative predictions of existing LVLMs and improves Recall by approximately 0.07.



    Medical Visual Question Answering, Large Vision-Language Models, Prompt Engineering




§ INTRODUCTION


Research on Large Language Models (LLMs) has yielded astonishing results in recent years. LLMs with billions of parameters have achieved outstanding abilities in a wide range of application scenarios <cit.>. The success of LLMs has quickly extended into the Vision-Language (VL) domain. Large Vision-Language Models (LVLMs) are built upon LLMs by training adapters that project visual features into features that can be interpreted by LLMs <cit.>. Visual Question Answering (VQA) is an essential skill of LVLMs, and VQA accuracy serves as a test metric for most of these models <cit.>. LVLMs have been pretrained on medical datasets <cit.> and they have been tested on medical VQA tasks <cit.>. These Medical LVLMs (MLVLMs) have been able to answer questions regarding the imaging modalities, organs, and abnormalities depicted by the input medical scans.

Unfortunately, however, “hallucination” has been a major problem for LVLMs. This refers to the generation of content that is contradictory to the input images. Hallucination can be measured via VQA. One may ask the model questions regarding the existence of objects in the input image(s) and the hallucination level is assessed as the percentage of correctly answered questions. VQA can also potentially serve for medical image diagnosis. Users pose questions regarding a pathology and the MLVLM responds based on its analysis of the medical scans. However, most of the available datasets involve simple questions such as “what is the modality of this image” and “what is the organ/tissue in this image”. MLVLMs have yet to be thoroughly evaluated on VQA accuracy across a broad spectrum of pathologies. Additionally, general VQA models are usually tested by the commonly known accuracy: the percentage of correctly answered questions, which is an unsuitable measure for medical VQA.
Medical image classification metrics such as the Precision, Recall, and F1 are more suitable for the evaluation of medical VQA models. Several strategies have been explored to enhance the question answering of LLMs/LVLMs, including chain-of-thought prompting <cit.>, self-consistency <cit.>, and retrieval-based augmentation <cit.>. All these methods involve fine-tuning the models, which is expensive. Training-free methods to improve the VQA accuracy of MLVLMs are desirable.

For MLVLMs, hallucination is exacerbated by imbalanced training data. Many pathologies are minority categories in medical datasets. Models trained on large-scale medical data may easily fail to learn the features of less common pathologies. Addressing data bias typically involves strategies such as including more data of better quality, but given the scarcity of medical data, significantly enlarging the dataset may not be feasible. Common remediations involve re-sampling the data such that the positive and negative cases are better balanced, but this poses challenges when the data involves multiple categories of pathology. Additionally, re-sampling may undermine the training needs of LVLMs, which generally require large quantities of data. These problems highlight the need of a cost-effective approach to navigate the problem of minority categories in datasets.

Our study focuses on the VQA abilities of MLVLMs. In particular, we test an existing MLVLM, LLaVA-Med <cit.> for chest X-ray VQA across 5 categories of pathologies. The results show that the model has low accuracy, especially on minority pathologies. To enhance its VQA accuracy, we propose two prompting strategies. The first involves enriching prompts with detailed explanations of the queried pathology. The explanations include how the queried pathology is defined and how it appears in images. Our second strategy involves introducing an auxiliary weak-learner model as another agent. We train a small image classifier and fine-tune it to identify negative images accurately. Then, the negative predictions of this classifier are appended to the prompt as a reference for the MLVLM.

We run our experiments on the MIMIC-CXR-JPG <cit.> and Chexpert <cit.> datasets. The results show that our prompt strategies improve the F1 score significantly in most pathology categories (highest +0.27). We also show that our weak-learner-prompting strategy is applicable to the general domain. It reduces the false negative predictions of general domain LVLMs and improves the Recall by around 0.07 according to POPE metrics <cit.>.

To summarize, our contributions include the following:

    
  * We improve the VQA accuracy of MLVLMs by prompting with detailed explanations of pathologies.
    
  * We introduce a low-cost weak learner model as a reference for LLaVA-Med, and this effectively reduces the false positive (FP) answers.
    
  * We show that our second prompting strategy can be extended to general domains to help models adapt to specialized accuracy needs.

<ref> reviews related work, <ref> describes our methodology, <ref> presents our empirical study and its results, and <ref> draws conclusions from our research.



§ RELATED WORK





  
LVLMs and VQA

LVLMs are built upon LLMs. A pretrained visual encoder extracts the visual features and an adapter module projects the extracted features to ones that can be understood by the LLM. Models of this type include those by <cit.>, <cit.>, and <cit.>. During training, the visual encoder and the LLM are usually fixed. VQA is an essential skill of LVLMs. Given an input image, the models should be able to answer questions correctly regarding that image.



  
Hallucination in LVLM VQA

The hallucination problem usually refers to the LVLM generating a response that is not consistent with the input image. For VQA, in their generated answers the models may make mistakes on object presence, location, attributes, or the mutual relationship between objects. <cit.> find that frequently occurring objects are easily hallucinated by LVLMs, in that they tend to mention such objects even if it they are absent in the image. <cit.> and <cit.> show that LVLMs sometimes presume the assumptions in questions are true and easily give wrong answers when asked about some objects not in the given image.



  
Causes of LVLM VQA Hallucination

Hallucination can result from bias in the training data, missing fine-grained visual features, and LLM decoding strategies <cit.>. For data bias, the imbalanced distribution of data is an important aspect. When most of the answers to a question in the training data are “Yes”, the model tends to answer “Yes” to that question. Missing fine-grained visual features usually result from the pretraining of the visual encoder. Most LVLMs use the visual encoder of CLIP trained through contrastive learning. The encoder mainly focuses on salient features and ignores fine-grained features <cit.>. LVLM decoding strategies mostly choose the next word as the one having maximum conditional probability given previous text and the input image. This can lead to hallucination when the model overly relies on the knowledge learned in its training texts. Other causes include model simplicity and insufficient attention <cit.>.



  
Mitigation of LVLM VQA Hallucination

Strategies to mitigate hallucination in LVLMs mainly fall into two categories: prompt engineering and model improvement. Regarding the former, <cit.> leverage visual instructions constructed from the bounding box information in the input image to prompt LLMs. <cit.> use a chain of thought scheme to prompt the models to perform step-by-step visual-language reasoning like humans, which eventually leads to the correct answers. <cit.> generate multiple chains of thought and use the one with the majority vote as the answer. <cit.> prompt the model with explanations of the terms in questions. With regard to the model improvement strategy for reducing hallucination, <cit.> improve the visual and text feature alignment through reinforcement learning. <cit.> propose a contrastive decoding strategy to reduce reliance on pretrained knowledge. <cit.> and <cit.> also focus on the inference stage and propose specialized decoding strategies to mitigate hallucination. Other strategies for reducing hallucination have been proposed. For example, <cit.> design a post-processing model to detect hallucinated objects and rephrase the generated answers, and <cit.> adapt a reinforcement learning strategy that uses human evaluation of the hallucination level to improve the model.



  
Assessment of LVLM Hallucination

There are two approaches to assessing hallucination in LVLMs. The first is VQA. The ground truth information of the input images is leveraged to construct questions regarding the existence of objects in the images (e.g., “Is there a black cat in the image?”), as well as questions about objects which do not exist in the images. The models are evaluated in terms of the percentage of correctly answered questions. Metrics of this type include POPE <cit.>, CIEM <cit.>, and NOPE <cit.>. The second  approach is to use pre-designed prompts from which the models produce various generations that are then evaluated. Examples include CHAIR <cit.>, which counts the hallucinated objects in generated image captions, and MMHAL-BENCH <cit.>, which uses GPT-4 <cit.> to compare the generations with human answers and determine the propensity toward hallucination.



  
VQA in MLVLMs

For MLVLMs, given a medical scan, models such as LLaVA-Med <cit.> and Med-PALM <cit.> are able to answer questions regarding the types of modalities, the scanned organs, and medical indicators such as opacity. They have demonstrated good performance on medical VQA datasets such as VQA-RAD <cit.>, SLAKE <cit.>, and Path-VQA <cit.>. However, most medical questions in existing datasets are simple. MLVLMs have not been tested on a broader range of complex pathologies.





§ METHODOLOGY



<ref> illustrates the structure of common LVLMs. They are based on a pretrained unimodal LLM such as Llama <cit.> and Vicuna <cit.>. A pretrained visual encoder, such as ViT <cit.> or conventional CNNs, is applied to extract image features that are projected to the text feature space by an adapter. The projected visual features are concatenated with the text prompt embeddings and fed to the LLM. The adapter usually consists of several linear layers with non-linear activations. The visual encoder and the LLM are usually frozen during training.

In our work, we choose the pretrained LLaVA-Med <cit.> as our model, which is a MLVLM built upon LLaVA <cit.>. The model structure resembles <ref>. It uses pretrained Vicuna <cit.> as the LLM and the pretrained ViT encoder from CLIP <cit.> as the visual encoder. The adapter is simply a trainable projection matrix. Both the visual encoder and LLM weights are frozen during training. LLaVA-Med fine-tunes LLaVA in two steps. First, it fine-tunes LLaVA to generate medical reports from input medical images. Second, it uses GPT-4 to generate various questions from the ground truth reports and fine-tunes the model to perform question answering. 

Most MLVLMs are currently trained by medical VQA such that medical diagnosis can be performed by asking questions related to various pathologies; e.g., “Does this image have lung lesion?”. To reduce model hallucination and improve VQA accuracy, we propose two prompting strategies at the inference stage: (1) providing the model with detailed explanations about the queried pathologies and (2) asking the model to consider the inferences of a weak learner.



 §.§ Prompting With Detailed Explanations


Given imbalanced training data, MLVLMs might not adequately be able to learn the features of the minority pathologies. To compensate for insufficient training, we provide a detailed explanation of the queried pathology as a prompt at the inference stage. The explanation briefly defines the pathology and lists several key findings in medical images that may indicate its existence. An example is shown in <ref>. The model is informed that Pulmonary Edema is defined as the accumulation of fluid in the lungs. Then several chest X-ray findings that may suggest its existence are provided. The model can determine if the given image has Pulmonary Edema by linking the given findings with the image features.

Prompt templates for a number of pathologies are listed in <ref>. 







 §.§ Prompting With Detailed Explanations and Weak Learners


Data re-sampling is a commonly-used strategy to deal with imbalanced datasets that are responsible for the tendency of traditional image classification models to return negative predictions for minority pathologies. Models trained on re-sampled datasets often exhibit improvements in Precision and Recall scores; however, this strategy may not be suitable to MLVLMs for two reasons. First, it is difficult to balance a dataset containing many categories of pathologies. Second, MLVLMs usually demand much larger datasets and fine-tuning is also expensive.

One can nevertheless enable MLVLMs to benefit by leveraging small models trained on re-sampled datasets. Our method resembles multiagent LLM systems, such as <cit.>, where multiple LLMs debate each other and hallucination can be corrected by referring to the generated outputs of other models. Given that traditional image classifiers are smaller, it is feasible to train multiple small classifiers each of which is trained on re-sampled datasets of a particular pathology. Those models can be further fine-tuned to optimize a single aspect, such as fewer false positives (FP) or fewer false negatives (FN). The classifiers are applied to the medical images and return preliminary predictions. These predictions are selectively included in the prompts as references for the MLVLM. Hence, MLVLMs can benefit indirectly from the nuanced understanding that these specialized models can provide. This method is meaningful because clinicians usually must balance the trade-off between overtreatment and undertreatment when making healthcare decisions. For instance, they may prefer models having a low FP rate if the cost of overtreatment is higher than that of undertreatment. 

 An example is shown in <ref>, which queries about the presence of Edema. We first provide the model with the detailed explanation of Edema. Then, we use the weak learner to suppress the FPs. The image is input to an Edema classifier that has been fine-tuned on a balanced dataset for high sensitivity and high true negative (TN) rate. If its prediction is negative, we append after the pathology explanation the prompt “For this image, another agent thinks the probability of Edema is 0.1”. Instead of using the actual predicted probability, the probability value is manually chosen because the decision threshold has been fine-tuned and is no longer 0.5. We do not use a zero probability value because we do not want the model overly to trust the weak learner. Although in this example our goal is only to reduce FPs, our strategy can also be applied to reduce FNs, simply by fine-tuning the classifier for a high true positive (TP) rate and applying the prompt in the case of positive predictions.




§ EMPIRICAL STUDY





 §.§ Datasets




LLaVA-Med is pretrained on the PMC-15M dataset <cit.>, which contains image-text pairs of multiple modalities; e.g., CT, MRI, X-ray, etc. In the first stage, 467,710 image-report pairs were selected for training. In the second stage, 56,708 question-answer pairs were created from the data of the first stage to fine-tune the model. <ref> shows the count of reports in the LLaVA-Med training data (second stage) that mention one of the five test pathologies as positive. Relative to the total amount of data, all five categories are minorities.



To assess the zero-shot performance of the MLVLM, we used the MIMIC-CXR-JPG <cit.> and Chexpert <cit.> chest X-ray test sets. They include 5,159 and 668 images, respectively. 
Neither dataset overlaps with PMC-15M.

MIMIC-CXR-JPG includes images and medical reports covering 13 categories of findings: Atelectasis, Cardiomegaly, Consolidation, Edema, 
Enlarged Cardiomediastinum, Fracture, Lung Lesion, Lung Opacity, Pleural Effusion, Pneumonia, Pneumothorax, Pleural Other, and Support Devices. The raw reports are parsed and rough image-level tags are automatically generated by a rule-based approach <cit.>. Each label contains four values: 1 (positive), 0 (negative), -1 (uncertain), and missing. For simplicity, we treat both uncertain and missing as negative. We also use the MIMIC-CXR-JPG training set, which contains 227,827 chest X-rays with reports, to train the weak learner models.

Chexpert covers the same 13 categories as MIMIC-CXR-JPG. However, it does not include medical reports and has only image-level labels. There is no overlap between MIMIC-CXR-JPG and Chexpert. 

<ref> shows the split of pathology categories (excluding normal) in the MIMIC-CXR-JPG and Chexpert test sets. Clearly, almost all pathology categories are minor classes with much fewer positive than negative occurrences.

For our main testing regimen, we selected the five pathologies in the Chexpert Competition <cit.>, Atelectasis, Cardiomegaly, Consolidation, Edema, and Pleural Effusion.



 §.§ Implementation Details




As was mentioned in <ref>, we use the pretrained LLaVA-Med MLVLM without any further fine-tuning. We convert the classification task into a VQA task by using the prompt template shown in Row 1 of <ref>, which we name Prompt Template 1 (PT1). We first run the pretrained LLaVA-Med with PT1. Next, we incorporate pathology explanations (Row 2 of <ref>), yielding Prompt Template 2 (PT2). Finally, we integrate the predictions of weak learners into the prompts (Row 3 of <ref>), resulting in Prompt Template 3 (PT3).

As will be justified by our experiments, our weak learner is designed to suppress FP predictions. To this end, we use the pretrained ResNet50 <cit.>. For each pathology, the training dataset was sampled such that the ratio of positive and negative cases is 2:1. The model was trained for 10 epochs with a 1e-4 learning rate. The training process was monitored using the AUC score and the one with the highest validation AUC was kept. Then, the decision threshold d was fine-tuned to optimize a weighted sum of Specificity and Negative Predictive Value (NPV); i.e.,

    d = w_1TN/TN+FP+w_2TN/TN+FN,

where weights w_1 and w_2 are preset to 0.2 and 0.8, respectively. The medical images were input to the weak learners to obtain preliminary predictions for each pathology and only the negative predictions were selected to craft the PT3 prompts. 

The responses returned by LLaVA-Med can take various forms, such as “This image has Edema”, “Edema is found”, “The fluid in the lung indicates Edema”, etc. An off-the-shelf Llama-7B <cit.> serves to summarize long responses into Yes/No answers such that accuracies could easily be computed.



 §.§ Results


To demonstrate the efficacy of our prompting strategies, starting from the PT1 baseline, the pathology explanations were provided first (strategy PT2) and then, based on the results, weak learners were introduced to improve performance on specific aspects, resulting in strategy PT3.





  
PT2: Adding Pathology Explanations
<ref> reports Precision, Recall, and F1 scores of the PT1 and PT2 strategies on the MIMIC-CXR-JPG and Chexpert test sets.[The AUC and ROC scores commonly reported in the literature to assess the performances of most medical image classification models on the MIMIC-CXR-JPG and Chexpert datasets are unsuitable in our context because MLVLMs output text rather than probabilities.] On MIMIC-CXR-JPG, after adding pathology explanations, the F1 scores increased for detecting Atelectasis, Cardiomegaly, Edema, and Pleural Effusion, albeit only minimally for Consolidation. On Chexpert, after adding pathology explanations, the F1 scores for detecting Atelectasis, Cardiomegaly, and Edema increased, whereas they did not for Consolidation and Pleural Effusion. The Precision and Recall scores reveal that adding explanations generally leads to a large increase in Recall, but only minimally influences Precision. For minority pathologies such as Consolidation whose F1 score is dominated by low Precision, improving the Recall would not have much effect. Thus, PT2's performance bottleneck is Precision.







  
PT3: Referring to Weak Learners
Going beyond our PT2 strategy, we applied our PT3 strategy to further improve diagnostic accuracy. <ref> provides the TP, FP, and FN prediction counts of LLaVA-Med on the Chexpert test set using the PT2 strategy. Note the large number of FP cases. Hence, we designed our weak learners to suppress FP predictions.
<ref> compares the performance on Chexpert before and after referring to the weak learner. It shows that the F1 prediction accuracy can be substantially increased by introducing weak learner predictions into the prompts. The F1 scores of Cardiomegaly, Edema, and Pleural Effusion increase by 0.115, 0.194 and 0.089, respectively. To further demonstrate the efficacy of our PT3 strategy, <ref> compares the FP predictions of the PT2 and PT3 strategies. The reduction of FP cases is noteworthy, especially on Edema, for which the FP count is reduced by 78.5% (322). 







  
Additional VQA Experiments
<ref> shows the results of applying the PT1, PT2, and PT3 strategies with LLaVA-Med on the MIMIC-CXR-JPG and Chexpert datasets across another five medical findings: Enlarged Cardiomediastinum, Lung Lesion, Lung Opacity, Pneumonia, and Pneumothorax. Providing pathology explanations (PT2) generally yields better results over the PT1 baseline, albeit inconsistently. Introducing weak learner references (PT3) yields only limited increases in Precision, but large decreases in Recall. Generally, it offers insignificant improvement. Enlarged Cardiomediastinum, Lung Lesion, Pneumonia, and Pneumothorax are minor categories and all our experimental settings, including for the weak learner, fail to learn them. Prompting is apparently unhelpful in such situations.







  
SOTA Benchmark
<cit.> report F1 scores for detecting Atelectasis, Cardiomegaly, Consolidation, Edema, and Pleural Effusion on the Chexpert dataset using their deep learning model, as well as for the performance of radiologists. Their work offers a state-of-the-art chest X-ray diagnosis benchmark. <ref> compares the F1 scores of radiologists, the model of <cit.>, and LLaVA-Med. It shows that LLaVA-Med's VQA performance of with the baseline PT1 strategy is unsatisfactory, rendering the model far from being deployable in clinical practice. However, while still underperforming radiologists, our PT3 strategy yields a significant improvement, especially on Atelectasis, Cardiomegaly, and Edema for which the F1 score increases by approximately 17% to 21%. 



  
Application to General Domain LVLMs
Our prompt strategies can also be applied to general domain LVLMs. We studied the performance of LLaVA <cit.> and MiniGPT-v2 <cit.> using POPE metrics <cit.>, which evaluate the hallucination of LVLMs by asking questions about the presence of objects. The POPE scores of LLaVA and MiniGPT-v2 have high Precision and low Recall. Hence, our weak learner strategy can be used to reduce the FN predictions. We selected an off-the-shelf Fast-RCNN <cit.> as the weak learner, fine-tuned the detection threshold of bounding box scores to achieve high Recall, and introduced the positive predictions of the weak learner into the prompts. The results in <ref> show that the Recall scores across three POPE categories increased by around 7% (Precision scores decrease slightly), thus improving the F1 scores.



§ CONCLUSIONS AND DISCUSSION



We have tested the visual question answering abilities of the LLaVA-Med medical large vision-language model when applied to the diagnosis of pathologies. Our results show that the model has unsatisfactory performance when asked questions regarding the presence of complex pathologies. We proposed two prompt engineering strategies to improve the visual question answering accuracy of the model: providing explanations of pathologies and referring to the predictions of weak learners. The first strategy helps the model understand minority pathologies that it does not learn well in the training stage. The second strategy can help improve diagnostic accuracy in specific ways; e.g., by suppressing false positives. This strategy can also be applied to LVLMs in other, non-medical domains.

However, our two strategies are not effective on pathologies with extremely scarce data. For example, providing text explanations for Consolidation, Fracture, Lung Lesion, Pneumonia, and Pneumothorax may not suffice since the visual encoder does not adequately learn meaningful visual features. Moreover, the data may not suffice to adequately train weak learners. A promising direction for future research would be to devise a strategy for handle these rare categories. Retrieval Augmented Generation (RAG) could be a potential solution. In addition to textual explanation of a pathology, for instance, typical example images can be provided to help the model make diagnostic decisions.





§ PATHOLOGY EXPLANATIONS



The explanations of the five pathologies are as follows:



  
Atelectasis: Atelectasis refers to the partial or complete collapse of a lung or a section of lung. The features of atelectasis on an X-ray can vary depending on the cause and extent of the collapse. Some common X-ray features include: 1. The affected area may appear denser or whiter than normal lung tissue due to the collapse, leading to increased opacity on the X-ray. 2. The affected portion of the lung may appear smaller or compressed compared to the surrounding healthy lung tissue. 3. Atelectasis can cause a shift or displacement of nearby structures, such as the trachea or heart, toward the affected area. 4. In obstructive atelectasis (caused by a blockage in the airways), there might be signs of hyperinflation in the unaffected areas of the lung and a visible blockage or narrowing in the affected bronchus. 5. Linear or band-like opacities may be visible, often referred to as plate or band atelectasis, which can occur due to the collapse of small airways. Given the information above, does this image have Atelectasis?



  
Cardiomegaly: Cardiomegaly is enlargement of the heart. The definition is when the transverse diameter of the cardiac silhouette is greater than or equal to 50% of the transverse diameter of the chest (increased cardiothoracic ratio) on a posterior-anterior projection of a chest radiograph or a computed tomography. Given the information above, does this image have Cardiomegaly?



  
Consolidation: Consolidation on an X-ray refers to the filling of the lung's air spaces with fluid inflammatory exudate, or cellular material. Typical X-ray findings suggesting consolidation include: 1. Areas of increased density in the lung tissue, appearing as an opaque or hazy patch on the X-ray. Given the information above, does this image have Consolidation?



  
Edema: Pulmonary edema is the accumulation of fluid in the lungs. Some common X-ray features include: 1. Increased density in the central lung fields resembling the shape of bat wings. 2. Thin, linear opacities at the lung periphery, often indicating interstitial edema. 3. Prominent blood vessel markings due to engorgement from increased pressure in the pulmonary vasculature. Given the information above, does this image have Edema?



  
Pleural Effusion: Pleural effusion is the accumulation of fluid in between the parietal and visceral pleura. Some common X-ray features include: 1. blunting of the costophrenic/cardiophrenic angle. 2. fluid within the horizontal or oblique fissures. 3. meniscus is seen. 4. mediastinal shift occurs away from the effusion.





