Paper Title: Diff-Cleanse: Identifying and Mitigating Backdoor Attacks in Diffusion Models.




Knowledge-Guided Prompt Learning for Lifespan Brain MR Image Segmentation
    Lin Teng 1 Zihao Zhao 1 Jiawei Huang 1 Zehong Cao 2 Runqi Meng 1 
Feng Shi 2^() 
Dinggang Shen 1,2,3^()
    August 5, 2024
===========================================================================================================




§ ABSTRACT 

  Diffusion models (DM) represent one of the most advanced generative models today, yet recent studies suggest that DMs are vulnerable to backdoor attacks. Backdoor attacks establish hidden associations between particular input patterns and model behaviors, compromising model integrity by triggering undesirable actions with manipulated input data. This vulnerability poses substantial risks, including reputational damage to model owners and the dissemination of harmful content. To mitigate the threat of backdoor attacks, there have been some investigations on backdoor detection and model repair. However, previous work fails to purify the backdoored DMs created by state-of-the-art attacks, rendering the field much underexplored. To bridge this gap, we introduce Diff-Cleanse, a novel two-stage backdoor defense framework specifically designed for DMs. The first stage employs a innovative trigger inversion technique to detect the backdoor and reconstruct the trigger, and the second stage utilizes a structural pruning method to eliminate the backdoor. We evaluate our framework on hundreds of DMs attacked by 3 existing backdoor attack methods. Extensive experiments demonstrate that Diff-Cleanse achieves nearly 100% detection accuracy and effectively mitigates backdoor impacts, preserving the model's benign performance with minimal compromise. Our code is avaliable at <https://github.com/shymuel/diff-cleanse>.





§ INTRODUCTION




Diffusion models (DM) have emerged as one of the most successful generative models due to their superior performance  <cit.> and are widely applied in image <cit.>, audio <cit.>, video <cit.>, text <cit.>, and text-to-speech <cit.>
generation. Moreover, DMs are increasingly used in safety-critical tasks, such as autonomous driving <cit.> and medical imaging <cit.>.


However, recent studies show that DMs are vulnerable to backdoor attacks <cit.>, where attackers poison a portion of the training dataset, causing the model to execute attacker-specified behaviors when encountering inputs embedded with specific trigger patterns. Given the common practice of using powerful pre-trained models available online or outsourcing the training process to third parties, such attacks could have disastrous consequences <cit.>.
For instance, a company using a backdoored DM for image generation service could unintentionally produce offensive content, damaging the company's reputation and leading to legal repercussions <cit.>. Additionally, DMs are often used to generate synthetic data for training downstream models when data is scarce or privacy is a concern <cit.>. Attackers could implant biases within these DMs, affecting subsequent model behaviors.


Despite these challenges, there is a significant lack of specialized backdoor defense strategies for DMs. Most existing backdoor defense methods are designed for traditional neural network in classification tasks and cannot be directly applied to DM due to different prediction objectives (class probability for classifier, estimation of mean or noise for DMs) and inference processes. Moreover, defense methods for DMs exhibit limitations, either restricted to input-level backdoor detection <cit.> or failing to effectively identify backdoors in our experiments <cit.>.


To address this gap, we study three types of backdoor attacks on DMs in the literature, noting that despite the diversity in semantics and textures of trigger and target patterns, attackers typically seek backdoored DMs with high utility and specifity <cit.>. High utility means that the DM generates high-quality images from clean inputs, while high specifity means that inputs containing trigger patterns force the DM to generate nearly identical target images. To repair the model, a natural approach is to reconstruct an effective trigger and remove neurons activated by it. Our experiments show that pruning a small number of neurons (typically 1% to 2%) is sufficient to remove the backdoor. Based on these insights, we propose a two-stage backdoor defense framework for DMs. First, we introduce a novel trigger inversion method that models trigger reconstruction as an optimization problem, aiming to find a trigger that causes the DM to generate identical images. We devise a metric to measure the effectiveness of the inverted trigger and to aid in backdoor detection. To eliminate the backdoor, we propose a pruning-based method to remove the backdoor-related neurons and design a loss function to improve the fine-tuning effect when using only a small number of clean samples.


The major contributions of this paper are as follows. 






  * We propose Diff-Cleanse, a two-stage backdoor defense framework designed for diffusion models, which can effectively mitigate the backdoor by targeted pruning backdoor-related neurons. To the best of our knowledge, this is the first study to explore the application of pruning methods in the field of DM backdoor defense.



  * We introduce a trigger inversion method based on the pursuit of high specificity in existing backdoor attacks in DMs. As part of this method, we design a metric, trigger effectiveness (TE), to measure the performance of the inverted trigger and facilitate backdoor detection.



  * We devise a simple yet effective pruning-based backdoor removal algorithm, including the design of a novel loss function, to eliminate backdoors.



  * We empirically show that in terms of 3 state-of-the-art backdoor attacks on DMs, Diff-Cleanse achieves close to 100% detection accuracy and reduces backdoor effects to nearly zero while largely maintaining the model's utility.







§ PRELIMINARIES








 §.§ Diffusion models



Utilizing diffusion models (DMs) generally involves three key procedures. (1) Diffusion process. Define a diffusion process that diffuses the data distribution q(x) into a certain distribution τ(x) over T timesteps. (2) Training. Train parameters θ to ensure that the generative process aligns with the reversed diffusion process, i.e., p_θ(x_t-1|x_t)=𝒩(x_t-1;μ_θ(x_t,t),β_θ(x_t,t))=q(x_t-1|x_t). (3) Sampling. Use the trained generative process p_θ(x_t-1|x_t) to sample from T_max down to T_min. 




  
DDPM. DDPM <cit.> is a representative diffusion model. It assumes τ(x)=𝒩(0,𝐈) and defines the Markov diffusion process q(x_t |x_t-1)=N(x_t; √(1-β_t)x_t-1,β_t𝐈), where {β_t}_t=1^T is a predefined variance schedule. Let α_t=1-β_t and α̅_i=∏_t=1^i α_t. Given x_0 ∼ q(x), t ∼Uniform({1, …, T}) and ϵ∼𝒩(0, I), the DM is trained to minimize the loss ϵ-ϵ_θ(√(α̅_t) x_0+..√(1-α̅_t)ϵ, t) ^2, thereby learning to predict the added noise given the x_t at timestep t. During sampling, given x_T ∼𝒩(0, I), DDPM samples from p_θ(x_t-1|x_t)=𝒩(x_t-1 ; μ_θ(x_t), β_θ(x_t)) from t=T to t=0, where μ_θ(x_t)=1/√(α_t)(x_t-1-α_t/√(1-α̅_t)ϵ_θ(x_t, t)) and β_θ(x_t)=(1-α̅_t-1) β_t/1-α̅_t.




  
Advanced samplers. The default samplers for DMs are relatively slow compared to other generative models such as GANs <cit.>, prompting the development of advanced samplers to accelerate the sampling.
Denoising Diffusion Implicit Models (DDIM) <cit.> introduces an abbreviated reverse process, reducing the steps from 1000 to as few as 50, based on a generalized non-Markovian forward chain derived from DDPM. Additionally, some researchers formulate the reversed diffusion process as Ordinary Differential Equations (ODEs) and employ higher-order approximations, enabling DMs to generate high-quality samples in fewer steps <cit.>.







 §.§ Backdoor attacks and defences on diffusion models




  
Backdoor attacks. Injecting backdoors to traditional models (e.g., classifiers) involves poisoning a small portion of the training data by stamping the pre-defined trigger and relabeling them as target classes <cit.>. After training, the model would produce the intended target label in the presence of the trigger. Backdoor attacks on DMs are more complex because the sampling process depends not only on the input but also on the transitional distribution. To the best of our knowledge, there are three existing noise-to-image DM backdoor attacks <cit.>. Their workflow is illustrated in <ref>. 


To inject backdoor into DMs, attackers first need to mathematically define a forward backdoor diffusion process x_0^b →x_T^b, where x_0^b is the target image, x_T^b is the backdoored noise, also know as Trojan noise <cit.>. Without loss of generality, we assume x_T^b=g+γ·x_T^c, where g∼𝒩(𝐫, 0) is the trigger pattern, x_T^c∼𝒩(0,𝐈) is the clean input. Attackers then need to define the corresponding backdoor reverse process and backdoor loss. During training, clean loss and backdoor loss are calculated separately for clean samples and backdoor samples, and the parameter θ learns both the clean transitional distribution q^c(x_t-1^c|x^c_t) and the backdoor transitional distribution q^b(x^b_t-1|x^b_t). During sampling, the backdoored DM will generate clean samples x_0^c∼ q^c(x_0^c) with clean input, and target image x_0^b with backdoored noise. In BadDiff <cit.> and VillanDiff <cit.>, the same sampling process can be used to generate clean samples from clean noise, as well as backdoor targets from backdoored noise. In TrojDiff <cit.>, generating clean samples and backdoor targets requires not only different inputs but also different sampling functions, which weakens the effectiveness of the injected backdoor. BadDiff is only applicable to the ancestral sampler DDPM, while VillanDiff proposes a general framework and can attack advanced samplers and more types of DMs, such as latent diffusion models and score-based diffusion models. Moreover, the target in TrojDiff can be a distribution, such as "horse" or handwritten digits, whereas the target in BadDiff and VillanDiff must be an image.





  
Backdoor defenses. Due to different prediction objectives and inference processes, backdoor defense methods designed for other tasks cannot be directly applied to backdoored DMs <cit.>. Recently, some efforts have been made to mitigate the effects of backdoor in DMs. Elijah <cit.> optimizes a trigger g that linearly decreases during the sampling process, uses the model's response to inputs containing g to detect backdoors and removes the backdoor by aligning the model's predictions for clean and backdoored inputs. UFID <cit.> introduces an input-level backdoor detection framework, which identifies backdoored inputs by analyzing the sensitivity of input-output pairs to input perturbations.






 §.§ Threat model


We adopt a similar setting as in <cit.>, which describes the goals and capacities of two parties, Attacker and Defender. A comprehensive explanation of each party is as follows.




  
Attacker. The attacker's goal is to develop a backdoored DM with two properties <cit.>. (1) High utility. The model has good performance in clean sampling.
(2) High specificity. Activation of the backdoor causes the model to generate the target image. In terms of capacity, we assume a strong attacker who can (1) design the backdoor diffusion process and backdoor sampling procedure, (2) modify all aspects of training, including the training loss, training data and hyperparameters.






  
Defender. The defender's goals are to (1) detect potential backdoors, and (2) remove them while maintaining the models' benign utility. We assume the defender has no prior knowledge of the attack settings but has white-box access to the DM and a limited amount of clean data (typically 10%).




§ METHODOLOGY





In this section, we start by introducing our two-stage backdoor defense framework, as illustrated in <ref> (<ref>). Given a risky DM, we first apply our trigger inversion method to reconstruct a trigger g^s and detect the backdoor (<ref>). If the DM is backdoored, we then execute our structural pruning-based backdoor removal algorithm to eliminate the injected backdoor (<ref>).







 §.§ Overview of our approach






  
Key intuition. We derive the intuition of our technique from the basic properties of backdoored DMs, i.e., clean DMs convert standard Gaussian noise to the dataset distribution through clean sampling, while backdoored DMs can additionally convert the backdoored noise to the target distribution through backdoor sampling. In attacks addressed in this paper, the backdoored noise differs significantly from clean noise, and the target distribution does not overlap with the dataset distribution. 


Intuitively, our approach reconstructs the trigger and use it to facilitate the backdoor detection and removal. Specifically, we hypothesize that the trigger eliciting target images indicates the presence of a backdoor and that neurons activated by the inverted trigger will also be activated by the attacker-specified trigger. We employ pruning-based method to remove the backdoor, which is common in backdoor defenses for discriminative models <cit.>. This leads to following observations. In following sections, we call g-X as trigger-target pair if clean inputs stamped with g cause the DM to generate highly similar outputs X. For simplicity, g and X have the same shape. 







  
Observation 1. Both clean and backdoored DMs contain trigger-target pairs. For the same target image, there are triggers with different levels of effectiveness.


We classify trigger-target pairs into two categories, natural and artificial. Natural triggers and targets emerge during normal training, which are typically images uncomprehensible to humans. In contrast, artificial triggers and targets are injected by backdoor training and contain intelligible semantics. Both clean and backdoored DMs contain natural trigger-target pairs, while backdoored DMs additionally contain artificial trigger-target pairs. The color, shape and texture of the artificial trigger can be arbitrarily configured. Therefore, we focus on the effectiveness of the trigger, i.e., the ability of the trigger to activate the backdoor. <Ref> presents examples of both natural and artificial targets, as well as the outcomes of backdoor sampling using triggers with varying levels of effectiveness.




  
Observation 2. Structural pruning method can completely eliminate the backdoor while removing very few neurons (typically 1% to 2%).


Existing prune-based backdoor defense methods typically focus on deep layers and utilize unstructural pruning <cit.>, which minimally impacts the benign performance of the model. However, we observe that the optimal pruning rate varies significantly across different attack settings when the algorithm only prunes deep layers, making the algorithm sensitive to hyperparameters. Moreover, we find that unstructural pruning usually fails to completely eliminate the backdoor. Backdoors tend to recover after fine-tuning, whereas structural pruning suppresses this phenomenon.


Based on the previous observations, we propose the two-stage backdoor defense framework Diff-Cleanse. In the first stage, we model the trigger inversion as an optimization problem and reconstruct effective triggers. We also use the effectiveness of the inverted triggers to determine the presence of backdoors. If a backdoor is detected, in the second stage, we apply structural pruning to remove backdoor-related neurons and restore the benign performance of the purified DM using our proposed loss function. Due to space limitations, the pseudocode of our algorithms are in <ref>.




 §.§ Trigger inversion (stage 1)





Recall that the backdoor diffusion maps the target image x_0^b to the backdoored noise x_T^b∼𝒩(𝐫, γ^2𝐈). Without loss of generality, we consider the unified form of backdoor diffusion as follows.

    x_t^b=α̂_t x_0^b+ρ̂_t 𝐫+β̂_t ϵ_t,  ϵ_t∼𝒩(0, 𝐈),  t∈[1,…,T]

where {α̂_t}_t=1^T is a decreasing sequence, {ρ̂_t}_t=1^T and {β̂_t}_t=1^T are increasing sequences. 


Since the defender does not know the specific parameters of the attack, let {β_t}_t=1^T represent a monotonically increasing variance sequence from 0 to 1, α_t = 1-β_t, and α̅_t=∏_i=1^t α_i. We derive a special case of the backdoor diffusion, denoted as {x_t^s}_t=1^T, which has the following form.

    x_t^s=√(α̅_t)x_0^s+(1-√(α̅_t)) r^s+√(1-α̅_t)γϵ_t,  ϵ_t∼𝒩(0, 𝐈),  t∈[1,…,T]

Using the fixed increasing variance sequence {β_t}_t=1^T in DDPM <cit.>, we have α̅_T ≈ 0. Hence, x_T^s≈r^s+γϵ_T ∼𝒩(r^s, γ^2 𝐈).  <Ref> provides a detailed derivation of x_t^s. 


Let g^s denotes the inverted trigger, X^s denotes the results of backdoor sampling with g^s. We model trigger inversion as an optimization problem that has two objectives. First, since attackers seek high specificity DMs <cit.>, artificial triggers should effectively activate the backdoor, leading to the generation of identical images, i.e., the targets. To quantify this, we propose the similarity loss ℓ_s as in <ref>, which evaluates the diversity of the generated samples.

    ℓ_s(g^s)=1-∑_i,j𝒜(X^s_i, X^s_j)/|X^s|^2

where (X^s_i, X^s_j) denotes pairs of samples selected from X^s. 𝒜(·) is the cosine similarity function.


The second objective is to find the artificial trigger-target pair rather than natural trigger-target pairs. Empirically, abnormal entropy values indicate samples without comprehensible semantics, and natural targets tend to have either too high (e.g., random noise) or too low entropy (e.g., pure black pictures). Thus, we propose entropy loss ℓ_e as follows.

    ℓ_e(g^s)=∑_i[ max(ℰ(X^s_i)-threshold_max, 0) + max(threshold_min-ℰ(X^s_i), 0) ]

where ℰ(·) is the entropy calculation function and X^s_i is the i-th results of backdoor sampling with g^s. The values of threshold_max and threshold_min are not strict and we use 0.1 and 0.9 in all experiments. Together, we formulate trigger inversion as a multi-objective optimization task by optimizing the weighted sum of the two objectives. The final formulation is as follows.

    min_g^sℓ(g^s)=ℓ_s(g^s)+λℓ_e(g^s)

where λ is the weight of ℓ_e. In experiments, we adjust λ dynamically so that the two loss terms are of similar magnitude. We use <ref> to calculate the derivative of x_0^s with respect to x_T^s, which speeds up the calculation, and use Adam optimizer to solve the above optimization.




  
Backdoor detection. We note that natural targets, sampled with a natural trigger, are visually similar but exhibit slight variations. In contrast, samples generated with an artificial trigger show a high degree of similarity. Therefore, we devise a metric, trigger effectiveness (TE), that measures the similarity between generated samples X to evaluate the effect of g^s and to detect the backdoor.

    TE(X)=|{X_i^s, X_j^s ∈X: 𝒜(X_i, X_j)> threshold }|/|X|

where |·| denotes the cardinality of the set. High TE indicates that g^s effectively activates the artificial backdoor in the model, while low TE indicates that there is no backdoor or the trigger is weak. The full algorithm is summarized in <ref>.




 §.§ Structural pruning-based backdoor removal (stage 2)




Upon detecting a backdoor, we can prevent the jeopardy by discard the model, which would also sacrifice its benign utility. To address this issue, we propose a structural pruning-based approach to eliminate the backdoor while preserving the model's utility. Our method utilize the inverted trigger to prune neurons that are important for backdoor sampling but less important for clean sampling. Following Diff-Pruning <cit.>, we use the taylor expansion of the loss to approximate neuron importance.


For simplicity, we consider the parameter θ of the backdoored DM as a 2-D matrix, where each row vector θ_i=[θ_i 0, θ_i 1, …, θ_i K] contains K scalar parameters and the goal of pruning is to obtain a sparse parameter matrix θ^'. A natural idea is that the loss disruption caused by pruning can indicate the changes in model performance, thus we formulate the objective as follows.



    min_θ^'|ℒ^c(θ^', x^c)-ℒ^c(θ,x^c)|,  max_θ^'|ℒ^b(θ^',x^b))-ℒ^b(θ,x^b))|,   s.t. θ^'_0 ≤ s

where ℒ^c(·) and ℒ^b(·) are loss functions of clean and backdoor training respectively, x^c and x^b denote clean and backdoored inputs respectively. s is the sparsity of θ^' and ·_0 is the L_0 norm, which counts the number of non-zero parameters. Since the defender does not know the attack settings, we calculate ℒ^b(·) using <ref>. Due to the iterative nature of diffusion process, ℒ^c and ℒ^b can be decomposed into sequences {ℒ^c_t}_t=1^T and {ℒ^b_t}_t=1^T.




  
Taylor importance score. Let ℒ_t denotes the loss at timestep t of either clean or backdoor training. Leveraging the Taylor expansion <cit.>, the loss disrupt can be approximated as follows.

    ℒ_t(θ^',x)-ℒ_t(θ,x)     =∇ℒ_t(θ,x)(θ^'-θ)+O(θ^'-θ^2)

Using the Taylor importance criterion, we can obtain the importance of an individual weight θ_i k by setting θ^'_i k=0 in <ref>. For structural pruning, the smallest unit of pruning is θ_i, so we adopt the aggregated Taylor importance <cit.> as the importance score, which is formulated as follows.

    ℐ_t(θ_i, x)=∑_k|ℒ_t(.θ,x|θ_i k=0)-ℒ_t(θ,x)|=∑_k|θ_i k·∇_θ_i kℒ_t(θ, x)|

We further explore the contribution of different loss terms {ℒ_1, …, ℒ_T} during backdoor diffusion. We observe that in backdoor diffusion, large t corresponds to large ℒ_b,t, while in clean diffusion, small t corresponds to large ℒ_c,t. <cit.> empirically demonstrates that small ℒ_t is uninformative. Thus, we introduce a threshold parameter 𝒯 and only use timestep t>𝒯 for gradient accumulation, which yields the finalized importance score.

    ℐ(θ_i, x)=∑_k|θ_i k·∑_{t |  t>𝒯.}∇_θ_i kℒ_t(θ, x)|





  
Implementation of model repair. After obtaining the Taylor importance scores of substructures, we firstly prune the top-p% channels that have the greatest difference in ℐ(θ_i, x^c) and ℐ(θ_i, x^b) to obtain θ^'. Then we fine-tune the pruned model.


Recall that the objective function of clean training is loss_c=ϵ-ϵ_θ(√(α̅_t)x+√(1-α̅_t)ϵ, t)^2. However, this loss alone is inefficient because defenders can only utilize a small number of clean samples. Therefore, we use the output of the unpruned model θ as a reference to align the updated model's distribution with the clean distribution obtained by training through the entire clean dataset. This loss can be expressed as loss_b=ϵ_θ^'(√(α̅_t)x+√(1-α̅_t)ϵ, t)-ϵ_θ(√(α̅_t)x+√(1-α̅_t)ϵ, t)^2, and the complete loss is loss_p=loss_c+loss_b. The full algorithm is summarized in <ref>.




§ EXPERIMENTS AND RESULTS















In this section, we first introduce the experiment setup (<ref>). We then evaluate the effectiveness of our method (<ref>, <ref>) and conduct ablation studies (<ref>). Finally, we discuss a potential countermeasure of the attacker (<ref>). Our main experiments run on a server equipped with Intel Golden 6240 CPUs and NVIDIA V100 GPUs.





 §.§ Experiment setup




  
Datasets, models and attacks. We use two benchmark vision datasets, CIFAR-10 (32× 32) <cit.> and CelebA-HQ (256× 256) <cit.>, since these are the datasets considered in the evaluated backdoor attack methods. Due to the instability of existing attack methods on latent diffusion models and score-based diffusion models, we limit our experiments to DDPM models. We employ both the default sampler and advanced samplers, including DDIM, PNDM, DEIS, DPMO1, DPMO2, DPMO3, DPM++O1, DPM++O2, DPM++O3, UNIPC, and HEUN. Clean and backdoored DMs are either downloaded from Hugging Face or trained by ourselves using official codes. We consider all the existing attacks in the literature, namely, BadDiff <cit.>, TrojDiff <cit.> and VillanDiff <cit.> in the context of unconditional image generation. Please check <ref> for more details.




  
Baselines. As we mentioned earlier, most existing defenses are not applicable for DMs. For trigger inversion, We use the method in Elijah <cit.> as the baseline, which is the only trigger inversion method for DMs in the literature. To the best of our knowledge, there is no other pruning-based model repair methods for DMs. So we propose two effective pruning strategies and evaluate them.




  
Evaluation metrics. Following existing literature <cit.>, we use the following metrics: (1) Detection Accuracy (ACC) evaluates the effectiveness of backdoor detection methods, with higher values indicating better performance. (2) Attack Success Rate (ASR) calculates the proportion of backdoor sampling results whose similarity to the target image exceeds a predefined threshold. (3) Structural Similarity Index Measure (SSIM) accesses the similarity between generated samples and the target image. (4) Fréchet Inception Distance (FID) evaluates the quality of the generated clean samples relative to the training data, with lower values indicating better utility. For trigger inversion methods, we use ACC, ASR and SSIM, where higher values indicate more effective triggers. To evaluate model repair methods, we focus on changes in metrics, i.e, ΔFID, ΔASR and ΔSSIM, where lower values are preferable. A low ΔFID score means better preservation of benign utility, while low ΔASR and ΔSSIM scores indicate more effective backdoor removal.










 §.§ Evaluation of trigger inversion




<ref> compares the performance of triggers inverted by Elijah and our method. Following the original paper, we use clean and backdoored DMs to train a random forest as the backdoor detector of Elijah. Our proposed method utilize TE scores for backdoor detection. Although both methods achieve high ACC scores, our method generates triggers with higher ASR and SSIM, especially for VillanDiff. 


In our experiments, small perturbations to clean inputs are sufficient to cause the models attacked by BadDiff to generate target images, which makes the backdoor easy to detect, indicating a large trigger space. In contrast, backdoors injected by VillanDiff are hard to activate, suggesting narrow trigger spaces. The effectiveness of the triggers inverted by Elijah is relatively low due to the lack of distinction between natural and artificial targets. Our method, which utilizes the properties of artificial targets, is empirically more effective. The small size of patch-based triggers affects the effectiveness of our inverted triggers, but our framework can still repair these DMs as shown in <ref>. 









 §.§ Evaluation of backdoor removal


In order to get better utility, we argue that the pruning algorithm should eliminate the backdoor while removing the minimum number of neurons. In all experiments, Diff-Pruing prunes 13% neurons for models trained on CIFAR-10 and 7% neurons for models trained on CelebA-HQ. Diff-cleanse pruning 1%∼3% neurons of the model until the inverted trigger becomes ineffective.

<ref> compares two pruning-based backdoor removal strategies. Both algorithms completely remove the backdoor across all attack settings, and Diff-Cleanse shows less impact on benign utility (lower ΔFID). In some cases, our approach does not achieve a smaller ΔSSIM because images with no visual resemblance to the target can still exhibit SSIM scores above zero. 








 §.§ Ablation study




  
Effect of loss_p.
<Ref> compares the results of fine-tuning using loss_c and loss_p. The fine-tuning method quickly reduces FID within 5 epochs, and our proposed loss_p always yields lower FID scores.




  
Effect of 𝒯. <Ref> shows the results of Diff-Cleanse with various threshold 𝒯 for gradient accumulation, where Diff-Cleanse's efficacy in removing the backdoor remains largely unaffected. 

Due to space constraints, we provide more results of ablation studies in <ref>



 §.§ Adaptive attack




After realizing the criteria of Diff-Cleanse, the attacker may train more concealed backdoors. One possible approach is to align the model's predictions for clean and backdoored inputs, which could be defended by our framework. The detailed method and experimental results are in <ref>.




§ CONCLUSION


In this paper, we present Diff-Cleanse, a novel two-stage backdoor defense framework. Extensive experiments show that Diff-Cleanse have a detection accuracy close to 100% and can completely remove backdoors without significantly affecting the model's benign performance.










10

an2024elijah
Shengwei An, Sheng-Yen Chou, Kaiyuan Zhang, Qiuling Xu, Guanhong Tao, Guangyu Shen, Siyuan Cheng, Shiqing Ma, Pin-Yu Chen, Tsung-Yi Ho, et al.
Elijah: Eliminating backdoors injected in diffusion models via distribution shift.
In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pages 10847–10855, 2024.

chen2023trojdiff
Weixin Chen, Dawn Song, and Bo Li.
Trojdiff: Trojan attacks on diffusion models with diverse targets.
In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 4035–4044, 2023.

chen2017targeted
Xinyun Chen, Chang Liu, Bo Li, Kimberly Lu, and Dawn Song.
Targeted backdoor attacks on deep learning systems using data poisoning.
CoRR, abs/1712.05526, 2017.

chou2023backdoor
Sheng-Yen Chou, Pin-Yu Chen, and Tsung-Yi Ho.
How to backdoor diffusion models?
In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 4015–4024, 2023.

chou2023villandiffusion
Sheng-Yen Chou, Pin-Yu Chen, and Tsung-Yi Ho.
Villandiffusion: A unified backdoor attack framework for diffusion models.
NeurIPS 2023 Workshop on Backdoors in Deep Learning, 2023.

dhariwal2021diffusion
Prafulla Dhariwal and Alexander Nichol.
Diffusion models beat gans on image synthesis.
Advances in neural information processing systems, 34:8780–8794, 2021.

fang2024structural
Gongfan Fang, Xinyin Ma, and Xinchao Wang.
Structural pruning for diffusion models.
Advances in neural information processing systems, 36, 2024.

goelimproving
Harsh Goel, Sai Shankar Narasimhan, and Sandeep P Chinchali.
Improving end-to-end autonomous driving with synthetic data from latent diffusion models.
In First Vision and Language for Autonomous Driving and Robotics Workshop.

gu2019badnets
Tianyu Gu, Kang Liu, Brendan Dolan-Gavitt, and Siddharth Garg.
Badnets: Evaluating backdooring attacks on deep neural networks.
IEEE Access, 7:47230–47244, 2019.

guan2024ufid
Zihan Guan, Mengxuan Hu, Sheng Li, and Anil Vullikanti.
Ufid: A unified framework for input-level backdoor detection on diffusion models.
arXiv preprint arXiv:2404.01101, 2024.

ho2020denoising
Jonathan Ho, Ajay Jain, and Pieter Abbeel.
Denoising diffusion probabilistic models.
Advances in neural information processing systems, 33:6840–6851, 2020.

huang2022fastdiff
Rongjie Huang, Max WY Lam, Jun Wang, Dan Su, Dong Yu, Yi Ren, and Zhou Zhao.
Fastdiff: A fast conditional diffusion model for high-quality speech synthesis.
IJCAI2022, pages 4132–4138, 2022.

jeong2021diff
Myeonghun Jeong, Hyeongju Kim, Sung Jun Cheon, Byoung Jin Choi, and Nam Soo Kim.
Diff-tts: A denoising diffusion model for text-to-speech.
Proc. Interspeech 2021, pages 3605–3609, 2021.

karras2022elucidating
Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine.
Elucidating the design space of diffusion-based generative models.
Advances in Neural Information Processing Systems, 35:26565–26577, 2022.

kong2020diffwave
Zhifeng Kong, Wei Ping, Jiaji Huang, Kexin Zhao, and Bryan Catanzaro.
Diffwave: A versatile diffusion model for audio synthesis.
arXiv preprint arXiv:2009.09761, 2020.

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

li2022diffusion
Xiang Li, John Thickstun, Ishaan Gulrajani, Percy S Liang, and Tatsunori B Hashimoto.
Diffusion-lm improves controllable text generation.
Advances in Neural Information Processing Systems, 35:4328–4343, 2022.

liu2018fine
Kang Liu, Brendan Dolan-Gavitt, and Siddharth Garg.
Fine-pruning: Defending against backdooring attacks on deep neural networks.
In International symposium on research in attacks, intrusions, and defenses, pages 273–294. Springer, 2018.

liu2019ppgan
Yi Liu, Jialiang Peng, JQ James, and Yi Wu.
Ppgan: Privacy-preserving generative adversarial network.
In 2019 IEEE 25Th international conference on parallel and distributed systems (ICPADS), pages 985–989. IEEE, 2019.

liu2015deep
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang.
Deep learning face attributes in the wild.
In Proceedings of the IEEE international conference on computer vision, pages 3730–3738, 2015.

lu2022dpm
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu.
Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps.
Advances in Neural Information Processing Systems, 35:5775–5787, 2022.

lu2022dpmplus
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu.
Dpm-solver++: Fast solver for guided sampling of diffusion probabilistic models.
arXiv preprint arXiv:2211.01095, 2022.

mei2023vidm
Kangfu Mei and Vishal Patel.
Vidm: Video implicit diffusion models.
In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pages 9117–9125, 2023.

moghadam2023morphology
Puria Azadi Moghadam, Sanne Van Dalen, Karina C Martin, Jochen Lennerz, Stephen Yip, Hossein Farahani, and Ali Bashashati.
A morphology focused diffusion probabilistic model for synthesis of histopathology images.
In Proceedings of the IEEE/CVF winter conference on applications of computer vision, pages 2000–2009, 2023.

molchanov2019importance
Pavlo Molchanov, Arun Mallya, Stephen Tyree, Iuri Frosio, and Jan Kautz.
Importance estimation for neural network pruning.
In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 11264–11272, 2019.

nguyen2020input
Tuan Anh Nguyen and Anh Tran.
Input-aware dynamic backdoor attack.
Advances in Neural Information Processing Systems, 33:3454–3464, 2020.

nichol2021glide
Alex Nichol, Prafulla Dhariwal, Aditya Ramesh, Pranav Shyam, Pamela Mishkin, Bob McGrew, Ilya Sutskever, and Mark Chen.
Glide: Towards photorealistic image generation and editing with text-guided diffusion models.
arXiv preprint arXiv:2112.10741, 2021.

song2020denoising
Jiaming Song, Chenlin Meng, and Stefano Ermon.
Denoising diffusion implicit models.
arXiv preprint arXiv:2010.02502, 2020.

NEURIPS2019_3001ef25
Yang Song and Stefano Ermon.
Generative modeling by estimating gradients of the data distribution.
In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.

song2020score
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole.
Score-based generative modeling through stochastic differential equations.
arXiv preprint arXiv:2011.13456, 2020.

wang2019neural
Bolun Wang, Yuanshun Yao, Shawn Shan, Huiying Li, Bimal Viswanath, Haitao Zheng, and Ben Y Zhao.
Neural cleanse: Identifying and mitigating backdoor attacks in neural networks.
In 2019 IEEE Symposium on Security and Privacy (SP), pages 707–723. IEEE, 2019.

wolleb2022diffusion
Julia Wolleb, Florentin Bieder, Robin Sandkühler, and Philippe C Cattin.
Diffusion models for medical anomaly detection.
In International Conference on Medical image computing and computer-assisted intervention, pages 35–45. Springer, 2022.

wu2021adversarial
Dongxian Wu and Yisen Wang.
Adversarial neuron pruning purifies backdoored deep models.
Advances in Neural Information Processing Systems, 34:16913–16925, 2021.

xue2020one
Mingfu Xue, Can He, Jian Wang, and Weiqiang Liu.
One-to-n & n-to-one: Two advanced backdoor attacks against deep learning models.
IEEE Transactions on Dependable and Secure Computing, 19(3):1562–1578, 2020.

yang2023video
Siyuan Yang, Lu Zhang, Yu Liu, Zhizhuo Jiang, and You He.
Video diffusion models with local-global context guidance.
arXiv preprint arXiv:2306.02562, 2023.

zhang2022fast
Qinsheng Zhang and Yongxin Chen.
Fast sampling of diffusion models with exponential integrator.
arXiv preprint arXiv:2204.13902, 2022.

zhao2024unipc
Wenliang Zhao, Lujia Bai, Yongming Rao, Jie Zhou, and Jiwen Lu.
Unipc: A unified predictor-corrector framework for fast sampling of diffusion models.
Advances in Neural Information Processing Systems, 36, 2024.

zheng2022data
Runkai Zheng, Rongjun Tang, Jianze Li, and Li Liu.
Data-free backdoor removal based on channel lipschitzness.
In European Conference on Computer Vision, pages 175–191. Springer, 2022.










§ DERIVATION OF X_T^S




In <ref>, we introduce a unified form of backdoor diffusion as presented in <ref>, and derive {x_t^s}_t=1^T, representing a special case of backdoor attacks on DMs. To guarantee that x_t^s could be represented by <ref>, we propose a corresponding backdoor diffusion process as shown in <ref>.

    q^b(x_t^s|x_t-1^s)=𝒩(x_t^s; μ_t^s,w_t^2𝐈),  t∈[1,…,T]

where μ_t^s=√(α)_t x_t-1^s+(1-√(α_t))r^s and w_t=γ√(1-α_t). Obviously, α̂_t=√(α̅_t). Let k_t=1-√(α_t), we can derive the coefficient of r^s at timestep t, ρ̂_t as follows using induction.

    ρ̂_t=k_t+√(α_t) k_t-1+√(α_t α_t-1) k_t-2+⋯+√(α_t …α_2) k_1=1-√(α̅_t),  t∈[1,…,T]

Combine <ref> and <ref>, we could derive the reversed conditional transition q^b(x^s_t-1|x^s_t, x^s_0) as follows. For simplicity and clarity, we omit the superscript s.




    q^b(x_t-1|x_t, x_0) 
        =q^b(x_t |x_t-1, x_0) q^b(x_t-1|x_0)/q^b(x_t |x_0)
       ∝exp(-1/2((x_t-√(α_t)x_t-1-k_t r)^2/w_t^2+(x_t-1-α̂_t-1x_0-ρ̂_t-1r)^2/β̂^2_t-1.. 
       ..-(x_t-α̂_t x_0-ρ̂_tr)^2/β̂^2_t)) 
        =exp(-1/2((x_t-√(α_t)x_t-1)^2-2(x_t-√(α_t)x_t-1) k_t r+k_t^2 r^2/w_t^2.. 
        +(x_t-1-α̂_t-1x_0)^2-2(x_t-1-α̂_t-1x_0) ρ̂_t-1r+ρ̂_t-1^2 r^2/β̂^2_t-1
       ..-(x_t-α̂_t x_0-ρ̂_tr)^2/β̂^2_t)) 
        =exp(-1/2((x_t-√(α_t)x_t-1)^2/w_t^2+(x_t-1-α̂_t-1x_0)^2/β̂^2_t-1-2(x_t-√(α_t)x_t-1) k_t r/w_t^2.. 
       ..-2(x_t-1-α̂_t-1x_0) ρ̂_t-1r/β̂^2_t-1 + k_t^2 r^2/w_t^2 + ρ̂_t-1^2 r^2/β̂_t-1^2-(x_t-α̂_t x_0-ρ̂_tr)^2/β̂^2_t)) 
        =exp(-1/2((α_t/w_t^2+1/β̂^2_t-1) x_t-1^2-2(√(α_t)/w_t^2x_t+α̂_t-1/β̂^2_t-1x_0... 
       ...+(ρ̂_t-1/β̂^2_t-1-√(α_t) k_t/w_t^2) r) x_t-1+C(x_t, x_0, r)))

where α̂_t=√(α̅_t), β̂_t is the fixed variance and C(x_t, x_0, r) is a function that contains x_t, x_0 and r.


Thus, the mean of x_t-1 can be derived as <ref>.


    μ_t(x_t, x_0, r) =    (√(α_t)/w_t^2x_t+α̂_t-1/β̂^2_t-1x_0+(ρ̂_t-1/β̂^2_t-1-√(α_t) k_t/w_t^2) r) /(α_t/w_t^2+1/β̂^2_t-1) 
    
    =    (√(α_t)/w_t^2x_t+α̂_t-1/β̂^2_t-1x_0+(ρ̂_t-1/β̂^2_t-1-√(α_t) k_t/w_t^2) r) w_t^2 β̂^2_t-1/α_t β̂^2_t-1+w_t^2
    
    =    √(α_t)β̂^2_t-1/α_t β̂^2_t-1+w_t^2x_t+α̂_t-1 w_t^2/α_t β̂^2_t-1+w_t^2x_0 +w_t^2 ρ̂_t-1-√(α_t) k_t β̂^2_t-1/α_t β̂^2_t-1+w_t^2r

Based on <ref>, we can obtain the reparametrization x_0=1/√(α̅_t)(x_t-(1-√(α̅_t)) r-√(1-α̅_t)γϵ_t), where ϵ_t∼𝒩(0, 𝐈). Plug the reparametrization into <ref>, we can drivethe reversed diffusion process as follows.

    μ_t(x_t, r) =    √(α_t)β̂^2_t-1/α_t β̂^2_t-1+w_t^2x_t+w_t^2 ρ̂_t-1-√(α_t) k_t β̂^2_t-1/α_t β̂^2_t-1+w_t^2r+√(α̅_t-1) w_t^2/α_t β̂^2_t-1+w_t^2x_0 
    
    =   √(α_t)β̂^2_t-1/α_t β̂^2_t-1+w_t^2x_t+w_t^2 ρ̂_t-1-√(α_t) k_t β̂^2_t-1/α_t β̂^2_t-1+w_t^2r
        +w_t^2/α_t β̂^2_t-1+w_t^21/√(α̅_t)(x_t-(1-√(α̅_t)) r-√(1-α̅_t)γϵ_t)

<Ref> can assist defenders in performing backdoor sampling to defend against attacks requiring specific sampling procedures, such as TrojDiff, without knowing the attacker's formulas.




§ ALGORITHMS







 §.§ Trigger Inversion

<Ref> summarizes our trigger inversion
and backdoor detection method. Line 2 samples a batch of clean noise from standard Gaussian noise. Line 3 randomly initializes the trigger g. Lines 4-9 optimize g based on our proposed objective, using Adam optimizer with a learning rate of 0.1 and a maximum of 100 iterations in all experiments. Line 10 calculates our proposed metric TE with threshold=0.999. Lines 11-14 determine whether the presence of a backdoor. Line 15 returns the inverted trigger and the backdoor indicator.






 §.§ Backdoor Removal

<Ref> outlines the procedure of removing the backdoor and restoring the benign utility of the pruned model. Line 1 randomly chooses clean samples from the clean dataset. Line 2 samples a clean noise. Lines 3-11 calculate the loss of θ against clean inputs and backdoored inputs, as well as the gradients with respect to the substructures θ_i,k. Lines 12-17 calculate the aggregated Taylor importance scores as shown in <ref>. Lines 18-20 remove parameters based on their importance scores and finetune the pruned model with our proposed objective loss_p. Line 21 returns purified model θ^'.






§ MORE EXPERIMENTAL RESULTS




 §.§ Detailed Configuration




We use 27 clean and 161 backdoored DMs, executing <ref> and <ref> once for each model. We consider three existing backdoor attacks for DMs and use all trigger-target pairs from the original paper. For BadDiff and VillanDiff, we use poison rates from 0 to 0.9 and triggers of different sizes to train the model. For TrojDiff, we implement the D2I attack described in the paper and train models with the same hyperparameters as those used in the original paper. Given the diverse attack settings employed, we believe the results should be reliable. To measure the utility of the DM, we sample 10k images from clean sampling to compute the FID. For metrics relative to specifity, ASR, SSIM, ΔASR and ΔSSIM, we sample 100 images from backdoored sampling since we find this sufficient to demonstrate the effectiveness of the trigger or backdoor.


The pruned model was fine-tuned using 10% of the clean dataset (6000 images for the CIFAR-10 dataset and 3000 images for the CelebA-HQ dataset). The DMs trained on CIFAR-10 are finetuned with 40 epochs at a learning rate of 2e-4, and DMs trained on CelebA-HQ are finetuned with 100 epochs at a learning rate of 2e-4.




 §.§ Effect of gradients accumulation threshold



<Ref> presents the results of Diff-Cleanse with various threshold 𝒯 against TrojDiff, demonstrating that our approach is not sensitive to 𝒯. Our experiments across the three attack methods yield consistent results, indicating that neurons sensitive to the backdoor remain stable during sampling. Considering that the selection of 𝒯 has little impact on the effectiveness of Diff-Cleanse, we use 𝒯=900 in experiments outside the ablation study to accelerate computation.





 §.§ Effect of the trigger size



We use BadDiff to backdoor DMs with various trigger sizes and test Diff-Cleanse on them. As shown in <ref>, Diff-Cleanse can detect and eliminate all the backdoors.







 §.§ Effect of the backdoor removal loss



<Ref> compares the performance of ordinary loss function loss_c and our proposed loss_p on a DM attacked by TrojDiff. Moreover, we test 10 DMs and report the average results as shown in <ref>. Our proposed method consistently leads to lower Δ𝐅𝐈𝐃. Sufficient experiments demonstrated that the proposed loss function can improve the effect of fine-tuning.














 §.§ Details of adaptive attack



After realizing the criteria of Diff-Cleanse, the attacker may
train more concealed backdoor diffusion models, where the
difference in activation values of clean noise and backdoored
noise is reduced. To simulate this situation, we design an adaptive version of BadDiffusion. Specifically, we calculate the Euclidean distance between the activation of clean noise and backdoored noise at the same time step, and use it as a weighted term in the loss. Such an adaptive attack aligns the activation in clean sampling and backdoor sampling, making the backdoor-related neurons more stealthy.


As <ref> shows, the proposed adaptive attack method adds a weighted item to the original loss L_θ.

    L^'_θ(x, t, ϵ, g)=L_θ(x, t, ϵ, g)+λΔ a_t ^2

where Δ a_t is the difference between model outputs of the clean inputs and that of the backdoored inputs at time step t and · is the L-2 norm. 


Then we attack DMs with the modified BadDiff, which adopts adaptive attack loss L^'_θ and various λ. We train the DMs 50 epochs and evaluate them every 10 epochs. <Ref> shows the utility of models during training process, indicating that the new loss term has limited effect on model's utility. However, when λ=1e-1, the model's FID score increases continuously during training, indicating that excessive alignment of the activations can lead to training failure. With λ=1e-2, 1e-3, the diffusion models lose little clean generation ability. Except for λ=1e-1, we implant backdoors into all other DMs and <ref> shows the defence performance of Diff-Cleanse. Although the adaptive attack method makes the backdoor-related neurons more concealed, our framework can still remove the backdoor and achieve good image quality in clean sampling.











§ VISUALIZED RESULTS





In this section, we provide additional visual examples. <Ref> and <ref> illustrate the sample results of backdoored DMs before and after purified by Diff-Cleanse. After model repair, the model loses its ability to respond to the trigger and generate target outputs, while its clean utility remains unaffected.


<Ref> depicts the patch-based trigger of TrojDiff and the intermediate results of our algorithm. In this scenario, the trigger is a 3× 3 white block and the target is "Micky". Our method does not invert a trigger that causes the DM to generate the target, but instead activates another backdoor in the model, the colored square in the lower right corner of the images as shown in <ref> (c). Since this pattern appears repeatedly on different samples, we believe Diff-Cleanse could still help the users detect the backdoor created by such attacks.











§ PARAMETER INSTANTIATION FOR ATTACKS




Recall that the unified form of backdoor diffusion is as shown in <ref>, this section describes the diffusion process of existing backdoor attacks. 


BadDiff only attacks DDPM, where β_t is the predefined scalar, α_t=1-β_t and α̅_t=∏_i=1^t α_i. For backdoor diffusion process of BadDiff, α̂_t=α̅_t, ρ̂_t=1-√(α̅_t) and β̂_t=1-α̅_t. When attacking DDPM models, VillanDiff the same α̂_t and β̂_t as BadDiff but doubles the ρ̂_t, so that the injected backdoor can be activated when sampling with advanced samplers. TrojDiff use α̂_t=α̅_t, ρ̂_t=√(1-α̅_t) and β̂_t=(1-α̅_t)γ^2. 




 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.





 Paper Title: Combining audio control and style transfer using latent diffusion.




[
    Amy Navarathna,1,2 Benjamin J. Carey,1,2 James S. Bennett,3 Soroush Khademi,1,2 and Warwick P. Bowen1,2,4*
    August 5, 2024
==============================================================================================================



§ ABSTRACT 

Deep generative models are now able to synthesize high-quality audio signals, shifting the critical aspect in their development from audio quality to control capabilities. Although text-to-music generation is getting largely adopted by the general public, explicit control and example-based style transfer are more adequate modalities to capture the intents of artists and musicians. 

In this paper, we aim to unify explicit control and style transfer within a single model by separating local and global information to capture musical structure and timbre respectively. To do so, we leverage the capabilities of diffusion autoencoders to extract semantic features, in order to build two representation spaces. We enforce disentanglement between those spaces using an adversarial criterion and a two-stage training strategy. Our resulting model can generate audio matching a timbre target, while specifying structure either with explicit controls or through another audio example. We evaluate our model on one-shot timbre transfer and MIDI-to-audio tasks on instrumental recordings and show that we outperform existing baselines in terms of audio quality and target fidelity. Furthermore, we show that our method can generate cover versions of complete musical pieces by transferring rhythmic and melodic content to the style of a target audio in a different genre. 

 




§ INTRODUCTION





Deep generative models are now particularly successful at synthesising high-quality, realistic audio signals. Hence, the major impediment to their broader use in creative workflows is not their audio quality anymore, but rather how end-users can have complete control over the generation process. Following early works on unconditional generation <cit.>, multiple methods proposed to enable control by conditioning generation on semantic tags or audio-descriptors <cit.>. However, such supervised approaches remain limited to the use of explicit descriptors and are constrained by their reliance on annotated datasets. The recent development of language models and representation learning led to impressive performance in text-conditioned generation, mainly relying on transformers <cit.> or diffusion models <cit.>. However, concepts such as timbre, musical style or genres boundaries are usually elusive and highly subjective. Hence, text descriptions might remain limited to common sounds and insufficient to precisely capture musical intentions. 





A parallel stream of research to alleviate those issues is to guide specific aspects of the generation process by providing audio examples. Most approaches in this audio-to-audio editing are focused on timbre transfer, where the timbre of a given sound is applied on the content of another. While some works can transfer any audio to the timbre of a given training set <cit.>; others achieve many-to-many timbre transfer but only between a small set of predefined instrument classes <cit.>. One-shot timbre transfer between different instrument recordings have been achieved using Variational Autoencoders (VAE) <cit.>, but these models rely on a latent bottleneck to enforce disentanglement between timbre and pitch, which hampers their ability to generate high-quality audio on real-world data. More recently, a text-inversion technique was proposed to perform musical style transfer between arbitrary content and style examples <cit.>, but it relies on a large pretrained text-to-music model and requires optimisation prior to each transfer, resulting in very slow inference. 



In this paper, we aim to unify explicit control through audio descriptors or MIDI sequences and style transfer within a single model. To do so, we separate local, time-varying factors of variations and global information, capturing musical structure and timbre in two separate representation spaces. We slightly abuse the terms structure and timbre: by structure, we designate time-varying features, e.g. melody, loudness; by timbre, we designate global features such as actual timbre, but also style or genre. The principle of our method is depicted in Figure <ref>.
Our approach is based on the recently proposed diffusion autoencoder <cit.>, which trains a semantic encoder to condition a diffusion model, in order to achieve both high-quality generation while being able to extract and control high-level features from the data. We extend this approach by building separate representations for timbre and structure, while enforcing their disentanglement with an adversarial criterion combined with a two-stage training strategy. Our method can generate audio matching a timbre target, while specifying the musical structure either with explicit controls (such as MIDI data input) or an audio example. For computation efficiency, our diffusion model operates in the latent space of pretrained autoencoders, resulting in faster than real-time inference on GPU[Experiments were conducted on a NVIDIA A5000 GPU]. 


First, we benchmark our model on a one-shot timbre transfer tasks and demonstrate that our model improves upon existing baselines in terms of audio quality, timbre similarity as well as note onsets and pitch accuracy. On the same dataset, we show that our model can also generate audio from MIDI input and a target timbre example with performances superior to a state-of-the-art MIDI-to-audio baseline. 
Finally, we show that our method can be applied to complete musical pieces and generate cover versions of a track by transferring its rhythmic and melodic content to the style of a target audio in a different genre. We provide audio examples, additional experiments and source code on a supporting webpage[ <https://nilsdem.github.io/control-transfer-diffusion/>]




§ BACKGROUND










 §.§ Diffusion models



Diffusion models (DMs) are a family of generative models that learn to reverse a stochastic process that gradually adds noise to the input data. These models benefit from high-quality generation, stable training and conditioning abilities, which led to their widespread adoption in image <cit.> and audio generation <cit.>. 

Formally, we define a forward process q(𝐱_1:T|𝐱_0) = ∏_t=1^T q(𝐱_t|𝐱_t-1), which is a Markov chain that increasingly adds noise to the data 𝐱_0 by relying on the conditional distribution 


    q(𝐱_t | 𝐱_t-1) = 𝒩(𝐱_t; √(1 - β_t)𝐱_t-1, β_t 𝐈),


where β_t are hyperparameters defining the noise levels at times t ∈{0,T}, with T ∈ℕ. We are interested in learning the reverse diffusion process p_θ(𝐱_t-1 | 𝐱_t), from which we can iteratively denoise a random sample 𝐱_T ∼𝒩(0, 𝐈) to a data sample 𝐱_0 ∼ p(𝐱_0). In a recenty study, <cit.> made a connection between diffusion and denoising score matching <cit.>, leading to a simplified formulation and improved experimental results. The authors show that the reverse process can be approximated by learning a denoising network ϵ_θ that predicts the noise ϵ∼𝒩(ϵ, 0, 𝐈) used to corrupt the data. This results in a simpler training objective 



    min_θ∈Θ  𝔼_t,𝐱_0,ϵ [ ϵ_θ(√(α_t)𝐱_0 + √(1 - α_t)ϵ, t) - ϵ ] ,


where α_t = ∏_s=1^t (1 - β_s), and ϵ_θ is usually parametrized by a UNet <cit.>. 



 §.§ Diffusion autoencoders


DMs naturally yield a series of latent variables 𝐱_1:T through their forward process. However, these stochastic variables built from increasingly adding noise do not capture much semantic information over the data. Although the more recent proposal of Denoising Diffusion Implicit Models (DDIMs) <cit.> extends the original diffusion formulation to a deterministic process allowing each data input to be mapped to a unique latent code 𝐱_T, it still fails to extract and organise high-level features from of the data. Diffusion autoencoders <cit.> alleviate this issue by employing a learnable encoder that compress the data to a semantic latent code 𝐳_sem = E_ϕ (𝐱_0), which then conditions a diffusion decoder. The semantic encoder and the DDIM decoder are trained jointly, following the objective


    min_θ, ϕ𝔼_t,𝐱_0,ϵ [ ϵ_θ(√(α_t)𝐱_0 + √(1 - α_t)ϵ, E_ϕ (𝐱_0), t) - ϵ ]


On image applications, the authors show that the semantic code captures high-level attributes such as person identity, smile or presence of glasses, and can be used for downstream tasks such as conditional generation and attributes manipulation, while achieving state-of-art reconstruction. This approach was also successfully applied to audio <cit.>, where the authors encode magnitude spectrograms into a semantic latent space, allowing to achieve high quality text-conditioned waveform generation. 




 §.§ Control in audio generation


A straightforward approach to extend unconditional generative models in order to provide instruments befitting artistic control is to introduce conditioning on explicit controls. The DDSP model <cit.> proposes explicit pitch and loudness conditioning, while FaderRave <cit.> extended explicit control to non-differentiable time-varying attributes, but both methods remain limited to explicit descriptors and annotated datasets.  

While recent text-to-music methods like MusicGen <cit.> and Music ControlNet <cit.>  have incorporated melody conditioning capabilities, their expressiveness remains constrained by the need to define subjective timbre properties through text prompts. Li et al. <cit.> proposed to use text-inversion in order to extract pseudo-words that represent timbre directly from audio, but their method is computationally intensive and requires to perform an optimisation for each new timbre target. Timbre conditioning directly from a waveform example was also recently proposed the context of bass accompaniment generation<cit.>.







 §.§ Unsupervised disentanglement in sequential data


Many works proposed to model sequential data as a combination of local (time-variant) and global (time-invariant) factors of variation. Notably, the disentangled sequential autoencoder <cit.> relies on simple architecture biases and parameter tuning to obtain disantangled local and global latent variables. Following this work, multiple methods improved the learned representation by explicitely minimizing the mutual information between the two learned variables <cit.>. It was shown that disentanglement can be further improved with contrastive learning and domain-specific transformations that preserve local or global attributes <cit.>.
























More specifically in audio generation, SS-VAE <cit.> employs a Vector-Quantized VAE to achieve disentanglement through compression on quantized structure latent codes, combined with timbre-preserving data augmentations. Luo et al. <cit.> follow a two-stage training strategy similar to ours, and improve disentanglement by enforcing the consistency of the global and local latent variables in style or content transfer. However, both disentanglement strategies degrades reconstruction accuracy, which on top of spectrogram inversion based synthesis leads in poor audio quality. 



§ METHOD


Our approach is based on the assumption that musical audio samples can be seen as specific instances of a set of latent features that are separated between global features that capture style, and time-varying features that capture the local evolution of the signal. Although diffusion models are capable of high-quality conditional generation, they are computationally expensive when dealing with high-dimensional data. Hence, we employ an invertible audio codec to first compress the audio into a low-dimensional latent space, onto which we can build an efficient generative model. We extend the DiffAE <cit.> architecture to two semantic encoders in order to extract separately timbre and structure features from input samples. To further disentangle the learned features and improve transfer as well as explicit control performances, we employ an adversarial training strategy during training. In this section, we detail our proposed model depicted in Figure <ref>.




 §.§ Audio codec


We build our audio codec as a convolutional autoencoder based on the RAVE model <cit.> architecture, featuring the adversarial discriminator recently proposed in <cit.>. The model compresses audio waveforms 𝐱 into an invertible latent sequence 𝐳∈ℝ^L × D, where D and L are the embedding space and time dimensions respectively. On top of the reconstruction and adversarial training objectives of RAVE, we introduce a penalty on the latent codes f(𝐳) = max(0, |𝐳| - 1) to enforce that most latent codes are distributed between -1 and 1. 



 §.§ Model structure



We extract a timbre representation  𝐯_T ∈ℝ^D_T from an audio target using an encoder E_ϕ^T applied to the latent sequence 𝐳 obtained with our audio codec. For structure, we extract a temporal representation 𝐯_S ∈ℝ^L× D_S from either an audio input or an explicit control signal 𝐜 (such as a MIDI sequence), using an encoder E_ψ^S. In the case of an audio input, it would be natural to also infer it from the latent sequence 𝐳. However, we found experimentally that it is particularly difficult to extract fine structure information from the highly-compressed representation 𝐳. Hence, we instead infer structure form the Constant Q Transform (CQT) <cit.> of the target signal, which has been shown to be a well-suited representation for pitch extraction tasks <cit.>. For explicit control, the sequence 𝐜 is directly used as input for E_ψ^S.

To generate audio, we sample a noise vector 𝐳_T and decode it to a latent code 𝐳_0 through latent diffusion conditioned on representations 𝐯_S and 𝐯_T. As diffusion formulation, we leverage the recent improvements introduced in <cit.> and parameterise our denoiser network D_θ to predict the data 𝐳_0 instead of ϵ. E_ψ^S, E_ϕ^T and D_θ are trained end-to-end to minimise the following loss function


    L_diff = 𝔼_t,𝐳_0,ϵ D_θ(√(α_t)𝐳_0 + √(1 - α_t)ϵ, 𝐯_S, 𝐯_T, t) - 𝐳_0


We parameterise D_θ as a 1D convolutional UNet with residual blocks and self-attention layers. The two encoders share a  similar architecture as the encoding half of the UNet, with the difference that the timbre encoder compresses the input temporally and applies average pooling on the time dimension of the last layer. We condition the UNet architecture on 𝐯_S through concatenation with each block inputs. For the timbre vector 𝐯_T we use Adaptative Group Normalisation (AdaGN) <cit.>. 





 §.§ Style and content disentanglement

Although splitting the semantic content between two vectors that are constrained on their dimensions already encourages disentanglement between timbre and structure information, there is no theoretical guarantee regarding their separation. Furthermore, the appropriate feature dimensions are highly dependent on the task and dataset at hand. Hence, to enforce disentanglement without constraints on the dimensions, we introduce a two-stage training combined with an adversarial strategy. First, we freeze the structure encoder and train the model to build an adequate timbre representation. To avoid 𝐯_T to encode all of the information required to reconstruct the target 𝐳, we extract timbre from a different sample 𝐳̃ coming from the same track, following the assumption that it shares the same timbre as 𝐳 but with a different structure. 

In the next stage, we introduce a discriminator D_ζ that tries to predict 𝐯_T from 𝐯_S and is trained to minimise 


    L_D = 𝔼_𝐯_𝕊, 𝐯_𝕋 [ 𝐯_T - D_ζ(𝐯_S)  ].


We train the discriminator alternatively with the encoders and denoising network, which try to minimize the following objective


    L_total =  L_diff - γ L_D,


where γ is an hyperparameter balancing between the reconstruction objective and disentanglement. Indeed, increasing L_D maximises the confusion of the timbre information in the structure space. This restrains the diffusion model from reconstructing 𝐳 solely from 𝐯_S and enables independant control of structure and timbre at inference. 



§ EXPERIMENTS


We aim to asses the ability of our model to generate high-quality audio samples that match characteristics of structure and timbre targets, with the structure being either taken from an audio example or through an explicit control signal. 

MIDI-to-audio For explicit structure control, we evaluate the capability of our model to generate audio from a MIDI score and a target recording for timbre. We compare it to a state-of-art baseline in MIDI-to-audio generation. 

Timbre transfer We evaluate the efficiency of our disentanglement strategy on a task of one-shot timbre transfer between polyphonic mono-instrument recordings. In this case, we consider that structure designates the notes being played (in terms of onset timing, pitch and loudness), while timbre corresponds to the remaining characteristics of the sound. We evaluate our model by randomly sampling timbre and structure examples and evaluate audio quality, timbre similarity as well as note accuracy. We compare our model with two example-based timbre transfer methods on synthetic and real recordings. 




 §.§ Dataset


Synthetic Data The Synthesized Lakh Dataset (SLAKH) <cit.> was generated from the LAKH MIDI collection using professional-grade sample-based virtual instruments. Synthesis parameters as well as audio effects settings were randomly chosen resulting in a very diverse set of timbres. We retain only the individual stems of non-percussive instruments, resulting in 400 hours of audio. 


Real Data To the best of our knowledge there is no multi-instrumental dataset of real recordings that contain a very large number of hours of audio. Hence, we combined the following three datasets to conduct our experiments :

    
  * MaestroV2 : Maestro <cit.> is a piano dataset recorded on Disklavier pianos, capturing both audio and notes played, resulting in approximately 200hours of annotated piano recordings. 

     
  * GuitarSet : Guitarset <cit.> is a collection of live guitar performances with solos and accompaniment from various genres and play styles, with a total audio duration of 6 hours.

     
  * URMP : The URMP dataset <cit.> is composed of pieces played by a large variety of classical instruments. For each piece we retain the mono-instrumental recordings, resulting in approximately 4 hours of audio.


As the GuitarSet and URMP are low sample-size datasets, we add synthetic data stems from SLAKH to the training set to facilitate learning. Furthermore, as the different datasets are greatly imbalanced in terms of sample size, we apply a sampling strategy to even the model performance on each dataset. Following <cit.>, if n_i is the number of samples in a given dataset, we draw examples from this dataset during training with probability (n_i / Σ_j n_j)^0.3.











 §.§ Evaluation metrics


We aim to evaluate how our method is able to match the timbre and structure targets characteristics, while maintaining high-quality audio. 

Audio quality We rely on the widely used Frechet Audio Distance (FAD) <cit.> to evaluate how the generated audio distribution matches the dataset distribution, both for reconstructed and transferred samples. We use the available reference implementation of FAD[https://github.com/gudgud96/frechet-audio-distance/tree/main] and use VGGish <cit.> embeddings of the samples to compute the distance


Timbre Similarity To evaluate timbre similarity we employ the metric proposed in the SS-VAE baseline <cit.>. It relies on a triplet network trained to predict if samples are played by the same instrument based on Mel-frequency cepstral coefficients 2-13. We use their implementation and train the metric on the mixing-secrets[https://www.cambridge-mt.com/ms/mtk/] dataset.

Structure To evaluate if our model is able to reproduce the notes of the structure target, we employ a transcription model <cit.> and compare its output to the ground-truth MIDI data. As metric, we use Onset F1 score from mir-eval, where two notes are considered identical if they have identical pitch and onsets within ± 50 ms of each other. 


 §.§ Baselines


For the timbre transfer experiments, we compare our method to SS-VAE <cit.> and Music Style Transfer <cit.> presented in Section <ref>. We train both models on the real and synthetic datasets, using the official implementation. For explicit control, we evaluate our method against a MIDI-to-audio model <cit.> that was also trained on the SLAKH dataset. 
We use the small configuration of the publicly available pretrained model, as larger models do not fit on our NVIDIA A5000 GPU. 



 §.§ Training details


We start by training our audio codec for 1M steps before training our diffusion model for 500k steps, with an initial timbre learning stage of 100k steps. The overall training takes one day on NVIDIA A5000 GPU. For all experiments we rely on the AdamW optimizer <cit.> with a constant learning rate of 1e^-4 and a batch size of 48. For inference we use the deterministic sampler proposed in <cit.> with 40 diffusion steps. 



§ RESULTS




 §.§ MIDI-to-audio


First, we evaluate our model performance in MIDI-to-audio generation in terms of audio quality, timbre similarity and Onset F1 score. We detail our results in Table <ref> for reconstruction and transfer setups, where the target timbre corresponds to either the instrument of the MIDI sequence or a different sample. In both cases, we obtain higher similarity, as our dedicated timbre embedding captures timbre much more precisely than simple label conditioning on instrument categories. Interestingly, we also obtain better F1 scores, although we did not design our model specifically for MIDI inputs as opposed to <cit.> where authors employ a dedicated note sequence embedding strategy. 

To assess the benefit of our disentanglement strategy, we experiment with bypassing the structure encoder and directly conditioning the UNet on the MIDI sequence (Ours w/o E_S entry in  Table <ref>). This results in overall better Onset F1 score, but degrades timbre similarity and FAD. This demonstrates that our disentanglement strategy improves the capability of the model to precisely render the timbre of the target recording. As described in Section <ref>, the Onset F1 score characterises the difference between the generated notes and the input MIDI sequence. The lower accuracy obtained with our full model in the transfer setup can be explained by the fact that some MIDI sequence are not plausible scores for some target instruments (such as playing chords with a flute). Through the disentangled structure encoding, our model is capable of adapting the input MIDI sequence to the range and capabilities of the target instrument, which results in more realistic sounding samples. We encourage the reader to listen to the examples on our supporting webpage that support this statement. 



 §.§ Timbre transfer


Synthetic data Here, we first evaluate timbre transfer on synthetic data and display the results in Table <ref>. The two baselines appear to provide low audio quality and timbre similarity, and both methods obtain lower Onset F1 scores indicating that they are not able to adequately control structure and timbre independently. Our method improve upon the baselines on all three evaluated aspects, and is interestingly able to reach a comparable performance as in the explicitly conditioned MIDI-to-audio setup.

We also performed an ablation study, by evaluating the effect of applying an information bottleneck on the structure latent space instead of using our adversarial strategy. As mentioned in Section <ref>, the model is able to transfer timbre when D_S is small but achieves a low F1 score. Increasing the latent dimensions improves structure fidelity at the cost of degrading timbre similarity. Using our disentanglement strategy, we are able to employ a 32-dimensional latent space and achieve higher timbre similarity with a slight decrease in note accuracy. Although multiple definitions of timbre transfer are possible, we argue that the most convincing timbre transfer do not necessarily imply a perfect note structure F1 score. In the case of a transfer between monophonic instruments playing in the same pitch range, we can expect all the notes from one recording to be transferred to the other. However, when performing transfer between recordings with very distinct timbre such as a piano playing in its high range and a bass, an interesting transfer would rather be the bass playing the main melodic line a few octaves lower than the piano, which would result in a low F1 score. The improvement in terms of FAD between the distribution of transferred samples and the original data obtained with our disentanglement strategy supports that our method generates more realistic transfers, as an instrument playing notes outside its usual range would be considered as out-of-distribution. 






Real data. We present our results for transfer between real instrumental recordings in Table <ref>. Our model improves even further on the existing baselines for which real instruments timbre seems particularly challenging. Even without adversarial regularisation, our model obtains better FAD, timbre similarity and note accuracy. Our disentanglement strategy further improves timbre match, although the relative decrease in note accuracy appears to be greater than on synthetic data. We believe this is mainly due to a necessary simplification of note structure when transferring complex piano recordings to the mainly monophonic URMP instruments. The improvement of transfer quality captured by the FAD and timbre similarity supports this interpretation. 

Qualitatively, we found that on top of generating realistic samples with the appropriate structure, the model is also able to add characteristic sound artefacts of the target instrument such as fret or hammer noises, as well as matching precise acoustic features of the original recording such as reverb or background noise. 



§ COMPLETE MUSIC STYLE TRANSFER





Finally, we apply our model to the task of creating cover versions of a song that match the style of an example in a different genre. We rely on an in-house dataset of 200 hours of jazz, dub, lofi hip-hop and rock. We use the same model architecture to extract structure from the original track and timbre from the cover targets, with two minor modifications. First, we introduce temporal compression in the structure encoder to avoid 𝐯_T to capture information that is too precisely located in time. Second, we condition the UNet on a BPM time series to help it generating coherent rhythms. Without those modifications, the rhythmic elements from timbre and structure targets are conflicting with each other, resulting in somewhat chaotic generations. We evaluate our model using the cover detection algorithm proposed in <cit.>, which outputs a cover probability based on melodic and harmonic similarities between tracks. To assess style transfer, we rely on the text and audio joint-embedding model CLAP <cit.> and classify genre based on the cosine similarity between the audio and genre label embeddings. 
We compare our method to MusicGen <cit.>, a text-to-music generation model with audio-based melody conditioning. We derive an input prompt from the target genre and use the structure target for melody. 

Our model without regularisation obtains a better cover identification than MusicGen, and our disentanglement strategy further improves transfer resulting in higher genre accuracy. Qualitatively, MusicGen seems to only extract the main melodic idea from the structure audio, whereas our method is able to capture most of the harmonic and melodic content. Furthermore, as our model extract style directly form audio rather than through a text prompt, it transfers the different structural elements towards the instruments actually present in the timbre target rather than just the typical instruments of the genre.





§ CONCLUSION


We presented a simple method to learn disentangled timbre and structure representations. To the best of our knowledge, this is the first model capable of generating realistic, high-quality audio through transfer and MIDI rendering. We leave for future works improvements on the trade-off between reconstruction and disentanglement and applications to more complex musical datasets. Furthermore, we aim to initiate a reflection on the characterisation of the elusive concept of musical style transfer, which we believe to be an exciting stream of research towards a broader use of deep generative models in artistic work-flows. 





























 Paper Title: CREW: Facilitating Human-AI Teaming Research.



Gradient dynamics model for drops of volatile liquid on a porous substrate
    Uwe Thiele
    August 5, 2024
==========================================================================



§ ABSTRACT 
  With the increasing deployment of artificial intelligence (AI) technologies, the potential of humans working with AI agents has been growing at a great speed. Human-AI teaming is an important paradigm for studying various aspects when humans and AI agents work together. The unique aspect of Human-AI teaming research is the need to jointly study humans and AI agents, demanding multidisciplinary research efforts from machine learning to human-computer interaction, robotics, cognitive science, neuroscience, psychology, social science, and complex systems. However, existing platforms for Human-AI teaming research are limited, often supporting oversimplified scenarios and a single task, or specifically focusing on either human-teaming research or multi-agent AI algorithms. We introduce CREW, a platform to facilitate Human-AI teaming research and engage collaborations from multiple scientific disciplines, with a strong emphasis on human involvement. It includes pre-built tasks for cognitive studies and Human-AI teaming with expandable potentials from our modular design. Following conventional cognitive neuroscience research, CREW also supports multimodal human physiological signal recording for behavior analysis. Moreover, CREW benchmarks real-time human-guided reinforcement learning agents using state-of-the-art algorithms and well-tuned baselines. With CREW, we were able to conduct 50 human subject studies within a week to verify the effectiveness of our benchmark.







§ INTRODUCTION













Over the past decade, significant progress in machine learning has increased the potential and necessity for humans to collaborate and interact with Artificial Intelligence (AI) agents. Human-AI teaming has emerged as a research paradigm to explore the dynamic interactions and collaborative processes between humans and AI. By leveraging the complementary strength of both humans and AI, advancements can significantly enhance the overall team performance.

Unlike traditional AI research, which typically focuses on machine learning algorithms in isolation, Human-AI teaming requires a multidisciplinary approach to incorporate insights from various scientific domains. Numerous studies have examined human-human teaming <cit.> with cognitive science, neuroscience, and psychology. Machine learning and robotics communities have extensively researched multi-agent machine learning <cit.>, while team dynamics <cit.> has been explored in complex systems, social science, and network science. Despite the importance and potential of this research paradigm, there is still a lack of a comprehensive and unified platform to benefit research on joint efforts across disciplines and scalable hypothesis verification.

Developing a comprehensive platform for Human-AI teaming research presents several unique challenges. Firstly, the platform needs to support diverse tasks with varying complexities with easy extensions for new tasks or modifications. While reinforcement learning research platforms <cit.> have widely adopted this practice, current Human-AI teaming platforms remain limited to single tasks <cit.>. Secondly, enabling real-time communication through various modalities between multiple humans and heterogeneous AI policies is essential for effective collaboration. However, existing solutions typically support human feedback only through replaying offline trajectories and do not implement real-time feedback mechanisms. Understanding how to build AI that can team with, learn from, be guided by, align with, and augment humans is as crucial as modeling human behavior during interactions with AI. Therefore, the platform must provide interfaces for recording human physiological data alongside agent data, tailored for cognitive science and neuroscience studies. Furthermore, current studies often involve small participant groups, making it difficult to derive generalizable conclusions. Lastly, the absence of a unified platform has limited fully open-sourced studies and the establishment of appropriate benchmarks.

We present CREW, a platform designed to facilitate Human-AI teaming research aiming to address the above challenges. We develop CREW around key principles such as extensible and open environment design, real-time communication support, hybrid Human-AI teaming modes, parallel sessions for scalable experiments, and comprehensive human and agent data collection. Additionally, CREW incorporates highly modular algorithm components compatible with practices in the machine learning community. We demonstrate CREW's potential by benchmarking real-time human-guided reinforcement learning (RL) algorithms alongside various RL baselines. With CREW, we were able to conduct 50 human subject studies within a week. Moreover, CREW includes a set of cognitive tests to explore how individual differences among humans impact their effectiveness in training AI agents. To our knowledge, CREW is the first platform to unify the desired features for Human-AI teaming research across multiple disciplines. We aim for CREW to serve as an infrastructural foundation for multidisciplinary, reproducible, and scalable Human-AI teaming research.




§ RELATED WORK



Human-AI Teaming Research Extensive research has explored Human-AI teaming across various domains. Machine learning studies have developed algorithms to leverage human expertise to improve the accuracy <cit.>, robustness <cit.>, and interpretability <cit.> of models. Integrating human feedback can not only improve performance <cit.> but also align the models with human preference <cit.>. Human-computer interaction research has created interfaces <cit.> and workflows <cit.> that enhance collaboration between humans and AI, combining their strengths to achieve superior performance. Ethical research focuses on understanding and mitigating the societal <cit.>, ethical <cit.>, and technical <cit.> challenges of the rapid advancement and wide adoption of AI. Many fields, including neuroscience, healthcare, robotics, transportation, education, security, and accessibility, have shown growing interest <cit.> in Human-AI teaming. Overall, the broad spectrum of interests highlights the need for multidisciplinary collaboration to drive further advancements.

Human-AI Teaming Platform While significant progress has been made in Human-AI teaming research, there remains an absence of a comprehensive research platform. Overcooked Environment <cit.> is a simplified version of the original popular game to challenge human agents and AI agents in tasks that require close coordination and strategic teamwork. StarCraft II Learning Environment (SC2LE) <cit.> supports adversarial settings to allow Human-AI interaction and learning from human demonstrations. Rapid Integration and Development Environment(RIDE) <cit.> focuses on defense-related scenarios, emphasizing rapid development and integration of AI systems for operational purposes. In addition to real-time decision-making tasks, previous research has developed platforms <cit.> that focus on offline preference or rating settings where humans can provide offline evaluations or corrections with imitation learning or reinforcement learning.



However, existing platforms have more than one of the following limitations that constrain Human-AI teaming research as summarized in Tab. <ref>. Most environments only support one type of task that can be difficult to extend, and the interactions between humans and machine learning agents are limited to mouse and keyboard operations. Moreover, most of the environments only support the data collection on the game data, such as state, action, or reward focusing on the machine learning algorithm development, but none of them support the collection and analysis of human physiological data (eye movement, pupillometry, brain activity, heart impulse, or natural language) to understand human cognitive states and different effects on human subjects. Furthermore, the scale of collaboration has been limited to two agents in a collaborative or competitive setting. Notably, most existing Human-AI teaming studies have only been evaluated on a very small number of subjects or among the authors, which greatly limits our understanding of the state-of-the-art performance. There remains a large gap between the existing platforms and the desired environment to facilitate future research.

Real-Time Human-Guided Learning Most real-world decision-making processes require rapid decisions and adaptation in real time. Therefore, real-time human-guided learning is an essential topic in Human-AI teaming research. Previous research has focused on algorithm design to integrate real-time human feedback into policy training. TAMER and Deep TAMER <cit.> learn to predict human discrete feedback and utilize the feedback model as the value function for policy learning. COACH and DeepCOACH <cit.> addressed the inconsistency of human feedback by modeling it as the advantage function under an actor-critic framework. Recent progress has further refined these algorithms, addressing various challenges in human feedback integration, such as different feedback modalities  feedback stochasticity <cit.>, and continuous action spaces <cit.>. Other works have explored indirectly inferring human objectives from feedback or preference <cit.>. CREW provides an extensive environment for real-time human-guided learning with online training and feedback integration, human physiological data collection, and parallel distributed experiment support.





§ CREW: DESIGN AND COMPONENTS






 §.§ Platform Vision and Design Philosophy




We design CREW to facilitate Human-AI teaming research. Our vision is to construct a unified platform for researchers from diverse backgrounds, allowing them to join forces from human-AI interaction, machine learning algorithms, workflow design, as well as human analysis and training. To achieve this, CREW incorporates the following capabilities, as illustrated in Fig. <ref>.

Extensible and open environment design. CREW provides built-in tasks for rapid development and allows users to integrate customized tasks to accommodate the limitless applications of Human-AI teaming.

Real-time communication. While some Human-AI interaction tasks, such as human preference-based fine-tuning, can be performed offline, many applications require online real-time interaction. Whether it is training decision-making models with real-time human guidance or general human-AI collaboration tasks, the ability to convey messages with minimum delay is essential. Synchronizing data flow between human interfaces, AI algorithms, and simulation engines necessitates the establishment of a real-time communication channel.

Hybrid Human-AI teaming support. Teaming is an essential aspect of our daily jobs. Our vision extends this concept to Human-AI teaming, where both humans and AI operate in teams. Increasing interest in the organization, dynamics, workflow, and trust in multi-human and multi-AI teams highlights the need for a platform capable of distributing and synchronizing tasks, states, and interactions across multiple environment instances and even across physical locations.

Parallel sessions support. A key bottleneck for human-involved AI research is the requirement to conduct experiments with dozens or hundreds of human subjects to obtain trustworthy and reliable conclusions. Such a process can be tedious and time-consuming. To enhance efficiency and scalability, CREW supports multiple independent parallel sessions of the same setting, unconstrained by geographical locations, to obtain the “crowd-sourcing” effects of large-scale experiments. This capability enables experimenters to collectively share experimental data and results.



Comprehensive human data collection. Though human plays an important role in Human-AI teaming, our understanding of human behaviors remains limited and under-explored in existing studies. Therefore, CREW offers interfaces to simultaneously collect multi-modal human data, ranging from active instructions and feedback to passive physiological signals.

ML community-friendly algorithm design. The choice of programming language and libraries should align with the customs and preferences of the ML community. The system design should be modular to allow seamless transitions between tasks and algorithms.





 §.§ Environments



Tasks We select Unity as the simulation engine for CREW due to its popularity in game design and AI research to allow extensible and open environment design. We have implemented four challenging tasks as examples. Multi-player tasks are designed for multi-agent and multi-human teaming research, and single-player tasks are designed for human-guided AI agent learning studies. For each task, we provide both visual and structured state input options. The detailed settings are summarized in Tab. <ref>.

Bowling is a modified version of Atari bowling where each round consists of 10 rolls and the score for each roll is the number of pins hit. Bowling is designed to have the simplest dynamics among our tasks to serve as a rapid test for training a single agent. Find Treasure (Fig. <ref>A) is a single-player embodied navigation task where the agent's goal is to explore a maze and reach the treasure with randomized initial and goal locations. 1v1 Hide-and-Seek <cit.> (Fig. <ref>B) is a one-on-one hide-and-seek task where the seeker learns to explore the maze and catch a moving hider that follows a heuristic policy for obstacle avoidance, and run away from the seeker within line of sight. We introduce this task as an adversarial competition setting. NvN Hide-and-Seek (Fig. <ref>C) is a multiplayer version where multiple seekers and hiders can coordinate, collaborate, and compete. The hiders and seekers can either be controlled by humans or heterogeneous AI policies.






Visual Observations Different visual observations create various challenges in visual embodiment learning for both humans <cit.> and AI agents <cit.>. We provide various camera configurations to support perceptual-motor research. As shown in Fig <ref>, for all our navigation and competitive tasks, we implemented visual observations from multiple views for the users' selection: a top-down fully observable view, a top-down accumulated partially observable view similar to SLAM in robotics <cit.>, a top-down egocentric view, and a first-person view.



Procedural Generation Learning robust, generalizable, and scalable AI agents requires diverse training environments. Procedural generation allows for the creation of a wide range of environment patterns and terrain types. We provide randomized maze patterns where the grids are guaranteed to be connected (Fig. <ref>A) and procedural-generated terrains for more complex simulations as in Fig. <ref>B.





 §.§ Human and Agent Role Assignment



Humans and AI agents often have complementary strengths. For example, humans are generally better at exploring and adapting to new situations, while AI agents are good at repetitive exploitation and precise calculation. Naturally, a team consists of humans and AI agents should have various roles to be effective. Different roles can also be assigned within AI agents to study multi-agent machine learning with heterogeneous policies. To facilitate these experiments, we provide the role assignment feature in CREW.

“Player” allows humans to directly control an agent. “Viewer” allows humans to observe and provide feedback signals as guidance to an AI agent. “Server” role allows humans to host the entire task by monitoring it without direct participation. “AI Agent” assigns different learning policies to each agent.

[13]r0.25


  

  
    < g r a p h i c s >

  

  

  (A) Discrete scalar feedback. (B) Option to take control of the agent and teleoperate. (C) Continuous scalar feedback: the human can hover the mouse over this window to provide per-step feedback.
  


Human Feedback Interface We provide a user interface to allow the Viewer role to provide direct feedback to AI agents shown in Fig. <ref>. Scalar feedback is the most common feedback type used in human-guided RL <cit.>. Conventionally, a human chooses to provide a positive or negative rating to the state-action pair of an agent policy. Additionally, CREW offers an interface that allows humans to provide feedback that is continuous in value and time, enabling granular feedback information. Moreover, our interface also allows humans to directly take control over agents and perform teleoperation.




 §.§ Multiplayer and Parallel Sessions



Enabling multi-human multi-agent sessions requires robust networking solutions (Fig. <ref>). We use Unity Netcode <cit.> for game state synchronization, and Nakama <cit.> as the networking server. In CREW, a server instance hosts the task, runs the simulation, and handles agent policy training, which can be executed on a powerful headless GPU server. Human participants can connect via client instances on less powerful machines, which display synchronized game states and collect human input. CREW is cross-platform, allowing participation from Linux, Windows or MacOS machines.







 §.§ Human and Agent Data Collection



Data collection is at the core of Human-AI teaming research. CREW includes a pipeline for thorough data collection on both the human side and AI agent side. 

Human data Besides the feedback interfaces that collect feedback signals of multiple modalities and teleoperation actions, we also provide interfaces for collecting audio, eye gaze, pupillometry, electroencephalogram (EEG), and electrocardiogram (ECG) physiological responses as in Fig. <ref>.

Agent data including the policy weights, observations, actions, rewards, feedback received, and loss values at every time step can all be saved for further analysis. Users also have the option to enable experiment monitoring and logging by Weights & Biases <cit.>. As all of our tasks include vision-based settings, we also provide implementations of a set of vision encoder architectures.






 §.§ Designing Algorithms



Algorithms research is crucial for Human-AI teaming. We designed the algorithm component of CREW to be extensible and accessible to the ML community. Algorithms are implemented in Python with PyTorch <cit.> and TorchRL <cit.> which is part of the PyTorch ecosystem, making it easy for researchers to design and deploy new algorithms. We provide implementations of state-of-the-art human-guided machine learning algorithms and strong reinforcement learning baselines. The API for communication between a Unity instance and Python algorithm is implemented with MLAgents <cit.>. All environments follow a uniform communication protocol for the observation and action data across tasks. Our modular design allows smooth switching between tasks and algorithms. Real human experiments are time-consuming and costly. To ease algorithm debugging, we implemented simulated human feedback providers as surrogates for real humans for complex tasks. These simulated feedback uses heuristics based on prior task knowledge and is not available in novel tasks in the real world; hence, they should be used solely as debugging tools, not as replacements for real human evaluations.

[16]r0.4


  

  
    < g r a p h i c s >

  

  

  Modular pipeline for experiment setup. Experimenters can freely select and organize the order of cognitive tests and Human-AI experiments with CREW's deployment pipeline.
  






 §.§ Modular Pipeline Design for Quantifying Human Characteristics



Individual differences among humans can significantly affect their teaming with AI agents. To support research along this line, CREW supports a set of cognitive tests to quantify these differences shown in Tab. <ref>. We provide a modular and convenient pipeline (Fig. <ref>) for executing cognitive tests and Human-AI experiments. The framework integrates various media files (e.g., instruction videos or pictures), inter-trial intervals, executable Python scripts, and Unity builds, ensuring a smooth and effective workflow. The pipeline requires minimal effort from researchers during proctoring, as a single click initiates the sequential execution of experiments. The pipeline allows restart from interruption points.








§ BENCHMARKING STUDY




As an example of running experiments with CREW, we benchmark a state-of-the-art real-time human-guided RL framework, Deep TAMER <cit.>, along with strong RL baselines. In the original Deep TAMER, the framework was only tested on  Atari Bowling with 9 human subjects. With CREW, this is the first time it is possible to systematically conduct human-guided RL benchmarking across multiple environments on a larger population. We summarize our findings as well as insights on the scalability of the framework in this section. We also discuss the relationship between human characteristics and guided agent performance.







 §.§ Experiment Setup
 


Tasks We selected 3 single-player games: Bowling, Find Treasure, and 1v1 Hide-and-Seek for this benchmark. For Find Treasure and Hide-and-Seek, each episode has a time limit of 15 seconds. All algorithms directly learn from visual inputs with the top-down accumulated partially observable view. 

Human Trainers We recruited 50 human subjects for the experiments under the approval of the Institutional Review Board. For each 1.5-hour session, the subject will receive $20 compensation. Benefiting from CREW's unique feature to host experiments with parallel sessions, we were able to conduct all 50 experiments within one week. Prior to guiding the agents, the participants were asked to complete all our cognitive tests in Tab. <ref>.


Deep TAMER leverages human feedback as value functions. During training, a human trainer provides positive or negative discrete feedback based on the agent's behavior. This feedback is assigned to relevant state-action pairs through a credit assignment window. A neural network is trained to predict human feedback given state-action pairs, and the policy chooses actions that maximize this predicted feedback. Originally, Deep TAMER was limited to discrete actions. We extend it to continuous action spaces using an actor-critic framework, termed as c-Deep TAMER. For bowling, the agent is trained for 5 minutes, and for 10 minutes in Find Treasure and Hide-and-Seek.

Baseline RL We selected two state-of-the-art RL algorithms: Deep Deterministic Policy Gradient (DDPG) <cit.> and Soft Actor-Critic (SAC) <cit.>. 

Heuristic feedback We also evaluate simulated feedback-guided RL. We simply add the feedback signals as additional dense rewards to the environment reward. DDPG is selected as the backbone RL algorithm as the state-of-the-art transitioned from SAC to DDPG <cit.> in visual control tasks.

Evaluation We checkpoint model weights every 2 minutes and evaluate on unseen test environments. For bowling, every checkpoint is evaluated for 1 game (10 rolls). For Find Treasure and 1v1 Hide-and-Seek, the checkpoints are evaluated for 100 episodes.




 §.§ Results



Agent training performance We hypothesize that subjects with higher cognitive tests can lead to higher-performing agents. Therefore, we show the agent performances guided by the 15 subjects who scored the highest in our cognitive tests side by side with the performance of all 50 subjects in Fig. <ref>. As shown in the results, the agents guided by the top 15 subjects exhibit higher performance than the overall average. In particular for the top 15 subjects, on the simple bowling task, c-Deep TAMER surpasses RL baselines by an average of 10 scores given the same training time. On Find Treasure, heuristic feedback achieved the highest performance as expected, showing the upper bound performance with accurate and non-delayed dense rewards. c-Deep TAMER also demonstrated strong performance with faster learning trends than RL baselines. On 1v1 Hide-and-Seek, c-Deep TAMER performed similarly to RL baselines, suggesting that c-Deep TAMER has difficulty scaling to tasks with higher complexity. Similar conclusions still hold for all 50 subjects.













Analysis of Individual Differences Due to the cognitive test feature and emphasis on Human-AI teaming, we can deepen our understanding of how individual human differences can affect the performance of human-guided agents. We calculated the correlation between human subjects' cognitive test scores and c-Deep TAMER training results in Fig. <ref>. The cognitive test scores are normalized by the mean and variance over the subjects through z-score, and outliers with scores 1.5× Interquartile Range above the third quartile or below the first quartile are removed. We found that non-outlier subjects performed equally well on the spatial mapping test so we do not include this for correlation analysis. The heatmap shows the sign of correlation coefficients from linear regression and the statistical significance (*).


[17]r0.5


  

  
    < g r a p h i c s >

  

  

  Correlation between cognitive test scores and c-Deep TAMER training performance. The darker the color, the more statistically significant the correlation. “+” or “-” means positive or negative correlation.
  

The Find Treasure score has a significant positive correlation with theory of behavior, mental rotation, and fitness scores, suggesting that subjects who performed better on these tests also trained better agents. This is likely due to the need for good spatial reasoning and future behavior estimation in Find Treasure. The 1v1 Hide-and-Seek score shows a significant positive correlation with reflex score, likely due to the rapid changes in environment and task dynamics in the complex adversarial setting. Subjects with faster reaction speeds provided timely feedback to align with relevant state-action pairs. Overall, the total score (i.e., the sum of three game scores) has significant positive correlations with rotation, fitness, and overall score, suggesting that these cognitive skills are most relevant to the performance of c-Deep TAMER-guided agents.




§ CONCLUSION, LIMITATION, AND FUTURE WORK



We introduce CREW for facilitating Human-AI teaming research from diverse human and machine learning scientific communities. CREW offers extensible environment design, enables real-time human-AI communication, supports hybrid Human-AI teaming, parallel sessions, multimodal feedback, and physiological data collection, and features ML community-friendly algorithm design. We also provide a set of built-in tasks and baseline algorithm implementations. Using CREW, we benchmarked a state-of-the-art human-guided RL algorithm against baseline methods involving 50 human subjects and provided insights into the relationship between individual human differences and agent-guiding performance.

CREW is still in the early efforts among several critical aspects. Future work will explore building more diverse and challenging tasks, including multiplayer tasks with complex strategies and robotics environments requiring an understanding of physics. While we have only benchmarked several algorithms, we hope CREW can help benchmark many existing algorithms that were not fully open-sourced in a unified environment. Finally, more supports on human physiological data processing and analysis shall be investigated and supported in CREW.



§ ACKNOWLEDGEMENTS

We thank the members at Duke General Robotics Lab, Sida Zhu and Harish Yerra for early development on the platform, and Jiaxun Liu and Easop Lee for beta testing. This work is supported in part by ARL STRONG program under awards W911NF2320182 and W911NF2220113.













































































































    

















§ ACCESSING CREW

CREW is a platform for cross-disciplinary human-AI teaming research. The codebase for our platform, including environments, interfaces, and algorithms can be accessed through <https://github.com/generalroboticslab/CREW>. Videos, documentation, and tutorial are hosted on <http://generalroboticslab.com/CREW>. Our platform is fully open-sourced for academic use.








§ PLATFORM DETAILS

The environments of CREW is implemented using , with packages  <cit.>,  <cit.> and  <cit.>. Algorithms are developed with  <cit.>.




§ ADDITIONAL RESULTS




 §.§ Feedback Visualization

In Fig <ref> we show examples of the state-action pairs and human assigned feedback value during c-Deep TAMER training. Since humans trainers do not see the action (next navigation destination) explicitly, we show consecutive frames instead.





 §.§ Full Cognitive Test Analysis

We include the full results for the cognitive test and c-Deep TAMER score analysis. All linear regression plots including coefficients and p-values and summarized in Fig <ref>.









§ COMPUTATIONAL RESOURCES


All human subject experiments were conducted on desktops with one NVIDIA RTX 4080 GPU. All evaluations were run on a headless server with 8 × NVIDIA RTX A6000 and NVIDIA RTX 3090 Ti.





§ BENCHMARKING EXPERIMENT DETAILS



 §.§ Implementation details


c-Deep TAMER Algorithm
The detailed c-Deep TAMER algorithm is summarized in Alg <ref>. It is modified from the original Deep TAMER integrated with an actor network for action selection. The actor is updated by gradient ascent that outputs the actions that maximizes the predicted human feedback.


Processing Inputs All tasks in our experiments directly learn from pixels. The frames rendered from the environment are first resized to 100×100 pixels. For Find Treasure and 1v1 Hide-and-Seek, we stack the past 3 frames as input. We then apply simple random shift to the frame stack, as it has been shown to be an effective augmentation strategy for visual reinforcement learning. The shifting factors along the height and width are uniformly sampled from [0, 0.08].


Model architecture We use a 3-layer CNN as the vision encoder, each having 64 channels and followed by batch normalization and ReLU activation function. All actor-critic frameworks uses a 3-layer MLP with 256 neurons in each layer for both the actor and critic network. 



 §.§ Hyperparameters

The hyperparameter settings for out experiments is summarized in Tab <ref>.






 §.§ Human subject experiment details


Overview
We recruit the human subject on campus by using flyers and emails. For every human subject, the experiment time is approximately 40 minutes without interruptions. The experiment starts with cognitive tests (10 minutes)  and is followed by the human guiding the agent using the c-Deep TAMER framework (30 minutes). The order of the cognitive tests is Eye Alignment, Reflex, Theory of Behavior, Mental Rotation, Mental Fitting, and Spatial Mapping. There are detailed instruction videos for each test before the test starts. As for the human guiding agent part, each human subject guides the agent to play 3 games for a total of 30 minutes(5 minutes for Bowling, 10 minutes for Find Treasure, and 10 minutes for 1v1 Hide-and-Seek). Before each game, there is a detailed instruction video that describes the rule of the game and how human subjects can give feedback to the agent, which we included below.



  §.§.§ Instruction Video Script


 General Introduction:  
 
 

Cognitive Test Introduction: 



Eye Alignment Instruction: 



Reflex Instruction: 



Theory of Behavior Instruction: 



Mention Rotation and Mental Fitting Instruction: 



Spatial Mapping Instruction: 



Human Guiding AI Introduction: 



Bowling Instruction: 



Find Treasure Instruction: 



1v1 Hide-and-Seek Instruction: 






  §.§.§ Compensation

We pay each human subject $20 for participation. 



§ AUTHOR STATEMENT

The authors bear all responsibility in case of violation of rights.




































 Paper Title: Exploiting Preferences in Loss Functions for Sequential Recommendation via Weak Transitivity.














  Omnious AI
  Seoul
  South Korea

hyunsoo.chung@omnious.com



  University of Pittsburgh
  Pittsburgh
  Pennsylvania
  United States

jungtaek.kim@pitt.edu



  Omnious AI
  Seoul
  South Korea

hyungeun.jo@omnious.com



  Omnious AI
  Seoul
  South Korea

hyungwon.choi@omnious.com


§ ABSTRACT 
A choice of optimization objective is immensely pivotal in the design of a recommender system as it affects the general modeling process of a user's intent from previous interactions. 
Existing approaches mainly adhere to three categories of loss functions: pairwise, pointwise, and setwise loss functions.
Despite their effectiveness, a critical and common drawback of such objectives is viewing the next observed item as a unique positive while considering all remaining items equally negative.
Such a binary label assignment is generally limited to assuring a higher recommendation score of the positive item, neglecting potential structures induced by varying preferences between other unobserved items.
To alleviate this issue, we propose a novel method that extends original objectives to explicitly leverage the different levels of preferences as relative orders between their scores.
Finally, we demonstrate the superior performance of our method compared to baseline objectives.








<ccs2012>
 <concept>
  <concept_id>10002951.10003317.10003347.10003350</concept_id>
  <concept_desc>Information systems Recommender systems</concept_desc>
  <concept_significance>500</concept_significance>
  </concept>
</ccs2012>


[500]Information systems Recommender systems














Exploiting Preferences in Loss Functions for Sequential Recommendation via Weak Transitivity
    Hyungwon Choi
    August 5, 2024
============================================================================================









§ INTRODUCTION




Distinguishing items by relevance to an each user's previous interaction history is the essence of most applied sequential recommender systems.
Conventional approaches <cit.> frame this task as learning the feature representations of a chronologically-ordered list of user's interacted items and a candidate item.
By computing the inner product between two vectors, we then obtain a recommendation score that quantitatively indicates the relevance.
On that basis, a prevalent training process reformulates the problem into a supervised learning framework with binary labels, where the next interacted item of each user becomes a unique positive to its sequence of previous items.
Such a procedure often accompanies negative sampling <cit.> among the rest of items since the score computation of all items during each training step inevitably results in severe inefficiency <cit.>.


In this line of research, numerous approaches have adopted advanced neural networks as feature encoders <cit.> to capture more complex correlations.
Apart from vast architectural improvements, the majority of models yet utilize one of three types of loss functions as an optimization objective:
pairwise, e.g., BPR <cit.>,
pointwise, e.g., BCE <cit.>,
and setwise, e.g., SSM <cit.> functions.
While the formulation of each function varies, they all share the desired result of increasing the score of the unique positive than the scores of the other items.
Notably, most recent studies on loss functions and negative sampling strategies <cit.> aim for improved training efficiency and robustness to false negatives within three core objectives.
Despite their effectiveness, the model trained with such objectives consequently learns to regard unobserved items as equally negative with their labels simply set to zero.
In real-world recommendation, however, a subset of negative items is occasionally more favorable than the other due to various side factors, e.g., item popularity.
Unfortunately, the current scheme of binary label assignments hinders from fully taking advantage of diverse levels of preferences among unobserved items.


To address the aforementioned issue, we first derive an inductive relation, dubbed weak transitivity, which represents preference-driven orders of item scores.
We then propose novel extensions of original loss functions that directly leverage this weak transitivity in their forms.
Consequently, the recommendation scores of unobserved items are aligned with respect to their preferences.
Figure <ref> highlights such a property of resulting recommendation policy trained with our proposed family of objectives.
Opposed to the results of BPR <cit.> and BCE <cit.>, item scores from our transitive extensions TransBPR_pop and TransBCE_pop (horizontal hatch) are generally proportional to their popularity.
Meanwhile, scores from TransBPR_niche and TransBCE_niche (diagonal hatch) are inversely proportional to their popularity, favoring more niche (i.e., unique) items.
To cap it all, we serve the predefined preferences of items as an additional supervisory factor for their recommendation scores.


It is noteworthy that in this work we do not propose any new distributions for negative sampling but instead introduce the modification of original objectives.
Hence, we solely utilize a combination of item popularity and uniform distributions for negative sampling, easily accessible in implicit settings.
We validate the effectiveness of our proposed extensions compared to the original loss functions and their renowned variants on four sequential recommendation benchmarks.
In all settings, our approaches substantially improve the recommendation performance compared to baseline methods.






§ BACKGROUND



In this section, we describe a sequential recommendation task and the details of representative loss functions for recommendation.




 §.§ Problem Statement


Let 𝒰 = {u_1, …, u_M} and ℐ = {i_1, …, i_N} denote a set of users and items where M and N are maximum numbers of users and items, respectively.
Given a chronologically ordered history h_u = {i_1^u, …, i_t^u} of observed items for a user u,
a goal of sequential recommendation is to recommend the most relevant next item i_t+1^u.
We first embed the history h_u and a candidate item i onto vectors h'_u and i', respectively.
Accordingly, a recommendation score ŝ_ui is calculated via an inner product between two embeddings.




 §.§ Training Objectives





  
Pairwise Objective 
Bayesian personalized ranking (BPR) <cit.> models personalized ranking of items.
It forces the score of the positive (i.e., next interacted) item to be higher than the scores of the rest unobserved items.
With a dataset 𝒟_s composed of a triplet (u,i,j) of which item i as a positive and item j as a sampled negative to a user u,
the corresponding loss is formulated as follows:

    ℒ_BPR = -
        ∑_(u,i,j) ∈𝒟_slogσ(ŝ_ui - ŝ_uj),

where σ(·) is a sigmoid function.
A resulting term σ(ŝ_ui - ŝ_uj) is the probability of a user u preferring an item i more than an item j. 




  
Pointwise Objective
Similar to the pairwise loss form, binary cross-entropy (BCE) <cit.> casts the recommendation problem into binary classification with a single sampled negative.
Hence with a dataset 𝒟_s of a triplet (u,i,j), the loss is defined as below:

    ℒ_BCE = -
        ∑_(u,i,j) ∈𝒟_slog(σ(ŝ_ui)) + 
        log(1 - σ(ŝ_uj)).





  
Setwise Objective
The sampled softmax loss (SSM) <cit.> turns the problem into a multi-class classification with a finite number of sampled negatives.
With a dataset 𝒟_m of a triplet (u,i, 𝒩_j) consisting of the positive i and a set of multiple negative samples 𝒩_j to the user u, the loss optimizes the probability of the positive as:

    ℒ_SSM = -
        ∑_(u,i,𝒩_j) ∈𝒟_mlogexp(ŝ_ui)/exp(ŝ_ui) + 
        ∑_j ∈𝒩_jexp(ŝ_uj).











§ OUR METHOD



Here, we introduce weak transitivity between unobserved items,
and then describe how to integrate it to original objectives.




 §.§ Weak Transitivity



Inducing orders on recommendation scores of unobserved items requires a model to sample two or more negatives differing in preferences.
In a basic form, let p_1 and p_2 represent two different sampling distributions for negatives.
Given a user u and its positive i,
we sample a negative j from p_1 and k from p_2 where i ≠ j ≠ k.
Since the positive i is the most preferred item to the user u,
a transitive relation ŝ_ui > ŝ_uj > ŝ_uk holds true when j is more preferred than k.
Strict transitivity then corresponds to a scheme where any j from p_1 is guaranteed to be more preferred than any k from p_2.
However, when the sampled negative j is actually less preferable than k, the relation ŝ_ui > ŝ_uj > ŝ_uk is violated.
We refer this scheme as weak transitivity which allows such occasional violations.


The concept of transitivity with multiple negatives is straightforward as well.
Instead of sampling a single negative from each distribution, we sample a set of negatives 𝒩_j = {j_1, …, j_n} from p_1 and 𝒩_k = {k_1, …, k_n} from p_2.
Consequently, with ŝ_𝒩_j = {ŝ_uj_1, …, ŝ_uj_n} and ŝ_𝒩_k = {ŝ_uk_1, …, ŝ_uk_n},
a transitive relation is modified as ŝ_ui > max(ŝ_𝒩_j) > min(ŝ_𝒩_j) > max(ŝ_𝒩_k) with the same criteria for weak and strict cases.




 §.§ Extensions with Weak Transitivity



We propose novel extensions of original training objectives that resolves the limitation of binary label assignments by incorporating the derived transitive relation to the loss formulation.
We first introduce two sampling schemes by utilizing an item popularity distribution p_pop and a uniform distribution p_unif:

    𝒟'_s(pop)    = {(u,i,j,k) | j ∼ p_pop, k ∼ p_unif},
        
    𝒟'_s(niche)    = {(u,i,j,k) | j ∼ p_unif, k ∼ p_pop},

where j is the more preferred negative and k is the less preferred one. 
The mini-batch 𝒟'_s(pop) assumes that a user generally prefers popular items whereas 𝒟'_s(niche) regards niche items more favored.
Our extension of the BPR objective with transitivity is then given by the following:

    ℒ_TransBPR
        = -
        ∑_(u,i,j,k)∈𝒟'_slogσ(ŝ_ui - ŝ_uj)_original + 
        γlogσ(ŝ_uj - ŝ_uk)_preference,

where γ is a balancing coefficient for two terms.
In essence, the preference term encourages the score of j to be higher than that of k while the original term assures it to be smaller than that of i.
As a consequence, our objective explicitly imposes the transitive relation ŝ_ui > ŝ_uj > ŝ_uk to recommendation scores. 
Combining the proposed formulation with the previously introduced sampling schemes, we obtain two distinct training objectives as follows:

    TransBPR_pop   =
        ℒ_TransBPR(𝒟'_s(pop); θ),
    TransBPR_niche   =
        ℒ_TransBPR(𝒟'_s(niche); θ),

where θ is learnable parameters of the model.
Similarly, the extension of the BCE function is formulated as below:

    ℒ_TransBCE =     -
        ∑_(u,i,j,k) ∈𝒟'_s[ log(σ(ŝ_ui)) + 
        log(1 - σ(ŝ_uj))
        
        + γ(
        log(σ(ŝ_uj)) + 
        log(1 - σ(ŝ_uk)) )].

Naturally, TransBCE_pop and TransBCE_niche are two consequent training schemes with 𝒟'_s(pop) and 𝒟'_s(niche), respectively.
For the setwise loss function, we sample a set of items 𝒩_j and 𝒩_k from each distribution instead of sampling a single item j or k:

    𝒟'_m(pop)    = {(u,i,𝒩_j,𝒩_k) |𝒩_j ∼ p_pop, 𝒩_k ∼ p_unif},
        
    𝒟'_m(niche)    = {(u,i,𝒩_j,𝒩_k) |𝒩_j ∼ p_unif, 𝒩_k ∼ p_pop}.

A corresponding extension for the SSM loss function is given by:

    ℒ_TransSSM =     -
        ∑_(u,i,𝒩_j,𝒩_k) ∈𝒟'_m[
        logexp(ŝ_ui)/exp(ŝ_ui) + 
        ∑_j ∈𝒩_jexp(ŝ_uj)
       + γ1/|𝒩_j|∑_j ∈𝒩_jlogexp(ŝ_uj)/exp(ŝ_uj) + 
        ∑_k ∈𝒩_kexp(ŝ_uk)].

As similar to previous extensions, TransSSM_pop and TransSSM_niche employs either 𝒟'_m(pop) or 𝒟'_m(niche) for negative sampling.


One thing to note is that our introduced sampling strategies 𝒟'_s and 𝒟'_m are both weak transitive.
Nonetheless, modifying them to a strict setting can be readily accomplished if we can quantify the preference of each item.
For instance, we can reformulate the mini-batch construction for pairwise and pointwise objectives as:

    𝒟^†_s(pop)    = {(u,i,j,k) | f(j) > f(k), j ∼ p_pop, k ∼ p_unif},
        
    𝒟^†_s(niche)    = {(u,i,j,k) | f(j) < f(k), j ∼ p_unif, k ∼ p_pop},

where f(·) denotes a function that measures the popularity of an item.
However, we argue that such strict transitivity rather hurts the quality of the resulting recommendation policy.







§ EXPERIMENTS




In this section, we conduct experiments to compare our proposed extensions to original objectives and their variants.




 §.§ Experimental Setup





  
Datasets
We employ public sequential recommendation tasks from different domains:
Beauty, Toys, and Sports,
which are product review datasets introduced by Amazon.com <cit.>,
and Yelp, which is a widely tested business recommendation dataset. 
Detailed statistics of preprocessed datasets are reported in Table <ref>.




  
Evaluation Settings and Metrics
For dataset partitioning, we adopt the conventional leave-one-out strategy <cit.> to assure the quality of the trained recommender system.
Then, we recommend K items with the highest recommendation scores from the entire item pool.
For evaluation, we adopt two common top-K metrics, HR@K and NDCG@K with K fixed to 10.




  
Baselines
We fix a model architecture to SASRec <cit.>
and switch only a training objective. 
For baselines, we compare our method against BPR <cit.>,
GBPR <cit.>,
BCE <cit.>,
SSM <cit.>,
InfoNCE <cit.>,
BPR-DNS <cit.>,
and gBCE <cit.>.
Here, gBCE_1 and gBCE_𝒩 denote gBCE with a single negative and multiple negatives, respectively.




  
Hyperparameters
For all objectives, We train with a fixed batch size of 256, a learning rate of 0.0003, and a maximum sequence length of 50.
The SASRec model is with 2 layers and 1 attention head with an embedding dimension of 256. 
For setwise objectives, we sample 100 negatives in total.
Our proposed TransSSM sets cardinality of 𝒩_j and 𝒩_k to 50 such that they sum to 100.
A balancing coefficient γ is selected from {0.5, 1.0, 1.5}.




 §.§ Performance Comparison



Table <ref> summarizes the performance of models trained with different optimization objectives. 
In general, our proposed objectives outperform three original objectives and their notable variants in both metrics within all benchmarks.
Particularly, we recognize TransSSM_pop obtain the highest performance among all baselines in all datasets.
Such results demonstrate that inducing preference order through weak transitivity consistently improves the performance regardless of base loss functions.
Typically, we observe our extensions with popularity preference (e.g., TransSSM_pop) achieve commonly improved metrics compared to niche preference (e.g., TransSSM_niche).
Hence, exploiting popularity as the preference indicator is particularly effective regardless of datasets.
Though, we find recommendations with niche preference performs on par with best-performing baselines or often even better.







 §.§ Transitivity Analysis



As seen in Table <ref>,
we observe comparatively higher performance with weak transitivity.
We hypothesize that the number of uninformative (i.e., easy) negatives <cit.> increases more within the strict transitivity as training proceeds.
To validate the claim, we report how average value of each term in the BPR extension,
logσ(x̂_uij) and logσ(x̂_ujk), changes throughout the training with different schemes.
Here, we compare three classes of transitivity:
Weak, i.e., (<ref>) and (<ref>),
Strict, i.e., (<ref>) and (<ref>),
and Disjoint, i.e., a specific case of Strict that sets 50% of total items with high preference as a support for p_pop and the rest for p_niche.
Results in Figure <ref> illustrate more steep decline of the preference term, logσ(x̂_ujk), as the transitivity becomes more strict.
Thus, weak transitivity provides the most informative gradients, increasing the training effectiveness.




§ CONCLUSION


In this work, we identified the common and crucial disadvantage of binary label assignments within conventional recommendation objectives.
To overcome this limitation, we have proposed novel extensions that directly exploit the preference differences of unobserved items.
Through extensive experiments, we demonstrated the effectiveness of our method with the thorough analysis.
As future work,
the uniformity of normalized embeddings <cit.> can be applied to our framework to improve recommendation performance more.


ACM-Reference-Format




 Paper Title: Enhancing Deep Hedging of Options with Implied Volatility Surface Feedback Information.


 

a]Pascal François
b]Geneviève Gauthier
c,d]Frédéric Godin 
c]Carlos Octavio Pérez MendozaCorresponding author. 
  Email addresses: 
mailto:pascal.francois@hec.capascal.francois@hec.ca (Pascal François),
mailto:genevieve.gauthier@hec.cagenevieve.gauthier@hec.ca (Geneviève Gauthier), mailto:frederic.godin@concordia.cafrederic.godin@concordia.ca (Frédéric Godin), mailto:carlos.octavio.perez92@gmail.comcarlos.octavio.perez92@gmail.com (Carlos Octavio Pérez Mendoza). 






[a]Fellow, Canadian Derivatives Institute, Department of Finance, HEC Montréal, Canada
[b]GERAD and Department of Decision Sciences, HEC Montréal, Canada
[c]Concordia University, Department of Mathematics and Statistics, Canada
[d]Quantact Laboratory, Centre de Recherches Mathématiques, Canada











Enhancing Deep Hedging of Options with Implied Volatility Surface Feedback InformationWe thank John Hull and Zissis Poulos for valuable feedback about this work. Gauthier is supported by the Natural Sciences and Engineering Research Council of Canada (NSERC, RGPIN-2019-04029), a professorship funded by HEC Montréal, and the HEC Montréal Foundation. Godin is funded by NSERC (RGPIN-2024-04593).
 
    [
     
August 5, 2024
=============================================================================================================================================================================================================================================================================================================================================================================================================

empty 







§ ABSTRACT 



We present a dynamic hedging scheme for S&P 500 options, where rebalancing decisions are enhanced by integrating information about the implied volatility surface dynamics. The optimal hedging strategy is obtained through a deep policy gradient-type reinforcement learning algorithm, with a novel hybrid neural network architecture improving the training performance. The favorable inclusion of forward-looking information embedded in the volatility surface allows our procedure to outperform several conventional benchmarks such as practitioner and smiled-implied delta hedging procedures, both in simulation and backtesting experiments. 





	


JEL classification: C45, C61, G32.



Keywords: Deep reinforcement learning, optimal hedging, implied volatility surfaces.




empty 


roman





arabic








§ INTRODUCTION



Since the advent of the <cit.> framework, dynamic hedging has become a standard financial risk management tool for managing the risk associated with options portfolios. The <cit.> framework has the remarkable property that delta hedging –a hedging strategy invested exclusively in the underlying asset and the money market account– achieves the perfect replication of a European-style contingent claim. In practice, this property is lost due to frictions. Most notably, considering the infrequent rebalancing of the hedging portfolio, classic delta hedging, which is inherently local, can no longer protect against infinitesimal shocks in the underlying asset price. Such an imperfect hedge inevitably yields a hedging error that has to be managed. 

Several works have subsequently extended the idealized setting of the Black-Scholes framework to account for imperfect hedging, incorporating features such as discrete-time rebalancing <cit.>, transaction costs <cit.>, trading constraints <cit.> or liquidity costs <cit.>. Another very important avenue for the development of a hedging procedure is the computation of the delta (and other "Greek" sensitivity parameters) which purely relies on market data and does not require stringent postulates about stochastic dynamics of associated risk factors: see among others, <cit.>, <cit.>, and <cit.>. 
This approach utilizes implied volatility (IV) surfaces to derive the Greeks, specifically the option delta and gamma, based on a mild scale-invariance assumption for the underlying asset return distribution. 

This paper builds on existing literature by developing a data-driven approach that incorporates IV surface information to derive optimal hedging positions. Unlike traditional methods focusing on local conditions, we adopt a global perspective, minimizing the risk metric associated with the terminal hedging error in a multi-period horizon framework.
We leverage developments in risk-aware reinforcement learning (RL) to find the sequence of hedge ratios optimizing the hedger's total risk until the option expiry, a setup analogous to <cit.>'s deep hedging approach. 
The novelty of our work lies in integrating factors that influence IV surface dynamics into the state variable set for determining hedging positions. The hedging agent therefore relies on market expectations for the underlying return distributions over several temporal horizons, producing genuine optimal forward-looking multi-stage hedging decisions informed by IV surfaces. To conduct our numerical experiments, we use the JIVR model of <cit.>, which is a parsimonious, tractable and data-driven econometric model representing the joint dynamics of the underlying asset return and five interpretable factors driving the IV surface of the S&P 500. Notably, the model has been calibrated on a data period spanning more than 25 years of daily option data, reflecting market behavior in a wide array of scenarios, including several financial crises.

As an additional contribution, we propose a novel hybrid neural network architecture combining feedforward and long-short-term memory (LSTM) layers, which is shown the improve training performance over conventional architecture.

Hedging procedures with risk-aware reinforcement learning procedures have received substantial attention from the literature recently, see for instance <cit.>, <cit.>, <cit.>, <cit.>, <cit.>, <cit.>, <cit.>, <cit.>, <cit.>, <cit.>, <cit.>, <cit.>, <cit.>, <cit.>, <cit.> and <cit.>. Nevertheless, to the best of our knowledge, our paper is the first to conduct multi-stage RL-based hedging that incorporates IV surface information directly as state variables.

Our RL approach is shown to substantially outperform delta-based hedging strategies acting as benchmarks. In particular, RL agents trained with asymmetric objective functions such as the conditional value-at-risk (CVaR) or the semi-mean squared-error (SMSE) offer superior tradeoffs between profitability and downside risk. 
The outperformance of RL agents is even more pronounced in the presence of transaction costs, as such agents manage to develop hedging strategies that remain efficient in terms of risk mitigation while generating lower turnover. 
A feature importance analysis highlights that the conditional variance of the underlying asset returns, the level of the IV surface and its slope all significantly influence hedging performance, irrespective of the risk metric employed within the objective function.





































The rest of the paper is organized as follows. Section <ref> frames the hedging problem in terms of a deep reinforcement learning framework. Section <ref> provides the components of the JIVR model. Section <ref> presents the numerical results, assessments, and global feature importance analysis.[The Python code to replicate the numerical experiments from this paper can be found at the following link: https://github.com/cpmendoza/DeepHedging_JIVR.githttps://github.com/cpmendoza/DeepHedging_JIVR.git.] Section <ref> concludes.














§ THE HEDGING PROBLEM


The mathematical formulation of the hedging problem considered herein, along with the solution approach based on deep reinforcement learning, are described in this section. 



 §.§ The hedging optimization problem


This paper introduces a dynamic hedging strategy for European-style options that leverages insights provided by the implied volatility surface. The approach aims to minimize some risk measure applied to the terminal hedging error.

The European option payoff Ψ(S_T) depends on the price of the underlying asset at maturity, denoted as T trading days. The hedging strategy involves managing a self-financing portfolio composed of both the underlying asset and a risk-free asset, with daily rebalancing.  The strategy is represented by the predictable process 
{ (ϕ_t, δ_t) }_t=1^T, where ϕ_t is the cash held at time t-1 and carried forward to the next period, and δ_t denotes the number of shares of the risky asset S held during the interval (t-1, t].
The time-t portfolio value is

    V_t^δ=ϕ_t^r_tΔ+δ_t S_t^q_tΔ

where r_t is the time-t annualized continuously compounded risk-free rate, q_t is the annualized underlying asset dividend yield, both on the interval (t-1,t], and Δ=1/252.
To account for transaction costs the self-financing condition entails 

    ϕ_t+1 +δ_t+1S_t= V_t^δ-κ S_t|δ_t+1 - δ_t| ,

where κ is the rate of proportional transaction costs.


The optimal hedging problem is an optimization task where an agent seeks to minimize the risk exposure associated with a short position in the option. More precisely, it is a sequential decision problem where the agent looks for the best sequence of actions δ= {δ_t}^T_t=1 that minimizes a penalty function ρ applied to the hedging error at maturity for a short position, defined as 

    ξ_T^δ = Ψ(S_T) - V_T^δ.

Note that ξ_T^δ is a loss variable, with profits being represented by -ξ_T^δ.
The problem is

    δ^*=min_δ ρ( ξ_T^δ) ,

where 
ρ is a risk measure, 
acting as the penalty function. Each time-t action (ϕ_t+1,δ_t+1) is of feedback-type, with such decision being a function of current available information on the market: δ_t+1= δ̃(X_t) for some function δ̃ with state variables vector X_t. 
Section <ref> further describes these state variables that include the underlying asset current value as well as some information about the implied volatility surface, among others. Due to Equation (<ref>), ϕ_t+1 is fully characterized when δ_t+1 is specified, and as such the time-t action to be chosen is simply δ_t+1.

In this paper we consider three penalty functions that are very popular in the literature:

    
  * Mean Square Error (MSE): ρ(ξ_T^δ)= 𝔼[(ξ_T^δ)^2].
    
  * Semi Mean-Square Error (SMSE): ρ(ξ_T^δ)= 𝔼[(ξ_T^δ)^21_{ξ_T^δ≥ 0}].
    
  * Conditional Value-at-Risk (CVaR_α): ρ(ξ_T^δ)=𝔼[ξ_T^δ|ξ_T^δ≥_α(ξ_T^δ)], where α∈ (0,1) and _α(ξ_T^δ) is the Value-at-Risk defined as _α(ξ_T^δ) = min_c{ c:ℙ(ξ_T^δ≤ c)≥α} and α=95% or 99% in this work.

The MSE was first proposed in the seminal variance-optimal hedging framework of <cit.>, which was later extended to the multivariate case by <cit.>. The SMSE is a particular case of the asymetric polynomial penalty considered for instance subsequently in <cit.>, <cit.> and <cit.>. It provides the advantage over the MSE to avoid penalizing hedging gains. Lastly, we consider CVaR as a standard metric for measuring potential catastrophic tail events, frequently mandated by financial regulators for use in financial institutions. This metric has also been explored in global hedging contexts in <cit.>, <cit.>, <cit.>, <cit.> or <cit.>, among others. 

Historically, an alternative method for addressing the problem described in Equation (<ref>) employs backward recursion within dynamic programming frameworks. However, this approach is hindered by the curse of dimensionality, limiting its practical applicability. 
To overcome these limitations, we tackle Problem (<ref>) using reinforcement learning.



 §.§ Reinforcement learning and deep hedging


The optimal hedging problem (<ref>) can be formulated as a reinforcement learning (RL) problem because it is a feedback sequential decision-making task. In this framework, an agent learns a policy (the investment strategy δ) which dictates trading actions to be applied as a function of state variables to minimize the hedging objective highlighted in Equation (<ref>). More precisely, the problem consists in learning the mapping δ̃.
 











Consistently with <cit.>, the problem established in (<ref>) is solved by direct estimation of the policy through a deep policy gradient approach that relies on the estimation of the mapping δ̃ by an Artificial Neural Network (ANN). Denoting by δ̃_θ the policy obtained when δ̃ is estimated by an ANN with parameters θ, the objective function we need to minimize is thus

    𝒪(θ) = ρ( ξ_T^δ̃_θ)




  §.§.§ Neural network architecture



We propose a nonstandard architecture which consists in a Recurrent Neural Network with a Feedforward Connection (RNN-FNN) that combines the traditional Long Short-Term Memory Network (LSTM) and the Feedforward Neural Network (FFNN) architectures. The inclusion of LSTM layers mitigates problems related to vanishing gradients.[Vanishing gradients arise when the gradients of the penalty function become extremely small, leading to slow or halted training (details can be found in <cit.>).] 
The RNN-FNN network is defined as a composition of LSTM cells { C_l}_l=1^L_1 and FFNN layers {ℒ_j}_j=1^L_2 under the following functional representation:

    δ̃_θ(X_t)=(ℒ_J∘ℒ_L_2∘ℒ_L_2-1∘ ... ∘ℒ_1_∘C_L_1∘ C_L_1-1...∘ C_1_)(X_t).


 The LSTM cell C_l is a mapping that transforms a vector Z^(C, l-1)_t of dimension d^(C,  l-1) into a vector Z^(C,  l)_t of dimension d^(C,  l) based on the following equations, considering Z^(C,  0)_t=X_t:

    i^(l)   = (W_i^(l)Z_t^(C,  l-1)+b_i^(l))
    
        o^(l)   = (W_o^(l)Z_t^(C,  l-1)+b_o^(l))
    
        c^(l)   = i^(l)⊙(W_c^(l)Z_t^(C,  l-1)+b_c^(l))
    
        Z_t^(C,  l)   = o_t^(l)⊙(c^(l))

where (·) and (·) are respectively the sigmoid and hyperbolic tangent functions applied element-wise and ⊙ is the Hadamard product. Layer ℒ_j represents a FFNN layer that maps the input vector Z^(ℒ,  j-1)_t of dimension d^(ℒ,  j-1) into a vector Z^(ℒ,  j)_t of dimension d^(ℒ,  j) by applying a linear transformation T_ℒ_j(Z^(ℒ,  j-1)_t)=W_ℒ_jZ^(ℒ,  j-1)_t + b_ℒ_j and, subsequently, an element-wise non-linear activation function g_ℒ_j, i.e., ℒ_j(Z^(ℒ,  j-1)_t)=(g_ℒ_j∘ T_ℒ_j)(Z^(ℒ,  j-1)_t) for j∈{1,...,L_2,J}, considering Z^(ℒ,  0)_t=Z^(C,  L_1)_t.

The trainable parameters θ of the RNN-FNN network are listed below:

    
  * If L_1≥ l≥ 1: W_i^(l), W_o^(l), W_c^(l)∈ℝ^d^(C,  l)× d^(C,  l-1) and b_i^(l), b_o^(l), b_c^(l)∈ℝ^d^(C,  l)× 1 with d^(C,  0) defined as the original dimension of the network input.
    
  * If L_2≥ j≥ 1: W_ℒ_j∈ℝ^d^(ℒ,  j)× d^(ℒ,  j-1) and b_ℒ_j∈ℝ^d^(ℒ,  j) with d^(ℒ,  0)=d^(C,  L_1).
    
  * If j=J: W_ℒ_J∈ℝ^1× d^(ℒ,  L_2) and b_ℒ_J∈ℝ.


The selected hyperparameter values for our experiments are detailed in Section <ref>.



    








  §.§.§ Neural network optimization


The RNN-FNN network is optimized with the Mini-batch Stochastic Gradient Descent method (MSGD). This training procedure relies on updating iteratively all the trainable parameters of the network  based on the recursive equation

    θ_j+1 = θ_j-η_j∇_θ𝒪̂(θ_j),

where θ_j is the set of parameters obtained after iteration j, η_j is the learning rate (step size) which determines the magnitude of change in parameters on each time step,[This parameter can be either deterministic or adaptive, i.e., it may be adjusted during the training period. For more details on this, please refer to <cit.>.] 
∇_θ is the gradient operator with respect to θ and 𝒪̂ is the Monte Carlo estimate of the objective function (<ref>) computed on a mini-batch. Additional details are provided in Appendix <ref>. 








§ JOINT MARKET DYNAMICS


This section describes the market dynamics that represent the joint evolution of the S&P 500 index price and its associated Implied Volatility (IV) surface. Such model is used to construct the state space of the hedging problem.





 §.§ Daily implied volatility surface representation


On any given day, the cross-section of option prices on the S&P 500 index is captured through the IV surface model introduced by <cit.> which characterizes the entire surface parsimoniously with a linear combination of five factors. 
More precisely, the time-t IV of an option with time-to-maturity τ_t=T-t/252 years and moneyness M_t=1/√(τ_t)S_te^(r_t-q_t)τ_t/K, where K is the strike price, is

    σ (M_t,τ_t,β_t) =  ∑_i=1^5β_t,if_i(M_t,τ_t)


where β_t=(β_t,1,β_t,2,β_t,3,β_t,4,β_t,5) stands for the IV factor coefficients at time t and the functions {f_i}_i=1^5 represent the long-term at-the-money (ATM) level, the time-to-maturity slope, the moneyness slope, the smile attenuation and the smirk, respectively (see Appendix <ref> for their specification).

Following the same data processing and estimation procedure outlined in the aforementioned study, we extract the daily time series of the IV factor coefficients spanning from January 4, 1996, to December 31, 2020.[The sample comes from OptionMetrics database, to which the conventional data exclusion filters are applied.]



 §.§ Joint Implied Volatility and Return (JIVR)


The JIVR model proposed by <cit.> leverages the IV representation (<ref>) and provides explicit joint dynamics for the IV surface and the S&P 500 index price.

The first building block of the JIVR model represents the daily underlying asset excess log-return, R_t+1=(S_t+1/S_t)-(r_t-q_t)Δ. It integrates an NGARCH(1,1) process with Normal Inverse Gaussian (NIG) innovations to capture volatility clusters and reproduce fat tailed asymmetric returns, while borrowing information from the volatility surface to anchor the evolution of the conditional variance process h_t,R. 
The second building block includes a multivariate heteroskedastic autoregressive processes with non-Gaussian innovations for all implied volatility factor coefficients. 

The multivariate time series representation of the JIVR model is presented in detail in .

The full model is characterized by the current market conditions (S_t,{β_t,i}^5_i=1, β_t-1,2, h_t,R, { h_t,i}^5_i=1), which are respectively the underlying S&P 500 index price, IV factor coefficients, and conditional variances for the S&P 500 return and for such coefficients.

Following <cit.>, the maximum likelihood estimation is applied to S&P 500 excess returns alongside the time series estimates of surface coefficients {β_t,i}^5_i=1. As a byproduct, we obtain daily estimates of conditional variance series h_t,R and { h_t,i}^5_i=1 for the time period extending between January 4, 1996 and December 31, 2020.  








§ NUMERICAL STUDY


In this section, simulation and backtesting experiments are conducted to evaluate the performance of the proposed hedging approach.



 §.§ Stochastic market generator




  §.§.§ Market simulator


The JIVR model is used in subsequent simulation experiments to generate paths of the variables pertaining to market dynamics (S_t,{β_t,i}^5_i=1, h_t,R, { h_t,i}^5_i=1), which drive the hedging decisions. The initial conditions ({β_0,i}^5_i=1, h_0,R, { h_0,i}^5_i=1) are randomly chosen among values prevailing in our data sample extending between January 4, 1996 and December 31, 2020. This constitutes a wide variety of states of the economy. Following the determination of initial values, the simulation progresses through all time steps in two distinct phases. First, a sequence of NIG innovations {ϵ_t}_t=1^T is simulated using the Monte Carlo method considering the dependence structure of contemporaneous innovations. Secondly, equations of Section <ref> are used to obtain values for ({ R_t, β_t,i}^5_i=1, h_t,R, { h_t,i}^5_i=1) for t=1,…,T based on the simulated innovations. 








  §.§.§ Market parameters for numerical experiments




 

The initial underlying asset value is normalized to S_0=100 for simplicity. The options being hedged are assumed to be European call options (Ψ(S_T)=(S_T-K,0)) with maturities T∈{21,63,126 } days for short-term, medium-term and long-term maturities, and strikes K∈{90,100,110 } for in-the-money (ITM), at-the-money (ATM) and out-of-the-money (OTM) options, respectively.[Unreported numerical results of European put options exhibit a similar pattern due to the Put-Call parity formula.]

Various levels of proportional transaction cost are considered, namely κ∈{ 0%,0.05%,0.5%,1%}. The initial value of the hedging strategy, V_0^δ, is set to the price of the option being hedged, which is provided by the prevailing implied volatility surface.

Other parameters of the model are specified in Section <ref>.






 §.§ Benchmarks


We benchmark RL hedging strategies against the performance of three standard models: (i) the practitioner's Black-Scholes delta hedging (BS) which applies the current implied volatility into the Black-Scholes formula to obtain the hedged option's delta; (ii)  the <cit.> delta (BSL),
which modifies the BS delta to reflect proportional transaction costs; and (iii) the Smile-implied delta (SI) introduced by <cit.> and whose closed-form expression for the IV model (<ref>) is provided by <cit.>. The explicit formulas for these three benchmarks are detailed in Appendix <ref>.



 §.§ Neural network settings




  §.§.§ Neural network architecture


We employ the RNN-FNN architecture from Section <ref> with two LSTM cells (L_1=2) of width 56  (d_i=56 for i=1,2) and two FFNN-hidden layers (L_2=2) of width 56 with ReLU activation function (i.e., g_ℒ_i(x)=(0,x) for i=1,2). In the context of the output FFNN layer ℒ_J, which maps the output of hidden layers Z^(ℒ,  L_2)_t into the position in the underlying asset δ_t, the standard deep hedging framework typically employs a linear activation function. However, this activation function tends to induce RL agents to adopt doubling strategies, increasing the position in the underlying asset several orders of magnitude over the current position on each period after any loss until the cumulative loss amount is fully recovered. Such strategies are definitely undesirable from the perspective of sound risk management. Therefore, we opt to introduce a leverage constraint through the output layer activation function.

This leverage constraint is a dynamic upper bound on the activation function, denoted as g_ℒ_J(Z,t)=(Z,B_t(Z)), with Z=T_ℒ_J(Z^(ℒ,  L_2)_t) representing the typical deep hedging underlying asset position, and B_t(Z) the highest position in the index

that can be held in the portfolio. 
Such an upper bound avoids excessive leverage and limits the borrowing capacity, i.e., the cash held satisfies ϕ_t+1(X_t)≥ -B for all X_t and t, and B>0.[This type of leverage condition has been previously employed in the literature. For instance, the seminal paper of  assumes a restricted borrowing capacity to maintain a positive wealth throughout the entire hedging period.] The latter, in conjunction with the self-financing constraint (<ref>), establishes the dynamic upper bound as

    B_t(Z) = {[               V_0+B/S_0                 if  t=0; V_t+B+κ S_tδ_t/S_t(1+κ)     if  t>0  and  Z≥δ_t; V_t+B-κ S_tδ_t/S_t(1-κ)   if  t>0  and  Z< δ_t.;                         ].


Appendix <ref> provides numerical experiments that illustrate the necessity of including a leverage constraint in the network architecture to alleviate the aforementioned issues.



Our numerical experiments demonstrate the superiority of the RNN-FNN architecture over standalone LSTM and FFNN networks, considering a leverage constraint of B=100 for all agents. These results are presented in Appendix <ref>. Additionally, agents are trained as described in Section <ref> on a training set of 400,000 independent simulated paths with mini-batch size of 1,000 and a learning rate of 0.0005 that is progressively adapted by the ADAM optimization algorithm.[ADAM is a dynamic learning rate algorithm engineered to accelerate training speeds in deep neural networks and achieve rapid convergence, details can be seen in <cit.>.] In addition, we include a regularization method called dropout with parameter p=0.5 to reduce the likelihood of overfitting and enhance the model performance on unseen data.[Full details of the dropout regularization method can be seen in <cit.>. The selection of p=0.5 stems from our numerical experiments detailed in Appendix <ref>, indicating that this value consistently outperforms others across all penalty functions.] The training procedure is implemented in Python, using Tensorflow and considering the <cit.> random initialization of the initial parameters of the neural network. Numerical results are obtained from a test set of 100,000 independent paths. 



  §.§.§ State space selection


The nature of the hedging problem, combined with the JIVR model, establishes a decision framework where the optimal decision is entirely defined by the state variables at time t, denoted as X_t=(V_t^δ,δ_t,τ_t,S_t,{β_t,i}^5_i=1, h_t,R, { h_t,i}^5_i=1). 

However, implementing numerical methods within the deep policy gradient approach across the entire state space may lead to overfitting. In fact, other studies utilizing the RL framework have reported optimal results with a reduced state space. For instance, <cit.> and <cit.> employ the deep hedging algorithm without including the portfolio value V_t^δ in the state space configuration. 
 Similarly, <cit.> employs the deep hedging framework in a frictionless market without including δ_t in the state space.

We explore three different state space configurations. The first configuration aims to replicate the full state space, denoted by 

    (V_t^δ,δ_t,τ_t,S_t,{β_t,i}^5_i=1, h_t,R, { h_t,i}^5_i=1),

typically considered in a dynamic optimization task with the portfolio value as a state variable. The second state space does not consider the portfolio value and also omits the variance of the IV coefficients { h_t,i}^5_i=1 under the intuition that these coefficients have a second-order effect, with the IV surface coefficients {β_t,i}^5_i=1 capturing most of the variability. Hence, the reduced state space is defined as

    (δ_t,τ_t,S_t,{β_t,i}^5_i=1, h_t,R).

The current position δ_t is required for the RL agent to learn about the transaction cost associated with the next rebalancing.  This state space component is no longer useful when κ=0.  For that reason, we remove this component in absence of transaction cost. In such case, the

reduced state space is denoted by 

    (τ_t,S_t,{β_t,i}^5_i=1, h_t,R).


Our numerical experiments consider the three aforementioned state spaces  to hedge a European ATM call option with maturity N=63 days, while taking into account four different transaction cost rates κ∈{ 0%, 0.05%, 0.5%, 1%}. The impact of the option moneyness and maturity on hedging performance are studied in later subsections. Additionally, our numerical experiments involve four RL agents: RL-CVaR_95%, RL-CVaR_99%, RL-MSE, and RL-SMSE. These agents aim to minimize different penalty functions.




<ref> illustrates the estimated value of the four penalty functions in proportion to that obtained with the BSL delta, ρ(ξ_T^δ)/ρ(ξ_T^BSL), across different transaction cost rates for all agents. These numerical results reveal that RL agents outperform BSL delta hedging for all considered state space configurations, as the relative values of the penalty functions are substantially lower than the reference line value of 1.

Moreover, these results confirm that including the portfolio value in the full state space is unnecessary for our approach, as metrics under the reduced state space (green curve) show lower values than those achieved by agents using the full state space (red curve). This is consistent with the work of <cit.>, <cit.>, <cit.>, and <cit.>, where RL techniques are applied in the hedging context without considering the portfolio value as a state variable.

In the absence of transaction costs, <ref> confirm that the inclusion of δ_t in the reduced state space is unnecessary, as the incremental performance with and without it is negligible. Conversely, in the presence of transaction costs, agents trained under the reduced state space (δ_t,τ_t,S_t,{β_t,i}^5_i=1, h_t,R) demonstrate superior performance (green curve), and the outperformance with respect to benchmarks becomes more pronounced as the transaction cost rate increases. This observation underscores the importance of considering previous positions in the underlying when optimizing rebalancing actions, as the former is indicative of transaction costs to be paid for the various possible actions.

The superiority of RL agents trained on the reduced state space can be explained by the bias-variance dilemma, where the informational content provided by some of the variables (reduction in bias) is insufficient to compensate for additional complexity and instability (variance) they cause during training. 
This is seemingly why removing the IV parameters' variances { h_t,i}^5_i=1 from the state space increases the performance in this experiment. Our numerical results indicate that the network architecture considered for the agents acting in our proposed high-dimensional environment does not require the full state space. In fact, the performance of the agents deteriorates for some transaction cost levels when the volatilities of the IV coefficients are included, as shown in Figure <ref>, and the time required to converge to optimal solutions during training increases by an average of 190% compared to RL agents under the reduced state space across all experiments. 


In all subsequent experiments, we consider the reduced state space (δ_t,τ_t,S_t,{β_t,i}^5_i=1, h_t,R) in the presence of transaction costs and the state space (τ_t,S_t,{β_t,i}^5_i=1, h_t,R) in the absence of transaction costs to enhance parsimony. 



 §.§ Benchmarking of hedging strategies


In this section, we compare the performance of RL agents with classic hedging approaches. Aggregated results are presented in Section <ref>, while the segmentation among various states of the economy is shown in Section <ref>. The following subsections break down the performance of the RL agents with respect to moneyness and maturity of the hedged option, or to transaction cost levels, demonstrating the robustness of the RL approach. Finally, Section <ref> outlines a comparison of performance over historical paths spanning from January 5, 1996, to December 31, 2020.



  §.§.§ Benchmarking over randomized economic conditions


We begin by comparing the hedging performance of the benchmarks and RL agents trained under the four penalty functions: CVaR_95%, CVaR_99%, MSE, and SMSE. This comparison is performed in terms of estimated values of all penalty functions and the average Profit and Loss, = 𝔼[ -ξ_T^δ], across all paths in a test set. Additionally, we employ the CVaR deviation measure, defined as _α(ξ_T^δ - E[ξ_T^δ]), as a deviation metric for the hedging error in the test set. In such test set, initial economic conditions (the initial value of state variables in the path) are sampled randomly among historical values from our sample. It therefore reflects aggregate performance across various economic conditions.
Our analysis focuses on hedging a European ATM call option with a maturity of N=63, under the assumption of no transaction costs, which is κ=0.






<ref> presents hedging performance metrics attained by both benchmarks and RL agents. 
Every RL agent achieves the lowest value for the corresponding metric which they used as objective function during training, which is expected.
Furthermore, the numerical results indicate that RL agents provide hedging strategies that are much less risky than benchmarks, as evidenced by metrics such as CVaR_95%, CVaR_99%, and SMSE. In particular, when computing the variation rate between the estimated value of each of these metrics obtained by RL agents relative to the value obtained by BS delta, we observe an average reduction rate across RL agents of 24%, 31%, and 38%, respectively. Similarly, when comparing with SI delta strategies, we observe average reductions across RL agents of 45%, 46%, and 67%, respectively. Moreover, in terms of the MSE metric, the MSE agent demonstrates significant superiority over benchmarks, reducing 26% and 63% compared to BS delta and SI delta strategies, respectively. 

Regarding average Profit and Loss (P&L), SI delta yields higher profitability; however, it also entails increased risk. In contrast, RL agents trained using CVaR_95%, CVaR_99%, and SMSE as penalty functions obtain lower values compared to SI delta strategies, but they reach a more favorable trade-off between profitability and risk management as shown by the ratio , which yields greater values for these agents. Additionally, RL agents achieve lower risk than both the BS and SI delta hedging strategies, which highlights enhanced risk management. 


 



  §.§.§ Impact of the state of the economy on performance






In this section, we analyze the performance metrics of hedging strategies across clusters of paths representing different economic states, as detailed in <ref>. A path is assigned to a cluster if its state variable initial values are drawn from the subset of dates corresponding to that cluster. This approach helps isolate the impact of economic conditions on performance and assess the robustness of RL hedging strategies.
Again, we hedge an ATM European call option with a maturity of N=63 days under the assumption of no transaction cost.



Our numerical results, illustrated in <ref>, indicate that both benchmarks and RL agents are sensitive to the environment, showcasing better performance during regular periods compared to financial crises in terms of risk, as shown by the two right plots showing the CVaR_95% and SMSE metrics. Additionally, both approaches tend to offer more profitable strategies during crisis periods and periods of high volatility (for instance, even though Regular period 2 is not labeled as a crisis periods, it is characterized by economic recovery with significant market fluctuations). This trend can be explained by the fact that higher volatility during these periods leads to a higher risk premium, increasing the initial portfolio value and thus the final P&L, as highlighted by the Avg P&L metric (see the top left panel of <ref>). Regarding the MSE metric, the results align with the Avg P&L metric, with both approaches showing superior performance during regular periods (except for Regular period 2 which exhibited high volatility), as higher profits are detrimental for the MSE metric penalizing upside risk.




The results that are segregated by time period align with the aggregated findings. Indeed, for any of the CVaR_95%, MSE, and SMSE risk metrics, the RL agent which was trained with such risk metric as its objective function consistently exhibits lower values than benchmarks across all periods. These differences are particularly notable during crisis periods, where BS delta and SI delta tend to be riskier, while RL agents demonstrate greater stability. The higher values in the MSE statistic of agents trained under CVaR and SMSE can be attributed to the fact that these objective functions only penalize hedging losses, allowing agents to seek positive returns on average. This is consistent with the Avg P&L panel, where CVaR and SMSE agents tend to display more profitable strategies than the MSE agent.

These results demonstrate the robustness of the RL approach across various environments, exhibiting greater stability in scenarios with extreme behavior and improving hedging performance in terms of the penalty functions considered in the hedging problem. Additionally, they support our initial findings, indicating that RL agents do not significantly sacrifice profitability even during extreme market conditions, a phenomenon that could be attributed to higher initial option prices during these periods, as shown in <ref>.



  §.§.§ Impact of moneyness level on performance


We investigate the robustness of our approach with respect to the moneyness level of the option being hedged by including OTM and ITM call options with a maturity of N=63 days in our analysis. Our findings, depicted in <ref>, demonstrate that RL consistently outperforms benchmarks with respect to the objective function considered during the training process, regardless of the option moneyness.



In line with our previous experiments, we observe that SI delta tends to offer more profitable strategies compared to benchmarks, while BS delta demonstrates a similar level of profitability for ATM and ITM options. However, this trend does not hold for OTM options, where RL agents not only exhibit better risk management but also yield more profitable strategies on average (refer to CVaR and SMSE agents for OTM options in <ref>). The discrepancy can be attributed to the fact that RL agents trained under CVaR and SMSE do not track the option value nor penalize gains at maturity, allowing agents to profit from OTM paths at maturity. In contrast, BS and BSL are option tracking methodologies that aim to replicate the option value at maturity, thereby reducing potential gains for OTM options at maturity. The latter is consistent with the RL agent trained under the MSE, which achieves the minimum MSE value for OTM options and the lowest Avg P&L, as it penalizes both gains and losses.




  §.§.§ Impact of option maturity on performance


As a third test to assess the robustness of our approach, we compare the performance of RL agents against benchmarks when hedging ATM call options with different maturities: 21, 63, and 126 days, in absence of transaction costs. Our results, illustrated in <ref>, show that the average profitability of hedging strategies increases with option maturity (see Avg P&L in the top left panel). However, the risk also increases with maturity for all hedging strategies (refer to CVaR_95% and SMSE depicted in the panels positioned at the top and bottom right corners), as expected.



These results also demonstrate consistent behavior across all maturities for all hedging metrics, displaying a uniform distribution among the different strategies across all maturity levels. In line with our previous experiments, RL agents consistently outperform benchmarks with respect to asymmetric penalty functions for all maturities, regardless of the penalty function used during training. Furthermore, in the case of the MSE metric (bottom left panel in <ref>), the RL agent trained under that penalty function displays the lowest value across all maturities.



  §.§.§ Impact of transaction costs on performance


We now investigate scenarios where hedgers encounter transaction costs, which is meant to reflect more realistic hedging scenarios. Specifically, we analyze the hedging effectiveness of RL agents in comparison to benchmarks, including the BSL delta hedging. Our examination focuses on hedging an ATM European call option with a maturity of N=63 days across varying levels of transaction costs assumed to be 0.05%, 0.5% and 1%. 

<ref> illustrates the performance of benchmarks and RL agents, with the four panels depicting the Avg P&L, CVaR_95%, MSE, and SMSE metrics. In general, as expected, the performance of each strategy tends to deteriorate as transaction costs increase. However, performance drops are more pronounced for benchmarks than for RL agents. As anticipated, RL agents consistently outperform benchmarks in terms of downside risk (observe CVaR_95% and SMSE metrics depicted in the panels positioned at the top and bottom right corners of <ref>) 
across all transaction cost levels, regardless of the penalty function used during training. Furthermore, while BSL delta hedging adjusts the strategies in terms of the transaction cost level and performs better than the other benchmarks, RL strategies display superior metrics, highlighting their better capacity to adapt to different transaction cost levels.





In contrast to our previous results where SI delta provides more profitability, the results regarding the Avg P&L metric (see top left panel in <ref>) indicate that this profitability is influenced by the inclusion of transaction costs, making the SI delta hedging less profitable strategy as transaction costs increase, whereas RL agents tend to exhibit much lower average losses due to their adaptability. In terms of MSE, benchmarks exhibit a more sensitive behavior than the RL agents regarding the increase in transaction costs, especially when the transaction costs are set at 0.5% and 1% (see the bottom-left panel of Figure <ref>), irrespective of the penalty function considered during training.

<ref> displays descriptive statistics about hedging costs, defined as the sum of discounted transaction costs over a path,

    ∑_t=0^T-1e^-rΔ tκ S_t|δ_t+1 - δ_t|,

 across different transaction cost levels. In the cases of BS delta and SI delta strategies, which overlook transaction costs, we observe a higher cost of hedging due to significant variations in the position of the underlying asset during rebalancing across different time steps, thereby leading to lower profitability as transaction costs increase. This trend is evident in the case of SI delta which incurs the highest average hedging cost, followed by BS delta and BSL delta, when the transaction cost rate is set to 0.5% and 1%. These results consistently align with the performance differences observed in <ref>,  where benchmarks are associated with the highest losses.




Furthermore, although none of the agents directly aim to minimize transaction costs, we observe that RL agents with asymmetric penalty functions tend to offer hedging strategies with lower costs, as shown by the highlighted values in <ref>. This underscores the robustness of RL agents in addressing transaction costs, as they minimize the penalty function and thereby indirectly reduce the cost of hedging. Conversely, the RL agent trained under the MSE criterion tends to exhibit higher average costs compared to the other agents. This can be attributed to the MSE agent penalizing both negative and positive returns, creating additional turnover in some situations to avoid large profits. 



 §.§ Global importance of IV surface features


We now investigate to what extent the risk factors characterizing the IV surface described in Equation (<ref>) contribute to total performance of our model. 


We employ the Shapley Additive Global Importance (SAGE) methodology, introduced by <cit.>, to evaluate their impact. 
Specifically, let ρ(ξ_T^δ̃_θ(X)) be the risk measure when the model is trained with the state space X. The amount of risk reduction achieved by adding the state variables ({β_t,i}^5_i=1, h_t,R) to a baseline model δ̃_θ with state variables (τ_t, S_t) is

    ρ(ξ_T^δ̃_θ(τ_t, S_t)) 
               - ρ(ξ_T^δ̃_θ(τ_t, S_t, {β_t,i}^5_i=1, h_t,R))= ∑_j ∈{{β_t,i}^5_i=1, h_t,R}𝒞_j

where 𝒞_j, the contribution of feature j to the total risk reduction, is

    𝒞_j= ∑_𝒳⊆{τ_t, S_t, {β_t,i}^5_i=1, h_t,R}∖{j}|X|! (5 - |X|)!/6![ ρ(ξ_T^δ̃_θ(X)) - ρ(ξ_T^δ̃_θ(X, j)) ].

Exact contributions 𝒞_j cannot be negative, although their estimates sometimes are. Negative values may occur for the contributions we present that are evaluated out-of-sample.







    



Our numerical experiments involve training RL agents using four penalty functions: CVaR_95%, CVaR_99%, MSE, and SMSE. We analyze the global importance of these state variables across different moneyness and maturities to comprehend their impact on the performance of the RL agents. The global importance of state variables is normalized by the risk reduction achieved by the respective RL agent to present contributions in the same order of magnitude: the relative global importance is


    𝒞_j/ρ(ξ_T^δ̃_θ(τ_t, S_t)) 
        - ρ(ξ_T^δ̃_θ(τ_t, S_t, {β_t,i}^5_i=1, h_t,R)),   j ∈{h_t,R}∪{β_t,i}^5_i=1.





The relative global importance of the IV characteristic β_t,i and the return conditional variance h_t,R to the risk reduction depends on the moneyness and the time-to-maturity of the option to be hedged, as well as on the choice of risk measure.

<ref> studies the case of 63-day-to-maturity call options. It illustrates the relative contribution of each state variable across moneyness levels. Overall, the conditional variance of the underlying asset returns, the long-term ATM level β_1 and the time-to-maturity slope β_2 of the IV surface play a major role, no matter what risk measure or moneyness is considered. This underscores that RL agents utilize both the historical variance process and market expectations of future volatility to adjust positions in the underlying asset. The moneyness slope, the smile attenuation and the smirk have a second order effect. 









<ref> presents contributions to hedging risk reduction for three option maturities, namely 21, 63 or 126 business days. It underscores the persistent significance of the conditional variance process {h_t,R} across the various option maturities. Its contribution is even more important for short-term maturities. This reflects the fact that h_t,R has a direct impact on immediate market shocks.


Both β_4 and β_5 are additional contributors mostly for short-term options when CVaR risk measures are considered. Intuitively, high values of the smile attenuation (β_4) and the smirk (β_5) indicate a steep slope and a strong smirk for the short term smile which, in turn, induces a high exposure to tail risk for short term options.











 §.§ Backtesting
 

In this section, we benchmark our approach using a historical path of the JIVR model spanning from January 5, 1996, to December 31, 2020, to assess the effectiveness of RL agents. Unlike our above experiments, which use simulated paths to test hedging strategies, this experiment examines their performance based on the historical series {R_t,β_t}. Specifically, we evaluate the hedging performance considering a new European ATM call option with a maturity of 63 days every 21 business days along this historical path. The option prices, serving as initial hedging portfolio values, are determined with the prevailing IV surface on the day the hedge is initiated.

To assess the robustness of the model under more general market conditions, we conduct a comparison of cumulative P&Ls, which are computed as the cumulative sums of the P&L achieved by each strategy at the maturity of each option during the analyzed period. As illustrated in <ref>, which depicts the evolution of the cumulative P&L across two panels, each representing a different transaction cost level, the gap in cumulative P&L between RL agents and benchmarks significantly widens as transaction cost rates increase. This, again, highlights the adaptability of the RL approach to various market conditions. 



In contrast to the findings presented in Section <ref>, where SI delta yielded the highest profitability among all strategies, the RL approach achieved the highest profitability in historical paths. Specifically, agents trained under the MSE and SMSE emerged as the most profitable ones.

Moreover, RL agents consistently outperform benchmarks across most of the transaction cost levels. However, when the transaction cost is small (left panel), Benchmarks exhibit a better cumulative P&L compared to the RL agent trained under the CVaR_99% penalty function at the end of the period, December 31, 2020. Nevertheless, Benchmarks also demonstrate more variability and larger losses, such as those observed between 2008 and 2012. In contrast, RL agents show a smaller decrease in cumulative P&L, indicating more resilience during crisis periods. In general, the observed performance of the RL agents under historical data is consistent with our findings under simulated data, indicating the robustness of our approach.









§ CONCLUSION


This study introduces a novel deep hedging framework that integrates forward-looking volatility information through a functional representation of the IV surface, combined with backward-looking conventional features. Our implementation employs deep policy gradient methods and utilizes a unique neural network architecture consisting of LSTM cells and FFNN layers to enhance training efficiency. Additionally, the architecture incorporates a budget constraint mitigating the incentive to gamble and enabling the agent to learn hedging strategies that are more effective for risk management.

Our approach consistently outperforms traditional benchmarks both in the absence and presence of transaction costs. The stability of our approach is assessed across various economic states using simulated data, demonstrating greater robustness than benchmarks in extreme scenarios such as economic crises.

The global importance analysis of IV features confirms the significant enhancement of hedging performance in terms of risk reduction relative to the penalty functions. Our analysis underscores the critical importance of key factors such as the underlying asset return conditional variance process (h_R), the long-term ATM IV level (β_1) and the time-to-maturity slope (β_2). RL agents utilize both the historical variance process and market expectations of future volatility to adjust positions in the underlying asset. 








apalike



  








§ DETAILS FOR THE MSGD TRAINING APPROACH


The MSGD method estimates the penalty function 𝒪(θ), which is typically unknown, through small samples of the hedging error called batches. Let 𝔹_j={ξ_T,i^δ̃_θ_j}_i=1^N_ be the j-th batch where N_ is the batch size and ξ_T,i^δ̃_θ_j denotes the hedging error of the i-th path in the j-th batch defined as

    ξ_T,i^δ̃_θ_j = Ψ(S_T,(j-1)N_+i) - V_T,i^δ̃_θ_j     i∈{1,…,N_},   j∈{1,...N },

where S_T,(j-1)N_+i is the price of the underlying asset at time T in the ((j-1)N_+i)-th simulated path, V_T,i^δ̃_θ_j is the terminal value of the hedging strategy for that path when θ = θ_j and the simulated states are X_i. 

The penalty function estimation for the batch 𝔹 is

    C^()(θ_j,𝔹_j)   =1/N_∑_i=1^N_(ξ_T,i^δ̃_θ_j)^2 
    C^()(θ_j,𝔹_j)   =1/N_∑_i=1^N_(ξ_T,i^δ̃_θ_j)^21_{ξ_T,i^δ̃_θ_j≥ 0} 
    C^()(θ_j,𝔹_j)   = _α(𝔹_j)+1/(1-α)N_∑_i=1^N_max( ξ_T,i^δ̃_θ_j-_α(𝔹_j),0 )


where _α(𝔹_j)=ξ_T,⌈α· N_⌉^δ̃_θ_j is the estimation of the VaR obtained from the ordered sample {ξ_T,[i]^δ̃_θ_j}_i=1^N_ and ⌈·⌉ is the ceiling function. These empirical approximations are used to estimate the gradient of the penalty function required in Equation (<ref>).[In particular, the gradient of these estimations has analytical expressions for FFNN, LSTM networks and thus for RNN-FNN networks. Details about gradient of the empirical objective function are provided in <cit.>.] 

The selection of batch size plays a key role in the MSGD training approach as we empirically measure tail risk. A larger batch size provides a more accurate gradient estimate by averaging gradients computed over more simulated paths, thereby promoting stable convergence during training. However, large batch sizes may introduce certain disadvantages such as increased memory requirements, slower convergence, and potential generalization issues. We adopt the batch size used in <cit.> (N_=1,000), which achieves a good balance between accuracy and convergence.





§ JOINT IMPLIED VOLATILITY AND RETURN MODEL




 §.§ Daily implied volatility surface


The functional representation of the IV surface model introduced by <cit.> is

    σ (M_t,τ_t,β_t) =  β_t,1_f_1 +  β_t,2e^-√(τ_t /T_conv)_f_2 + β_t,3( M_t1_{ M_t≥ 0 } + e^2M_t-1/e^2M_t+11_{M_t< 0})_f_3
        + β_t,4( 1-e^-M_t^2)(τ_t/T_max)_f_4+β_t,5( 1-e^(3M_t)^3)(τ_t/T_max)1_{M_t< 0}_f_5  ,   τ_t∈ [T_min,T_max]


where T_max is set to 5 years, T_min=6/252 and T_conv=0.25. As in <cit.>,  a minimum threshold of 0.01 is applied to the volatility surface to prevent negative values.



 §.§ Joint Implied Volatility and Return


The JIVR model proposed by <cit.> has 6 components, one for the underlying asset excess returns and the other 5 for fluctuations of the IV surface coefficients. The S&P 500 excess return follows:
 
    R_t+1    = ξ_t+1-ψ(√(h_t+1,RΔ)) + √(h_t+1,RΔ)ϵ_t+1,R, 
    
         h_t+1,R    = Y_t+κ_R(h_t,R-Y_t)+a_Rh_t,R(ϵ_t,R^2-1-2γ_Rϵ_t,R), 
    
         Y_t   = (ω_R σ(0,1/12,β_t))^2,

where the equity risk premium is

    ξ_t+1=ψ(-λ√(h_t+1,RΔ))- ψ((1-λ)√(h_t+1,RΔ)) + ψ(√(h_t+1,RΔ)),

the process {ϵ_t,R}^T_t=0 is a sequence of iid standardized NIG random variables with parameters ζ_R and φ_R,[The standard NIG random variable ϵ follows the probability density function with parameters ζ and φ: 
    f(x)=B_1( √(φ^6/φ^2+ζ^2+(φ^2+ζ^2)(x+ φ^2ζ/φ^2+ζ^2)^2))/π√(1/φ^2+ζ^2+φ^2+ζ^2/φ^6(x+ φ^2ζ/φ^2+ζ^2)^2) e^( φ^4/φ^2+ζ^2+ζ(x+ φ^2ζ/φ^2+ζ^2) ),
 where B_1(·) denotes the modified Bessel function of the second kind with index 1. The common (α, β, δ, μ)-specification can be obtained by replacing β and γ (γ = √(α^2-β^2)), with ζ and φ, respectively, and imposing a null mean and unit variance to express δ and μ interms of α, β.] 
and ψ is their cumulant generating function.[For -√(ζ^2+φ^2)-ζ<z<√(ζ^2+φ^2)-ζ, the cumulant generating function is given by 

    ψ(z)=φ^2/φ^2+ζ^2( -ζ z+φ^2-φ√(φ^2+ζ^2-(φ+ζ)^2)).
]
Parameters of the excess return component of the model are thus Θ_R = (λ,κ_R,γ_R,a_R,ω_R,ζ_R,φ_R). 

The evolution of the long-term factor (β_1) is

    β_t+1,1   = α_1+∑_i=1^5θ_1,jβ_t,j+√(h_t+1,1Δ)ϵ_t+1,1, 
    
         h_t+1,1   = U_t+κ_1(h_t,1-U_t)+a_1h_t,1(ϵ_t,1^2-1-2γ_1ϵ_t,1), 
    
         U_t   = (ω_1·σ(0,1/12,β_t))^2.


For the other 4 coefficients, i∈{2,3,4,5 }, the time evolution satisfies

    β_t+1,i   = α_i+∑_j=1^5θ_i,jβ_t,j+νβ_t-1,21_{i=2}+√(h_t+1,iΔ)ϵ_t+1,i, 
    
        h_t+1,i   = σ_i^2+κ_i(h_t,i-σ_i^2)+a_ih_t,i(ϵ_t,i^2-1-2γ_iϵ_t,i),


where {ϵ_t,i}_i=1^5 are time-independent standardized NIG random variables with parameters { (ζ_i,φ_i) }_i=1^5, respectively. Parameters are {ω_1,ν, Θ_1, Θ_2, Θ_3, Θ_4, Θ_5} with Θ_i = (α_i,{θ_i,1}_i=1^5,σ_i,κ_i,a_i,γ_i,ζ_i,φ_i)}_i=1^5.

The JIVR model also imposes a dependence structure on contemporaneous innovations ϵ_t = (ϵ_t,R,ϵ_t,1,...,ϵ_t,5) through a Gaussian copula parameterized in terms of a covariance matrix Σ of dimension 6× 6. 

Estimates of all JIVR model parameters are taken from Table 5 and Table 6 of <cit.>. 
In all experiments, the annualized continuously compounded risk-free rate and dividend yield are assumed to be constant with values set to r=2.66% and q=1.77%, respectively.



§ BENCHMARKS


The three benchmarks outlined in this appendix operate under the premise that the implied volatilities follow the IV model introduced in Equation (<ref>).



 §.§ Black-Scholes model



The <cit.> pricing formula for an European call option is:

    _t=S_t·^-q_tτ_tΦ(d_t)-K· e^ -r_tτ_tΦ(d_t - σ (M_t,τ_t,β_t)√(τ_t))


where d_t=( S_t/K)+( r_t-q_t+1/2σ (M_t,τ_t,β_t)^2 )τ_t/σ (M_t,τ_t,β_t) √(τ_t), σ_t is the implied volatility of the option and Φ is the cumulative distribution function of the standard normal distribution. Moreover, the Black-Scholes delta is

    Δ_t^BS = ^-q_tτ_tΦ(d_t).




 §.§ Leland model


The Leland delta, as proposed by <cit.>, presents a variation of the option replication approach outlined in the work of <cit.>, incorporating parameters such as the transaction cost proportion κ and the rebalancing frequency λ.


    Δ_t^L= ^-q_tτ_tΦ(d̃_t),


where d̃_t=( S_t/K)+( r_t-q_t+1/2σ̃_t^2 )τ_t/σ̃_t√(τ_t) with σ̃_t^2 = σ (M_t,τ_t,β_t)^2[ 1+√(2/π)2κ/σ (M_t,τ_t,β_t)√(λ)].



 §.§ Smile-implied model


<cit.> provide the delta associated to the IV surface under model (<ref>):

    Δ_t^SI = e^-q_tτ_t( Φ(𝒹_t,1)+ϕ(𝒹_t,1)∂σ/∂ M)

with

    𝒹_t,1=M_t/σ(M_t,τ_t,β_t)+1/2σ(M_t,τ_t,β_t)√(τ_t),

and

    ∂σ/∂ M   =β_t,31_{M_t≥ 0}+β_t,3(1-  (e^2M_t-1/e^2M_t+1)^2)1_{M_t< 0}+β_t,42M_te^-M_t^2(T/T_max)
       -β_t,581M_t^2e^27M_t^3(T/T_max)1_{M_t< 0}.






§ NETWORK FINE-TUNING




 §.§ Bounded strategies - leverage constraint


This section presents a comparison of the hedging performance of RL agents trained under the full state space considering the CVaR_95% as penalty function. The comparison is conducted considering two RL agents: (i) RL-CVaR_95%-LC, agent subject to a leverage constraint  equivalent to the initial value of the underlying asset, set at B=100, and (ii) RL-CVaR_95%, agent operating without leverage constraints. This comparison evaluates the estimated values of all penalty functions and the average Profit and Loss (Avg P&L). This analysis aims to elucidate the impact of boundary conditions on RL agent behavior and performance when hedging of a European ATM call option with a maturity of N=63 days in the absence of transaction costs, i.e., κ = 0%).

<ref> outlines that the agent without a leverage constraint shows more profitable strategies; however, it also leads to very large losses. For instance, out-of-sample CVaR_99% and SMSE metrics exhibit higher values for this agent, compared to its counterpart with a leverage constraint. 



Moreover, our numerical experiments demonstrate that the agent without leverage constraints learns doubling strategies, which are incompatible with sound risk management practices. For instance, <ref> illustrates such behavior through three panels associated with the hedging process of a deep OTM path (first panel).[For the purpose of this experiment, a deep OTM path refers to the trajectory of the underlying asset that keeps the option significantly out-of-the-money.] We observe that the agent, without a leverage constraint, tends to increase its position in the underlying (third panel) when a loss in the portfolio value is observed (second panel), aiming to recover the loss over the long run with doubling strategies. Conversely, agents trained with a leverage constraint control their position in the underlying asset, resulting in the learning of different and less risky strategies. 







 §.§ Network architecture selection


In this section, we investigate the superiority of the RNN-FNN architecture compared to conventional architectures introduced in deep hedging literature, such as the FFNN and the LSTM architectures. In line with our previous experiments, we consider four penalty functions to train the agents under the full state space: CVaR_95%, CVaR_99%, MSE, and SMSE. Again, our experiment focuses on hedging a short position of a European ATM call option with a maturity of N=63 days, assuming no transaction costs.

The superiority of the RNN-FNN over the LSTM and the FFNN is evaluated based on the optimal value of each penalty function. <ref> illustrates the optimal values of each penalty function for each architecture, normalized by the estimated value obtained with the BS delta. Notably, the RNN-FNN setup considered in this paper significantly outperforms both the benchmark and other architectures. Conversely, the FFNN does not surpass the benchmark for the MSE, and the LSTM exhibits almost the same performance as the benchmark.



A second test to demonstrate the superiority of our architecture involves computing the optimal values of the penalty functions over various clusters of paths. The objective is to isolate the impact of different state of the economy on performance and assess the robustness of our architecture. <ref> displays the optimal values of each architecture normalized by the estimated values obtained by the BS delta for all penalty functions. Notably, RNN-FNN agents outperform both the benchmark and other architectures across all penalty functions, regardless of the economic conditions under which the simulations were conducted. Conversely, FFNN and LSTM agents fail to outperform the benchmark across all states of the economy when the agents are trained under the CVaR_95%, MSE and SMSE penalty functions, as shown in the top-left, bottom-left and bottom-right panels, respectively.


 The periods aim to approximate different states of the economy considering the time frames specified in <ref>.

Results obtained under various economic conditions unequivocally demonstrate the superiority of the RNN-FNN network across all performance metrics. Moreover, the RNN-FNN not only offers better performance in terms of risk management but also reduces computational costs and implementation complexity, with an average reduction in training time of 46%.







 §.§ Dropout parameter selection


The process of selecting the dropout parameter for the regularization method involved evaluating the performance of four agents across a range of potential parameter values. These agents are trained using the full state space, considering four penalty functions: CVaR_95%, CVaR_99%, MSE, and SMSE. The performance of the agents is measured in terms of the estimated values of the penalty functions after hedging a short position of a European call option with a maturity of N=63 days, with no transaction costs.


<ref> illustrates a comparison of the optimal values of penalty functions normalized by the estimated values obtained using the BS delta. It is noteworthy that the RNN-FNN network consistently outperforms the benchmark regardless of the dropout parameter value. Additionally, the selection of the dropout parameter remains consistent across all objective functions, with a dropout probability (p) of 50% yielding optimal performance. Consequently, for all our experiments, a dropout regularization parameter of p=50% is adopted. This parameter drives the likelihood of randomly dropping out a fraction of the units within a neural network during training, thereby generating varied architectures for each training iteration. As demonstrated in <cit.>, this method not only effectively mitigates overfitting but also boosts performance, which is consistent with our numerical results.



OnlineAppendix



 Paper Title: Temporal Subspace Clustering for Molecular Dynamics Data.






Both authors contributed equally to this research.

  University of Vienna
  Vienna
  Austria
anna.beer@univie.ac.at


[1]

  LMU Munich
  Munich
  Germany

heinrigs-martin@t-online.de



  University of Vienna
  Vienna
  Austria
plant@univie.ac.at




  Aarhus University
  Aarhus
  Denmark
ira@cs.au.dk



§ ABSTRACT 
We introduce MOSCITO (MOlecular Dynamics Subspace Clustering with Temporal Observance), a subspace clustering for molecular dynamics data. 
MOSCITO groups those timesteps of a molecular dynamics trajectory together into clusters in which the molecule has similar conformations.
In contrast to state-of-the-art methods, MOSCITO takes advantage of sequential relationships found in time series data. Unlike existing work, MOSCITO does not need a two-step procedure with tedious post-processing, but directly models essential properties of the data. 
Interpreting clusters as Markov states allows us to evaluate the clustering performance based on the resulting Markov state models. 
In experiments on 60 trajectories and 4 different proteins, we show that the performance of MOSCITO achieves state-of-the-art performance in a novel single-step method.
Moreover, by modeling temporal aspects, MOSCITO obtains better segmentation of trajectories, especially for small numbers of clusters.





Temporal Subspace Clustering for Molecular Dynamics Data
    Ira Assent
    August 5, 2024
========================================================





§ INTRODUCTION

Proteins play an important role in every living organism.
Analyzing the transition between different shapes of a protein, also known as folding, is important for understanding the structure and function of proteins, e.g., for drug design. 
Many diseases like Alzheimer, Parkinson, and different types of cancer are connected to misfolding of certain proteins <cit.>. Thus, understanding such folding processes may help to prevent or cure diseases. 

Molecular dynamics data contains trajectories of proteins that change their shape over time. It is usually very high-dimensional, with thousands of time steps and hundreds of atoms for a single protein. 
Analyzing molecular dynamics data with traditional clustering is typically not meaningful due to the curse of dimensionality. To mitigate issues with high dimensionality, subspace clustering methods detect clusters in lower-dimensional subspaces <cit.>.
However, existing subspace clustering methods do not cater for the characteristics of molecular dynamics data.

In this paper, we introduce MOSCITO (Molecular Dynamics Subspace Clustering with Temporal Observance), a subspace clustering algorithm that uses temporal regularization to exploit sequential relationships found in molecular dynamics trajectories that describe folding activities of single molecules. MOSCITO is the first one-step approach to handle this complex data type by identifying the best subspace projection. It finds features tailored specifically to molecular dynamics data. Leveraging methods from video analytics, we show that MOSCITO finds temporally contiguous groups of of protein conformations.  

In trajectories, data points that are near in time are usually more closely related than those that are further apart. 
MOSCITO allows users to choose different sizes of time windows and different weighting methods for temporal regularization: binary, Gaussian, exponential, or logarithmic. 
Interestingly, the resulting clusters can be interpreted as states in a Markov State Model (MSM), which we use for comparing MOSCITO to approaches like SSC, PCA + k-Means, and TICA + k-Means.
As we demonstrate in thorough empirical evaluation, MOSCITO successfully clusters molecular dynamics data in a model that directly captures the essential features and temporal relationships.
Our main contributions are as follows:

    
  * We introduce MOSCITO, a subspace clustering method for molecular dynamics data that makes use of temporal relations between time steps that are close
    
  * MOSCITO finds temporally contiguous groups of protein conformations that are similar 
    
  * In contrast to currently used methods that require several steps, the clusters found by MOSCITO can directly be used as states of a Markov State Model describing the trajectory, yielding state-of-the-art quality with a one-in-all holistic approach



§ CLUSTERING TRAJECTORIES OF MOLECULAR DYNAMICS DATA
 

In this paper, we study the trajectories of a single molecule over time. A trajectory is given by the three-dimensional Cartesian coordinates of all atoms of the molecule for a number of time steps with fixed intervals.
Where most traditional clustering methods target numerical data in ℝ^d, clustering MD data brings some unique challenges as we outline in the following. 

Due to the size of MD data, automated methods for analyzing them are required. These often involve clustering methods to automatically group similar groups of atoms or timesteps.

In this paper, we focus on simulations of very long molecular dynamics trajectories that show folding processes of single molecules <cit.>. 
Based on the starting positions of a molecule's atoms, physical forces on all other atoms can be calculated, yielding simulated trajectories with a temporal resolution that usually lies in the femtosecond range.
While the overall shape of the molecules is decisive for its function, several features are shown to be relevant, e.g., the torsions between C^α atoms that form the backbone of the protein, distances between specific atoms, or the solvent accessible surface area.


There are different use cases for which a clustering of MD trajectories is required. 
E.g., clusters of atoms that move similarly over time can be used to find dynamic domains, i.e., fragments of a protein that are internally rigid <cit.>. 
Clustering time steps in which the conformation of the protein, i.e., the arrangement of its atoms in space, is similar is used for detecting relevant states of the protein. 
Those states can be linked to (local) valleys in the energy landscape of the protein and are e.g. used for developing Markov models that describe a protein trajectory. 
Thus, one goal is to separate metastable states of a protein and the transitions between those states. 





§ MARKOV STATE MODELING
 

Markov state models (MSMs) can be used to describe the dynamics of a MD trajectory by using a n × n transition matrix. To generate an MSM from an MD trajectory X={x_1,...,x_t}, 
discrete states of the trajectory need to be identified. 
State-of-the-art methods use a two-step approach in which they first reduce the dimensionality of the data by using PCA or TICA and then clustering it with k-Means into n discrete states S={S_1,...S_n}<cit.>. These states are used to calculate the transition probability matrix for a given lag time τ by counting the transitions from time step t to t+τ: 


    p_ij(τ)=ℙ(x_t+τ∈ S_j|x_t∈ S_i)


The resulting MSM can be used for further analysis of the trajectory. In this paper, we analyze the clustering performance of our method in comparison to state-of-the-art methods for clustering MD trajectories based on their suitability as Markov states.


The variational approach for Markov models (VAMP)<cit.> is a generalization of the Variational approach to conformational dynamics (VAC) <cit.>, <cit.> and provides a tool to evaluate MSMs. The VAC allows the approximation of eigenvalues and eigenfunctions with statistical observables <cit.>. It also defines a family of scoring functions, called the VAMP-r, which can be calculated from data. This score can be used to compare the choices of hyper-parameters, like the featurization of the data or the used clustering, used for constructing the MSM. The score is calculated from the singular value decomposition of the Koopman operator. The Koopman operator is explained below. It is important to note that the VAMP score can only be compared for MSMs constructed with the same lag time τ and cannot be used to find the optimal lag time τ itself.

The Koopman operator gives the conditional expected value of an arbitrary observable g at time t+τ for a given x_t. The Koopman operator K_τ of a Markov process is defined as follows <cit.>:


    K_τ g(x)=𝔼[g(x_t+τ|x_t=x]


When choosing the Dirac delta function δ_y centered at y, applying the Koopman operator evaluates the transition density of the dynamics <cit.>.


    K(τ)=C_00^-1/2C_01C_11^-1/2≈ U_mΣ_mV_m^T,

where Σ_m=diag(σ_1,...,σ_m) are the first m singular values. 



  
*VAMP-r Score
The VAMP-score <cit.> defines a function that can be used to score the quality of MSMs for a set lag-time τ.  This score can be used to compare the methods used to create the discrete trajectories from the same featured trajectory <cit.>.

With Σ_m as defined above (see <Ref>), the VAMP-r score can be written as:

    VAMP-r=∑_i=1^m σ_i^r


The VAMP-2-score is based on the sum of squared eigenvalues of the transition matrix and can be interpreted as kinetic content <cit.>. In this paper, we use the VAMP-2 score to score the performance of different discretization methods for MD trajectories since there usually is no given ground truth to compare to. The most commonly used scores are the VAMP-1, which is analogous to the Rayleigh trace, and the VAMP-2, which is analogous to the kinetic variance <cit.>. 




§ SUBSPACE CLUSTERING
 
In high-dimensional datasets where the curse of dimensionality renders traditional clustering infeasible, subspace clustering finds clusters in a union of lower dimensional subspaces <cit.>. For a dataset X=[x_1, x_2, ..., x_n]∈ℝ^d it is often assumed that vectors x_j are drawn from a union of k subspaces {S_i}_i=1^k of unknown dimensionality <cit.>. As MD data is typically very high dimensional, subspace clustering is a suitable approach.
For details on subspace clustering, see <cit.>.
In this paper, we focus on spectral-based subspace clustering.



The advantages of spectral subspace clustering methods include a relatively simple implementation and robustness regarding data corruption and initialization <cit.>. Spectral subspace clustering methods usually perform two steps. First, a subspace representation of the original data is learned and used to construct an affinity matrix. Then spectral clustering is used on the affinity matrix to obtain the final clustering. Many different spectral algorithms like sparse subspace clustering (SSC) <cit.> or Low-Rank Representation (LRR) <cit.> have been proposed, for an overview see <cit.>.

While spectral clustering has been successfully applied to numeric data, it lacks the ability to handle temporal data. We, thus, turn our attention to methods that have been developed with a different application goal in mind, but, as we argue, with similarities that make them attractive for the design of our MD subspace clustering method.

Temporal subspace clustering (TSC) <cit.> is a subspace clustering algorithm designed to segment video sequences into separate human motions. 
Video sequences are high-dimensional time series, with one frame of the video at each timestep. 
Significant changes in a video sequence usually are not immediate, but take place over many frames. 
Neighboring time frames are more correlated to each other than more distant frames. 
In that sense, molecular dynamics data is similar to video data. 
Unlike classical subspace clustering methods, which do not take any relations between the dimensions into account, TSC is designed to take advantage of the sequential relationship in time series data. 
TSC uses a dictionary learning approach together with a temporal Laplacian regularization function to encode those relationships in the dataset. Dictionary learning is a method that learns a set of basis vectors, a so-called dictionary <cit.>. The original data can then be represented as linear combinations of the basis vectors. TSC combines the dictionary learning approach with a temporal regularization function to encode the sequential relationships in the data. 
A non-negative dictionary and a corresponding coding are learned from the given data creating a robust affinity graph.



§ MOSCITO: TEMPORAL SUBSPACE CLUSTERING FOR MD DATA
  

We are now ready to describe our MOSCITO (MOlecular dynamics Subspace Clustering with Temporal Observance) approach, the first subspace clustering method for molecular dynamics data that captures temporal relations. 
Existing methods <cit.> often rely on a two- or three-step approach consisting of 1) dimensionality reduction, 2) clustering, and 3) finding macrostates by combining several clusters. As these steps are handled in isolation, their separate results may not be optimal overall. To avoid such undesirable side effects,  MOSCITO takes a holistic approach: 
Firstly, in subspace clustering, it inherently combines clustering with dimensionality reduction. 
Secondly, MOSCITO incorporates temporal relation by weighting sequential neighbors in a time window. In contrast, existing subspace clustering algorithms cannot capture e.g., correlation between consecutive time steps.

For our new method MOSCITO, we exploit the similarity between MD data and video data (see <Ref>) as well as the advances in the field of data mining on videos: TSC <cit.> clusters video data while incorporating the temporal relations between time steps as well as allowing clusters in subspaces of the data. Thus, we base MOSCITO on TSC. However, TSC is not able to handle MD data, because despite some similarities between the videos and MD data (e.g., smooth temporal transitions between time steps and potential clusters in subspaces), they also have strong differences: for every time step , MD data describes a chain of atoms with its 3d coordinates, while video data contains colors for every pixel, where pixels are in a 2d grid. 
Thus, we need to change several steps and include knowledge from the field of MD as shown in Figure <ref>, which provides an overview of MOSCITO. 
The blue parts are equivalent to those of TSC <cit.>, see <Ref>. 
The green parts are adapted or changed completely. 
At first, essential MD features 
are extracted from the MD trajectory. 
From the features and a temporal regularization function, we learn a dictionary and the corresponding representation coding matrix. 
We then construct an affinity graph, and generate temporal segments by clustering the affinity graph.






  
*Feature extraction 
From the input MD trajectors, we extract features that capture essential properties. 

Cartesian coordinates: 
The C^α atoms determine the shape and function of a protein, thus we use the common approach to focus solely on their 3d coordinates, rather than considering all atoms. 
For each frame,  coordinates are centered and aligned to the first frame of the trajectory.

Backbone torsions: 
The backbone torsions determine the main structure of a protein. The torsion angles ϕ and ψ describe the dihedral angles between atoms of the backbone. Backbone torsions better describe the structure of a protein than absolute coordinates <cit.> as internal angles are independent of the global rotation of the protein, removing the need to align all frames to a fixed rotation.  

Distance based: Heavy atoms are all atoms that are not hydrogen. We extract the minimal distance between heavy atoms <cit.> (we exclude the pairs at indices(i, i+1) and (i, i+2), as their distance is largely determined by their chemical bonds). Following <cit.>, distances d are transformed to the negative exponential function e^-d  

Flexible torsions: The flexible torsion feature consists of the χ_1 to χ_5 side chain torsions. The side chains of a protein are chains of atoms, connected to the backbone. The angles are transformed into sine and cosine angles.

Solvent accessible surface area: For each frame of the trajectory, the solvent accessible surface area (SASA) of each residue is calculated using the Shrake-Rupley algorithm <cit.>. The SASA of a residue is calculated by summing up the SASA of each of the residues atoms

3D shape histogram: Based on <cit.>, a combined 3D shape histogram for each frame of the trajectory is created. The state space is separated into five equally spaced shells around the origin of the coordinate system. The shells are split up further by combining them with 12 sectors into 60 cells. Histograms are created by counting the number of atoms in each cell. Unlike <cit.>, all atoms instead of only the surface atoms were used as we are not only interested in the shape and function of the protein at a certain time step, but also in its internal folded structure. Final histograms are normalized between 0 and 1.


  
*Temporal regularization
For trajectories, data points at neighboring time steps are usually more closely related than those at more distant time steps. Let the i-th column of the coding matrix Z a representation of vector x_i, then the encodings of its sequential neighbors, e.g., z_i-1, z_i+1 should be similar to z_i. To achieve this, we use a temporal Laplacian regularization function f(Z). 


    f(Z)=1/2∑_i=1^n ∑_j=1^n w_ij ||z_i - z_j||_2^2 = tr(ZL_T Z^T)

L_T is the temporal Laplacian matrix defined as L_T = D̃ - W. W is the weight matrix, capturing the sequential relationship in X, and D̃_ii = ∑_j=1^nw_ij is a diagonal matrix with the sum of the rows of W on its diagonal. 


The temporal regularization in MOSCITO weighs the sequential neighbors using a weight matrix W∈ℝ^n× n for a trajectory with n timesteps. 
Row i of W represents the weights of x_i's neighbors. 
In the base case, binary weights are used. 
This weighs all s sequential neighbors of x_i the same, independent of the actual temporal distance between the data points. 
In contrast to TSC, MOSCITO offers several 
Naturally, points in close temporal proximity are more strongly correlated than points further apart in time. To model the decrease in correlation, we introduce different decaying weightings: 

Binary:  weights of 1 for points in the neighborhood and 0 outside

Gaussian: Gaussian curve fitted to the neighborhood

Logarithmic: Logarithmic decay with a value of 1 for the direct neighbors dropping logarithmically to 0 for x_i+s and x_i-s

Exponential: Exponential decay w.r.t. the temporal distance 



  
*Dictionary learning Similar to TSC <cit.>, MOSCITO learns to represent dataset X by a dictionary D and a coding matrix Z, such that X≈ DZ <cit.>. 

Adopting a least square regression approach to dictionary learning, we use the following objective function <cit.>

    min_Z,D||X-DZ||_F^2+λ_1||Z||_F^2
    s.t.  Z ≥ 0, D ≥ 0, ||d_i||_2^2 ≤ 1,

where ||Z||_F^2 denotes the Frobenius norm, and we use λ_1 for balancing both terms. 

Importantly, MOSCITO integrates temporal regularization to capture relationships between nearby time steps. We achieve this by adding temporal regularization function f(Z) such that our new objective function is the same as in TSC <cit.>

    min_Z,D||X-DZ||_F^2+λ_1||Z||_F^2 + λ_2f(Z) 
    s.t.  Z ≥ 0, D ≥ 0, ||d_i||_2^2 ≤ 1,

with λ_2 tradeoff parameter <cit.>.

To solve the dictionary learning problem efficiently, we follow the alternating direction method of multipliers (ADMM) <cit.>. It decomposes the problem into subproblems that are more easy to solve. Using auxiliary variables U, V, Eq. <ref> is equivalent to

    min_Z,D,U,V||X-UV||_F^2+λ_1||V||_F^2 + λ_2f(V) 
    s.t.  U=D, V=Z, Z ≥ 0, D ≥ 0, ||d_i||_2^2 ≤ 1.

The augmented Lagrangian of Eq. <ref> is:

    ℒ    = 1/2||X-UV||_F^2 + λ_1||V||_F^2 + λ_2 tr(VL_TV^T) 
        + ⟨Λ, U_D⟩ + ⟨Π, V-Z⟩ + α/2||U-D||_F^2 + β/2||V-Z||_F^2,

with Λ and Π Lagrangian multipliers <cit.>, which can now be solved by alternatingly minimizing ℒ respective to V, U, Z and D. 

Updating V: To find the minimum of ℒ, its derivative with respect to V is set to zero, resulting in the following equation <cit.>:

    [I⊗ (U^TU+(λ_1+α)I) + λ_2L_T⊗ I] vec(V)=vec(U^TX-Π+α Z)


Updating U: As for V, set derivative of ℒ wrt. U to zero <cit.>:

    U=(XV^T-Λ+α D)(VV^T+α I)^-1


Updating Z and D: The update rules for Z and D are <cit.>: 

    Z=F_+(V+Π/β)
    ,     
        D=F_+(U+Λ/α)

where (F_+(A))_ij=max{A_ij, 0}, non-negativity constraint. Additionally, each column of D is normalized to unit length  ||d_i||_2^2≤1.



  
*Clustering
Using the learned coding matrix Z, an affinity graph is constructed. To capture sequential relationships, the following similarity measure is used <cit.>: 

    G(i, j)=z_i^T z_j/||z_i||_2||z_j||_2

Spectral clustering on graph G yields the final clusters.





§ EXPERIMENTS
 



  
Implementation 
We extract features of MD trajectories using libraries PyEMMA <cit.> and MDTraj[<https://www.mdtraj.org>].
Matrices are represented as NUMPY arrays and sparse matrices by scipy sparse matrices. 
We efficiently solve linear equations of the form Ax=b using the PyPardiso 0.4.2[<https://github.com/haasad/PyPardisoProject>] (version 2021.4) wrapper for the Intel oneMKL PARDISO library. 
For spectral clustering, we use the scikit-learn implementation. 
Our Python code for MOSCITO is https://github.com/Mhae98/MOSCITOpublicly available [<https://github.com/Mhae98/MOSCITO>]. 
For PCA, TICA and k-Means we use the implementations in PyEMMA.
The implementation of SSC is based on <cit.> and a Python translation of SSC [<https://github.com/abhinav4192/sparse-subspace-clustering-python>] 
All experiments run on a virtual machine with 64 CPUs and 512GB of RAM running Ubuntu 22.04. 





  
Datasets and parameter settings
We use MD trajectories of four proteins: 2F4K, Ace, 2WXC, and Savinase, which differ in the number of atoms | A|, number of Cα atoms | Cα|, and number of residues | R|, as summarized in Table <ref>, together with number of  trajectories and timesteps in each trajectory.

If not specified otherwise, parameter default values are 
 d= 60, s= 3, λ_1= 0.01, λ_2= 15, a= 0.1, b= 0.1.





 §.§ MOSCITO settings and features
 
In this section, we analyze the parameter sensitivity and impact of feature extraction for MOSCITO. 
For the following experiments, 10 trajectories of the 2WAV-protein are used. The Markov state models are constructed for k∈{10, 50, 100, 200} clusters with a lag time τ∈{1, 10, 100}. The average VAMP2-score of the 10 trajectories is calculated for each choice of clusters k and lag time τ. The VAMP2-score is calculated for five eigenvalues.
It is important to keep in mind that the VAMP2-score for different MSM lag times is not directly comparable, only for different numbers of clusters.




  §.§.§ Dictionary size
  d defines the dimensionality of the learned dictionary and thus the number of subspaces. As such, it is crucial for the clustering and runtime performance of MOSCITO, where larger dictionary size d provides more subspaces to represent the data.
We study d∈{10, 20, ..., 100}, and additionally s∈{2, 3, 4} number of sequential neighbors. Figure <ref> shows VAMP2-scores for different numbers of clusters (rows) and MSM lag times (columns): clustering performance improves for larger dictionary size d until it converges. For smaller MSM lag times, the VAMP2-score converges faster than for larger lag times, independent of the number of clusters. With an increasing number of sequential neighbors s, the overall shape of the curves is highly similar, with a slight preference for more neighbors that capture more of the temporal relationships.






  §.§.§ Sequential neighbors
 
Setting the number of sequential neighbors s in the temporal regularization to 0, the temporal aspect is not considered, with growing s longer temporal relationships are captured. The scores for s∈{0,...,12} are shown in Figure <ref> for different MSM lags (columns) and number of clusters (rows). 
The number of sequential neighbors s clearly affects clustering performance, with too few or too neighbors reducing performance, but with a broad stable range where temporal relationships are successfully captured.

Figure <ref> shows the clustering of a 2F4K protein trajectory with a long folded state into 5 clusters represented by different colors. With no (top) or few (center) sequential neighbors, time intervals comprised by a cluster can be very short.
Increasing the number of sequential neighbors (bottom) leads to continuous clusters with mostly uninterrupted blocks.  





  §.§.§ λ parameters
 
We study settings of the tradeoff parameters  λ_1 ∈{0.01, 0.05, …, 15} for sparsity  of the coding matrix Z, λ_2 ∈{0.01, 0.05, …, 35} for the temporal regularization function.
Figure <ref> shows that best performance is achieved for low values of λ_1  and higher values of λ_2 (light areas in the heatmap). This shows that temporal regularization is the more important factor for optimal clustering performance. 








  §.§.§ Learning rates
 
α, β∈{0.01, 0.05, …, 15} 
define how much Lagrangian multipliers Π and Λ change in each iteration of the algorithm.
Heat maps in Figure <ref> show that the results are stable across parameter settings, with VAMP2-scores in each heatmap relatively close, meaning our default values work well.





  §.§.§ Weighting mode
 
All weighting modes for the Laplacian regularization function: binary, Gaussian, logarithmic, and exponential, are tested for different numbers of sequential neighbors.
In Figure <ref> we see that for few sequential neighbors, the performance for all weighting modes is similar. For many sequential neighbors, performance starts to decrease. In most cases, the performance when using the binary weighting mode decreases earlier than for the other methods, but all reliably reach the same maximal score.






  §.§.§ Feature comparison
 
We analyze the above described features and their combinations via simple concatenation. 
Figure <ref> 
shows the average VAMP2-score for 10 trajectories of the 2WAX-protein, using the five largest eigenvalues for scoring. 






Best results for individual features are achieved with backbone torsions. Overall, combinations of backbone torsions, SASA, and shape histograms yield the best results. Figure <ref> shows the importance that MOSCITO gives to each of the features in its dictionary when clustering the full-dimensional feature space. While, e.g., the feature "flex" (flexible torsions) yields the worst VAMP2 score (left), it still has a high influence on the results (right), explaining the suboptimal VAMP2 score when using the full feature set.




 §.§ Clustering performance
 
In this section, the clustering performance of MOSCITO is compared to that of state-of-the-art approaches that are currently used in the field. We compare to PCA with k-Means and to TICA with k-Means because of their widespread use <cit.>. Additionally, we compare against the sparse spectral clustering algorithm <cit.>, a subspace clustering approach for molecular dynamics data.
For PCA and TICA, the default parameter settings of PyEmma <cit.> were used. SSC is a subspace clustering algorithm, thus it does not need any dimensionality reduction in the preprocessing. 

A general challenge in this kind of evaluation is the fact that no ground truth is given. As such, clustering performance is compared by constructing a MSM and scoring it using the VAMP2-score, in order to obtain quantitative insights. Still, visual inspection plays an important role in assessing the clustering performance, in particular, the temporal connectivity of clusters.



  
*Quantitative analysis
Figure <ref> shows the average VAMP2-score for 20 trajectories of the 2F4K protein in the first row. For MSM lag time 1, the performance of all methods is comparable with the exception of 10 clusters. There, MOSCITO and TICA + k-Means deliver the best performances. For longer MSM lag times, the scores separate with SSC performing better.




In the second row of Figure <ref>, we see the average VAMP2-score for 10 trajectories of the 2WXC protein. For the MSM lag times 1 and 10, the performance for all methods is relatively similar for large cluster sizes. For small cluster sizes, SSC performs poorly, while the other methods already reach close to their maximal score for only 10 clusters. For MSM lag time of 100, PCA + k-Means drastically drops in performance, while the other methods maintain comparable scores to each other.


The third row of Figure <ref> shows the average VAMP2-score for 28 trajectories of the Ace protein. Similar to the results of 2WXC, the performance for a MSM lag time of 1 and 10 is comparable between all methods for a large number of clusters. For few clusters, SSC again performs significantly worse. For MSM lag time of 100, the gap between the methods grows, with SSC performing best.

The last row of Figure <ref> shows the average VAMP2-score for 2 trajectories of Savinase. For all MSM lag times, SSC performs best while TICA + k-Means performs worst. Unlike for the other proteins, all methods require a large number of clusters to reach the maximal VAMP2-score. This might be due to the large size of Savinase and shorter trajectories compared to the other proteins.



  
*Qualitative analysis
The purpose of modeling temporal aspects is to obtain clusters with larger continuous segments, ideally retrieving macrostates directly.  Figure <ref> visually compares the segmentation of clustering methods, either by directly clustering the trajectory into five clusters or, alternatively by coarse-graining 50 clusters into five clusters using PCCA+ <cit.> (see <Ref>). We regard 
a 2F4K trajectory with a known long folded segment.


Directly clustering, i.e. without PCCA+, PCA + k-Means and SSC struggle to find the large folded state as one large block. MOSCITO and TICA+k-Means find similar clusterings and succeed to uncover the large folded state. When applying PCCA+, PCA + k-Means, TICA + k-Means, and MOSCITO generate similar clusterings, comparable to the ones of MOSCITO and TICA + k-Means without PCCA+. In this specific case, SSC fails to find the large segment, which is reflected in the respective VAMP2-score for this specific trajectory which is unusually low.



 §.§ Clustering runtime
 
In this section, the runtime of MOSCITO is studied and compared to the runtime of PCA + k-Means, TICA + k-Means, and SSC. The runtimes for MOSCITO are measured for 20 iterations. The following runtimes are the average of 10 2WXC-trajectories clustered into ten clusters, using all atom coordinates as features. 
The main factors influencing the runtime of MOSCITO, besides the size of the input data, are the dictionary size d and the number of sequential neighbors s. The effect of different dictionary sizes [10, ..., 100] on the runtime is shown left in Figure <ref>. The effect of different numbers of sequential neighbors [0, ..., 25] is on the right. 




The runtime of MOSCITO grows with increasing dictionary sizes. For optimal cluster performance, a large enough dictionary size has to be selected. Large values do not affect clustering performance negatively, but directly impact runtime. A dictionary size between 60 and 80 provides a good tradeoff.
The runtime generally grows with increasing number of regarded sequential neighbors. 
The overall run-to-run variance for the different trajectories is close to zero, indicating a stable runtime for a fixed number of iterations. 

For the comparison of the different clustering methods, backbone torsions were used. The comparison between the different clustering algorithms is shown in Figure <ref>.
PCA + k-Means and TICA + k-Means are both much faster than the subspace clustering methods.  Compared to SSC, MOSCITO is substantially faster. 



 §.§ Discussion
 
We conclude that MOSCITO provides at least comparable performance to state-of-the-art clustering methods for MD data. We evaluated the clustering performance by the clusters' suitability as Markov states using the VAMP2-score. The visual analysis of MOSCITOs clustering results shows a better clustering for higher numbers of sequential neighbors. Thus, MOSCITO is a very promising approach for creating MSMs in a one-step-approach.

MOSCITO's parameter robustness shows that there are good tradeoffs in dictionary size between runtime and clustering performance. The number of sequential neighbors matters, and our experiments indicate a good default value is between 3 and 5 to capture nearby temporal relations. 
In terms of feature extraction, MOSCITO on backbone torsions leads to best results, followed by Cα coordinates and SASA. Combining features does not improve the overall clustering performance. 
While simpler non-subspace methods are even faster, MOSCITO provides a faster runtime than SSC, and state-of-the-art quantitative results, with qualitative analysis showing the promise of the approach.




§ RELATED WORK
 
Distance-based clustering algorithms suffer from the curse of dimensionality when it comes to high-dimensional data like MD data.
Thus, the dimensionality is often reduced in a preprocessing step, e.g., with principal component analysis (PCA) or time-lagged independent component analysis (TICA) <cit.>. Both are widely used for analyzing MD data <cit.>. 
PCA transforms the data onto the axes with the highest variance, also known as principal components. 
PCA is a commonly used for MD data, e.g., in PyEMMA <cit.>, MSMBuilder [<http://www.msmbuilder.org>], or MDAnalysis [<https://www.mdanalysis.org/>].
TICA (Time-lagged independent component analysis) <cit.> linearly projects the data to independent, uncorrelated components (IC). 
The ICs are computed by maximizing the autocovariance for a fixed lagtime τ, and are frequently used for MD data (<cit.>, <cit.>).


State-of-the-art methods often follow a two-step approach, where the trajectories are discretized into a large number of microstates, and subsequently combined to the relevant macrostates.
Since the transitions between the metastable states are taking place over multiple timesteps
it is hard to assign time steps that are part of a transition to a metastable state clearly. 
Perron-Cluster Cluster Analysis (PCCA) (<cit.>,<cit.>) and a more robust version PCCA+ (<cit.>, <cit.>) fuzzily assign the microstates to their corresponding macrostates based on spectral methods for transition matrices:


The transition matrix P of a decomposable MSM with n_c macrostates can be reordered to a block-diagonal matrix P̃ with n_c blocks. 
For the matrix X containing the n_c dominant eigenvectors of P̃, the rows belonging to states in the same block are identical <cit.>. 
These rows represent the n_c corners of a (n_c-1) dimensional simplex.

For MSMs with transition states, i.e., nearly decomposable MSMs, the structure of the simplex is only slightly perturbed. To assign the microstates to the macrostates, they are fuzzily assigned to the n_c corners of the simplex based on their position in the simplex. The goal is to find a linear transformation A∈ℝ^n_c× n_c so that χ=XA where χ_i are the membership vectors. Those memberships all have to be positive and sum up to 1. 
The resulting macrostates generally cannot be considered to be a Markov chain <cit.>.


Time series data can be found in many fields, like traffic analysis, geology, computer vision, or molecular dynamics data. Many existing subspace clustering algorithms do not take the sequential properties of time series data into account <cit.>. One of the first subspace clustering methods taking advantage of sequential relationships in data was SpatSC introduced by <cit.>. A penalty term used to preserve sequential relationships in the learned affinity graph was added to the subspace clustering problem. Based on the same idea, OSC, a more robust subspace clustering method for sequential data was introduced by <cit.>. 
TSC <cit.> further improved on the idea by defining a Laplacian regularization to encode the temporal information. 
Our method MOSCITO builds upon TSC. 

To the best of our knowledge, only one subspace clustering method has been used on molecular dynamics data, namely SSC <cit.>, which was first introduced in the field of image processing <cit.>. We include SSC in our experiments as comparative method, thus we introduce it in detail in the following. 
Note that SSC does not take advantage of any properties that distinguish time series from other high-dimensional data, leading to suboptimal use of the available information in MD data.
The main idea behind SSC is to take advantage of the self-expressiveness of the data, i.e., assuming that every point y_i∈ Y can be rewritten as a linear combination of other points in the dataset. Each datapoint y∈∪_i=1^nS_i can be written as

    y_i=Yc_i,   c_ii=0,

where Y=[y_1,...,y_N] and c_i:=[c_i1  c_i2  ...  c_iN]^T with the constraint c_ii=0 eliminating the trivial solution of representing the point as itself. 
Since each subspace usually has more data points than dimensions the self-expressive dictionary Y is generally not unique. 
Among all those possible solutions for Equation <ref> there is a sparse solution so that the non-zero entries of c_i correspond to datapoints in the same subspace as y_i. This solution can be found by minimizing the l_1-norm of c_i, which can be solved efficiently and is known to prefer sparse solutions <cit.>.

These sparse representations are transferred to an undirected graph where points that are part of another point's sparse representation are connected. Based on this graph, spectral clustering detects the desired subspace clusters. 




§ CONCLUSION
 
We introduced MOSCITO, a subspace clustering method for high-dimensional molecular dynamics data. It uses temporal regularization to capture the relationship of sequential neighbors in time-series data.
While traditional and state-of-the-art methods rely on coarse-graining using PCCA+ in a two-step approach, MOSCITO achieves similar overall performance in a single clustering step.

Especially for small numbers of clusters, the temporal regularization in MOSCITO yields clusters consisting of connected time steps, fitting the real world more closely than alternative methods without relying on any postprocessing steps. 
Compared to current subspace clustering methods on MD data, MOSCITO provides a better runtime.
In future work, the integration of a multi-view approach into MOSCITO might take advantage of the additional information provided by using multiple features.


ACM-Reference-Format




 Paper Title: AutoPV: Automatically Design Your Photovoltaic Power Forecasting Model.















AutoPV: Automatically Design Your Photovoltaic Power Forecasting Model
    Dayin Chen,
        Xiaodan Shi1,
        Mingkun Jiang,
        Haoran Zhang, Senior Member, IEEE,       
        Dongxiao Zhang1,
        Yuntian Chen, Member, IEEE,
        and Jinyue Yan1


* Xiaodan Shi Dongxiao Zhang and Jinyue Yan are the corresponding authors.
Dayin Chen is with Department of Building Environment and Energy Engineering, The Hong Kong Polytechnic University, Hong Kong SAR, China; International Centre of Urban Energy Nexus, The Hong Kong Polytechnic University,  Hong Kong SAR, China; Research Institute for Smart Energy, The Hong Kong Polytechnic University, Hong Kong SAR, China; Eastern Institute for Advanced Study, Eastern Institute of Technology, Ningbo, China. E-mail: 23038748r@connect.polyu.hk
Xiaodan Shi is with School of Business, Society and Technology, Mälardalens University, 72123 Västerås, Sweden; Center for Spatial Information Science, the University of Tokyo, 5-1-5 Kashiwanoha, Kashiwa-shi, Chiba 277-8568, Japan. E-mail: xiaodan.shi@mdu.se

Mingkun Jiang is with PV Industry Innovation Center, State Power Investment Corporation, 710061 Xi'an, Shaanxi, China. E-mail: jiangmk_spic@163.com
Haoran Zhang is with School of Urban Planning and Design, Peking University, Shenzhen, China. E-mail: h.zhang@pku.edu.cn
Dongxiao Zhang and Yuntian Chen are with Ningbo Institute of Digital Twin, Eastern Institute of Technology, Ningbo, China. E-mail: dzhang@eitech.edu.cn; ychen@eitech.edu.cn
Jinyue Yan is with Department of Building Environment and Energy Engineering, The Hong Kong Polytechnic University, Hong Kong SAR, China; International Centre of Urban Energy Nexus, The Hong Kong Polytechnic University,  Hong Kong SAR, China; Research Institute for Smart Energy, The Hong Kong Polytechnic University, Hong Kong SAR, China. E-mail: j-jerry.yan@polyu.edu.hk



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



§ ABSTRACT 
Photovoltaic power forecasting (PVPF) is a critical area in time series forecasting (TSF), enabling the efficient utilization of solar energy. With advancements in machine learning and deep learning, various models have been applied to PVPF tasks. However, constructing an optimal predictive architecture for specific PVPF tasks remains challenging, as it requires cross-domain knowledge and significant labor costs. To address this challenge, we introduce AutoPV, a novel framework for the automated search and construction of PVPF models based on neural architecture search (NAS) technology. We develop a brand new NAS search space that incorporates various data processing techniques from state-of-the-art (SOTA) TSF models and typical PVPF deep learning models. The effectiveness of AutoPV is evaluated on diverse PVPF tasks using a dataset from the Daqing Photovoltaic Station in China. Experimental results demonstrate that AutoPV can complete the predictive architecture construction process in a relatively short time, and the newly constructed architecture is superior to SOTA predefined models. This work bridges the gap in applying NAS to TSF problems, assisting non-experts and industries in automatically designing effective PVPF models.





Neural Architecture Search, AutoML, Photovoltaic power forecasting, Time series forecasting




§ INTRODUCTION




Photovoltaic (PV) power forecasting (PVPF) involves predicting future PV power generation based on historical power data and corresponding weather features. The variable and non-controllable nature of PV power generation necessitates precise forecasting, which is crucial for maintaining grid balance and offering technical assistance during trading activities <cit.>. With the advent of deep learning, 
technologies such as Multi-Layer Perceptron (MLP) <cit.>, Convolutional Neural Networks (CNN) <cit.>, Long Short-Term Memory (LSTM) <cit.>, and Temporal Convolutional Network (TCN) <cit.> have been extensively applied in the field of PVPF. Even in physics-based PVPF models, deep learning has been leveraged in various applications, such as enhancing numerical weather prediction (NWP) <cit.> and sky image analysis <cit.>. These advancements highlight the growing mainstream adoption of deep learning techniques in the field of PVPF.



As known to us, PVPF is inherently a time series forecasting (TSF) problem <cit.>, which is a prominent area within data mining. The recent surge in transformer architectures <cit.> has led to the development of numerous transformer-based models such as informer <cit.>, Autoformer <cit.>, and FEDformer <cit.>, which have demonstrated exceptional performance in TSF problems. Concurrently, several non-transformer-based TSF models, including DLinear <cit.> and TSMixer <cit.>, have also been developed and shown strong performance. These typical TSF models have now started to be applied in the context of PVPF tasks as well.



Given the plethora of PVPF models available, a significant challenge lies in selecting an optimal model for a specific task. This issue is particularly challenging for non-experts and industries, as building a PVPF model requires expertise across multiple domains such as Energy Engineering and Computer Science. Factors including data granularity, data scale, and data distribution can significantly impact the final model design, thus making the model design process a non-trivial endeavor. A model that performs excellently on one task may not perform as well on another, which can lead to suboptimal decision-making and introduce various risks. Therefore, an efficient and automatic model search and construction strategy is essential in the current landscape of proliferating PVPF models.


Neural Architecture Search (NAS) <cit.> is an emerging technology aimed at automating the process of model structure discovery and design, which can address the aforementioned issue. It comprises three fundamental elements: the search space, the evaluator, and the search strategy. The search space encompasses a variety of module choices that can be assembled into a cohesive architecture. The search strategy employs a specific algorithm to navigate through and combine these modules, with the goal of identifying an optimal architecture. Meanwhile, the evaluator is tasked with rating each generated architecture.






However, despite the significant advancements in NAS for many years, there has been limited exploration of NAS specifically for TSF tasks. Instead, most current NAS efforts are predominantly focused on visual tasks. This focus is largely due to the relatively straightforward nature of vision model components, such as convolutional and pooling layers, which lack complex interdependencies and can be seamlessly integrated with other standard modules, including linear layers, activation functions, and softmax layers. In contrast, TSF models are often cohesive, indivisible entities with distinct design principles. For instance, some TSF models employ seasonal-trend decomposition <cit.>, while others analyze data in the frequency domain <cit.>. These models are designed from multiple and varied perspectives, making it challenging to decompose them into modular components that can be used to construct an effective search space.






In this study, we introduce AutoPV, a novel NAS system for the automated search and design of predictive architectures tailored to specific PVPF tasks. We construct a novel NAS search space that encompasses various typical data processing techniques from state-of-the-art (SOTA) TSF models, as well as the foundational deep learning models commonly used in PVPF tasks. This meticulously designed search space encompasses four different stages with twelve optional parameters, enabling the exploration of approximately 2.5 × 10^6 potential architectures.






To evaluate the effectiveness of AutoPV, we designed two distinct PVPF tasks: one utilizing only historical data and the other incorporating future weather information. These tasks further encompass six sub-tasks with different forecasting lengths, ranging from half a day ahead to one month ahead. All the evaluation experiments were conducted using the dataset from the Daqing PV Station in China, covering the period from 2022 to 2023. Experimental results show that AutoPV exhibits outstanding performance across different PVPF tasks compared to other benchmarks. We believe this work could help bridge the gap in applying NAS to PVPF problems and potentially extend to TSF problems.



The principal contributions of this work are outlined as follows:

  
  * We construct a detailed knowledge graph and NAS search space for the domains of PVPF. This helps address the research gap in the automated design of effective PVPF models.


  
  
  * Based on the constructed search space, we develop an automated PVPF model design system called AutoPV, which can assist non-experts in searching for effective models for specific PVPF tasks.

  
  
  * We conduct extensive experiments on a PV power dataset from the Daqing Photovoltaic Station in China to demonstrate the effectiveness of AutoPV across various PVPF tasks and provide a comprehensive analysis of its characteristics. Results show the new searched architecture improves performance by 4.52% compared to the best baseline model. Furthermore, when incorporating future weather information, it could boosts the improvement by 9.88%.

  






§ RELATED WORK



 §.§ Photovoltaic Power Forecasting (PVPF)


PVPF tasks involve using historical PV power generation data along with weather information to model future PV power generation. Existing works involve a wide array of different models and also explore various forecasting lengths. For example, Gao et al. <cit.> use LSTM for day-ahead PV power prediction based on weather classification. Mellit et al. <cit.> evaluate short-term PVPF (less than 8 time steps ahead) using deep learning models like LSTM, GRU, and CNN. Li et al. <cit.> employ a TCN model for hours-ahead utility-scale PV forecasting. However, no single model has consistently outperformed others across different prediction scales or datasets. Efficiently designing an optimal model for a specific PVPF task remains a challenging problem.












 §.§ Time Series Forecasting (TSF) Models



The following content introduces two main architectural categories of TSF models.


Transformer-based Models. The core idea of the transformer <cit.> is to allow the model to dynamically focus on different parts of the input by computing attention scores. In recent years, to address the computational complexity of the vanilla transformer <cit.> and enhance performance, researchers have proposed numerous transformer variants <cit.>. Among these, some variants such as Autoformer <cit.>, Informer <cit.>, FEDformer <cit.>, iTransformer <cit.> and PatchTST <cit.> are specifically designed for TSF problems. These methods have made various modifications to the vanilla transformer from different perspectives, such as incorporating frequency domain information <cit.>, dividing time tokens into patches <cit.>, and introducing data decomposition <cit.>, to better capture the characteristics of time series data.



















Non-transformer-based Models. Although transformer-based structures are currently popular, their validity and necessity for TSF problems have been questioned by some researchers <cit.>. They maintain that non-transformer-based models can achieve equally good performance while having a lower computational burden. Typical non-Transformer-based models include DLinear <cit.>, FreTS <cit.>, MICN <cit.>, TSMixer <cit.>, and TimesNet <cit.>, among others. These models employ various time series data processing methods, such as decomposition, multi-scale decomposition, and frequency domain analysis, to extract high-level information from the data. The extracted information is often projected to the output using a fully connected (FC) layer. These models have demonstrated superior performance on a variety of publicly available datasets.





In recent years, a new approach of applying diffusion models to TSF has also emerged <cit.>. However, considering that the diffusion model is fundamentally different in principle from the other two types of methods (i.e., transformer-based models and non-transformer-based models), we have decided not to include diffusion models in the construction of the search space for this study.




 §.§ Neural Architecture Search (NAS)



NAS was first proposed by Google in the work <cit.>, where they used recurrent neural networks (RNNs) and reinforcement learning to automatically search and construct an image classification model on the CIFAR-10 dataset. The search space included parameters such as the length, width, and stride of the convolution operations. In their follow-up work the next year <cit.>, they refined the approach by searching the structure of a single cell and then repeating it multiple times to build the final network. While this method improved the search efficiency, it still required training on 500 GPUs for 4 days, which is a huge burden for most industries. Additionally, tuning the hyperparameters for the reinforcement learning process is also a complex and tedious task. This motivates people to seek more efficient search algorithm and certain specialized tricks to accelerate the search process of NAS.

Weight sharing is a typical improvement strategy which sharing weights among the different architectures <cit.>. However, the stability and effectiveness of this method have been questioned <cit.>. Beyond that, several studies have proposed different search strategies that eliminate the need for tedious hyperparameter tuning, such as evolutionary algorithm <cit.>. In the work <cit.>, the BANANAS algorithm is introduced, which utilizes Bayesian optimization techniques to search for optimal architectures. Furthermore, in the work<cit.>, the BANANAS algorithm has been extended for multi-objective optimization. This extension, referred to as MoBananas, has been demonstrated to outperform many other existing NAS search algorithms












One significant challenge in NAS is that most research efforts have focused on searching CV models, as the modules and operations in the CV search space generally do not have complex relationships or restrictions <cit.>. However, the need for automated network design extends to many other domains as well. In recent years, researchers have gradually explored the application of NAS to other areas, particularly by constructing domain-specific search spaces. For example, Seng et al. apply NAS to search for the federated architecture and hyperparameters in their work <cit.>. Wang et al. <cit.> use NAS to search the structure of transformer models for natural language processing (NLP) tasks. Additionally, Wang et al. <cit.> constructed a NAS search space based on five selected models for trajectory prediction.


To the best of our knowledge, there is limited prior work on extending NAS to the TSF or PVPF domain. In this study, we aim to address this gap by building a new search space and implementing an NAS system in the specific context of PVPF.





§ METHODOLOGY







 §.§ Problem Description






PVPF. PVPF task can be formally defined as follows:
The input is a time series of historical PV power data, denoted as 𝐗 = {x_1, x_2, …, x_t}∈ℝ^T_s × D, where x_t represents the data sample at time t, T_s is the length of the historical time series, and D is the dimensionality of each data point. A single data point comprises 1 dimension representing PV power generation data and (D-1) dimensions corresponding to weather information features. The searched predictive model, ℳ(𝐗), takes 𝐗 as input and outputs a series of future PV power generation predictions, denoted as Ŷ = {ŷ_t+1, ŷ_t+2, …, ŷ_t+n}∈ℝ^T_p × 1. Here, ŷ_t+n indicates the predicted value at step n in the future, and T_p is the length of the predicted time series. Each data point in Ŷ has only 1 dimension, representing the forecasted PV power generation.





Given the current availability of future weather information, in addition to the standard form of PVPF tasks, we have also designed another PVPF senario that utilizes future weather information, which is named PVPF Task 2. As depicted in Fig. <ref>, the future weather data is appended to the input historical sequence data along the time dimension. To handle the unknown future PV power data, we employ zero padding as a placeholder. In this case, the new input data length becomes T_s^' which equals T_s + T_p. In real-world scenarios, future weather information is typically obtained from weather reports, which can inherently contain biases. To account for these uncertainties, dynamic Gaussian noise is often added to the appended future weather data. The variance of this noise increases exponentially as the forecasting length grows longer.





















NAS. NAS requires first building a search space, which we denote as 𝕊. By selecting modules from 𝕊, we can construct various PVPF architectures ℳ. The overall goal of the NAS framework is to discover and construct an optimal architecture ℳ^* that minimizes the future PV power prediction error. In this work, we uniformly use the L1 loss, also known as mean absolute error (MAE), as the training and evaluation criterion. Let the ground truth of PV power generation from t+1 to t+n be 𝐘 = {y_t+1, y_t+2, …, y_t+n}∈ℝ^T_p × 1. The objective function is then:








    ℳ^*   = min_ℳ∈𝕊 MAE_𝔻_val(𝐖^*_ℳ, ℳ) 
    s.t.     𝐖^*_ℳ = min_𝐖ℒ_𝔻_train(𝐖, ℳ)




Where ℒ_𝔻_train(𝐖, ℳ) represents the training loss of the selected model ℳ on the training dataset 𝔻_train with the model parameters 𝐖. 










 §.§ Search Space



















The search space is the most critical component of the NAS framework. As detailed in Table <ref>, we divide the overall AutoPV search space into four stages and employ a total of 12 parameters to describe the entire searched architecture. This encompasses the complete model-building and training pipeline, including feature selection, data processing, model construction, and training process settings. The optional modules in each stage are derived from either typical PVPF works or SOTA TSF models. To ensure that the constructed models remain relatively lightweight, we exclusively extract processing methods from non-transformer-based models to build our search space. Experimental results demonstrate that the models constructed using these simple modules can outperform predefined SOTA transformer-based models.








In the following sections, we illustrate each stage of the AutoPV search space in detail.








Stage 1: Feature Selection. In this initial stage, we refine the data by selecting relevant features, as including irrelevant ones can impair the model's effectiveness. We use two common feature selection methods (FSM) in the PVPF domain: Maximum Relevance and Minimum Redundancy (mRMR) <cit.>, and Pearson Correlation <cit.>. The Pearson Correlation coefficient measures the linear association between a feature and the target variable, while mRMR evaluates both feature relevance and inter-feature redundancy. We also optimize the feature selection threshold (FST) for these methods, treating it as an additional hyperparameter.












Stage 2: Data Processing. In this phase, a series of preprocessing operations are applied to the input data to prepare it for the core predictive model.

The first parameter, Data Generalization Method (DGM), determines whether Gaussian noise should be added to the original dataset. This technique is known for enhancing the model's training stability and generalization capability <cit.>.

The second parameter involves two common Stationarization Methods (SM) to address distributional shifts. The first method, RevIN <cit.>, is an adaptable instance normalization technique that normalizes the input data to a mean-centered distribution before training. After training, RevIN uses a reverse denormalization procedure to restore the output distribution to its original state. The second method, DAIN <cit.>, employs nonlinear neural networks for dynamic adaptive normalization of the data series.

We also introduce the Feature Adding Method (FAM) parameter to determine whether to include additional timestamp embedding features. Time series data often exhibit seasonal and periodic variations, making it beneficial to derive predictive insights from timestamps. This kind of feature has been widely used in various transformer-based models, such as Transformer <cit.>, Informer <cit.>, and FEDformer <cit.>.









The Feature Extraction Methods (FEM) parameter plays a crucial role in enhancing the learning efficacy of TSF models. Within our AutoPV search space, we incorporate a diverse array of feature extraction methods from well-established non-transformer-based TSF models to process the input data, thereby obtaining distinct high-level information. These feature extraction methods have demonstrated their efficiency and have been adapted and evolved within certain transformer-based architectures, leading to the creation of novel structures. The methodologies corresponding to the entries in Table <ref> are introduced as follows:











  * 𝐅𝐄𝐌_1: This method leaves the data in its original state without applying any modifications.

  * 𝐅𝐄𝐌_2: This method employs a single linear layer to embed the input data. It is a common preprocessing step before introducing data into deep learning architectures such as LSTM, representing one of the simplest strategies for extracting multidimensional information.

  * 𝐅𝐄𝐌_3: This method performs a seasonal-trend decomposition on the original time series, segregating it into trend and seasonal components that represent the long-term progression and cyclic patterns of the series, respectively. Such a technique is extensively utilized in several prominent models, including Autoformer <cit.>, FEDformer <cit.>, and DLinear <cit.>. Following their methodologies, we apply an additional FC layer to process the trend component, while the seasonal component is handled by the core predictive structure selected through our search. The outputs of these two components are then aggregated. Fig. <ref> shows the complete process.


  * 𝐅𝐄𝐌_4: Unlike the FEM_3 approach, which employs a fixed-size moving average window for data decomposition, FEM_4 introduces a multi-scale hybrid decomposition block that employs multiple moving windows with varying kernel sizes. This approach, inspired by MICN <cit.>, aims to capture a broader and more detailed range of patterns within the time series.


  * 𝐅𝐄𝐌_5: Inspired by the TSMixer model <cit.>, FEM_5 employs the MLPs on both the time and feature dimensions to integrate information. The similar concept has also been applied in the iTranformer <cit.>.





  * 𝐅𝐄𝐌_6: FEM_6 is similar to FEM_5, with the primary distinction being its application of MLPs in the frequency domain. Specifically, FEM_6 utilizes a frequency channel learner to model inter-series dependencies through a frequency-domain MLP applied across channel dimensions. Concurrently, a frequency temporal learner models temporal dependencies by applying a frequency-domain MLP across time dimensions. This method is from the FreTS model <cit.>.

























Stage 3: Model Construction Stage. In this stage, we focus on selecting a series of parameters to build the core predictive structure. We incorporate four of the most representative deep learning architectures commonly used in PVPF tasks: LSTM, MLP, CNN, and TCN. The choice of architecture is determined by the Core Predictive Structure (CPS) parameter. Additionally, the detailed structure is defined by the Layer Number (LN) and Hidden Size (HS), which are also essential components of our search space. It is important to note that within this framework, a 1-layer MLP is conceptualized as a linear layer.


Our AutoPV framework introduces a minor modification to the typical TSF architecture. As illustrated in Fig. <ref>, conventional TSF models typically generate the future sequence for all input features, with the sequence corresponding to PV power subsequently extracted as the final forecast result. Instead of directly extracting the PV power sequence, AutoPV incorporates an FC layer at the end to aggregate the forecasting results of all features. This design choice is motivated by the nature of the PVPF task, where all input features contribute to PV power generation. By aggregating information from all features, the AutoPV framework can better leverage the interdependencies between the features to produce a more accurate PV power forecast.






Stage 4: Model Training Stage. This stage involves the calibration of several standard training parameters, including the Learning Rate (LR), Optimization Function (OF), and Batch Size (BS). To improve the search efficiency of the overall architecture, we have constrained the search range for these parameters.





Other Hyper-parameters. In addition to the parameters such as layer number and hidden size discussed earlier, there are other hyper-parameters specific to certain modules in Stage 2 and Stage 3. To maintain the universality of the search space and enhance the efficiency of the search process, we assign standardized fixed values to these hyper-parameters. The details of all fixed hyper-parameters are provided in the Appendix.













 §.§ Architecture Construction Pipeline






Fig. <ref> illustrates the construction of the entire architecture using the 12 parameters defined within the search space. In the initial stage, the parameters FSM and FST execute an input selection operation across the feature dimensions. Subsequently, the DGM, SM, and FAM parameters implement a suite of data augmentation operations, including the addition of noise, data stationarization, and the incorporation of supplementary timestamp embedding features.


The next step involves the FEM parameter. Among the included methods, 𝐅𝐄𝐌_3 and 𝐅𝐄𝐌_4 decompose the original time series data into two separate streams, whereas 𝐅𝐄𝐌_5 and 𝐅𝐄𝐌_6 output a single data series with a modified shape. Taking 𝐅𝐄𝐌_3 as an example in Fig. <ref>, the seasonal component is fed into the core predictive structure which is constructed in the subsequent stage, while the trend component is processed by an FC layer. This processing method is based on prior researches that considers the trend component to exhibit relatively simple patterns that do not necessitate a complex model for processing. The outputs of these two components are then aggregated to reconstruct the original sequence. For other feature extraction methods that yield a single processed output, we treat it similarly to the seasonal component in Fig. <ref> and omit the flow path for the trend component.

The core predictive structure is defined by the CPS, LN, and HS parameters. It is a straightforward deep learning model that accepts the seasonal component as input. As mentioned earlier, an additional FC layer is added to aggregate the forecasting results of all input features. The generated seasonal prediction result then is added with the trend prediction as the final output. To enhance clarity, all non-essential processes are depicted with dotted lines in Fig. <ref>.


Finally, the BS, OF, and LR parameters are searched and govern the overall training configuration of the architecture.



























 §.§ Evaluator



Within the NAS framework, the evaluator is responsible for determining the efficacy of the architectures identified through the search process. In our AutoPV system, we assess the performance of an architecture based on its prediction error on the validation dataset.

Specifically, after a neural architecture is constructed, it is initially trained on the training dataset. Here we employ an early stopping strategy with a patience of 3 epochs to control the training process. Subsequently, the trained model is used to perform prediction on the validation dataset. The MAE across the entire validation dataset constitutes the final evaluation score, with a lower value indicating a superior architecture. This metric will be referred to as Measured Error in subsequent sections.

During this evaluation process, a hash map is maintained to prevent the redundant assessment of identical structures. This helps improve the efficiency of the overall NAS framework by avoiding unnecessary re-evaluations.


Additionally, given the growing interest in device-dependent NAS <cit.>, our evaluator also calculates and optimizes the number of parameters within the searched architecture. This aspect will be elaborated upon in the experimental section.






 §.§ Search Strategy







Architecture Encoding Method. The majority of existing NAS research treats the search space as a graph, employing an adjacency matrix representation or path encoding <cit.> to characterize different architectures. This is because these NAS frameworks are primarily tailored to CV tasks, which necessitate the flexible splicing, repetition, and skipping of modules. However, such encoding methods are meaningless to our search space, as the modules for TSF tasks cannot be arbitrarily combined as in visual tasks. Moreover, the option to omit operations (denoted as 'None' in the table) is already incorporated as a selectable module in specific stages, obviating the need to accommodate architectures of dynamic lengths. 

Based on these considerations, we adopt a one-hot encoding sequence with a fixed length to represent each architecture. This encoding method enhances the efficiency of the search process and facilitates the inclusion of non-architectural hyperparameters such as learning rate and batch size <cit.>.




































            

		




MoBananas Search. To evaluate the effectiveness of the constructed search space, we employ the MoBananas algorithm <cit.> in this work to search for the optimal architecture while simultaneously maintaining a relatively small model scale. It is a multi-objective optimization variant of the Bananas algorithm <cit.>, which utilizes Bayesian optimization to conduct an architecture search.




The MoBananas algorithm is characterized by three key innovative components. First, throughout the process of searching and evaluating different architectures, MoBananas trains a Deep Neural Network (DNN) predictor f̂, which is capable of forecasting the performance of new architectures. Second, MoBananas selects the top k_p parent architectures for mutation by employing non-dominated sorting <cit.>. The third innovation is the application of Thompson Sampling <cit.> by MoBananas to choose the subsequent generation from a pool of new candidate individuals. The comprehensive procedure of MoBananas Search is detailed in Algorithm <ref>. To save space, we do not repeat the whole process in the main text.










































            

		












§ EXPERIMENTS





 §.§ Experimental Setup


Datasets. The dataset used in this paper is collected from a PV array located in Daqing city, China. Each data point in the dataset contains 11 dimensions, including PV power output, solar irradiation, wind speed, humidity, and other relevant weather parameters. The data was recorded from August 1, 2022 to October 28, 2023, with a sampling frequency of 1 minute, resulting in a total of 612,059 data points. However, the original dataset contained some missing values and outliers. To address these issues, we performed the following data cleaning and completion procedures:








  * Data Deletion: We first group the data by day. The following two cases of daily recordings are considered abnormal, and the entire day's recordings will be removed: 1) the ratio of missing values and zero values in a given day exceeds 0.8; 2) the same outlier occurs ten or more times consecutively in one day.














  * Data Insertion: For the remaining days of recording in the dataset, we first set the recorded values with timestamps between 21:00 and 24:00, as well as those between 0:00 and 3:00, to 0. This is because there is no solar irradiation during these periods throughout the entire year. Next, starting from one side, for each missing data point, we apply a moving average window of size 3 to calculate the average of the surrounding non-missing values. We then use this average value to impute the missing data.






Following data cleaning and imputation, we retain a complete dataset encompassing 369 days, totaling 8,856 samples. To make the variations between data points more pronounced and to improve the meaningfulness of short-term predictions, we down-sample the original dataset to a resolution of one sample per hour. Detailed information regarding this dataset is presented in Table <ref>.

We partition the dataset in a 6:2:2 ratio. Specifically, the initial 60% of the dataset is allocated for training, the subsequent 20% for validation, and the final 20% for testing. Each training sample was segmented using a sliding window approach, with a window size of (T_s + T_p) and a step size of 1. We ensured that the (T_s + T_p) values within each sample were continuous in time by skipping any discontinuous parts during the sampling process. The first T_s data points within each window serve as the historical sequence for input, while the subsequent T_p data points represent the actual future PV power output, serving as the ground truth.




























 §.§ Evaluation Metrics




Two evaluation metrics are employed in the experiments: MAE as shown in Eq. <ref> and weighted mean absolute percentage error (WMAPE) as shown in Eq. <ref>. The reason why we do not use the standard mean absolute percentage error (MAPE) metric is that the PV power dataset contained numerous small values, which would lead a to a significant impact on the evaluation result. In these two equations, Ŷ represents the prediction results, and 𝐘 denotes the ground truth.







    MAE(Ŷ, 𝐘) = 1/N∑_n=1^N| ŷ_n - y_n |














    WMAPE(Ŷ, 𝐘)    = ∑_n=1^N w_n| ŷ_n - y_n |/∑_n=1^N w_n | y_n |
    
    w_n    = | y_n |




 §.§ Performance Evaluation



The effectiveness of the AutoPV framework is evaluated on both PVPF Task 1 and PVPF Task 2, as described in Section <ref> and illustrated in Fig. <ref>. For each task, six different sub-tasks with varying forecasting lengths are defined: 12-step (half-day) ahead forecasting, 24-step (one-day) ahead forecasting, 48-step (two-day) ahead forecasting, 72-step (three-day) ahead forecasting, 168-step (one-week) ahead forecasting, and 336-step (one-month) ahead forecasting. To facilitate the discussion, we will refer to these sub-tasks as DQH-12, DQH-24, DQH-48, DQH-72, DQH-168, and DQH-336, respectively. The historical data sequence length for all sub-tasks is fixed at 96.



Search Efficiency Visualization. In our search experiments, all hyperparameters for the MoBananas algorithm are set as follows: k_ini = 10, k_p = 10, p_m = 0.2, k_m = 5, k_l = 10, T_max = 120. Each search task is executed on an NVIDIA RTX 3090 GPU. Fig. <ref> illustrates the entire architecture search process, where the ordinate represents the measured error of the optimal model at each iteration. The results indicate that most search tasks could converge at around 70 iterations. By comparing the search results of PVPF Task 1 and PVPF Task 2, we observe that incorporating future weather information significantly reduces the measured error of the predictive architecture.



Fig. <ref> presents the average search time per iteration and the total wall time required to find the optimal architecture across different PVPF tasks. Due to the increased input size from integrating future weather data, PVPF Task 2 requires more time for each iteration of search and evaluation. On average, the AutoPV system takes 6.42 hours to construct an optimal forecasting architecture for a specific task. This time can be further reduced in practice by increasing computational resources and reducing the dataset size.




In the Appendix, we also demonstrate the multi-objective optimization process for different search tasks. In each iteration, the system constructs an optimal architecture set, known as the Pareto Frontier, which includes architectures that either exhibit good forecasting performance or have a small model scale. The results show that as the search progresses, the optimal architectures consistently move towards smaller scales and higher performance. In practical applications, a suitable architecture can be selected from the final Pareto Frontier based on the constraints imposed by the computing hardware.




































Baselines Comparison.  
In this experiment, we select the architecture with the highest performance score (i.e., the lowest measured error) from the final Pareto frontier of each task as the result architecture of AutoPV for evaluation. This selected architecture is then compared with 13 well-recognized baseline forecasting models, which include: (1) Basic PVPF deep learning models: LSTM, MLP, CNN, and TCN; (2) Transformer-based models: Transformer, Autoformer, Informer, FEDformer, iTransformer and PatchTST; (3) Non-transformer-based models: DLinear, TimesNet and TSMixer. The implementation of the models refers to the Time Series Library <cit.>.



For all evaluated benchmarks, the learning rate is set to 0.001, the optimizer used is Adam, and the batch size is 64. Each of the four basic deep learning models has 3 layers with a hidden size of 512. All transformer-based models also have 3 layers, consisting of 2 encoder layers and 1 decoder layer, with a model dimension fixed at 512. It is important to note that some TSF models, particularly the transformer-based models, require setting a portion of the historical sequence as the labeled sequence. For these models, the length of the labeled sequence is set to 48.























For each task on each model, we fix a global random seed and conduct the training-evaluation experiments five times. The evaluation results including 
average values and the standard deviations are presented in Table <ref> for PVPF Task 1 and Table <ref> for PVPF Task 2.


1) Analysis of PVPF Task 1. For PVPF Task 1, as shown in Table <ref>, in the first five sub-tasks, the AutoPV system successfully finds an optimal architecture that outperforms other predefined models. In the remaining task DQH-336, It is only surpassed by iTransformer. This demonstrates that within the current search space, the AutoPV system is able to construct effective architectures for PVPF tasks with varying forecasting lengths. Based on the average MAE across six sub-tasks, we select the top-performing model from each benchmark category: CNN, PatchTST, and DLinear. When compared to these three models, the AutoPV system demonstrates relative performance improvements in MAE of 9.02%, 7.21%, and 4.52%, respectively.


Additionally, it is evident that non-Transformer-based models generally outperform most Transformer-based models, especially in short-term and medium-term prediction tasks. This help to suggest that constructing the PVPF search space based on non-Transformer models is reasonable and feasible. Regarding the long-term forecasting tasks (DQH-168 and DQH-336), the latest Transformer variants, PatchTST and iTransformer, begin to show performance advantages, likely because transformer-based models are better at handling of long sequences. Despite this, the AutoPV system still achieves relatively close performance by combining some simple non-transformer elements.






2) Analysis of PVPF Task 2. Table <ref> presents the forecasting results in PVPF Task 2 where the future weather information is incorporated. The results demonstrate that the weather data appending method described in Section <ref> significantly enhances the forecasting performance for most models. Notably, for DLinear and PatchTST, this data appending method does not yield performance improvements. We attribute this to the following reasons: For PatchTST, the attention mechanism between each patch (multiple tokens) is disrupted by zero padding, which affects the attention calculation for the future weather data. In the case of DLinear, since it primarily considers the periodic time series trends of each feature, appending zeros to the power feature does not introduce new trend information.



In comparison, AutoPV has the lowest forecasting error in PVPF Task 2 across all six sub-tasks. We also calculated the performance improvements of AutoPV over three top-performing models: LSTM, iTransformer, and TSMixer. The relative improvements in MAE are 9.88%, 28.79%, and 18.76%, respectively. Compared to PVPF Task 1, AutoPV shows an even greater performance advantage in PVPF Task 2.




The standard deviations in Table <ref> reveal that the forecasting performance becomes highly unstable for long-term tasks in many models, especially the transformer-based models. This instability is due to the future weather data often being noisy and becoming increasingly unreliable as the forecasting length extends. In contrast, the prediction results of AutoPV exhibit significantly greater stability than most baseline models.



3) Comprehensive Analysis. After analyzing the results in Table <ref> and Table <ref>, we observe that the predefined models TSMixer and LSTM also exhibit relatively good performance in PVPF Task 1 and PVPF Task 2, respectively. However, there are two reasons that underscore the necessity of AutoPV. Firstly, as shown in Table <ref>, TSMixer loses its performance edge in Task DQH-168 and even utterly fails in Task DQH-336. Similarly, in Table <ref>, the performance of LSTM deteriorates significantly. In contrast, AutoPV consistently maintains high-quality performance across different tasks. Additionally, prior to evaluation, it is uncertain which model will perform well. Even though TSMixer is good at PVPF Task 1, it is outperformed by other models at PVPF Task 2. This indicates that for a new specific task, evaluation experiments are still necessary to select the best model. AutoPV not only automates the process of model construction, effectively reducing labor and time costs, but also provides better-performing models, thereby increasing the accuracy of forecasts.


















Architectures Visualization. Figs. <ref> and <ref> illustrate the structures of all the searched architectures. The results indicate that for PVPF Task 1, the distribution of modules within the architectures is scattered. Different sub-tasks employ various training settings, data processing methods, and core predictive models. In contrast, for PVPF Task 2, the FEM_5 (Time & Feature Mixing) module appears with high frequency. Additionally, a 1-layer MLP or 1-layer LSTM structure is predominantly chosen. To demonstrate the effectiveness of module selection in constructing the search space, we present the module proportions of all 12 parameters in Appendix, based on statistics from the optimal architecture set (Pareto frontier) for all tasks. The result shows that, except for the OF_2 (SGD optimizer) module, all modules have a non-zero occurrence rate. In future work, we may optimize the search space by removing less frequently used modules and adding more powerful ones.














































Forecasting Performance Visualization. In Fig. <ref>, we present the forecasting results of AutoPV across all sub-tasks, both with and without future weather data utilization. The results clearly demonstrate that the periodic pattern of photovoltaic power output is effectively captured. For PV power forecasting on sunny days, the inclusion of future weather data does not significantly impact the results. While on rainy and cloudy days, forecasting PV power generation proves to be more challenging without future weather data. Even when the future weather data is contaminated with noise, it still contributes to a notable improvement in the accuracy of future PV power forecasting.























































































































































§ CONCLUSION

In this study, we propose AutoPV, an NAS-based framework designed to automatically search and construct efficient TSF models for specific PVPF tasks. We develop a novel search space that incorporates many state-of-the-art TSF and PVPF models. Utilizing a multi-objective optimization method named MoBananas, we search for the best-performing architecture while also controlling the model scale. To the best of our knowledge, this is the first work that bridges the gap between NAS and TSL tasks. Extensive experiments demonstrate the efficiency and effectiveness of AutoPV across various PVPF tasks. In future work, we plan to expand the search space and include more valuable modules to further enhance the performance of AutoPV.





§ ACKNOWLEDGMENTS

This work is supported by the RISUD Projects (No. P0042845 and No. 1-BBWW) and the RISE Project (No. P0051003) of The Hong Kong Polytechnic University, International Centre of Urban Energy Nexus of The Hong Kong Polytechnic University (No. P0047700), and KKS Synergy project, Energy flexibility through synergies of big data, novel technologies & systems, and innovative markets (20200073). It is also supported by High Performance Computing Center at Eastern Institute of Technology, Ningbo. 


























 
 




IEEEtran










§ BIOGRAPHY SECTION







 



[
    < g r a p h i c s >
]Dayin Chen is currently pursuing his Ph.D. in Building Energy and Environment Engineering at The Hong Kong Polytechnic University. He holds both a Bachelor's and a Master's degree in Computer Science from the Southern University of Science and Technology. His research interests include mobile computing, crowdsourcing, and neural architecture search. Additionally, Chen is exploring the application of computer science technologies in enhancing building energy efficiency.


[
    < g r a p h i c s >
]Xiaodan Shi received the B.E. and M.S. degrees in photogrammetry and remote sensing from Wuhan University, China. She  received the Ph.D. degree with the Center for Spatial Information Science, The University of Tokyo, Kashiwa, Japan. She is working as a researcher with the Center for Spatial Information Science, The University of Tokyo. Her current research interests include computer vision and its applications in pedestrian trajectory prediction, human flow monitoring, multi-objects tracking and data mining in GPS.



[
    < g r a p h i c s >
]Mingkun Jiang received the B.E. degree in process equipment and control engineering from Yanshan University, China. He received the M.S. and Ph.D. degrees with the School of Mechanical and Power Engineering, The East China University of Science and Technology, Shanghai, China. He is working as a research engineer with the PV Industry Innovation Center, The State Power Investment Corporation. His current research interests include renewable energy potential evaluation, PV and storage system modelling and its power output prediction.











[
    < g r a p h i c s >
]Haoran Zhang is a researcher with an interdisciplinary background, holding dual Bachelor's degrees in Engineering and Economics, complemented by two Ph.D. degrees in Oil & Gas Storageand Transportation Engineering ond Environmental science. His research is at the forefront of sustainable development, focusing on clean energy supply chains, green transportation, and smart urban energy systems. His innovative research has garnered prestigious accolades, including the Energy Globe Award, the R&D 100 Award, the John Tiratsoo Award for Young Achievement, and the Smart 50 Award among others.






[
    < g r a p h i c s >
]Dongxiao Zhang received the M.S. and Ph.D. degrees in hydrology from the University of Arizona, Tucson, AZ, USA, in 1992 and 1993, respectively. From 1996 to 2003, he was a Technical Staff Member and the Team Leader with the Los Alamos National Laboratory, Los Alamos, NM, USA. From 2004 to 2007, he was the Miller Chair Professor of petroleum and geological engineering with the University of Oklahoma, Norman, OK, USA. From 2007 to 2010, he was the Marshall Professor of the
Petroleum Engineering Program with the University of Southern California, Los Angeles, CA, USA. From 2010 to 2019, he was the Chair Professor and the Dean of the College of Engineering, Peking University, Beijing, China. He was the Provost and the Chair Professor with the Southern University of Science and Technology, Shenzhen, China. He is currently the Provost and the Chair Professor with the Eastern Institute of Technology, Ningbo, China. His research interests include stochastic uncertainty quantification and inverse modeling, mechanisms for shale–gas and coalbed–methane production, and geological carbon sequestration. Dr. Zhang is a member of the U.S. National Academy of Engineering, a fellow of the Geological Society of America, and an Honorary Member of the Society of Petroleum Engineers.



[
    < g r a p h i c s >
]Yuntian Chen is an assistant professor at Eastern Institute of Technology, Ningbo. He received the B.S. degree from Tsinghua University, Beijing, China, in 2015, the dual B.S. degree from Peking University, Beijing, China, in 2015, and the Ph.D. degree with merit from Peking University, Beijing, China, in 2020. His research field includes scientific machine learning and intelligent energy systems. He
is interested in the integration of domain knowledge and data-driven models.




[
    < g r a p h i c s >
]Jerry Yan is a member of the European Academy of Sciences and Arts and the Fellow of Hong Kong Academy of Engineering Sciences, currently serving as a Chair-Professor at the Hong Kong Polytechnic University. With a PhD from the Royal Institute of Technology (KTH), he has held chair professor positions at Luleå University of Technology, Mälardalen University, and KTH, Sweden. Prof. Yan's research focuses on renewable energy, advanced energy systems, climate change mitigation, and environmental policies. He has a publication record of over 500 papers in renowned journals such as Science, Nature Energy, Nature Climate Change etc and holds more than 10 patents. Having supervised nearly 200 post-doctoral researchers and 50 doctoral candidates, Prof. Yan has secured substantial external grants exceeding 20 million Euros. Prof. Yan's contributions have been acknowledged through prestigious awards, including the Global Human Settlements Award of Green Technology, the EU Energy Islands' Award, Research2Business Top100 and the IAGE Lifetime Achievement Award.





 Paper Title: Contrastive Factor Analysis.




Ge-based Clinopyroxene series: first principles and  experimental local probe study
    João P. Araújo
    August 5, 2024
===================================================================================




§ ABSTRACT 
Factor analysis, often regarded as a Bayesian variant of matrix factorization, offers superior capabilities in capturing uncertainty, modeling complex dependencies, and ensuring robustness.
As the deep learning era arrives, factor analysis is receiving less and less attention due to their limited expressive ability. 
On the contrary, contrastive learning has emerged as a potent technique with demonstrated efficacy in unsupervised representational learning.
While the two methods are different paradigms, recent theoretical analysis has revealed the mathematical equivalence between contrastive learning and matrix factorization, providing a potential possibility for factor analysis combined with contrastive learning.
Motivated by the interconnectedness of contrastive learning, matrix factorization, and factor analysis, this paper introduces a novel Contrastive Factor Analysis framework, aiming to leverage factor analysis's advantageous properties within the realm of contrastive learning. 
To further leverage the interpretability properties of non-negative factor analysis, which can learn  disentangled representations, contrastive factor analysis is extended to a non-negative version.
Finally, extensive experimental validation showcases the efficacy of the proposed contrastive (non-negative)  factor analysis methodology across multiple key properties, including expressiveness, robustness, interpretability, and accurate uncertainty estimation.





§ INTRODUCTION

Before the era of deep learning, factor analysis (FA) <cit.> had received widespread attention and achieved impressive results.
Generally, FA is a statistical method used to identify underlying factors or latent variables that explain
patterns of correlations among observed variables. 
With its stochastic latent variables, FA often enjoys good properties for modeling complex dependence and enhancing robustness <cit.>. 
Meanwhile, its non-negative extensions, such as Poisson factor analysis <cit.>, can not only learn meaningful latent representation from high-dimension observation unsupervised but also can discover interpretable semantic concepts 
 <cit.>. 
With these appealing properties, FA and non-negative FA have been widely used in many applications, including text analysis <cit.> image representation learning <cit.>. 

With the deep learning era arriving <cit.>, FA is receiving less and less attention due to its limited expressiveness. 
Correspondingly, recent years have witnessed a surge of interest in self-supervised learning, such as contrastive learning (CL) <cit.>, owing to its remarkable capacity to learn high-quality representations unsupervised, leading to notable advancements in various fields <cit.>.
Generally, CL methodologies aim to extract discriminative features by training models to distinguish positive and negative samples. 
The distinction between FA and CL lies primarily in the latter's ability to harness effective training objectives <cit.> and expressive deep neural networks <cit.>. 


Given that FA and CL are regarded as two distinct research fields, integrating CL to enhance FA or vice versa might appear challenging. Fortunately, recent advancements in the theoretical underpinnings have provided insights into this integration.  Specifically, <cit.> leverages spectral contrastive loss to derive the mathematical equivalent of CL and matrix factorization (MF) <cit.>.

Considering the relationship between MF and FA <cit.>, the integration of FA with CL becomes feasible. 
Beyond feasibility, merging FA and CL offers numerous attractive potential benefits. 
In general, from the perspective of FA, enhancing FA with CL can alleviate its limitation in expressiveness, leading to richer and more informative representations <cit.>.
From the perspective of CL, considering FA's capacity to capture complex dependencies and its stochastic latent variables, incorporating CL with FA not only enhances performance but also facilitates improved uncertainty estimation and interpretability <cit.>.










 
Drawing inspiration from the mathematical equivalence established between MF, CL, 
and FA, as illustrated in Figure. <ref>, in this paper, we try to develop a Contrastive Factor Analysis (CFA) framework. 
Unlike FA, which directly models observation data, CFA follows CL methods to model the relation matrix, which represents the relationship between positive and negative samples.
More specifically, in CFA, the relation matrix is factorized into Gaussian latent variables. 
Considering the target matrix is non-negative, the CL can also be modified as a contrastive non-negative FA, which factorizes the relation matrix into gamma latent variables.
Compared with CFA, its non-negative extension can enjoy good properties for learning disentangled representation <cit.>. 

To approximate the posterior distribution of the latent variables, we further develop a variational inference network.
While it's straightforward to employ the Gaussian distribution to appropriate the posterior of Gaussian latent variables,  considering the non-reparameterizability of the gamma distribution, we instead utilize the Weibull distribution to approximate the posterior of gamma latent variables <cit.>. 
Finally, both models' parameters can be effectively trained within a variational inference framework <cit.> with reparameterize technique. 


To assess the properties of contrastive factor analysis, we conducted a series of experiments targeting various aspects. 
These properties encompass the attainment of more expressive representations, robustness against out-of-distribution data, interpretability, and uncertainty evaluation. 




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



  * 
Drawing inspiration from the interplay between contrastive learning, matrix factorization, and factor analysis, we introduce a novel contrastive (non-negative) factor analysis framework. 


  * To approximate the posterior distribution of latent variables, we propose two variational inference networks: Gaussian and Weibull variational inference network.


  *  Extensive experimental validation showcases the efficacy of the proposed contrastive factor analysis methodology across multiple key properties, including expressiveness, robustness, interpretability, and accurate uncertainty estimation.





§ PRELIMINARY ON CONTRASTIVE LEARNING



In this section, we briefly review the standard contrastive learning modules and the relationship between contrastive learning and matrix factorization.


  
Contrastive learning. 
In representation learning, the aim is to learn an encoder function  f: ℝ^D→ℝ^d that maps high-dimensional input data  x ∈ℝ^D to low-dimensional representations z ∈ℝ^d. 
To perform contrastive learning, we start by building pairs of positive samples ( x,x_ + ) by augmenting the same natural sample x̅∼𝒫( x̅),

where 𝒜( ·|x̅) represents the augmentation distribution. Additionally, negative samples x̅ are independently augmented samples, following the marginal distribution 𝒫( x ) = 𝔼_x̅𝒜( x |x̅).

The objective of contrastive learning is to encourage the model to map positive samples closer together in the feature space while pushing negative samples further apart.
A well-known CL objective is the InfoNCE loss <cit.>. 









    0.95ℒ_(f ) = -𝔼_x, x_+, {x_i^ - }_i = 1^Mlogexp( f( x )^Tf( x_ + ))/exp( f( x )^Tf( x_ + )) + ∑_i = 1^M exp( f( x )^Tf( x_i^ - ))


where {x_i^ - }_i = 1^M are M negative samples independently drawn from 𝒫(x). <cit.> proposed the spectral contrastive loss that is more amenable for the theoretical analysis: 

    0.95ℒ_(f) =  - 2𝔼_x,x + f( x )^Tf( x_ + ) + 𝔼_x,x + ( f( x )^Tf( x_i^ - ))^2





  
Contrastive learning as matrix factorization. 
Notably, <cit.>
 show that the spectral contrastive loss can be equivalently written as the following matrix factorization objective: 

    ℒ_ = A̅ - FF^T^2,   F_x,: = √(P( x )) f( x )



Here, A̅ = D^-1/2 A D^-1/2 denotes the normalized version of the co-occurrence matrix A ∈ℝ_+^N × N of all augmented samples x ∈𝒳 (assuming |𝒳| = N for ease of exposition): ∀ x, x' ∈𝒳, A_x, x' := 𝒫(x, x') = 𝔼_x̅[ 𝒜(x|x̅) 𝒜(x'|x̅) ].


 



§ CONTRASTIVE (NON-NEGATIVE) FACTOR ANALYSIS

This section furnishes a comprehensive overview of the proposed contrastive (non-negative) factor analysis, comprising the generative model  (Sec.<ref>) and the variational inference network (Sec.<ref>).
It further elaborates on variational inference (Sec.<ref>) and discusses techniques for uncertainty evaluation (Sec. <ref>). 


 §.§ Bayesian generative model



In classical FA, an observation matrix   X 
is factorized into two latent variables X ∼(θ Φ), where both θ and Φ are drawn from distributions such as Gaussian.
While FA is intuitive for feature learning, it may not be effective for learning meaningful representations with its reconstruction training objects <cit.>. 
Motivated by the connection between CL and MF, as indicated by Eq. <ref>, we build contrastive factor analysis (CFA) by instead modeling the co-occurrence matrix A̅, which represents the relationship among positive and negative samples. 



Unlike classical FA, which focuses solely on modeling the generative process between the latent variable θ and target data X, CFA requires modeling the generative process among observation data samples X, the latent variable θ, and the corresponding normalized co-occurrence matrix A̅, as depicted in Figure. <ref>.
To address this, we draw inspiration from conditional VAE <cit.> to devise a modified FA. 
Specifically, given data samples and augmentations X ∈ℝ^N × D and the corresponding normalized co-occurrence matrix A̅∈ℝ+^N × N, CFA factorize A̅∈ℝ+^N × N under the Gaussian likelihood as follows:

    A̅_i,j x_i,x_j∼𝒩(  θ_i θ_j), θ_i x_i∼𝒩( f(x_i),   1 ), θ_j x_j∼𝒩( f(x_j),   1 ).

where θ_i,  θ_j  ∈ℝ_+^d is the factor score
matrix, each column of which encodes the relative importance of each atom in a sample. 

Further, recent studies suggest that modeling latent features using non-negative vectors can harness the advantageous properties of non-negative matrix factorization (MF), which facilitates the acquisition of  disentangled representations <cit.>.
Hence, we draw inspiration from relevant research <cit.> to develop a contrastive non-negative FA (CNFA) approach.
Specifically, the CNFA factorizes A̅ into sparse and non-negative gamma latent variables, which can be described as follows:

    A̅_i,j x_i, x_j∼𝒩(  θ_i θ_j), θ_i x_i∼( f(x_i), 1 ), θ_j x_j∼( f(x_j), 1 ).


In our formulation, the ideal scenario involves modulating the prior of the latent variables θ_i by the input x_j. However, we can readily relax this constraint to ensure that the latent variables are statistically independent of the input variables <cit.>. Thus, for contrastive factor analysis (CFA), we simply define the data-independent prior as f( x_i) = 0, while for contrastive non-negative factor analysis (CNFA), we set f(x_i) = 1.









 §.§ Reparameterizable latent variable distributions






  
Variational Gaussian posterior for CFA. 
The Gaussian distribution θ∼𝒩(μ, σ^2) has probability density function f( θμ ,σ ^2) = 1/√(2σ ^2π)e^ - ( x - μ)^2/2σ ^2, where θ∈ℝ. It is reparmeterizable as drawing 
θ∼𝒩(μ, σ^2) us equivalent to letting θ = g(ε) := μ + εσ ,   ε∼𝒩(0,1). The KL divergence is analysis as ( 𝒩( μ _1,σ _1^2) || 𝒩( μ _2,σ _2^2)) = logσ _2/σ _1 + σ _1^2 + ( μ _1 - μ _2)^2/2σ _2^2.
As shown in Fig. <ref>, the variational inference network utilize the obtained latent features h_i from a deep feature extractor, such as ResNet <cit.>, to 
construct the variational posterior:

    q(θ_i  x_i ) = 
    	𝒩(μ_i, σ^2_i ), 
    	μ_i = f_μ(h_i),
    	σ^2_i = (f_σ( h_j)),
where f_·^(l)(·) denotes the neural network, and  applies (1+(·)) non-linearity to each element to ensure positive Gaussian variance parameters. 
It's evident that the CFA will reduce to matrix factorization 
as σ of the Gaussian distribution goes to zero. Therefore, the proposed CFA can be viewed as a generalization of
vanilla contrastive learning <cit.>.



  
Variational Weibull posterior for CNFA: 
While the gamma distribution is ideal for the posterior due to its encouragement of sparsity and satisfaction of non-negativity constraints, its non-reparameterizability and incompatibility with gradient descent optimization lead to the adoption of the Weibull distribution as an approximation for the posterior of gamma latent variables.

The reason for choosing the Weibull distribution is threefold <cit.>:  ) the Weibull distribution is
similar to a gamma distribution, capable of modeling sparse,
skewed and positive distributions;
) the Weibull distribution has a simple
reparameterization so that it is easier to optimize; That is, to sample θ∼( k,λ) with probability density function p( θ k,λ) = k/λ ^kθ ^k - 1e^ - ( S/λ)^k, where θ∈ℝ_+, it is equivalent to letting 
θ = g(ε) := λ( - ln (1 - ε ))^1/k,  ε∼(0,1).
) The KL divergence from the gamma to Weibull distributions has an analytic expression as
( (k,λ)  || (α,β) )  = γα/k - αλ+ 
k + βλΓ(1+1/k) - γ -1 -α(β) + Γ(α).
where γ is the Euler-Mascheroni constant and  Γ is the gamma function. This provides an efficient way to estimate the training
objective, which we will discuss in detail in the next subsection.
Similarly, the variational inference network employs the obtained latent features _i to 
construct the variational posterior:










    0.95q(θ_i | x_i ) = Weibull(k_i, λ_i ), 
        k_i = Softplus(f_k(h_i)), 
        λ_i = ReLU(f_λ(h_i))/Γ(1 + 1 /k_i )


where f_·(·) denotes the neural network, and  applies (1+(·)) non-linearity to each element to ensure positive Weibull shape and scale parameters. 
With above inference network, there are 𝔼( θ) = ( f_λ( h_i)).
It's evident that the CNFA will reduce to non-negative matrix factorization if the parameter k of the Weibull distribution goes to infinity.
Therefore, the proposed CNFA can be viewed as a generalization of
vanilla non-negative contrastive learning <cit.>.




 §.§ Variational inference


For the contrastive factor analysis, given the model parameters referred to as W, which consist of the parameters in the generative model and inference network, marginal likelihood is defined as:

    0.95p ( A̅W)  = ∫∏_i = 1^N ∏_j = 1^N p( A̅_i,j  |  θ_i, θ_j)∏_j = 1^N p( θ _j   |   x _j) d θ _j = 1^N .

The inference task is
to learn the parameters of the generative model and the inference network.
Similar to VAEs,
the optimization objective of contrastive factor analysis can be achieved by
maximizing the evidence lower bound (ELBO) of log-likelihood as

    0.95ℒ(A̅)   =  ∑_i = 1^N ∑_j = 1^N 𝔼_q(θ_j  x_j)[ ln p(A̅_i,j  |   θ _i, θ _j)]  -  ∑_j = 1^N 𝔼_q(θ_j  x_j)[ lnq(θ _j  |   x_j)/p(θ _j) ] .
 
where the first term is the expected log-likelihood of the generative model, which ensures reconstruction performance, and the second term is the Kullback–Leibler (KL) divergence that constrains the variational distribution q(θ_j) to be close to its prior p(θ_j).
Unfortunately, Equation. <ref> becomes intractable with factor analysis algorithms as the input space X grows exponentially large with N →∞. To address this challenge, we reformulate the factor analysis problem (Equation. <ref>) equivalently as a tractable contrastive learning (CL) objective that only necessitates sampling positive and negative samples from the joint and marginal distributions, respectively:

    0.95ℒ = 
        - 2𝔼_x,  x + (θ_+^T  θ_+) 
     + 𝔼_x,   x_+( θ_-^T θ_-)^2
       -  𝔼_q(θ_+ x_+)[ lnq(θ _+  |   x_+)/p(θ _+) ] 
    -  𝔼_q(θ_- x_-)[ lnq(θ _-  |   x_-)/p(θ _-) ]
 
 The parameters in contrastive (non-negative) factor analysis can be directly optimized by advanced gradient algorithms, like Adam <cit.>.



  
Complexity analysis 
 Our framework is computationally and memory-efficient because only a few parameters need to be added to infer the distribution parameters. 
 The extra memory cost comes from the posterior inference network, which only adds d × 1 parameters due to σ_i^2 ∈ℝ^1 and k_i∈ℝ^1, where d is the dimension of feature h_i.
 Meanwhile, the additional computations involve the sampling process and computing the KL term, which is of scale O(d). 
 Both of above
 is inconsiderable compared to the memory and computational scale of the deep backbone neural network <cit.>. 



 §.§ Uncertainty-aware representation
 
To quantify the inferred latent variable uncertainty, we adopt Shannon's entropy <cit.> of the variational posterior. Specifically, The Shannon's entropy of Gaussian variational posterior with parameters (μ, σ) is defined as: 

    0.95H ( q_ϕ( θ _ix_i)) = ∑_m = 1^d ln( σ _m√(2π e))
. And Shannon's entropy of the Weibull variational posterior with parameters (λ, k) is defined as: 

    0.95H( q_ϕ( θ _ix_i)) = ∑_m = 1^d ( k_m - 1)γ/k_m + lnλ _m/k_m + 1

where γ stands for Euler's constant. 
It should be noted that the higher the entropy, the higher the uncertainty.



§ RELATED WORK

Our work is mainly related to factor analysis and probabilistic representation learning. Here, we introduce the related research.

Factor analysis. 
Before the era of deep learning, factor analysis was widely used in representation learning. And its non-negative extension <cit.> has achieved great progress in text analysis. Further, there is increasing interest in hierarchical factor analysis, which is aimed at exploring hierarchical data representation. 
But all this work is different from ours,  which mainly focuses on improving factor analysis with contrastive learning. 
Besides, <cit.> proves contrastive learning is capable of revealing latent variable posterior information in FA. 
Nonetheless, our approach diverges from this earlier work as it places more emphasis on proposing a novel model that can effectively learn representations.

Probabilistic representation learning. 
Probabilistic representation learning has indeed made significant strides in recent years and has garnered considerable attention in the machine learning community. 
Its ability to capture uncertainty and model complex data distributions has proven immensely useful in various applications such as natural language processing <cit.>, computer vision <cit.>, and graph processing <cit.>.
In contrast to the previous efforts, this paper develops a general framework for (non-negative)  contrastive factor analysis that benefits from both (non-negative) factor analysis and contrastive learning. 
Besides, there are efforts to improve contrastive learning with the Bayesian method, which is aimed at evaluating the uncertainty <cit.>. 
Unfortunately, most of these techniques rely on a decoder, which frequently degrades the performance of downstream tasks <cit.>. 
On the opposite hand, our research is an encoder-only approach that can enhance downstream tasks' performance. 


§ EXPERIMENT



 §.§ Feature learning



  
Baselines & Datasets. 
To evaluate the expressiveness of the proposed model, we compare it with strong baselines, including SimCLR(CL) <cit.> and non-negative contrastive learning(NCL) <cit.>. 
More specifically, we compare CFA with CL and CNF with NCL.
It should be noted that CFA is built on SimCLR codebase [<https://github.com/vturrisi/solo-learn>] with a minor modification to the hidden projector output.
We evaluate linear probing and fine-tuning performance on three benchmark datasets: CIFAR-10, CIFAR-100<cit.>, and Imagenet-100<cit.>. 
Apart from in-distribution evaluation, we also compare models on three out-of-distribution datasets, including stylized ImageNet<cit.>, ImageNet-Sketch<cit.>, ImageNet-C<cit.>(restricted to ImageNet-100 classes). 
For all experiments, we run 3 random trials and report their mean and standard deviation. 
All experiments are performed on
workstation equipped with a CPU i7-10700 and accelerated
by four GPU NVIDIA RTX 4090 with 24GB VRAM. 




  
Experiment Setting.
In our experiments, Resnet-18<cit.> is used as the backbone, and the model is trained for 200 epochs on CIFAR-10, CIFAR-100, and 100 epochs for ImageNet-100. 






  
In distribution experiment results. The experiment results are  exhibited in Table <ref>. 
Firstly, compared to the FA model, CFA has achieved a significant improvement, primarily attributed to the effective contrastive learning training strategies and corresponding expressive deep neural networks.
Secondly, acting as the Bayesian extension of contrastive learning (CL) and non-negative contrastive learning (NCL), respectively, CFA and CNFA demonstrate varied degrees of improvement across different datasets. This result indicates the efficacy of stochastic latent variables in capturing complex data dependencies and enhancing overall performance <cit.>.
Finally, when compared with CFA, CNFA consistently exhibits superior performance across all datasets, as NCL performs better than CL. This improvement may stem from the ability of non-negative stochastic latent variables to learn  disentangled representations <cit.>.



Out-of-distribution experiment results. Apart from the expressiveness evaluation, we further evaluate the model's robustness on the out-of-distribution dataset, and the results are depicted in Table <ref>. 
Firstly, compared to deterministic contrastive learning methods, CFA and CNFA show significant improvement in out-of-distribution generalization downstream tasks.
The improved OOD performance can be partially attributed to the generalization introduced by the stochastic latent variables, which allows the model to handle unseen and diverse data better <cit.>.
By inheriting the advantages of non-negative factor analysis, CNFA can learn disentangled representation <cit.> and exhibit better transfer ability compared with CFA.





 §.§ Uncertainty evaluation

In addition to their enhanced expressiveness and robustness on out-of-distribution (OOD) data, another notable strength of CFA and CNFA is their ability to capture data uncertainty.


Evaluation Metric. 
 With the entropy of the testing result and a given entropy threshold, we can determine whether the model is certain or uncertain about one prediction.  
 To evaluate the
uncertainty estimates, we use Patch Accuracy vs Patch Uncertainty (PAvPU) <cit.>, which is
defined as PAvPU = (n_ac + n_iu)/(n_ac + n_au + n_ic + n_iu), where n_ac, n_au, n_ic, n_iu  represent the numbers of accurate and certain, accurate and uncertain, inaccurate and certain, and inaccurate and uncertain samples, respectively. Compared to Top-1 Accuracy= (n_ac + n_au)/(n_ac + n_au + n_ic + n_iu),
 PAvPU metric would be higher if the model tends to
generate accurate predictions with low uncertainty and inaccurate predictions with high uncertainty leading to a more trustworthy algorithm.


Evaluation Setting.
Firstly, utilizing Equation <ref> and Equation <ref>, we calculate the information entropy of each sample in the test set. Subsequently, we select the M samples with the highest information entropy as uncertain samples, while the remaining samples are considered definite. In practice, we conducted experiments with different values of M, specifically selecting 100, 200, and 400 samples as uncertain samples to compute the test indices.





Experiment results:
Table. <ref>  show the PAvPU results of CFA and CNFA on ImageNet-100. CFA and CNFA achieves better performance on PAvPU metrics compared to Top-1 accuracy. 
This suggests that CFA and CNFA are effective not only in making correct predictions but also in identifying when it is likely to be wrong, which is crucial for applications requiring high reliability. 

 


Qualitative analysis:
To further validate the efficacy of the proposed model in uncertain evaluation, we conducted qualitative analysis. Specifically, we selected 20 images with the highest and lowest uncertainty predicted by CNFA from the ImageNet-100 dataset, and the experimental results are illustrated in Figure <ref>.
Figure. <ref> visually indicates that images with noisy backgrounds exhibit higher levels of uncertainty. 
In contrast, Figure.  <ref> shows that a clean background alongside high contrast is prone to displaying low uncertainty. 
This clear distinction highlights the effectiveness of CNFA in identifying the factors contributing to uncertainty in image predictions. 

Ablation study:
To further investigate the relationship between prediction uncertainty and downstream performance, we conducted an ablation study. Specifically, we sorted the samples in the test set based on their evaluation uncertainty and divided the set into five subsets accordingly. Subsequently, we calculated the average accuracy and corresponding average uncertainty for each subset. We then plotted histograms depicting the average uncertainty and accuracy for these five subsets, as illustrated in Figure. <ref>. The histogram clearly indicates a negative correlation between uncertainty and accuracy: higher uncertainty correlates with lower accuracy. It's worth noting that these findings are supported by qualitative analysis, where higher uncertainty suggests noisier and more challenging-to-classify data.










































































































    


















 


 §.§ Disentanglement representation learning

Learning a robust representation involves extracting explanatory factors that are sparse, disentangled, and semantically meaningful <cit.>. By incorporating non-negative constraints, NCL can learn disentangled features <cit.>. On the other hand, the Gamma distribution promotes sparsity due to its probability density function, which helps filter out irrelevant features and focus on the most significant ones. Thus, we evaluate CFA's performance in disentangled representation learning.


  
Evaluation Metric. 
Many disentanglement metrics like MIG <cit.> are supervised and require ground-truth factors which are often unavailable in practice. Consequently, current disentanglement evaluation is often limited to synthetic data<cit.>. To validate the disentanglement on real-world data, we adopt an unsupervised disentanglement metric SEPIN@k<cit.>. SEPIN@k measures how each feature f_i(x) is disentangled from others f_≠ i(x) by computing their conditional mutual information with the top k features, i.e., SEPIN@k= 1/k∑_i=1^kI(x, f_r_i(x)|f_≠ r_i(x)), which are estimated with InfoNCE lower bound <cit.> implemented following <cit.>.



  
Experiment Result. Table <ref> demonstrates that CNFA features show better disentanglement than CL and NCL in all top-k dimensions. The advantage is larger by considering the top features, since learned features may contain noisy dimensions <cit.>. Incorporating proper prior distribution (e.g., Gamma distribution) will effectively mitigate this problem. Thanks to the inherent properties of the Gamma distribution, CNFA tends to produce disentangled features, making it more effective in distinguishing the underlying factors of variation within the data. In conclusion, CNFA achieves better feature disentanglement on real-world data and exhibits exceptional performance in estimating uncertainty. 






§ CONCLUSION

In conclusion, this paper introduces a novel framework called Contrastive Factor Analysis (CFA), which bridges the gap between traditional Factor Analysis (FA) and the powerful technique of contrastive learning. By leveraging the mathematical equivalence between contrastive learning and matrix factorization, CFA capitalizes on the strengths of both methods, offering enhanced expressiveness, robustness, and interpretability.
Further, by extending CFA to a non-negative version (CNFA), we exploit the interpretability properties of non-negative factor analysis, enabling the learning of disentangled representations. Through extensive experimental validation, our methodology demonstrates efficacy across various dimensions, including expressiveness, robustness, interpretability, and accurate uncertainty estimation.
In essence, CFA and CNFA represent promising approaches that amalgamate the insights from factor analysis and contrastive learning, paving the way for more effective unsupervised representational learning in the deep learning era.




unsrtnat










































 Paper Title: Leaf Angle Estimation using Mask R-CNN and LETR Vision Transformer.







Leaf Angle Estimation using Mask R-CNN and LETR Vision Transformer


    Venkat Margapuri
Department of Computing Sciences 

Villanova University

Villanova, PA, USA 


Prapti Thapaliya
Department of Computing Sciences 

Villanova University

Villanova, PA, USA 


Trevor Rife
Pee Dee Research and Education Center  

Clemson University

Florence, SC, USA 



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



§ ABSTRACT 
Modern day studies show a high degree of correlation between high yielding crop varieties and plants with upright leaf angles. It is observed that plants with upright leaf angles intercept more light than those without upright leaf angles, leading to a higher rate of photosynthesis. Plant scientists and breeders benefit from tools that can directly measure plant parameters in the field; i.e., on-site phenotyping. The estimation of leaf angles by manual means in a field setting is tedious and cumbersome. We mitigate the tedium using a combination of the Mask R-CNN instance segmentation neural network, and Line Segment Transformer (LETR), a vision transformer. The proposed Computer Vision (CV) pipeline is applied on two image datasets, Summer2015-AmesULA and Summer2015-AmesMLA, with a combined total of 1,827 plant images collected in the field using FieldBook, an Android application aimed at on-site phenotyping. The leaf angles estimated by the proposed pipeline on the image datasets are compared to two independent manual measurements using ImageJ, a Java-based image processing program developed at the National Institutes of Health and the Laboratory for Optical and Computational Instrumentation. The results, when compared for similarity using the Cosine Similarity measure, exhibit ≈0.98 similarity scores on both independent measurements of Summer2015-AmesULA and Summer2015-AmesMLA image datasets, demonstrating the feasibility of the proposed pipeline for on-site measurement of leaf angles.



Fieldbook, ImageJ, Leaf Angle Estimation, Line Transformer (LETR), On-site Phenotyping, Mask R-CNN




§ INTRODUCTION

The yields and planting densities of maize in the United States have increased concurrently during the past 50 years <cit.>. A comparative analysis of U.S. commercial maize hybrids released since the 1960s revealed that by selecting high yielding hybrids under high planting densities, breeders indirectly selected hybrids with upright leaf angles (LAs) <cit.>. The discovery was that upright LAs combined with higher planting densities improve light distribution within the canopy. For instance, modern hybrids intercept 14% more light than older hybrids <cit.>. Leaf Angle Distribution (LAD) is a key parameter that describes the structure of horizontally homogeneous vegetation canopies. It is defined as the probability of a leaf element of unit size to have its normal within a specified unit solid angle <cit.>. LAD affects the manner in which incident photosynthetically active radiation is distributed on plant leaves, thus directly affecting plant productivity. The traditional technique to measure LAD involves the use of mechanical inclinometer, a precision instrument that measures the angle of slope of an object with respect to its gravity by creating an artificial horizon. The use of the mechanical inclinometer makes the process of determining LAD laborious and time-consuming. Other techniques such as 3-D digitizing of individual plant elements using specialized instrumentation <cit.> and laser scanning <cit.> are available. However, these techniques are resource-demanding. A feasible, and inexpensive alternative is the use of digital imagery to estimate LA since image analysis preserves the state of the canopy while providing enough insight into the morphometry. Image processing tools such as ImageJ <cit.> and Adobe Photoshop have built-in angle estimation tools that aid in the estimation of LA. The pitfall is that the operation of the tools is manual which means that the estimations may be subject to user bias and users may spend extended amounts of time to estimate LAs, in case of large datasets. In order to overcome the pitfalls of the manual setup, this paper proposes an automated technique for LA estimation using Mask R-CNN <cit.> instance segmentation neural network, and Line Segment Transformer (LETR) <cit.>, a vision transformer. 



§ RELATED WORK

Dzievit <cit.> conducted genetic mapping and a meta-analysis to dissect genetic factors controlling LA variation on maize. Genetic mapping populations were developed using inbred lines B73 (Iowa Still Stalk Synthetic), PHW30 (lodent, expired plant variety protection inbred), and Mo17 (Non-Stiff Stalk) that have distinct LA architectures. The leaf angles were estimated using the ImageJ image processing tool.   Herbert <cit.> described a technique for the estimation of leaf angles using stereo-photogrammetry. The technique permitted accurate measurement of leaf angle and position from several meters away and had sufficient resolution to permit the analysis of complex phenomena such as the effect of leaf shape upon interception of light and photosynthesis.   Sinoquet et al. <cit.> proposed a method to measure light interception by vegetation canopies using a 3-D digitizer and image processing software. The 3-D digitizer allowed for simultaneous acquisition of the spatial coordinates of leaf locations and orientations. The software for image synthesis also had the ability to make virtual photographs of the real canopy. The information on light interception was derived from virtual images by using the simple features of image analysis software.   Hasoi et al. <cit.> estimated the LAD of wheat canopy at different growth stages such as tillering, stem elongation, flowering, and ripening stages by using a high-resolution portable scanning lidar. The canopy was scanned three-dimensionally by optimally inclined laser beams emitted from several measuring points surrounding the canopy and 3-D point cloud images in each stage were obtained. After co-registration of lidar images between different measurement positions, leaves were extracted from the images and each leaf was divided into small pieces along the leaf-length direction. Each of the pieces was approximated as a plane, to which normal (perpendicular) was estimated. The distribution of the leaf inclination angles was derived from the angles of the normal with respect to the zenith.   Ryu et al. <cit.> examined the feasibility of seven techniques such as litterfall, allometry, LAI-2000, TRAC, digital hemispheric photography, digital cover photography, and traversing radiometer system to determine leaf area index across a 9-ha domain in an oak-savanna ecosystem in California, USA. It was shown that the combination of digital cover photography and LAI-2000 could provide spatially representative leaf area index, gap fraction and element clumping index.



§ MATERIALS AND METHODS

Researchers at Iowa State University captured the images of several plants in a field using the Fieldbook android application <cit.>. As for the current work, two datasets namely, Summer2015-AmesULA and Summer2015-AmesMLA wherein the former consists of 872 images and latter, 955 images, are used as the experimental datasets. Each of the images in either dataset shows the intersection of the leaf and stem. A sample image is as shown in Fig. <ref>.



The proposed leaf angle estimation procedure contains two broad steps: 1.	Extraction of region of interest (ROI) 2.	Estimation of Leaf Angle



 §.§ Extraction of Region of Interest (ROI)

The first step in the extraction of ROI is the detection of leaf and stem within the image. The images captured in the field contain numerous plants that look similar. Hence, it is important to determine the plant of interest and perform a curated extraction of it for the estimation of LA. Conventionally, a feasible technique for foreground extraction is OpenCV’s GrabCut <cit.> algorithm. However, the results obtained from the application of the GrabCut algorithm on each of the datasets are not satisfactory. While the algorithm captures the foreground, it also captures part of the plants in the background. The captured plants in the background are essentially noise and detrimental to LA estimation. Fig. <ref> shows the botched foreground extraction performed on a plant image from the Summer2015-AmesMLA dataset.



In order to better perform foreground extraction, the CNN model of Mask R-CNN <cit.> <cit.> is employed to extract foreground and identify the region of interest appropriately. Mask R-CNN is a deep neural network aimed to solve the problem of instance segmentation in computer vision. Instance Segmentation <cit.> is the task of precisely identifying the pixels of each of the objects in the image. It is perhaps the hardest of the computer vision (CV) tasks of classification <cit.>, semantic segmentation <cit.>, object detection <cit.>, and instance segmentation.


    
  * 175 images belonging to each of the Summer2015-AmesMLA and Summer 2015-AmesULA datasets are manually annotated using makesense.ai tool. The tool is free to use and generates annotations in the COCO format that Mask R-CNN is able to consume. A polygonal annotation is drawn where the stem and leaf meet each other. Fig. <ref> shows a polygonal annotation drawn using makesense.ai on an image with leaf and stem.

    


  * The annotated images are trained on the Mask R-CNN network for a total of 20 epochs. The loss of the network after 20 annotations is 0.034. However, the loss is not significant since Mask R-CNN is used to extract the portion of the image where the leaf and stem meet, but not to extract masks that are precise in shape and size. All the annotated images are used for training only. The dataset is not split into training, validation, and test datasets since the goal is not to evaluate the performance of Mask R-CNN, but to extract regions of interest in the image that are later processed to estimate the angle between leaf and stem. Upon training, the trained weights are used to extract the masks of each of the images in the Summer2015-AmesMLA and Summer2015-AmesULA datasets. Fig. <ref> shows the mask extracted by Mask R-CNN for the plant image shown in Fig. <ref>.




  * The predicted mask indicates the region of interest that corresponds to the leaf and stem. In order to extract the region from the original image, the 'bitwise AND' operation is applied on the mask predicted by Mask R-CNN and original image. The resulting image gives the region of interest in the original image. Fig. <ref> shows the region of interest extracted from the original image. Please note that the image is zoomed in for clarity.


    






 §.§ Estimation of Leaf Angle

The basis for the estimation of leaf angle is provided by <cit.>. The work describes the above-ear and below-ear leaf angles that are measured by plant scientists. Fig. <ref>, an excerpt from <cit.>, shows the exact position at which the angle between leaf and stem is to be measured. While leaves are typically long and tilt in different directions owing to environmental factors, the angle between leaf and stem is measured exactly where the leaf and stem conjoin.   The estimation of leaf angle is performed using Line Segment Transformer (LETR), a vision transformer <cit.>. Derived from transformers in the world of Natural Language Processing (NLP), a Vision Transformer (ViT) is a type of neural network that is used to perform CV tasks, such as image classification, and object detection. The basic principle of the transformer used for NLP is the representation of words as sequences. Drawing inspiration from it, ViT works by representing an image as a sequence of patches. A patch is a small rectangular size of the image, typically 16x16 pixels. Each patch that is extracted from the image is represented as a feature vector. The feature vectors are extracted using a convolutional neural network (CNN), such as ResNet-50 or VGG-16. The extracted feature vectors are input to an encoder network. The encoder network is a stack of self-attention <cit.> layers, where self-attention is a mechanism that allows the model to identify and establish long-range dependencies between the patches of images extracted from the original image. The encoder network outputs a sequence of vectors that may be fed into a multi-layer perceptron (MLP) for CV tasks such as classification, and detection.   The LETR vision transformer is an adaptation of the Detection Transformer (DETR). The DETR is an object detection model developed by Facebook (Meta) that leverages the transformer architecture. The ResNet architecture is typically preferred as the feature extractor within DETR to extract features from patches of an image. The DETR model follows the encoder-decoder architecture similar to the conventional transformer, with the use of multi head self attention being the key difference between the two. Multi head self attention is the idea where the Attention module comprises N (multiple) attention heads and each attention head repeats its computations multiple times in parallel. All the attention computations are combined to yield a final attention score. Bipartite Matching computed using the Hungarian Algorithm <cit.> is used to compute the difference between the ground truth bounding boxes and predicted bounding boxes of each of the objects in the image. The LETR vision transformer follows the DETR model, but is designed specifically for the identification of line segments. The LETR model improves upon the DETR model's ability to detect geometric structures by augmenting the loss term. 




    
  * Xu et al. trained the LETR model on two datasets namely, Wireframe <cit.> and YorkUrban <cit.> using the feature extractors of ResNet-101 and ResNet-50 respectively. The pretrained weights of the LETR model trained on the Wireframe dataset are used for the experiment.
    
  * The images from the Summer2015-AmesULA and Summer2015-AmesMLA datasets are evaluated using the pretrained weights. The line segments output by the model are plotted on evaluation images, as shown in Fig. <ref>.

    
    
  * However, it is the case that several of the detected line segments indicate the stem rather than the leaf, as shown in Fig. <ref>. The lines on the stem are irrelevant and must be ignored for the estimation of leaf angle. Two constraints, one on the slope and another on proximity to the image boundary are placed on the detected line segments to filter out the ones that don’t belong to the leaf.
    
        
  * Slope Constraint: The leaf is generally oriented in comparison to the stem that is almost perpendicular to the ground, in the majority of cases. As a result, line segments whose orientations are between 80 and 90 degrees are ignored, since they are likely to belong to the stem but not the leaf.
        
  * Boundary Proximity Constraint: It is empirically determined that the line segments which belong to the leaf are at least 100 pixels apart from the image boundary. The line segments that are closest to the image boundary are those of the stem. Hence, the constraint to ignore line segments that are within 100 pixels of the image boundary works to ignore line segments that belong to the stem. Both constraints are required to be satisfied for a line to be ignored. In certain images, the orientation of the leaf is between 80 and 90 degrees. In such cases, without the boundary constraint, the lines pertinent to the leaf are ignored. The line segments that pass the filter are assumed to be present on the leaf.
    
    
  * Compute the slope of each of the line segments and find the mode of the slopes. Mode refers to the most frequently occurring item among a list of items.
    
  * Retrieve the line segments with the most frequently occurring slope for further processing. Only the line segments with the orientation of the mode are considered to eliminate any stray lines that may be detected on the leaf. In the absence of mode i.e. no two line segments have the same slope, the line segment whose slope is the median of all slopes is considered.
    
  * Compute the orientation of the line segment made with the horizontal as the tan inverse of slope i.e., angle = tan^-1(slope). The orientation made by the line segment is deemed the orientation of the leaf with the stem.
    





§ IMPACTING FACTORS




 §.§ Impact of Image Sharpness

Image Sharpness describes the clarity of detail in a photo. The aspects of resolution and acutance primarily impact the sharpness of an image. It is imperative that the images be sharp for the line segments to be identified precisely by LETR. Multiple failures are observed with the detection of line segments in images of low sharpness. However, when the sharpness is increased, the line segments are well detected.



 §.§ Multiple Leaf and Stem Detections

In a few cases, the trained Mask R-CNN model detects multiple instances on a given image, some not pertinent to the leaf and stem. Fig. <ref> shows two images where multiple instances are detected. In such cases, the detection of the greatest size (contour) is considered the instance of interest, i.e. leaf and stem. However, it may be flawed in cases where the identified contour is larger than the actual leaf and stem instance in the image. 





§ RESULTS


The algorithm is applied to each of the images from the Summer2015-AmesULA and Summer2015-AmesMLA datasets, and the orientation between leaf and stem is estimated. The results are compared with the leaf angle measurements made by two graduate students at Iowa State University. The students worked independently of each other as they arrived at their estimates using the ImageJ tool. The measurements by the students and the proposed algorithm are compared using a metric known as Cosine Similarity <cit.>. Cosine Similarity is the cosine of the angle between two vectors that are typically non-zero and within an inner product space.  It is useful to compare the similarity between two vectors represented in a higher-dimensional vector space. Mathematically, it is defined as the division between the dot product of vectors and the product of the magnitude of each vector and is expressed as, similarity = A · B/|A||B| where A and B are the two vectors compared for similarity. The measure is expressed as a value between 0 and 1.   On the Summer2015-AmesMLA dataset, the cosine similarity between the proposed algorithm and student 1 is 0.98 which indicates a ten-degree difference in orientation between vectors, and the proposed algorithm and student 2 is 0.99 which indicates an eight-degree difference in orientation between the vectors. On the Summer2015-AmesULA dataset, the cosine similarity between the proposed algorithm and student 1 is 0.998 which indicates a four-degree difference in orientation between vectors, and the proposed algorithm and student 2 is 0.992 which indicates a seven-degree difference in orientation between the vectors.   Furthermore, an intuitive outlier analysis is performed on the results to gain insight into the similarity of estimation between the different techniques. The rule used to identify an outlier is, “the difference in estimates between the proposed algorithm and a student is greater than eight degrees”. The analysis assumes that the measurements made by the students are accurate. However, the estimations by the students also have a variance between them. The variance in student estimates, given by Σ(e1-e2)/n where e1 and e2 are estimates made by students 1 and 2 respectively, and n is the total size of the data sample, are 2.57 and 2.07 on Summer2015-AmesMLA dataset and Summer2015-AmesULA dataset respectively. Hence, the eight-degree range to identify outliers is reasonable. Applying the outlier rule, a total of 31 outliers on the Summer2015-AmesMLA dataset and 16 outliers on the Summer2015-AmesULA dataset are detected. Among the estimates that are not outliers, a variance of 3.73 degrees and 3.11 degrees is identified between the proposed algorithm and student 1, and the proposed algorithm and student 2 respectively on the Summer2015-AmesMLA dataset whereas a variance of 2.67 degrees and 2.72 degrees is identified between the proposed algorithm and student 1, and proposed algorithm and student 2 respectively on the Summer2015AmesULA dataset. The results show that the proposed algorithm outputs results that highly co-relate with the manual measurements, demonstrating the feasibility of the technique for leaf angle estimation.



§ FUTURE WORK AND CONCLUSION

Moving forward, the proposed algorithmic pipeline will be implemented as part of a real-time Android application so plant scientists and individual enthusiasts may leverage the application on a day-to-day basis. The bottleneck for the mobile application is the use of neural network models. The current Mask R-CNN network will have to be adapted so it works on mobile devices using networks such as Mobilenet SSD v1/v2 <cit.> as the backbone in place of the current ResNet-101 model. The annotation of images is another bottleneck that is laborious and time-consuming. The use of a screen to act as the common background over the plants as images are captured is an efficient way to help identify the location of the plant in the image.   Overall, the proposed technique demonstrates the idea of estimating the angle between leaf and stem using instance segmentation neural networks and vision transformers, a task that is laborious and cumbersome when performed manually. We will implement the leaf angle estimation as a feature in the FieldBook <cit.> Android application. The estimation of leaf angles using automation is novel and has the potential to turn into a one-of-a-kind utility for the agricultural community. 

00
Dzievit
M Dzievit and X Li and J Yu,
Dissection of leaf angle variation in maize through genetic mapping and meta‐analysis,
The plant genome,
12,
1

Zou
X Zou and M Mottus and P Tammeorg and C L Torres and T Takala and J Pisek and P Pellikka,
Photographic measurement of leaf angles in field crops,
Agricultural and Forest Meteorology,
184,
137-146

Sinoquet
H Sinoquet and S Thanisawanyangkura and H Mabrouk and P Kasemsap,
Characterization of the light environment in canopies using 3D digitising and image processing,
Annals of Botany,
 82,
2,
203 - 212

Hasoi 
F Hasoi and Y Nakai and K Omasa,
Estimating the leaf inclination angle distribution of the wheat canopy using a portable scanning lidar,
Journal of Agricultural Meteorology,
65,
3,
297-302

Herbert 
T J Herbert,
Stereophotogrammetry of plant leaf angles,
Photogrammetric Engineering and Remote Sensing,
61,
1,
89-92

Ryu 
Y Ryu and O Sonnentag and T Nilson and R Vargas and H Kobayashi and R Wenk and D D Baldocchi,
How to quantify tree leaf area index in an open savanna ecosystem: a multi-instrument and multi-model approach,
Agricultural and Forest Meteorology,
150,
1,
63 - 76

Kiryati
N Kiryati and Y Eldar and A M Bruckstein,
A probabilistic hough transform,
Pattern Recognition,
24,
4,
303 - 316

He
K He and G Gkioxari and P Dollar and R Girshick,
Mask r-cnn,
Proceedings of the IEEE conference on Computer Vision,
2017,
2961 - 2969

Bharati
P Bharati and A Pramanik,
Deep learning techniques-R-CNN to mask R-CNN,
Computational Intelligence in Pattern Recognition,
657 - 668,
2020

Tian
D Tian and Y Han and B Wang and T Guan and H Gu and W Wei,
2021,
Review of object instance segmentation based on deep learning,
Journal of Electronic Imaging,
31,
4

Lu
D Lu and Q Weng,
A survey of image classification methods and techniques for improving classification performance,
International journal of Remote Sensing,
28,
5,
823 - 870


Lateef
F Lateef and Y Ruichek,
2019,
Survey on semantic segmentation using deep learning techniques,
Neurocomputing,
338,
321 - 348

Liu
L Liu and W Ouyang and X Wang and P Fieguth and J Chen and X Liu and M Pietikainen,
2020,
Deep learning for generic object detection: A survey,
International journal of computer vision,
128,
2,
261 - 318

Schneider
C A Schneider and W S Rasband and K W Eliceiri,
NIH Image to ImageJ: 25 years of image analysis,
2012,
Nature Methods,
9,
7, 
671 - 675

Xu
Xu, Yifan and Xu, Weijian and Cheung, David and Tu, Zhuowen,
Line segment detection using transformers without edges,
Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition,
4257–4266,
2021

Ramachandran
Ramachandran, Prajit and Parmar, Niki and Vaswani, Ashish and Bello, Irwan and Levskaya, Anselm and Shlens, Jon,
Stand-alone self-attention in vision models,
Advances in neural information processing systems,
32,
2019

Sahbani
Sahbani, Bima and Adiprawita, Widyawardana,
Kalman filter and iterative-hungarian algorithm implementation for low complexity point tracking as part of fast multiple object tracking system,
2016 6th international conference on system engineering and technology (ICSET),
109–115,
2016,
IEEE

Huang
Huang, Kun and Wang, Yifan and Zhou, Zihan and Ding, Tianjiao and Gao, Shenghua and Ma, Yi,
Learning to parse wireframes in images of man-made environments,
Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition,
626–635,
2018

Denis
Denis, Patrick and Elder, James H and Estrada, Francisco J,
Efficient edge-based methods for estimating manhattan frames in urban imagery,
Computer Vision–ECCV 2008: 10th European Conference on Computer Vision, Marseille, France, October 12-18, 2008, Proceedings, Part II 10,
197–210,
2008,
Springer

Talbot
J F Talbot and X Xu,
2006,
Implementing grabcut,
Brigham Young University

Vijaymeena
M K Vijaymeena and K Kavitha,
2016,
A survey on similarity measures in text mining,
Machine Learning and Applications: An International Journal,
3,
2,
19 - 28

Pedregosa
F Pedregosa and G Varoquaux and A Gramfort and V Michel and B Thirion and O Grisel and E Duchesnay,
2011,
Scikit-learn: Machine learning in Python,
The journal of machine learning research,
12,
2825 - 2830

Chiu
Y C Chiu and C Y Tsai and M D Ruan and G Y Shen and T T Lee,
Mobilenet-SSDv2: An improved object detection model for embedded systems.,
2020,
International conference on system science and engineering (ICSSE),
1 - 5

Rife
T Rife and J Poland,
Crop Science,
Fieldbook: an open-source application for field data collection on android,
54,
4,
2014,
1624 - 1627





 Paper Title: Large Language Monkeys: Scaling Inference Compute with Repeated Sampling.



Information Scrambling at Quantum Hall Interfaces and Their Analog to Black Hole Event Horizon
    Ken K. W. Ma and Kun Yang
    August 5, 2024
==============================================================================================



§ ABSTRACT 

Scaling the amount of compute used to train language models has dramatically improved their capabilities. However, when it comes to inference, we often limit the amount of compute to only one attempt per problem. Here, we explore inference compute as another axis for scaling by increasing the number of generated samples. Across multiple tasks and models, we observe that coverage – the fraction of problems solved by any attempt – scales with the number of samples over four orders of magnitude. In domains like coding and formal proofs, where all answers can be automatically verified, these increases in coverage directly translate into improved performance. 
When we apply repeated sampling to SWE-bench Lite, the fraction of issues solved with DeepSeek-V2-Coder-Instruct increases from 15.9% with one sample to 56% with 250 samples, outperforming the single-attempt state-of-the-art of 43% which uses more capable frontier models.
Moreover, using current API pricing, amplifying the cheaper DeepSeek model with five samples is more cost-effective and solves more issues than paying a premium for one sample from GPT-4o or Claude 3.5 Sonnet. Interestingly, the relationship between coverage and the number of samples is often log-linear and can be modelled with an exponentiated power law, suggesting the existence of inference-time scaling laws. 
Finally, we find that identifying correct samples out of many generations remains an important direction for future research in domains without automatic verifiers. 
When solving math word problems from GSM8K and MATH, coverage with Llama-3 models grows to over 95% with 10,000 samples. 
However, common methods to pick correct solutions from a sample collection, such as majority voting or reward models, plateau beyond several hundred samples and fail to fully scale with the sample budget.

[0]^* Equal Contribution. Work done by BB as a visiting researcher at Stanford.





§ INTRODUCTION



The ability of large language models (LLMs) to solve coding, mathematics, and other reasoning tasks has improved dramatically over the past several years <cit.>. Scaling up model training has been a consistent driver of these gains. Investments in larger models, larger pre-training datasets, and more extensive post-training (e.g. through collecting human preference labels) has led to remarkably capable generalist systems <cit.>.

In contrast, a comparatively limited investment has been made in scaling the amount of computation used during inference. Larger models do require more inference compute than smaller ones, and prompting techniques like chain-of-thought <cit.> can increase answer quality at the cost of longer (and therefore more computationally expensive) outputs. However, when interacting with LLMs, users and developers often restrict models to making only one attempt when solving a problem. In this work, we investigate repeated sampling (depicted in Figure <ref>) as an alternative axis for scaling inference compute to improve LLM reasoning performance.



The effectiveness of repeated sampling is determined by two key properties:


    
  * Coverage: As the number of samples increases, what fraction of problems can we solve using any sample that was generated?
    
  * Precision: In settings where we must select a final answer from the collection of generated samples, can we identify correct samples?



With unlimited attempts, any model that assigns a non-zero probability to every sequence will achieve perfect coverage. However, repeated sampling is only practical if we can improve coverage with a feasible budget.
Moreover, without the ability to decide between samples, the applications of repeated sampling are limited. Existing work provides encouraging evidence along both of these directions, showing examples of repeated sampling improving performance in math, coding, and puzzle-solving settings <cit.>. Notably, AlphaCode <cit.>, a state-of-the-art system for competitive programming, finds that performance continues to improve with a million samples per problem. 


Here, we show that repeated sampling is an effective methodology for improving coverage across a range of tasks, models, and sample budgets.
For example, when solving CodeContests <cit.> programming problems using Gemma-2B <cit.>, we increase coverage by over 300x as we scale the number of samples, from 0.02% with one attempt to 7.1% with 10,000 attempts.  
Interestingly, the relationship between log(coverage) and the number of samples often follows an approximate power law. With Llama-3  <cit.> and Gemma models, we observe that coverage grows nearly log-linearly with the number of samples over several orders of magnitude.


In settings where all of a model’s solutions can be automatically verified, such as with proof checkers or unit tests, these increases in coverage translate directly into improved task performance. When applying repeated sampling to competitive programming and writing Lean proofs, models like Llama-3-8B-Instruct can exceed the single-attempt performance of much stronger ones like GPT-4o <cit.>. This ability to amplify weaker models extends to the challenging SWE-bench Lite dataset of real-life GitHub issues <cit.>, where the current single-attempt state-of-the-art (SOTA), achieved by a mixture of GPT-4o and Claude 3.5 Sonnet models, is 43% <cit.>. When given a single attempt, DeepSeek-Coder-V2-Instruct <cit.> achieves only 15.9% on the benchmark. By simply increasing the number of attempts to 250, we increase the fraction of solved problems to 56%, exceeding the single-attempt state-of-the-art by 13%. 

In addition to improving model quality, repeated sampling provides a new mechanism for minimizing LLM inference costs. When holding the total number of inference FLOPs constant, we find that on some datasets (e.g. MATH) coverage is maximized with a smaller model and more attempts, while on others (e.g CodeContests) it is better to use a larger model. We also compare API prices between DeepSeek-V2-Coder-Instruct, GPT-4o, and Claude Sonnet 3.5 in the context of solving SWE-bench Lite issues. When keeping the agent framework (Moatless Tools <cit.>) constant, sampling five times from the weaker and cheaper DeepSeek model solves more issues than single attempts from Claude or GPT, while also being over 3x cheaper.

Finally, in math word problem settings, where answers cannot be automatically verified by existing tools, we identify a large gap between coverage and the performance of common methods for deciding on a final answer. When solving MATH <cit.> problems with Llama-3-8B-Instruct, coverage increases from 79.8% with 100 samples to 95.3% with 10,000 samples. However, methods such as majority voting and using reward models plateau with a lower sample budget, scaling only from 38.7% to 39.8% over the same range. These results highlight that building robust verifiers remains an open problem.

In summary, our primary observations are:


    
  * We demonstrate that scaling inference compute through repeated sampling leads to large improvements in coverage across a variety tasks, models, and sample budgets. This makes it possible, and sometimes cost-effective, to amplify weaker models with many samples and outperform single attempts from more capable models. Notably, we are able to solve 56% of issues from SWE-bench Lite by sampling 250 times from DeepSeek-V2-Coder-Instruct, exceeding the single-attempt SOTA of 43%.

    
  * We show that the relationship between coverage and the number of samples can often be modelled using an exponentiated power law, suggesting a form of scaling laws for inference-time compute.  

    
  * In domains without automatic verifiers, we show that common approaches to verification like majority voting and reward model scoring plateau beyond approximately 100 samples. This leads to a growing gap between the performance achieved with these methods and the coverage upper bound.




§ SCALING REPEATED SAMPLING



We focus on pass-fail tasks where a candidate solution can be scored as right or wrong. The primary metric of interest for these tasks is the success rate: the fraction of problems that we are able to solve. With repeated sampling, we consider a setup where a model can generate many candidate solutions while attempting to solve a problem. The success rate is therefore influenced both by the ability to generate correct samples for many problems (i.e. coverage), as well as the ability to identify these correct samples (i.e. precision).

The difficulty of the precision problem depends on the availability of tools for sample verification. When proving formal statements in Lean, proof checkers can quickly identify whether a candidate solution is correct. Similarly, unit tests can be used to verify candidate solutions to coding tasks. In these cases, precision is handled automatically, and improving coverage directly translates into higher success rates. In contrast, the tools available for verifying solutions to math word problems from GSM8K and MATH are limited, necessitating additional verification methods that decide on a single final answer from many (often conflicting) samples.

We consider the following five tasks:


    
  * GSM8K: A dataset of grade-school level math word problems <cit.>. We evaluate on a random subset of 128 problems from the GSM8K test set.
    
  * MATH:  Another dataset of math word problems that are generally harder than those from GSM8K <cit.>. Similarly, we evaluate on 128 random problems from this dataset's test set.
    
  * MiniF2F-MATH: A dataset of mathematics problems that have been formalized into proof checking languages <cit.>. We use Lean4 as our language, and evaluate on the 130 test set problems that are formalized from the MATH dataset. 
    
    
  * CodeContests: A dataset of competitive programming problems <cit.>. Each problem has a text description, along with a set of input-output test cases (hidden from the model) that can be used to verify the correctness of a candidate solution. We enforce that models write their solutions using Python3. 
    
  * SWE-bench Lite: A dataset of real world Github issues, where each problem consists of a description and a snapshot of a code repository <cit.>. To solve a problem, models must edit files in the codebase (in the Lite subset of SWE-bench that we use, only a single file needs to be changed). Candidate solutions can be automatically checked using the repository’s suite of unit tests.
 

Among these tasks, MiniF2F-MATH, CodeContests, and SWE-bench Lite have automatic verifiers (in the form of the Lean4 proof checker, test cases, and unit test suites, respectively).
We begin by investigating how repeated sampling improves model coverage. Coverage improvements correspond directly with increased success rates for tasks with automatic verifiers and in the general case provide an upper bound on the success rate. 
In coding settings, our definition of coverage is equivalent to the commonly-used pass@k metric <cit.>, where k denotes the number of samples per problem. We use this metric directly when evaluating on CodeContests and SWE-bench Lite. For MiniF2F the metric is similar, with a “pass” defined according to the Lean4 proof checker. For GSM8K and MATH, coverage corresponds to using an oracle verifier that checks if any sample “passes” by outputting the correct final answer. To reduce the variance when calculating coverage, we adopt the unbiased estimation formula from <cit.>. In each experiment, we first generate N samples for each problem index i and calculate the number of correct samples C_i. We then calculate the pass@k scores at each k ≤ N of interest according to:


    pass@k   = 1/# of problems∑_i = 1^# of problems1 - N - C_ik/Nk


We use the numerically stable implementation of the above formula suggested in <cit.>.
Results and code will be available at <https://scalyresearch.stanford.edu/pubs/large_language_monkeys/>.



 §.§ Repeated Sampling is Effective Across Tasks






Here, we establish that repeated sampling improves coverage across multiple tasks and a range of sample budgets. We evaluate Llama-3-8B-Instruct and Llama-3-70B-Instruct on CodeContests, MiniF2F, GSM8K, and MATH, generating 10,000 independent samples per problem. For SWE-bench Lite, we use DeepSeek-V2-Coder-Instruct <cit.>, as the required context length of this task exceeds the limits of the Llama-3 models. As is standard when solving SWE-bench issues, we equip our LLM with a software framework that provides the model with tools for navigating through and editing codebases. In our work, we use the open-source Moatless Tools library <cit.>. Note that solving a SWE-bench issue involves a back-and-forth exchange between the LLM and Moatless Tools. One sample/attempt for this benchmark refers to one entire multi-turn trajectory. To minimize costs, we restrict the number of attempts per issue to 250, with all attempts made independently of one another.

We report our results in Figure <ref>. We also include the single-attempt performance of GPT-4o on each task, as well the single-attempt state-of-the-art for SWE-bench Lite (CodeStory Aide <cit.> which uses a combination of GPT-4o and Claude 3.5 Sonnet). Across all five tasks, we find that coverage smoothly improves as the sample budget increases. When all LLMs are given a single attempt, GPT-4o outperforms the Llama and DeepSeek models at every task. However, as the number of samples increases, all three of the weaker models exceed GPT-4o's single-attempt performance. In the case of SWE-bench Lite, we solve 56% of problems, exceeding the single-attempt SOTA of 43%. 



 §.§ Repeated Sampling is Effective Across Model Sizes and Families





The results from Section <ref> indicate that repeated sampling improves coverage. However, we only show this trend for three recent, instruction-tuned models with 8B or more parameters. We now show that these trends hold across other model sizes, families, and levels of post-training. We expand our evaluation to include a broader set of models:


    
  * Llama 3: Llama-3-8B, Llama-3-8B-Instruct, Llama-3-70B-Instruct.
    
  * Gemma: Gemma-2B, Gemma-7B <cit.>.
    
  * Pythia: Pythia-70M through Pythia-12B (eight models in total) <cit.>.


We restrict evaluation to the MATH and CodeContests datasets to minimize inference costs, reporting results in Figure <ref>. Coverage increases across almost every model we test, with smaller models showing some of the sharpest increases in coverage when repeated sampling is applied. On CodeContests, the coverage of Gemma-2B increases by over 300x, from a pass@1 of 0.02% to a pass@10k of 7.1%. 
Similarly, when solving MATH problems with Pythia-160M, coverage increases from a pass@1 of 0.27% to a pass@10k of 57%.

The exception to this pattern of increasing coverage across models is with the Pythia family evaluated on CodeContests. All Pythia models achieve zero coverage on this dataset, even with a budget of 10,000 samples. We speculate that this due to Pythia being trained on less coding-specific data than Llama and Gemma. 



 §.§ Repeated Sampling Can Help Balance Performance and Cost




One takeaway from the results in Sections <ref> and <ref> is that repeated sampling makes it possible to amplify a weaker model's capabilities and outperform single samples from stronger models. 
Here, we demonstrate that this amplification can be more cost-effective than using a stronger, more expensive model, providing practitioners with a new degree of freedom when trying to jointly optimize performance and costs.

We first consider FLOPs as a cost metric, examining the Llama-3 results from Section <ref>. We re-plot our results from Figure <ref>, now visualizing coverage as a function of total inference FLOPs instead of the sample budget. Since Llama-3 models are dense transformers where the majority of parameters are used in matrix multiplications, we approximate inference FLOPs with the formula:


    FLOPs per token   ≈ 2 * num parameters + 2 * num layers * token dim * context length
    total inference FLOPs   ≈num prompt tokens * FLOPs per token
       + num decoded tokens * FLOPs per token * num completions



We present our re-scaled results for MiniF2F, CodeContests, MATH, and GSM8K in Figure <ref>. Interestingly, the model that maximizes coverage varies with the compute budget and task. On MiniF2F, GSM8K and MATH, Llama-3-8B-Instruct always obtains higher coverage than the larger (and more expensive) 70B model when the FLOP budget is fixed. 
However for CodeContests, the 70B model is almost always more cost effective.
We note that examining FLOPs alone can be a crude cost metric that ignores other aspects of system efficiency <cit.>. 
In particular, repeated sampling can make use of high batch sizes and specialized optimizations that improve system throughput relative to single-attempt inference workloads <cit.>. 
We discuss this in more detail in Section <ref>.



We also examine the dollar costs of repeated sampling when solving SWE-bench Lite issues using current API pricing. Keeping the agent framework (Moatless Tools) constant, we consider making a single attempt per issue with Claude 3.5 Sonnet and GPT-4o, as well as repeated sampling using DeepSeek-V2-Coder-Instruct. We report the average cost per issue and issue resolution rate with each approach in Table <ref>. While the DeepSeek model is weaker than the GPT and Claude models, it is also over 10x cheaper. In this case, repeated sampling provides a cheaper alternative to paying a premium for access to strong models while achieving a superior issue solve rate.






§ CHARACTERIZING THE BENEFITS OF REPEATED SAMPLING




The relationship between an LLM’s loss and its training compute has been well-characterized with training scaling laws <cit.>. 
These laws have empirically held over many orders of magnitude and inspire confidence in model developers that large investments in training will pay off. 
Inspired by training scaling laws, here we aim to better characterize the relationship between coverage and the sample budget (i.e. the amount of inference compute), presenting two interesting observations:


    
  * The relationship between coverage and the number of samples can often be modelled with an exponentiated power law.
    
  * For a given task, the coverage curves of different models from the same family resemble S-curves with similar slopes but distinct horizontal offsets.





 §.§ Scaling Laws for Repeated Sampling





Here, we develop an explicit model for the relationship between coverage and the number of samples. The GPT-4 technical report <cit.> finds that the relationship between a model’s mean-log-pass-rate on coding problems and its training compute can be modelled well using a power law. We start by adopting the same function class, but now modelling the log of coverage c as a function of the number of samples k: 

    log(c) ≈ a k^-b

where a,b ∈ are fitted model parameters. 
In order to directly predict coverage, we exponentiate both sides, ending up with the final model of:

    c ≈exp(a k^-b)

We provide examples of fitted coverage curves in Figure <ref>, and additional curves in Appendix <ref>. While these laws are not as exact as training scaling laws (most strikingly on MiniF2F-MATH), they provide encouraging early evidence that the benefits of inference scaling can be characterized.



 §.§ Similarities in Coverage Curves Across Models







Interestingly, when comparing the coverage curves (with a logarithmic x-axis) of different models from the same family on the same task (see Figure <ref>), it appears that the traced S-curves have the same slope, but unique horizontal offsets. To investigate this further, we overlay the coverage curves of different models from the same family in Figure <ref>. We do this by picking an anchor coverage value c, and shifting every curve leftward (in log-space) so that each passes through the point (1, c). This corresponds to a leftward shift by log(pass@k^-1(c)), where pass@k^-1(c) denotes the closest natural number k such that pass@k = c. We pick c to be the maximum pass@1 score over all models from the same family. These similarities demonstrate that across models from the same family, the increase in the log-sample-budget (or equivalently, the multiplicative increase in the sample budget) needed to improve coverage from c to c' is approximately constant.



§ HARNESSING REPEATED SAMPLING REQUIRES PRECISION



So far, we have focused on measuring model coverage, characterizing the benefits of repeated sampling under the best-case scenario where we can always identify correct model samples. We now turn to the complementary problem of precision: given a collection of model samples, can we identify the correct ones? In Section <ref>, we evaluate two common verification methods (majority voting and reward model scoring) on GSM8K and MATH. Additionally, in Section <ref>, we discuss potential pitfalls when relying on unit tests to identify correct software programs.



 §.§ Common Verification Methods Don't Always Scale with the Sample Budget







Of the five tasks we evaluate, only GSM8K and MATH lack tools for automatically verifying solutions. Here, we evaluate two common approaches to deciding on a final answer: calculating a majority vote across samples and using a reward model to assign a score to each sample. We test these techniques on their ability to identify correct solutions from the 10,000 samples we generated with Llama-3-8B-Instruct and Llama-3-70B-Instruct in Section <ref>. We benchmark three methods:


    
  * Majority Vote: We pick the most common final answer <cit.>.
    
  * Reward Model + Best-of-N: We use a reward model <cit.> to score each solution, and pick the answer from the highest-scoring sample.  
    
  * Reward Model + Majority Vote: We calculate a majority vote where each sample is weighted by its reward model score.


We use ArmoRM-Llama3-8B-v0.1 <cit.> as a reward model, which currently has the highest reasoning score on the RewardBench leaderboard <cit.> among open-weight models. We report our results in Figure <ref> as we increase the number of samples. While success rate initially increases with the number of samples for all three methods, it plateaus around 100 samples. Meanwhile, coverage continues to increase with the number of samples and exceeds 95%.

In the case of majority voting, this success rate saturation is easily explainable. As the number of samples increases, the proportion of votes allocated to each answer stabilizes, and therefore the success rate plateaus. For some GSM8K and MATH problems, correct solutions are sampled with a probability of 1% or lower (see Figure <ref>), making them a minority of samples. As the number of samples increases, rare correct solutions will appear for more problems, increasing coverage but not the success rate with majority voting. In order to fully benefit from repeated sampling, sample identification methods must be able to solve these “needle-in-a-haystack” cases and identify rare, correct samples.

Given the poor performance of the reward model, it is reasonable to wonder how “hard’’ it is to verify a candidate solution. With GSM8K and MATH, only a sample’s final answer is used for assessing correctness, with the intermediate chains of thought being excluded. If models generated only non-sensical chains of thought before guessing a correct final answer, verification may not be any easier than solving the problem in the first place. We investigate this question by manually evaluating 105 chains of thought from correct Llama-3-8B-Instruct solutions to GSM8K problems, reporting our results in Table <ref>. 
We find that over 90% of the chains of thought that we graded are faithful, even among problems where correct answers are generated infrequently.
These correct reasoning steps indicate that there is signal for a verifier to exploit when identifying correct samples. Interestingly, during this process we also identified one GSM8K problem that has an incorrect ground truth answer (see Appendix <ref>). 
This incorrect GSM8K problem is also the only one that Llama-3-70B-Instruct did not generate a “correct” sample for across 10,000 attempts.








 §.§ Verifiers and Software Tasks: Two Cautionary Tales



Software development tasks can occupy a middle-ground with respect to available verification tools.
On one hand, the ability to execute and test code allows for a higher degree of automatic verification than is possible with unstructured language tasks.
However, tools like unit tests take a black-box approach to verifying a piece of code and are not as comprehensive as methods like proof checkers.
These imperfections in the verification process can lead to false positives and/or false negatives that are important to consider when applying repeated sampling. Below we provide two examples of software verifier imperfections that we encountered when generating our results from Section <ref>.  



  §.§.§ Flaky Tests in SWE-bench Lite


 
When producing our results on SWE-bench Lite, we identified that 11.3% of problems have flaky test suites that do not produce consistent results when running them on the same candidate solution.
These flaky tests occasionally classify even the dataset's ground-truth issue solutions as incorrect. Additionally, the test suites for some issues can be non-determinstic depending on the candidate solution.
For example, two SWE-bench Lite issues involve manipulating Python sets, which are naturally unordered.
The gold solutions for these issues explicitly order the items in the set and pass the test suites reliably. However, some model-generated candidate solutions do not impose such an ordering, and therefore pass the tests on some “lucky’’ runs and not others.
In Appendix <ref>, we list all of the problem IDs where we identified flaky tests. We also report our SWE-bench Lite results from Figure <ref> with the problematic issues removed, finding similar results to our evaluations on the whole dataset. 



  §.§.§ False Negatives in CodeContests


Each problem from the CodeContests dataset comes with a set of input-output test cases used to asses the correctness of solutions. These test cases are more comprehensive than those from earlier coding benchmarks like APPS <cit.>, cutting down on the frequency of false positive solutions that pass all test cases but do not fully solve the described problem. However, the construction of the CodeContests test suites leads to false negative solutions that are correct but fail the tests.

For some CodeContests problems, the problem description allows for multiple distinct correct outputs for a given test input. However, the corresponding test cases do not handle these scenarios, instead requiring that one particular correct output is emitted. Additionally, many CodeContests test cases have been programmatically generated by mutating original test cases from the problem. Some mutated inputs violate the problem’s input specifications (e.g. a mutated input being zero when the description promises a positive integer). 
These malformed test cases can lead to inconsistent behaviour between different correct solutions.

We assess the prevalence of these issues by running each problem’s test suite on the list of correct solutions that CodeContests provides. Of the 122 problems in the test set that have Python3 solutions, we find that 35 problems have “correct’’ solutions that fail the corresponding tests. Since we do not allow models to view all of a problem’s test cases (and their peculiarities), applying repeated sampling to these problems contains an element of “rolling the dice” to generate a solution that is not only correct, but emits the particular outputs that pass the tests.







§ DISCUSSION AND LIMITATIONS



In this work, we explore repeated sampling as an axis for scaling compute at inference time in order to improve model performance. Across a range of models and tasks, repeated sampling can significantly improve the fraction of problems solved using any generated sample (i.e. coverage). When correct solutions can be identified (either with automatic verification tools or other verification algorithms), repeated sampling can amplify model capabilities during inference. This amplification can make the combination of a weaker model and many samples more performant and cost-effective than using fewer attempts from a stronger, more expensive model.

 

Improving Repeated Sampling: In our experiments, we explore only a simple version of repeated sampling where all attempts to a problem are generated independently of one another using the exact same prompt and hyperparameters. We believe that this setup can be refined to improve performance, particularly along the following directions: 


    
  * Solution Diversity: We currently rely on a positive sampling temperature as the sole mechanism for creating diversity among samples. Combining this token-level sampling with other, higher-level approaches may be able to further increase diversity. For example, AlphaCode conditions different samples with different metadata tags.

    
  * Multi-Turn Interactions: Despite automatic verification tools being available when solving CodeContests and MiniF2F problems, we use only a single-turn setup where models generate a solution without any ability to iterate on it. Providing models with execution feedback from these tools should improve solution quality. We are interested in the tradeoffs associated with multi-turn interactions, since each attempt becomes more expensive, but also may be more likely to succeed.
    
  * Learning From Previous Attempts: Currently, our experiments fully isolate attempts from each other. Access to existing samples, particularly if verification tools can provide feedback on them, may be helpful when generating future attempts.


  

Repeated Sampling and Inference Systems: Repeated sampling is a distinct LLM inference workload from serving chatbot requests. Production chatbot deployments place an emphasis on low response latencies, and adhering to latency targets can force a lower per-device batch size and reduce hardware utilization. In contrast, when sampling many completions to a single prompt, a larger emphasis can be placed on overall throughput and maximizing hardware utilization. Additionally, repeated sampling can benefit from specialized attention optimizations that exploit overlaps in prompts across sequences <cit.>. Repeated sampling inference can therefore be accomplished at a lower cost than naively making many parallel requests to a chatbot-oriented API. These cost savings can further motivate choosing to sample many times from a cheaper model instead of fewer times from a more expensive one. 

 

Verifiers: Our results from Section <ref> highlight the importance of improving sample verification methods when tools for automatically doing so are unavailable. Equipping models with the ability to assess their own outputs will allow repeated sampling to be scaled to far more tasks. Of particular interest is applying repeated sampling to unstructured tasks like creative writing, which can require a more subjective comparison between different samples than the pass-fail tasks we consider. An alternative direction to developing model-based verifiers is to design converters that can make an unstructured task verifiable, for example by formalizing an informal math statement into a language like Lean so that proof checkers can be applied.




§ RELATED WORK


Scaling Inference Compute: Methods that perform additional computation during inference have been successful across many areas of deep learning. Across a variety of game environments, state-of-the-art methods leverage inference-time search that examines many possible future game states before deciding on a move <cit.>. Similar tree-based methods can also be effective in combination with LLMs, allowing models to better plan and explore different approaches <cit.>. Another axis for increasing LLM inference compute allows models to spend tokens deliberating on a problem before coming to a solution <cit.>. Additionally, multiple models can be ensembled together at inference time to combine their strengths <cit.>. Yet another approach involves using LLMs to critique and refine their own responses <cit.>.

 

Repeated Sampling: Previous work has demonstrated that repeated sampling can improve LLM capabilities in multiple domains. One of the most effective use cases is coding <cit.>, where performance continues to scale up to a million samples and verification tools (e.g. unit tests) are often available to automatically score every candidate solution. 
Recently, <cit.> shows that repeated sampling is effective when solving puzzles from the ARC challenge <cit.>, observing log-linear scaling as the number of samples increases.
In chat applications, repeated sampling combined with best-of-N ranking with a reward model can outperform greedily sampling a single response <cit.>.
In domains without automatic verification tools, existing work shows that using majority voting <cit.> or a trained model-based verifier <cit.> to decide on a final answer can improve performance on reasoning tasks relative to taking a single sample. Concurrent with our work, <cit.> finds that using the best available sample improves LLM performance on chat, math, and code tasks, sweeping up to a max of 128 samples.


 

Scaling Laws: Characterizing how scaling affects model performance can lead to more informed decisions on how to allocate resources. Scaling laws for LLM training find a power law relationship between loss and the amount of training compute and provide estimates for the optimal model and dataset size given a fixed compute budget  <cit.>.
<cit.> finds scaling laws in the context of the board game Hex, observing that performance scales predictably with model size and the difficulty of the problem.
Interestingly, they also show that performance scales with the amount of test-time compute spent while performing tree search.
Recently, <cit.> observe scaling laws when augmenting LLMs with external retrieval datasets, finding that performance on retrieval tasks scales smoothly with the size of the retrieval corpus.






§ ACKNOWLEDGEMENTS


We thank Together AI for partially sponsoring the compute for this project, as well as Rahul Chalamala and Ben Athiwaratkun for their help managing this infrastructure. We thank John Yang for his advice and support when running our SWE-bench experiments. Finally, we are grateful to Mayee Chen, Neel Guha, Quinn McIntyre, Jon Saad-Falcon, and Benjamin Spector for  their helpful discussions and feedback throughout this project.


We gratefully acknowledge the support of NIH under No. U54EB020405 (Mobilize), NSF under Nos. CCF2247015 (Hardware-Aware), CCF1763315 (Beyond Sparsity), CCF1563078 (Volume to Velocity), and 1937301 (RTML); US DEVCOM ARL under Nos. W911NF-23-2-0184 (Long-context) and W911NF-21-2-0251 (Interactive Human-AI Teaming); ONR under Nos. N000142312633 (Deep Signal Processing); Stanford HAI under No. 247183; NXP, Xilinx, LETI-CEA, Intel, IBM, Microsoft, NEC, Toshiba, TSMC, ARM, Hitachi, BASF, Accenture, Ericsson, Qualcomm, Analog Devices, Google Cloud, Salesforce, Total, the HAI-GCP Cloud Credits for Research program,  the Stanford Data Science Initiative (SDSI), and members of the Stanford DAWN project: Meta, Google, and VMWare. The U.S. Government is authorized to reproduce and distribute reprints for Governmental purposes notwithstanding any copyright notation thereon. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views, policies, or endorsements, either expressed or implied, of NIH, ONR, or the U.S. Government. 

This work was completed with the support of the Clarendon Fund Scholarships.


plainnat





§ SAMPLING EXPERIMENTAL SETUP




 §.§ Lean Formal Proofs

We report results on the 130 questions in the test set of the https://github.com/rah4927/lean-dojo-mew/blob/main/MiniF2F/Test.leanlean4 MiniF2F dataset that correspond to formalized MATH problems. This dataset is derived from the https://github.com/facebookresearch/miniF2Ffixed version of the original MiniF2F dataset created by <cit.>. We sample with a temperature of 0.5 and do not use nucleus sampling. We generated 10,000 samples per problem. We use proofs of the following 5 theorems from the https://github.com/rah4927/lean-dojo-mew/blob/main/MiniF2F/Validation.leanvalidation set as few-shot examples:

    
  * 
    
  * 
    
  * 
    
  * 
    
  * 



Our prompt consists of:

    
  * Few shot examples.
    
  * Header imports present in each problem in the HuggingFace dataset  dataset, an upload of the lean4 MiniF2F dataset. 
    
  * The theorem definition. In order to avoid leaking information about how to solve the theorem from its name, we replace the name of the theorem with . i ∈{1,2,3,4,5} for the few-shot examples and i=6 for the current problem. 


We set 200 as the max token length for the generated solution. To grade solutions, we use the  library with lean version  . We set a timeout of 10 seconds for every tactic step.

[
    colback=gray!5,
    colframe=gray!75,
    title=Few-Shot Example,
    fonttitle=]
Write a lean4 proof to the provided formal statement. You have access to the standard mathlib4 library.

```import Mathlib.Algebra.BigOperators.Basic

import Mathlib.Data.Real.Basic

import Mathlib.Data.Complex.Basic

import Mathlib.Data.Nat.Log

import Mathlib.Data.Complex.Exponential

import Mathlib.NumberTheory.Divisors

import Mathlib.Data.ZMod.Defs

import Mathlib.Data.ZMod.Basic

import Mathlib.Topology.Basic

import Mathlib.Data.Nat.Digits


open BigOperators

open Real

open Nat

open Topology

theorem theorem1

  Int.floor ((9:ℝ) / 160 * 100) = 5 :=

by (

  rw [Int.floor_eq_iff]

  constructor

  all_goals norm_num

)```

[
    colback=gray!5,
    colframe=gray!75,
    title=Example Prompt,
    fonttitle=]
Write a lean4 proof to the provided formal statement. You have access to the standard mathlib4 library.

```import Mathlib.Algebra.BigOperators.Basic

import Mathlib.Data.Real.Basic

import Mathlib.Data.Complex.Basic

import Mathlib.Data.Nat.Log

import Mathlib.Data.Complex.Exponential

import Mathlib.NumberTheory.Divisors

import Mathlib.Data.ZMod.Defs

import Mathlib.Data.ZMod.Basic

import Mathlib.Topology.Basic

import Mathlib.Data.Nat.Digits


open BigOperators

open Real

open Nat

open Topology

theorem theorem6

  (b h v : )

  (h_0 : 0 < b ∧ 0 < h ∧ 0 < v)

  (h_1 : v = 1 / 3 * (b * h))

  (h_2 : b = 30)

  (h_3 : h = 13 / 2) :

  v = 65 :=

by (






 §.§ CodeContests


We report results on the 140 test set questions that do not include image tags in the problem description. We sample with a temperature of 0.6 and a top-p value of 0.95 following the experiments in CodeLlama <cit.>. We generate 10,000 samples per problem. We use two few-shot examples from the training set that are randomly sampled per-problem. We set 1024 as the max token length for the generated solution. We use the same answer comparison function as <cit.> and use the concatenation of public, private, and generated tests to validate correctness of solutions.

[
    colback=gray!5,
    colframe=gray!75,
    title=Few-Shot Example,
    fonttitle=]
Q: Write python code to solve the following coding problem that obeys the constraints and passes the example test cases. The output code needs to read from and write to standard IO. Please wrap your code answer using ```:

Petya has equal...

...

(Rest of question and input/output examples)

...

A: ```
```

[
    colback=gray!5,
    colframe=gray!75,
    title=Example Prompt,
    fonttitle=]
Q: Write python code to solve the following coding problem that obeys the constraints and passes the example test cases. The output code needs to read from and write to standard IO. Please wrap your code answer using ```:

This is the...

...

(Rest of question and input/output examples)

...

A:






 §.§ MATH
 We report results on 128 randomly selected test-set problems. We sample with a temperature of 0.6 and do not use nucleus sampling. We use the fixed 5 few-shot example from <cit.> for each problem.  We generate 10,000 samples per problem. We set 512 as the max token length for the generated solution. To grade solutions, we use the  functions from LMEval <cit.>.


[
    colback=gray!5,
    colframe=gray!75,
    title=Few-Shot Example,
    fonttitle=]
Problem:

If 𝐀 = 2 and 𝐁 = 12, then find (𝐀𝐁).

Solution:

We have that (𝐀𝐁) = (𝐀)(𝐁) = (2)(12) = 24.
Final Answer: The final answer is 24. I hope it is correct.


[
    colback=gray!5,
    colframe=gray!75,
    title=Example Prompt,
    fonttitle=]
Problem:

What is the domain of the function 
    f(x) = (2x-3)(2x+5)/(3x-9)(3x+6) ?
 Express your answer as an interval or as a union of intervals.

Solution:




 §.§ GSM8K
 We report results on 128 randomly sampled test-set problems. We sample with a temperature of 0.6 and do not use nucleus sampling. We use 5 few-shot examples from the training set that are randomly sampled per-problem. We generate 10,000 samples per problem. We set 512 as the max token length for the generated solution. To grade solutions, we follow LMEval <cit.>.


[
    colback=gray!5,
    colframe=gray!75,
    title=Few-Shot Example,
    fonttitle=]
Question: James decides to replace his car.  He sold his $20,000 car for 80% of its value and then was able to haggle to buy a $30,000 sticker price car for 90% of its value.  How much was he out of pocket?

Answer: He sold his car for 20000*.8=$20000*.8=1600016,000
He bought the new car for 30,000*.9=$30000*.9=2700027,000
That means he was out of pocket 27,000-16,000=$27000-16000=1100011,000

#### 11000

[
    colback=gray!5,
    colframe=gray!75,
    title=Example Prompt,
    fonttitle=]
Question: Mary has 6 jars of sprinkles in her pantry. Each jar of sprinkles can decorate 8 cupcakes. Mary wants to bake enough cupcakes to use up all of her sprinkles. If each pan holds 12 cupcakes, how many pans worth of cupcakes should she bake?

Answer:




§ SWE-BENCH LITE





 §.§ Experimental Setup


For our experiments, we use DeepSeek-V2-Coder-Instruct with the Moatless Tools agent framework (at commit ).
We use Voyage AI <cit.> embeddings for retrieval, the default used by Moatless Tools. We make no modifications to the model or framework, using them entirely as off-the-shelf components.

With this setup, we sample 250 independent completions for each problem using standard temperature-based sampling. To determine the optimal sampling temperature, we conducted a sweep on a random subset of 50 problems from the test set, testing temperatures of 1.0, 1.4, 1.6, and 1.8. Based on these results, we selected a temperature of 1.6 for our main experiments. 



 §.§ Test Suite Flakiness


During our analysis, we identified 34 problems in SWE-bench Lite whose test suites had flaky tests. Using the SWE-bench testing harness provided by the authors of SWE-bench, we tested each solution repeatedly: for some solutions, sometimes the solution was marked as correct, and other times it was marked as incorrect. In 30 of these 34 cases, we observed flakiness even on the correct solutions provided by the dataset authors. Table <ref> lists the problem IDs of the 34 instances with flaky tests.



An additional instance, , was flaky on some machines and not others. The authors of SWE-bench were able to reproduce the flakyness; however, we were unable to. Our prelimiary investigation indicates this specific issue is due to unpinned versions of dependencies in the docker environments that run the unit tests.



Here, we include results on a subset with the problems in Table <ref> removed (266 problems). For the full dataset evaluation, on any problem that has flaky tests, we run the test suite 11 times and use majority voting to determine whether a solution passed or failed. For the evaluation on the subset without flaky tests, all baselines we compare against release which problems they correctly solve, so we simply removed the problems with flaky tests and recomputed their scores.









§ SCALING LAW DETAILS




 §.§ Experimental details


To fit exponentiated power laws to coverage curves, we first sample 40 points spaced evenly along a log scale from 0 to 10,000 and remove duplicates. We then use SciPy's <cit.>  function to find the a and b parameters from Equation <ref> that best fit these points.



 §.§ Additional results



In Figure <ref>, we show additional results fitting power laws to coverage curves for an expanded set of datasets and models. 






§ PRECISION DETAILS



To calculate the Majority Vote, Reward Model + Best-of-N and Reward Model + Majority Vote metrics, we use the same 128 problem subsets for both MATH and GSM8K datasets introduced in Section <ref>.
Each problem corresponds to 10,000 samples for each model we test.
For each verification method, we take 100 random subsets of size k and calculate the success rate using each subset.
We report the mean and standard deviation across subsets in Figure <ref>. 
To calculate the Majority Vote answer, we take the plurality answer in each subset.
For the Reward Model + Best-of-N, we take the answer with the highest score assigned by the reward model.
For the Reward Model + Majority Vote metric, we sum the reward model score across all the samples with the same final answer, and use the final answer with the highest sum.




§ GSM8K INCORRECT ANSWER


As discussed in <ref>, we identify that  https://huggingface.co/datasets/openai/gsm8k/viewer/main/test?row=1042a problem in the GSM8K test set (index 1042 on HuggingFace) has an incorrect ground truth solution.


[
    colback=gray!5,
    colframe=gray!75,
    title=Question,
    fonttitle=]
Johnny's dad brought him to watch some horse racing and his dad bet money. On the first race, he lost $5. On the second race, he won $1 more than twice the amount he previously lost. On the third race, he lost 1.5 times as much as he won in the second race. How much did he lose on average that day?


[
    colback=gray!5,
    colframe=red!75,
    title=Answer,
    fonttitle=]
On the second race he won $11 because 1+ 5 × 2 = <<1+5*2=11>>11 

On the third race he lost $15 because 10 × 1.5 = <<10*1.5=15>>15 

He lost a total of $20 on the first and third races because 15 + 5 = <<15+5=20>>20 

He lost $9 that day because 11 - 20 = <<11-20=-9>>-9 

He lost an average of $3 per race because 9 / 3 = <<9/3=3>>3 

#### 3


The mistake is in the second line of the answer: on the third race, Johnny's dad lost $16.5, not $15, meaning he made $11 and lost $16.5 + $5 = $21.5. So, the answer is an average loss of $3.5 per race, not $3 per race (the answer in the dataset).












 Paper Title: You Can't Ignore Either: Unifying Structure and Feature Denoising for Robust Graph Learning.















Both authors contributed equally to this research.

Peking University
Beijing
China

youngtimmy@pku.edu.cn


[1]

Peking University
Beijing
China

mengjiahao@stu.pku.edu.cn



Huawei Cloud
Shenzhen
China

zhoumin27@huawei.com



Peking University
Beijing
China

yamingyang@stu.pku.edu.cn




Peking University
Beijing
China

yujwang, lxtpku@pku.edu.cn



Peking University
Beijing
China

yhtong@pku.edu.cn













§ ABSTRACT 
Recent research on the robustness of Graph Neural Networks (GNNs) under noises or attacks has attracted great attention due to its importance in real-world applications. Most previous methods explore a single noise source, recovering corrupt node embedding by reliable structures bias or developing structure learning with reliable node features. However, the noises and attacks may come from both structures and features in graphs, making the graph denoising a dilemma and challenging problem. In this paper, we develop a unified graph denoising (UGD) framework to unravel the deadlock between structure and feature denoising. Specifically, a high-order neighborhood proximity evaluation method is proposed to recognize noisy edges, considering features may be perturbed simultaneously. Moreover, we propose to refine noisy features with reconstruction based on a graph auto-encoder. An iterative updating algorithm is further designed to optimize the framework and acquire a clean graph, thus enabling robust graph learning for downstream tasks. Our UGD framework is self-supervised and can be easily implemented as a plug-and-play module. We carry out extensive experiments, which proves the effectiveness and advantages of our method. Code is avalaible at https://github.com/YoungTimmy/UGDhttps://github.com/YoungTimmy/UGD.








<ccs2012>
<concept>
<concept_id>10010147.10010257.10010293.10010294</concept_id>
<concept_desc>Computing methodologies Neural networks</concept_desc>
<concept_significance>500</concept_significance>
</concept>
</ccs2012>


[500]Computing methodologies Neural networks














You Can't Ignore Either: Unifying Structure and Feature Denoising for Robust Graph Learning
    Yunhai Tong
    Received ...; accepted ...
===========================================================================================





§ INTRODUCTION
 
By extending deep learning on graphs, graph neural networks have achieved great success and been applied to many important scenarios, such as social networks <cit.>, recommendation <cit.> and financial transaction <cit.>. Most GNNs <cit.> follow a message passing framework <cit.> to learn node representations for downstream tasks, relying on a favorable neighborhood aggregation including both structure connections and node features.

Unfortunately, real-world graphs are often noisy or incomplete due to error-prone data measurement or collection. For example, in a user-item graph, users may have imperfect profiles (feature missing) or misclick some unwanted items (error link) <cit.>. Malicious accounts may camouflage themselves, including adjusting their behaviors (feature camouflage) and connecting with several benign entities (structure camouflage) <cit.>. What is worse, recent studies show that GNNs are vulnerable to adversarial attacks <cit.>. Disturbing graph structure or node features would greatly degrade the performance of GNNs and may lead to severe consequences for critical applications requiring high safety and privacy <cit.>. 

To ensure robust graph learning, many previous efforts are devoted to refining corrupted node features or identifying problematic edges. According to the homophily assumption, AirGNN <cit.> designs an adaptive message passing scheme with residual connection to boost the resilience to abnormal features. MAGNET <cit.> automatically detects locally corrupted feature attributes and reconstructs robust estimations with sparsity promoting regularizer. Meanwhile, structure learning methods, such as ProGNN <cit.> and GSR <cit.>, aim to learn an optimal graph structure from noisy input and thus alleviate the structure noise or attacks.

However, most of these approaches assume a single noise source of edge perturbation or feature poisoning, purify the graph structures based on feature similarity, or refine node features based on reliable structures. They would face heavy performance degradation when feature noise and structure perturbation exist simultaneously. In addition, we find that simply combining the two processes into a pipeline cannot achieve a satisfactory performance and sometimes even gets worse (detail at Sec <ref>). Therefore, reconciling this dilemma remains an unexplored challenge.

To unravel the deadlock, we present a Unified Graph Denoising framework (termed as UGD) that considers both structure and feature noises. Specifically, we propose to conduct structure denoising by developing high-order neighborhood proximity instead of computing the directed similarity of pair nodes, which may be easily disturbed by feature noise. For feature denoising, we leverage a graph auto-encoder to reconstruct node features guided by local smoothness. Furthermore, an iterative updating algorithm is adopted to infer the preserved edges and jointly optimize the node features step by step.

We conduct extensive experiments on various datasets to evaluate the proposed UGD framework. The results validate its superior performance with noise and poisoning attacks over state-of-the-art methods. In particular, numerical comparison on downstream node classification tasks shows that UGD improves accuracy by 4.0% on the Citeseer dataset and 6.7% on the AComp dataset over the best previous methods, respectively. Moreover, detailed analyses under different ratios of perturbations and visualizations are also provided to show the strength and advantages of our UGD.




§ THE PROPOSED UGD FRAMEWORK




 §.§ Problem Formulation

Let 𝒢=(𝒱,ℰ) be an graph with node set 𝒱 and edge set ℰ. The feature matrix of the graph is X ∈ℝ^n × d, where n = |𝒱| denotes the number of nodes and d denotes the number of input features. A ∈{0,1}^n × n is the adjacency matrix, where A_uv=1 means the node u and node v are connected. N(u) is the direct neighbors of node u.
A graph data can also be denoted as 𝒢 = {ℰ,X}. 

Given a noisy or poisoned graph 𝒢 = {ℰ, X}, the unified denoising problem in this paper is to produce an optimized graph 𝒢^* = {ℰ^*, X^*} considering both structure and feature, and 𝒢^* will be used for the following downstream graph learning tasks. The overall framework of UGD is illustrated in Figure <ref>.



 §.§ High-order Neighborhood Proximity


Most existing works in graph structure learning focus on direct relations between interconnected nodes, with metrics like similarity measures or attention mechanisms to remove anomalous edges. However, in the presence of nodes with high feature noise, these pairwise metrics may become less reliable, potentially leading to the erroneous removal of normal edges.  
An intuitive way is to make a node observe more neighborhood information and keep up with more friends. From the perspective of a central node, its identity can be better represented by the neighbors and thus be resistant to feature disturbance. Thus, we define the neighborhood prototype of node u as an aggregating vector P_u of its neighbors with a permutation-invariant readout function. For simplicity, we apply the mean pooling as P_u = ∑_v ∈ N(u)x_v/|N(u)|.

Subsequently, for each edge (u, v) in the noisy graph 𝒢, we introduce a function sim(·,·) to measure high-order neighborhood proximity D_(u,v) between node v and prototype P_u, and distinguish whether the interaction of edge (u, v) is positive or negative. 
In this work, the cosine similarity score D_(u,v)=P_u· x_v/P_u_2x_v_2 is adopted for its better performance. Considering the asymmetry of D_(u,v) and D_(v,u) in undirected graphs, the smaller of the two values is retained as weight to filter more noise and formulated as:

    weight_(u,v) = weight_(v,u)= min {D_(u,v),D_(v,u)},

where weight_(u,v) goes beyond simple consideration of relations between adjacent nodes, incorporating the features of second-order neighbors for a more robust judgment of noisy edges. A threshold θ is employed to selectively remove edges with lower weights, thereby preserving a credible graph structure.





 §.§ Feature Reconstruction

To alleviate the impact of abnormal features and achieve local smoothness, prior approaches <cit.> based on graph signal denoising often formulate it as a complex convex optimization problem. Nevertheless, graph anomaly detection methods, such as CONAD <cit.>, have observed that feature compression and reconstruction can effectively identify anomalous nodes.  
Motivated by this, we develop a graph auto-encoder to ensure that the reconstructed features deviate from the original noise while maintaining similarity with its neighbors. In addition, we introduce a residual connection, combining the reconstructed results with the original features in a weighted manner to form the denoised features. With a fundamental reconstruction loss function ℒ_recon, the forward propagation process of node feature X can be defined as follows:

    X̂=β· X + (1-β) · Dec(Enc(X |ℰ ) |ℰ ), 
    ℒ_recon=|| X̂-X||_2 = 1/| 𝒱|∑_v ∈𝒱 || x̂_v - x_v ||_2,

where the encoder Enc and decoder Dec are two-layer GCNs <cit.>, X̂ is the denoised feature and β is a residual weight parameter.

To better achieve local smoothness in the reconstructed features, we also introduce a neighborhood smoothness loss (ℒ_smooth) to enhance the representational capacity of denoised features. Formally, the objective is to force a neighborhood proximity:

    ℒ_smooth=1/2∑_v ∈𝒱∑_u∈ N(v) ||x̂_̂û/√(d_u)-x̂_̂v̂/√(d_v)||_2^2 = tr(X̂^TLX̂),

where L is the normalized Laplacian matrix derived on the structure of graph 𝒢, d_v is the degree of node v. 
The reconstructed feature X̂ can be optimized by minimizing the two losses:

    ℒ=ℒ_recon (X̂ | X)+γ·ℒ_smooth(X̂ | ℰ, X) ,

with γ being a balance weight of the reconstruction and smoothness.  




 §.§ Optimization with Iterative Updating

Since the graph structure and node features are interrelated, the overall objective function ℒ is hard to optimize. To address this, we build the optimization of structure denosing (SD) and feature denoising (FD) with an iterative updating (IU) algorithm.

In each iteration i, the SD-step aim to find a credible edge subset while minimizing the ratio of noisy edges. At this point, node features are fixed, and the edges are updated by our defined high-order neighborhood proximity in section <ref>:

    update   ℰ^i= { (u,v) |  (u,v) ∈ℰ^i-1 and  weight_(u,v)^i-1≥θ} ,

where ℰ^i is the updated structure.

Then the target is redirected to update node feature in FD-Step. Model parameters of the graph auto-encoder are optimized with the objective function in section <ref>:

    minimize  ℒ= ℒ_recon(X^i | X^0) + γ·ℒ_smooth(X^i | ℰ^i, X^0).

These two steps would alternate until the estimation ℰ^i achieves convergence that the difference between two iterations is not more than a stop number ϵ. 

With the designs above, we build our proposed UGD framework and purify the graph data to support subsequent learning of downstream tasks. 
UGD is self-supervised and can be easily implemented as a plug-and-play module for various models and tasks.



§ EXPERIMENTS




 §.§ Experimental Setup

Datasets. We evaluate the semi-supervised node classification performance of UGD on five widely used benchmark datasets <cit.>. The dataset statistics are summarized in Table <ref>.

Baselines. We compare the proposed UGD with representative GNNs and defense models, including: 
    (1) Vanilla GCN <cit.> that ignores noises.
    (2) Feature denoising methods: MAGNET <cit.> uses a sparsity promoting regularizer to recover robust estimations of features; AirGNN <cit.> and ElasticGNN <cit.> focus on local smoothness by combining l_1 and l_2 loss.
    (3) Structure learning methods: Pro-GNN <cit.> directly takes the adjacency matrix as parameters for learning; IDGL <cit.> and GSR <cit.> utilize metric learning to refine the structure.
    (4) Combination of feature denoising and structure learning methods, including different models and orders.


Settings.
For a fair comparison, we closely follow the experiment setting in previous works <cit.>. We employ attribute injection and metattack <cit.> to inject the feature noise and modify the graph structure, respectively. For UGD, the optimal threshold θ for structure denoising varies across different datasets. We employ a smaller threshold during the early iterations to prevent removing excessive edges due to noisy features. Other hyper-parameters are tuned from: (1) β: {0, 0.5}; (2) γ: {1e-5, 5e-4, 3e-4, 1e-3}; (3) learning rate η: {5e-4, 1e-3};  (4) stop number ϵ: {0, 2, 10, 30};
The backbone model for downstream node classification is a two-layer GCN <cit.> for all methods, with a learning rate of 0.01, weight decay of 1e-3, and trained for 100 epochs. Adam optimizer <cit.> is used for all experiments. Our methods are implemented using PyTorch and PyTorch Geometric <cit.>. All experiments are conducted on a machine with an NVIDIA 3090 GPU (24GB memory).







 §.§ Numerical Comparison


We evaluate the performance of all methods against two types of noise, with an injection of 10% structure noise and 50% feature noise into the data. The ratio settings are referenced from prior works <cit.>, but we are pioneering in simultaneously considering both of them. Tabel <ref> summarizes the results of downstream node classification tasks after denoising. We report the mean test accuracy, together with a standard deviation over five different random seeds (OOM means out of memory). We can make the following observations from the results:



    (1) The performance of feature denoising methods can be heavily affected by structure noise. Specifically, on the AComp and Coauthor datasets, the elaborate approaches perform even worse than vanilla GCN. Besides, structure learning methods have shown large fluctuations in performance with feature noise. GSR performs much worse than GCN on most datasets (except for Cora). These results demonstrate the necessity of considering both types of noise.
    
    (2) Simply connecting the methods cannot achieve a satisfying result. For example, AirGNN+GSR leads to performance degradation compared with a single AirGNN on all graphs, which confirms our motivation to develop a unified denoising framework.
    

    (3) With the iterative updating algorithm, UGD jointly optimizes structure and feature denoising, and significantly improves the performance of GNNs. Concretely, UGD consistently outperforms the best baselines across five datasets, obtaining improvements of 0.9%, 4.0%, 1.3%, 6.7%, and 2.2%, respectively.



 §.§ Ablation Study

To examine the contributions of different components in UGD, we also consider variants of UGD that serve as ablation: (1)w/o HNP, which focuses solely on features denoising without refining the graph structure, and (2)w/o FR, which applies structure denoising without feature reconstruction. (3)w/o IU, which removes the iterative updating algorithm and performs graph denoising in order of feature first (F+S) or structure first (S+F). 

In Table <ref>, the bottom part summarizes the results, from which we have two observations. First, the performance of all UGD variants drops distinctly when compared to the full model, demonstrating that each designed component contributes to the success of UGD. Second, UGD without high-order proximity still outperforms most feature denoising methods, and UGD without feature reconstruction is still comparable to most structure learning methods, verifying the effectiveness of our designs under mixed noises.




 §.§ Model Analysis

We have performed comprehensive analysis to verify the advantages and effectiveness of UGD. Limited by the page, we report main results on Citeseer as following:



  §.§.§ Performance on different feature noise ratios.

Maintaining a consistent structure noise ratio of 10%, we conduct experiments with varying feature noise ratios of 0% to 50% on Citeseer dataset. The experimental results are presented in Figure <ref>(a). In this study, we compare our UGD model with GCN and other GSL methods to analyze how feature noise affects the structure denoising capability of the models. It can be observed that when the feature noise ratio increases, the baselines' node classification accuracy drops fast. UGD consistently outperforms baselines with a significant margin. Concretely, the accuracy on Citeseer drops by 18% and 21% for GSR and IDGL, respectively. For UGD, the accuracy decreases from 0.57 to 0.55, resulting in a slight decay of 3.5%, indicating that UGD is more robust against feature noise.



  §.§.§ Performance on different structure noise ratios.
 Maintaining a consistent feature noise ratio of 50%, we also compare UGD with graph denoising methods with different ratios of structure noise varying from 0% to 25% on Citeseer dataset. As shown in Fig. <ref>(b), structure noise can severely affect the graph denoising methods. Nevertheless, our UGD's performance is superior to baselines in most cases, showing better robustness against structure noise.



  §.§.§ Efficiency.

Another advantage of UGD is its high efficiency in graph denoising, which could be essential for large graphs. We report the end-to-end runtime of denoising methods with the same epochs in Fig.<ref>(c). 
Compared to MAGNET+GSR, UGD can achieve a speedup of 5.0x and 3.5x on Cora and Citeseer, respectively. In fact, UGD performs faster than single MAGNET and GSR, showing potential in applications with high-efficiency requirements. Moreover, UGD is model-free and can be readily combined with scalable GNNs to handle large graphs.



§ CONCLUSION

In this paper, we focus on robust graph learning with challenging structure and feature perturbations and introduce a UGD framework to handle structure and feature denoising for downstream tasks simultaneously. In UGD, high-order neighborhood proximity and feature reconstruction are jointly optimized with an iterative updating algorithm. Experiments on various datasets demonstrate the significant effectiveness of UGD. Extensive analyses further verify UGD's advantages against different noise and high efficiency.












This work was supported by National Key R&D program of China under Grant No.2023YFC3807603.





ACM-Reference-Format













 Paper Title: MIST: A Simple and Scalable End-To-End 3D Medical Imaging Segmentation Framework.




MIST



A. Celaya et al.


The University of Texas MD Anderson Cancer Center, Houston TX 77030 Rice University, Houston TX 77005


MIST: A Simple and Scalable End-To-End 3D Medical Imaging Segmentation Framework
    
Adrian Celaya1,2 Evan Lim1 Rachel Glenn1 Brayden Mi1 Alex Balsells1,2 Tucker Netherton1 Caroline Chung1 Beatrice Riviere2 David Fuentes1

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


§ ABSTRACT 
Medical imaging segmentation is a highly active area of research, with deep learning-based methods achieving state-of-the-art results in several benchmarks. However, the lack of standardized tools for training, testing, and evaluating new methods makes the comparison of methods difficult. To address this, we introduce the Medical Imaging Segmentation Toolkit (MIST), a simple, modular, and end-to-end medical imaging segmentation framework designed to facilitate consistent training, testing, and evaluation of deep learning-based medical imaging segmentation methods. MIST standardizes data analysis, preprocessing, and evaluation pipelines, accommodating multiple architectures and loss functions. This standardization ensures reproducible and fair comparisons across different methods. We detail MIST's data format requirements, pipelines, and auxiliary features and demonstrate its efficacy using the BraTS Adult Glioma Post-Treatment Challenge dataset. Our results highlight MIST's ability to produce accurate segmentation masks and its scalability across multiple GPUs, showcasing its potential as a powerful tool for future medical imaging research and development.





§ INTRODUCTION

Medical imaging segmentation is a highly active area of research. Since its introduction in 2015, the U-Net architecture has received nearly 90,000 citations on Google Scholar <cit.>. Other architectures and frameworks like nnUNet have achieved state-of-the-art accuracy in several medical imaging benchmarks like the Brain and Tumor Segmentation (BraTS) and Medical Segmentation Decathlon (MSD) <cit.>. Since nnUNet's introduction in 2018, several innovative approaches for medical imaging segmentation have emerged. These innovations include new architectures like vision transformers <cit.> and loss functions like the boundary and generalized surface losses <cit.>. Despite these recent advances, there remains a lack of standardized tools for testing and evaluating different approaches for medical imaging segmentation. For example, several works like <cit.> report superior performance to nnUNet while <cit.> refutes these results. This inconsistency makes it difficult to evaluate and assess claims of state-of-the-art performance for new research in deep learning-based medical imaging segmentation.


To address this inconsistency, we propose the Medical Imaging Segmentation Toolkit or MIST, a simple, modular, and end-to-end deep learning-based medical imaging segmentation framework that allows researchers to seamlessly train, test, and evaluate existing and new methods on their data in a consistent and reproducible way. MIST has a rule-based data analysis and preprocessing pipeline that it makes available to multiple architectures, loss functions, and training parameters. Additionally, MIST has a standardized and flexible evaluation pipeline that computes multiple metrics like the Dice, Haudorff distance, and surface Dice for any user-defined segmentation class. This standardization of several pipelines allows for fair comparisons between methods and their effect on segmentation accuracy for any given dataset.



§ METHODS
 
In this section, we describe the required data format for MIST, the different pipelines in MIST, and the auxiliary features, such as evaluation, postprocessing, and test-time inference. We also provide the specific settings we use to train on the BraTS challenge data.



 §.§ Data Format
 
MIST uses the same file structure as the BraTS challenge datasets <cit.>. That is, MIST assumes that the train and test (optional) data directories are formatted so that each sub-directory corresponds to one patient and that each imaging modality (i.e., T1, T2, or CT) is its own NIfTI file. The labels should all be in a single NIfTI file in the same sub-directory for each patient. Once the dataset is in the correct format, the next step is to prepare a small JSON file containing the details of the dataset. These include information like the imaging modality, naming conventions for the images and masks, labels, and the final classes that we use for evaluation (i.e., labels 1, 2, and 3 for whole tumor). An example of this JSON file is provided in the documentation page https://mist-medical.readthedocs.io/en/latest/getting_started/here. MIST also provides support for converting datasets that are in CSV or MSD format. 



 §.§ Pipelines
 
MIST contains three main pipelines: Data Analysis, Preprocessing, and Training. These pipelines can all run at once with a single command or individually. We also provide several auxiliary pipelines for evaluation, postprocessing, and test-time inference.



  §.§.§ Data Analysis
 
The data analysis pipeline gathers parameters about the dataset including cropping to foreground, target spacing, patch size selection, normalization parameters, and is described further below. The results of this analysis are saved in a  file.



  
Cropping to Foreground Often, zeros or useless background information surround the region of interest. For example, zeros surround the region outside of the brain in the BraTS datasets, and the table and air surrounding the body in CT imaging do not provide relevant information. From a computational perspective, it is desirable to crop the images to the foreground to make them smaller. Additionally, cropping to the foreground also reduces the sparsity of the dataset. MIST crops each image by first windowing it with its 33rd and 99.5 percentile values, then using an Otsu filter to produce a foreground threshold. We then crop the image according to the binary mask resulting from this threshold (i.e., one if the voxel intensity is greater or equal to the threshold). To determine if we need to use cropping, we check to see if, on average, the resulting volume reduction is at least 20%.



  
Target Spacing To determine the target spacing for a dataset, we first look at the distribution of the voxel spacings along each axis. Our initial target spacing is the median spacing in each direction. However, suppose the resulting target spacing is anisotropic (i.e., the maximum and minimum spacing ratio is greater than three), then we take the 10th percentile spacing along the lowest resolution axis in that case. However, if resampling with this selected target spacing results in any image taking up more than 2 GB of memory, then we increase all components of our target spacing by 25%. We repeat this process until none of the resampled images take up more than 2 GB of memory. This last step is important because it conserves computational resources. Larger images take longer to load from disk and require significantly more resources to manipulate (i.e., patch extraction and augmentation).



  
Patch Size Selection We set the patch size for training from the median resampled image size by taking the nearest power of two less than or equal to each dimension in the median image size up to a maximum patch size. The default value for this maximum patch size is 256×256×256. However, users can adjust this maximum patch size based on their knowledge of their computing resources. 



  
Normalization Parameters If the dataset consists of MR or other non-CT images, then our windowing parameters are the 0.5 and 99.5 percentile values of either the entire image or the non-zero values of the image if, on average, the ratio of the number of non-zero to zero-valued voxels is less than 0.2. We use the mean and standard deviation with the same rules regarding the non-zero values for normalization. On the other hand, if we are dealing with CT images, then we compute the 0.5 and 99.5 percentile values, mean, and standard deviation of all of the voxels over the entire dataset corresponding to regions labeled as non-zero in the ground truth segmentation masks.




  
Other Checks In addition to the previous analysis steps, MIST checks to see that the header information in the images and masks agree and, if there are multiple modalities, that the headers in each image agree. If there is a disagreement between the headers, MIST prints a warning with the patient ID to the console and excludes the example from training. MIST also computes equal-weighted class weights for each class based on the amount of voxels belonging to each class over the entire dataset. These class weights are primarily informational but can be used for weighted loss functions. 



  §.§.§ Preprocessing
 
The preprocessing pipeline takes the information from the analysis pipeline to crop the images based on their foreground mask (if called for), reorient to right-anterior-inferior (RAI), resample to the target spacing, and window and normalize the intensity values. This pipeline can also compute the distance transform maps (DTMs) for each class in the ground truth masks or apply bias correction for MR images. Users also have the option to skip this pipeline if, for example, their NIfTI files contain preprocessed images already. Figure <ref> illustrates the workflow for this pipeline.





  
Reorient & Resample After optionally cropping or bias correcting each image, the preprocessing pipeline will reorient each image and mask to the RAI orientation. For resampling images to the target spacing, we use third-order spline interpolation. We apply one-hot encoding and use linear interpolation to each channel for masks. If the current image spacing is anisotropic, we first resample along the low-resolution axis using nearest-neighbor interpolation and then apply the appropriate interpolation along the other two axes <cit.>. We implement these operations using SimpleITK to optimize the speed and computational performance of these operations <cit.>. 



  
Window & Normalize Before normalizing the voxel intensities in each image, the preprocessing pipeline windows each image by either the image's 0.5 and 99.5 percentile or, if we are using CT images, by the precomputed windowing parameters. After windowing, MIST applies z-score normalization with either the image's mean and standard deviation or the precomputed normalization parameters for CT images. Again, in the case that the ratio of non-zero to zero-valued voxels is less than 0.2, then the mean and standard deviation are computed from the non-zero voxels only, and we multiply the resulting images by the non-zero mask (different from foreground mask) to preserve the zero-valued voxels.



  
Distance Transform Maps A DTM represents the distance between each voxel and an object's closest boundary or edge. The values in a DTM are positive on the exterior, zero on the boundary, or negative in the object's interior. As a result, they can provide helpful information during training. Several loss functions and shape-aware networks use DTMs to improve the accuracy of medical imaging segmentation methods <cit.>. MIST allows users to pre-compute the DTMs for each label in the ground truth segmentation masks. Each channel in the resulting DTM represents the DTM for that label (i.e., channel one is the DTM for label one). If a label does not exist in the image, the DTM will be a 3D array with all values being the diagonal distance of the image. Additionally, MIST allows users to output normalized DTMs whose values range from -1 to 1. 



  §.§.§ Training

Once preprocessing is complete, MIST uses a five-fold cross-validation to train and evaluate a model on the dataset. In other words, for each fold we set aside 20% of the data as an independent test set and use the other 80% of the data for training. By default, MIST sets aside 10% of the training data as a validation set, but users can adjust this percentage. Additionally, users have the option to adjust the number of folds and can even specify custom folds if, for example, they are running leave-one-institution-out cross-validation.


MIST's default behavior is to use all available GPUs. When MIST runs multi-GPU training, it uses data parallelism via PyTorch's DistributedDataParallel package <cit.>. Additionally, MIST uses the Nvidia Data Loading Library (DALI) to handle data loading, patch extraction, and random augmentation. This feature further optimizes its computational performance by running most of the operations on the GPUs. 



  
Network Architectures MIST supports six different network architectures. These architectures include nnUNet (default) <cit.>, U-Net <cit.>, Swin UNETR <cit.>, and PocketNet <cit.>. MIST also provides users with options like adding deep supervision <cit.> or variational autoencoder regularization <cit.> to most of these architectures. Other regularization features like L2 and L1 regularization and gradient clipping are available for all architectures. For more details about the supported architectures (and their differences), please refer to our https://mist-medical.readthedocs.io/en/latest/advanced_topics/#network-architecturesdocumentation page. Finally, MIST's modular design allows users to implement their own architectures and integrate them into MIST with minimal effort. 




  
Loss Functions Like with the choice of network architecture, MIST allows users to choose from a variety of loss functions like Dice with Cross Entropy (default) <cit.>, the Generalized Dice Loss <cit.>, and boundary-based loss functions like the Boundary, Hausdorff, and Generalized Surface losses <cit.>. Again, MIST's modular design allows researchers to implement custom loss functions and seamlessly use them within the MIST framework. 



  
Other Training Options Finally, MIST gives users the flexibility to choose different optimizers like Adam (default) <cit.>, stochastic gradient descent, and AdamW <cit.>. MIST also supports several learning rate schedulers, with a constant learning rate of 0.0003 being the default. Additional features including transfer learning are included in our documentation page https://mist-medical.readthedocs.io/en/latest/here.



  §.§.§ Auxiliary Pipelines




  
Evaluation Once training is complete, the evaluation pipeline launches to assess the accuracy of the predictions with respect to the final classes defined by the user in their dataset JSON file. By default, these metrics are the Dice similarity coefficient and the 95th percentile Hausdorff distance (HD95). Other metrics like the surface Dice and average surface distance are also available. The output of the evaluation pipeline is a CSV file called  where each row shows the accuracy for each class and metric for one patient. The last five rows give the mean, standard deviation, median, and 25th and 75th percentiles for all of the metrics and classes. The evaluation pipeline is also available as a stand-alone pipeline that can run on any set of predictions. 



  
Postprocessing Postprocessing is available as an optional stand-alone pipeline. This pipeline allows users to specify which labels to apply a user-specified set of operations. These operations include the removal of small objects, only keeping the top-k components (adjustable value of k), morphological cleaning, and filling holes with a specified value. Once the postprocessing pipeline finishes, the evaluation pipeline launches and compares the change in accuracy to the baseline results. MIST computes this change in accuracy as a weighted average of the specified metrics. 



  
Test-Time Inference Finally, MIST provides a stand-alone inference pipeline that takes as input one or more MIST models and a CSV or JSON file with the paths to test-time data. This pipeline writes test-time predictions to a user-specified output directory. 



 §.§ BraTS 2024 Training Protocols
 
Our choice of architecture is the Pocket nnUNet with deep supervision (two supervision heads) and with residual convolution blocks <cit.>. MIST automatically selects a patch size of 128×128×128. We use eight NVIDIA H100 GPUs with a batch size of 32 uniformly distributed across the GPUs. Additionally, we use L2 regularization with a penalty parameter equal to 1e-5. Our choice of loss function is the Dice with Cross Entropy loss. While the default number of epochs per fold for MIST is 1,000, we train for 15,000 epochs per fold. We use a cosine learning rate schedule with an initial learning rate of 0.001 <cit.>. Within each fold, we set aside 2.5% of the training data (27 patients) as a validation set to select the best model. Once training is complete, MIST runs inference on the challenge validation data using test time augmentation (flipping along each axis and averaging each prediction) for each model. Inference uses sliding windows with an overlap of 0.5 with Gaussian blending (σ=0.125). The predictions from all five models are averaged to produce a final prediction. All other options and hyperparameters are left at their default values. Please refer to MIST's https://mist-medical.readthedocs.io/en/latest/documentation for these values.



§ RESULTS
 


 §.§ Accuracy

We use MIST and the parameters described in Section <ref> to perform a five-fold cross-validation (CV) with the BraTS Adult Glioma Post-Treatment Challenge data <cit.>. Table <ref> gives the results of this five-fold CV in terms of the Dice score and HD95. Tables <ref> and <ref> give the accuracy of the predictions in the validation set in terms of the same metrics and their lesion-wise versions.









 §.§ Computational Performance

In addition to the accuracy of our predictions, we also report the scalability of MIST for different numbers of A100 and H100 GPUs for multiple batch sizes. Figures <ref> and <ref> show these results for A100 and H100 GPUs, respectively. In these figures, we see that MIST effectively scales for multiple GPUs. This scaling is especially evident for the H100 GPUs. It is also important to note that we observe a roughly two times speed up with the H100 GPUs, which is in line with what Nvidia reports in their comparisons between the two types of graphics cards. 







§ DISCUSSION

Our results indicate that models trained using MIST can produce accurate adult post-operative gliomas, achieving median Dice scores of at least 0.9 for all of the segmentation classes in our five-fold CV. We choose to submit the raw output of our MIST pipeline to the BraTS 2024 competition, omitting the use of post-processing. However, future submissions to the continuous evaluation and competitions will explore different post-processing strategies with MIST by using its built-in post-processing pipeline. Indeed, previous BraTS competitions cite replacing small objects belonging to the ET class with the SNFH label as beneficial <cit.>. Additionally, we may see improvements in accuracy by using information from DTMs using boundary-based losses like the Generalized Surface Loss. Finally, different architectures like the multi-grid inspired FMG-Net or vision transformer-based architectures like Swin UNETR are also available on MIST <cit.>. These architectures may also improve our results.

MIST is able to utilize multiple GPUs to speed up training efficiently. This efficient use of multiple GPUs is especially true for H100s, where we see a nearly six times speed up when going from one to eight GPUs. The ability to take advantage of multiple GPUs and compute nodes effectively will be essential to keep training times manageable as larger datasets like TotalSegmentor become available <cit.>. As the size and quality of publicly available 3D medical imaging datasets increase, the creation of foundation models for medical imaging segmentation will be beneficial to speed up training through transfer learning <cit.>. MIST already supports transfer learning with other MIST models. Given this capability and its scalability, MIST can be instrumental in the creation of future foundation models.

In conclusion, our results show that MIST is capable of achieving accurate segmentation results while also scaling well to multiple GPUs by effectively taking advantage of technologies like PyTorch's DDP and Nvidia's DALI loader. MIST's simplicity, accuracy, and scalability have the potential to make MIST a valuable tool in the development of foundation models for medical imaging tasks like segmentation and radiation treatment planning. MIST's modular design makes it easy for researchers to add new loss functions or architectures to its pipelines. This modularity and easy integration into MIST's existing pipelines allow for new research to be evaluated and compared to existing methods fairly and consistently. MIST is still under active development. We encourage researchers and the broader medical imaging community to contribute to its development by using it, providing feedback, or contributing new features. 



§ AVAILABILITY & ATTRIBUTION

MIST is an open-source package (Apache 2.0 license) and is available on https://github.com/mist-medical/MISTGitHub or https://pypi.org/project/mist-medical/PyPI. Please cite this manuscript, <cit.>, and <cit.> if you use MIST for your own publications.



§ ACKNOWLEDGMENTS

The Department of Defense supports Adrian Celaya through the National Defense Science & Engineering Graduate Fellowship Program. This research was partially supported by the Tumor Measurement Initiative through the MD Anderson Strategic Research Initiative Development (STRIDE), NSF-2111147, NSF-2111459, and NIH R01CA195524


splncs04




 Paper Title: Who should I trust? A Visual Analytics Approach for Comparing Net Load Forecasting Models.
































Who should I trust?A Visual Analytics Approach for Comparing Net Load Forecasting Models



Parts of this work were supported by the U.S. Department of Energy Solar Energy Technologies Office and the Office of Electricity Sensors Program, and performed jointly at the Pacific Northwest National Laboratory under Contract DE-AC05-76RL01830 and at the Lawrence Livermore National Laboratory under Contract DE-AC52-07NA27344).

    Kaustav Bhattacharjee1, Soumya Kundu2, Indrasis Chakraborty3, and Aritra Dasgupta1

1Department of Data Science, New Jersey Institute of Technology, USA. Email: {kb526, aritra.dasgupta}@njit.edu 
2Optimization and Control Group, Pacific Northwest National Laboratory, USA. Email: soumya.kundu@pnnl.gov
3Center for Applied Scientific Computing, Lawrence Livermore National Laboratory, USA. Email: chakraborty3@llnl.gov 

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



§ ABSTRACT 
Net load forecasting is crucial for energy planning and facilitating informed decision-making regarding trade and load distributions. However, evaluating forecasting models' performance against benchmark models remains challenging, thereby impeding experts' trust in the model's performance. 

In this context, there is a demand for technological interventions that allow scientists to compare models across various timeframes and solar penetration levels. 

This paper introduces a visual analytics-based application designed to compare the performance of deep-learning-based net load forecasting models with other models for probabilistic net load forecasting. 

This application employs carefully selected visual analytic interventions, enabling users to discern differences in model performance across different solar penetration levels, dataset resolutions, and hours of the day over multiple months. We also present observations made using our application through a case study, demonstrating the effectiveness of visualizations in aiding scientists in making informed decisions and enhancing trust in net load forecasting models.



visual analytics, AI/ML, net load forecasting




§ INTRODUCTION



Comparing multiple computational models for their performance is crucial for enhancing trust in model outcomes because it provides a basis for evaluating the reliability and consistency of each model <cit.>. By assessing how different models perform under various conditions and scenarios, stakeholders can better understand their strengths, weaknesses, and overall effectiveness. Such comparative analysis helps to identify the most suitable model for specific tasks or applications, thereby instilling confidence in the reliability of the chosen model. This becomes more important while predicting the net load of an electric grid, which is defined as the difference between total electricity demand and generation from behind-the-meter resources like solar and distributed generators, and is influenced by various factors such as weather conditions and time of day <cit.>. 

Accurate net load forecasting enables grid operators, policymakers, and energy providers to make informed decisions regarding energy trade, load distribution, and resource allocation. However, the rise of solar energy generation sources in residential settings has significantly impacted the performance of traditional net load forecasting models, highlighting the need for robust time-series forecasting techniques <cit.>. Collaborating with scientists, we developed a deep-learning model that integrates variables such as temperature, humidity, apparent power, and solar irradiance to achieve strong predictive performance and resilience in the face of missing data <cit.>. However, the model's sensitivity to seasonal variations and noisy inputs underscores the need for further exploration with domain experts. Hence, we developed an interactive tool that empowers users to investigate the model's performance across diverse time periods and input scenarios <cit.>.





However, while we initially recognized the importance of comparing the model's performance with traditional models, feedback from domain experts further emphasized its significance in building trust in model outcomes. Visual analytics can play a pivotal role here, as evidenced by prior research demonstrating its importance in enhancing trust in machine learning models <cit.>.

This is exemplified by the interactive tool we developed, which enabled domain experts to extract valuable insights concerning the model's sensitivity toward temperature and humidity. Moreover, recent discourse, as highlighted in <cit.>, emphasizes the critical role of visual analytics in fostering trust-augmented applications of artificial intelligence and machine learning (AI/ML) within the energy sector. Building upon this, we designed  , incorporating carefully selected visual analytic interventions. These interventions facilitate the comparison of multiple models across various parameters, including solar penetration levels, dataset resolutions, and different hours of the day, enhancing stakeholders' confidence in model performance.



The aim of our application is to build trust through model comparison, primarily comparing our net load forecasting model with a reference model. We enhanced our model to process inputs at varying resolutions and then devised the reference model, which generates predictions by averaging net load ground truths for the last 30 days at the same time point. Despite lacking predictive ML components, this reference model serves as a benchmark in various net load forecasting competitions, including the recent Net Load Forecasting Prize by the National Renewable Energy Laboratory (NREL) and the U.S. Department of Energy Solar Technologies Office (SETO) <cit.>. Through the web interface of , we were able to uncover patterns in the performance that help improve trust in the model outcomes. In this work, we identify the visual analytic tasks that are required to compare these models. These tasks can be extended to compare other similar net load forecasting models, enabling users to make well-informed decisions based on the outcomes of these models.

While previous research has predominantly focused on developing interfaces during the model development phase, tailored to assist model developers, our emphasis lies in the post-hoc evaluation of model performance, specifically addressing the needs of energy scientists and grid operators <cit.>. Other studies have explored the performance of probabilistic net load forecasting models through different visualization charts but have often lacked an integrated interactive interface <cit.>. In contrast, our interactive visual analytic tool provides carefully designed visual cues for comparing model performance across different factors like dataset resolutions and solar penetration levels, thus offering a novel approach in this domain.

In this paper, we first introduce the different models used in this work and the rationale behind choosing them. Subsequently, we outline the identified visual analytic tasks and detail the design decisions guiding the development of our interactive interface. This is followed by some of the observations made by our power scientist collaborator through that demonstrate its efficacy in comparing model performance across multiple facets. Finally, we conclude by sharing insights gained from this development and discussing some of the future research opportunities in this domain. Additionally, a short demonstration video is available https://bit.ly/trustvishere.



§ MODEL DESCRIPTION




We start with a deep learning-based probabilistic model tailored for net load forecasting in high behind-the-meter solar scenarios <cit.>. This model has three key components: a kernelized probabilistic forecasting (kPF) module, an autoencoder (AE), and a long short-term memory (LSTM) network. This model effectively captured complex temporal dependencies and uncertainties inherent in net load data, which is crucial for reliable forecasting in environments with high solar penetration. By incorporating kernel methods into probabilistic forecasting, the model handled non-linear relationships and captured subtle variations in net load influenced by solar energy fluctuations. The autoencoder component enhanced feature extraction and dimensionality reduction, facilitating the LSTM network's ability to capture long-term dependencies and predict future net load values accurately. Experimental results demonstrated the superior performance of this model compared to traditional forecasting models, showcasing its efficacy in addressing the challenges posed by high solar scenarios and advancing the state-of-the-art in net-load forecasting methodologies. 

This model was used in the Net Load Forecasting Prize competition hosted by NREL and SETO <cit.>. However, during this competition, we observed its underperformance on the data provided by the organizers, which had lower resolutions. We discovered that while the autoencoder component excelled with high-resolution datasets (such as 15-minute intervals), it struggled to effectively capture temporal dependencies in lower-resolution datasets (e.g., 1-hour intervals provided by the organizers). Consequently, this limitation led to subpar outcomes generated by the LSTM component. In light of this, we opted to remove the kPF and autoencoder components and instead developed a version of the model solely utilizing the LSTM component. Additionally, fine-tuning the number of layers in the LSTM component yielded significantly improved results in the competition.

The competition employed a reference model to assess model performance across various probability levels. As previously mentioned, this reference model simply utilizes historical input data from the past 30 days to generate probabilistic forecasts for a specific time point. This model can serve as the initial benchmark for assessing the effectiveness of other models. Therefore, in this work, we developed a reference model following the same principles. Since these forecasts are probabilistic in nature, we calculated the Continuous Ranked Probability Score (CRPS) for both the reference model and our model. Subsequently, we computed the Continuous Ranked Probability Skill Score (CRPSS) based on these CRPS scores, evaluating whether our forecast presents an improvement or deterioration compared to the reference forecast <cit.>. A positive CRPSS indicates that the forecast outperforms the reference forecast, whereas a negative value suggests inferior performance. Utilizing these CRPSS values, we compare the performance of our model against the reference model across multiple dates throughout the year, and present these values in . 



§ VISUAL ANALYTICS-BASED DESIGN


Our application employs a visual analytics-based design featuring coordinated views enhanced with visual cues to help users compare model performance. It combines interactive visualization with comparison metrics like CRPSS to improve trust in the model outcomes and also allows the users to probe the model and understand its performance across different solar penetration levels and months. In this section, we highlight the two tasks performed by and how its visual analytics-based design aids in executing these tasks:
T1: Compare model performance across different solar penetration levels and data resolutions: Model performance may vary across different solar penetration levels due to the increased variability in net load data caused by intermittent solar generation. LSTM-based models might struggle to capture and predict these dynamic behaviors accurately. Furthermore, datasets with higher resolutions, such as sub-hourly intervals, enable models to more effectively capture short-term fluctuations and dependencies. Hence, this task essentially involves comparing model performance at different solar penetration levels and data resolutions.
T2: Identify patterns across different timeframes: By analyzing performance across multiple timeframes, power scientists can assess the net load forecasting models' robustness and consistency in capturing both short-term fluctuations and long-term trends. This evaluation helps to identify whether a model's performance is consistent across various temporal scales or exhibits variability or biases at specific time periods. This task relates to identifying patterns in the model's performance across different months, various hours of the day, and different time periods.

Our application implements multiple coordinated views and components in order to fulfill these tasks. Next, we discuss the design of these views and components along with the rationale behind them:
Comparison View: As comparing model performance is the main objective of , we begin with the Comparison View, which utilizes modified box plots to compare models. Figure <ref>a depicts CRPSS values on the y-axis and different data resolutions on the x-axis. The box plots illustrate the distribution of CRPSS values, with the median typically exceeding zero, indicating superior performance of our model over the reference in most cases. Users can utilize the solar penetration level filter to compare performance across various levels (20%, 30%, 50%) (Figure <ref>c) (T1). However, based on initial feedback from domain experts, we recognized the importance of displaying the distribution of CRPSS values. Therefore, we integrated dots representing the CRPSS values, jittered along the x-axis to illustrate their distribution over the entire year. These dots align with insights from the box plot and additionally reveal instances where our model performs worse than the reference on certain dates. This led to the development of the Patterns View, where we can identify the timeframes where the model underperforms compared to the reference.

Patterns View: In this view, utilizes a heatmap to depict performance patterns for each month across different hours of the day (T2). The x-axis represents the 24 hours of the day (0-23), while the y-axis displays the months of the year (Feb - Dec) (Figure <ref>b). Each box in the heatmap denotes the average CRPSS value for each month at each hour, indicated by the color. Darker blues signify more positive CRPSS values, indicating the superior performance of our model compared to the reference at that time. Conversely, darker reds indicate more negative CRPSS values, signifying poorer performance of our model compared to the reference at that time. This diverging color scale aids in easily identifying performance patterns across different time points and facilitates the identification of instances where our model does not outperform the reference. Users can filter specific months to focus on particular time periods and analyze performance accordingly (Figure <ref>e). Additionally, based on feedback from domain experts, we implemented the Sidebar component to allow users to select specific date ranges and analyze performance patterns within those ranges.
Sidebar: The Sidebar facilitates user selection of start and end dates to filter results across all views, enabling focus on specific date ranges (T2) for comparing model performance within those periods (Figure <ref>d). Additionally, based on feedback from domain experts expressing interest in comparing model performance across all solar penetration levels simultaneously, offers a comparison mode toggle in the Sidebar. Enabling this mode updates both views to display box plots and heatmaps for all solar penetration levels side by side, aiding in point-to-point comparison and identification of performance patterns across different solar penetration levels.



§ RESULTS FROM A CASE STUDY

The efficacy of an application can be validated if the application is able to perform the intended tasks effectively. In this section, we show the results through a case study that demonstrates how can be used to compare and select net load forecasting models effectively.

This case study involved a power scientist with over 10 years of experience in power and grid systems. With expertise in nonlinear dynamics, large-scale networks, and distributed control, he played a crucial role in developing the model. His main objective was to assess the model's performance relative to the reference model across various time points and solar penetration levels, which are critical factors to consider before deploying it for a project. We informed him that we had integrated CRPSS values for both the model and the reference model for all dates throughout the year. He then accessed through a browser and examined the distribution of CRPSS values across different data frequencies at a 20% solar penetration level (Figure <ref>a). Notably, he discovered that the model performed better with the lower resolution dataset (1-hour), contrary to expectations for an LSTM-based model (T1). Upon further examination, he noted a marginal difference in the median CRPSS between high and low-resolution datasets (0.85 and 0.89, respectively). Therefore, he enabled the comparison mode through the Sidebar, enabling him to assess the model's performance as solar penetration increased. He noted that the model consistently performed well with the higher resolution dataset (15-min) across all other solar penetration levels, confirming the notion that our model excels with high-resolution datasets in high solar penetration scenarios (Figures <ref>b and  <ref>c). Next, he observed that although the median CRPSS was significantly above zero, the minimum value was negative. Upon inspecting the dots adjacent to the box plots in the Comparison View, he observed that while most dots clustered around the median line, there were a few outliers with negative CRPSS values.

Hence, the scientist sought to understand temporal patterns to identify if any specific hour of the day contributed to the negative values. Consequently, he navigated to the Patterns View within to analyze the CRPSS value distribution in the heatmap at a 20% solar penetration level (Figure <ref>d). Consistent with observations from the Comparison View, most heatmap boxes displayed varying shades of blue, indicating superior model performance compared to the reference across most months and hours. On closer look, the scientist noted a box-like pattern in the heatmap, revealing enhanced performance during morning hours (8 am to 10 am) from April to September, followed by a decline during midday (11 am to 4 pm), and another spike in the evening (4 pm to 6 pm) during these months. This pattern suggested that the model effectively captured diurnal variations in net load data and adjusted predictions accordingly (T2). While similar box-like patterns were evident across other solar penetration levels, the intensity of model performance varied with increased solar penetration. This insight holds practical significance for model selection during deployment, as it suggests the potential use of different models or model ensembles tailored to distinct times of the day, leveraging their respective performances on diurnal patterns. Thus, the power scientist was satisfied that could yield valuable insights regarding the model, aiding informed decision-making. A short demonstration video is available at https://bit.ly/trustvisthis link.


§ CONCLUSION

Efficient model selection for net load forecasting plays a pivotal role in energy planning and grid operations. In this work, we delve into this process, integrating visual analytics with input from domain experts to identify key tasks for comparative model selection. Subsequently, we translate these tasks into an interactive interface, enabling users to assess model behavior across various factors such as solar penetration levels, data resolution, and time of day.


Throughout this endeavor, we gained valuable insights into the model behavior and the challenges posed during the multi-way comparison of models. This collaborative effort underscored the need for interactive tools that facilitate the seamless translation of model insights into actionable decisions. We can argue that our application is a first step towards this direction. As a next step, our plan is to incorporate multiple net load forecasting models into the application and integrate additional metrics for effective comparison of their performance.

Looking ahead, we aim to enhance our application in several ways. Incorporating economic planning and analysis options will allow stakeholders to assess the cost-benefit ratio before model selection, thereby enhancing trust in the outcomes. Additionally, as the energy landscape evolves, our application's flexibility will be pivotal in effectively comparing and selecting forecasting models for real-world applications. In summary, our collaborative endeavor has produced a robust tool for multi-faceted model comparison and has paved the way for informed decision-making through visual analytics in energy planning.








abbrv-doi




