Paper Title: Graph Representation Learning via Causal Diffusion for Out-of-Distribution Recommendation.






Graph Representation Learning via Causal Diffusion for Out-of-Distribution Recommendation





    Chu Zhao
Northeastern University 

Shenyang, China 

chuzhao@stumail.neu.edu.cn
Enneng Yang
Northeastern University 

Shenyang, China 

ennengyang@stumail.neu.edu.cn
Yuliang Liang
Northeastern University 

Shenyang, China 

liangyuliang@stumail.neu.edu.cn
Pengxiang Lan
Northeastern University 

Shenyang, China 

2390222@stu.neu.edu.cn
Yuting Liu
Northeastern University 

Shenyang, China 

yutingliu@stumail.neu.edu.cn
Jianzhe Zhao
Northeastern University 

Shenyang, China 

zhaojz@swc.neu.edu.cn
Guibing Guo 
Northeastern University 

Shenyang, China 

guogb@swc.neu.edu.cn 
Xingwei Wang
Northeastern University 

Shenyang, China 

wangxw@mail.neu.edu.cn

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













§ ABSTRACT 
Graph Neural Networks (GNNs)-based recommendation algorithms typically assume that training and testing data are drawn from independent and identically distributed (IID) spaces. However, this assumption often fails in the presence of out-of-distribution (OOD) data, resulting in significant performance degradation. In this study, we construct a Structural Causal Model (SCM) to analyze interaction data, revealing that environmental confounders (e.g., the COVID-19 pandemic) lead to unstable correlations in GNN-based models, thus impairing their generalization to OOD data. To address this issue, we propose a novel approach, graph representation learning via causal diffusion (CausalDiffRec) for OOD recommendation. This method enhances the model's generalization on OOD data by eliminating environmental confounding factors and learning invariant graph representations. Specifically, we use backdoor adjustment and variational inference to infer the real environmental distribution, thereby eliminating the impact of environmental confounders. This inferred distribution is then used as prior knowledge to guide the representation learning in the reverse phase of the diffusion process to learn the invariant representation. In addition, we provide a theoretical derivation that proves optimizing the objective function of CausalDiffRec can encourage the model to learn environment-invariant graph representations, thereby achieving excellent generalization performance in recommendations under distribution shifts. Our extensive experiments validate the effectiveness of CausalDiffRec in improving the generalization of OOD data, and the average improvement is up to 10.69% on Food, 18.83% on KuaiRec, 22.41% on Yelp2018, and 11.65% on Douban datasets. Our implementation code is available at <https://github.com/user683/CausalDiffRec>.



Graph Neural Networks, Out-of-distribution, Invariant Learning, Recommender Systems















Graph Representation Learning via Causal Diffusion for Out-of-Distribution Recommendation





    Chu Zhao
Northeastern University 

Shenyang, China 

chuzhao@stumail.neu.edu.cn
Enneng Yang
Northeastern University 

Shenyang, China 

ennengyang@stumail.neu.edu.cn
Yuliang Liang
Northeastern University 

Shenyang, China 

liangyuliang@stumail.neu.edu.cn
Pengxiang Lan
Northeastern University 

Shenyang, China 

2390222@stu.neu.edu.cn
Yuting Liu
Northeastern University 

Shenyang, China 

yutingliu@stumail.neu.edu.cn
Jianzhe Zhao
Northeastern University 

Shenyang, China 

zhaojz@swc.neu.edu.cn
Guibing Guo 
Northeastern University 

Shenyang, China 

guogb@swc.neu.edu.cn 
Xingwei Wang
Northeastern University 

Shenyang, China 

wangxw@mail.neu.edu.cn

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





§ INTRODUCTION


Graph Neural Networks <cit.>, due to their exceptional ability to learn high-order features, have been widely applied in recommendation systems.

GNN-based recommendation algorithms <cit.> learn user and item representations by aggregating information from neighboring nodes in the user-item interaction graph and then computing their similarity to predict user preferences. In addition, researchers have introduced various other techniques to continuously improve GNN-based recommendation algorithms. For example, integrating attention mechanisms <cit.> with knowledge graphs <cit.> led to improving recommendation accuracy. Furthermore, the introduction of contrastive learning aims to improve the robustness of recommendation algorithms <cit.>.

Despite the significant progress these methods have made in improving recommendation accuracy, most of them assume that the distribution of the test dataset and the training dataset is independently and identically distributed (IID) and focus on enhancing recommendation performance based on this assumption. Unfortunately, methods based on the above assumption fail to generalize excellent recommendation performance to out-of-distribution (OOD) data <cit.>, where the distribution of the test data significantly differs from that of the training data. In Figure <ref>, we present a simple example to illustrate how the popularity of medical supplies changes with the environmental factor of COVID-19. Specifically, due to the pandemic, people may be required by the government to stay at home, leading to reduced outdoor activities. During this period, while purchasing medical supplies, people may also increase their demand for fitness equipment and electronic products. The recommender system might learn that users who purchase masks also frequently buy fitness equipment and electronic products, significantly increasing the popularity of these items. This correlation is driven by the common factor of the `pandemic' rather than a direct causal relationship between the items. When the pandemic ends (i.e., the environmental factors change), the popularity of masks decreases, and people's demand for fitness equipment and electronic products diminishes. The recommender system, relying on the unstable correlations learned during the training phase, may incorrectly recommend fitness equipment and electronic products to users who purchase masks, resulting in poor performance under the new distribution. Furthermore, we construct IID and OOD test sets on the Yelp2018 dataset to illustrate the changes in model performance. As shown in Figure <ref>, LightGCN's <cit.> performance on the OOD dataset, compared to the IID dataset, shows an average decline of 29.03% across three metrics. Such issue of GNN-based models lacking robustness on OOD datasets inspires us to propose a recommendation framework with strong generalization capabilities for distribution shifts.



Several works <cit.> have focused on improving the generalization ability of recommender systems on OOD datasets. Some researchers use causal inference to counteract shifts in the data distribution. For example, CausPref <cit.> 
is based on the NeuMF<cit.> method and designs invariant user preference causal learning and anti-preference negative sampling methods to improve model generalization.
COR <cit.>  uses the
Variational Auto-Encoder for
causal modeling by incorporating an encoder to infer unobserved
user features from historical interactions.
Although the aforementioned methods have improved the performance of recommendation models on out-of-distribution datasets to some extent, they are not specifically designed for GNNs, making direct migration to GNN-based methods difficult. 

Other researchers employ techniques such as graph contrastive learning and graph data augmentation to enhance the robustness of GNN-based recommendation algorithms, such as SGL <cit.>, SimGCL <cit.>, and LightGCL <cit.>. These methods mainly address noise or popularity bias in the data, but they struggle to achieve good recommendation performance when the test data distribution is unknown or has multiple distributions. As shown in the experimental results of Tables <ref>, these methods perform worse than other baselines in other types of data distribution shifts. In recent literature, few GNN-based methods <cit.>  have been proposed to improve generalization when faced with multiple distributions. However, these methods lack solid theoretical support. 

Given these limitations, there is an urgent need to design theoretically grounded GNN-based methods to address distribution shifts. In this paper, we use invariant learning to improve the generalization of the OOD dataset. 
Utilizing insights from the prior knowledge of environment distribution and invariant learning <cit.>  enhances model stability across varied environments. This is achieved by acquiring invariant representations, which in turn boosts the model's generalization capabilities and overall robustness.
However, designing models based on invariant learning still faces the following two challenges:

    
   1) How to infer the distribution of underlying environments from observed user-item interaction data? 
    
   2) How to recognize environment-invariant patterns amid changing user behaviors and preferences?


To address the aforementioned challenges, in this paper, we first construct the Structural Causal Model (SCM) to analyze the data generation process in recommender systems and investigate the learning process of GNN-based recommendation algorithms under data distribution shifts. We conclude that latent environmental variables can lead GNN-based algorithms to capture unstable correlations related to the environment, which is the key reason for the failure of GNN-based models to generalize on OOD data. Furthermore, we propose a novel approach called graph representation learning via Causal Diffusion (CausalDiffRec) for OOD recommendation, which leverages causal inference to eliminate unstable correlations caused by environmental variables. This approach aims to learn invariant representations across different environments, thereby achieving OOD generalization in recommender systems. Specifically, CausalDiffRec consists of three main components: an environment generator, an environment inference, and a diffusion module. The environment generator is used to create K significantly different graphs to simulate data distributions under various environments. The environment inference module then infers the environment components from these generated graphs and uses them as input for the diffusion reverse stage to guide invariant graph representation learning. Finally, we provide theoretical proof that CausalDiffRec, under the conditions of invariant learning theory, can identify invariant graph representations across different environments, thereby improving generalization performance on OOD data. The contributions of this paper are concluded as follows:

    
  * Causal Analysis. We construct the SCM and analyze the generalization ability of GNN-based recommendation models on OOD data from the perspective of data generation. Based on our analysis and experimental results, we conclude that environmental confounders lead the model to capture unstable correlations, which is the key reason for its failure to generalize under distribution shifts.
    
  * Methodology. We propose a novel GNN-based method, CausalDiffRec, for OOD recommendation. CausalDiffRec primarily consists of three modules: environment generation, environment inference, and the diffusion module. The environment generation module simulates user data distributions under different conditions/environments; the environment inference module employs causal inference and variational approximation methods to infer the environment distribution; and the diffusion module is used for graph representation learning.
    Our theoretical analysis guarantees that optimizing the objective function of CausalDiffRec enables the model to achieve great generalization. 
    
  * Experimental Findings. We constructed three common types of distribution shifts across four datasets and conducted comparative experiments. The experiments demonstrate that CausalDiffRec consistently outperforms baseline methods. Specifically, when dealing with OOD data, CausalDiffRec exhibits enhanced generalization capabilities, achieving a maximum metric improvement rate of 36.73% compared to the baseline methods.






§ PRELIMINARY



 §.§ GNN-based Recommendation


Given the observed implicit interaction matrix ℛ∈{0,1}^m × n, in which 𝒰 = {u_1, u_2, …, u_m} represents the set of users, ℐ = {i_1, i_2, …, i_n} represents the set of items, m and n denote the number of users and items, respectively. For the elements in the interaction matrix, r_ui = 1 indicates an interaction between user u and item i, otherwise 0. In GNN-based recommendation algorithms, the user-item interaction matrix ℛ is first transformed into a bipartite graph G={𝒱,ℰ}. We employ 𝒱 to represent the node set and ℰ={(u,i)|u ∈𝒰, i ∈ℐ,r_ui = 1} denotes the edge set.
Given a user-item interaction graph 𝒢_u and the true user interactions y_u with respect to (w.r.t) user u, the optimization objective of GNN-based methods can be expressed as:

    arg min_θ 𝔼_(𝒢_u,y_u) ∼ P(G,Y)[l(f_θ(𝒢_u;θ),y_u)],

where f_θ(·) is a learner that learns representations by aggregating high-order neighbor information from the user-item interaction graph. l denotes the loss function and P(G,Y) represents the joint distribution of the interaction graph G and true label Y. 





 §.§ Denoising Diffusion Probabilistic Models


Denoising Diffusion Probabilistic Models (DDPM) <cit.> have been widely used in the field of image and video generation. The key idea of DDPM is to achieve the generation and reconstruction of the input data distribution through a process of gradually adding and removing noise. It leverages neural networks to learn the reverse denoising process from noise to real data distribution.


The diffusion process in recommender system models the evolution of user preferences and item information through noise addition and iterative recovery. Initially, data 𝐱_0 sampled from q(𝐱) undergo a forward diffusion to generate noisy samples 𝓍_1, …, 𝐱_T over T steps. Each step adds Gaussian noise, transforming the data distribution incrementally <cit.>:

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

where β_t ∈ (0, 1) controls the level of the added noise at step t. In the reverse phase, the aim is to restore the original data by learning a model p_θ to approximate the reverse diffusion from 𝐱_T to 𝐱_0. The process, governed by p_θ(𝐱_t-1|𝐱_t), uses the mean μ_θ and covariance Σ_θ learned via neural networks:

    p_θ(𝐱_t-1|𝐱_t) = 𝒩( 𝐱_t-1; μ_θ(𝐱_t, t), Σ_θ(𝐱_t, t) ).

The reverse process is optimized to minimize the variational lower bound (VLB), balancing the fidelity of reconstruction and the simplicity of the model <cit.>:

    ℒ_VLB    = 𝔼_q(𝐱_1:T|𝐱_0)[ ∑_t=1^T D_KL(q(𝐱_t-1|𝐱_t, 𝐱_0) || p_θ(𝐱_t-1|𝐱_t)) ] 
         - log p_θ(𝐱_0|𝐱_1),

where D_KL denotes the Kullback-Leibler (KL) divergence. Following <cit.>, we mitigate the training instability issue in the model by expanding and reweighting each KL divergence term in the VLB with specific parameterization. Therefore, we have the following mean squared error loss:

    ℒ_simple = 𝔼_t, 𝐱_0, ϵ_t[ ϵ_t - ϵ_θ( √(α)_t 𝐱_0 + √(1 -α_t)ϵ, t ) ^2 ],

where ϵ_t ∼𝒩 (0,I) is the noise for injection in forward process, ϵ_θ(·) denotes a function approximator that neural networks can replace, and α_t = ∏_t=1^T 1 - β_t.
This framework allows the model to effectively learn noise-free representations and improve recommendation accuracy.





§ METHODOLOGY

In this section, we first construct SCM and identify environmental confounders as the key reason for the failure of GNN-based models to generalize on OOD (out-of-distribution) data. Subsequently, we introduce the variational inference to infer the true distribution of the environment. We use the diffusion model to learn the representation based on invariant learning. Finally, we provide rigorous theoretical proof of CausalDiffRec that can achieve great generalization. The model framework is illustrated in Figure <ref>.




 §.§ SCM of GNN-based Recommendation


To explore the reasons behind the failure of GNN-based models to generalize on OOD data, we follow previous works <cit.> <cit.> and first construct the SCM for data generation and data modeling in recommendation systems, as shown in Figure <ref> (a) and Figure <ref> (b). We find that environmental confounding factors are the key reason for the generalization failure of GNN-based methods. Finally, we design an intervention model in Figure <ref> (c) to eliminate the impact of environmental confounding factors.



  §.§.§ Causal View in GNN-based Recommendation

In Figure <ref>, the three causal relations are derived from the definitions of data generation. The detailed causal analysis behind them is presented as follows:




  * E denotes the unobserved environmental factors (i.e., sudden hot events or policies). G and Y denote the user-item interaction graph and the true label, respectively. I is the invariant attribute of users and items unaffected by environmental factors, such as user gender and item category information. Previous work <cit.> has indicated that leveraging these invariant features can effectively enhance the model's generalization capability in OOD environments.

  * E → G. The direct influence of the environment on user-item interactions can be defined as P(G|E). For example, if the environmental variable is the weather, users might interact more frequently with warm clothing in a cold environment.

  * G → Y. This represents the influence of the user-item interaction graph G on the user behavior label Y. The GNN-based recommendation model Y = f_θ(G) defines the relation. When the model parameters θ are fixed, the mapping between G and Y is also deterministic.

  * I → Y. This indicates that the invariant user-item attributes directly influence the user behavior label Y. For example, a user may consistently prefer to dine at a particular restaurant, and attributes such as the restaurant's location and name generally do not change.

  * E → Y. This denotes that the environment directly influences the user behavior label Y, independent of user interactions between users and items. For example, in a specific holiday environment, users may be more inclined to purchase holiday-related items, regardless of whether they have interacted with them. 
 
In real-world scenarios, training data is collected from heterogeneous environments. Therefore, the environment directly influences the distribution of the data and the prediction result, which can be explicitly represented as P(Y, G|E)=P(G|E)P(Y|G, E). If we employ D_tr(E) to represent the training data distribution for unobserved environments, the GNN-based model, when faced with OOD data, can rewrite Eq. (<ref>) as:

    arg min_θ 𝔼_e ∼ D_tr(E), (𝒢_u,y_u) ∼ P(G,Y|E=e)[ l (f_θ(𝒢_u;θ), y_u)|e],

Eq. (<ref>) shows that environment E affects the data generation used for training the GNN-based recommendation model.





  §.§.§ Confounding Effect of E

Figure <ref> (a) and Figure <ref> (b)  illustrate the causal relationships in data generation and model training for graph-based recommendation algorithms. E acts as the confounder and directly optimizing P(Y|G) leads the GNN-based recommendation model to learn the shortcut predictive relationship between 𝒢_u and y_u, which is highly correlated with the environment E. During the model training process, there is a tendency to use this easily captured shortcut relationship to model user preferences. However, this shortcut relationship is highly sensitive to the environment E. When the environment of the test set is different from that of the training set (i.e., D_tr(E) ≠ D_ts(E)), this relationship becomes unstable and invalid. The recommendation model that excessively learns environment-sensitive relationships in the training data will struggle to accurately model user preferences when faced with out-of-distribution data during the testing phase, resulting in a decrease in recommendation accuracy.



  §.§.§ Intervention

Through the above analysis, we can improve the generalization ability of GNN-based recommendation models by guiding the model to uncover stable predictive relationships behind the training data, specifically those that are less sensitive to environmental changes. Thus, we can eliminate the influence of environmental confounders on model predictions.
Specifically, we learn stable correlations between user item interaction 𝒢_u and ground truth y_u by optimizing P_θ(Y | do(G)) instead of P_θ(Y | G). In causal theory, the do-operation signifies removing the dependencies between the target variable and other variables. As shown in Figure <ref> (c), by cutting off the causal relationship between the environment variables and the user interaction graph, the model no longer learns the unstable correlations between 𝒢_u and y_u. The do-operation simulates the generation process of the interaction graph G, where environmental factors do not influence the user-item interactions. This operation blocks the unstable backdoor path G ← E → Y, enabling the GNN-based recommendation model to capture the desired causal relationship that remains invariant under environmental changes.  

Theoretically, P_θ(Y|do(G)) can be computed through randomized controlled trials, which involve randomly collecting new data from any possible environment to eliminate environmental bias. However, such physical interventions are challenging. For instance, in a short video recommendation setting, it is impossible to expose all short videos to a single user, and it is also impractical to control the environment of data interactions. 
In this paper, we achieve a statistical estimation of P_θ(Y|do(G)) by leveraging backdoor adjustment.
The derivation process is shown as follows:

    P_θ(Y|do(G)) 
       
        = ∑_e P_θ(Y|do(G), E=e, I)P_θ(E=e|do(G))P_θ(I) 
       
        = ∑_e  P_θ(Y|G, E=e, I)P_θ(E=e|do(G))P_θ(I) 
       
        = ∑_e  P_θ(Y|G, E=e, I)P_θ(E=e)P_θ(I)   
       
        = ∑_e P_θ(Y|G, E=e, I)P_θ(E=e, I)   
       
        = 𝔼_e ∼ D_tr(E)[P_θ(Y|G, E, I)],

through the aforementioned backdoor adjustment, the influence of the environment E on the generation of G can be eliminated, enabling the model to learn correlations independent of the environment. However, in recommendation scenarios, environmental variables are typically unobservable or undefined, and their prior distribution P(E=e) cannot be computed. Therefore, directly optimizing the Eq. (<ref>) is challenging.




 §.§ Model Instantiations



  §.§.§ Environment Inference

This work introduces a variational inference method and proposes a variational inference-based environment instantiation mechanism. The core idea is to use variational inference to approximate the true distribution of environments and generate environment pseudo-labels as latent variables. The following tractable evidence lower bound (ELBO) can be obtained as the learning objective:

    log P_θ(Y|do(G)) ≥ℒ_envInf = 𝔼_Q_ϕ(E|G, I)[logP_θ(Y|G, E, I)] 
    - D_KL (Q_ϕ(E|G, I)∥ P_θ(E)),

where Q_ϕ (E|G, I) denotes environment estimation, which draws samples from the true distribution of the environment E. D_KL represents the Kullback–Leibler (KL) divergence of the volitional distribution Q_ϕ(E|G, I) and the prior distribution P_θ(E).
P_θ(Y|G, E, I) is the graph representation learning module that employs the user-item interaction graph and the node attributes of users and items as input to learn invariant representations. Section <ref> will provide a detailed introduction to the graph representation learning module. We present the derivation process of Eq. (<ref>) as follows:

Taking the logarithm on both sides of Eq. (<ref>) and according to Jensen’s Inequality, we have:

    logP_θ(Y|do(G)) 
        
        =log 𝔼_e ∼ D_tr(E)[P_θ(Y|G, E, I)] 
        =
        log∑_e P_θ(Y|G, E=e, I)P_θ(E=e, I)Q_ϕ(E=e|G, I)/Q_ϕ(E=e|G, I)
       ≥∑_e Q_ϕ(E=e|G, I) 
           log P_θ(Y|G,E=e,I)
        P_θ(E=e,I)1/Q_ϕ(E=e|G, I)
       
        = ∑_e [Q_ϕ(E=e|G, I) log P_θ(Y|G,E=e, I) -
           log Q_ϕ(E=e|G, I)P_θ(E=e,I)/Q_ϕ(E=e|G, I)] 
       
        = 𝔼_Q_ϕ(E=e|G, I)[logP_θ(Y|G, E=e, I)] 
       -  D_KL (Q_ϕ(E=e|G, I)∥ P_θ(E=e)).




  §.§.§ Invariant Pattern Recognition Mechanism


Following previous invariant learning studies <cit.>, this work proposes an invariant pattern recognition mechanism to explore invariant graph representations, encouraging the model to learn invariant correlations under distribution shifts. We make the following assumption:


Assumption: For a given user-item interaction graph (i.e., data distribution D), these interaction data are collected from K different environments E. User behavior patterns exist independently of the environment and can be used to generalize out-of-distribution user preference prediction. There exists an optimal invariant graph representation learning F^*(·) satisfying:

    
  * Invariance Property. ∀ e ∈ D(E), P_θ(Y|F^*(G),E=e, I) = P(Y|F^*(G), I). 
    
    
  * Sufficiency Condition. Y = F^*(G) + ϵ, ϵ E, where  indicates statistical independence and ϵ is random noise.


The invariance property assumption indicates that a graph representation learning model exists capable of learning invariant user-item representations across different data distribution environments. The sufficiency condition assumption means that the learned invariant representations enable the model to make accurate predictions.



  §.§.§ Invariant Representation Learning
  This section mainly consists of an environment generator, a diffusion-based graph representation learning module, and a recommendation module. Next, we will detail how they collaborate to enhance the generalizability of GNN-based models on OOD data and improve recommendation accuracy.

Environment Generator.
In real-world recommendation scenarios, training datasets are collected in various environments. However, for a single user-centric interaction graph, the training dataset comes from a single environment. We need to learn environment-invariant correlations from training data originating from different environments to achieve the generalization capability of GNN-based recommendation models under distribution shifts. To circumvent this dilemma, this paper designs an environment generator g_ω_k(·) (1 ≤ k ≤ K ), which takes the user's original interaction graph G as input and generates a set of K interaction graphs {G_i}_i=1^K to simulate training data from different environments. The optimization objective is expressed as follows:

    ℒ_generator = [ Var(ℒ(g_ω_k(G)): 1 ≤ k ≤ K )],

where Var(·) denotes the variance and ℒ(·) is the loss function. Following existing work <cit.>, we modify the graph structure by adding and removing edges. Given a Boolean matrix B_k, the adjacency matrix A of the graph, and its complement A', the k-th generated view for the original view is A_k = A + B_k ⊙ (A - A'). Since B_k is a discrete matrix and not differentiable, it cannot be optimized directly. To address this issue, we borrow the idea from <cit.> and use reinforcement learning to treat graph generation as a decision process and edge editing as actions. Specifically, for view k, we consider a parameter matrix θ_k = {θ_nm^k}. For the n-th node, the probability of exiting the edge between it and the m-th node is given by:

    h(a_nm^k) = exp(θ_nm^k)/∑^m'=m_m'=1exp(θ_nm'^k).


We then sample s actions {b_nmt^k}_t=1^s from a multinomial distribution M(h(α_n1^k), …, h(α_nm^k)), which give the nonzero entries in the n-th row of B_k. The reward function R(G_k) can be defined as the inverse loss. We can use the reinforcement algorithm to optimize the generator with the gradient:

    ∇_θ_klog h_θ_k(A_k) R(G_k),

where θ_k is the model parameters  and h_θ_k(A_k) = ∏_n ∏_t=1^s h(b_nmt^k). Optimizing Eq. (<ref>) ensures that the generated graphs have large differences.



Causal Diffusion. Given the generated interaction graph G_k=(A_k, I) where A is the adjacency matrix, and I is the feature matrix of users or items, instead of directly using G_k as input for diffusion, we use the encoder from the Variational Graph Autoencoder (VGAE) to compress G_k to a low-dimensional vector 𝐱^0_k ∼𝒩(μ_k,σ_k) for subsequent environment inference and graph invariant representation leanring. The encoding process is as follows:
 
    q_ψ(𝐱^0_k|A_k,I) = N(𝐱^0_k|μ_k,σ_k),

where μ_k = GCN_μ(A_k, I) is matrix of mean vectors and σ_k = GCN_σ(A_k, I) denotes standard deviation. The GCN(·) is the graph convolution network in the graph variational autoencoder. According to the reparameterization trick, x^0_k can be calculated  as follows:

    𝐱^0_k = μ_k + σ_k ⊙ϵ,

where ϵ∼ N(0,I) and ⊙ is the element product. Latent embedding 𝐱^0_k will be used as the input for the environment inference module to generate environment pseudo-labels. Meanwhile, 𝐱^0_k will do the forward and reverse processes in the latent space to learn the user/item embeddings in DDPM.  The forward process can be calculated as:

    q(𝐱^1:T_k |𝐱^0_k) = ∏_t=1^T q(𝐱^t_k |𝐱^t-1_k).

After obtaining the environment approximation variable  z_causal = Q(E|G_k, I) according to Eq. (<ref>), the pair of latent variables (z_causal, 𝐱^T_k) to learn the invariant graph representation. We approximate the inference distribution by parameterizing the probabilistic decoder through a conditional DDPM p_θ(𝐱^t-1_k | 𝐱^T_k, z_causal). Using DDPM, the forward process is entirely deterministic except for t = 1.  We define the joint distribution of the reverse generative process as follows:

    p_θ(𝐱^0:T_k | z_causal) = p(𝐱^T_k) ∏_t=1^T p_θ(𝐱^t-1_k |𝐱^t_k, z_causal).



    p_θ(𝐱^t-1_k|𝐱) = 𝒩( 𝐱^t-1_k; μ_θ(𝐱^t_k,z_causal, t), Σ_θ(𝐱^t_k, z_causal, t) ).

The loss function in Eq. (<ref>) can be rewritten as:

    ℒ_Invsimple = 𝔼_t, 𝐱_0, ϵ_t[ ϵ_t - ϵ_θ( 𝐱^t_k, z_causal, t ) ^2 ],

where 𝐱^t_k = √(α)_t 𝐱^0_k + √(1 -α_t)ϵ. 

After obtaining the reconstructed output vector ℛ = 𝐱^0'_k from DDPM, it will be used as the input for the decoder of the variational graph decoder, which then reconstructs the input graph. The entire process is illustrated as follows:

    Â_k = ϕ(ℛℛ^⊤),



where ϕ(·) is the activation function (the sigmoid function is used in this paper). The VGAE is optimized by the variational lower bound:

    ℒ_VGAE = 𝔼_q_ψ(𝐱^0_k|A_k,I)[ log p_θ(Â_k|𝐱^0_k) ]    
    - D_KL( q_ψ(𝐱^0_k|A_k,I) ∥ p(𝐱^0_k) ).


Prediction and Joint Optimization. Using the well-trained diffusion model to sample the final embeddings for user preference modeling:

    r̂_u,i =  e^⊤_ue_i,

where e_u and e_i denote the final user embedding w.r.t u-th user and item embedding w.r.t i-th item, respectively. Without loss of generality, LightGCN is used as the recommendation backbone, and Bayesian Personalized Ranking (BRP) loss is employed to optimize the model parameters:

    ℒ_rec = ∑_u, v^+, v^- -logσ(r̂_u, v^+ - r̂_u, v^-),

where (u, v^+, v^-) is a triplet sample for pairwise recommendation training. v^+ represents positive samples from which the user has interacted, and v^- are the negative samples that are randomly drawn from the set of items with which the user has not interacted, respectively. 

We use a joint learning strategy to optimize CausalDiffRec:

    ℒ =    ℒ_rec + λ_1 ·ℒ_generator
        +λ_2 · (ℒ_VGAE + ℒ_Invsample) +  λ_3 ·ℒ_envInf,

where λ_1, λ_2, and λ_3 are hyper-parameters. 

In section <ref>,
we provide the theoretical proof that our proposed CausalDiffRec can get rid of the unstable correlation and learn OOD generalized representations when satisfying the Assumption. 

The pseudocode for the model training process is shown in Algorithm <ref>.



§ THEORETICAL ANALYSIS
 
CausalDiffRec aims to learn the optimal generator F^*(·) as stated in the assumption in section <ref>, thereby obtaining invariant graph representations to achieve OOD generalization in recommendation performance under data distribution shifts. 
Before starting the theoretical derivation, let's do some preliminary work. For the convenience of theoretical proof, we rewrite Eq. (<ref>) as:

    argmin_θ  (ℒ_task +  ℒ_infer),

where ℒ_task = ℒ_rec + ℒ_generator + ℒ_VGAE+ ℒ_Invsample and  ℒ_infer = ℒ_envInf, and for the derivation convenience, we temporarily ignore the penalty coefficient. The ℒ_task and ℒ_infer can be further abstracted as:

    ℒ_task  = arg min_θ𝔼_e ∼ D_tr(E), (𝒢_u,y_u) ∼ P(Y,G|E=e)
          [  l (f_θ(𝒢_u;θ), y_u)] 
       ℒ_infer 
            = min_q(Y|z_causal)Var{𝔼_e ∼ D_tr(E), (𝒢_u,y_u) ∼ P(Y,G|E=e)
          [l (f_θ(𝒢_u;θ), y_u)|do(𝒢_u)]}.

We follow the proof technique from <cit.> and show the optimality of the Eq. (<ref>) with the following two propositions that can achieve OOD recommendation.

 
Minimizing Eq. (<ref>) promotes the model's adherence to the Invariance Property and the Sufficient Condition outlined in Assumption (in Sec. <ref>).



 
Optimizing Eq. (<ref>) corresponds to minimizing the upper bound of the OOD generalization error described in Eq. (<ref>).



Proposition <ref> and Proposition <ref>, respectively, avoid strong hypotheses and ensure that the OOD generalization error bound of the learned model is within the expected range. In fact, this can also be explained from the perspective of the SCM model in Figure <ref>. Optimizing Eq. (<ref>) eliminates the negative impact of unstable correlations learned by the model, which are caused by latent environments, on modeling user preferences. At the same time, it enhances the model's ability to learn invariant causal features across different latent environments. Proofs for Proposition <ref> and Proposition <ref> are shown as follows:

Before starting the proof, we directly follow previous work <cit.> to propose the following lemma, using information theory to interpret the invariance property and sufficient condition in Assumption and to assist in the proof of Proposition <ref>.


Using the Mutual Information 𝕀(;), the invariance property and sufficient condition in Assumption can be equivalently expressed as follow lemma:

(1) Invariance: ∀ e ∈ D(E), P_θ(Y|𝒫^*_Inv, E=e, I) = P(Y|𝒫^*_Inv, I) ⇔𝕀(Y; E|𝒫^*_Inv, I) = 0 where 𝒫^*_Inv =  F^*(G). 
(2) Sufficiency: 𝕀(Y; 𝒫^*_Inv, I) is maxmized.





    For the invariance property, it is easy to get the following equation:
    
    𝕀(Y;E|𝒫^*_Inv, I) 
       = 𝔼_𝒫^*_Inv, I[ 𝔻_KL( P(Y, E | 𝒫^*_Inv, I)  P(Y | 𝒫^*_Inv, I) P(E | 𝒫^*_Inv, I) ) ]

    For the sufficient condition, we use the method of contradiction and prove it through the following two steps:
    
First, we prove that for Y, 𝒫^*_Inv, and I satisfying 𝒫^*_Inv = max_𝒫_Inv𝕀(Y; 𝒫_Inv, I), they also satisfy that 𝕀(Y; 𝒫^*_Inv, I) is maximized. We use the method of contradiction to prove this.
Assume 𝒫^*_Inv≠max_𝒫_Inv𝕀(Y; 𝒫_Inv, I), and there exists 𝒫^'_Inv = max_𝒫_Inv𝕀(Y; 𝒫_Inv, I), where 𝒫^'_Inv≠𝒫^*_Inv. We can always find a mapping function M such that 𝒫^'_Inv = M(𝒫^*_Inv, R), where R is a random variable. Then we have:

    𝕀(Y; 𝒫^'_Inv, I)     = 𝕀(Y; 𝒫^*_Inv, R, I) 
        = 𝕀(Y; 𝒫^*_Inv, I) + 𝕀(Y; R | 𝒫^*_Inv, I).

Since R is a random variable and does not contain any information about Y, we have 𝕀(Y; R | 𝒫^*_Inv, I) = 0. Therefore:

    𝕀(Y; 𝒫^'_Inv, I) = 𝕀(Y; 𝒫^*_Inv, I).

This leads to a contradiction. 

Next, we prove that for Y, 𝒫^*_Inv, and I satisfying 𝒫^*_Inv = max_𝒫_Inv𝕀(Y; 𝒫_Inv, I), they also satisfy that 𝕀(Y; 𝒫^*_Inv, I) is maximized. Assume 𝒫^*_Inv≠max_𝒫_Inv𝕀(Y; 𝒫_Inv, I), and there exists 𝒫^'_Inv = max_𝒫_Inv𝕀(Y; 𝒫_Inv, I), where 𝒫^'_Inv≠𝒫^*_Inv. We have the following inequality:

    𝕀(Y; 𝒫^*_Inv, I) ≤𝕀(Y; 𝒫^'_Inv, I).

From this, we can deduce that:

    𝒫^'_Inv = max_𝒫_Inv𝕀(Y; 𝒫_Inv, I),
 where contradicts 𝒫^*_Inv = max_𝒫_Inv𝕀(Y; 𝒫_Inv, I).
Since the assumption leads to a contradiction, and the assumption does not hold. Therefore, 𝒫^*_Inv = max_𝒫_Inv𝕀(Y; 𝒫_Inv, I) holds. This proves that 𝕀(Y; 𝒫^*_Inv, I) is maximized. The lemma <ref> is complicated proven.


Proof of Proposition 
<ref>. 

    First, optimizing the first term ℒ_task in Eq. (<ref>) enables the model to satisfy the sufficient condition. Analyzing the SCM in Figure <ref>(c), we have the fact that max_q(z|G, I)𝕀(Y,z_causal) is equivalent to min_q(z|G, I)𝕀(Y, G|Z_causal), as we use do(G) to eliminate the unstable correlations between Y and G caused by the latent environment. We have:
    
    𝕀(Y,G|z_causal)     = D_KL(p(Y|G,E)p(Y|z_causal,E)) 
        = D_KL(p(Y|G,E)p(Y|z_causal)) 
        - D_KL(p(Y|z_causal, E)  q(Y|z_causal)) 
       ≤ D_KL(p(Y|G,E)p(Y|z_causal)).

    Based on the above derivation, we have:
    
    𝕀(Y,G|z_causal) ≤
       min_q(Y|z_causal) D_KL(p(Y|G,E)p(Y|z_causal)).

    Besides, we have:
    
    D_KL(p(Y|G,E)p(Y|z_causal))
               
       = 𝔼_e∈ D_tr(E)𝔼_(G,Y) ∼ p(G,Y|e)
         𝔼_z_causal∼ q(z_causal|G,I)[logq(Y|G,e)/p(Y|z_causal)] 
       ≤𝔼_e∈ D_tr(E)𝔼_(G,Y) ∼ p(G,Y|e)
       [logp(Y|G,e)/𝔼_z_causal∼ q(z_causal|G,I)q(Y|z_causal)] (Jensen Inequality).

    Finally, we reach:
    
    min_q(Y|z_causal) D_KL(p(Y|G,E)p(Y|z_causal)) 
             
       ⇔arg min_θ𝔼_e ∼ D_tr(E), (𝒢_u,y_u) ∼ P(Y,G|E=e) [  l (f_θ(𝒢_u;θ), y_u)].

    Thus, we have demonstrated that minimizing the expectation term (ℒ_task) in Eq. (<ref>) is equivalent to minimizing the upper bound of 𝕀(Y; G | z_causal). This results in maximizing 𝕀(Y; 𝒫^*_Inv, I), thereby helping to ensure that the model satisfies the Sufficient Condition.
    
    Next, we prove that optimizing the first term ℒ_task in Eq. (<ref>) enables the model to satisfy the Invariance Property.
    Similar to Eq. (<ref>), we have:
    
    𝕀(Y; E =e | z_causal) 
       = D_KL(p(Y | z_causal, e) ∥ p(Y | z_causal)) 
       = D_KL(p(Y | z_causal, E) ∥𝔼_e ∈ D(E)[p(Y | z_causal, e)]) 
       = D_KL(q(Y | z_causal) ∥𝔼_e ∈ D(E)[q(Y | z_causal)]) 
        - D_KL(q(Y | z_causal) ∥ p(Y | z_causal, e)) 
        - D_KL(𝔼_e ∈ D(E)[p(Y | z_causal, e)] ∥𝔼_e ∈ D(E)[q(Y | z_causal)]) 
       ≤ D_KL(q(Y | z_causal) ∥𝔼_e ∈ D(E)[q(Y | z_causal)]).

    Besides, the last term in Eq. (<ref>) can be further expressed as:
    
    D_KL(q(Y | z_causal) ∥𝔼_e ∈ D(E)[q(Y | z_causal)]) 
        = 𝔼_e∈ D_tr(E)𝔼_(G,Y) ∼ p(G,Y|e)𝔼_z_causal∼ q(z_causal|G,I)
         [logp(Y|G,e)/𝔼_e ∈ D(E)q(Y|z_causal)] (Jensen Inequality)
       ≤𝔼_e∈ D(E)[| l (f_θ(𝒢_u;θ), y_u)-  𝔼_e ∈ D(E)[l (f_θ(𝒢_u;θ), y_u)]|]

    where the last term in Eq. (<ref>) is the upper bound for the D_KL(q(Y | z_causal) ∥𝔼_e ∈ D(E)[q(Y | z_causal)]). 
    Finally, we have:
    
    min_q(Y|z_causal) D_KL(q(Y | z_causal) ∥𝔼_e ∈ D(E)[q(Y | z_causal)])  
       ⇔min_q(Y|z_causal)Var{𝔼_e ∼ D_tr(E), (𝒢_u,y_u) ∼ P(Y,G|E=e)
        [l (f_θ(𝒢_u;θ), y_u)|do(𝒢_u)]}.

    Hence, minimizing the variance term (ℒ_risk) in Eq. (<ref>) effectively reduces the upper bound of 𝕀(Y; E =e | z_causal). Thereby ensuring the model adheres to the Invariance Property.


Proof of Proposition <ref>. Optimizing Eq. (<ref>) is tantamount to reducing the upper bound of the OOD generalization error in Eq. (<ref>).
   Let q(Y | G) represent the inferred variational distribution of the true distribution p(Y | G, E). The OOD generalization error can be quantified by the KL divergence between these two distributions:
   
    D_KL(p(Y|G,E)q(Y|G)) 
        =𝔼_e ∈ D_tr(E)𝔼_(Y,G) ∼ p(Y,G|E)logp(Y|G,E=e)/q(Y|G).

   Following previous work, we use information theory to assist in the proof of Proposition <ref>.
   We propose the lemma <ref> to rewrite the OOD generalization, which is shown as follows:
    
       The out-of-distribution generalization error is limited by:
       
    D_KL(p(Y|G,E)q(Y|G)) ≤ D_KL[p(Y|G,E)q(Y|z_causal)],

   
   where q(Y|z_causal) is the inferred variational environment distribution. The proof of Lemma <ref> is shown as:
   
    D_KL(p(Y|G,E=e)q(Y|z_causal)) 
        = 𝔼_e ∈ D(E)𝔼_(Y,G) ∼ p(G, Y|E=e)[logp(Y|G,E=e)/q(Y|G)] 
        = 𝔼_e ∈ D(E)𝔼_(Y,G) ∼ p(G, Y|E=e)
         [logp(Y|G,E=e)/𝔼_z_causal∼ q(z_causal|G,I)q(Y|z_causal)] 
       ≤𝔼_e ∈ D(E)𝔼_(Y,G) ∼ p(G, Y|E=e)
         𝔼_z_causal∼ q(z_causal|G,I)logp(Y|G,e)/q(Y|z_causal)
        = D_KL[p(Y|G,E)q(Y|z_causal)],

   The Lemma <ref> has been fully proven. Based on Lemma <ref> and Proposition <ref>, the Eq. (<ref>) can be adapted as:
   
    min_q(z_causal|G,I),q(Y,z_causal)
        D_KL(p(Y|G,E=e)q(Y|z_causal)) + 𝕀(Y,E=e|z_causal)

   Hence, according to Lemma <ref>, we confirm that minimizing Eq. (<ref>) is tantamount to minimizing the upper bound of the OOD generalization error in Eq. (<ref>), meaning that:
   
    argmin_θ(ℒ_task + ℒ_infer)  ⇔min_q(z_causal|G,I),q(Y,z_causal)
          D_KL(p(Y|G,E=e)q(Y|z_causal)) 
        + 𝕀(Y,E=e|z_causal) (𝕀(Y,E=e|z_causal) is  non-negative)  
       ≥min_q(z_causal|G,I),q(Y,z_causal)
          D_KL(p(Y|G,E=e)q(Y|z_causal)) 
       ≥  D_KL(p(Y|G,E)q(Y|G)).

   The Proposition <ref> is completely proven.








§ EXPERIMENTS

In this section, we conducted extensive experiments to validate the performance of CausalDiffRec and address the following key research questions:


  * RQ1: How does CausalDiffRec compare to the state-of-the-art strategies in both OOD and IID test evaluations?

  * RQ2: Are the proposed components of CausalDiffRec effective for OOD generalization?

  * RQ3: How do hyperparameter settings affect the performance of CausalDiffRec?








 §.§ Experimental Settings


Datasets. We evaluate our proposed CausalDiffRec on four real-world datasets: Food[https://www.aclweb.org/anthology/D19-1613/], KuaiRec[https://kuairec.com] Yelp2018[https://www.yelp.com/dataset], and Douban[https://www.kaggle.com/datasets/] comprise raw data from the Douban system. Detailed statistics of the datasets are presented in Table 1. More data processing details can be seen in the code implementation link.
Following exiting work <cit.>,
we constructed OOD test sets for three common scenarios of distribution shift:


  * Temporal shift. We sort the dataset in descending order by timestamp and used each user's most recent 20% of interactions as the OOD test set. 
The first 60% of the data based on interaction time is used for the training set. Food is used for this type of shift.

  * Exposure shift. In KuaiRec, the small matrix is fully exposed and used for the OOD test set, and the big matrix is collected from the online platform and used for the training dataset, resulting in a distributional shift. 

  * Popularity shift. We randomly select 20% of the interactions to create the OOD test set, ensuring that item popularity follows a uniform distribution. The training set maintains a long-tail distribution. This type of shift was applied to the Yelp2018 and Douban datasets.




Baselines. We compare the CusalDiffRec with the state-of-the-art models:
LightGCN<cit.>, 
SGL <cit.>,
SimGCL<cit.>, LightGCL<cit.>, InvPref <cit.>, 
InvCF <cit.>, 
AdvDrop <cit.>, AdvInfo <cit.>, and DRO <cit.>.



Hyperparameter Settings. We implement our CausalDiffRec in Pytorch. All experiments are conducted on a single RTX-4090 with 24G memory. Following the default hyperparameter search settings of the baselines, we expand their hyperparameter search space and tune the hyperparameters.
For our CausalDiffRec, we tune
the learning rates in {1e-3, 1e-4, 1e-5}. The number of diffusion steps varies between 10 and 1000, and the diffusion embedding size is chosen in {8,16,32,64}. Additional hyperparameter details are available in our released code.





























































 §.§ Overall Performance (RQ1)

This section compares CausalDiffRec's performance and baselines under various data shifts and conducts a performance analysis.




Evaluation on temporal shift: Table <ref> shows that CausalDiffRec significantly outperforms SOTA models on the Food dataset, with improvements of 1.99%, 24.89%, 6.03%, and 9.68% in Recall and NDCG. This indicates CausalDiffRec's effectiveness in handling temporal shift. DRO also excels in this area, with a 15% improvement over LightGCN in NDCG@20, due to its robust optimization across various data distributions. CDR surpasses GNN-based models thanks to its temporal VAE-based architecture, capturing preference shifts from temporal changes.

Evaluation on exposure shift: In real-world scenarios, only a small subset of items is exposed to users, leading to non-random missing interaction records. Using the fully exposed KuaiRec dataset, CausalDiffRec consistently outperforms baselines, with improvements ranging from 6.90% to 28.83%, indicating its capability to handle exposure bias. DRO and AdvInfoNce also show superior performance in NDCG and Recall metrics, enhancing the generalization of GNN-based models and demonstrating robustness compared to LightGCN.


Evaluation on popularity shift. We compare model performance on the Yelp2018 and Douban datasets, showing that our model significantly outperforms the baselines. On Douban, CausalDiffRec achieves 8.22% to 17.96% improvement, and on Yelp2018, the improvements range from 11.24% to 36.73%.
Methods using contrastive learning (e.g., SimGCL, LightGCL, AdvInfoNce) outperform other baselines in handling popularity shifts. This is because the InfoNCE loss helps the model learn a more uniform representation distribution, reducing bias towards popular items. InvPref performs best among the baselines on Yelp2018, using clustering for contextual labels, unlike our variational inference approach. Our method, tailored for graph data, aggregates neighbor information for better recommendation performance than matrix factorization-based methods.

Additionally, in Table <ref>, we report the performance of CausalDiffRec compared to several baseline models that use LightGCN as the backbone. From the table, we observe the following: 1) These baseline models outperform LightGCN on IID datasets; 2) CausalDiffRec outperforms all baseline models across all metrics. This indicates that CausalDiffRec also performs well on IID datasets. We attribute the performance improvement to our use of data augmentation and the incorporation of auxiliary information in modeling user preferences.

In summary, the analysis of experimental results demonstrates that our proposed CausalDiffRec can handle different types of distribution shifts and achieve good generalization.










 §.§ In-depth Analysis (RQ2)

In this section, we conduct ablation experiments to study the impact of each component of CausalDiffRec on recommendation performance. The main components include the environment generator module and the environment inference module. Additionally, we use t-SNE to visualize the item representations captured by the baseline model and CausalDiffRec, to compare the models' generalization capabilities on OOD data.



Ablation studies. Table <ref> presents the results of the ablation study that compares LightGCN, CausalDiffRec, and its two variants: 'w/o Gen.' (without the environment generator) and 'w/o Env.' (without the environment inference). The results show that removing these modules causes a significant drop in all metrics across four datasets. For example, on Yelp2018, Recall, and NDCG decreased by 148.15% and 126.83%, respectively, demonstrating the effectiveness of CausalDiffRec based on invariant learning theory for enhancing recommendation performance on OOD datasets.
Additionally, even without the modules, CausalDiffRec still outperforms LightGCN on popularity shift datasets (Yelp2018 and Douban) due to the effectiveness of data augmentation and environment inference. However, on the Food and KuaiRec datasets, removing either module results in worse performance than LightGCN, likely due to multiple biases in these datasets. Without one module, the model struggles to handle multiple data distributions, leading to a performance drop.
Overall, the ablation experiments highlight the importance of all modules in CausalDiffRec for improving recommendation performance and generalizing on OOD data.

Visualization analysis. In Figure <ref> and Figure <ref>, we used t-SNE to visualize the item representations learned by LightGCN, SimGCL, and CausalDiffRec on Douban and Yelp2018 datasets to better observe our model's ability to handle distribution shifts. Following previous work <cit.>, we recorded the popularity of each item in the training set and designated the top 10% most popular items as 'popular items' and the bottom 10% as 'unpopular items'. It is obvious that the embeddings of popular and unpopular items learned by LightGCN still exhibit a gap in the representation space. In contrast, the embeddings learned by CausalDiffRec are more evenly distributed within the same space. This indicates that CausalDiffRec can mitigate the popularity shift caused by popular items.  Additionally, we found that the embeddings of popular and unpopular items learned by SimGCL are more evenly distributed compared to LightGCN. This is because contrastive learning can learn a uniform representation distribution.

















 §.§ Hyperparameter Investigation (RQ3)

Effect of Diffusion Step T.
We conducted experiments to investigate the impact of the number of diffusion steps on performance; in CausalDiffRec, we used the same number of steps in the forward and reverse phases.
we compare the performance with T changing from 10 to 500. We present the results in Figure <ref>, and we have the following findings:

    
  * When the number of steps is chosen within the range {10, 50, 100}, CausalDiffRec achieves the best performance across all datasets. We find that appropriately increasing the number of steps significantly improves Recall@20 and NDCG@20. These performance enhancements are mainly attributed to the diffusion enriching the representation capabilities of users and items.
    
  * Nevertheless, as we continue to increase the number of steps, the model will face overfitting issues. For example, on the food and yelp2018 datasets, Recall@20 and NDCG@20 consistently decrease. Although there is an upward trend on KuaiRec, the optimal solution is not achieved. Additionally, it is evident that more steps also lead to longer training times. We should carefully adjust the number of steps to find the optimal balance between enhancing representation ability and avoiding overfitting.


Effect of Penalty Coefficient.
Figure <ref> reports the experimental results of varying the coefficients λ_1, λ_2, and λ_3 on food and KuaiRec. We compare the performances with all coefficients changing from 1e-1 to 1e-4. Further, we have the following conclusions: 

    
  * λ_1 ∈ [1e-1, 1e-2] typically strikes a good balance and delivers excellent performance. When further reducing the value of λ_1, the model's performance declined. We speculate that this led to overfitting. Similarly, we found that larger values of λ_3 can also achieve better recommendation performance. Therefore, the value range of λ_3 is the same as that of λ_1.
    
  * λ_3 directly regularizes the capability of the graph representation learning module. We found that for different datasets, values that are too large or too small do not allow the model to learn optimal representations. Taking everything into consideration, we set λ_3 to 1e-3 to achieve good generalization.



Effect of the number of Environments. Figure <ref> shows the impact of the number of environments on the model's performance. We can see that as the number of environments increases, the performance of CausalDiffRec improves across the three datasets. This indicates that more environments help enhance the model's generalization on OOD (Out-of-Distribution) data. However, as the number of environments further increases, performance declines, which we believe is due to the model overfitting to too many environments.




§ RELATED WORK



 §.§ GNN-based Recommendation

Recent developments in graph-based recommender systems have leveraged graph neural networks to model user-item interactions as a bipartite graph, enhancing recommendation accuracy through complex interaction capture <cit.>. Notably, LightGCN focuses on neighborhood aggregation without additional transformations, while other approaches employ attention mechanisms to prioritize influential interactions <cit.>. Further research explores non-Euclidean spaces like hyperbolic space to better represent user-item relationships <cit.>. Knowledge graphs also enhance these systems by integrating rich semantic and relational data directly into the recommendation process <cit.>. Despite these advancements, graph-based systems often struggle with out-of-distribution data due to the IID assumption and are challenged by multiple distribution shifts <cit.>. Additionally, contrastive learning methods in these systems rely on a fixed paradigm that lacks robust theoretical support, limiting adaptability to varied data shifts. However, the aforementioned models are trained on datasets where the training and test data distributions are drawn from the same distribution, leading to generalization failure when facing OOD data.



 §.§ Diffusion based Recoomendation

The integration of diffusion processes into recommender systems leverages diffusion mechanisms to model dynamic propagation of user preferences and item information through interaction networks, enhancing recommendation accuracy and timeliness <cit.>. These models capture evolving user behaviors and have shown potential in various recommendation contexts, from sequential recommendations to location-based services. For instance, DiffRec <cit.> applies diffusion directly for recommendations, while Diff-POI <cit.> models location preferences. Furthermore, approaches like DiffKG <cit.> and RecDiff <cit.> utilize diffusion for denoising entity representations in knowledge graphs and user data in social recommendations, respectively, enhancing the robustness and reliability of the systems. These studies underscore diffusion's suitability for advanced representation learning in recommender systems. However, these methods cannot solve the OOD problem.




 §.§ Out-of-Distribution Recommendation


Researchers have focused on recommendation algorithms for out-of-distribution (OOD) data. COR <cit.> infers latent environmental factors in OOD data. CausPref <cit.> learns invariant user preferences and causal structures using anti-preference negative sampling. CaseQ <cit.> employs backdoor adjustment and variational inference for sequential recommendations. InvPref <cit.> separates invariant and variant preferences by identifying heterogeneous environments. However, these methods don't directly apply to graph-based recommendation models and fail to address OOD in graph structures. AdaDrop <cit.> uses adversarial learning and graph neural networks to enhance performance by decoupling user preferences. DRO <cit.> integrates Distributionally Robust Optimization into Graph Neural Networks to handle distribution shifts in graph-based recommender systems. Distinct from these GNN-based methods, this paper explores how to use the theory of invariant learning to design GNN-based methods with good generalization capabilities.





§ CONCLUSION

In this paper, we propose a novel GNN-based model for OOD recommendation, called CausalDiffRec, designed to learn environment-invariant graph representations to enhance the generalization of recommendation models on OOD data. CausalDiffRec employs the backdoor criterion from causal inference and variational inference methods to eliminate the influence of environmental confounders. It uses a diffusion-based sampling strategy to learn graph representations. Our approach is grounded in invariant learning theory, and we provide theoretical proof that by optimizing the objective function in CausalDiffRec, the model is encouraged to identify invariant graph representations across environments, thereby improving its generalization performance on OOD data. Experiments conducted on four real-world datasets demonstrate that the proposed CausalDiffRec framework outperforms baseline models. In addition, ablation studies further validate the effectiveness of the model.



00
wu2022graph
S. Wu, F. Sun, W. Zhang, X. Xie, and B. Cui,
"Graph neural networks in recommender systems: a survey,"
ACM Computing Surveys, vol. 55, no. 5, pp. 1–37, 2022.

wang2019neural
X. Wang, X. He, M. Wang, F. Feng, and T.-S. Chua,
"Neural graph collaborative filtering,"
in SIGIR, 2019, pp. 165–174.

he2020lightgcn
X. He, K. Deng, X. Wang, Y. Li, Y. Zhang, and M. Wang,
"Lightgcn: Simplifying and powering graph convolution network for recommendation,"
in SIGIR, 2020, pp. 639–648.

wang2020disenhan
Y. Wang, S. Tang, Y. Lei, W. Song, S. Wang, and M. Zhang,
"Disenhan: Disentangled heterogeneous graph attention network for recommendation,"
in CIKM, 2020, pp. 1605–1614.

chang2021sequential
J. Chang, C. Gao, Y. Zheng, Y. Hui, Y. Niu, Y. Song, D. Jin, and Y. Li,
"Sequential recommendation with graph neural networks,"
in SIGIR, 2021, pp. 378–387.

wang2021session
J. Wang, K. Ding, Z. Zhu, and J. Caverlee,
"Session-based recommendation with hypergraph attention networks,"
in SDM, 2021, pp. 82–90, SIAM.

wang2020disentangled
X. Wang, H. Jin, A. Zhang, X. He, T. Xu, and T.-S. Chua,
"Disentangled graph collaborative filtering,"
in SIGIR, 2020, pp. 1001–1010.

niu2021review
Z. Niu, G. Zhong, and H. Yu,
"A review on the attention mechanism of deep learning,"
Neurocomputing, vol. 452, pp. 48–62, 2021.

fukui2019attention
H. Fukui, T. Hirakawa, T. Yamashita, and H. Fujiyoshi,
"Attention branch network: Learning of attention mechanism for visual explanation,"
in CVPR, 2019, pp. 10705–10714.

zhang2022geometric
Y. Zhang, C. Li, X. Xie, X. Wang, C. Shi, Y. Liu, H. Sun, L. Zhang, W. Deng, and Q. Zhang,
"Geometric disentangled collaborative filtering,"
in SIGIR, 2022, pp. 80–90.

sun2021hgcf
J. Sun, Z. Cheng, S. Zuberi, F. Pérez, and M. Volkovs,
"Hgcf: Hyperbolic graph convolution networks for collaborative filtering,"
in WWW, 2021, pp. 593–601.

wang2019kgat
X. Wang, X. He, Y. Cao, M. Liu, and T.-S. Chua,
"Kgat: Knowledge graph attention network for recommendation,"
in KDD, 2019, pp. 950–958.

cao2019unifying
Y. Cao, X. Wang, X. He, Z. Hu, and T.-S. Chua,
"Unifying knowledge graph learning and recommendation: Towards a better understanding of user preferences,"
in WWW, 2019, pp. 151–161.

wu2021self
J. Wu, X. Wang, F. Feng, X. He, L. Chen, J. Lian, and X. Xie,
"Self-supervised graph learning for recommendation,"
in SIGIR, 2021, pp. 726–735.

yu2022graph
J. Yu, H. Yin, X. Xia, T. Chen, L. Cui, and Q. V. Hung,
"Are graph augmentations necessary? simple graph contrastive learning for recommendation,"
in SIGIR, 2022, pp. 1294–1303.

jiang2023adaptive
Y. Jiang, C. Huang, and L. Huang,
"Adaptive graph contrastive learning for recommendation,"
in KDD, 2023, pp. 4252–4261.

xia2023automated
L. Xia, C. Huang, C. Huang, K. Lin, T. Yu, and B. Kao,
"Automated self-supervised learning for recommendation,"
in WWW, 2023, pp. 992–1002.

li2023graph
C. Li, L. Xia, X. Ren, Y. Ye, Y. Xu, and C. Huang,
"Graph transformer for recommendation,"
in SIGIR, 2023, pp. 1680–1689.

wang2022causal
W. Wang, X. Lin, F. Feng, X. He, M. Lin, and T.-S. Chua,
"Causal representation learning for out-of-distribution recommendation,"
in WWW, 2022, pp. 3562–3571.

he2022causpref
Y. He, Z. Wang, P. Cui, H. Zou, Y. Zhang, Q. Cui, and Y. Jiang,
"CausPref: Causal preference learning for out-of-distribution recommendation,"
in WWW, 2022, pp. 410–421.

yang2022towards
C. Yang, Q. Wu, Q. Wen, Z. Zhou, L. Sun, and J. Yan,
"Towards out-of-distribution sequential event prediction: A causal treatment,"
NIPS, vol. 35, pp. 22656–22670, 2022.

wang2022invariant
Z. Wang, Y. He, J. Liu, W. Zou, P. S. Yu, and P. Cui,
"Invariant preference learning for general debiasing in recommendation,"
in KDD, 2022, pp. 1969–1978.

zhang2024general
A. Zhang, W. Ma, P. Wei, L. Sheng, and X. Wang,
"General Debiasing for Graph-based Collaborative Filtering via Adversarial Graph Dropout,"
in WWW, 2024, pp. 3864–3875.

wang2024distributionally
B. Wang, J. Chen, C. Li, S. Zhou, Q. Shi, Y. Gao, Y. Feng, C. Chen, and C. Wang,
"Distributionally Robust Graph-based Recommendation System,"
in WWW, 2024, pp. 3777–3788.

ho2020denoising
J. Ho, A. Jain, and P. Abbeel,
"Denoising diffusion probabilistic models,"
NIPS, vol. 33, pp. 6840–6851, 2020.

wang2023learning
Y. Wang, X. Wang, A.-D. Dinh, B. Du, and C. Xu,
"Learning to schedule in diffusion probabilistic models,"
in KDD, 2023, pp. 2478–2488.

wu2022handling
Q. Wu, H. Zhang, J. Yan, and D. Wipf,
"Handling distribution shifts on graphs: An invariance perspective,"
ICLR, 2022.

yuan2024environment
H. Yuan, Q. Sun, X. Fu, Z. Zhang, C. Ji, H. Peng, and J. Li,
"Environment-Aware Dynamic Graph Learning for Out-of-Distribution Generalization,"
NIPS, vol. 36, 2024.

wangDro
B. Wang, J. Chen, C. Li, S. Zhou, Q. Shi, Y. Gao, Y. Feng, C. Chen, and C. Wang,
"Distributionally Robust Graph-based Recommendation System,"
in WWW, 2024.

zhang2024robust
A. Zhang, W. Ma, J. Zheng, X. Wang, and T.-S. Chua,
"Robust collaborative filtering to popularity distribution shift,"
TOIS, vol. 42, no. 1, pp. 1–23, 2024.
cai2023lightgcl
X. Cai, C. Huang, L. Xia, and X. Ren,
"LightGCL: Simple yet effective graph contrastive learning for recommendation,"
arXiv preprint arXiv:2302.08191, 2023.

wang2023causal
W. Wang, X. Lin, L. Wang, F. Feng, Y. Ma, and T.-S. Chua,
"Causal disentangled recommendation against user preference shifts,"
TOIS, vol. 42, no. 1, pp. 1–27, 2023.

zhang2023invariant
A. Zhang, J. Zheng, X. Wang, Y. Yuan, and T.-S. Chua,
"Invariant collaborative filtering to popularity distribution shift,"
in WWW, 2023, pp. 1240–1251.

zhang2024empowering
A. Zhang, L. Sheng, Z. Cai, X. Wang, and T.-S. Chua,
"Empowering Collaborative Filtering with Principled Adversarial Contrastive Loss,"
NIPS, 2024.

zhou2020graph
J. Zhou, G. Cui, S. Hu, Z. Zhang, C. Yang, Z. Liu, L. Wang, C. Li, and M. Sun,
"Graph neural networks: A review of methods and applications,"
AI open, vol. 1, pp. 57–81, 2020, Elsevier.

wu2020comprehensive
Z. Wu, S. Pan, F. Chen, G. Long, C. Zhang, and P. S. Yu,
"A comprehensive survey on graph neural networks,"
TNNLS, vol. 32, no. 1, pp. 4–24, 2020, IEEE.

zheng2022graph
X. Zheng, Y. Wang, Y. Liu, M. Li, M. Zhang, D. Jin, P. S. Yu, and S. Pan,
"Graph neural networks for graphs with heterophily: A survey,"
arXiv preprint arXiv:2202.07082, 2022.

guo2020survey
Q. Guo, F. Zhuang, C. Qin, H. Zhu, X. Xie, H. Xiong, and Q. He,
"A survey on knowledge graph-based recommender systems,"
TKDE, vol. 34, no. 8, pp. 3549–3568, 2020, IEEE.

wang2021graph
S. Wang, L. Hu, Y. Wang, X. He, Q. Z. Sheng, M. A. Orgun, L. Cao, F. Ricci, and P. S. Yu,
"Graph learning based recommender systems: A review,"
arXiv preprint arXiv:2105.06339, 2021.

guo2022attention
M.-H. Guo, T.-X. Xu, J.-J. Liu, Z.-N. Liu, P.-T. Jiang, T.-J. Mu, S.-H. Zhang, R. R. Martin, M.-M. Cheng, and S.-M. Hu,
"Attention mechanisms in computer vision: A survey,"
Computational visual media, vol. 8, no. 3, pp. 331–368, 2022, Springer.

ji2021survey
S. Ji, S. Pan, E. Cambria, P. Marttinen, and P. S. Yu,
"A survey on knowledge graphs: Representation, acquisition, and applications,"
TNNLS, vol. 33, no. 2, pp. 494–514, 2021, IEEE.

jing2023contrastive
M. Jing, Y. Zhu, T. Zang, and K. Wang,
"Contrastive self-supervised learning in recommender systems: A survey,"
TOIS, vol. 42, no. 2, pp. 1–39, 2023.

yu2023self
J. Yu, H. Yin, X. Xia, T. Chen, J. Li, and Z. Huang,
"Self-supervised learning for recommender systems: A survey,"
TKDE, vol. 36, no. 1, pp. 335–355, 2023, IEEE.

creager2021environment
E. Creager, J.-H. Jacobsen, and R. Zemel,
"Environment inference for invariant learning,"
in ML, 2021, pp. 2189–2200, PMLR.

li2022learning
H. Li, Z. Zhang, X. Wang, and W. Zhu,
"Learning invariant graph representations for out-of-distribution generalization,"
NIPS, vol. 35, pp. 11828–11841, 2022.

zhang2024spectral
Z. Zhang, X. Wang, Z. Zhang, Z. Qin, W. Wen, H. Xue, H. Li, and W. Zhu,
"Spectral invariant learning for dynamic graphs under distribution shifts,"
NIPS, vol. 36, 2024.

yang2024generalized
J. Yang, K. Zhou, Y. Li, and Z. Liu,
"Generalized out-of-distribution detection: A survey,"
IJCV, pp. 1–28, 2024, Springer.

liu2021towards
J. Liu, Z. Shen, Y. He, X. Zhang, R. Xu, H. Yu, and P. Cui,
"Towards out-of-distribution generalization: A survey,"
arXiv preprint arXiv:2108.13624, 2021.

hsu2020generalized
Y.-C. Hsu, Y. Shen, H. Jin, and Z. Kira,
"Generalized odin: Detecting out-of-distribution image without learning from out-of-distribution data,"
in CVPR, 2020, pp. 10951–10960.

yang2022learning
N. Yang, K. Zeng, Q. Wu, X. Jia, and J. Yan,
"Learning substructure invariance for out-of-distribution molecular representations,"
NIPS, vol. 35, pp. 12964–12978, 2022.

lee2023exploring
S. Lee, J. Jo, and S. J. Hwang,
"Exploring chemical space with score-based out-of-distribution generation,"
in ICML, 2023, pp. 18872–18892, PMLR.

chen2022learning
Y. Chen, Y. Zhang, Y. Bian, H. Yang, K. Ma, B. Xie, T. Liu, B. Han, and J. Cheng,
"Learning causally invariant representations for out-of-distribution generalization on graphs,"
NIPS, vol. 35, pp. 22131–22148, 2022.

wang2023diffusion
W. Wang, Y. Xu, F. Feng, X. Lin, X. He, and T.-S. Chua,
"Diffusion recommender model,"
in SIGIR, 2023, pp. 832–841.

li2023diffurec
Z. Li, A. Sun, and C. Li,
"Diffurec: A diffusion model for sequential recommendation,"
TOIS, vol. 42, no. 3, pp. 1–28, 2023.

ma2024plug
H. Ma, R. Xie, L. Meng, X. Chen, X. Zhang, L. Lin, and Z. Kang,
"Plug-in diffusion model for sequential recommendation,"
in AAAI, vol. 38, no. 8, pp. 8886–8894, 2024.

wu2023diff4rec
Z. Wu, X. Wang, H. Chen, K. Li, Y. Han, L. Sun, and W. Zhu,
"Diff4rec: Sequential recommendation with curriculum-scheduled diffusion augmentation,"
in CIKM, 2023, pp. 9329–9335.

qin2023diffusion
Y. Qin, H. Wu, W. Ju, X. Luo, and M. Zhang,
"A diffusion model for poi recommendation,"
TOIS, vol. 42, no. 2, pp. 1–27, 2023.

jiang2024diffkg
Y. Jiang, Y. Yang, L. Xia, and C. Huang,
"Diffkg: Knowledge graph diffusion model for recommendation,"
in WSDM, 2024, pp. 313–321.

li2024recdiff
Z. Li, L. Xia, and C. Huang,
"RecDiff: Diffusion Model for Social Recommendation,"
arXiv preprint arXiv:2406.01629, 2024.

X.HE2017
 X. He,  L. Liao,  H. Zhang, L, Nie,  X. Hu, and T.S. Chua. "Neural collaborative filtering," in WWW, 2017, pp. 173–182.

Lin2024
X. Lin, W. Wang, J. Zhao, Y. Li, F. Feng, and T.S. Chua. "Temporally and distributionally robust optimization for cold-start recommendation," in AAAI, 2024, pp. 8750-8758.







 Paper Title: DiM-Gesture: Co-Speech Gesture Generation with Adaptive Layer Normalization Mamba-2 framework.











        












DiM-Gesture: Co-Speech Gesture Generation with Adaptive Layer Normalization Mamba-2 framework
    Fan Zhang^ 0000-0002-9534-1777,
Naye Ji ^() ^ 0000-0002-6986-3766,
Fuxing Gao^ 0009-0008-5586-4734,
Bozuo Zhao ^ 0009-0008-2152-0087,
Jingmei Wu ^ 0009-0000-9221-8784,
Yanbing Jiang ^ 0009-0003-3020-598X,
Hui Du^ 0000-0001-6737-9420, 
Zhenqing Ye^ 0009-0003-4341-2734, 
Leyao Yan ^ 0009-0008-8797-176,
Jiayang Zhu ^ 0009-0005-3072-651X,
WeiFan Zhong^ 0009-0006-3251-1392,
Xiaomeng Ma ^ 0009-0008-0112-5354 


Fan Zhang, Naye Ji, Zhenqing Ye, Jiayang Zhu, WeiFan Zhong, Leyao Yan, Yanbing Jiang, Hui Du, Fuxing Gao, are with the College of Media Engineering, Communication University of Zhejiang, China; (e-mail: fanzhang@cuz.edu.cn;jinaye@cuz.edu.cn; 1767641184@qq.com; zjy190127 26077@foxmail.com; 2633252395@qq.com; yanleyao8@gmail.com; duhui@ cuz.edu.cn; fuxing@cuz.edu.cn;)

Jingmei Wu, Xiaomeng Ma, are with the School of Broadcast Announcing Arts, Communication University of Zhejiang; (e-mail: 20190095@cuz.edu.cn; 1256914539@qq.com)
Zhao Bozuo is with Changjiang Academy of Art and Design, Shantou University, China (e-mail: bzzhao@stu.edu.cn)

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



§ ABSTRACT 
Speech-driven gesture generation is an emerging domain within virtual human creation, where current methods predominantly utilize Transformer-based architectures that necessitate extensive memory and are characterized by slow inference speeds. In response to these limitations, we propose DiM-Gestures, a novel end-to-end generative model crafted to create highly personalized 3D full-body gestures solely from raw speech audio, employing Mamba-based architectures. This model integrates a Mamba-based fuzzy feature extractor with a non-autoregressive Adaptive Layer Normalization (AdaLN) Mamba-2 diffusion architecture. The extractor, leveraging a Mamba framework and a WavLM pre-trained model, autonomously derives implicit, continuous fuzzy features, which are then unified into a singular latent feature. This feature is processed by the AdaLN Mamba-2, which implements a uniform conditional mechanism across all tokens to robustly model the interplay between the fuzzy features and the resultant gesture sequence. This innovative approach guarantees high fidelity in gesture-speech synchronization while maintaining the naturalness of the gestures. Employing a diffusion model for training and inference, our framework has undergone extensive subjective and objective evaluations on the ZEGGS and BEAT datasets. These assessments substantiate our model's enhanced performance relative to contemporary state-of-the-art methods, demonstrating competitive outcomes with the DiTs architecture (Persona-Gestors) while optimizing memory usage and accelerating inference speed. Code can be accessed at https://github.com/zf223669/DiMGestures https://github.com/zf223669/DiMGestures.




Speech-driven, Gesture synthesis, Fuzzy inference, AdaLN,  Diffusion, Mamba, SSMs.




§ INTRODUCTION

Recent advancements have significantly broadened the scope of 3D virtual human technology, with its applications permeating diverse fields such as animation, gaming, human-computer interaction, and digital reception. Central to this research is the development of credible, personalized co-speech gestures. Speech-driven gesture generation, facilitated by deep learning, offers an efficient alternative to traditional motion capture systems, which typically require extensive manual input.

However, a fundamental challenge within this domain is the identification and integration of numerous input conditions essential for effective gesture synthesis. This complexity stems from a variety of factors, such as acoustic nuances, semantic content, emotional expressions, personality traits, and demographic elements like gender and age. These components collectively influence the dynamics of co-speech gestures, necessitating sophisticated models that can interpret and synthesize naturalistic human movements based on verbal communication.

Previous approaches<cit.> have employed manual labels and diverse feature inputs to facilitate the synthesis of personalized gestures. However, these methods rely heavily on varied unstructured feature inputs and necessitate complex multimodal processing, posing significant barriers to the practical application and broader adoption of virtual human technologies.

The fuzzy inference strategy, derived from the concept of fuzzy logic<cit.>, has proven particularly useful in fields requiring the management of uncertain or imprecise information. This approach is noted for its effectiveness in applications such as speech-emotion recognition<cit.> and audio classification<cit.>. Unlike methods that rely on explicit classification outputs, fuzzy inference provides a spectrum of feature information, transitioning from a limited explicit discrete space to an expansive implicit continuous fuzzy space. This fuzzy space aligns more closely with real-world scenarios. Psychological research underscores the importance of various factors in speech<cit.>, which, when considered as fuzzy features, are intricately linked to co-speech gestures. Zhang et al.<cit.> pioneered the successful integration of fuzzy inference strategies for generating personalized co-speech gestures.

Achieving high gesture-speech synchronization while maintaining naturalness poses a significant challenge in the field of speech-driven gesture generation. Recent advancements have pivoted towards employing Transformer and Diffusion-based models, enhancing the efficiency and flexibility of gesture-generation technologies. Notable innovations in this domain include Diffuse Style Gesture<cit.>, Diffuse Style Gesture+<cit.>, GestureDiffuClip<cit.>, and LDA<cit.>. However, these methods sometimes struggle with maintaining an optimal correlation between gesture and speech, which can detract from the naturalness of the gestures produced.

The emergence of Diffusion Transformers (DiTs) in fields such as text-to-image generation <cit.> and text-to-video tasks like Sora[https://openai.com/sora], which utilize Adaptive Layer Normalization (AdaLN), represents a substantial leap forward. These models introduce a uniform conditional mechanism across all tokens, significantly improving the representation of conditional and output features. This enhancement is poised to effectively model the complex relationship between speech and gestures. Despite the Persona-Gestor <cit.> architecture's ability to generate high-quality co-speech gestures using DiT, it faces challenges such as high memory demands and slower inference speeds due to its convolutional and transformer-based structure. 

The Mamba-bases frameworks <cit.>, a relatively new entrant, has garnered attention for its simplicity, efficiency, and flexibility. These findings also underscore the Mamba architecture as a formidable competitor to the Transformer architecture, offering compelling alternatives in handling complex tasks efficiently. 

In this research, we introduce DiM-Gesture, a groundbreaking model designed to synthesize personalized gestures exclusively from raw speech audio. This approach employs a fuzzy feature inference strategy within its condition extractor, utilizing the Mamba architecture, and integrates an AdaLN Mamba-2 in a diffusion-based module. DiM-Gesture transitions from relying on explicit conditions to a refined, continuous representation of fuzzy features, effectively capturing a wide range of stylistic nuances and specific audio details. These features are amalgamated into a unified latent representation, facilitating the synthesis of intricate 3D full-body gestures. The implementation of AdaLN significantly bolsters the model's ability to articulate the complex relationship between speech and gestures. Leveraging a diffusion process, the model is capable of generating diverse, high-fidelity gesture outputs, showcasing its effectiveness in gesture synthesis.

For clarity, our contributions are summarized as follows:


  * Introduction of a pioneering infer Mamba-based fuzzy features strategy: This strategy allows the synthesis of a broader range of personalized gestures solely from speech audio, eliminating the need for style labels or additional inputs. The Mamba-based fuzzy feature extractor significantly enhances the system's usability capabilities.


  * Integration of the AdaLN Mamba-2 architecture within the diffusion model: This architectural choice improves the modeling of the intricate interplay between speech and gestures. Our findings validate that the Mamba architecture is a formidable competitor to traditional Transformer architectures, demonstrating an optimal balance of naturalness and synchronization in gesture generation.


  * Extensive subjective and objective evaluations: These evaluations confirm that our model surpasses current state-of-the-art approaches, highlighting the exceptional capability of our method to generate credible, speech-appropriate, and personalized gestures.





§ RELATED WORK

The discussion presented provides a concise overview of Transformer-based and diffusion-based generative models within the realm of speech-driven gesture generation.

DiffMotion <cit.> marks an innovative use of diffusion models in gesture synthesis, integrating an LSTM to generate diverse gestures.  Alexanderson et al. <cit.> have refined DiffWave by replacing dilated convolutions to optimize the potential of transformer architectures. Conformers <cit.> employ classifier-free guidance to enhance style expression. GestureDiffuCLIP <cit.> utilizes transformers and AdaIN layers to integrate style guidance within the diffusion process. DiffuseStyleGesture (DSG) <cit.> and its extension DSG+ <cit.> incorporate cross-local attention and layer normalization within their transformer models. However, these methodologies often face challenges in striking an optimal balance between gesture and speech synchronization, potentially leading to gestures that seem either underrepresented or excessively matched to the speech. Persona-Gestor <cit.> introduces a fuzzy feature extractor that leverages a 1D convolution to process raw speech audio for feature extraction, combined with an Adaptive Layer Normalization (AdaLN) transformer <cit.> to model the intricate correlation between these features and the resulting gesture sequence. While this method achieves superior motion quality, the significant memory requirements and slower inference speeds of convolutional and transformer architectures remain a challenge.

In this study, we implement a fuzzy feature inference strategy to implicitly capture nuanced features within speech audio. We synthesize natural, personalized co-speech gestures exclusively based on raw speech audio, leveraging the Mamba architecture. Furthermore, we employ the AdaLN Mamba-2 architecture in place of traditional AdaLN transformers. This adaptation preserves the model’s ability to capture the complex interplay between speech and gestures while ensuring the generation of high-quality actions. Notably, this approach significantly reduces memory requirements and enhances inference speed, offering a more efficient alternative for gesture synthesis in virtual human interactions. 






§ SYSTEM OVERVIEW

DiM-Gesture, an end-to-end Mamba and diffusion-based architecture, processes raw speech audio as its sole input to synthesize personalized gestures. This model adeptly balances naturalness with synchronized alignment to the corresponding speech, ensuring the gestures it generates are both realistic and timely, enhancing the interaction experience in virtual environments. 




 §.§ Problem Formulation

In this study, we frame the problem of co-speech gesture generation as a sequence-to-sequence translation challenge, where the objective is to map a sequence of speech audio features, denoted as a = a_1:T = [a_1, …, a_t, …, a_T] ∈ℝ^T, to a corresponding sequence of full-body gesture features, represented as g^0 = g^0_1:T = [g^0_1, …, g^0_t, …, g^0_T] ∈ℝ^T × (D+6). Here, each g^0_t ∈ℝ^(D+6) comprises 3D joint angles alongside root positional and rotational velocities at frame t, with D indicating the number of joint channels.

We define the probability density function (PDF), p_θ(·), to approximate the actual gesture data distribution p(·), facilitating efficient sampling. Our aim is to generate a non-autoregressive whole pose sequence (g^0) from its conditional probability distribution given the audio sequence (a) as a covariate:


    g^0 ∼ p_θ(g^0 | a) ≈ p(·) := p(g^0 | a)


This formulation underscores the use of a denoising diffusion model trained to approximate the true conditional distribution of gestures given speech, illustrating the direct relationship between the audio inputs and gesture outputs, thus setting the stage for our model to learn this complex mapping effectively.




 §.§ Model Architecture

The architecture of DiM-Gesture, illustrated in Figure <ref>, is composed of four main components designed to streamline the synthesis of personalized gestures from speech audio. These components include:

1) Mamba-based Style Fuzzy Feature Extractor, 2) AdaLN Mamba-2 Architecture, 3) Gesture Encoder and Decoder, and 4) Diffusion Network.

Together, these components form a cohesive system that not only captures the complexity of human gestures relative to speech but also enhances the quality and personalization of the generated gestures.



  §.§.§ Mamba-based Fuzzy Feature Extractor


This module employs a fuzzy inference strategy, which means it does not rely on explicit classification outputs. Instead, it provides implicit, continuous, and fuzzy feature information, automatically learning and inferring the global style and specific details directly from raw speech audio. Illustrated in Figure <ref> and Figure <ref>, this module is a dual-component extractor comprising both global and local extractors. The local extractor utilizes the WavLM large-scale pre-trained model<cit.> to process the audio sequence into discrete tokens. We selected WavLM for its proficiency in capturing the complex attributes of speech audio, which allows it to effectively represent universal audio latent features, denoted as z_a. This sophisticated approach ensures a nuanced translation of audio features into gesture-relevant data, enhancing the fidelity and personalization of the generated gestures.

We implement a global style extractor within our system, utilizing the Mamba module <cit.> to process z_a, the universal audio latent representations. This global extractor is adept at automatically capturing and embedding global fuzzy style information from z_a, producing an output token z_last∈ℝ^1× D', which is the last output of the Selective State Space Model(SSM). This token is subsequently broadcasted and amalgamated with z_a ∈ℝ^T'× D' to forge a unified latent representation z_l ∈ℝ^T× D”. By merging both local and global insights, our architecture enhances the representational fidelity of the entire sequence, tailored for co-speech gesture generation. This unified latent representation is then channeled to the downsampling module for further refinement. 

The downsampling module, crucial for aligning each latent representation with its corresponding sequence of encoded gestures, is seamlessly integrated into the condition extractor. Initially, we explored linear alignment strategies, similar to those employed in DSG<cit.> and DSG+<cit.>. However, these methods often led to issues such as foot-skating. To address this, we have implemented a Conv1D layer with a kernel size of 201 within our architecture. This configuration enables the mapping of every 201-length target token output from the WavLM to a single gesture frame, thereby enhancing the precision of gesture synchronization. The output of this module, c_1, serves as a unified latent representation that encapsulates both encoded audio features and the diffusion time step n, ensuring a coherent and accurate gesture generation process. 





  §.§.§ AdaLN Mamba-2

The AdaLN's fundamental purpose is to incorporate a conditional mechanism that uniformly applies a specific function across all tokens, thereby significantly improving the model's capacity for representing both conditional and output features with enhanced efficiency. It offers a more sophisticated and nuanced approach to modeling, enabling the system to capture and articulate the complex dynamics between various input conditions and their corresponding outputs. Consequently, this leads to an improvement in the model's predictive accuracy and ability to generate outputs more aligned with the given conditions.  

Diffusion Transformers (DiTs) exemplify a sophisticated advancement in diffusion model architectures, incorporating an AdaLN-infused transformer framework primarily for text-to-image synthesis. This methodological enhancement has substantially lowered Fréchet Inception Distance (FID) scores, demonstrating improved image quality. The utility of DiTs has recently expanded to include text-conditional video generation, illustrating their versatility. Furthermore, DiTs have shown potential in co-speech gesture generation <cit.>, marking a significant step in applying these models to sequence-based tasks. However, the inherent quadratic space complexity associated with Transformers results in substantial memory consumption and slower inference speeds, presenting practical challenges in real-time applications. 

In contrast to traditional Diffusion Transformers, our approach integrates the Mamba-2 architecture <cit.> as a replacement for the conventional transformer module. This strategic adaptation leverages the minimal memory footprint and enhanced processing efficiency of Mamba-2, significantly accelerating inference speeds without sacrificing output quality. This novel substitution is pivotal in addressing the challenges associated with real-time, speech-driven gesture synthesis, ensuring efficient and high-quality performance.

The module involves regressing the dimension-wise scale and shift parameters (γ and β), derived from the fuzzy feature extractor output c_1:T, rather than directly learning these parameters, as depicted in Figure <ref>. In each AdaLN Mamba-2 stack, a latent feature z^n_1:T,m is generated, combining condition information and gesture using AdaLN and the Mamba-2 architecture. The index m ranges from 1 to M, where M represents the total number of AdaLN Mamba-2 stacks. Furthermore, as illustrated in Figure <ref>, the final layer utilizes the same fuzzy features, supplemented by a scale and shift operation to fine-tune the gesture synthesis.

This method enables the creation of detailed gesture sequences directly from speech audio, obviating the need for discrete style labels or additional inputs. It significantly enhances the model's ability to generate personalized and contextually aligned gestures, providing a refined and context-sensitive gesture synthesis capability. Moreover, compared with the AdaLN transformer, the Mamba-2 architecture reduces the memory footprint and improves inference speed, while still ensuring high generation quality. 



  §.§.§ Gesture Encoder and Decoder

The architecture of the gesture encoder and decoder is designed to process the gesture sequence, as illustrated in Fig. <ref> and Fig. <ref>. The gesture encoder employs a Convolution1D layer with a kernel size of 3 to encode the initial sequence of gestures g into a hidden state h ∈ℝ^T × D”. Our experimental results indicate that a kernel size of 1 tends to produce animation jitter. In contrast, a kernel size of 3 effectively mitigates this issue by capturing the spatial-temporal relationships inherent in gesture sequences.

The gesture decoder then reduces the feature dimension of the transformer's output D” back to the original dimension D, which corresponds to the number of channels representing skeleton joints, resulting in the output of the predicted noise (ϵ_θ). Employing a 1D kernel of size 1 in the input sequence allows the model to effectively extract meaningful features and relationships between adjacent joint channels, enhancing the quality and coherence of the generated gestures.

 



 §.§ Training and Inferencing with DDPM
 
The diffusion process in our architecture is crucial for reconstructing the conditional probability distribution between gestures and fuzzy features. This involves employing a systematic sampling approach from this reconstructed distribution, facilitating the generation of diverse gestures.

Following our previous work, Persona-Gestor <cit.>, incorporating the Denoising Diffusion Probabilistic Model (DDPM) into our approach. 

The model operates through two principal processes: the diffusion process and the generation process. In the diffusion process during training, the model gradually transforms the original gesture data (g^0) into white noise (g^N) by optimizing a variational bound on the data likelihood. This transformation is characterized by progressively adding noise to the data in a controlled manner.

During inference, the generation process endeavors to reverse this transformation. It recovers the original data from the noise by reversing the noising process through a Markov chain, employing Langevin sampling <cit.>. This technique ensures the effective and accurate reconstruction of the gesture data from its noised state. The Markov chains utilized in both the diffusion and generation processes ensure a coherent and systematic transition between stages, thereby maintaining the integrity and quality of the generated gestures. This dual-process framework allows the model to efficiently handle and synthesize complex gesture data, reflecting the dynamic nature of human movement. The Markov chains in the diffusion process and the generation process are:

    p(g^n|g^0) = 𝒩(g^n; √(α^n) g^0, (1-α^n)I)      and
       p_θ(g^n-1|g^n, g^0) = 𝒩(g^n-1; μ̃^n(g^n, g^0), β̃^n I),

where α^n := 1 - β^n and α^n := ∏_i=1^n α^i. As shown by<cit.>, β^n is a increasing variance schedule β^1,...,β^N with β^n ∈ (0,1), and β̃^n := 1-α^n-1/1-α^nβ^n.

The training objective centers on optimizing the parameters θ by minimizing the Negative Log-Likelihood (NLL). This is achieved via the Mean Squared Error (MSE) loss, which quantifies the discrepancy between the true noise, represented by ϵ∼𝒩(0, I), and the predicted noise ϵ_θ:


    𝔼_g^0_1:T, ϵ, n[||ϵ - ϵ_θ(√(α^n g^0)+√(1-α^n)ϵ , a_1:T,n)||^2],
  Here ϵ_θ is a neural network (see figure <ref>), which uses input g_t^0 , a_t-1 and n that to predict the ϵ, and contains the similar architecture employed in <cit.>. The complete training procedure is outlined in Algorithm <ref>.



After the training phase, we employ variational inference to generate a complete sequence of new gestures that match the original data distribution (g^0 ∼ p_θ(g^0 | a)). This is facilitated by the sampling procedure detailed in Algorithm <ref>. During this process, we sample the entire sequence g^0 from the learned distribution.

The term σ_θ represents the standard deviation of the conditional distribution p_θ(g^n-1| g^n), which is crucial for accurately capturing the variability and intricacies of the transition between different diffusion stages. For our model, we set σ_θ := β̃^n, where β̃^n is a predetermined scaling factor that adjusts the noise level during each diffusion step, allowing for a controlled smoothing and detail enhancement in the generation process.



During inference, we send the entire sequence of raw audio to the condition extractor component. The output of this component is then fed to the diffusion model to generate the whole sequence of accompanying gestures (g^0).



§ EXPERIMENTS

To validate our approach, we utilized two co-speech gesture datasets: ZEGGS <cit.> and BEAT <cit.>. Our experiments focused on producing full 3D body gestures, including finger motions and locomotion.



 §.§ Dataset and Data Processing




  §.§.§ Datasets

The ZEGGS dataset features a comprehensive collection of emotional expressions, enabling the exploration of gesture generation across a spectrum of emotions. Conversely, the BEAT dataset is distinguished by its focus on personalized movements, capturing the unique gesture styles of various individuals. This diversity in datasets allows for robust testing and enhancement of models aimed at generating nuanced and contextually appropriate co-speech gestures.



  §.§.§ Speech Audio Data Process

Due to the WavLM large model being pre-trained on speech audio sampled at 16 kHz, we uniformly resampled all audio from the ZEGGS and BEAT datasets down from 44.1 kHz to match this frequency, ensuring compatibility and optimal performance.



  §.§.§ Gesture Data Process

We concentrate exclusively on full-body gestures, employing the data processing techniques detailed by Alexanderson et al. <cit.>. This includes capturing translational and rotational velocities to accurately delineate the root's trajectory and orientation. The datasets are uniformly downsampled to a frame rate of 20 fps. To ensure precise and continuous representation of joint angles, we utilize the exponential map technique <cit.>. All data are segmented into 20-second clips for training and validation purposes. For user evaluation, the generated gesture sequences are divided into 10-second clips to enhance evaluation efficiency.



 §.§ Model Settings

Our experiments utilize three AdaLN Mamba-2 blocks, with each Mamba-2 configured with a 256 SSM state expansion factor, a local convolution width of 4, and a block expansion factor of 2. This encoding process transforms each frame of the gesture sequence into hidden states h ∈ℝ^1280. We employ the pre-trained WavLM Large model[https://github.com/microsoft/unilm/tree/master/wavlm] for audio processing.

The diffusion model uses a quaternary variance schedule, starting from β_1 = 1 × 10^-4 to β_N = 8 × 10^-2 with a linear beat schedule, and a total of N = 1000 diffusion steps. The training batch size is set to 32 per GPU.

Our model was tested on an Intel i9 processor with a 4090 GPU, in contrast to the A100 GPU used by Persona-Gestor. The training times were approximately 6 hours for ZEGGS and 22 hours for BEAT. 





 §.§ Visualization Results

Our system excels in generating personalized gestures that are contextually aligned with speech. By leveraging the Mamba-based fuzzy inference strategy, it autonomously derives fuzzy features directly from speech audio. The results demonstrate that our AdaLN Mamba-2 framework produces gesture actions comparable to those generated by Persona-Gestor, which uses convolution global fuzzy feature extractor and AdaLN transformers. 



Figure <ref> illustrates the visual outcomes of gestures aligned with the emotional valence conveyed by the audio. For example, the system generates gestures of joy in response to happy audio cues (see Figure <ref>) and gestures of sadness for sorrowful audio (as shown in Figure <ref>). Additionally, the system can infer age-related characteristics or other nuanced states from the speech audio (as depicted in <ref> and <ref>). These visual comparisons demonstrate that DiM-Gesture can produce co-speech gestures of comparable quality to those generated by Persona-Gestor. 



 §.§ Subjective and Objective Evaluation

In line with established practices in gesture generation research, we conducted a series of subjective and objective evaluations to assess the co-speech gestures generated by our proposed DiM-Gesture (DiM) model.

We adopted slightly varied baselines for different datasets. For the ZEGGS dataset, we employed LDA<cit.>, and DiffuseStyleGesture (DSG)<cit.>. For the BEAT dataset, we utilized the same baseline models as in ZEGGS but replaced DSG with DSG+<cit.> and introduced GestureDiffuCLIP (GDC)<cit.> as an additional baseline model. All baseline models employed are based on transformer and diffusion architectures.

In our experiments with the ZEGGS and BEAT datasets, we extended the original LDA, DSG, and DSG+ models to encompass all styles within these datasets. Efforts to adapt LDA to include finger motions encountered significant challenges, resulting in unsatisfactory gesture-generation outcomes. Consequently, we utilized gestures generated by the LDA model, excluding finger movements, for our analysis. 





  §.§.§ Subjective Evaluation


For comprehensive subjective evaluations, we employ three metrics: human-likeness, appropriateness, and style-appropriateness. Human-likeness assesses the naturalness and resemblance of gestures to authentic human movements, independent of speech. Appropriateness evaluates the temporal alignment of gestures with speech rhythm, intonation, and semantics, ensuring a natural flow. Style-appropriateness measures the similarity between generated gestures and their original human counterparts.

We conducted a user study employing pairwise comparisons, as recommended by <cit.>. In each trial, participants were presented with two 10-second video clips side by side, generated by different models, including the Ground Truth (GT), for direct comparison. Participants were instructed to select the clip they preferred based on specified evaluation criteria. Preferences were quantified on a scale from 0 to 2, with the non-selected clip in each pair receiving an inverse score (e.g., a -2 score for the non-chosen clip if the chosen one received a score of 2). A score of zero indicated a neutral preference.

Given the extensive range of styles in the ZEGGS (19) and BEAT (30) datasets, evaluating each style individually was deemed impractical. Therefore, we employed a random selection method, assigning each participant a subset of five styles from ZEGGS and six characters from BEAT for evaluation. Importantly, none of the selected audio clips were included in the training or validation sets, ensuring the integrity of the assessments.

Thirty volunteer participants—14 males and 16 females aged between 18 and 35—were recruited for this study. All participants demonstrated a high level of English proficiency, essential for accurately interpreting and responding to the tasks involved.

One-way ANOVA and posthoc Tukey, multiple comparison tests, were conducted to determine if there were statistically significant differences among the models' scores across the three evaluation aspects. The results are presented in Table <ref>, offering detailed insights into the performance variances observed among different models regarding human-likeness, appropriateness, and style-appropriateness. 

The results from the ZEGGS and BEAT datasets show that the Ground Truth (GT) achieves the highest scores (0.93±1.22 for ZEGGS and 0.65±1.16 for BEAT), exhibiting statistically significant differences (p<0.001) in human-likeness evaluations compared to model-generated gestures. The GT features a diverse yet limited selection of gestures, each characterized by distinct traits that enhance the realism of movements. However, these specific gestures fall into the long-tail distribution of the datasets, presenting substantial challenges to the learning capabilities of the models. Furthermore, the uniqueness of these gestures significantly influences the appropriateness and style-appropriateness scores. 

The evaluations of the ZEGGS and BEAT datasets showed no statistically significant differences (p>0.05) between our DiM model and Persona-Gestor (PG) across all three metrics. However, PG scored slightly higher in Human-likeness on the ZEGGS dataset, whereas DiM attained marginally higher scores in Appropriateness and Style-appropriateness. In contrast, on the BEAT dataset, DiM scored slightly higher across all subjective evaluation metrics. These outcomes demonstrate that DiM is capable of generating co-speech gestures of comparable quality to those produced by PG, affirming the robustness and effectiveness of our approach in synthesizing natural and well-aligned gestures across diverse datasets. 



  §.§.§ Objective Evaluation

We employ three objective evaluation metrics to assess the quality and synchronization of generated gestures: Fréchet Gesture Distance (FGD) in both feature and raw data spaces<cit.>, and BeatAlign<cit.>. Inspired by the Fréchet Inception Distance (FID)<cit.>, FGD evaluates the quality of generated gestures and has shown moderate correlation with human-likeness ratings, surpassing other objective metrics<cit.>. BeatAlign, on the other hand, assesses gesture-audio synchrony by calculating the Chamfer Distance between audio beats and gesture beats, thus providing insights into the temporal alignment of gestures with speech rhythms.

Table <ref> presents our results, underscoring the state-of-the-art performance of our method in objective evaluations using the Fréchet Gesture Distance (FGD) metrics. Our model achieves superior performance (28.16 for ZEGGS) compared to other architectures, generating gestures that closely align with the Ground Truth (GT). It also matches the BeatAlign scores (0.67 for ZEGGS) of other models, except for GestureDiffuClip (GDC) which scores slightly higher (0.69 for BEAT), highlighting its efficacy in producing co-speech gestures that synchronize accurately with speech rhythms. Despite GDC's high BeatAlign score, corroborating user feedback indicates its overemphasis on prosodic cues results in frequent high-frequency gestures that, while technically accurate, reduce the naturalness of the gestures. On the BEAT dataset, our model's FGD scores (276.32) are marginally lower than Persona-Gestor (276.25), reflecting competitive performance in gesture quality.



 §.§ Ablation Studies


Ablation studies were conducted to assess the influence of crucial components within our model, particularly focusing on the Mamba-based fuzzy feature extractor and the AdaLN Mamba-2 architecture. These studies aimed to elucidate how each component contributes to the overall performance and effectiveness of the gesture synthesis process. 


  §.§.§ Ablation of Mamba-based Style Fuzzy Feature Extractor

For the Mamba-based style fuzzy feature extractor, we conducted an investigation into the implications of replacing the Mamba architecture with a convolutional style extractor (similar to the approach in PG), which we termed DiM-ConvSE. This experiment aimed to assess the differential impacts of the Mamba and convolutional methodologies on the effectiveness and fidelity of the synthesized gestures. 

The results, as detailed in Table <ref>, reveal no statistically significant differences (p>0.05) in the style-appropriateness metrics between DiM and DiM-ConvSE for both the ZEGGS and BEAT dataset experiments. These findings underscore the equivalence of the Mamba-based style fuzzy feature extractor in performance with the convolution-based extractor, highlighting its efficacy in gesture synthesis. Additionally, as indicated in Table <ref>, the Mamba architecture demonstrates less parameter count and lower memory consumption than DiM-ConvSE.





  §.§.§ Ablation of AdaLN Mamba-2

In our ablation study, we explored the impact of different versions of the Mamba architecture by replacing the AdaLN Mamba-2 with Mamba-1 <cit.>. This experiment was designed to assess the efficacy of the AdaLN Mamba-2 in handling the intricate dynamics of gesture generation compared to its predecessor, Mamba-1. The main focus was to determine whether the advancements in Mamba-2, particularly the integration of adaptive layer normalization (AdaLN), provide substantial improvements in gesture synthesis quality, efficiency, and model responsiveness. The results of this comparison are detailed in Table <ref> and Table <ref>, illustrating differences in gesture quality, computational efficiency, and performance metrics between the two Mamba versions.

In our ablation studies on the AdaLN Mamba-2, substituting the Mamba-2 module with Mamba-1 led to a significant decline in performance across all metrics. This deterioration can be attributed to the earlier architecture's inability to precisely synchronize speech rhythm and capture stylistic nuances like its successor. These findings highlight the integrated enhancements in Mamba-2, particularly its advanced capability to effectively handle the intricate dynamics of speech-driven gesture generation. However, it is worth noting that DiM-AdaLN-Mamba-1 exhibits a faster inference time.




§ DISSCUSTION AND CONCLUSION

In this study, we present DiM-Gesture, an innovative network architecture designed to generate personality-specific gestures directly from raw speech audio using a Mamba-based architecture. DiM-Gesture incorporates a Mamba-based style fuzzy feature extractor and an AdaLN Mamba-2 diffusion architecture, facilitating the seamless synthesis of nuanced, personality-driven gestures. Our approach achieves a quality of action generation comparable to that of the AdaLN Transformer architecture (Persona-Gestor) while requiring significantly less memory and substantially enhancing generation speed than Persona-Gestor. This demonstrates DiM-Gesture's capability to provide efficient and high-quality gesture synthesis.

The Mamba-based fuzzy feature extractor employs a fuzzy feature inference strategy within its dual-component module to autonomously infer both fuzzy stylistic features and specific audio details. These elements are merged into a unified latent representation, enabling the generation of speaker-aware personalized 3D full-body gestures. This approach integrates a pivotal innovation in synthesizing personality-driven gestures by leveraging automatically inferred fuzzy features, thereby eliminating the need for explicit style labels or additional features. Such advancements facilitate an end-to-end gesture generation process that authentically reflects the speaker's unique characteristics directly from raw speech audio. The integration of fuzzy feature inference streamlines the creation process, enhancing both generalization capabilities and user accessibility.

The AdaLN Mamba-based mechanism, a conditional architecture, uniformly applies a specific function across all sequence tokens, significantly enhancing the model's ability to capture and represent both conditional dependencies and output characteristics efficiently. Like AdaLN transformers, AdaLN Mamba-2 enhances the understanding and processing of the complex interactions between continuous fuzzy features as conditional inputs and resultant gesture synthesis. This leads to improved model performance and output fidelity. Ultimately, DiM-Gesture employs a diffusion mechanism to produce a diverse spectrum of gesture outputs, showcasing its capability to handle varied and nuanced gesture synthesis effectively.

Our study highlights key areas for enhancement: There remains a noticeable disparity between the current batch generation form and the desired real-time generation capability. Bridging this gap is crucial for applications requiring immediate gesture synthesis, such as live interactions or performances.  
 



§ ACKNOWLEDGMENTS

This work was partially supported by the "Pioneer" and "Leading Goose" R&D Program of Zhejiang (No.2023 C01212), the National Key Research and Development Program of China (No.2022YFF 0902305), the Public Welfare Technology Application Research Project of Zhejiang (No.LGF21F020002, No.LGF22F020008), the Key Program and development projects of Zhejiang Province of China (No.2021C03137), and the Key Lab of Film and TV Media Technology of Zhejiang Province (No.2020E10015).




IEEEtran



 Paper Title: Automated Software Vulnerability Static Code Analysis Using Generative Pre-Trained Transformer Models.



Extracting doubly-excited state lifetimes in helium directly in the time domain with attosecond noncollinear four-wave-mixing spectroscopy
    Stephen R. Leone
    August 5, 2024
==========================================================================================================================================



§ ABSTRACT 
Generative Pre-Trained Transformer models have been shown to be surprisingly effective at a variety of natural language processing tasks – including generating computer code. However, in general GPT models have been shown to not be incredibly effective at handling specific computational tasks (such as evaluating mathematical functions). 
In this study, we evaluate the effectiveness of open source GPT models, with no fine-tuning, and with context introduced by the langchain and localGPT Large Language Model (LLM) framework, for the task of automatic identification of the presence of vulnerable code syntax (specifically targeting C and C++ source code). This task is evaluated on a selection of 36 source code examples from the NIST SARD dataset, which are specifically curated to not contain natural English that indicates the presence, or lack thereof, of a particular vulnerability (including the removal of all source code comments). The NIST SARD source code dataset contains identified vulnerable lines of source code that are examples of one out of the 839 distinct Common Weakness Enumerations (CWE), allowing for exact quantification of the GPT output classification error rate. A total of 5 GPT models are evaluated, using 10 different inference temperatures and 100 repetitions at each setting, resulting in 5,000 GPT queries per vulnerable source code analyzed. 
Ultimately, we find that the open source GPT models that we evaluated are not suitable for fully automated vulnerability scanning because the false positive and false negative rates are too high to likely be useful in practice. However, we do find that the GPT models perform surprisingly well at automated vulnerability detection for some of the test cases, in particular surpassing random sampling (for some GPT models and inference temperatures), and being able to identify the exact lines of code that are vulnerable albeit at a low success rate. The best performing GPT model result found was Llama-2-70b-chat-hf with inference temperature of 0.1 applied to NIST SARD test case 149165 (which is an example of a buffer overflow vulnerability), which had a binary classification recall score of 1.0 and a precision of 1.0 for correctly and uniquely identifying the vulnerable line of code and the correct CWE number. 
Additionally, the GPT models are able to, with a rate quantifiably better than random sampling, identify the specific line of source that contains the identified CWE for many of the NIST SARD test cases. 









§ INTRODUCTION



Generative Pre-Trained Transformer (GPT) models, built using the attention architecture <cit.>, have been shown to be surprisingly effective at a variety of information processing tasks including computer code generation and natural language text summarization <cit.>. GPT models are intended to work very well with language, and are typically referred to as a class of Large Language Model (LLM). Notably however, GPT models can, in certain settings, also perform algorithmic tasks such as learning the greatest greatest common divisor algorithm <cit.>. The surprising capability of GPT models to perform well at tasks involving computer code leads to the natural question of whether GPT models could be used for determine security properties of static software samples. In this study, we examine whether current open source GPT models can be used to correctly identify software vulnerabilities and weaknesses in source code correctly. 


The evaluation of our study is based on a set of 36 labeled examples from the highly credible NIST SARD dataset <cit.>. Each source code example in this dataset is meticulously mapped to the Common Weakness Enumeration (CWE) framework <cit.>, providing a solid foundation for our research on software vulnerability detection. The Common Weakness Enumeration (CWE) is a categorization of known types of software weaknesses, hardware weaknesses, and software vulnerabilities. 



Static code vulnerability scanners are an important field of study because they speed up the process of identifying and thus patching software and even potentially hardware vulnerabilities and flaws that cause security problems in computer systems. There are a wide variety of static code analysis scanners that have been developed, usually focusing on specific domains or specific languages <cit.>. A challenge with static vulnerability analysis is that the outputs typically have high false positive rates because the tools usually use string pattern matching to identify known syntax representing an instance of a specific code flaw. The main advantage of static code analysis is that it is much faster and easier to apply than dynamic program analysis. 



Ref. <cit.> studies the task of repairing CWE examples using GPT models. Ref. <cit.> examined source code vulnerability detection using several code language models (namely BERT variants) on a large dataset of software examples. 


Our study introduces novel methods compared to existing research, particularly in the use of highly curated datasets that include the exact vulnerable line of source code. The structured GPT prompt we employ with these GPT models enables the automated extraction of vulnerability detection data from the output, facilitating large-scale, repeatable, accuracy analysis. This dataset combined with structured JSON prompting and parsing allows us to fully leverage the capabilities of GPT models for automated vulnerability detection, requiring no manual processing of the GPT output. In total, we generate and process 100 · 10 · 36 · 5 = 180000 discrete GPT text outputs (each of these being a separate, independent, and complete text inference computations. All of the GPT model computations are performed locally, on a local GPU compute server, using open source GPT models. 












§ METHODS



The five open-source GPT models used in this study are summarized in Table <ref>. These GPT models were selected for their general performance among available GPT models - and all of these models can generate and parse computer code. Although a large number of different GPT models could be utilized for this study - due to the high computation time required to perform these experiments, a small number of GPT models were selected. These GPT models are from the huggingface GPT repository <cit.> and the GPT models are run using the Python3 modules  <cit.> and , using a local compute server that has four Nvidia A100 GPU's <cit.> (each with 82 Gigabytes of memory), and CUDA Version 12.4. 


GPT models can be used to predict the most likely token, given the previous N tokens (where N is the context window for current GPT models on the order of 4096-32768). When used in this manner, the text produced by the GPT models is deterministic, given that the previous tokens are the same. However, GPT models can learn and generate probability distributions of the most likely subsequent tokens - and we can add stochasticity in the text generation process by sampling from that distribution, weighted by their respective probabilities. The parameter that controls this sampling is known as temperature. This stochastic text generation is part of what makes generative machine learning models able to produce a wide variety of outputs. In the context of investigating the usage of these models as code-scanning tools, these models will only perform somewhat well at this task. Moreover, we need outputs to vary somewhat to determine the possible range of characteristics of the output since the text generation process of GPT models can be stochastic. Therefore, we vary the text generation inference temperature from 0.1 to 1.0 in increments of 0.1. The huggingface <cit.> python 3 module used to execute this GPT model inference allows temperatures up to 100, where 0 denotes deterministic output (with no sampling), and 100 denotes closer to uniform token sampling. We choose relatively low temperatures because we need the output to conform to parsable data structures. Critically, because the outputs vary significantly, we view this task as a stochastic sampling problem, where we characterize the proportion of samples with specific required properties (the most basic being whether the data could be automatically parsed). 100 samples (GPT outputs) are generated for each inference temperature, GPT model, and source code test case. 










To ground the GPT model output in the definitions of the Common Weakness Enumeration (CWE) framework, we utilize a technique known as retrieval augmented generation <cit.>. Specifically, we use a codebase known as  <cit.> that uses langchain <cit.> to augment the capabilities of GPT models by using additional word embeddings of reference documents (not as part of the model pre-training). Specifically, the source documents are split into character maximum chunk sizes of 1000, with chunk overlap of 200 characters. 

The localGPT prompt template is a combination of context and then the user prompt question. No chat history was used, meaning that each text generation call is independent. The prompt construction uses a system prompt <cit.> followed subsequently by the context of the generated word embeddings and then the user prompt. The system prompt used in this study, which is from a version of the localGPT codebase <cit.>, is given below:



You are a helpful assistant, you will use the provided context to answer user questions.
Read the given context before answering questions and think step by step. If you can not answer a user question based on the provided context, inform the user. Do not use any other information for answering user. Provide a detailed answer to the question.



The reference documents that are used to provide context for the GPT model inference calls are the following: one text file for each of the 839 CWEs that are currently defined, where the text file contains the CWE number, the name of the software weakness, and the complete description of the CWE from Mitre [<https://cwe.mitre.org>]. Next is a text file with a brief definition of the CWE framework. And lastly, two books (PDF files of the books) on software, security, and vulnerabilities (specifically for the C family languages) <cit.>. 



Table <ref> outlines the exact code test cases that are analyzed in this study. Note that NIST SARD test case 148871 is the vulnerable code that was tracked by CVE-2009-3549, and test case 149199 is an example of an implementation of CVE-2002-1869. All of the selected NIST SARD test cases have a few specific properties that we manually selected these for. First, each test case is an example of a different CWE, which makes the results representative of a variety of different possible software weaknesses. Second, each test case contains exactly one line of source code that is an example of a particular CWE; this simplifies the parsing and accuracy metrics. Third, each test case contains no variable names, text, or function names that indicate whether the source code contains a vulnerability or not. In cases where there is a string or function name that indicates what type of example the source code is, we manually remove, or change, that string before parsing by the GPT models. Fourth, for no test cases does the vulnerable line of code occur more than once in the entire source code; this eliminates ambiguity for measuring the accuracy rate of the GPT models for identifying specific lines of source code. Any comments in the source code examples are removed before parsing by the GPT models so as to not give the GPT models natural language that describes the test cases. 

Appendix <ref> shows the exact source code for 10 of these NIST SARD test cases. The phrase vulnerable line of source will be used throughout the text to refer to the exact syntax of the vulnerable line of code – not line number. Line numbers can be somewhat ambiguous, especially for these C syntax test cases, and moreover having the GPT model output the exact source code makes checking the validity of that output much clearer. 


The exact prompt that is used for all of the GPT models is given in the text box below. Note the use of the triple backticks to delineate where the source code is, and where the JSON datastructure is. The use of the triple backticks in the prompt allows the GPT outputs to be automatically parsed by using the triple backticks, and possibly the markdown code identifier, as the delimiter. Note that this triple backtick parsing does not always work, in particular if the JSON is not correctly formatted or the triple backticks are not correctly used in the GPT output. The prompt is intended to be very source code agnostic; there are no details given of the source code in the prompt. The prompt is also intended to not bias the output towards specifying a specific number of vulnerabilities. In the case of the 36 NIST SARD test cases used in this study, the number of CWEs is always exactly 1, however the goal is for this method to more general than only identifying a single CWE and moreover we want the GPT model to be able to fail in a variety of ways in order to rigorously evaluate its accuracy (in this case, identifying more CWEs than are present would constitute a high false positive rate). 



I am seeking assistance to analyze provided source code for the purpose of detecting potential Common Weakness Enumeration (CWE) vulnerabilities. 
The key requirement is for the analysis results to be exclusively formatted in JSON enclosed by triple backticks, listing only the CWEs that are detected as 'Present' in the code, thereby omitting any 'Not Present' entries for a cleaner output. 
The exact line, or lines, of code which have this CWE vulnerability needs to be specified in the returned JSON structure, along with a succinct description of what the vulnerability category is for this source code. 
The JSON output should enumerate each identified CWE with its identification number, the exact source code syntax that causes the CWE, and a description of the vulnerability, as shown in the example structure below:




Source code:


Re-check that your output contains the required JSON data-structure before emitting text. 


The actual GPT prompt is generated by replacing the vulnerable source code in place of the text . 


The usage of requesting the vulnerable code syntax and the description of the vulnerability has two uses:


    
  * Ground the GPT response to only respond based on the provided source code
    
  * To check whether the provided syntax is correct. Because of potential ambiguity with line numbers, only exact syntax string matches are checked for, not integer line number responses. 


The description field is not parsed for accuracy or correctness, but the text contained in that field is an interesting indicator of how correct (or incorrect) the GPT output is. For this reason, when reporting results (in Section <ref>), we provide specific instances of GPT outputs, including the description fields. 


Some of the raw GPT output includes characters that are not ascii and can cause issues with saving and parsing the output. Therefore, the raw GPT output is encoded as utf-8 where character parsing errors are ignored and decoded into a parsable string using utf-8. The field in the JSON output  has repeated excess whitespace removed, as is done with the reference vulnerable line of source code from each of the NIST SARD samples; this ensures that when the string equality check is performed, the only relevant characteristic with respect to whitespace in the string that is being checked is that the whitespace is delineating characters (for example, repeated spaces, tabs, newlines do not impact this accuracy check). 

The GPT output is (attempted to be) parsed in several stages, as follows:


    
  * First, if there are triple backtick delimiters in the output, then the output is split at those characters and the first enclosed set of strings is attempted to be parsed as a valid JSON structure. Interestingly, occasionally the GPT model will add in markdown code identifiers immediately after the triple backticks - in these cases we automatically check for these and remove them if they occur so that the structure can be parsed. These markdown identifiers are given in Appendix <ref>. 
    
  * If the previous step failed, then we attempt an unstructured JSON parsing approach, where output strings are incrementally deleted from the beginning and from the end until a JSON structure can be parsed (or, not). 


Using this parsing, any valid JSON strings that are produced as part of the GPT output will be automatically parsed. This allows this study to perform a systematic evaluation of the accuracy rates of the GPT inference output. This type of data parsing is also highly scalable to a large number of static code samples - and importantly does not require any manual input, tuning, or processing. In other words, this analysis pipeline is entirely automated. 





 §.§ CWE Identification Accuracy Measures



This section lists several measures for quantifying the success rate of the GPT models for identifying known vulnerable source code. 


The first four metrics that are checked are defined as follows:


    
  * Count of parsable GPT outputs (being able to extract a JSON structure and parse it into Python, with not necessarily correct JSON fields. For example a JSON with an empty list is a valid parse. However, this step does require that there is at least one JSON key, which is correctly named as . The subsequent value is allowed to be empty or have any sort of incorrect structure (which could denote that the model is returning no found CWEs) but we need at least this key to validate that a JSON datastructure was parsed. 
    
  * The CWE field is labeled as  is correct (meaning, the correct CW number is listed), but the identified source code is not correct (meaning, it is not the labeled line of code that has the vulnerability) and does not contain the vulnerable line of code either. 
    
  * The CWE is correct and identified as  and the vulnerable code is contained within the identified code syntax but also includes other lines of code from the test case that are not themselves vulnerable as defined by the NIST SARD dataset. These instances can be interpreted as being partially correct. 
    
  * Count of fully correct instances where the CWE is correct (and identified as ) and the identified line of code is correctly identified as the vulnerable code. 


This last metric (metric number 4) is the most important accuracy measure because it counts out of the 1,000 GPT outputs how many correctly performed the computation that we are attempting to prompt the GPT model to perform. 

For none of these four metrics is the Description field, or any other additional field in the JSON, parsed for correctness or even considered. All of these measures are integers, where the larger the integer is the better the GPT model performed, but the absolute correct metric is measure 4. Note that the quantities measured in metrics 2, 3, and 4 (in the immediately preceding list of measures) are always disjoint; there are no instances that are counted in two or more of these measures. An important, and informative, aspect of these metrics is finding cases where the GPT output correctly identified the vulnerable line of code. However, since line numbers of the code are not checked or requested in the prompt (since they are somewhat ambiguously defined especially for C-family languages), a potential source of error in this quantification is if there exist multiple lines of code with the same syntax, but only one is vulnerable. This does happen, but only in one of the test cases. This is NIST SARD test case 149107, which is CWE-415 (Double Free) which by its definition contains two lines of memory free calls. Therefore, for this one test instance the identification of the vulnerable line of code is ambiguous, but for all others the unique identification is valid and correct. 


Another set of metrics that are measured are binary classification error rates. This is computed by considering the GPT output for identified CWEs as a vector of True or False entries, with length equal to the number of valid CWEs (which is currently 839). This measure disregards both the requested  and  fields (and in particular does not consider whether these entries even exist in the JSON structure. For all  entries (that identifies a valid CWE number), and whose  is exactly  (allowing for different capitalizations), that is measured as a binary state of True. All other entries (including instances of identified CWE numbers that are not valid) are measured as binary classifications of False. JSON dictionaries that are completely empty are a valid GPT output that is identifying no CWEs (meaning a binary classification vector of length 839 that is all False entries). These measures are only computed on the outputs that can be parsed; all outputs that can not be parsed are not counted towards any classification accuracy.[In the case where a JSON entry contains multiple entries with the same CWE number, that case is treated as the GPT model having identified that CWE number only once; in other words, repeated entries are ignored. This specifically works for this accuracy metric of binary classification since all other fields (except  are ignored. ] These binary classification measures allow rates of false negatives, false positives, true negatives, and true positives to be measured, and then standard machine learning classification measures such as recall and precision can thus be computed. Precision is intuitively defined as the classifier (which in this case is a GPT language model) being able to not label as positive a sample that is negative. Recall is intuitively defined as the classifier (which in this case is a GPT language model) being able to find all of the positive samples. These metrics being 0 corresponds to the highest possible error rate, and 1 is the best measure these metrics can be. 

All parsing steps require the field names in the JSON datastructure to be correct (correct spelling, no ancillary text), but differing capitalizations are always allowed. 









§ RESULTS



For each of the 36 NIST SARD test cases in Table <ref>, a total of 5,000 inference calls are performed using the methods described in Section <ref> for the 5 different open-sourced GPT models. The accuracy measures from these outputs are then summarized using several high-level metrics. Tables <ref> and <ref> present two different accuracy measure summaries. Table <ref> counts how many of the GPT outputs were parsable and correctly identified the CWE number and vulnerable line of code. Table <ref> shows the binary classification measures recall and precision. Table <ref> reports the absolute integer count of how many of the GPT outputs correctly identified the single line of vulnerable source code; importantly, these counts are independent of correct identification of the relevant CWE. These counts, as a proportion out of 1000 should be compared against uniform random sampling of the total lines of source code (shown in Table <ref>). Table <ref> shows that the correct identification of vulnerable lines of code is much higher than the correct identification of the CWE number and the vulnerable line of source code, as measured in Table <ref>. Notably in Table <ref>, only for a few of the NIST SARD test cases is the count of ever correctly identifying the syntax of vulnerable line of source code zero across all GPT models and settings; namely test cases 149203, 148871, and 2046. For some of the test cases, the proportion of correctly identified vulnerable line of source code is very high. The test cases where the vulnerable line of source was not uniquely identified in Table <ref> does not necessarily mean that the correct CWE identification was always 0 (when analyzing strictly whether the identified CWE was correct) - an example of this is test case 2046. 

Table <ref> shows that for the code test cases where the accuracy is 0, the 5 different language models do not perform uniformly bad – for some of the models the accuracy is always 0 for a given test case, but then when analyzed with a different model the accuracy is non-zero (an example of this is NIST SARD test case 500757). This means that some of the GPT models do perform better than other GPT models for specific test cases. And moreover, there is not consistently a single model that always outperforms the other models. 


Table <ref> shows the best performing GPT model and inference temperature in terms of recall and precision measures. Finally, Table <ref> counts the most frequently identified false positives. The best performing GPT model and inference temperature shows that there is no specific model, or temperature setting, that consistently performs the best - and in in particular, every one of the tested GPT models was the best performer for at least one of the CWE test cases. This shows that this task does require diversity of the model architecture, and very likely the type of text on which the models were trained. Lastly, even with these best performing GPT models and inference temperatures some of the CWE test cases were never correctly identified, meaning that the accuracy rates for all GPT models and settings were always 0 - that happened for 11 out of 36 of the test cases as denoted by asterisks in Table <ref>. 


Section <ref> then gives several examples of correct JSON output from the GPT models where the vulnerable line of source code was identified, and where the description is correct in some cases and incorrect in other cases. 



Table <ref> quantifies the exact counts of false positive CWE id occurrences for the full matrix of the 5 GPT models and NIST SARD test cases. Specifically, for each NIST SARD test case and GPT model, Table <ref> reports the most frequently mis-identified CWE number. Table <ref> demonstrates two important trends. First, the proportion of false positives is very high, and this is consistent across the different GPT models. Second, for most test cases and GPT models the distribution of incorrect CWE ids is not uniformly distributed across the possible CWE numbers - it is heavily biased towards a few specific CWE numbers. For example,  shows a clear bias towards predicting CWE-124 for multiple test cases. Similarly,  shows a consistent bias for (incorrectly) predicting CWE-123, and  shows a consistent bias for (incorrectly) predicting CWE-476. 



































 §.§ Example Correctly Identified Vulnerable Lines of Code



This section enumerates several parsed JSON outputs from the various GPT models and code test cases. These outputs were hand selected to be representative of cases that are correct (in the sense that the CWE number is correctly identified, as is the vulnerable line of source code), but whose description fields are in some cases correct and in other cases incorrect. Showing explicit examples of the output is intended to primarily show the variability of the generated Description field, and also to show concretely what successful static code analysis identification looks like from the GPT models. 

No text changes were made to these GPT outputs except for formatting to fit within the text boxes. Note that these outputs are specifically examples of text from the enclosed triple backtick fields - any additional prose or output from the model is not shown. 



[title=Example of correct CWE number and code line identification in a JSON datastructure format from NIST SARD Test Case 500843. Produced from the Turdus GPT model with temperature 0.8. Note that the generated description is partially correct - CWE-476 is NULL pointer dereference but the description also mentions Divison-by-Zero which is not correct. ]






[title=Example of correct CWE number and code line identification in a JSON datastructure format from NIST SARD Test Case 500843. Produced from the zephyr-7b-alpha GPT model with temperature 0.7. Note that the generated description is entirely incorrect.]



[title=Example of correct CWE number and code line identification in a JSON datastructure format from NIST SARD Test Case 500843. Produced from the Turdus GPT model with temperature 1.0. Note that the generated description is entirely incorrect. ]





[title=Example of correct CWE number and code line identification in a JSON datastructure format from NIST SARD Test Case 500843. Produced from the Llama-2-70b-chat-hf GPT model with temperature 0.8. Note that the provided description is mostly correct in that the correct vulnerability name was identified but there is not a function named f in the source code. ]




[title=Example of correct CWE number and code line identification in a JSON datastructure format from NIST SARD Test Case 1792. Produced from the Turdus GPT model with temperature 0.3. Note that the generated description is correct.]




[title=Example of correct CWE number and code line identification in a JSON datastructure format from NIST SARD Test Case 1792. Produced from the Turdus GPT model with temperature 0.4. Note that the generated description is correct.]




[title=Example of correct CWE number and code line identification in a JSON datastructure format from NIST SARD Test Case 1792. Produced from the Turdus GPT model with temperature 0.4. Note that the generated description is correct except that the identified print statement in the description text is incorrect.]





[title=Example of correct CWE number and code line identification in a JSON datastructure format from NIST SARD Test Case 1645. Produced from the Turdus GPT model with temperature 0.5. The generated Description field in this example is correct. ]




[title=Example of correct CWE number and code line identification in a JSON datastructure format from NIST SARD Test Case 1645. Produced from the zephyr-7b-alpha GPT model with temperature 0.9. The generated Description field in this example is correct. ]




[title=Example of correct CWE number and code line identification in a JSON datastructure format from NIST SARD Test Case 149165. Produced from the Llama-2-70b-chat-hf GPT model with temperature 0.9. The generated Description field in this example is entirely incorrect. ]




[title=Example of correct CWE number and code line identification in a JSON datastructure format from NIST SARD Test Case 149165. Produced from the Llama-2-70b-chat-hf GPT model with temperature 0.1. The generated Description field in this example is correct. ]



[title=Example of correct CWE number and code line identification in a JSON datastructure format from NIST SARD Test Case 149165. Produced from the Mistral-7B-Instruct-v0.1 GPT model with temperature 0.6. The generated Description field in this example is correct. ]



[title=Example of correct CWE number and code line identification in a JSON datastructure format from NIST SARD Test Case 149165. Produced from the zephyr-7b-beta GPT model with temperature 0.2. The generated Description field in this example is correct. ]




[title=Example of correct CWE number and code line identification in a JSON datastructure format from NIST SARD Test Case 149165. Produced from the zephyr-7b-beta GPT model with temperature 0.4. The generated Description field in this example is not correct. ]




[title=Example of correct CWE number and code line identification in a JSON datastructure format from NIST SARD Test Case 500757. Produced from the Turdus GPT model with temperature 0.8. The generated Description field in this example is correct. ]




[title=Example of correct CWE number and code line identification in a JSON datastructure format from NIST SARD Test Case 1989. Produced from the Turdus GPT model with temperature 0.9. The generated Description field in this example is correct. ]






[title=Example of correct CWE number and code line identification in a JSON datastructure format from NIST SARD Test Case 149111. Produced from the Llama-2-70b-chat-hf GPT model with temperature 0.3. The generated Description field in this example is correct. ]




[title=Example of correct CWE number and code line identification in a JSON datastructure format from NIST SARD Test Case 149111. Produced from the Mistral-7B-Instruct-v0.1 GPT model with temperature 0.7. The generated Description field in this example is correct. ]



[title=Example of correct CWE number and code line identification in a JSON datastructure format from NIST SARD Test Case 149241. Produced from the zephyr-7b-alpha GPT model with temperature 0.2. The generated Description field in this example is correct. ]



[title=Example of correct CWE number and code line identification in a JSON datastructure format from NIST SARD Test Case 313. Produced from the Llama-2-70b-chat-hf GPT model with temperature 0.6. The generated Description field in this example is correct. ]









§ DISCUSSION AND CONCLUSION




This study has examined the task of correctly identifying software weaknesses and vulnerabilities in software using open source Generative Pre-Trained Transformer (GPT) models. This was able to be examined because of two specific reasons - the first is the existence of the NIST SARD dataset, which is a very well labelled dataset that includes the vulnerable line of code and the relevant CWE. The second is the specific prompting template usage that allows fully automated extraction of JSON output for the vulnerability detection - and therefore error rates can be quantified for a large range of NIST SARD test cases, GPT models, and GPT model settings. Thus, this study reports a systematic analysis of automated vulnerability detection using GPT models. 

In general, our findings are that the correct CWE identification rate varies significantly depending on the specific code test case that is analyzed. For some test cases the correct CWE identification (and relevant line code of syntax) is surprisingly high, but for other test cases the identification accuracy rate is always 0 for all GPT models and temperature settings. In general, the error rates of this vulnerability detection is too high – in particular, the false positive rate is very high. However, it is notable that for some cases, seemingly the more popular classes of CWEs that are mentioned in the type of internet text on which these GPT models were trained, the identification rates are quite good. As specific examples of the lowest error rate results (from Table <ref>), we list several specific examples of CWE identification binary classification accuracy rates (without considering rates for correctly identifying the vulnerable line of source code). Test case 1989 had a best recall of 0.928, a best precision of 0.709 (with a parsability rate of 0.42) when analyzed with  GPT model and temperature 0.1. Test case 149165 had a best recall of 1.0, a best precision of 1.0 (with a parsability rate of 1.0) when analyzed with  GPT model and temperature 0.1. Test case 149111 had a best recall of 0.727, a best precision of 0.675 (with a parsability rate of 0.77) when analyzed with  GPT model and temperature 0.5. Test case 149241 had a best recall of 1.0, a best precision of 0.5 (with a parsability rate of 1.0) when analyzed with  GPT model and temperature 0.1. Given the high accuracy rates of at least some of these test cases, it seems plausible that some of the current open source GPT models, if thorough verification is used, could correctly identify certain classes of vulnerabilities, such as buffer overflows, with very low error rates. 


It should be emphasized that this particular software security analysis task requires a number of semantic connections to be made in the GPT model inference process (and of course, the correct specification of the JSON syntax) in order for the (automated) vulnerability identification to be correct. It is not clear to what extent a pre-trained language model can learn these semantic connections - but it is possible least for some of the CWEs, these connections were learned in the standard open source GPT models that we tested. This is is a non-trivial finding, and suggests that GPT models do have some capability to identify security vulnerabilities. However, it also needs to be emphasized that such usage of current GPT models should not be used in cases where the accuracy can not be tested. In other words, these GPT models and vulnerability identification prompting should not be used in software engineering tasks – they should be used in cases where careful and precise testing of the GPT output can be performed for the purpose of examining specific characteristics of source code of interest. The primary reason for this is that the false positive rate is high - and not only is it high but the false positives are biased towards incorrectly identifying specific CWEs. In other words, simply measuring which CWE is predicted most frequently in a GPT output distribution (for a non-zero inference temperature) will typically not be correct. 


The primary limitation with increasing the scope of more code test cases to more thoroughly test the capabilities of GPT models for software security analysis is the total amount of compute time that is required for performing these inference calls, especially in a systematic study of varying inference parameters. Future studies that increase the number of labeled source code examples (e.g., from the extensive amount of data available in the NIST SARD dataset) that are used for testing the accuracy GPT automated vulnerability identification would be very relevant. 


An aspect of this study that is unclear is what details of the source code are the GPT models using to infer relevant CWE's in the source code. In particular, what semantic details are used to parse whether a CWE is more likely or less likely. In some cases, it could be that context, such as what types of libraries are used or what types of strings are being constructed (such as SQL query strings), is sufficient to reduce the total number of possibly applicable CWEs so that random guessing among those feasible candidates accounts for some of the GPT model's performance. The metrics in this study do not determine what components of the source code are important; this would require quantifying the explainability of GPT models <cit.> for this task. We leave this study open to future research, and emphasize that this type of detailed explainability analysis is absolutely a critical component of study before these GPT models are incorporated into any type of vulnerability detection in production software development and testing. 






§ ACKNOWLEDGMENTS



Sandia National Laboratories is a multi-mission laboratory managed and operated by National Technology & Engineering Solutions of Sandia, LLC (NTESS), a wholly owned subsidiary of Honeywell International Inc., for the U.S. Department of Energy’s National Nuclear Security Administration (DOE/NNSA) under contract DE-NA0003525. This written work is authored by an employee of NTESS. The employee, not NTESS, owns the right, title and interest in and to the written work and is responsible for its contents. Any subjective views or opinions that might be expressed in the written work do not necessarily represent the views of the U.S. Government. The publisher acknowledges that the U.S. Government retains a non-exclusive, paid-up, irrevocable, world-wide license to publish or reproduce the published form of this written work or allow others to do so, for U.S. Government purposes. The DOE will provide public access to results of federally sponsored research in accordance with the DOE Public Access Plan.














§ MARKDOWN COMPUTER CODE LANGUAGE IDENTIFIERS











§ NIST SARD TEST CASES



This section shows the exact source code for 10 of the NIST SARD test cases that are used in this study. The single vulnerable line of source code is highlighted as red text for each test case. 


.49
[caption= NIST SARD Test case 500843, instance of CWE-476 ,captionpos=b,label=SARD_source_code:500843,language=C,showstringspaces=false,style=base,keywordstyle=]

struct Token 
    const Token* nextArgument() const;
    const Token* next() const;
    int varId() const;
    void str() const;;
void f(const Token *first) 
    first = first->nextArgument();
    if (first)
        first = first->next();
    @first->str();@



.49
[caption= NIST SARD Test case 1792, instance of CWE-79. Note that some of the strings in this test case have been modified so that terms related to cross site scripting are removed. ,captionpos=b,label=SARD_source_code:500843,language=C,showstringspaces=false,style=base,keywordstyle=]

#include <stdio.h>
#include <cgic.h>
#include <string.h>
#include <stdlib.h>

int cgiMain() 

	cgiHeaderContentType("text/html");
	 
	fprintf(cgiOut, "<html><head>");
	fprintf(cgiOut, "<title>Text: 1</title></head>");
	fprintf(cgiOut, "<body><h1>Text</h1>");
	 
	char q[4][1024];
	unsigned int i = 0;
	for (; i < 4; ++i)
		char name[2];
		sprintf(name,"q
		cgiFormString(name, q[i], sizeof(q[i]));
		if (strlen(q[i]))
		
			@fprintf(cgiOut, "Value number 
		
			
	 
	fprintf(cgiOut, "</body></html>");
	return 0;





.49
[caption= NIST SARD Test case 1779, instance of CWE-463. ,captionpos=b,label=SARD_source_code:1779,language=C,showstringspaces=false,style=base,keywordstyle=]

main ()

char *foo;
int counter;
foo=malloc(sizeof(char)*10);
for (counter=0;counter!=14;counter++)
@foo[counter]='a';@
printf("









.49
[caption=NIST SARD Test case 1645, instance of CWE-20. ,captionpos=b,label=SARD_source_code:1645,language=C,showstringspaces=false,style=base,keywordstyle=]

#include <stdio.h>
#include <stdlib.h>

#define	MAXSIZE		40

void test(char *str)
	char buf[MAXSIZE];
	snprintf(buf, sizeof buf, "/bin/echo 
	buf[MAXSIZE-1] = 0;
	@system(buf);@


int main(int argc, char **argv)
	char *userstr;
	if(argc > 1) 
		userstr = argv[1];
		test(userstr);
	
	return 0;









.49
[caption=NIST SARD Test case 149165, instance of CWE-121. ,captionpos=b,label=SARD_source_code:149165,language=C,showstringspaces=false,style=base,keywordstyle=]

#include <stdlib.h>
#include <stdio.h>
#include <time.h>
#include <string.h>

unsigned int getRand()

	unsigned int r;
	FILE *f;
	f = fopen("/dev/urandom", "rb");
	if(f == NULL)
	
		fprintf(stderr, "Error opening file");
		exit(-1);
	
	if(fread(   r, sizeof r, 1, f) != 1)
	
		fprintf(stderr, "Error reading file");
		fclose(f);
		exit(-1);
	
	if(fclose(f) != 0)
		fprintf(stderr, "Error closing file");

	return r;


unsigned plop() 
	return getRand() 


int main(int argc, char *argv[])

	char buffer[256];
	memset(buffer, 0, sizeof(buffer));
	@buffer[plop()] = '!';@		 
	printf("
	return 0;









.49
[caption=NIST SARD Test case 2015, instance of CWE-329. ,captionpos=b,label=SARD_source_code:2015,language=C,showstringspaces=false,style=base,keywordstyle=]

#include <openssl/evp.h>

main()
EVP_CIPHER_CTX ctx;
char key[EVP_MAX_KEY_LENGTH];
char iv[EVP_MAX_IV_LENGTH];
int b=8;
RAND_bytes(key, b);
memset(iv,0,EVP_MAX_IV_LENGTH);
@EVP_EncryptInit(   ctx,EVP_bf_cbc(), key,iv);@
return;









.49
[caption= NIST SARD Test case 500757, instance of CWE-787 ,captionpos=b,label=SARD_source_code:500757,language=C,showstringspaces=false,style=base,keywordstyle=]
#include <cstdlib>
void f()

  char *p; p = (char *)malloc(10);
  @p[10] = 7;@
  free(p);





.49
[caption=NIST SARD Test case 149135, instance of CWE-489. ,captionpos=b,label=SARD_source_code:149135,language=C,showstringspaces=false,style=base,keywordstyle=]

#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>
#include <string.h>

static bool debug = false;

void promote_root() 
	if (debug) 
		 
		printf ("# You are root now...");
	


int main(int argc, char *argv[])

	if (argc > 1)
	
		unsigned i;
		const unsigned nbArgs = argc;

		for (i = 1; i < nbArgs; ++i)
		@
			if (!strncmp(argv[i],"-debug",6)) @
			
				debug = true;
				printf("Move to debugging mode");
				 
				if (strlen(argv[i]) >= 11        !strncmp(argv[i]+6,":root",5))
				
					promote_root();
				
			
			
	
	else
	
		printf("No args...");
	

	return 0;










.49
[caption=NIST SARD Test case 149203, instance of CWE-416. ,captionpos=b,label=SARD_source_code:149203,language=C,showstringspaces=false,style=base,keywordstyle=]

#include <stdlib.h>
#include <stdio.h>
#include <string.h>

typedef struct str_t str;
struct str_t 
	union 
		int   a;
		char *b;
	 foo;
;

int main()

	str container;
	char *str1;
	container.foo.a = 0;
	
	if ((container.foo.b = (char *)malloc(256*sizeof(char))) != NULL)
	
		strcpy(container.foo.b, "Falut!");
		container.foo.b[0] = 'S';
		printf("
		free(container.foo.b);
		str1 = malloc(65536); 
 		if (str1 == NULL)
        		printf("Memory allocation problem"); return 1;

		free(str1);
		
		if (container.foo.b)
			@container.foo.b[0] = 'S'; printf("
	
	return 0;









.49
[caption=NIST SARD Test case 149111, instance of CWE-134. ,captionpos=b,label=SARD_source_code:149111,language=C,showstringspaces=false,style=base,keywordstyle=]

#include <stdlib.h>
#include <stdio.h>
#include <string.h>

#define MAX_SIZE 255

typedef struct cont_t cont;
struct cont_t 
	char fmt[MAX_SIZE];
;

int main(int argc, char *argv[])

	if (argc > 1) 
		cont container = .fmt="";
		strncpy(container.fmt, argv[1],MAX_SIZE-1);
		container.fmt[MAX_SIZE - 1] = '';
		@printf(container.fmt);		@		 
	
	return 0;


















 Paper Title: TransferTOD: A Generalizable Chinese Multi-Domain Task-Oriented Dialogue System with Transfer Capabilities.


Kramers-Kronig relations via Laplace formalism and L^1 integrability
    Leonardo Ricci
    August 5, 2024
====================================================================

§ ABSTRACT 
Task-oriented dialogue (TOD) systems aim to efficiently handle task-oriented conversations, including information gathering. How to utilize ToD accurately, efficiently and effectively for information gathering has always been a critical and challenging task. Recent studies have demonstrated that Large Language Models (LLMs) excel in dialogue, instruction generation, and reasoning, and can significantly enhance the performance of TOD through fine-tuning. However, current datasets primarily cater to user-led systems and are limited to predefined specific scenarios and slots, thereby necessitating improvements in the proactiveness, diversity, and capabilities of TOD. In this study, we present a detailed multi-domain task-oriented data construction process for conversations, and a Chinese dialogue dataset generated based on this process, TransferTOD, which authentically simulates human-machine dialogues in 30 popular life service scenarios. Leveraging this dataset, we trained a TransferTOD-7B model using full-parameter fine-tuning, showcasing notable abilities in slot filling and questioning. Our work has demonstrated its strong generalization capabilities in various downstream scenarios, significantly enhancing both data utilization efficiency and system performance. The data is released in <https://github.com/KongLongGeFDU/TransferTOD>.




§ INTRODUCTION


The Task-Oriented Dialogue System (TOD) is a human-computer interaction system aims to aid users in accomplishing specific tasks or acquiring particular information, which has found extensive use in daily life and commercial applications. At present, TOD systems have displayed the capability to adapt effectively to diverse tasks, domains, and user behaviors. Nonetheless, they continue to encounter various challenges related to generality, deep understanding, proactive questioning, and other aspects.


To gather the necessary information, the system must proactively ask questions or guide users to provide the required information for filling specific slots, known as slot filling (SF) <cit.>. 
Although various approaches <cit.> has been explored to maxmize data efficiency (e.g. transfer learning and fine-tuning), traditional SF methods still rely on expert-labeled data <cit.>, which is costly and inefficient, and are limited to predetermined scenarios and time periods. These methods cannot be generalized to more general scenarios easily (certain methods even rely on external databases <cit.>), and it is difficult to ensure accurate, real, and diverse responses to user needs. Furthermore, existing datasets primarily revolve around user-driven systems, where the focus is on constructing systems that primarily respond to user inquiries and requests <cit.>. Recently, Large Language Models (LLMs) have exhibited promising performance in dialogue participation, instruction generation, and zero-shot reasoning <cit.>, which brought new ideas to solving the above problems. Research has confirmed that fine-tuned LLMs on dialogue corpora of different sizes can achieve enhanced performance across diverse tasks, domains, and even languages <cit.>. Hence, we can use LLMs to drive TOD and solve some problems that were difficult to solve in the small model era.








In the paper, we introduce TransferTOD: a multi-domain, task-oriented Chinese dialogue dataset encompassing more complex and diverse dialogue tasks, and simulating more realistic conversation scenarios. Inspired by real-world questionnaire-style information-gathering scenarios, TransferTOD facilitates interactions between users and systems to assist in information acquisition and record updating. The dataset includes 35965 turns of statements and 5460 dialogues across 30 popular life service scenarios, providing researchers with a more challenging and practically significant dataset. Considering potential human errors and the variability of the Chinese context in practical applications, we have incorporated perturbed data and data polished through various methods into the dataset.



By selecting appropriate base models and fine-tuning methods, we have successfully demonstrated that training the TransferTOD-7B model using our dataset can achieve high accuracy. This model not only can proactively ask users for missing slots and accurately fill them based on their answers but also performs efficiently in guiding fluency and generating responses. Additionally, we evaluated the quality of the models in terms of slot filling ability and semantic accuracy in guiding user responses. The results indicate that our dataset can significantly improve model performance by handling noise, increasing question diversity, and optimizing language fluency.

Summarizing, the principal contributions of our paper are as follows:



1. We construct a new dataset called TransferTOD for task-oriented dialogue generation in various lifestyle service scenarios. It consists of 30 scenarios with 5460 dialogues, and ablation experiments have demonstrated that this dataset exhibits good noise resistance, diversity, and fluency.

2. We present a comprehensive dataset construction pipeline with high generalizability and transferability, enabling fellow researchers to effectively apply the methodology for creating datasets across various languages or in multilingual contexts.

3. We have utilized TransferTOD as our SFT dataset and trained the TransferTOD-7B model through full-parameter fine-tuning, achieving better slot filling and questioning capabilities compared to GPT-4. Additionally, with appropriate secondary fine-tuning techniques, our model demonstrates superior performance in out-of-domain testing compared to GPT-3.5-Turbo fine-tuned with an equivalent amount of data.





§ RELATED WORK



  
Task-oriented Dialogue Datasets The performance of intelligent dialogue systems is profoundly influenced by the quality of the dialogue datasets, making dataset construction an active research area. Initial generations of task-oriented dialogue datasets often focused on a single task or even a single scenario, ATIS <cit.>, DSTC2 <cit.>, WOZ2.0 <cit.>, etc. included. 
The emergence of these databases not only enhanced the conversational fluency of conversational agents but also made task completion through natural dialogues between machines and humans possible. Considering that user dialogues often involve domain transitions, datasets Multi-WoZ <cit.>, CrossWoZ <cit.> etc. encompassing more scenes and larger volumes of data were subsequently proposed. 

However, these dialogues are user-led discussions on relevant topics, requiring a user to pose questions or set tasks for the dialogue agent to respond accordingly. 



  
TOD System Enhancement Methodology
Enhancing the performance and data utilization of TOD systems and strengthening their ability to understand specific tasks expressed by users remain hot research topics. To complete tasks and improve accuracy, <cit.> proposed an end-to-end neural dialogue system based on reinforcement learning. TOD gradually started to realize across tasks <cit.>, domains <cit.>, and even languages <cit.>. TOD-BERT <cit.>, MinTL <cit.>, Soloist <cit.>, etc. has been successively proposed improving the success rate of tasks. However, as task complexity increases, these methods still rely heavily on large-scale datasets and lack competitiveness in handling noise robustness.




  
LLM-based TOD System
Existing research <cit.> has demonstrated LLMs' exceptional capabilities in natural language understanding, zero-shot reasoning, and command generation. With their advent and deep utilization, dialogue systems have entered the LLM-based era <cit.>. Utilizing LLMs, many dialogue tasks have achieved significant breakthroughs. On one hand, through internal dialogues with users, systems can be equipped with human-like perception and reasoning abilities, including intent classification, semantic parsing, dialogue state tracking, and reply generation. On the other hand, the integration of external information sources, such as specific databases, memory knowledge sources, the internet, etc., ensures the system provides the latest, rich, accurate, personalized, and necessary information to complete tasks.



§ TRANSFERTOD



 §.§ DataSet






TransferTOD aims to construct a cross-disciplinary task-oriented information collection multi-turn dialogue dataset, encompassing tasks such as goal-oriented questioning, dialogue state maintenance, information collection, and parsing. Existing task-oriented Wizard of Oz (WoZ) datasets are typically user-driven systems with relatively single domains. Departing from scenarios in the real world where questionnaire-style information collection may occur, we have curated dialogues spanning 30 different domains. We have enhanced the data in terms of robustness, diversity, and fluency, ensuring that the data closely mirrors real-world situations.



Figure <ref> illustrates the 4 steps of data collection and processing: 1. Original slot construction and dialogue generation; 2. Introduction of perturbed data; 3. GPT-enhanced dialogue diversity; 4. Manual refinement of dialogue content for fluency. Overall statistics of are shown in Table <ref>.



  §.§.§  Field Selection and Slot Collection


We crawl the most popular 30 life service offerings from local lifestyle applications (such as Meituan and Yelp) to construct the domain for our dialogue system. Specifically, we analyzed the submitted forms of each service, abstracting the information that the system would require users to provide as slots. 


After constructing the slots, we built a corpus containing all possible values for each slot. For string-type slots, we adopted a method of collecting publicly available information from the internet and generating rules. 

During the collection process, we kept the information for each slot separate, ensuring that no real personal information was involved. For number-type data, we described its range and distribution, generating it in real-time during the dialogue construction process. 


Human experts[Details of the human experts are shown in the appendix <ref>] manually created a set of high-quality dialogues as test data across 30 domains; three of these domains were selected for constructing an out-of-domain test set due to their minimal overlap in slots with the other domains. The remaining data is used as the in-domain test set. For the training dataset, the following steps will be undertaken to generate it on a large scale.



  §.§.§ Dialog Construction


Based on existing slot type descriptions and vocabularies, we have implemented the first version of a dialogue dataset using a script-generated approach. Specifically, we constructed a template library for each domain [Examples of our templates are shown in the appendix <ref>].

Each dialogue round consists of a user response, a system question, or a summary, forming the values before and after the dialogue state changes.

 For the number of slots k that could potentially be extracted in a single dialogue, we experimented with four scenarios: k=1, 2, 3, 4. Specifically,  when k=3, the system simultaneously asks the user for information on 3 slots, and the user needs to respond to these three corresponding aspects. 
 
 The statistical information of the original dialogue data is detailed in Table <ref>.
The dataset obtained after this step is TransferTOD-v1.
























  §.§.§ Noisy Data Construction

In real-world scenarios, users may provide information that does not conform to standards or common sense. Therefore, a comprehensive dialogue system should possess the capability to scrutinize the responses provided by users and, when necessary, seek clarification to obtain accurate information. To address this, a portion of the data is delineated to incorporate rounds of interaction specifically designed to handle incorrect responses from users. 

There are two types of data disturbances: 1. Non-responsive answers, where the content of the user's reply significantly deviates from the system's query. This dialogue alteration is achieved by replacing the user's response with an irrelevant answer; 2. Illogical responses, where the user's reply may contradict basic common sense. This data segment necessitates the introduction of non-factual content into the slot value lexicon to accommodate such instances.

During rounds with erroneous responses, the system will identify the user's mistake, repeat the original question, and maintain the dialogue state without updating it. We constructed 3013 noise dialogue data, with each dialogue containing at least one of the aforementioned errors, where the first type of error represented more than 90% of the cases. The dataset obtained after this step is TransferTOD-v2. Data examples are shown in Appendix <ref>.



  §.§.§  Dialogue Diversity and Fluency Polish

Dialogue data generated by static script schemes exhibit a shortfall in the diversity of questioning and answering modes. Each slot is confined to merely 5-6 variations of queries and responses, which fails to mirror the spectrum of linguistic preferences encountered in real-life scenarios. Consequently, we have leveraged the GPT-3.5 model to reformulate the texts of questions and answers, ensuring fidelity to the original intents while adjusting the temperature coefficient to 0.5 for an enriched array of textual content. The dataset obtained after this step is TransferTOD-v3.

Furthermore, we have refined the fluidity of dialogues that encompass inquiries about multiple slots within a singular exchange. Initially, dialogue data were essentially composed of simplistic amalgamations of disjointed questions or responses, not aligning with conventional spoken habits. Through the application of GPT-4 for sentence amalgamation and enhancement of coherence, coupled with rule-based scrutiny to pinpoint instances of fragmented sentences, we have engaged human annotators for the revision of overlooked or non-compliant sentences, thereby assuring dialogue smoothness. The dataset obtained after this step is TransferTOD-v4, which is our final dataset. Prompts are shown in Appendix <ref>.



 §.§ Models

Upon acquiring the TransferTOD dataset, we opted for the Baichuan2-7B-Base <cit.> as the foundational model for fine-tuning. During the model training process, we employed two methods: full-parameter fine-tuning <cit.> and LoRA (Low-Rank Adaptation) fine-tuning <cit.>.



  §.§.§ Supervised Fine-tuning

To equip the model with basic conversational abilities, we initially combined the training subset of the TransferTOD dataset with the general Chinese conversational dataset BELLE <cit.> in equal proportions to construct the SFT (Supervised Fine-Tuning) dataset. This dataset was utilized for full-parameter fine-tuning of the Baichuan2-7B-Base model to derive the TransferTOD-7B model.



  §.§.§ Secondary Fine-tuning

Following the development of our TransferTOD-7B model, we aimed for our model to achieve commendable performance in specific downstream tasks, necessitating that our model possesses superior generalization capabilities. In three external domain test sets, we adopted a limited-sample secondary fine-tuning approach to further enhance the accuracy of TransferTOD-7B in external domain test sets. Research <cit.> indicates that compared to full-parameter fine-tuning, LoRA fine-tuning achieves better generalization. Consequently, we employed LoRA fine-tuning for secondary fine-tuning. Experimental evidence demonstrates the effectiveness of our methodology in scenarios where data availability for downstream tasks is significantly constrained.



§ EXPERIMENT


In this section, we detail the experiments conducted. The primary experiments were carried out on the test set of TransferTOD. Additionally, we conducted ablation studies on the dataset construction phase, as well as supplementary experiments to further investigate the effects of secondary fine-tuning.



 §.§ Experimental Setup


For the in-domain test, we evaluated various methods known for their effectiveness in slot extraction. Traditional TOD systems divide the task into several modules <cit.>, each managed by a distinct model, forming a system pipeline. However, LLMs can reduce the reliance on task decomposition, thereby allowing us to directly evaluate the core competency of slot filling through information extraction.


For the out-of-domain test, a model's ability to adapt and generalize is paramount. Consequently, our initial evaluation centered on a selection of open-source LLMs with parameter counts comparable to our base model (7 billion), all of which demonstrated strong performance in Chinese benchmarks. To further enhance our analysis, we incorporated two powerful, near-source models from OpenAI.



  §.§.§ Baseline

For the in-domain test, we select 4 models as baseline: BertNLU <cit.>, SoftLexicon(LSTM) <cit.>, LEBERT+CRF <cit.> and W2NER <cit.>.

For the out-of-domain test, we select 6 Large Language Models as baseline: Baichuan2 <cit.>, ChatGLM3 <cit.>, Qwen <cit.>, Yi[<https://github.com/01-ai/Yi>], GPT-3.5-Turbo[<https://platform.openai.com/docs/models/gpt-3-5>], GPT-4 <cit.>. Please refer to the appendix <ref> for details.




  §.§.§ Implementation Detail




For evaluating the slot filling capability, we have annotated user utterances with BIO tags and trained 4 models for the in-domain test. A detailed system prompt was designed when inferencing with those LLMs in out-of-domain test. Please refer to the appendix <ref> for details.




  §.§.§ Evaluation Metrics

For the out-of-domain test, we assess the model's capabilities in two main aspects: slot filling ability and semantic accuracy during the phase of guiding user responses. To evaluate the slot filling ability, we employ F1 and Joint Accuracy, which are widely used in the TOD systems for slot extraction tasks. To evaluate the semantic accuracy of model-generated questions, we use a manual evaluation approach. Please refer to the appendix <ref> for details.
It is worth noting that we use the Dialog Act F1 as the evaluation metric in this context. While the Dialog Act F1 and SlotF1 metrics appear to be calculated in the same way, they differ slightly in essence. For a detailed explanation of these subtle differences, please refer to the Appendix <ref>.










 §.§ Results on TransferTOD

This section shows the results of our main experiment. 


  §.§.§ Results on In-Domain Test

Table <ref> presents the results of the in-domain test. Compared with traditional methodologies including W2NER, the State-Of-The-Art model in several ChineseNER tasks, our model significantly outperforms others on the in-domain test set in terms of the Dialogue Act F1 Score. This underscores the exceptional slot-filling accuracy of our model within domain-specific data.






  §.§.§ Results on Out-Of-Domain Test




Table <ref> showcases the results for the out-of-domain test set. The findings affirm that the average joint accuracy of TransferTOD-7B reached 75.09%, with a Slot F1 Score of 96.20%, surpassing other large-scale models, including the most advanced GPT-4, which only achieved a joint accuracy of 41.68%. In terms of query selection, GPT-4 leads the performance compared to other open-source models. TransferTOD's performance in this aspect scored 75, trailing just behind GPT-4. However, TransferTOD surpassed other models in terms of the fluency of queries. Besides, we have conducted a further experiment to compare our TransferTOD-7B to both open-source and close-source model with In-Context Learning 5-shot setting, reducing the probability of poor score caused by wrong format, the results are presented in Table <ref>, showing our TransferTOD's superior performance.

The experimental results validate that our TransferTOD model possesses robust generalization capabilities, achieving nearly 80% accuracy in specific downstream tasks. With appropriate secondary fine-tuning, the overall score can be further enhanced.

 


 §.§ Secondary Fine-Tuning Study



  §.§.§ Secondary Fine-Tuning

In this section, we primarily discuss our experiments on performing secondary fine-tuning on TransferTOD-7B. The objective was to simulate enhancing our model's slot filling and question-asking capabilities in external scenarios using a small subset of downstream scenario data. We fine-tuned GPT-3.5-Turbo[<https://platform.openai.com/docs/guides/fine-tuning>] as our baseline and conducted fine-tuning with 50, 100, and 200 pieces of data across three out-of-domain scenarios, respectively. The remaining data served as the test set for this experiment.

In the third scenario (Courier), we undertook multiple experiments employing various fine-tuning strategies, such as adding BELLE <cit.> dataset, incorporating in-domain data, and upsampling out-of-domain scenario data. This research aimed to identify methods that could further enhance the TransferTOD-7B model's slot filling capabilities.


  §.§.§ Result



Table <ref> shows the results of fine-tuning GPT3.5 and TransferTOD-7B in scenarios. The secondary fine-tuning can improve the model's out-of-domain capability. After fine-tuning, TransferTOD-7B still outperform GPT-3.5 (especially SlotF1) in most cases.



 §.§ Ablation Studies

Based on the TransferTOD-v1, v2, v3, and v4 mentioned in <ref>, we trained models TransferTOD-7B-v1 to v4 individually. To ascertain the efficacy and trustworthiness of our data construction methodologies, we rigorously assessed their performance in terms of robustness, diversity and fluency. The method we employed, which combines GPT-based assessment with expert review, is a widely adopted approach for evaluating the language fluency of models <cit.>. For details on the GPT assessment instructions and the expert review process, please refer to the appendix <ref>, <ref>.




  
Noise Injection To strengthen the model's resilience to noise, we augmented the standard dataset with a controlled amount of noisy data and trained the TransferTOD-7B-v2 model on it. As shown in Table <ref>, the improvement in joint accuracy substantiates the hypothesis that incorporating noisy data indeed strengthens the model's resistance to noise.






  
Language Augmentation To enhance the diversity of model interrogation techniques, we expanded our dataset by leveraging GPT, followed by a comprehensive assessment of the diversity in the questions generated by the newly developed models. Evaluators were provided with four assessment options: model A exhibits superior diversity, model B exhibits superior diversity, both models demonstrate comparable diversity, or neither model exhibits satisfactory diversity. The assessment results collected are shown in the Figure <ref>. Both the outcomes of expert review and GPT review affirm that v3 model surpasses v2 model in linguistic diversity.












  
Fluency Enhancement To enhance the fluency of the model's inquiries, we manually revised the dataset and employed a hierarchical scoring system to evaluate the models' query smoothness. The findings, as delineated in Table <ref>, underwent normalization to a 100-point scale, unequivocally demonstrate an improvement in the model's questioning fluency. The calculation of fluency score is given by equation <ref>.The experimental results demonstrate that the v4 model outperforms v3 model in both the high score rate on the GPT-based review and the expert review, as well as in terms of fluency score. Thus, our method effectively models query fluency.











§ CONCLUSION

Empirical evidence substantiates that our TransferTOD dataset possesses substantial noise resilience and superior linguistic performance. Utilizing this dataset for supervised fine-tuning, the resultant model, designated TransferTOD-7B, attains a joint accuracy of 75.09% in out-of-domain evaluations, accompanied by a Slot F1  of 96.20%. When it comes to question-asking ability, the accuracy of TransferTOD-7B is only slightly inferior to GPT-4, whereas its fluency in generating questions surpasses all other models we tested.

Furthermore, our findings suggest that appropriate secondary fine-tuning of the TransferTOD-7B model can further enhance its generalization capabilities. By employing a small portion of the out-of-domain test set for secondary fine-tuning, the resulting model surpasses the performance of GPT-3.5-Turbo, which was fine-tuned with an equivalent amount of data.

In summary, we have proposed a highly versatile data construction process that enhances the quality of task-oriented dialogue data for information gathering tasks. The models fine-tuned with this data exhibit strong generalization capabilities, performing well in out-of-domain scenarios.



§ LIMITATIONS

Our research presents a comprehensive set of experiments, yet it is not without limitations. One significant constraint stems from our dataset being primarily in Chinese, which precluded the testing of other major English-language open-source models due to their suboptimal performance on tasks in Chinese. Additionally, our assessment of question-asking accuracy employed manual evaluation methods, potentially introducing a degree of subjectivity despite our efforts to minimize such bias.






§ EXPERIMENTAL DETAILS




 §.§ Baselines


For the in-domain test, we select 4 models as baseline:


  
BertNLU <cit.>
is a BERT-based NLU model, initialized with Chinese pre-trained BERT and fine-tuned on tagged training data. For the input word embeddings, utilize MLP to generate BIO-tagged outputs.


  
SoftLexicon(LSTM) <cit.> is an effective method for incorporating the word lexicon into the character by categorizing the matched words, condensing the word sets and combining them with character representation.


  
LEBERT+CRF <cit.>
Lexicon Enhanced BERT for Chinese sequence labeling, utilizing a Lexicon adapter layer to integrate external lexicon knowledge into BERT layers.


  
W2NER <cit.>
is a modeling method of neighboring relations between entity words with Next-Neighboring-Word and Tail-Head-Word-* relations.

For the out-of-domain test, we select 6 Large Language Models as the baseline:


  
Baichuan2 <cit.> is an open-sourced large language model trained on 2.6 trillion tokens, achieving top performance in various Chinese and multilingual benchmarks. We utilized Baichuan2-7B-chat for our experiments.



  
ChatGLM3 <cit.> 
is Jointly developed by Zhipu AI and Tsinghua University, is the strongest in its class for datasets across multiple disciplines, supporting complex tasks like function calls and code interpretation. We utilized ChatGLM3-6B for our experiments.



  
Qwen <cit.>
Trained on 3 trillion tokens across multiple languages, Qwen models show competitive performance, excelling in tasks like chatting, text generation, and information extraction. We utilized Qwen-7B-chat for our experiments.



  
Yi[<https://github.com/01-ai/Yi>]
A powerful bilingual model, demonstrating significant potential in language cognition and reasoning, ranking highly on the SuperCLUE leaderboard and surpassing other large models in Chinese language proficiency. We utilized Yi-6B-chat for our experiments.



  
BlueLM <cit.> is
a large-scale model from vivo AI Global Research Institute, trained on a 2.6 trillion token corpus, showing leading results in Chinese benchmarks, indicating strong competitiveness. We utilized BlueLM-7B-chat for our experiments.



  
GPT-3.5-Turbo[<https://platform.openai.com/docs/models/gpt-3-5>] stands out as the most potent and cost-efficient model within the GPT-3.5 series. Tailored for conversations, it excels in comprehending and generating natural language. 



  
GPT-4 <cit.> is an advanced language model with enhanced understanding and generation capabilities. Trained on diverse internet text, it excels in various tasks, including text generation, translation, and problem-solving. We utilized GPT-4-1106-preview for our experiments.





 §.§ Implementation Details





  
Settings When training TransferTOD-7B, we use Baichuan-7B-base as base model, formatting the data to adapt to the Baichuan training format. Training cost about 8 hours on 8 A800-80GB GPUs and some hyper-parameters of our training are shown in Table <ref>, each version of our TransferTOD-7B adopted the same hyper-parameters when training.






  
In-Domain test When training in-domain models with dataset TransferTOD-v4, we tokenize the user utterance with Chinese pre-trained BERT <cit.> and annotate it with sequence labels using BIO tagging scheme.



  
Out-Of-Domain test 
For the first part, evaluating the model's capability of slot filling. When inferencing with the LLMs in out-of-domain test, we meticulously designed a system prompt, describing the task and desired output format in detail, to get the best result from each LLM, while some chat models may still perform fairly bad for the slots in their output don't match JSON format. The system prompt used has been translated to English and showed in Table <ref>.

For the second part, evaluating the semantic accuracy of model-generated questions, we use a manual evaluation approach. For detailed evaluation metrics, please refer to Appendix <ref>.







 §.§ Evaluation Metrics




  
Joint Accuracy measures the accuracy of dialogue states, considering a state correctly predicted only if all values of given slots are exactly matched.

Given the formula for Joint Accuracy is defined as:

    JA = N_cds/T_ds

where JA denotes Joint Accuracy,N_cds stands for the Number of dialog states correctly predicted,and T_ds represents the Total number of dialog states.









  
Slot F1 calculates the F1 score of (slot, value) pairs, deeming a tuple correctly predicted if the slot's value is exactly matched.

Given the formula for Slot F1 is defined as:










    SlotF1 = 1/N_Slots∑_i=1^N_SlotsF1 Score_i

where N_Slots represents the total number of (slot, value) pairs.



  
Dialogue Act F1 calculates the F1 score of (intent, slot, value) dialogue acts, where intent are always "inform", deeming a dialogue act correctly predicted if the slot and value extracted from user utterance is exactly matched.
Given the formula for Dialogue Act F1 is defined as:

    Dialogue Act F1 = ∑_i=1^N_DialogueActsF1 Score_i/N_DialogueActs

where N_DialogueActs represents the total number of (intent, slot, value) dialogue acts.





  
Ask Accuracy measures the model's ability to correctly select the corresponding number of slots from empty slots or to correctly point out errors in user answers and ask questions that correspond to the correct slots and will not cause misunderstandings. 


    Ask Accuracy = ∑_i=0^3 i × A_i/N × 3× 100

where A_i represents the number of the dialogues that got a score of accuracy i which ranks from 0 to 3 and N represents the total number of the dialogues.

For question accuracy scores, the scoring rules are as follows:


- 0 points represent that the model's questions are ambiguous, or it fails to correctly select fields from the empty slots for questioning, and the number of questioned fields does not match {extract_slot} (if the number of remaining empty fields is less than {extract_slot} and the number of questions asked does not equal the total of all remaining empty fields while meeting the previous condition, it should also be categorized here).

- 1 point represent that the model's questions might cause ambiguity, but it can correctly select fields from the empty slots for questioning, yet the number of questioned fields does not match {extract_slot} (if the number of remaining empty fields is less than {extract_slot}) and the number of questions asked does not equal the total of all remaining empty fields while meeting the first two conditions, it should also be categorized here).

- 2 points represent that the model's questions are precise, unambiguous, and it can correctly select fields from the empty slots for questioning, but the number of questioned fields does not match {extract_slot}) (if the number of remaining empty fields is less than {extract_slot} and the number of questions asked does not equal the total of all remaining empty fields while meeting the first two conditions, it should also be categorized here).

- 3 points represent that the model's questions are precise, unambiguous, and it can correctly select fields from the empty slots for questioning, and the number of questioned fields matches {extract_slot}) (if the number of remaining empty fields is less than {extract_slot}) and the number of questions asked equals the total of all remaining empty fields while meeting the first two conditions, it should also be categorized here). If all slots are filled and the model does not initiate a question or says "I have obtained all the information," the message content "" should also fall into this category.





  
Ask Fluency measures the fluency of the model’s questions and the degree to which they are consistent with natural language features. 

    Ask Fluency = ∑_i=0^3 i × F_i/N × 3× 100

where F_i represents the number of the dialogues that got a score of fluency i whick ranks from 0 to 3 and N represents the total number of the dialogues.

For the fluency score, experts rate the model's questions on fluency across a scale of 0 to 3 points. 


- 0 points represent that the representative's questioning style is rigid and awkward, completely deviating from the characteristics of natural language.

- 1 point represent that the representative's questioning style is somewhat rigid, yet the language is relatively natural, aligning with certain characteristics of natural language.

- 2 points represent that the representative's questioning style is relatively natural, and the language used is also quite consistent with the characteristics of natural language.

- 3 points represent that the representative's questioning style is very natural, and the language fully complies with the characteristics of natural language.





We referred to the work of LLMEval <cit.> to design the evaluation criteria for accuracy and fluency.



§ TEMPLATES FOR SCRIPT-GENERATED DATA


Table <ref> shows an example of a question and answer template when Domain set to ‘Hotel' and Slot set to ‘Hotel Type'.





§ PROMPTS




 §.§ Prompt GPT-3.5 to Polish the Data

The prompt showed in Table <ref> and Table <ref> are used when using GPT-3.5 to polish the text, rewriting questions and answers respectively, in our dataset.







 §.§ Prompt GPT-4 to Evaluate the Results











































































































































The prompt in Table <ref> is used when using GPT-4 to conduct comparative evaluation of diversity in ablation experiments, while the prompt in Table <ref> is used when scoring the fluency of model questions in ablation experiments with GPT-4.


















































§ DATA EXAMPLES


Examples of our supervised-finetuning data are showed in Figure <ref> and Figure <ref>, also we provide examples of data with noise in Figure <ref> and Figure <ref> as well as raw TransferTOD data in Figure <ref> and Figure <ref>.






§ HUMAN EXPERTS



 §.§ Experts in Constructing Datasets




During the dataset construction phase, we relied on 5 students from our institute to participate in this work as human experts. These students possessed basic computer knowledge and coding skills, which enabled them to perform the task effectively.


Their primary responsibility is to generate dialogue data for test sets. We assign tasks based on different scenarios, ensuring they are familiar with the entire dataset construction process and principles. They work professionally, providing human support for the dataset creation and ensuring smooth project execution. Additionally, we fairly compensate their efforts to show respect and recognition for their contributions.


 Another task for human experts involves refining non-fluent content. Given the potential for incoherence and unnaturalness in rule-based generation in <ref>, characterized by the lack of appropriate connective words and inconsistent tone, we prioritize addressing this issue.Thus, human experts are employed to revise dialogue content, such as transforming "What's your name? Please tell me your phone number." into a more coherent and natural structure like "Please provide your name and phone number."



Compared to rule-based mass generation, expert-crafted data exhibits significant advantages. The work of domain experts enhances the linguistic fluency, naturalness, and brevity of the generated dialogues. This high-quality, manually constructed data boasts greater authenticity and representativeness, more effectively emulating real-world conversation scenarios. Consequently, it serves as a more reliable foundation for subsequent fine-tuning tasks.



 §.§ Experts in Ablation Experiment



During the ablation experiment phase, we invited 12 students from our institution to conduct comparative evaluations of the results. Each student was assigned to complete the full assessment tasks for one or more large models. This entailed each student conducting a comprehensive evaluation of the designated model to ensure a thorough understanding of its performance.


Specifically, we selected 200 data points from the inference results of TransferToD-7B-v2 and TransferToD-7B-v3, and conducted 200 random samples. 5 data points were sampled each time, resulting in a total of 40 evaluations for each model's inference results. This random sampling method contributed to ensuring the objectivity and reliability of the assessment, minimizing potential biases.


Subsequently, the evaluators rated the sampled data based on the questioning style, diversity, and fluency. They provided an overall score for each set of data by considering factors such as the model's questioning approach, sentence completeness, clarity of questioning, diversity, and fluency. These scores provided quantitative data on the model's performance in various aspects, facilitating a more comprehensive assessment and comparison of the models' strengths and weaknesses.




















 Paper Title: Unsupervised Pairwise Causal Discovery on Heterogeneous Data using Mutual Information Measures.


          








A]Alexandre Trilla,
B]Nenad Mijatovic

[A]Alstom, Santa Perpètua de la Mogoda, Barcelona, 08130, Spain
	
[B]Alstom, Saint Ouen, Paris, 93482, France

§ ABSTRACT 

A fundamental task in science is to determine the underlying
causal relations because it is the knowledge of this functional
structure what leads to the correct interpretation of an effect 
given the apparent associations in the observed data.
In this sense, 
Causal Discovery is a technique that tackles this
challenge by analyzing the statistical properties of the 
constituent variables.

In this work, we target the generalizability of the 
discovery method by following a reductionist approach that 
only involves two variables, i.e., the pairwise or bi-variate
setting.

We question the current (possibly misleading)
baseline results on the basis that they were obtained through 
supervised learning, which is arguably contrary
to this genuinely exploratory endeavor. In consequence, we
approach this problem in an unsupervised way, using robust Mutual
Information measures, and observing the impact of the different variable
types, which is oftentimes ignored in the design of solutions.
Thus, we provide a novel set of standard unbiased results
that can serve as a reference to guide future discovery tasks 
in completely unknown environments.



Causal Discovery, Functional Causal Model, Unconditional Independence Test, Mutual Information, Unsupervised Learning, Heterogeneous Data








§ INTRODUCTION

Causal Discovery methods are able to identify the causal structure
from the joint distribution of the data by introducing assumptions 
that restrict the model of their generating process <cit.>.
In a multivariate setting, the traditional constraint-based and 
score-based methods exploit conditional independence 
relationships in the data <cit.>.
These approaches have found a great deal of success in 
challenging environments such as biology and the Earth
sciences <cit.>. Nevertheless,
they do not necessarily provide complete causal information 
because they output Markov Equivalence Classes, i.e., a set of
causal structures that satisfy the same conditional independence
statements. Moreover, they cannot handle isolated cause-effect
settings that lack the diversity of other variables for the 
conditional tests. Therefore, the single reduced cause-effect association,
also known as the pairwise or bivariate scenario, constitutes an 
essential building block of more complex causal structures. 

To distinguish cause from effect in this pairwise setting, 
one needs to find a way to capture the asymmetry between the
two variables <cit.>. In this sense,
computational methods based on properly defined Functional Causal 
Models (FCMs) are able to distinguish different structures in the 
same equivalence class.
A FCM represents the effect variable Y as a function of its direct
causes X and some noise term ε, i.e., 
Y = f(X, ε), where ε is
independent of X. Thanks to the restricted functional classes,
the causal direction between X and Y is generally identifiable because
the independence between the noise and the cause holds
only for the true causal direction and is violated for the wrong
direction <cit.> (unless the variables are jointly Gaussian,
which render the orientation unidentifiable). Additionally, the
identifiability of pairwise FCM generalizes to the identifiability 
of multivariate FCM <cit.>.

In this work, we argue that some reference benchmark results for pairwise 
causal discovery could be misleading because they tackle this genuinely 
exploratory objective as a classical supervised classification problem, 
inducing them to possibly learn specific contextual details beyond causality 
that tend to overestimate their actual performance <cit.>.
Moreover, there exists a tight relationship between the nature of data 
and the specific tools from statistics that could also accrue the estimation 
of skewed results, which was not originally addressed. Consequently, we 
propose a reevaluation of the benchmark driven by the heterogeneity 
of the data and also by following an unsupervised 
learning approach. Our method leverages linear functional causal models and 
combines different unconditional independence tests based on robust
Mutual Information measures.
The paper is organized as follows: Section <ref>
introduces the fundamental concepts for pairwise causal
discovery, Section <ref> describes the enhanced method
proposed in this work, Section <ref> shows the novel
benchmarked results, Section <ref>
discusses the approach from a causal perspective, along with its
limitations, and
Section <ref> concludes the article.






§ BACKGROUND
 

This section reviews the fundamental concepts that support the
discovery of causality in the bivariate or pairwise setting.



 §.§ Identifiability of a Linear Model with Non-Gaussian Noise

The fundamental principle that enables the identification of cause
and effect lies in the asymmetry of the association between them: 
the causal direction, i.e., from cause to effect, is functionally 
simpler than the anticausal one, i.e., from effect to cause, which
requires a more complex application <cit.>.
Thus, a given model with limited capacity will be able to 
faithfully capture the expressiveness 
of only one of these two causal orientations.

Linear additive noise models assume that the effect Y is a linear 
function of the cause X plus a noise term ε that
is independent of the cause, formally defined as: Y = bX + ε.
The linear FCM is learnable if at most one of X and ε
is Gaussian <cit.>. Pairwise discovery approaches such as this one are regarded to be very flexible for identifying causal models in the general case because they don't require a third variable for conditioning.
In specific scenarios such as multivariate time series, they can even
improve the performance of traditional constraint-based
approaches such as the Peter-Clark algorithm and others that
utilize Granger causality <cit.>.

The linear model with non-Gaussian noise can be estimated from 
(unconfounded) observational data by exploiting the inherent 
asymmetry between cause and effect through a regression as 
follows: for both directions (i.e., causal and anticausal), the 
FCM is fit; then, the amount of association between the estimated 
noise term (i.e., the regression residual) and the hypothetical 
cause is computed; finally, the direction assignment which gives an 
independent noise term is considered plausible <cit.>.

Obviously, the cornerstone of this regression-based method is the
measure of dependence between the regression residual and its
associated potential cause. The next section 
describes a solution to estimate the strength of this association.





 §.§ Mutual Information as a Measure of Association

The amount of association between the regression residual RR and
the hypothetical cause HC signals the causal direction. However,
these variables are uncorrelated by construction of the regression. To 
properly quantify this dependence, the Mutual Information (MI) measure 
provides a reliable indicator, which is defined as


    MI(RR,HC) = ∑_RR,HC P_RR,HClog ( P_RR,HC/P_RR  P_HC )  ,


where P_RR,HC represents the (discrete) joint probability between
RR and HC, and P_RR and P_HC represent their marginal
distributions, respectively.

For illustrative purposes, the following variable dependencies 
are defined. First, a random Uniform noise sample U is
independently assigned to Z, X_ind and Y_ind. Then,
the following structural variable associations are created: 
Y_dep← X_ind + Z, X_conf← X_ind + Z,
and Y_conf← Y_ind + Z, which yield the subsequent relationships:


  Causal Between X_ind and Y_dep: X → Y

  Anticausal Between Y_dep and X_ind: Y ← X

  Independent Between X_ind and Y_ind: X  Y

  Confounded Between X_conf and Y_conf: X ↔ Y


Figure <ref> shows the distribution of MI between
the regression residual and the hypothetical cause considering
the preceding types of structural variable associations for a random
sample of 1000 instances. The different density plots
show how the causal orientation can be reliably recovered using
MI when the model assumptions are respected. Note that
the overlap between Causal and Independent structures 
will be disambiguated when both experiments are conducted regarding 
the different causal direction hypotheses.
Finally, MI is also able to detect latent confounding, which is 
oftentimes assumed (and required) not to occur for 
discovery <cit.>, i.e., the causal 
sufficiency principle <cit.>. This feature 
adds value to the robustness of the MI indicator.



While the distribution of MI is a useful tool to visualize the
different causal structures, there is still a missing criterion to 
make a classification decision in order to construct the underlying
causal structure from the data. The next section addresses this point.




 §.§ Pearson's χ^2 as an Unconditional Independence Test


In a real-word scenario, one cannot expect to get a null measure
of association because of sampling noise, the sample size, etc.
In this case, the use of an (unconditional) independence
test is the proper way to handle this situation. In general,
this is known as the REgression with Subsequent Independence Test
(RESIT) procedure <cit.>. 
Once the independence test is in place, a confidence interval needs 
to be observed to obtain actionable results. To this end, Pearson's 
χ^2 test is here introduced as a fundamental 
method to inform the decision process <cit.>.

The Pearson's χ^2 test is a statistical hypothesis test used 
in the analysis of contingency tables for categorical variables.
It is used here to determine whether there is a statistically 
significant difference between the expected frequencies E
for independence (i.e., the marginal probabilities)
and the observed frequencies O in one or more categories i
of the contingency table (RR,HC), formally defined as

    χ^2(RR,HC) = ∑_i ∈ (RR,HC)(O_i - E_i)^2/E_i .


The test is valid as long as its statistic is χ^2 distributed 
under the null hypothesis.
Note that χ^2(RR,HC) is in its essence related to the MI measure of 
association as is used for the discovery of causal structure.

While the χ^2 test is most suitable for discrete variables <cit.>,
a general
causal discovery method should be able to seamlessly deal with 
continuous variables. Typically, a uniform grid is applied to discretize
a real-valued space prior to conducting the test.
Following the former illustrative example, Table <ref>
shows the p-values for the different causal structures.
Once a threshold is introduced through the confidence interval,
e.g., p<0.05 for a given number of degrees of freedom related to
the resolution of the grid,
passing the test becomes the rule for making further decisions.
Note that in all cases, there is a unique combination of results
for the two hypothesized causal directions that identifies the true
underlying causal structure.
Also note that the threshold may be adjusted to require further decision strength because the closer the variables get to a Gaussian distribution, the harder it is to distinguish the direction of causation <cit.>.



While these results
may be sufficient as an indication, they may also have some 
limitations for discovering causal associations between numerical variables.
The next section addresses this shortcoming and describes
an enhanced solution for this type of variables.







 §.§ Total Information Coefficient as a Robust Independence Test



The Total Information Coefficient (TIC) is a robust independence
test based on MI for real-valued variables 
that features the two important heuristic properties of generality and 
equitability <cit.>, which are defined as follows:



  Generality With sufficient sample size the statistic 
should capture a wide range of interesting associations, not limited 
to specific function types. 

  Equitability The statistic should give similar scores to 
equally noisy relationships of different types.


For a pair of (RR,RC) variables, the TIC algorithm applies a search
procedure to partition their joint probability function and find the
grid with the highest induced MI.
TIC operates by summing over optimal grids G on the joint 
density distribution of the two variables, such that it exhibits 
a stronger power 
against independence <cit.>, formally expressed as

    TIC(RR,HC) = ∑_G MI((RR,HC)|_G)/log∥ G ∥ ,

where ∥ G ∥ denotes the minimum of the number 
of rows and columns of G. This can be viewed as a regularized 
version of MI that penalizes complicated grids.
Figure <ref> shows qualitatively the impact of the
optimum grids G.




The performance of TIC has been positively rated with
other successful approaches for causal discovery dealing with
continuous variables, such as the
Hilbert-Schmidt independence criterion <cit.>.
Finally, an additional motivation for
considering TIC in this work is to increase the odds of success
in case of model hypothesis violations, i.e., the functional form
of the relationship Y = f(X), which is crucial for causal
identification.






§ METHOD


In the pursuit of a general approach to causal discovery
combining several methods <cit.>,
one common point of success is found in the use
of information-theoretic measures, such as Mutual Information,
to quantify the amount of regularity in the 
data <cit.>. 

However, the nature of the variables, i.e., discrete or continuous, 
can be problematic in an heterogeneous context if the statistical 
methods are confused <cit.>.


Moreover, supervised learning approaches are excluded from the
general objective because these methods do not yet work as standalone 
techniques for causal learning <cit.>. Therefore,
the desideratum in causal discovery is to have methods that work on a 
broad range of problems under different conditions with relaxed
assumptions <cit.>, ideally showing a 
certain degree of robustness regarding violations of the model 
hypothesis <cit.>.


In this work we propose a general method to discover the causal 
structure in heterogeneous data based on a bivariate linear FCM by
implementing a flexible RESIT procedure, see 
Section <ref>, where the
unconditional independence test, e.g., χ^2 or TIC,
also see Section <ref>,
is driven by the nature of the variables involved. The proposed 
strategy is described with comments in Algorithm <ref>.
Note that the orientation decision rules follow from the evidence 
given by the illustrative setting in Section <ref>, 
which is regarded as self-evident.
Such explicitly-stated logic rules are meant to increase the 
interpretability and explainability of the proposed method.
Also note that discrete variables are assimilated to real-valued variables for computing the regression residuals on which the RESIT method is based. Finally,
since our contribution targets the problem of applying one single discovery technique on data with different types of variables, and we focus our effort at the integration level rather than at the fundamental level, we refer the interested reader to the references for the specific methods being integrated for further details about their comparison to other methods from the literature.
We rely on the advice from Peters, Reshef, and 
colleagues <cit.>, to select the best methods that
we integrate in RESFIT in order to cover a broad range of techniques, 
and we provide an analysis on how these separate tools perform when
they are put together.















§ RESULTS


This section describes the reference benchmark dataset that was used
to evaluate the introduction of the flexible independence test selection, 
the results that were obtained, and the tools for the implementation 
of the research.



 §.§ Reference Benchmark Dataset



The “ChaLearn cause-effect pair (SUP2)” is taken for reference as
the recommended benchmark dataset to evaluate the performance
of pairwise causal discovery <cit.>.
It comprises pairs of artificially generated dependent 
variables with different types 
(numerical, categorical and binary) for the full 
causal discovery task (orientation, independence and confounding).
The dataset features a balanced number of unique values across all 
classes and includes around 6000 pairs. In terms of data type balance,
the majority is comprised of numerical and mixed variable pairs. Finally, 
the average median length of an instance is around 2000 values.



 §.§ Performance Scores


The accuracy classification score, i.e., the total rate of correct
predictions statistically given by the overall amount of
true positives and true negatives,
is used in this research following the previous benchmark
evaluation approaches <cit.>, yielding a baseline around 
65±2% for the supervised learning setting <cit.>.
To further assess the stability of our scores in the unsupervised
scenario, statistical
bootstrapping is introduced in line with this former work.
Moreover, Gaussianity is asserted with the Lilliefors normality
test <cit.>,  and the main descriptive
statistics are extracted despite the low amount of available performance
samples, well under 30, which are commonly
required to obtain reliable estimations <cit.>.
Finally, in terms of statistically significant comparisons, the t-test has been
used <cit.>.
Table <ref> shows the results of the experiments.
Note that they are presented separately by data type and independence
test, whereas the proposed RESFIT algorithm automatically integrates
them. This is done for illustrative purposes and for enriching 
the ensuing discussion.




The first straightforward conclusion that the results show is that,
in all cases, there is a statistically significant difference between
the accuracy averages for the two unconditional independence tests
within each data type stratum. Therefore, introducing a selection 
action on the test function is expected to impact 
on the discovery of causal associations.
Additionally, TIC yields a slightly smaller variance in all
the scores. Finally, when numerical data types are present,
the standard deviation in accuracy drops an order of magnitude.

The code for this research is available 
here[Code repository: https://github.com/atrilla/ccia24].
The next section discusses the global results with respect to the
perspective of a causal effect, and addresses the limitations
of the proposed approach.





§ DISCUSSION


This section delves into the finer details of the results that
were obtained, and sheds light on the actual value introduced by
the flexibility in selecting the unconditional independence test
for causal discovery.



 §.§ Average Causal Effect of the Flexible Test Selection


To properly quantify the impact of introducing the flexibility
on the unconditional independence test, this section 
treats this selection action as a “treatment” variable X 
and studies its impact on the “outcome” discovery accuracy
score Y.
The Average Causal Effect (ACE) of X → Y is
expressed using the following counterfactual notation 


    ACE =    𝔼 [Y_1 - Y_0] = 𝔼 [Y_1] - 𝔼 [Y_0] = 𝔼 [Y(X=1)] - 𝔼 [Y(X=0)] 
    
    	=    𝔼 [Y|X=1] - 𝔼 [Y|X=0]  ,


where Y_x refers to the value the Y accuracy result would have
if X was set to x, i.e., Y(X=x). Here, X
could take the values x=1 to indicate
flexibility of independence test, and x=0 to indicate no
preference (i.e., the null random choice). The conditionals that
eventually follow are the values that are actually observed in the results.
Note that all these expected quantities can be exactly calculated 
because the different experiments can be conducted on the same data
(also dispelling any doubts about latent confounding).
This setting is thus not subject to the fundamental problem of 
causal inference where only one of the potential outcomes
can be observed <cit.>.

Equation (<ref>) is shown to be an unbiased estimator for the
ACE <cit.>. For computing the expectations in the
causal effect comparison, a weighted mixture of Gaussian random
variables is required.
The quantity of the reference potential outcome 
𝔼 [Y_0], where no smart selection occurs, is given by
the arithmetic mean value of the total results, and it yields
an accuracy score of 0.3531±0.0357. Alternatively, the potential outcome
𝔼 [Y_1], where the smart flexible test selection is introduced, 
is given by the average of the best performance results among 
the different variable
types, i.e., a sum weighted by their balance in the dataset,
and it yields a value of 0.3866±0.0690. Therefore, the ACE
is of 3.36%, and this estimated difference is statistically significant.
However, the absolute results are far from the 65% baseline, which
suggests that the original supervised approach may have leaked
statistical patterns beyond causation. Although the unsupervised
learning approach does not seem to offer any pragmatic quantifiable
advantage, we believe that its adoption is a must for tackling the
inherently hard exploratory objective of causal discovery and avoid
any qualms with regards to learned biases.




 §.§ Limitations

The realization that TIC, which was conceived for
numerical data, performed better on discrete data is surprising.
Also, χ^2, which is defined for discrete data, performed 
better on numerical data. This counterintuitive behavior 
where theory and practice disagree may
suggest that the independence tests that were used are subject 
to some inherent limitations such as the sample size. 
We regard the shortage of samples to be the main limitation
of our approach on real-world data applications such as the
Tuebingen dataset <cit.>.

Finally, while the approach that we propose is able to welcome
any testing technique, we selected RESIT flexibly paired with TIC and 
χ^2 following the advice of Peters, Reshef, and 
colleagues <cit.>. To the best
of our knowledge, this limited choice should be sufficient to 
generally cover the spectrum of approaches.





§ CONCLUSION



Environments with heterogeneous data pose challenging questions to
the discovery of causal structure, and leveraging one single method may lead to erroneous results in general.
In this work, we propose an
unsupervised pairwise approach using linear functional causal 
models and different unconditional independence tests based on 
Mutual Information measures, the selection of which 
is driven by the nature of data and the empirical results 
from a reference benchmark.
The introduction of this independence test selection flexibility is
estimated to have a positive and statistically significant average
causal effect over 3% in accuracy.
These scores may establish the first 
standard baseline for this kind of flexible focus, but more
research is needed because the limitations of the
data size and the statistical techniques can result in 
counterintuitive results.







vancouver





 Paper Title: AgentGen: Enhancing Planning Abilities for Large Language Model based Agent via Environment and Task Generation.





[NO \title GIVEN]
    [NO \author GIVEN]
    August 5, 2024
======================



§ ABSTRACT  




Large Language Model (LLM) based agents have garnered significant attention and are becoming increasingly popular. 


Furthermore, planning ability is a crucial component of an LLM-based agent, involving interaction with the environment and executing actions to complete a planning task, which generally entails achieving a desired goal from an initial state.

This paper investigates enhancing the planning abilities of LLM-based agents through instruction tuning, referred to as agent training.
Recent studies on agent training have demonstrated that utilizing expert-level trajectory data (sequences of action-observation pairs) for instruction-tuning LLMs effectively enhances their planning capabilities.
However, existing work primarily focuses on synthesizing trajectories from manually designed planning tasks and environments. The labor-intensive nature of creating these environments and tasks impedes the generation of sufficiently varied and extensive trajectories for agent training.
To address this limitation, this paper explores the automated synthesis of diverse environments and a gradual range of planning tasks, from easy to difficult. 
In response, we introduce a framework, , that leverages LLMs first to generate environments and subsequently generate planning tasks conditioned on these environments.


Specifically, to improve environmental diversity, we propose using an inspiration corpus composed of various domain-specific text segments as the context for synthesizing environments. 
Moreover, to increase the difficulty diversity of generated planning tasks, we propose a bidirectional evolution method, , that evolves planning tasks from easier and harder directions to synthesize a task set with a smoother difficulty curve, thereby enhancing the learning process of LLMs more effectively.
These methods collectively contribute to the generation of diverse trajectory data for instruction-tuning.


Based on , we greatly expanded the number of environments and planning tasks available for agent training.


The evaluation results derived from AgentBoard show that greatly improves LLMs' planning ability, e.g., the instruction-tuned Llama-3 8B surpasses GPT-3.5 in overall performance. Moreover, in certain tasks, it even outperforms GPT-4.




§ INTRODUCTION




Recently, owing to advancements in Large Language Models (LLMs) <cit.>, the LLM-based Agents have garnered widespread attention from the artificial intelligence community. 
Generally, an LLM-based agent refers to utilizing LLMs to perceive the environment, make decisions, and execute actions to substitute or help people accomplish some specific tasks <cit.>.
Furthermore, planning is often regarded as one of the most important applications of LLM-based agents, such as robotic planning <cit.>, travel planning <cit.>, etc. 
In this study, planning is conceptualized as the systematic process of identifying a sequence of executable actions within a given environment to complete a planning task, defined as the transition from an initial state to achieve specified goal conditions, considering constraints and available resources <cit.>.




Improving planning capabilities through instruction-tuning LLMs is a significant research problem, referred to as agent training.

As shown in Figure <ref>, similar to imitation learning <cit.>, a typical agent training process can be divided into three stages:
(i) Preparing environments and planning tasks.
(ii) Synthesizing expert-level trajectories (sequences of action-observation pairs) on these planning tasks. For example, utilizing state-of-the-art LLMs (e.g., GPT-4 <cit.>) as the agent and filtering trajectory based on reward score <cit.>.
(iii) Instruction-tuning LLMs with the synthesized trajectory data.
Recently, the effectiveness of enhancing the planning capabilities of LLMs through agent training has been demonstrated by many studies <cit.>.

Despite their success, one key limitation of these works is that they primarily rely on manually designed environments and planning tasks. 
The labor-intensive nature of creating environments and planning tasks hinders the generation of diverse and extensive trajectory data.

More explicitly, designing diverse environments requires defining a range of rich and practical scenarios, and implementing these environments typically involves the participation of human experts with programming skills.
Additionally, formulating tasks often demands creating a task set with a gradual difficulty progression.
Due to this constraint, existing agent training studies typically use only a few environments for data synthesis.





To address the aforementioned deficiencies, this paper introduces an automatic framework  that utilizes LLMs to construct diverse environments and planning tasks for agent training, expanding the available environments from a few to hundreds.




More specifically, is structured around two stages:

(1) Environment Generation: 
Achieving sufficient environmental diversity is essential for creating diverse planning tasks, which involves covering a broad range of scenarios and domains. 
To ensure this, we use an inspiration corpus composed of diverse text segments as context for generating environment specifications with LLMs, where actions, restrictions, and other details are defined using natural language.

For example, in Figure <ref>, we randomly selected a text segment from the inspiration corpus: “How to boost your diet with peanut butter powder?” This prompted the generation of a related environment specification: “You are a nutritionist tasked with creating a new healthy recipe book that incorporates peanut butter powder as a key ingredient”.
Subsequently, we prompt the LLM to produce the corresponding code based on this specification, which may be composed of Python, Planning Domain Definition Language (PDDL) <cit.>, or other domain-specific languages.
Furthermore, we constructed an environment library to serve as in-context examples and iteratively expanded it by incorporating high-quality newly generated environments.


(2) Task Generation: Conditioned on the generated environment, we aim to create multiple planning tasks. 
In this stage, it is crucial to have a gradual set of tasks ranging from easy to difficult, i.e., difficulty diversity.
To achieve greater difficulty diversity, we propose a bidirectional evolution method, , where the LLM first generates random planning tasks and then evolves these tasks by applying constraints towards both simplification and increased difficulty. 
The created task set with has a smooth difficulty curve that facilitates LLMs' smoother acquisition of planning skills.




















To verify the effectiveness of our method, we synthesized environments and planning tasks based on PDDL <cit.> and constructed a dataset comprising 592 environments, each with 20 tasks. 
We then used a domain-independent planner to obtain 7,246 high-quality trajectories.

Subsequently, we used this trajectory data for instruction-tuning a series of LLMs and demonstrated the trained model on AgentBoard <cit.>. 

Since our instruction-tuning dataset is composed of trajectory synthesized from PDDL-based planning tasks, we refer to evaluation tasks implemented in PDDL as in-domain tasks and tasks implemented in other programming languages as out-of-domain tasks. 

Importantly, this evaluation was conducted in a zero-shot manner without utilizing any trajectory data from these tasks. 

Experimental results show that achieved over five times improvement compared to the raw LLama3-8B on in-domain tasks (11.67 vs. 1.67), and overall performance surpassed GPT-3.5. Additionally, it outperformed GPT-4 in specific tasks.
In out-of-domain tasks, also demonstrated similar experimental outcomes. Specifically, it significantly improved success rates compared to the raw LLama3, achieving 29.1 and 4.47 improvement on Alfworld and BabyAI, respectively. On Alfworld, even surpassed the performance of GPT-3.5 (29.1 vs. 17.2).
In summary, the proposed environment and planning task generation method can help improve planning ability. Moreover, not only can in-domain tasks benefit from this, but out-of-domain tasks also improve, which confirms both the effectiveness and generalization.

Our contributions can be summarized as follows:

    
  * We introduce , which, as far as we know, is the first framework for automatically generating diverse planning tasks and environments targeted for LLM-based agent training.
    
  * We propose utilizing an inspiration corpus as the context for generating environments with LLMs, resulting in 592 diverse environments that encompass a broad range of scenarios.
    
  * We propose a bidirectional evolution method that evolves seed planning tasks in both simpler and more complex directions, thereby constructing a task set with a smoother difficulty curve.
    
  * We constructed an agent instruction-tuning dataset with 7246 high-quality trajectories through . LLMs instruction-tuned with this dataset achieved massive improvement in both in-domain and out-of-domain planning tasks, which validated the effectiveness and generalization of .







§ PRELIMINARY























 §.§ Planning Problems


We consider goal-directed deterministic planning problems <cit.>, which are formally defined as a tuple 𝒫 = (𝕋, 𝔼), where 𝔼 denotes the environment in which the agent interacts and 𝕋 denotes the task that the agent needs to complete.
Specifically, an environment 𝔼 typically models a world, encompassing the definitions of the action space 𝒜 and state space 𝒮, as well as the transition function 𝒯: 𝒮×𝒜→𝒮.
Task 𝕋 is further defined by the tuple 𝕋 = (𝒢, ℐ), where 𝒢 refers to the goal conditions and ℐ refers to initial states of the agent.
The initial states ℐ are a subset of the state space 𝒮_i that specifies the starting conditions of the agent.
The goal 𝒢 is a subset of the state space 𝒮_g that specifies the desired outcomes or conditions. Specifically, 𝒢 can be expressed as 𝒢 = { s ∈𝒮_g|ϕ(s) = true}. Here, ϕ(s) is a boolean-valued function representing conditions or propositions that must be satisfied for the state s to be considered part of the goal set. 





 §.§ Planning Problem Implementation





















A planning problem can be implemented with programming languages such as Python or domain-specific languages such as Planning Domain Definition Language (PDDL) <cit.>.
For example, in a PDDL-based planning problem, the domain PDDL file can be regarded as the environment 𝔼, defining states (predicates) and actions and specifying the transition function using preconditions and effects of each action. The problem PDDL file, on the other hand, can be seen as the task 𝕋. Both initial states and goal conditions are typically defined as combinations of predicates.
Another widely used programming language for constructing planning problems is Python.

For example, in OpenAI gym[<https://www.gymlibrary.dev/index.html>], a planning problem will be implemented as a Python class, where the transition function is implemented as a method of the class, usually named the "step" or "update" function. Meanwhile, the goal 𝒢 is typically represented as a reward function that indicates the objective of the task, and the initial states ℐ are defined in a method named "reset."







 §.§ Large Language Model based Agent

An LLM-based agent leverages a pre-trained language model to operate within the defined environment 𝔼 and complete the given task 𝕋.

Given an environment 𝔼, the LLM-based agent perceives its state 𝒮 and takes actions 𝒜 based on its understanding and processing of the input. 
The transition function 𝒯: 𝒮×𝒜→𝒮 remains consistent, where the LLM-based agent determines the next state by generating appropriate actions through natural language processing.
The goal G guides the LLM-based agent in selecting actions that maximize the reward. The agent utilizes the language model to interpret the task requirements and generate actions that align with achieving the specified goal.
In essence, the LLM-based agent forms a policy π: 𝒮→𝒜 using the LLM, where π(s) is the action taken in state s based on the LLM's understanding and processing of the task.



















































§ METHODOLOGY





  
Problem Definition


 
The process of generating planning tasks can be formalized as a function f: I → (𝕋, 𝔼), where I is the input space (e.g., instructions or prompts) and tuple (𝕋, 𝔼) is the space of all possible planning tasks and environments.
Based on the definition in Section <ref>, we can express this as f(i) = (𝕋_i, 𝔼_i),    i ∈ I, where 𝕋_i is the generated planning task and 𝔼_i is the generated environment for a given input i.
Our two-stage approach can be further decomposed as follows: 

i) Environment Generation (<ref>):
In the first stage, we generate the environment 𝔼_i based on the input instruction i. This can be represented as 𝔼_i = g_𝔼(i),  where g_𝔼 is the environment generation function that takes the instruction i as input and produces the environment 𝔼_i.
ii) Task Generation (<ref>):
In the second stage, we generate the task 𝕋_i, conditioned on the environment 𝔼_i generated in the first stage. This can be expressed as:
𝕋_i = g_𝕋(i, 𝔼_i),
where g_𝕋 is the task generation function that takes both the original instruction i and the generated environment 𝔼_i as inputs to produce the task 𝕋_i.
We will detail the implementation of these two stages in the following section.



 §.§ Environment Generation







  
Overview


As is shown in Figure <ref>, we propose a sophisticated framework for environment generation structured around three main components: (1) an environment specification generation module where an LLM first generates a specification of the environment, typically including a general overview of the environment, descriptions of the state space and action space, and definitions of the transition functions; (2) an environment implementation module that generates corresponding code based on the environment specification; and (3) an environment library that stores previously generated high-quality environments, serving as a comprehensive environment dataset and providing in-context examples for generating new environments. Each component will be elaborated on in the following paragraph.



  
Environment Specification
We initially prompt the LLM to generate an environment specification, which typically includes an overall depiction of the environment, specific actions and their corresponding preconditions and effects, and certain restrictions within the environment.
The environment specification will serve as the basis for generating specific environment codes. This two-stage approach, similar to the Chain-of-Thought <cit.>, can better assist the LLM in creating high-quality environments.
For generating environment specifications, One direct approach is to prompt LLMs to generate random environments. 
However, due to the inherent inductive bias of LLMs, they struggle to generate diverse environments in this way.
Therefore, to address this issue, we build an inspiration corpus D={t_0,t_1,·,t_n}, containing sufficiently diverse text segments used to serve as the "inspiration" for generating environment specification with LLMs. 
More specifically, when generating an environment, we first sample a text segment t_i from D, then prompt the LLM to generate a related environment based on t_i.
Taking the example in Figure <ref>, we first sample a text segment "How to boost your diet with peanut butter powder?" from D. Then we prompt an LLM to generate a related environment where the agent is defined as a nutritionist tasked with creating a new healthy recipe book that prominently features peanut butter powder as a key ingredient. 
This approach significantly enhances the diversity of generated environments, thereby empowering more generalized agent training.
The inspiration corpus can be implemented in various ways, such as using a large-scale pre-trained corpus like Common Crawl. Alternatively, a domain-specific corpus, such as a code generation dataset <cit.>, can be used to generate environments for a specific domain. This paper uses LIMA <cit.> as the inspiration corpus, an instruction-tuning dataset with sufficient diversity.



  
Environment Implementation
Conditioned on the generated environment specification, we generate its corresponding code, i.e., implementing the environment.
This can be formulated as a typical code-generation problem with LLMs.
We also introduce a validation tool capable of capturing syntax errors to provide feedback during the code generation process, thereby iteratively refining it.



  
Environment Library
We define the library at iteration t as: L_t = L_0 ∪⋃_k=1^t{𝔼_i | 𝔼_i = g_𝔼(i, L_k-1), i ∈ I_k, v(𝔼_i) = true},
where L_0 is the initial seed library, and the union represents all verified environments generated up to iteration t.
This iterative process allows continuous expansion and refinement of the environment library, potentially leading to increasingly complex and diverse environments over time.





 §.§ Task Generation








  
Overview
As depicted in Figure <ref>, conditioned on the generated environments, we prompt LLMs to generate corresponding planning tasks.
We employ a two-stage generation approach for creating a diverse range of planning tasks in terms of difficulty. 

We begin by prompting the LLM with a specific environment, enabling it to generate an initial set of planning tasks in a zero-shot way.
Subsequently, we adjust these tasks to make them simpler or more challenging, forming a comprehensive set of planning tasks.




  
Bidirectional Evolution
Many studies have proposed evolving instructions, primarily focusing on making instructions more difficult <cit.>. The effectiveness of this approach relies heavily on the assumption that LLMs inherently possess the ability to follow simple instructions. 
However, according to findings from some studies <cit.>, LLMs often exhibit poor performance even in simple planning tasks. 
Therefore, we propose , which introduces evolution in two directions: easy-evol and hard-evol. 
Easy-evol typically involves simplifying the goal conditions. The motivation is that easier tasks can facilitate learning when the agent performs poorly and cannot directly learn from typically difficult goals. 
Conversely, hard-evol usually involves making the goal conditions more complex, increasing the number of steps required for the agent to complete the task. This can further enhance the agent's capability to perform the planning task.
To our knowledge, we are the first to introduce bidirectional evolution in the agent data generation scenario.
The prompt examples are shown in Figure <ref>.












§ EXPERIMENTS





To evaluate the effectiveness of the proposed framework, we synthesize environments and planning tasks using the Planning Domain Definition Language (PDDL), a widely adopted programming language for planning.
Subsequently, we evaluate its performance across various unseen planning tasks in a zero-shot manner. To validate the effectiveness and generalizability of , we categorized the evaluated tasks into two distinct groups: 
i) In-Domain Tasks: Planning tasks implemented using PDDL.
ii) Out-of-Domain Tasks: These comprise tasks developed using other programming languages, such as Python.



 §.§ Experimental Setup




  
Evaluation Tasks

For In-Domain Tasks, we select four widely used PDDL-based planning tasks: Blocksworld, Gripper, Tyreworld, and Barman <cit.>.

More explicitly, Blocksworld requires an agent to achieve a target configuration by moving blocks, while Gripper involves moving objects between different rooms.
Tyreworld simulates changing a car tire, including removing the flat tire, replacing it with a spare, and installing the new tire. 
Barman emulates a bartender's tasks in mixing cocktails, which include combining various ingredients, using shakers, and garnishing drinks.

For Out-of-Domain Tasks, we select two challenging partial-observable planning tasks: Alfworld <cit.> and BabyAI <cit.>. 
Alfworld is an environment designed to test agents' abilities to perform everyday household tasks.
While in BabyAI, the agent interprets and executes natural language instructions in a grid-world setting.




  
Evaluation Metrics

We utilized two evaluation metrics to evaluate planning ability: success rate and progress rate <cit.>. 
During each interaction round, we assigned a progress rate, denoted as r_t, to measure the progression towards the goal state g. 
As the agent transitions through states s_t = [s_0, …, s_t], its progress is assessed using a matching score f(·, g) → [0, 1], which quantifies the similarity between the current state and the goal state. 
Initially, r_t is set to 0, indicating no progress. 
Only when the progress rate reaches 1 does the success rate attain 1; all other scenarios yield a 0 outcome. 
The success rate reflects the agent's capacity to complete a comprehensive task.



  
Baselines


We compare with a series of widely-used multipurpose foundation models that exhibit state-of-the-art performance, such as GPT-3.5 <cit.> (gpt-3.5-turbo) and GPT-4 <cit.> (gpt-4-turbo), CodeLlama-7B <cit.>, Mistral-7B <cit.>, Llama2-7B <cit.>, and Llama3-8B <cit.>. We use their instruct-tuned versions for all multipurpose foundation models (<ref>). Additionally, some models have undergone specialized training on agent trajectory data, such as AgentLM-7B <cit.>.
AgentLM is instruction-tuned using a combination of a lightweight trajectory dataset across 6 environments and a larger general instruction-tuning dataset.









  
Implementation Details

We followed the environment and task implementation of AgentBoard <cit.>. For the configuration of evaluation tasks, we employ act-only prompting <cit.>, setting the maximum step length for the LLM agent to 30.
We selected LIMA <cit.> as the text corpus D for generating environments, which leverages various data manipulation techniques to ensure a diverse range of instructions. 
For environment generation and task generation, we employ GPT-4 [We applied the gpt-4-20230321 API from Azure OpenAI service.], configuring the inference parameters with a temperature of 0 and a top_p value of 0.95. 

Based on , we generated a total of 592 environments. For each environment, we generated ten unconditioned tasks, which were then evolved into ten refined tasks using . To generate trajectory data for training, we utilized the domain-independent planner FastDownward[<https://www.fast-downward.org/>], ensuring optimal trajectory data. This process ultimately led to 7246 trajectories. 
More details of the dataset can be found in Appendix <ref> and <ref>.

Since the trajectory data is structured, such as "pickup(o1)", we employ GPT-4 to generate a natural language mapping, for example, "pick up object {arg1}", to transform structured actions into natural language actions. We detailed the generation of natural language mapping are presented in <ref>.

During the training process, we employed Llama3-8B [https://huggingface.co/meta-llama/Meta-Llama-3-8B-Instruct] as our foundation model. The hyperparameters were configured as follows: batch size of 64, 10 epochs, a context length of 4096 tokens, no warmup steps. Checkpoints from epochs 5 through 10 were preserved and subsequently evaluated on In-Domain Tasks. The model exhibiting optimal performance was selected for further assessment on Out-of-Domain Tasks.
We conducted all experiments utilizing a cluster of eight V100 GPUs.



 §.§ Evaluation on In-Domain Tasks
























































































































Despite its relatively small size, overall outperforms GPT-3.5 in success rate (11.67 vs. 5.0). Furthermore, in the barman task, even surpassed GPT-4's performance (15 vs. 10). also achieved a comparable level to GPT-4 in tyreworld.
Compared to other models with similar parameter scales, consistently outperforms them across four distinct tasks.


Compared to Llama3, our model shows an overall success rate and progress rate improvement of 10 and 9.95, respectively. Notably, in several tasks where the success rate of Llama3 is zero (gripper, blockworld, tyreworld), achieves significant breakthroughs, further demonstrating the effectiveness of the dataset.






We can draw several conclusions from previous discussions:
i) surpasses GPT-3.5 in overall performance and achieves results that either exceed or are comparable to GPT-4 on certain specific tasks;

ii) fine-tuned Llama3 has achieved a significant improvement in success rate;
iii) consistently surpasses other models with similar parameter scales in performance.




 §.§ Robustness





















To validate the robustness of the constructed dataset with , we conducted a series of experiments to evaluate its performance across different foundation models. 
We selected several widely used 7-8B foundation models, including Llama3-8B, CodeLLama-7B, and Mistral-7B, to test the versatility and effectiveness of .
As is shown in Table <ref>, all three models exhibited significant improvements after training, with Llama3-8B showing the highest success rate increase of 10.0 and CodeLLama-7B demonstrating a maximum progress rate increase of 9.9. These experimental results prove that the dataset constructed with for agent training is highly effective across different models.



 §.§ Evaluation on Out-of-Domain Tasks










§ RELATED WORK




  
Large Language Model based Agent.

Large Language Models have demonstrated exceptional reasoning capabilities <cit.>. Owing to such abilities, over the past two years, LLM-based agents have experienced significant development <cit.>.
Unlike the traditional method of using LLMs for text-based reasoning, such as Chain-of-Thought <cit.>, LLM-based agents typically involve interaction with the environment, adjusting the output in a closed-loop manner based on environmental information.
These LLM-based agents, now fortified with capabilities like Memorizing <cit.>, Tool-use <cit.>, and Planning <cit.>, exhibit a marked enhancement in their overall efficacy. 
Although this paper mainly focuses on the planning capability of LLM-based agents, we believe has the potential to generalize to other scenarios of LLM-based agents.



  
Planning with Large Language Models.

Planning is one of the key applications of LLM-based agents, applicable in various scenarios such as robotic planning <cit.>, travel planning <cit.>, calendar scheduling <cit.>, code generation <cit.> and others <cit.>. 
It is typically defined as the process of systematically determining a sequence of actions or steps required to achieve a desired goal from an initial state, considering constraints and available resources. 
This definition primarily differentiates from studies that utilize LLMs to generate ungrounded plans as guidance for problem-solving <cit.>, rather than directly producing executable actions.
Planning can be categorized into two types: open-loop planning, where the LLM outputs an entire action sequence before execution <cit.>, and closed-loop planning, where the LLM-based agent decides the next action based on real-time environmental interaction after executing a previous action <cit.>. 
This paper mainly focuses on close-loop planning, which is more adaptable for error correction, human interaction, and environmental grounding. 
Recent studies on close-loop planning have integrated chain-of-thought reasoning into the planning process <cit.>. 
Additionally, some papers have explored the use of tree-search methods to enhance the performance of LLM planning <cit.>.
Instead of designing novel frameworks or engaging in prompt engineering, this paper explores how training can enhance the planning capabilities of LLM-based agents.




  
Agent Training.

Recently, numerous studies have aimed to enhance LLM-based agent capabilities by incorporating agent trajectory data into their training <cit.>. 
Advanced works such as AgentTuning <cit.> utilize GPT-4 to generate trajectory data across six distinct environments. Subsequently, this data is filtered and employed in training Large Language Models, enhancing the agent capabilities of base models.
Another work, FireAct <cit.>, proposes training with both CoT data and ReAct format data, enabling the model to discern when to use reasoning to solve problems and when to call external tools. 
Agent LUMOS <cit.> suggests separately training Planning and Grounding models, enabling LLM-based agents to learn to decompose complex problems before execution.
LLM-Modulo framwork <cit.> proposes to leverage LLMs generating candidate plans and verify them with an external verifier. Then, use the verified trajectories for fine-tuning LLMs.
Similarly,  <cit.> takes a generate-test loop to synthesize trajectories for LLM training.
Unlike previous papers on all agent training, goes beyond merely generating trajectory data using Large Language Models. Instead, we utilize Large Language Models to generate agent environments, which can be considered a more foundational application. As a result, we have constructed over 500 environments for training, whereas previous works typically use fewer than 10 environments to synthesize agent data.




  
Environment and Task Generation with Large Language Models.


The utilization of LLMs to generate environments and tasks is an emerging application.
Some studies have explored utilizing LLMs to generate layouts in robotic simulations, typically involving the creation of configuration files <cit.>. 
While these methods can construct numerous scene-level environments, they often struggle to achieve diversity at the underlying mechanism level. 
Agenttuning <cit.> employs a task generation approach similar to the Self-instruct <cit.> method, using the test set as seed data. This approach not only poses a risk of data leakage but also leads to insufficient diversity in task difficulty.
ByteSized32 <cit.> uses LLMs to generate Python-based games based on predefined task specifications automatically. 
Similarly, other works <cit.> leverage LLMs to automatically construct PDDL domains based on a task specification.
In contrast to these studies, this paper proposes using a diverse text corpus to generate environment code automatically. 
This approach facilitates the creation of a wide range of rich environments without predefined definitions.



§ CONCLUSION




In this paper, we explore using LLMs to automatically generate environment and planning tasks for LLM-based agent training. 
Specifically, for generating diverse environments, we propose utilizing an inspiration corpus composed of various domain-specific text segments as the context for environment synthesis. To enhance the difficulty diversity of generated planning tasks, we introduce a bidirectional evolution method, , which evolves planning tasks from both easier and more challenging directions to create a task set with a more gradual difficulty curve, thereby improving the effectiveness of LLM learning. Based on , we developed a dataset consisting of 592 environments and 7246 trajectories and trained it on a series of LLMs. The trained model outperformed GPT-3.5 across multiple tasks and, in certain specific instances, exceeded the performance of GPT-4.

plainnat








colframe = blue!50!black, 
colback = white, 
toptitle=1mm, 
bottomtitle=1mm, 
fonttitle=10pt12pt,  breakable, enhanced jigsaw, finish broken=colframe=white, bottom=0mm, bottomrule=0mm,before upper=colframe=white, toprule=0mm,after=colframe=blue!40!black, toprule=1mm, bottomrule=1mm,



§ MORE IMPLEMENTATION DETAILS




 §.§ Models



We applied the instruct version for models. Specifically, the detailed version for each model is presented in Table <ref>.





 §.§ Natural Language Mapping



We leverage GPT-4 to generate the natural language mapping that converts structured actions into its natural language format. When the mapping failed to yield, we heuristically serialized the structured actions.
The prompt for generating natural language mapping with GPT-4 is as follows:

[colback=gray!10!white,colframe=gray!50!white,title=Natural Language Mapping Generation]
I would like you to create natural language mapping for PDDL.   

The form of the natural language mapping is a Python dictionary, wherein

1. The key corresponds to the name of a predicate or action within the domain PDDL.

2. The value is its equivalent in natural language, with parameters presented in "{argn}", where n is the index of its parameters in the PDDL expression.

3. You must ensure that the number of "{}" corresponds precisely to the number of parameters in predicates or actions.  

4. You should very carefully check the order of {argn}. 



Your output must strictly follow the provided example.



Example:

PDDL Domain:

“`pddl

(define (domain hanoi)

  (:requirements :strips)

  (:predicates

  (clear ?x)

  (on ?x ?y)

  (smaller ?x ?y)

  )



  (:action move

    :parameters (?disc ?from ?to)

    :precondition (and (smaller ?to ?disc) (on ?disc ?from)

               (clear ?disc) (clear ?to))

    :effect  (and (clear ?from) (on ?disc ?to) (not (on ?disc ?from))

          (not (clear ?to))))

  )

“`

Specification:

Your goal is to solve the Tower of Hanoi puzzle, which involves moving a stack of discs from one peg to another, with the restriction that no disc may be placed on top of a smaller disc. The puzzle is solved when all the discs are moved to the target peg following these rules.



The actions defined in this domain include:

- move <disc> <from> <to>: This action allows moving a disc from one peg to another. The preconditions for this action are that the target peg is smaller than the disc being moved, the disc is on the source peg, and both the disc and the target peg are clear (i.e., there is no disc on top of them). The effect of this action is that the source peg becomes clear, the disc is now on the target peg, the disc is no longer on the source peg, and the target peg is no longer clear.



You have the following restrictions on your actions:

- A disc can only be moved if it is clear, meaning there is no other disc on top of it.

- A disc can only be placed on another disc or peg that is larger than itself.

- A disc can only be moved to a peg that is clear.

- Once a disc is moved from a peg, that peg becomes clear.

- Once a disc is placed on a peg, that peg is no longer clear.

 Natural Language Mapping:

“`python

{

    "clear": "{arg1} is clear.",

    "on": "{arg1} is on {arg2}.",

    "smaller": "{arg1} is smaller than {arg2}.",

    "move": "Move {arg1} from {arg2} to {arg3}."

}

“`



You need to generate the corresponding natural language mapping for the following pddl domain.

    

PDDL Domain:

{PDDL_Domain}

Specification:

{PDDL_Description}

Natural Language Mapping:






§ MORE STATISTICS ON ENVIRONMENT





 §.§ Environment Specification




We analyzed the token distribution within the environmental specifications. Among the 592 environmental specifications, the average token count is 473.55, with a median of 467.00. The minimum token count is 207, and the maximum is 934. As depicted in Figure <ref>, the number of specification tokens for the environment is predominantly concentrated within the range of 300 to 699.



 §.§ Environment Implementation




























The scale of action space and state space in an environment typically dictates its complexity, with a greater number of actions and states generally indicating a more complex environment. 
An environment library with a greater variety of difficulty levels is preferable for a training set. 
As shown in Figure <ref>, there is a significant diversity in the number of actions and predicates.






 §.§ Diversity Analysis



We evaluate the diversity of generated environments using cosine similarity. More specifically, we randomly sampled 100 environment specifications for better visualization and converted them into TF-IDF vectors. After calculating the cosine similarity matrix between all pairs of specifications, we visualize the matrix using heatmap as is shown in Figure <ref>.
The computed average cosine similarity of the sampled environment specifications is 0.176, indicating that the corpus exhibits a high degree of diversity, reflecting a rich tapestry of distinct semantic features and thematic elements.



§ EXAMPLES



In this section, we present the specific details of the cases depicted in Figure <ref> and Figure <ref>.



 §.§ Environment Specification


[colback=gray!10!white,colframe=gray!50!white,title=Environment Specification]
You are a nutritionist tasked with creating a new healthy recipe book that incorporates peanut butter powder as a key ingredient. Your environment consists of your office, a computer for research and writing, a kitchen for testing recipes, and a collection of ingredients including peanut butter powder. The computer can be used to write and save recipes, but it needs to be charged after every three hours of use.



The actions defined in this domain include:

- research_ingredient <location> <ingredient> <nutritionist>: This action allows the nutritionist to research peanut butter powder at the office. The effect of this action is that the nutritionist gains knowledge about the nutritional benefits and potential uses of peanut butter powder in recipes.

- develop_recipe <location> <computer> <nutritionist> <ingredient>: This action allows the nutritionist to create a new recipe using peanut butter powder. The precondition for this action is that the nutritionist has researched peanut butter powder. The effect of this action is a draft of a new recipe that incorporates peanut butter powder.

- test_recipe <location> <kitchen> <nutritionist> <recipe>: This action allows the nutritionist to test the drafted recipe in the kitchen. The preconditions for this action are that the nutritionist has developed a recipe and has all the necessary ingredients, including peanut butter powder. The effect of this action is feedback on the recipe's taste, nutritional value, and overall feasibility.

- finalize_recipe <location> <computer> <nutritionist> <tested_recipe>: This action allows the nutritionist to finalize the recipe after testing and receiving feedback. The preconditions for this action are that the nutritionist has tested the recipe and the computer is charged. The effect of this action is a completed recipe that is ready to be included in the healthy recipe book.



You have the following restrictions on your actions:

- You can only research peanut butter powder if you are in the office where your research materials are located.

- You can only develop a recipe after researching peanut butter powder and must have a charged computer.

- You can only test a recipe in the kitchen if you have developed a recipe and have all the necessary ingredients.

- You can only finalize a recipe after testing it in the kitchen and receiving feedback, and if your computer is charged.





 §.§ Environment Implementation


[colback=gray!10!white,colframe=gray!50!white,title=Environment Implementation, sharp corners=south]
[t].95
[commandchars=
{}]
(define (domain healthy-recipe-book)
  (:requirements :strips)
  (:predicates
    (in-office ?nutritionist)
    (in-kitchen ?nutritionist)
    (researched-peanut-butter ?nutritionist)
    (has-recipe-draft ?nutritionist ?recipe)
    (has-tested-recipe ?nutritionist ?recipe)
    (computer-charged)
    (has-ingredients ?recipe)
  )

  (:action research_ingredient
    :parameters (?nutritionist ?ingredient)
    :precondition (and 
                    (in-office ?nutritionist))
    :effect (and 
              (researched-peanut-butter ?nutritionist))
  )

  (:action develop_recipe
    :parameters (?nutritionist ?ingredient)
    :precondition (and 
                    (researched-peanut-butter ?nutritionist) 
                    (computer-charged))
    :effect (and 
              (has-recipe-draft ?nutritionist ?ingredient) 
              (not (computer-charged)))
  )

  (:action test_recipe
    :parameters (?nutritionist ?recipe)
    :precondition (and 
                    (in-kitchen ?nutritionist) 
                    (has-recipe-draft ?nutritionist ?recipe) 
                    (has-ingredients ?recipe))
    :effect (and 
              (has-tested-recipe ?nutritionist ?recipe))
  )

  (:action finalize_recipe
    :parameters (?nutritionist ?tested_recipe)
    :precondition (and 
                    (has-tested-recipe ?nutritionist ?tested_recipe) 
                    (computer-charged))
    :effect (and 
              (not (computer-charged)) 
              (not (has-tested-recipe ?nutritionist ?tested_recipe)))
  )






 §.§ Trajectory Data


[colback=gray!10!white,colframe=gray!50!white,title=Trajectory Data]
(User) You are a nutritionist tasked with creating a new healthy recipe book that incorporates peanut butter powder as a key ingredient. Your environment consists of your office, a computer for research and writing, a kitchen for testing recipes, and a collection of ingredients including peanut butter powder. The computer can be used to write and save recipes, but it needs to be charged after every three hours of use.



The actions defined in this domain include:

- research_ingredient <location> <ingredient> <nutritionist>: This action allows the nutritionist to research peanut butter powder at the office. The effect of this action is that the nutritionist gains knowledge about the nutritional benefits and potential uses of peanut butter powder in recipes.

- develop_recipe <location> <computer> <nutritionist> <ingredient>: This action allows the nutritionist to create a new recipe using peanut butter powder. The precondition for this action is that the nutritionist has researched peanut butter powder. The effect of this action is a draft of a new recipe that incorporates peanut butter powder.

- test_recipe <location> <kitchen> <nutritionist> <recipe>: This action allows the nutritionist to test the drafted recipe in the kitchen. The preconditions for this action are that the nutritionist has developed a recipe and has all the necessary ingredients, including peanut butter powder. The effect of this action is feedback on the recipe's taste, nutritional value, and overall feasibility.

- finalize_recipe <location> <computer> <nutritionist> <tested_recipe>: This action allows the nutritionist to finalize the recipe after testing and receiving feedback. The preconditions for this action are that the nutritionist has tested the recipe and the computer is charged. The effect of this action is a completed recipe that is ready to be included in the healthy recipe book.



You have the following restrictions on your actions:

- You can only research peanut butter powder if you are in the office where your research materials are located.

- You can only develop a recipe after researching peanut butter powder and must have a charged computer.

- You can only test a recipe in the kitchen if you have developed a recipe and have all the necessary ingredients.

- You can only finalize a recipe after testing it in the kitchen and receiving feedback, and if your computer is charged.

Goal: The goal is to satisfy the following conditions: The computer is charged. jordan has tested the recipe almond_butter_bars.

Observation:The computer is charged. The recipe almond_butter_bars has all the necessary ingredients. jordan has a recipe draft for almond_butter_bars. jordan has researched peanut butter. jordan is in the kitchen.

(Assistant) Action: jordan tests the recipe almond_butter_bars.

(User) Observation: The computer is charged. The recipe almond_butter_bars has all the necessary ingredients. jordan has a recipe draft for almond_butter_bars. jordan has researched peanut butter. jordan has tested the recipe almond_butter_bars. jordan is in the kitchen.

(Assistant) Action: jordan develops a recipe using almond_butter_bars.










 Paper Title: Gradient Harmonization in Unsupervised Domain Adaptation.


























































Journal of  Class Files, Vol. 14, No. 8, August 2015
Shell et al.: Bare Demo of IEEEtran.cls for Computer Society Journals





































§ ABSTRACT 
Unsupervised domain adaptation (UDA) intends to transfer knowledge from a labeled source domain to an unlabeled target domain. Many current methods focus on learning feature representations that are both discriminative for classification and invariant across domains by simultaneously optimizing domain alignment and classification tasks. However, these methods often overlook a crucial challenge: the inherent conflict between these two tasks during gradient-based optimization.
In this paper, we delve into this issue and introduce two effective solutions known as Gradient Harmonization, including GH and GH++, to mitigate the conflict between domain alignment and classification tasks. GH operates by altering the gradient angle between different tasks from an obtuse angle to an acute angle, thus resolving the conflict and trade-offing the two tasks in a coordinated manner. Yet, this would cause both tasks to deviate from their original optimization directions. We thus further propose an improved version, GH++, which adjusts the gradient angle between tasks from an obtuse angle to a vertical angle. This not only eliminates the conflict but also minimizes deviation from the original gradient directions.
Finally, for optimization convenience and efficiency, we evolve the gradient harmonization strategies into a dynamically weighted loss function using an integral operator on the harmonized gradient.
Notably, GH/GH++ are orthogonal to UDA and can be seamlessly integrated into most existing UDA models. Theoretical insights and experimental analyses demonstrate that the proposed approaches not only enhance popular UDA baselines but also improve recent state-of-the-art models.





Unsupervised Domain Adaptation, Gradient Harmonization, Transfer Learning, Image Classification.




Gradient Harmonization in Unsupervised Domain Adaptation
    Fuxiang Huang, 
         Suqi Song,
         Lei Zhang

        
This work was partially supported by National Key R&D Program of China (2021YFB3100800), National Natural Science Fund of China (62271090), Chongqing Natural Science Fund (cstc2021jcyj-jqX0023). This work is also supported by Huawei computational power of Chongqing Artificial Intelligence Innovation Center. (Corresponding author: Lei Zhang)

Fuxiang Huang, Suqi Song and Lei Zhang are with the School of Microelectronics and Communication Engineering, Chongqing University, Chongqing 400044, China. (E-mail: huangfuxiang@cqu.edu.cn, songsuqi@stu.cqu.edu.cn, leizhang@cqu.edu.cn).








Manuscript received April 19, 2005; revised August 26, 2015.

    
================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================






























§ INTRODUCTION




























Deep convolutional neural networks and transformers, driven by extensive labeled samples, have achieved remarkable success in various computer vision tasks such as classification, semantic segmentation and object detection.  However, these models often demonstrate high vulnerability when deployed in novel application scenarios due to data distribution discrepancies. The process of collecting and annotating data across various domains is expensive, labor-intensive and time-consuming. Consequently, Unsupervised Domain Adaptation (UDA) arises to transfer the knowledge from a labeled source domain to an unlabeled target domain <cit.>. 



In recent years, UDA algorithms have made significant progress in enhancing classification performance <cit.>. The primary focus of these approaches is to acquire domain-invariant feature representations, thereby achieving domain alignment and narrowing the probability distributions across domains. Currently, domain alignment methods fall into two main categories: distance metrics-based methods <cit.> and adversarial learning-based methods <cit.>. The distance metrics-based approaches align the source and target domains by mapping them into a shared feature space while minimizing the distribution disparities between them. Inspired by generative adversarial networks <cit.>, adversarial learning techniques have been introduced to tackle domain adaptation challenges <cit.>. For instance,  <cit.> is among the pioneering methods to achieve domain alignment by introducing a game between the feature extractor and domain classifier. <cit.> adjusts the adversarial domain adaptation models on discriminative information conveyed in the classifier predictions. These methods usually jointly optimize domain alignment and classification tasks in the course of training.

However, due to their objective differences between the two tasks, the optimal gradient descent directions from the two tasks may be uncoordinated or imbalanced. Therefore, directly sharing network parameters (i.e., feature generator) may result in optimization conflict between tasks and affect domain-invariant feature learning. During optimization, the angle between the gradients of the two tasks is sometimes an obtuse angle, indicating an obvious gradient conflict. As shown in Fig. <ref>, the black arrow and red arrow point to the optimal gradient descent direction of the domain alignment task and classification task, respectively. In the joint optimization process, obtuse angle formed by the two gradients of both tasks leads to optimization conflict, which reveals the sub-optimality of the two tasks.

To further verify the above observations, we conduct some experiments on real cross-domain classification scenarios. Fig. <ref> displays inner product distributions of two gradients between two tasks (alignment vs. classification) throughout the training process on task of MNIST → USPS. From Fig. <ref>, we can obtain the following two observations. First, it implies that the inner product between two gradients is approximately normally distributed with a mean slightly around above zero. Second, both acute and obtuse angles exist on the two models, where obtuse angles account for about 42% of the total number in MCD <cit.> and 37% in DWL <cit.> (DWL relieves the imbalance by trade-offing the transferability and discriminability to some extent but ignores optimization conflicts, whereas MCD does not take into account the conflict problem). In the joint optimization process, the aggregated gradient direction generated by two obtuse gradients will be seriously deviated from their respective optimal gradient descent directions, thereby resulting in sub-optimal alignment and classification. Further, we empirically observe that the obtuse angle is always present during training over time and cannot be eliminated automatically. Therefore, how to reasonably eliminate optimization conflicts while maintaining the task-specific optimality during training is a challenging and practical work. 

Multi-task learning (MTL), as explored in studies such as <cit.>, focuses on simultaneously optimizing multiple conflicting criteria. These approaches typically seek one or more Pareto optimal solutions that offer different trade-offs. Inspired by MTL, ParetoDA <cit.> introduces a target-classification-mimicking (TCM) loss based on held-out target data and dynamically seeks a desirable Pareto optimal solution for the target domain. However, Pareto optimization relies on additional information to make decisions and requires significant computational resources and time in high-dimensional or complex systems.

In this paper, to address the aforementioned challenges, we propose an idea of de-conflict on the gradients. Technically, we introduce two intuitive yet highly effective approaches called Gradient Harmonization, involving GH and GH++, which aim to alleviate optimization conflicts that emerge between the domain alignment and classification tasks by leveraging insights from geometric awareness in model optimization. Specifically, we first calculate the original gradients of two tasks to find out if there is a conflict, i.e., the angle between the gradients of the tasks is obtuse (negatively correlated). Subsequently, we employ GH to turn the angle between the original gradients of two tasks from an obtuse angle to an acute angle, i.e., positively correlated. This harmonization of gradients enables both tasks to progress in a coordinated manner. However, this process may cause both tasks to deviate from their original directions to compromise performance.
To further mitigate gradient deviation while maintaining the task-specific optimality, we introduce an improved version, GH++, which turns the angle between the original gradients of two tasks from an obtuse angle to a vertical angle, i.e., orthogonal. Then both tasks can evolve harmoniously during joint training while preserving their individual task-specific optimality.


Detailed theoretical support is also provided for the proposed approaches.

Furthermore, in order to facilitate optimization and improve computation efficiency, we derive an equivalent but more efficient model of GH/GH++ for unsupervised domain adaptation. The equivalent model is derived as a dynamic objective function, namely UDA+GH/GH++, which can achieve fast elimination of optimization conflicts without sacrificing the task-specific objective. The proposed GH/GH++ is a universal plug-and-play approach for between-task balance learning and can be easily embedded in most alignment-based unsupervised domain adaptation methods. This work is a new upgrade and more general form of our previous CVPR paper <cit.> from the perspective of gradient harmonization.

The main contributions and novelties of this paper are summarized as follows:


  * We verify and visualize the gradient conflict in UDA methods and develop two novel Gradient Harmonization approaches, including  GH and GH++. The proposed approaches alleviate the gradient conflict problem between the domain alignment task and classification task by adjusting their gradient angle from obtuse to acute/vertical angle adaptively. Fig. <ref> is a representative and illustrative example to depict the UDA model based on GH/GH++.

  * The proposed approaches are orthogonal to most existing UDA methods. To facilitate optimization and improve computation efficiency, we further derive the equivalent models, i.e., UDA with GH/GH++. Specifically, the proposed approaches can be evolved into dynamically weighted loss functions via an integral operation, and promote optimization.

  * Extensive experiments validate the effectiveness and universality of our approaches on various benchmarks in several mainstream UDA models. More insights and analyses justify the reasonability of the proposed approaches.





§ RELATED WORK




 §.§ Unsupervised Domain Adaptation (UDA)




Unsupervised Domain Adaptation (UDA) aims to leverage the knowledge learned from a labeled source dataset to solve similar tasks in a new unlabeled domain. Recent works <cit.> have focused on UDA based on domain alignment and discriminative feature learning methods.



Domain Alignment. The mainstream approach for UDA aims to realize domain alignment by learning domain-invariant feature representations across domains, which can be mainly summarized into two categories: distance metric based methods and adversarial learning based methods. First, the methods based on distance metrics mostly use several existing distance measurement indicators to assess the alignment degree across domains, such as MMD <cit.>, MDD <cit.> and their variants. With the popularity of deep learning methods, more and more researchers apply deep neural networks to domain adaptation. Compared with traditional non-deep domain adaptation methods, deep methods directly improve the learning effect on different classification tasks. Second, some domain adaptation methods based on adversarial training can learn the domain-invariant feature representations well and better implement knowledge transfer. For example, CDAN <cit.> proposes an adversarial adaptation model for the discriminative information transmitted in the prediction of the classifier. GVB <cit.> learns the domain-invariant feature representations by applying the gradually vanishing bridge mechanism on the feature generator.
Besides, considering the equilibrium problem of adversarial learning, FGDA <cit.> reduces the distribution discrepancy by constraining feature gradient. CGDM <cit.> explicitly minimizes the discrepancy of gradients generated by source samples and target samples to improve the accuracy of target samples. FixBi <cit.> introduces a fixed ratio-based mixup to augment multiple intermediate domains between the source and target domain.

Classification. Other approaches focus on improving the performance of the classification task by enhancing class discriminability feature learning. ETD <cit.> puts forward an attention-aware optimal transport distance to measure the domain discrepancy under the guidance of the prediction feedback and enables the model to learn distinguished feature representations. MCD <cit.> proposes to maximize the prediction discrepancy of two classifiers to obtain strong discriminant feature representations. BSP <cit.> penalizes the largest singular values so that other eigenvectors can be relatively strengthened to boost the feature discriminability. CAN <cit.> optimizes the metric for minimizing the domain discrepancy, which explicitly models the intra-class domain discrepancy and the inter-class domain discrepancy. RSDA <cit.> introduces a spherical classifier for label prediction and a spherical domain discriminator for discriminating domain labels and utilizes robust pseudo-label loss in the spherical feature space. DMRL <cit.> focuses on guiding the classifier to enhance consistent predictions between samples and enriches the intrinsic structures of the latent space. To combine the source and target domains, DMRL introduces two mixup regularizations based on randomness. Recently, SSRT <cit.> and TVT <cit.> learn both transferable and discriminative features by applying Vision Transformer (ViT).

Balancing Domain Alignment and Classification.
Previous approaches concentrate on learning domain-invariant and class-discriminative feature representations by jointly optimizing domain alignment and classification task but ignore imbalance or uncoordinated optimization problem between tasks.

Recently, some methods have started to pay attention to this problem and adopt strategies to alleviate it. To adapt to different cross-domain classification scenarios, DWL <cit.> proposes dynamic weighted learning to avoid the discriminability vanishing problem caused by excessive alignment and domain misalignment problem caused by excessive discriminant learning from a macro perspective, which ignores gradient conflict between two tasks during optimization.
Meta <cit.> maximizes inner product of the gradients of the two tasks during training. However, it cannot guarantee the adjusted gradient direction is close to the original optimal gradient descent direction, and even appears serious deviation. Since Meta processes the gradients of the two tasks in an extreme way, it may produce some negative feedback, such as sacrificing partial alignment or classification performance.
Recently, in order to mitigate the effects of conflicting gradients, ParetoDA <cit.> adopts Pareto optimization <cit.> to cooperatively optimize all training objectives, which searches for the desirable Pareto optimal solution on the entire Pareto front. However, it is computationally inefficient. In this paper, we propose two intuitive yet efficient approaches, including GH and GH++, to handle the gradient conflict.

Differences. The proposed approaches represent a fundamental departure from previous methods like DWL <cit.>, Meta <cit.>, and ParetoDA <cit.>. Our primary goal is to implicitly address the optimization conflict throughout the entire training process by active gradient harmonization, thereby ensuring a balanced learning between tasks. However, as depicted in Figure <ref>, the optimization conflict still persists in DWL. Additionally, ParetoDA introduces Pareto optimization and an additional target-classification-mimicking (TCM) loss, which necessitates the involvement of all class-wise discriminators. In contrast, the proposed GH/GH++ is evolved into a dynamically weighted loss function via an integral operator on the harmonized gradient. 






 §.§ Multi Task Learning

Multi-task learning (MTL) <cit.> aims at learning multiple tasks in a unified model to achieve mutual improvement among tasks considering their shared knowledge. By sharing parameters across tasks, MTL methods learn more efficiently with an overall smaller model size compared to learning with separate models.



Prior MTL approaches formulate the total objective as a weighted sum of task-specific objectives, such as DWA <cit.> and GradNorm <cit.>, which optimize weights based on task-specific learning rates or by random weighting. However, these weighted optimization based MLT may not achieve satisfactory performance due to the presence of gradient conflicts among different tasks. Gradient optimization based MLT <cit.> overcome this limitation, mitigating effects of conflicting or dominating gradients. MGDA <cit.> proposes to simply update the shared network parameters along a descent direction which leads to solutions that dominate the previous one. PCGrad <cit.> proposes a “gradient surgery” to avoid gradient conflicts. CAGrad <cit.> seeks the gradient direction within the neighborhood of the average gradient and maximizes the worst local improvement of any objective. <cit.> merge the gradients of auxiliary tasks and applied a scaling factor to adaptively adjust their impact on the main tasks, followed by applying gradient sharing between the main tasks and the merged auxiliary task. GradDrop <cit.> proposes a probabilistic masking procedure, which samples gradients at an activation layer based on their level of consistency. Aligned-MTL <cit.> eliminates instability in the training process by aligning the orthogonal components of the linear system of gradients.



In this study, we introduce two simple yet highly efficient approaches, GH and GH++,  to effectively address gradient conflicts between any two distinct tasks. Detailed theoretical derivations and fundamental insights regarding the proposed approaches are elaborated in Sections <ref>. Additionally, under thorough theoretical analysis and a wealth of experiments, we have substantiated the soundness and effectiveness of the proposed methodologies.







§ PROPOSED APPROACH





 §.§ Problem Definition

Given a labeled source domain 𝒟_s={x_i^s,y_i^s}_i=1^n_s with n_s samples and an unlabeled target domain 𝒟_t={x_j^t}_j=1^n_t with n_t samples, where y_i^s is the class label of the i^th source sample x_i^s. 𝒟_s and 𝒟_t share the same feature space and category space, but have different data distributions. Our purpose is to utilize the labeled data 𝒟_s and unlabeled data 𝒟_t to learn a deep model, which can accurately predict the class label of samples in the target domain.




 §.§ A General Framework of UDA


Adversarial learning has proven to be an effective method for domain alignment, starting from Domain Adversarial Neural Network (DANN) <cit.>. The basic idea is to trick the domain discriminator D by generating features via a feature generator G. Then the domain discriminator predicts whether the generated feature by G is from the source domain or the target domain. The training of domain alignment is achieved through the game between generator and domain discriminator. The parameter θ_g of generator G and the parameter θ_d of domain discriminator D are optimized by the following domain alignment objective function.

    ℒ_dom(θ_g,θ_d)=    𝔼_x_i^s∼𝒟_slog[D(G(x_i^s))]+
            𝔼_x_j^t∼𝒟_tlog[1-D(G(x_j^t))].


In order to improve the classification performance of the target domain samples, we must first ensure that the classifier C can correctly classify the samples from the source domain. Thus, the supervised classification loss can be described as

    ℒ_cls(θ_g,θ_c)=1/N_s∑_i=1^N_sℒ_ce(C(G(x_i^s;θ_g);θ_c),y_i^s),

where ℒ_ce is the standard cross-entropy loss function.

During training stage, the existing methods usually jointly optimize the two objective functions (ℒ_dom and ℒ_cls) to obtain domain-invariant and class-discriminant feature representation. The overall minimax objective function is

    θ_g,θ_cmin  θ_dmax  ℒ_dom+ℒ_cls,

where θ_g, θ_d, θ_c denote the parameters of feature generator, domain discriminator and classifier, respectively.






 §.§ Gradient Harmonization (GH)


Domain alignment and classification are two different tasks. Their optimal gradient descent directions may not be coordinated, which results in the optimization conflict of the two loss functions in the training process and deteriorates the final domain adaptation performance.

In order to ensure that the two target tasks can be optimized in a coordinated manner, we propose an idea of de-conflict on the gradients, which then formulates the proposed GH technique, i.e., altering the gradient angle between different tasks from an obtuse angle to an acute angle. The de-conflict process for the gradients (i.e., g_1 and g_2) is schematically shown in Fig. <ref>. Specifically, we first provide three Lemmas to support our idea. Then the proposed GH is summarized as Theorem 1. For convenience, we define ℒ_1(Θ) and ℒ_2(Θ) as the general form of any two conflicted loss functions. In UDA, ℒ_1(Θ) and ℒ_2(Θ) represent the domain alignment loss ℒ_dom(θ_g,θ_d) and the classification loss ℒ_cls(θ_g,θ_c), respectively. Θ={θ_g, θ_d,θ_c} indicates model parameters of generator, discriminator and classifier. In fact, the three Lemmas aim to deduce the gradient harmonization formula during model optimization (training) phase, by exhausted mathematical solving process.


Given two objective functions ℒ_1(Θ) and ℒ_2(Θ), we define g_1 and g_2 as their gradient, respectively, and g̃_1 is the result of harmonizing the gradient g_1. For minimizing the objective ℒ_1(Θ-g̃_1)+ℒ_2(Θ-g̃_1), we consider to first minimize the objective ℒ_1(Θ-g̃_1), and then there is,

    g̃_1=g_1-δ(g_1^Tg_2<0)g_1^Tg_2/‖ g_2 ‖^2g_2,

where δ(·) represents the indicator function whose value is 0 or 1 and the mathematical expression is:

    δ(A) =
    
    1,      if A is true 
    
    0,      if A is false







In order to minimize the loss function ℒ_1(Θ-g̃_1), the problem is transformed into solving the following model:

    g̃_1min  ℒ_1(Θ-g̃_1)  
    
    s.t.    ℒ_1(Θ-g̃_1)≤ℒ_1(Θ),  ℒ_2(Θ-g̃_1)≤ℒ_2(Θ).

Due to the gradient g_1=∇_Θℒ_1(Θ) is the fastest descent direction of the objective function ℒ_1(Θ), if the harmonic gradient g̃_1 is the same as g_1 as much as possible, the loss function ℒ_1(Θ) will be optimized along the gradient descent direction. If the angle between the harmonic gradient g̃_1 and g_1 is an acute angle, then g̃_1 points to the direction of gradient descent. Naturally, the angle between g̃_1 and g_2=∇_Θℒ_2(Θ) is also expected to be an acute angle. Therefore, the above model is transformed into:

    g̃_1min  1/2‖ g_1-g̃_1‖^2  
       s.t. g̃_1^Tg_1≥0, g̃_1^Tg_2≥0.

The problem (<ref>) is a convex optimization problem. To solve the problem (<ref>), we define the Lagrangian function










    ℒ(g̃_1,α_1,α_2)=1/2‖ g_1-g̃_1‖^2
        -α_1g̃_1^Tg_1-α_2g̃_1^Tg_2,

where α_1, α_2 are the Lagrangian multipliers associated with the inequality constraint of problem (<ref>). To simplify the solution, first note that

    P=    max_α_1≥0,α_2≥0ℒ(g̃_1,α_1,α_2) 
    
    =   {[        1/2‖ g_1-g̃_1‖^2, g̃_1^Tg_1≥0, g̃_1^Tg_2≥0;                      +∞,                otherwise;                          ].

This means that we can express the optimal value of the primal problem (<ref>) as:

    p^∗ = min_g̃_1  max_α_1≥0,α_2≥0  ℒ(g̃_1,α_1,α_2).

By the definition of the dual function, we can express the optimal value of the dual problem:


    d^∗=α_1≥0,α_2≥0max  g̃_1min  ℒ(g̃_1,α_1,α_2),

where α_1≥0,α_2≥0 is the dual feasibility condition. It is easy to get the weak duality d^∗≤ p^∗. When the strong duality holds, i.e., d^∗ = p^∗, we have

    α_1g̃_1^Tg_1=0, α_2g̃_1^Tg_2=0.

Eq. (<ref>) is known as complementary slackness condition. To summarize, for convex  optimization problem (<ref>), we have











    {[        ∇_g̃_1ℒ(g̃_1,α_1,α_2)=0,                               1;       g̃_1^Tg_1≥0, g̃_1^Tg_2≥0,                               2;                    α_1≥0,α_2≥0,                               3; α_1g̃_1^Tg_1=0, α_2g̃_1^Tg_2=0.                               4;                                 ].

Conditions (<ref>) are called the Karush-Kuhn-Tucker (KKT) conditions<cit.>, which is the the necessary condition for pair of primal and dual optimal points. 

We can solve the dual problem Eq. (<ref>). As for the inner optimization problem, i.e., D=g̃_1min  ℒ(g̃_1,α_1,α_2), we then obtain its optimal solution by setting its derivative with respect to g̃_1 as zero. Then,

    g̃_1=g_1+α_1g_1+α_2g_2.


We can calculate the extremum of D by substituting Eq. (<ref>) into Eq. (<ref>) :

    D   =g̃_1min  ℒ(g̃_1,α_1,α_2)
    
           =1/2g̃_1^Tg̃_1 -g̃_1^Tg_1-α_1g̃_1^Tg_1-α_2g̃_1^Tg_2+const
       =1/2g̃_1^Tg̃_1- g̃_1^T(g_1+α_1g_1+α_2g_2)+const
       = -1/2g̃_1^Tg̃_1+const
       =-1/2‖ g_1+α_1g_1+α_2g_2‖^2+const.

We then get the outer optimization problem by substituting Eq. (<ref>) into Eq. (<ref>) as follows:

    α_1≥0,α_2≥0max  -1/2‖ g_1+α_1g_1+α_2g_2‖^2.

For the problem (<ref>), we can take the derivative of α_1 and α_2 respectively, and set the derivative to 0. Then we can obtain the optimal solution: α_1^*=-1, α_2^*=0. Obviously, the obtained α_1^* does not satisfy the dual feasibility condition. Therefore, it is necessary to find the boundary solution that satisfies the constraints, which is established under two possible cases:



Case 1: Suppose α_1^*=0, the Eq. (<ref>) is transformed into:

    α_2min  1/2‖ g_1+α_2g_2‖^2.

We can set the derivative of α_2 to 0, i.e., (g_1+α_2 g_2)^T g_2 = 0. Then we can obtain the optimal solution:




    α_2^*=-g_1^Tg_2/‖ g_2 ‖^2.

Obviously, Eq. (<ref>) satisfies the KKT conditions 124 in Eq. (<ref>). According to dual feasibility, i.e., 3 in Eq. (<ref>), α_2^*=-g_1^Tg_2/‖ g_2 ‖^2≥0, then we have g_1^Tg_2≤0. Therefore, we can obtain a boundary solution

    α_1^*=0,
    α_2^*=-g_1^Tg_2/‖ g_2 ‖^2,    g_1^Tg_2≤0


Case 2: Suppose α_2^*=0, the Eq. (<ref>) is transformed into:

    α_1min  1/2‖ g_1+α_1g_1‖^2.

Then we can obtain α_1^*=0 in nature. Obviously, it satisfies the KKT conditions 134 in Eq. (<ref>). According to primal feasibility, i.e., 2 in Eq. (<ref>), we have g_1^Tg_2≥0. Therefore, the boundary solution is

    α_1^*=0,
    α_2^*=0,     g_1^Tg_2≥0






In summary, the optimal solution of problem (<ref>) is:

    α_1^*=0
    α_2^*=-δ(g_1^Tg_2<0)g_1^Tg_2/‖ g_2 ‖^2.

Then the optimal value of the primal problem (<ref>) g̃ can be obtained by substituting Eq. (<ref>) into Eq. (<ref>).


    g̃_1=g_1-δ(g_1^Tg_2<0)g_1^Tg_2/‖ g_2 ‖^2g_2,

where g̃_1 is the result of harmonizing the gradient g_1 of ℒ_1. Then the proof of Eq. (<ref>) in Lemma 1 is completed.

The above process not only solves the problem (<ref>) but also harmonizes the gradient g_1 of ℒ_1(Θ), i.e., Fig. <ref> (b1).



Define g̃_2 as the result of harmonizing the gradient g_2, we consider to further minimize the optimization objective ℒ_2(Θ-g̃_2), then there is,

    g̃_2=g_2-δ(g_1^Tg_2<0)g_2^Tg_1/‖ g_1 ‖^2g_1.







In order to minimize the loss function ℒ_2(Θ-g̃_2), similar to Lemma 1, we need to solve the following minimization problem.

    g̃min  ℒ_2(Θ-g̃_2)  
    
    s.t.    ℒ_1(Θ-g̃_2)≤ℒ_1(Θ), ℒ_2(Θ-g̃_2)≤ℒ_2(Θ).

The above model can be transformed into

    g̃_2min  1/2‖ g_2-g̃_2‖^2  
       s.t. g̃_2^Tg_1≥0, g̃_2^Tg_2≥0.

Similar to problem (<ref>), by solving (<ref>), we can obtain

    g̃_2=g_2-δ(g_1^Tg_2<0)g_2^Tg_1/‖ g_1 ‖^2g_1,

where g̃_2 is the result of harmonizing the gradient g_2 of ℒ_2(Θ), i.e., Fig. <ref> (b2). 





Define the overall loss function ℒ=ℒ_1(Θ)+ℒ_2(Θ) of two tasks, and g̃ denotes the overall harmonized gradient. We consider the whole optimization objective ℒ_1(Θ-g̃)+ℒ_2(Θ-g̃), then there is,

    g̃=   g̃_1+g̃_2
    
        =   g_1+g_2-δ(g_1^Tg_2<0)g_1^Tg_2/‖ g_2 ‖^2g_2-δ(g_1^Tg_2<0)g_2^Tg_1/‖ g_1 ‖^2g_1.












According to Lemma 1 and Lemma 2, the results of harmonizing the gradients g_1 and g_2, i.e., g̃_1 and g̃_2, can be obtained. By aggregating Eq. (<ref>) in Lemma 1 and Eq. (<ref>) in Lemma 2, the whole gradient after harmonization can be obtained as Eq. (<ref>), and proof of Lemma 3 is completed.



For brevity, we summarize the above process of GH in Theorem 1.


For any two different tasks, optimization conflicts can be eliminated by Gradient Harmonization (GH). Define the overall loss function ℒ (Θ)=ℒ_1 (Θ)+ℒ_2 (Θ), composed of two sub-objectives (refer to domain alignment and classification in this paper), and g_1 and g_2 as the gradient of ℒ_1 (Θ) and ℒ_2 (Θ). Then the overall harmonized gradient g̃ of the whole loss ℒ (Θ) with GH can be formulated as:

    g̃=g_1+g_2-δ(g_1^Tg_2<0)g_1^Tg_2/‖ g_2 ‖^2g_2-δ(g_1^Tg_2<0)g_2^Tg_1/‖ g_1 ‖^2g_1.


















Obviously, in the optimization process of standard SGD algorithm, the overall gradient of the loss ℒ is g=g_1+g_2. In this paper, two cases 
are considered based on the correlation between gradients g_1 and g_2.

Case 1: When g_1 and g_2 are positively correlated or unrelated, i.e., cos(g_1,g_2)≥0 or g_1^Tg_2≥0, the angle between the two gradients is an acute or vertical angle. Then, we believe that there is no conflict between the two objectives in training, that is, two gradients are coordinated or balanced. Thus g_1 and g_2 do not need to be harmonized.



Case 2: When g_1 and g_2 are negatively correlated, i.e., cos(g_1,g_2)<0 or g_1^Tg_2<0, the angle between the two gradients is an obtuse angle. Then, we believe that there is a gradient optimization conflict between the two objectives. In this case, g_1 and g_2 need to be harmonized. For the overall loss function ℒ (Θ)=ℒ_1(Θ)+ℒ_2(Θ) of an arbitrary model, the proposed UDA+GH is to find a gradient g̃ that satisfies the following problem:



















    g̃min  ℒ_1(Θ-g̃)+ℒ_2(Θ-g̃),

where Eq. (<ref>) can be solved by aforementioned three lemmas. Firstly, the harmonic gradient g̃_1 of g_1 can be obtained by Lemma 1. Then the harmonic gradient g̃_2 of g_2 can be obtained by Lemma 2.  Finally, as can be seen from Lemma 3, it is easy to acquire the whole aggregated gradient g in Eq. (<ref>) after harmonization and the gradient aggregation process is visualized in Fig. <ref>.











 §.§ Essence and Insights of GH


To understand the essence of GH more intuitively, we demonstrate the proposed gradient harmonization by backward inference in this section. Fig. <ref> (a) and (b) show the essential analysis by performing gradient harmonization on original gradients g_1 and g_2, respectively. Note that we mainly consider the case where gradient conflict exists, that is, the angle between g_1 and g_2 is an obtuse angle.

First, we analyze the essence of gradient harmonization for g_1 (i.e., Fig. <ref> (a)). θ denotes the angle between original gradients g_1 and g_2. n is the projection of g_1 on the reversely extended line of g_2. m is perpendicular to g_2. Note that each letter except θ on the figure represents a vector. Naturally, m can be derived as Eq. (<ref>), from which we can see that the expression of the vector m is the same as Eq. (<ref>) when g_1^Tg_2<0. That is, the vector m is the result of gradient harmonization of g_1 (i.e., g̃_1).

    m=g_1-n=   g_1-|g_1|· cos(π-θ)· (-g_2/|g_2|)
    
        =   g_1-|g_1|· cosθ·g_2/|g_2|
    
        =   g_1-|g_1|·⟨ g_1,g_2⟩/|g_1|· |g_2|·g_2/|g_2|
    
        =   g_1-⟨ g_1,g_2⟩/g_2^2g_2.

where ⟨·⟩ is the inner product operator.



Second, we analyze the essence of gradient harmonization for g_2 (i.e., Fig. <ref> (b)). Similar to (a), the vector e is the projection of g_2 on the reversely extended line of g_1. The vector f is perpendicular to g_1. Through the derivation of Eq. (<ref>), it can be found that the expression of the vector f is the same as Eq. (<ref>) when g_1^Tg_2<0. That is, the vector f is the result of harmonizing the gradient g_2 (i.e., g̃_2).

    f=g_2-e=   g_2-|g_2|· cos(π-θ)· (-g_1/|g_1|)
    
        =   g_2-|g_2|· cosθ·g_1/|g_1|
    
        =   g_2-|g_2|·⟨ g_1,g_2⟩/|g_1|· |g_2|·g_1/|g_1|
    
        =   g_2-⟨ g_1,g_2⟩/g_1^2g_1.




From the above derivations, the nature of GH is summarized as follows. 1) The harmonic gradient g̃_1 is essentially the projection of the raw gradient g_1 on the vertical line of g_2. 2) The harmonic gradient g̃_2 is essentially the projection of the raw gradient g_2 on the vertical line of g_1.


Therefore, with the above nature, the following three observations can be obtained. 1) The proposed gradient harmonization method not only ensures that the harmonic gradient (g̃_1/g̃_2) is close to the original gradient (g_1/g_2), but also ensures that the aggregated gradients before and after applying GH are close to each other. 2) If the angle between the original gradients g_1 and g_2 is θ, then the angle between harmonized gradient g̃_1 and g̃_2 becomes π-θ after applying GH. That is, GH really realizes the transformation from obtuse angle to acute angle. 3) The proposed GH aims to move the harmonic gradients towards the direction favorable for optimization, rather than those arbitrary gradient directions with acute angles.







 §.§ Improved Version: GH++

In this section, we propose an improved version called GH++, which aims to adjust the gradient angle between the two tasks from an obtuse angle to a vertical angle, i.e., making them orthogonal. This is to eliminate the conflict but simultaneously relieve the gradient deviation. Fig. <ref> visually illustrates the concepts of the proposed GH and GH++, with a primary focus on scenarios where gradient conflict exists, i.e., when the angle between g_1 and g_2 is obtuse.


Gradient deviation. For clarity, we designate the original gradients as g_1=OA and g_2=OB. Let θ represent the angle between g_1 and g_2. As illustrated in Fig. <ref> (a), we denote the gradients after applying GH as g̃_1=OC and g̃_2=OD, where OC⊥ OB and OD⊥ OA. Apparently, the sum of the angles deviated from the original directions is 2(θ - π/2), i.e., ∠ AOC+∠ DOB. In other words, although GH can promote the positive correlation between the two gradients, its optimization gradient is seriously deviated from the original gradient. Therefore, we propose an improved version, GH++, to eliminate the conflict and minimize the sum of the gradient deviations. As shown in Fig. <ref> (b), GH++ adjusts the gradient angle between the two tasks from an obtuse angle to a vertical angle, i.e., ∠ EOF. The sum of the gradient deviation angles is (θ - π/2) i.e., ∠ AOE+∠ FOB, which is half of the sum of the gradient deviations of GH.





Specifically, as shown in Fig. <ref> (b), let denote the harmonized gradients of GH++ as g̃_1=OE and g̃_2=OF.  In order to resolve the conflict and relieve the deviation from the original gradient directions, we designate OE⊥ OF,  i.e., ∠ EOF = π/2, where points E and F  move along arcs AC and BD, respectively. Intuitively, △ EOF forms a rotatable right triangle. We define the direction of rotation from g_1 to g̃_1 as positive. β represents the rotating angles from g_1 to g̃_1, which is a positive angle and β̅ represents the rotating angles from g_2 to g̃_2, which is a negative angle. According to Fig. <ref> (b), the harmonization gradients of GH++, i.e., g̃_1 and g̃_2, can be represented as



    g̃_1 = OE = OA+AE= g_1+ AE,
    g̃_2 = OF = OB+BF= g_2+ BF.

Since the △ AOE and △ BOF are isosceles triangles, we can obtain AE and BF according to the trigonometric formula, whose expression can be written as

    AE = 2 ·OA·sin∠ AOE/2= 2 · g_1 ·sinβ/2, 
    BF = 2 ·OB·sin∠ BOF/2= 2 · g_2 ·sinβ̅/2.

Substituting Eq. (<ref>) into Eq. (<ref>), we can derive the harmonized gradients of GH++.

    g̃_1 = (1+2 ·sinβ/2)· g_1,
    g̃_2 = (1+2 ·sinβ̅/2)· g_2.


Referring to Figure <ref> (b), we have a variable angle β = ∠ AOE ∈ [0, θ - π/2], where θ = ∠ AOB = arccosg_1^Tg_2/|g_1|·|g_2|∈ (π/2, π]. To simplify, we can express β as follows:

    β = λ∠ AOC = λ (θ-π/2) = λ(arccosg_1^Tg_2/|g_1|·|g_2|-π/2),

where λ∈ [0, 1] serves as a trade-off parameter, controlling the magnitude of deviation from g̃_1 to g_1. When λ = 0, g̃_1 = g_1= OA remains unchanged, and g_2 = OD. When λ = 1, g_1= OC, and g̃_1 = g_2 = OB remains constant. Since OE⊥ OF,  i.e., ∠ EOF = π/2, β̅ can be represented as

    β̅    =∠ BOF = - ∠ FOB = -(θ-π/2 - β) = β-(θ-π/2).

Substituting Eq. (<ref>) into Eq. (<ref>), we have

    β̅    =  λ (θ-π/2) + (θ - π/2)= (λ-1)(θ - π/2)
       =(λ-1)(arccosg_1^Tg_2/|g_1|·|g_2|-π/2).

Finally, substituting Eq. (<ref>) and Eq. (<ref>) into Eq. (<ref>), we can obtain the harmonization gradients of GH++ as follows:

    g̃_1 = (1+2 ·sinλ(arccosg_1^Tg_2/|g_1|·|g_2|-π/2)/2)· g_1,
       g̃_2 = (1+2 ·sin(λ-1)(arccosg_1^Tg_2/|g_1|·|g_2|-π/2)/2)· g_2.


Similar to GH, when there is no conflict, i.e., g_1^Tg_2≥ 0, no action. The proposed GH++ is indicated as Theorem 2.





For any two different tasks, optimization conflicts can be eliminated by GH++. Define the overall loss function ℒ (Θ)=ℒ_1 (Θ)+ℒ_2 (Θ), composed of two sub-objectives (refer to domain alignment and classification in this paper). Define g_1 and g_2 as the gradient of ℒ_1 (Θ) and ℒ_2 (Θ), respectively, then the general expression of the overall harmonized gradient g̃ of the whole loss ℒ (Θ) with GH++ can be formulated as:

    g̃   =g̃_1 + g̃_2 
       =(1+2 δ(g_1^Tg_2<0) sinλ(arccosg_1^Tg_2/|g_1|·|g_2|-π/2)/2) g_1
       +(1+2 δ(g_1^Tg_2<0) sin(λ-1)(arccosg_1^Tg_2/|g_1|·|g_2|-π/2)/2) g_2.




It is worth noting that GH++ presents a more favorable optimization scenario for the original task compared to GH. This advantage arises from the fact that the sum of the gradient deviation from the original direction of GH++ is only half of GH.



































 §.§ Equivalent Model of UDA with GH/GH++



By combining the general UDA model (Eq. (<ref>), described in Section <ref>) with gradient harmonization strategies (Eq. (<ref>)/Eq. (<ref>), a well-balanced UDA model can be trained and implemented. However, the gradient aggregation operator in Eq. (<ref>)/Eq. (<ref>) is intricate and has an impact on optimization efficiency. Therefore, we introduce a computation-efficient alternative model that is functionally equivalent to UDA with GH/GH++. For convenience, we can express the proposed gradient harmonization approaches, i.e., Eq. (<ref>)/Eq. (<ref>), as follows:

    g̃=τ_1g_1+τ_2g_2,

where τ_1 and τ_2 are constants that can be calculated by using the original gradients g_1 and g_2. Notably, the gradient g_1 of the original loss ℒ_1(Θ) and the gradient g_2 of the original loss ℒ_2(Θ) can be easily computed, only if the UDA model is fixed.  Then, if GH is chosen to resolve the conflict, τ_1 and τ_2 can be calculated as follows.

    τ_1=1-δ(g_1^Tg_2<0)g_2^Tg_1/‖ g_1 ‖^2,  
       τ_2=1-δ(g_1^Tg_2<0)g_1^Tg_2/‖ g_2 ‖^2.

If GH++ is used to eliminate the conflict, τ_1 and τ_2 can be calculated as follows.

    τ_1=(1+2 δ(g_1^Tg_2<0) sinλ(arccosg_1^Tg_2/|g_1|·|g_2|-π/2)/2),
       τ_1=(1+2 δ(g_1^Tg_2<0) sin(λ-1)(arccosg_1^Tg_2/|g_1|·|g_2|-π/2)/2).

Ultimately, we can derive the equivalent UDA model embedded GH/GH++ by conducting the integral operation on the gradient in Eq. (<ref>). The overall loss of UDA with GH/GH++ model can be represented as



    ℒ̃=∫(τ_1g_1+τ_2g_2)d Θ=τ_1ℒ_1(Θ)+τ_2ℒ_2(Θ),

where ℒ_1(Θ) and ℒ_2(Θ) can be the domain alignment loss and classification loss described in Eq. (<ref>) and Eq. (<ref>), respectively, for unsupervised domain adaptation.
Finally, the objective function of the UDA with GH/GH++ is:

    θ_g,θ_cmin  θ_dmax  ℒ̃=τ_1 ℒ_dom(θ_g,θ_d)+τ_2 ℒ_cls(θ_g,θ_c).



It can be seen that the proposed approaches essentially reweight the original loss function ℒ_1(Θ) and ℒ_2(Θ). The computation of the weights τ_1 and τ_2 needs to use the gradient of the original loss function, as presented in  Eq. (<ref>)/Eq. (<ref>). Finally, a balanced and efficient UDA model is formulated. Without loss of generality, when g_1^Tg_2>0, we can observe τ_1 =1 and τ_2 =1, and the overall loss function is degenerated into ℒ̃=ℒ_1(Θ)+ℒ_2(Θ), i.e., the general UDA model. We describe the training procedure in Algorithm <ref>.

Notably, the proposed approaches are  plug-and-play modules and can be deployed in almost all models with different losses ℒ_1(Θ) and ℒ_2(Θ). In order to clearly observe the computational cost of the equivalent model, we discuss the training speed and recognition accuracy before and after applying GH in Table <ref>.





§ EXPERIMENTS




 §.§ Datasets






Office-31 <cit.> is a mainstream benchmark dataset for visual domain adaptation, which consists of three distinct domains: Amazon (A), DSLR (D), Webcam (W). It totally contains 4,652 images from 31 categories. We evaluate our method in all 6 different transfer tasks across domains.


Office-Home <cit.> is a more challenging and harder benchmark than Office-31. It contains 15.5K images across 65 object categories from 4 different domains: Artistic images (Ar), Clip Art (Cl), Product images (Pr), and Real-World images (Rw). We evaluate our method in all 12 different transfer tasks across domains.

Digits Datasets. We mainly study three datasets: MNIST (M) <cit.>, USPS (U) <cit.> and SVHN (S) <cit.>. MNIST and USPS are two general handwriting recognition datasets involving 10 categories. SVHN is obtained from house numbers in Google Street View images. We conduct experiments in 3 universal tasks, including M→U, U→M and S→M.

VisDA-2017 <cit.> is a simulation-to-real dataset for domain adaptation with two distinct domains: synthetic object images rendered from 3D models and real object images. It contains over 280K images across 12 categories.



DomainNet <cit.> is the largest domain adaptation dataset consisting of about 600K images from six distinct domains, including Clipart (clp), Infograph (inf), Painting (pnt), Quickdraw (qdr), Real (rel), Sketch (skt). There are 48K-172K images categorized into 345 classes per domain. We evaluate our method in all 30 cross-domain tasks.








 §.§ Implementation Details

We compare our method with the following state-of-the-art unsupervised domain adaptation methods: ResNet <cit.>, DAN (Deep Adaptation Networks) <cit.>, DANN (Domain-adversarial Neural Networks) <cit.>, JAN (Joint Adaptation Networks) <cit.>, DRCN (Deep Reconstruction-Classification Networks) <cit.>, CoGAN (Coupled Generative Adversarial Networks) <cit.>, ADDA (Adversarial Discriminative Domain Adaptation) <cit.>, CyCADA ( Cycle-consistent Adversarial Domain Adaptation) <cit.>, CAT (Cluster Alignment with a Teacher) <cit.>, TPN (Transferrable prototypical networks) <cit.>, LWC (Light-weight Calibrator) <cit.>, ETD (Enhanced Transport Distance) <cit.>, TAT (Transferable Adversarial Training) <cit.>, TADA (Transferable Attention for Domain Adaptation) <cit.>, SymNets (Symmetric Networks) <cit.>, BNM (Batch Nuclear-norm Maximization) <cit.>, ALDA (Adversarial-learned Loss for Domain Adaptation) <cit.>, MIMTFL (Mutual Information Maximisation
and Transferable Feature Learning) <cit.>, Meta (Metaalign) <cit.>, MDD (Margin Disparity Discrepancy) <cit.>, CDAN <cit.>, MCD (Maximum Classification Discrepancy) <cit.>, GVB (Gradually Vanishing Bridge) <cit.> and DWL (Dynamic Weighted Learning) <cit.>, SAFN (Stepwise Adaptive Feature Norm) <cit.>, SWD (Sliced Wasserstein Discrepancy) <cit.>, DMRL (Dual Mixup Regularized Learning) <cit.>, CAN (Contrastive Adaptation Network) <cit.>, FixBi <cit.>, FGDA (Feature Gradient Distribution Alignment) <cit.> and ParetoDA (Pareto Domain Adaptation) <cit.>, CGDM (Cross-domain Gradient Discrepancy Minimization) <cit.>, SSRT (Safe Self-Refinement for Transformer-based domain adaptation) <cit.> and TVT (Transferable Vision Transformer) <cit.>.  SSRT and TVT are transformer-based approaches, while other approaches are CNN-based. We use the results in their original papers for fair comparison. When there is no relevant experimental result in the original papers, we implement them according to their released source codes.




 
Furthermore, we investigate the effect of GH/GH++ on five popular UDA approaches, including four CNN-based approaches (i.e., CDAN <cit.>, MCD <cit.>, GVB <cit.> and DWL <cit.>) and one transformer-based approach (i.e., SSRT <cit.>). For CNN-based approaches, we use pre-trained ResNet-50 <cit.> on ImageNet-1K <cit.> as backbone network for Office-31 and Office-Home, and adopt the pre-trained ResNet-101 <cit.> as backbone network for VisDA-2017 and DomainNet. We employ LeNet <cit.> as backbone for Digits datasets. For transformer-based approach, SSRT uses the ViT-base with 16 × 16 patch size pre-trained on ImageNet-21K and ImageNet-1K (denoted by SSRT^∗) as the transformer backbones, respectively. In terms of optimization, we follow the original protocol of the baselines. Note that, in this paper, CDAN means CDAN+E in <cit.>) and GVB refers to GVB-GD in <cit.>. For GH++, we empirically set λ to 0.5.
MCD achieves domain alignment by the game between two classifiers and the generator. 
Therefore, the gradient of the discrepancy loss between the two classifiers is just the gradient of domain alignment loss. We implement our method in PyTorch by following UDA protocol <cit.> that source domain is labeled and target domain is unlabeled. 

For fair comparisons, we independently run all experiments five times and report the average target accuracy. 



 §.§ Results on UDA


























































Tables <ref>, <ref>, <ref>, <ref> and <ref> present evaluation results on Office-31, Office-Home, VisDA-2017, Digits and DomainNet, respectively. Generally, the transformer-based results, such as TVT and SSRT, are much better than CNN-based results, which has been validated in previous work. For TVT and SSRT, the ViT backbone pre-trained on ImageNet-1K is slightly inferior than ImageNet-21K (i.e., TVT^∗&TVT, SSRT^∗&SSRT). Considering that the CNNs are pre-trained on ImageNet-1K, we take SSRT^∗ for fair analysis as default in the following.  

In order to demonstrate the effectiveness and universality of the proposed approaches in UDA, we investigate the effect of GH/GH++ on five popular UDA approaches, including CNN-based approaches (i.e., CDAN, MCD, GVB and DWL) and transformer-based approach (i.e., SSRT). We can observe that the proposed approaches can further improve the baseline, whether based on transformer or CNN. This is attributed to the proposed approaches in alleviating gradient conflict between domain alignment task and classification task in optimization. Meanwhile, it verifies that the proposed approaches can be easily plugged and played in the existing UDA methods. In addition, GH++ usually performs better than GH. The reason is that during removing the gradient conflict between the two tasks, GH++ minimizes the gradient deflection, such that the two tasks can evolve harmoniously during joint training while preserving their individual task-specific optimality.

Specifically, as shown in Table <ref>, when GH/GH++ is applied, CDAN, MCD, GVB, DWL and SSRT outperform the original models by 1.5%/2.2%, 1.0%/1.4%,  0.2%/0.7%, 0.7%/0.8% and 0.7%/0.9% on Avg^ of Office-31, respectively.
As shown in Table <ref>, CDAN, MCD, GVB, DWL and SSRT after applying GH/GH++ outperform the original models by 2.6%/3.2%, 4.5%/5.1%, 0.8%/1.2%, 0.4%/0.6% and 0.6%/0.8% on Avg^ of Office-Home, respectively. As shown in Table <ref>, we can observe that CDAN, MCD, GVB, DWL and SSRT after applying GH/GH++ outperform the original models by 0.9%/0.9%, 2.2%/2.4%,  1.1%/1.3%, 0.6%/0.9% and 0.3%/0.3% on Avg of Digits, respectively.
As shown in Table <ref>, CDAN, MCD, GVB, DWL and SSRT after applying GH/GH++ outperform the original models by 1.4%/3.9%, 4.1%/4.5%, 0.5%/0.8%, 0.7%/0.3% and 0.7%/0.6% on VisDA-2017, respectively.
From the results in Table <ref>, we can observe that CDAN, MCD, GVB, DWL and SSRT after applying GH/GH++ outperform the original models by 3.2%/3.7%, 7.1%/7.3%, 0.7%/1.0%, 1.0%/1.3% and 1.1%/1.4% on Avg of DomainNet, respectively. The proposed approaches generally improve CDAN and MCD more than other baselines because the original models of CDAN and MCD have more obvious optimization conflicts than other baselines. Fig. <ref> also confirms this perspective. Therefore, the proposed approaches have better performance when the inherent gradient conflict is serious. Besides, it is worth noting that the proposed approaches usually achieve greater performance gains on large-scale datasets, i.e., VisDA-2017 and DomainNet, which further indicate the effectiveness and versatility of our method.

























































§ MODEL ANALYSIS AND DISCUSSION



 §.§ Feature Visualization
  Fig. <ref> describes the t-SNE <cit.> visualizations of features learned by MCD (baseline) and MCD+GH on the tasks of U → M and M → U. Fig. <ref> (a) and (c) are visualization features generated by MCD. Fig. <ref> (b) and (d) are visualization features generated by MCD+GH. It can be observed that both features learned by MCD and MCD+GH achieve well-performed global alignment effect with 10 clusters under two tasks. Further, the visualization feature distributions with GH deployed have better clustering effect and have fewer samples distributed across class boundaries, which intuitively boosts the feature discriminability. In addition, visualization results further validate the balance learning ability of our GH approach.



 §.§ Convergence Analysis
  We present the convergence curves of test error with respect to the number of iterations on tasks of U → M and Synthetic → Real as shown in Fig. <ref>. For each subfigure, the blue line represents the test error of different baselines, and the red line represents the test error for baseline+GH (e.g., CDAN+GH). Obviously, compared with baselines, the introduction of GH can further improve the test performance and convergence. This fully indicates that GH plays an active coordination role in practical optimization process, which promotes the cooperation of domain alignment task and classification task towards a benign direction as expected. 

















 §.§ Confusion Matrix Visualization


Fig. <ref> displays the visualizations of confusion matrix for the classifier trained by DWL and DWL+GH. DWL obtains several uncertain predictions with small values while DWL+GH obtains more confident predictions. Comparing with Fig. <ref> (a) and (b), the confusing “Class 1, 2, 7, 8, and 9” are correctly recognized in DWL+GH. From Fig. <ref> (c) and (d), the confusing “Class 8” is corrected in DWL+GH. It implies that optimization conflict during training deteriorates the discriminability in the target domain. Our method preserves the feature discriminability of target samples in the course of harmonic training. DWL+GH improves the between-task balance by coordinating the explicit external task conflict and the implicit internal optimization conflict.





 §.§ Balance Analysis of GH-based UDA
  Fig. <ref> shows MMD distance <cit.> and the max J(W) <cit.> values based on the feature representation learned by MCD+GH. The left vertical axis corresponding to the red curve represents the MMD distance used to measure the alignment degree across domains. The smaller value means the better domain alignment. The right vertical axis corresponding to the blue curve represents the degree of discriminability based on Linear Discriminant Analysis (LDA) <cit.>. The larger max J(W) implies better discriminability, which directly affects the classification result of samples. Fig. <ref>(a) and (b) reflect the degree of domain alignment and discriminability on task of M→U and U→M, respectively. From Fig. <ref>, MCD+GH has a smaller MMD distance and a larger max J(W) value during training. This shows that GH facilitates the coordinated optimality of alignment and classification tasks while maintaining their task-specific optimality. Enhanced and balanced domain-invariant and class-discriminative feature representations can be obtained.



 §.§ Training Speed and Accuracy
 In order to observe the efficiency of the proposed equivalent model more clearly, we present the training speed and classification accuracy before and after applying GH. As shown in Table <ref>, for tasks M→U and U→M on digits datasets, the training speed of MCD+GH is 0.14254 s/epoch and 0.09809 s/epoch longer than MCD, respectively. In other words, MCD+GH takes less training time than MCD in training process, but can get 2.5%/2.7% classification accuracy gain. The computational cost of employing GH is quite low, and thus GH is a powerful and efficient auxiliary tool to facilitate those popular domain adaptation baselines towards more outstanding classification performance.







 §.§ Gradient Inner Product Visualization
  Fig. <ref> presents inner product distributions of two gradients between domain alignment and classification tasks before and after applying Gradient Harmonization for MCD. From Fig. <ref> (a), we observe the acute and obtuse angles between gradients of the two tasks before coordination. The obtuse angles account for about 40% of the total number, which exhibits the identical property as Fig. <ref>. In other words, there exists between-task conflict in the model optimization process but paid less attention. After Gradient Harmonization, as shown in Fig. <ref> (b), the inner products of the two gradients are all positive. That is, the gradient angles between the two tasks are coordinated into acute angles. The proposed GH avoids the optimization conflict by separately adjusting the gradients of the two tasks to achieve the purpose of optimal coordination. Experimental results fully illustrate the effectiveness of the proposed GH.







 §.§ Rationality and Comparison to Other Alternatives



The proposed GH/GH++ aims at altering the gradient angle between two different tasks from an obtuse angle to an acute/vertical angle. 
Whether the same effect can be achieved with other intuitive gradient correction alternatives or not? For instance, one might consider convert an obtuse angle into an acute angle by simply applying a negative sign to one of the two gradients, i.e., transforming either g_1 into -g_1 or g_2 into -g_2. As depicted in Fig. <ref>, it is apparent that the results reveal a significant decline in classification accuracy when using the -g_1&g_2 and g_1&-g_2 methods in contrast to the proposed approaches. While applying a negative sign can indeed transit from an obtuse angle to an acute angle and force the two gradients to be positively correlated, it cannot replicate the excellent classification performance achieved by the proposed GH/GH++.

The rationale behind this observation is that directly applying a negative sign to one of the gradients alters the gradient towards the opposite direction of the original optimal gradient descent direction. This significantly undermines the primary objective of the original task. In other words, the performance deteriorates sharply when the gradient harmonization direction significantly deviates from the original direction. GH/GH++ reasonably adjusts the two gradients to achieve the purpose of optimization coordination on the premise of keeping the harmonic gradient as close as possible to the original gradient descent direction. Therefore, the proposed approaches do not break the task-specific optimality, but pursues a better cooperation, which verifies the rationality of the proposed GH/GH++.








 §.§ Parameter Sensitivity Analysis

To investigate the effect of the parameter λ in GH++, we conduct experiments on three tasks (i.e., M→U, U→M and S→M) based on MCD and GVB by varying λ∈{0, 0.1, 0.3, 0.5,0.7, 0.9, 1}. The results are presented in Fig. <ref>. We can observe GH++ is little sensitive to the scale variety of λ, which indicates that GH++ is robust across different baselines and tasks. Besides, we observe that when λ=0.5, models generally achieve the best performance. In other words, when the gradient deviation of the two tasks is relieved, the performance can be largely improved.
























 §.§ Scalability to Other Multi-task Problems

To further demonstrate the universality and scalability of the proposed approaches, we evaluate GH/GH++ in the object detection and multi-modal interactive retrieval fields, which also involves optimization of multiple objectives.



Dataset. We select widely used benchmarks, i.e., PASCAL VOC 2007 <cit.> and CSS <cit.> for object detection and multi-modal interactive retrieval, respectively.
PASCAL VOC 2007 consists of about 5K trainval images and 5K test images over 20 object categories.
CSS consists of 38K synthesized images with different colors, shapes and sizes. It contains about 19K training image-text pairs and 19K testing image-text pairs, respectively. 



Implementation Details. For object detection, we use the DSSD <cit.> model with 321× 321 inputs as the baseline. We follow the same experimental settings and protocol as <cit.> and adopt Average Precision (averaged AP at IoUs from 0.5 to 0.9) to measure performance. For multi-modal interactive retrieval, we use TIRG <cit.> and the adversarial training version(AT-TIRG) follows the paper <cit.> as the baseline. For evaluation, we follow the same protocols as <cit.> and use retrieval accuracy R@N as our evaluation metric, which computes the percentage of test queries where at least one target or correctly matched image is within the top N retrieved images. Note that we exactly follow the original experimental setups of the baselines. In other words, DSSD+GH/GH++ and AT-TIRG+GH/GH++ indicate that GH/GH++ are inserted directly into the DSSD and AT-TIRG, respectively, without any change in either the backbone or hyper-parameters.





Results on Object Detection. Object detection models <cit.> usually consider two tasks: classification and regression. In this section, we use the proposed GH/GH++ to mitigate the conflict between these two tasks. Table <ref> presents the Average Precision (AP) on the PASCAL VOC 2007 test set. We can observe that GH/GH++ can improve the AP with different bounding boxes, which verifies the effectiveness of the proposed approaches in object detection.








Results on Multi-modal Interactive Retrieval. Recently, <cit.> introduces adversarial training <cit.> into multi-modal interactive retrieval model (abbreviated as AT-TIRG) to improve model robustness, which aims to train a generalized and robust model suitable for both clean samples and adversarial attack samples. During adversarial training, there is conflict in training the clean samples and adversarial attack samples, which results in the limitation of the final performance. In this section, we try to adopt the proposed approaches to solve this problem. The experimental results are provided in Table <ref>. Comparing AT-TIRG+GH with AT-TIRG, we can see the proposed GH can enhance the AT. For the adversarial attack samples, AT-TIRG+GH outperforms the AT-TIRG by 3.4%, 2.9%, and 2.2% on R@1, R@5, and R@10, respectively. AT-TIRG+GH++ outperforms the AT-TIRG by 5.2%, 4.9%, and 3.9% on R@1, R@5, and R@10, respectively. The performance of the clean samples is also improved in most cases. Note that the performance of clean samples is usually close to 100%, so there is little space for improvement.





































§ CONCLUSION

In this paper, we pay attention to the optimization conflict (i.e., imbalance or incoordination) problem between different tasks (i.e., the alignment task and the classification task) in alignment-based unsupervised domain adaptation models. To mitigate this problem, we propose two simple yet efficient Gradient Harmonization approaches, including GH and GH++, which take measures to de-conflict between the gradients of both tasks in optimization. Besides, to facilitate the harmonization during adaptation, we derive the equivalent but more efficient model of UDA with GH/GH++, which becomes a dynamically reweighted loss function of most existing unsupervised domain adaptation models. Further, the essence and insights of the proposed approaches are provided to indicate its rationality. Exhaustive experiments and model analyses demonstrate that the proposed approaches significantly improve the existing UDA models and contribute to achieving state-of-the-art results.
In addition, we have verified that the proposed approaches can be adapted to other problems and areas, such as object detection and multi-modal retrieval, to de-conflict between the gradients of any two tasks in optimization and improve model performance.




























 

  


































IEEEtrans






























































 Paper Title: Universal Approximation Theory: Foundations for Parallelism in Neural Networks.



Turbulence, Waves, and Taylor's Hypothesis for Heliosheath Observations
    M. Nakanotani
    
=======================================================================



§ ABSTRACT 

Neural networks are increasingly evolving towards training large models with big data, a method that has demonstrated superior performance across many tasks. However, this approach introduces an urgent problem: current deep learning models are predominantly serial, meaning that as the number of network layers increases, so do the training and inference times. This is unacceptable if deep learning is to continue advancing. Therefore, this paper proposes a deep learning parallelization strategy based on the Universal Approximation Theorem (UAT). From this foundation, we designed a parallel network called Para-Former to test our theory. Unlike traditional serial models, the inference time of Para-Former does not increase with the number of layers, significantly accelerating the inference speed of multi-layer networks. Experimental results validate the effectiveness of this network.






§ INTRODUCTION

In today's era, deep learning models have undoubtedly taken the lead in the fields of computer vision (CV) and natural language processing (NLP), demonstrating exceptional performance across numerous tasks. In CV, models such as ResNet<cit.> and ViT<cit.> for image recognition, U-Net <cit.> for image segmentation, GAN <cit.> for image generation, and diffusion models like SORA <cit.> for video generation highlight the extensive application of deep learning in CV. These models are primarily based on convolutional neural networks (CNNs), Transformers, or a combination of both. Simultaneously, large language models (LLMs) based on Transformers have become dominant in natural language processing. Overall, deep learning models have shown vast application potential and powerful capabilities so far.

However, this progress has brought forth some challenges: the development trend in deep learning is moving towards larger models and larger datasets. This trend is fundamentally driven by the UAT<cit.>. From the papers UAT2LLM<cit.> and UAT2CV<cit.>, we know that residual-based CNNs and Transformer-based models are specific implementations of UAT, both exhibiting the characteristic of dynamically approximating functions based on inputs. Larger models and datasets are indeed requirements of the UAT theory: deeper network layers can more precisely approximate target functions dynamically based on inputs, while larger and more diverse datasets help ensure the training data closely represents real-world conditions, making the model's approximated functions more reflective of real-world complexities.

However, this trend inevitably leads to increased demands for computational resources and extended training times. As more powerful models are needed, the UAT theory mathematically necessitates larger networks, which in turn require more computational resources. While technological advancements, especially in hardware performance, provide potential ways to alleviate this issue, the increase in time costs remains a fundamental obstacle to scaling serial networks. The improvement in computational resources such as GPUs and TPUs is limited by the speed of electrical signal transmission, whereas the theoretical scale of networks has no such upper limit. When networks reach tens or even hundreds of thousands of layers, inference times can extend to several minutes or even hours for each batch. Current solutions include model optimization techniques like quantization and pruning <cit.>, hardware acceleration with specialized accelerators and edge computing <cit.>, and parallel and distributed computing <cit.> with model parallelization, data parallelization, and distributed computing frameworks. However, these do not fundamentally address the inference delay caused by the increasing number of serial network layers.

We believe the root cause lies in the serial computation mechanism of existing deep learning models, where the output of one layer serves as the input for the next. This design stems from the early development of deep learning networks in the field of CV, where image data exhibit local correlations and spatiotemporal invariance. To accommodate these features, early CV networks primarily employed convolutional layers with small receptive fields. To achieve a larger receptive field, deeper network layers were necessary, leading to early deep-learning architectures like AlexNet and VGG. The introduction of ResNet brought a paradigm shift, as residual networks have the ability to dynamically fit functions based on input, influencing almost all subsequent network designs and solidifying the tradition of multi-layer stacking.



While this design was manageable with fewer layers, the increasing depth and number of parameters in modern networks have created a critical issue: slow computation efficiency due to the need for sequential layer-by-layer calculations during inference. This computation method inevitably leads to delays, which become more severe as the network depth increases, significantly constraining the progress of CV and NLP technologies.

Therefore, it is imperative to explore parallelization techniques for deep learning networks, enabling parallel computation across multiple layers. This approach is key to accelerating the inference process, breaking the time bottleneck, and advancing CV and NLP technologies towards greater efficiency and speed. To develop such parallel computation, we must start from the fundamental theory of deep learning: the UAT. The network must be designed in accordance with UAT principles and must satisfy the requirements outlined in the papers UAT2LLM and UAT2CV, which emphasize the ability to dynamically fit functions based on data. With this in mind and leveraging existing network structures, we designed a parallel network, Para-Former, to validate the feasibility of our theory. Our contributions include:






  * Proposing the theory of parallel deep learning networks and experimentally validating its feasibility.

  * Analyzing the characteristics, advantages, and disadvantages of both serial and parallel networks.

  * Investigating the issues present in deep learning from a data perspective.

 


The structure of this paper is as follows: In Section <ref>, we first present the general form of the Universal Approximation Theorem (UAT). Based on the mathematical form of UAT, we define the parallel network Para-Former in Section <ref>. In Section <ref>, we validate the effectiveness of Para-Former in various ways: we assess its feasibility in Section <ref>, explore the impact of network depth in Section <ref>, and analyze the influence of different datasets in Section <ref>.





§ UAT



Deep learning networks have demonstrated their powerful capabilities in numerous CV and NLP tasks. In the papers UAT2LLMs and UAT2CVs, it has been proven that both Transformer-based networks and residual-based CNNs are specific implementations of the UAT. Therefore, we consider UAT to be the fundamental theory for deep learning. To design parallelized networks, we must start from the principles of UAT. Before formally designing parallel networks, we will first review the basic form of the UAT theory and then propose a network parallelization solution based on this mathematical theory. The UAT was initially proposed by  <cit.>, and its basic form is as follows:





    G(𝐱)=∑_j=1^N α_j σ(𝐖_j^T𝐱+θ_j)

is dense in C(𝐈_n). Here, 𝐖_j ∈ℝ^n and α_j, θ∈ℝ are parameters σ is a sigmoid function. For any f ∈ C(𝐈_n) and ε>0, there exists a function G(𝐱):


    |G(𝐱)-f(𝐱)|<ε   for all   𝐱∈𝐈_n .



This indicates that when N is sufficiently large, a neural network can approximate any continuous function on a closed interval.   <cit.> further proved that this function can approximate any Borel measurable function. Observing Equation (<ref>), where the function G(𝐱) produces scalar outputs in the real number field ℝ, the scenario naturally extends when G(𝐱) maps to ℝ^m, requiring approximation in each dimension. To accommodate such multi-dimensional outputs, we simply need to extend Equation (<ref>): the transformation matrix 𝐖_j is revised to the space ℝ^n × m, the bias term θ_j is recast as a vector in ℝ^m, and α_j is reshaped into a matrix. We call the parameters 𝐖_j and α_j in UAT as weight and θ_j as bias.








§ UAT FOR PARALLEL NETWORK




 §.§ Pre-request for Basic Block of Parallel Network


From Equation (<ref>), we can observe a significant characteristic: if we define different indices j as the indices of layers and the corresponding layer's mathematical form is α_j σ(𝐖_j^T𝐱+θ_j), it is evident that the parameters between different layers do not interact with each other. Therefore, based on (<ref>), we can design networks where parameters across multiple layers are independent, achieving parallel computation, like Figure <ref>. Figure <ref>.a represents a typical sequential network, which clearly requires computation to proceed in order from front to back. In contrast, Figure <ref>.b depicts a parallel network, where there is no parameter interaction between different blocks, allowing them to be considered independent and thus suitable for parallel computation. The blocks in this network should satisfy the formula given in Figure <ref>.b. Clearly, under this condition, the sum of the multiple blocks' outputs adheres to the UAT theory.

If we design the parallel network solely based on Eq. (<ref>), the network will face a limitation once training is complete: the function that can be approximated by the network is fixed, as the parameters in the corresponding UAT are fixed. So to ensure the network can dynamically approximate functions based on input, some or all of the parameters α_j, W_j and θ_j must be influenced by the input (referencing to UAT2LLM and UAT2CV), which can be simply written as α_j(𝐱), W_j(𝐱)  and θ_j(𝐱). 

This allows the network to dynamically adapt and fit the function based on the input while enabling parallel computation. So we could design the basic block in Figure <ref>, like Block_1⋯Block_n are the implementation of α_j(𝐱)) σ(𝐖_j^T(𝐱) 𝐱+θ_j(𝐱)), i=1⋯ n. Some or all the parameters α_j, W_j  and θ_j will change according to input. 






 §.§ Para-former



Based on the previous request on the basic module, we designed the Para-Former which is a parallel network by using the Transformer module as the basic module. We chose Transformer for this design because it is a well-established and widely-used network architecture, excelling in both CV and NLP. Another reason is that its basic mathematical form aligns with the requirements of Section <ref>.

Before formally introducing the structure of Para-Former, let's briefly review the Transformer structure. The basic structure of a Transformer is shown in Figure <ref>, and for convenience, we refer to it as a Transformer module. Typically, it is represented as follows:


    𝐇=   MultiHead(𝐐, 𝐊, 𝐕) 
    
    =   Concat(𝐇̂_1, …, 𝐇̂_h) 𝐖_O



    𝐇̂_i   =Attention(𝐱_i 𝐖_Qi=𝐐_i, 𝐱_i 𝐖_Ki=𝐊_i,
       𝐱_i𝐖_Vi=𝐕_i)
       =softmax(𝐐_i 𝐊_i^T/√(M)) 𝐕_i
       =𝐇_i𝐕_i =𝐇_i[𝐱_i𝐖_Vi]




    FFN(𝐱_i) = 𝐖_i,3σ (𝐖_i,2𝐱_i+𝐛_i,2)+𝐛_i,3


The above representation mainly describes the engineering process rather than its true mathematical form. According to UAT2LLM, the Multi-Head Attention (MHA) and Feed-Forward Network (FFN) components of Transformers can be expressed in matrix-vector form. We denote the matrix-vector form of variables with a prime symbol (e.g., 𝐱_i becomes 𝐱'_i). For more details, refer to UAT2LLM.


    MHA(𝐱'_i) = 𝐖'_i,1𝐱'_i



    FFN(𝐱'_i) = 𝐖'_i,3σ (𝐖'_i,2𝐱'_i+𝐛'_i,2)+𝐛'_i,3


The parameter 𝐖'_i,1 is influenced by the input. As shown in Figure <ref>, the output of a single Transformer layer, 𝐱'_i+1, can be described by Eq. (<ref>) (Details could be found in Section <ref>):



    𝐱'_i+1   = 𝐖'_i,1𝐱'_i + 𝐖'_i,3σ (𝐖̂'̂_̂î,̂1̂𝐱'_i + 𝐛'_i,2) + 𝐛'_i,3



The mathematical form in Eq. (<ref>) meets the requirements of a basic module in UAT. Therefore, we can use the structure in Figure <ref> as the basic module in Para-Former, corresponding to Figure <ref>.a. Besides, we add a symbol of a hat on some parameters to represent the parameters will be affected by the input, like 𝐖̂'̂_̂î,̂1̂. 




    𝐱'_i+2   = 𝐖̂'̂_̂î+̂1̂,̂2̂𝐱'_i + 𝐖̂'̂_̂î+̂1̂,̂3̂σ (𝐖̂'̂_̂î,̂1̂𝐱'_i + 𝐛'_i,2) 
       + 𝐖'_i+1,3σ{𝐖̂'̂_̂î+̂1̂,̂4̂𝐱'_i + 𝐛̂'̂_̂î+̂1̂,̂4̂} + 𝐛̂'̂_̂î+̂1̂,̂5̂


The output of a two-layer Transformer is shown in Eq. (<ref>) (Details could be found in Section <ref>). From Eq. (<ref>), it is evident that a serial configuration of multi-layer Transformers also conforms to the mathematical form of UAT. Therefore, we can extend the Para-Former structure. That is, within one block in Figure <ref>, there can be multiple serially connected Transformer modules, while multiple blocks operate in parallel, as shown in Figure <ref>.b. We refer to the number of Transformer modules within a block as the depth and the total number of Transformer modules across all blocks as the network layer count.





 §.§ Speed-up Ratio


We have proposed a network parallelization scheme, and in this section, we will compare the running speed of this parallelized network with that of traditional serial networks. It is important to note that due to differences in hardware and implementation methods, providing an objective comparison through experiments can be challenging. Therefore, we will focus on a theoretical analysis. Based on the network structure shown in Figure <ref>, the inference time of Para-Former is only affected by its network depth. 

Assume we use a Para-Former with a depth of M and a layer count of L, compared to a traditional Transformer network with a depth of N layers, where L >> N. In this case, the running speed of the parallel network would be N/M times faster than that of the serial network. This is because the speed of the parallel network is solely related to the depth of Para-Former. As the network depth increases, the acceleration effect becomes more pronounced.



§ EXPERIMENTS


Our experimental approach is to first verify the feasibility of the parallel network theory and its adherence to the UAT approximation properties using several commonly used public image classification datasets (Section <ref>). We then explore the impact of network depth on the UAT theory (Section <ref>), followed by an investigation into the influence of the dataset (Section <ref>).





 §.§ Dataset


In pursuit of a comprehensive understanding of the efficacy and practicality of parallel models, we meticulously selected several widely-acknowledged and representative image classification datasets for an integrated comparative analysis. These datasets encompass a diverse array of visual content and span various levels of difficulty, thereby furnishing us with a fertile ground to probe the performance boundaries of parallel models. Specifically, our study centers on the following pivotal datasets:

1. CIFAR-10 & CIFAR-100 <cit.>: These were constructed by Alex Krizhevsky et al. CIFAR-10 encompasses 60,000 32x32 pixel color images across 10 classes, while CIFAR-100 extends this to 100 categories.

2. STL-10 <cit.>: STL-10 comprises 5,000 training and 8,000 testing images distributed among 10 classes, each at a 96x96 pixel resolution. 

3. Flowers-102 <cit.>: Focused on floral categorization, this dataset houses 102 distinct flower species, totalling 8,189 images with a per-species count ranging from 40 to 258. 

4. Oxford-IIIT Pets <cit.>: Tailored for the recognition of pet breeds (cats and dogs), this dataset contains over 7,300 images across 37 categories, characterized by varied animal poses, backgrounds, and lighting conditions. 

In Table <ref>, we provide more details for these dataset. 



Through experimentation on these meticulously chosen datasets, we aspire to thoroughly validate the feasibility and strengths of parallel models across a spectrum of scales, complexities, and domains in image classification tasks. Simultaneously, we seek to elucidate potential limitations and chart avenues for future research endeavors.






 §.§ The Feasibility of Para-Former


To validate the feasibility and characteristics of Para-Former, we conducted training and validation on several commonly used image classification datasets, with all models trained for 500 epochs. The experimental design involved preprocessing inputs using the common ViT techniques of patch extraction and positional encoding. We then used Transformer as the basic model architecture, setting all network depths to 1 and configuring Para-Former from single-layer to multi-layer setups, as shown in Figure <ref>.a. A single-layer setup corresponds to a ViT with a depth of 1. The results of the models are presented in Table <ref>.

It is evident from Table <ref> that as the number of parallel network layers increases, the model accuracy improves progressively (Except for the result of Para-F-1-24 being worse than Para-F-1-12 on the Flower-102 dataset, this variation is within a reasonable range.). According to UAT theory, increasing the network layers is equivalent to increasing N in Equation (<ref>), thereby enhancing the approximation capability of UAT. The results in Table <ref> are in complete agreement with UAT theory. Therefore, we believe that this parallel network approach is entirely feasible.





However, the overall prediction accuracy in Table <ref> is quite low. What could be the reasons for this? We believe there are three main factors affecting the model's predictive performance:

1. Model's Fitting Capability: This is influenced by the model's size and design. Specifically, the value of N in UAT and the degrees of freedom in the parameters. The larger the N value, the stronger the UAT's fitting capability. Parameter freedom refers to the number of parameters affected by the input in UAT, such as a freedom of 1 in Eq. <ref> and a freedom of 6 in Eq. <ref>. The degrees of freedom correspond to the model's ability to dynamically fit based on input. Merely increasing N enhances the model's ability to fit specific functions, but images are often diverse. Therefore, the network must have the ability to dynamically fit the corresponding function based on the input. Higher degrees of freedom mean the model can fit a wide range of features with significant differences. Additionally, the bias term in the network can also dynamically fit based on input using UAT, which we refer to as bias freedom. Models with bias freedom have a stronger dynamic fitting capability. The mathematical model generated by the residual network is an example of UAT with bias freedom, achieving the fitting capability of a multi-layer network with just shallow layers. We will talk about this problem in Section <ref>.

2. Data Influence: Data is another critical factor in deep learning training. Since natural images are usually diverse, limited data can lead to overfitting. For example, in a cat-dog classification task, if dogs in the training set mostly appear against backgrounds of blue skies and grass, these backgrounds might be mistakenly considered part of the dog's features, a phenomenon known as feature coupling. If cats rarely appear in such backgrounds in the training set but do in the test set, they might be misclassified as dogs. Training on a large, diverse dataset means that common backgrounds frequently appear with various objects or animals, such as cats, dogs, and elephants against blue skies, white clouds, houses, and trees. Thus, even when classifying objects, the model inherently learns background information. In this context, fine-tuning on the cat-dog dataset yields better results because the background does not interfere with dog feature recognition. Without sufficient data, even a network with strong fitting capabilities cannot perform well and is prone to overfitting. This problem will be considered in Section <ref>.

Considering these factors, it is clear that simply improving the model architecture is not enough. A robust training dataset is equally important for achieving high predictive accuracy.








 §.§ Model Depth Effect


Based on the conclusions from Section <ref>, we will explore the impact of network depth on predictive performance in this section. According to Eq. <ref> and Eq. <ref>, we know that model depth primarily affects the degrees of freedom of the model parameters in UAT. Therefore, we designed Para-Former models with 24 layers and varying depths of 2, 3, and 6, and compared their predictive performance with that of a ViT with a depth of 6. The results are shown in Table <ref>.

It is evident that all results surpass those of Para-F-1-24 in Table <ref> because the networks in Table <ref> have higher degrees of freedom in their parameters. Notably, Para-F-6-6 consistently outperforms Para-F-2-24 and Para-F-3-24, indicating that the higher the degrees of freedom of the parameters, the stronger the model's performance. However, the model cannot be too deep, so we must balance network depth and layer count. Another issue is that almost all models hit a performance bottleneck. Particularly, the prediction accuracy for OxfordIIITPet and Flower-102 is very poor. This is because these datasets are relatively small and exhibit high data diversity.











 §.§ Dataset Effect







Current deep learning models generally perform well on the aforementioned datasets. However, as shown in Table <ref>, the prediction accuracy of Para-Former is not very high, especially on the OxfordIIITPet and Flower-102 datasets. What is the fundamental reason for this poor performance? In this section, we will explain the relationship between deep learning and data.

First, the results were obtained using the validated ViT model (Para-F-1-1 and Para-Former-6-6), which performed similarly to Para-Former. Given that the ViT model has been extensively validated and is a commonly used model architecture, the effectiveness of the models is unquestionable. So, why is the performance so poor? The root cause lies in the data itself: the model's performance is generally directly proportional to the amount of available data. Take the Flower-102 dataset, for example, which includes 102 classes of flowers but has only 1,020 images in the training set, with some reserved for validation. This is far from sufficient to train deep learning models. Why, then, does the OxfordIIITPet dataset, which has more data per class on average than Flower-102, perform even worse? This is due to the diversity of the data (i.e., the similarity of features within the same class) and the variety of backgrounds.

Figures <ref> and <ref> show some results from the OxfordIIITPet and Flower-102 datasets. Both figures indicate that prediction results are worse for categories with diverse backgrounds or inherent diversity, as seen in Figures <ref>.b and <ref>.b. Conversely, results are relatively better when the background is consistent but the species vary, as shown in Figures <ref>.a and <ref>.a. The best results are achieved when both the background and species variation are minimal, as shown in Figures <ref>.c and <ref>.c. Although the Flower-102 dataset has fewer images, its backgrounds are more consistent and there is less variation among the species. In contrast, animals in the OxfordIIITPet dataset can appear in various forms, while flowers in the Flower-102 dataset are usually photographed from relatively fixed angles, leading to better performance.

Does this mean that deep learning models are not powerful enough or lack generalization and reasoning capabilities? On the contrary, achieving around 25% accuracy on OxfordIIITPet and about 40% on Flower-102 with such limited data actually demonstrates the strength and generalization capability of deep learning models. A common misconception about deep learning is that humans, even if unfamiliar with certain species, can accurately distinguish OxfordIIITPet and Flower-102 images by carefully examining them, thus believing that deep learning lacks generalization or human-like inferencing ability. This viewpoint is flawed. Assuming the human brain functions as a UAT model, it can automatically disregard background information before seeing these images, as this background is frequently encountered in daily life, and the UAT in the human brain has already fitted this information. Therefore, humans only need to focus on specific information and then use it to train their brains to find similarities and differences. For instance, a 10-year-old child’s brain, observing the world at 60 frames per second, would have processed about 1 × 10^10 high-resolution images by the age of 10, many of which include common objects and backgrounds in daily life. Therefore, criticizing deep learning models for being unable to train on small datasets and lacking universality and generalization is unfair.

Thus, we should train on large datasets and then fine-tune on smaller ones. This is also the current deep learning training method, showcasing the powerful generalization capability of deep learning models, especially when there is some correlation between datasets. Following this fine-tuning approach, we fine-tuned the pre-trained ViT model on the aforementioned datasets for just 10 epochs. The results, as shown in Table <ref>, indicate a significant improvement in prediction performance.

Figure <ref> shows the rationality behind fine-tuning. Since we essentially adjust parameters based on the input to produce target results using UAT, if there is a correlation in the data, the corresponding parameters derived from the input (α_i, 𝐖_i, and θ_i; i=1,·,N) will be similar. Therefore, fine-tuning is feasible.








§ CONCLUSION

In this paper, we propose a parallel network based on the Universal Approximation Theorem (UAT) and validate its effectiveness on large datasets. We analyze the characteristics of both serial and parallel networks. Serial networks, especially those based on residual connections, often have higher degrees of parameter freedom, leading to stronger generalization capabilities. However, as the network depth increases, the inference time also becomes longer. On the other hand, parallel networks trade computational resources for time, meaning their runtime is not affected by the number of network layers. This allows for the design of multi-layer parallel networks that can replace serial networks to achieve fast inference. Additionally, we analyze the impact and reasons data have on deep learning. From the UAT perspective, training on large datasets essentially enhances the model's ability to fit data.




















§ THE UAT FORMAT OF TRANSFORMER










 §.§ The UAT Format of Transformer-based ViTs

Eq.(<ref>) and (<ref>) give an example of the matrix-vector format of MHA and FFN. Eq. (<ref>) and  Eq. (<ref>) give the corresponding UAT format of 𝐱_i+1 and 𝐱_i+2. It is obvious that the bias and weight terms in the UAT corresponding to Transformed can change dynamically based on the input.

Due to the complexity of the variables involved, we will combine some terms to align the resulting equations with the standard UAT form. To distinguish these parameter variables, we will keep the original variables unchanged. If a parameter variable is influenced only by other variables, we will denote it with a bar on top, such as 𝐖̅. If it is influenced by the input 𝐱_i, we will mark it with a hat, such as 𝐖̂.

















 Paper Title: Point-supervised Brain Tumor Segmentation with Box-prompted MedSAM.






Point-supervised Brain Tumor Segmentation with Box-prompted MedSAM
    Xiaofeng Liu, Jonghye Woo, Chao Ma, Jinsong Ouyang, and Georges El Fakhri
X. Liu, C. Ma, J. Ouyang and G. El Fakhri are with the Department of Radiology and Biomedical Imaging, Yale University, New Haven, CT 06519.
J. Woo is with the Massachusetts General Hospital and Harvard Medical School, Boston, MA 02114.

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



§ ABSTRACT 
Delineating lesions and anatomical structure is important for image-guided interventions. Point-supervised medical image segmentation (PSS) has great potential to alleviate costly expert delineation labeling. However, due to the lack of precise size and boundary guidance, the effectiveness of PSS often falls short of expectations. Although recent vision foundational models, such as the medical segment anything model (MedSAM), have made significant advancements in bounding-box-prompted segmentation, it is not straightforward to utilize point annotation, and is prone to semantic ambiguity. In this preliminary study, we introduce an iterative framework to facilitate semantic-aware point-supervised MedSAM. Specifically, the semantic box-prompt generator (SBPG) module has the capacity to convert the point input into potential pseudo bounding box suggestions, which are explicitly refined by the prototype-based semantic similarity. This is then succeeded by a prompt-guided spatial refinement (PGSR) module that harnesses the exceptional generalizability of MedSAM to infer the segmentation mask, which also updates the box proposal seed in SBPG. Performance can be progressively improved with adequate iterations. We conducted an evaluation on BraTS2018 for the segmentation of whole brain tumors and demonstrated its superior performance compared to traditional PSS methods and on par with box-supervised methods.


 




§ INTRODUCTION
 



The costly labeling effort in medical image delineation significantly hinder the development of data-driven AI models. There are increasing interests on weakly supervised segmentation to utilize bounding box or even a single point as label for training supervision <cit.>. However, due to the absence of accurate size and boundary guidance, there is a large performance gap between point-supervised medical image segmentation (PSS) and mask/box-supervised counterparts <cit.>. 

The recent progress of vision foundational models has achieved breakthroughs in several weakly supervised tasks, benefited by their strong zero-shot generalizability. For instance, the point-prompt natural image segment anything model (SAM) has been applied to enhance pseudo labels for point-supervised segmentation <cit.>. Recently, the medical image version of SAM (MedSAM) <cit.> has been trained with 1.5 million segmented medical images to achieve generalizable box-prompt segmentation. MedSAM takes both an image slice and a bounding box, i.e., prompt, to predict the possible segmentation mask within the box. Notably, the natural image SAM itself follows point-prompt design. Therefore, it is not straight forward to integrate box-prompt MedSAM to PSS task as <cit.>. In addition, a significant limitation of SAM/MedSAM is the lack of classification ability, resulting in class / structure-agnostic segmentation results <cit.>. They are designed for general use, while fails to accurately delineate the specific lesion or structure as desired <cit.>.  



 

To our knowledge, this is the first attempt to integrate MedSAM to facilitate semantic-aware PSS. We adopt an iterative framework to achieve coarse-to-fine progression of the bounding box. The point prompt is first converted to a box-proposal seed. Then, we configure a semantic box-prompt generator (SBPG) to propose and pick the reasonable box according to semantic similarity as <cit.>. It is followed by a box prompt guided spatial refinement (PGSR) to utilize the generalizable MedSAM to predict the segmentation mask. In addition, the smallest box that covers the mask is further used as the box seed in the next round of SBPG. We do not need to fine-tune the MedSAM to be semantic aware, which is prone to catastrophic forgetting. Notably, only the prompt refiner θ with the encoder part of the segmentor model will be trained, which involves about half the parameters of mask-supervised UNet training. 

We demonstrated its effectiveness in BraTS2018 for PSS of the whole brain tumor segmentation from T2-weighted MRI slices. We show that in testing, with 3 to 5 rounds of iteration, the point-prompt can achieve superior performance to approximate the box-supervised MedSAM <cit.>.








§ METHODOLOGY

In PSS, we are given tuplet {x^i,ρ_c^i,s^i}, in which x^i can be an MR slice, while the point prompt ρ^i_c={ρ^i_x,ρ^i_y,c} indicates the 2D spatial coordinates and the interested class c, e.g., tumor or normal tissue. We would expect the segmentation mask prediction ŝ^i_c to approximate the expert annotation s^i_c. 

To catering the box-prompt MedSAM <cit.>, an initial box proposal seed of class c, i.e., b_c:0^i, is generated from ρ^i_c, which is centered on point location (ρ^i_x,ρ^i_y) with the size of X× Y. We do not expect b_c:0^i fit the region of interests very well as there is no prior information about the size in the initial step. Then, the proposal bag ℬ^i_c={b_c:n^i}_n=1^N is created by scaling b_c:0^i with N different scales. Therefore, the boxes in group ℬ^i_c has strong spatial correlation with ρ_c^i, which avoid the inefficient random proposal in whole image <cit.>. 

To enable the model be aware of the specific semantic class, for example tumor in our task, a parameterized prompt refiner θ with a segmentor encoder and two fully connected layers is applied. Specifically, for the training sample with segmentation mask label s^i_c, we convert s^i_c to the corresponding smallest rectangle bounding box label b^i_c. Then, we store a set of feature f^i_c:n=θ(x^i,b_c:n^i) in the most recent M batches to a memory buffer, and calculate the mean of {f^i_c:nm}_n,m=1^N,M as the prototype 𝒱_c. As conventional multiple instance learning (MIL) based methods <cit.>,
the instance level probability indicating the likelihood of b^i_c:n matches 𝒱_c can be approximated with 
    p(b^i_c:n,𝒱_c)=e^cos(f^i_c:n,𝒱_c)/∑_c=1^C e^cos(f^i_c:n,𝒱_c),
where cos(·) indicates the cosine similarity, and C is the number of segmentation category. With the ground truth b^i_c, θ is trained with the binary cross entropy loss of 

    ℒ=-∑_c=1^C {c^i_1log∑_n^N p(b^i_c:n,𝒱_c) + (1-c^i_1)log(1-∑_n^N p(b^i_c:n,𝒱_c))},
where c^i_1 is the one-hot class label. Therefore, the prompt refiner is optimized to be semantic aware for the specific anatomical structure. The prototype in buffer is also updated with the new model θ. With p(b^i_c:n,𝒱_c) for each b^i_c:n, we pick the highest probability one  in ℬ^i_c as our optimal proposal b^i_c:* in current stage, which is also our bounding box inference.


Then, in the PGSR module, the off-the-shelf MedSAM <cit.> takes both x^i and b^i_c:* to predict the possible segmentation mask ŝ^i_c. Notably, no information about the interested class can be directly informed in MedSAM. Although the MedSAM has strong ability of zero-shot segmentation to delineate the relatively accurate mask within the current box prompt b^i_c:*, the alignment of ŝ_c^i is highly dependent on the precise b^i_c:*. This motivated us to refine b^i_c:* with a more suitable proposal seed b_c:0^i in SBPG. 

In practice, we can simply apply the smallest bounding box b̂^i_c:0 for ŝ^i_c:0, and use b̂^i_c:0 as our box proposal seed in the next round of SBPG. With T rounds of iteration, we expect that the box-prompt can be progressively refined to inform the MedSAM to predict an accurate segmentation mask. 

In particular, in the last round of testing, we do not need to generate b̂^i_c:0. Instead, the T-th round ŝ^i_c:0 is used as our final prediction.









§ EXPERIMENTS AND RESULTS

We evaluated on BraTS dataset for whole tumor segmentation as <cit.>, and use T2-weighted slices as our input. We used the 80%/20% split for training or testing. We generate the smallest bounding box for the mask and use its center as ρ^i_c. Our prompt refiner adopted the encoder part of the segmentor in <cit.>, which is followed by two fully connected layers of 1024 and 256 dimensions, respectively. We empirically set X× Y=21×21 and search proper T∈{1,2,3,5,10}. We followed <cit.> to compare with the conventional point-supervised method of WISE-Net without the use of the vision fundation model. Also, we directly input the ground truth of the bounding box into MedSAM, which can be “upper bound" of the performance of pointwise method. In Tab.<ref>, we can see that our model with T≥5 outperforms WISE-Net by a large margin. The Dice score is close to the box-supervised MedSAM.   





§ CONCLUSION

In this work, we proposed an efficient iterative framework to enable box-prompt MedSAM to point-supervised medical image segmentation. A lightweight prompt refiner with only encoder is specifically trained for the interested structural class. We show that the fixed off-the-shelf large model can flexibly target the specific downstream task without the need of large scale fine-tuning. The medical PSS performance can be largely improved by the advance of vision foundation model, e.g., MedSAM, which show great promise for facilitating the image-guided interventions. 
 


 


§ ACKNOWLEDGMENT

This work is supported by R21EB034911, R01CA165221, R01DC018511, P41EB022544, and NAIRR240016. 

ieee





 Paper Title: Assessing Programming Task Difficulty for Efficient Evaluation of Large Language Models.




Describe, Transform, Machine Learning: Feature Engineering for Grain Boundaries and Other Variable-Sized Atom Clusters
    [
    
======================================================================================================================




§ ABSTRACT 
  Large Language Models (LLMs) show promising potential in Software Engineering, especially for code-related tasks like code completion and code generation. LLMs' evaluation is generally centred around general metrics computed over benchmarks. While painting a macroscopic view of the benchmarks and of the LLMs' capacity, it is unclear how each programming task in these benchmarks assesses the capabilities of the LLMs. In particular, the difficulty level of the tasks in the benchmarks is not reflected in the score used to report the performance of the model. Yet, a model achieving a 90% score on a benchmark of predominantly easy tasks is likely less capable than a model achieving a 90% score on a benchmark containing predominantly difficult tasks. This paper devises a framework, HardEval, for assessing task difficulty for LLMs and crafting new tasks based on identified hard tasks. The main idea of our framework is to use a diverse array of prompts for a single task across multiple LLMs in order to obtain a difficulty score for each task of a benchmark. Using two code generation benchmarks, HumanEval+ and ClassEval, we show that HardEval can reliably identify the hard tasks within those benchmarks, highlighting that only 21% for HumanEval+ and 27% for ClassEval of the tasks are hard for LLMs. Through our analysis of task difficulty, we also characterize 6 practical hard task topics which we used to generate new hard tasks. Orthogonal to current benchmarking evaluation efforts, HardEval can assist researchers and practitioners in fostering better assessments of LLMs. The difficulty score can be used to identify hard tasks within existing benchmarks. This, in turn, can be leveraged to generate more hard tasks centred around specific topics either for evaluation or improvement of LLMs. Moreover, while we focus on code generation tasks in our experiments, HardEval generalistic approach can be applied to other domains such as code completion or Q/A. 





§ INTRODUCTION

Large Language Models (LLMs) have been widely used for various code-related tasks in Software Engineering (SE) <cit.>. Although the performance of LLMs in generating code for different programming languages, such as Python, Java, and C <cit.> is promising, LLM-generated code like human-written one is not flawless. Researchers already acknowledged that LLM-generated code are prone to bugs <cit.>, and may suffer from vulnerability <cit.>, and quality issues <cit.>. 


Generally, to assess LLMs on code generation tasks and compare them to each other, well-known benchmarks such as HumanEval <cit.>, MBPP <cit.>, CoderEval <cit.>, and ClassEval <cit.> are used. These benchmarks are generally obtained by handcrafting tasks or mining existing public repositories over GitHub and filtering them. This way, the code generation capability of LLMs on programming tasks is assessed indirectly at the benchmark level based on general characteristics such as dependency level (e.g., function level being less difficult than class level) and through general metrics such as Pass@k over the benchmark <cit.>. Moreover, the assessment of those tasks is conducted via a single prompt, which is a single formulation of the task. Yet, recent studies <cit.> showed that the way prompts are formulated has a high impact on the output quality of LLMs. Thus, current benchmarks' usage and assessment paint a general yet incomplete picture as they fail to provide a more fine-grained analysis of individual tasks within a benchmark. In particular, such benchmarks' usage does not provide a way of quantifying the difficulty of a programming task for LLMs. 

As a motivating example, consider the following task to implement the function cycpattern_check from HumanEval benchmark: 

[language=python, caption=A task example from HumanEval+., label=list:example]
def cycpattern_check(a , b):
    """ Write a function named 'cycpattern_check' that checks if the second word or any of its rotations is a substring of the first word. """


To assess an LLM, one would generally run this prompt either using greedy decoding or temperature sampling <cit.>, and then test the outcome against a set of predefined tests to evaluate correctness. Multiple LLMs could even be used to offset the potential impact of the particular training/architecture on the evaluated correctness. Doing so with a greedy decoding using the 5 LLMs we use in our experiments, results in all models except one failing on this task. In this context, should the task be considered hard? To tackle this issue, we propose HardEval, a framework for the evaluation, identification, and generation of difficult programming tasks for LLMs. We show with HardEval that this task is not difficult. Indeed, using our framework, all models except one indicate a low difficulty score for the task with a low difficulty averaged over all LLMs. As it turns out, rephrasing this basic prompt slightly or providing small additional contextual information helped the LLMs to easily address the task. As such, while the above prompt could be considered hard, the task per se is not.










To guide our study, we formulate the following Research Questions (RQs):


    
  RQ1 
  RQ2 
  RQ3 


HardEval is orthogonal to the existing benchmark's usage and works on top of available benchmarks. HardEval starts by generating a variety of prompts for each task in two well-known code generation benchmarks, HumanEval+ <cit.> and ClassEval <cit.>, representing multiple possible formulations per task. For each programming task, we craft prompts with different levels of context information and distinctive phrasing to represent diverse possible wording of the task using GPT-4, resulting in 18 different prompts per task. The framework then queries an array of 5 LLMs (i.e., CodeLLama 7B <cit.>, MagiCoder 6.7B <cit.>, DeepSeekCode 7B <cit.>, CodeGemma 7B <cit.> and GPT-3.5 <cit.>) using the formulated prompts to obtain different outputs for every single task. Those outputs are then used to compute a task difficulty score by averaging overall LLMs a combination of the functional correctness (passing all available tests for the task) and the syntactic correctness (using code similarity metrics) obtained on each code sample obtained. This difficulty score separates hard from easy tasks. Such hard tasks uncover particular shortcomings of LLMs that can then be used for improving LLMs. Then, we cluster hard tasks to recognize similar ones using a topic modeling technique called BERTopic <cit.>. To further generate new tasks, we leverage hard tasks from those topics in two ways: either by transferring constraints from hard tasks onto an easy task, or by modifying an already hard task by adding an additional constraint. Our results show that while benchmark-wide metrics such as accuracy using pass/fail code can give a general trend of a benchmark difficulty, they fail at assessing the difficulty of individual tasks contrary to the difficulty score proposed in HardEval. Moreover, using the identified topics we managed to generate 15 new tasks for both benchmarks for which the difficulty score is high as an example of HardEval's capacity. To the best of our knowledge, this paper is the first attempt to investigate the difficulty of code generation tasks for LLMs, suggest metrics, and provide a way to generate new targeted hard tasks based on identified hard programming tasks.



In summary, this paper makes the following contributions:

    
  * We present HardEval, a framework for evaluating the difficulty of code generation tasks for LLMs, identifying hard tasks, and generating novel hard coding tasks for LLMs,
    
  * We show the limit of assessing difficulty based on general benchmark metrics such as Pass@1 accuracy.
    
  * We make the data and code used in our study publicly available for researchers to duplicate our results <cit.>. In particular, we also provide the hard tasks identified for further research.


The rest of this paper is organized as follows. In Section <ref>, we describe the methodology we followed to propose HardEval. Our experimental setup for running HardEval is detailed in Section <ref>. We report and then analyze the obtained results in Section <ref>. A discussion of our results and insights for the community is presented in Section <ref>. We conclude the paper in Section <ref> after discussing threats to the validity of our findings in Section <ref>.



§ HARDEVAL


In this section, we discuss the methodology followed to propose HardEval. After an overview of the framework, we describe how we generate different prompts for programming tasks. We then delve into the details of leveraging LLMs to generate code for our prompts, measuring the difficulty of tasks for LLMs, identifying hard coding tasks, and generating new hard tasks.



 §.§ General Overview of HardEval

HardEval uses existing benchmarks to operate. Such benchmarks are generally composed of an ensemble of tasks which are represented by a single prompt per task. In our study, we make the distinction: the task is the essence of the functionality to be implemented. For instance, Quicksort or Fibonacci sequence are examples of such tasks. A prompt is a way to formulate the task so it can be understood by and fed to an LLM. For instance, Write a function to calculate the n-th element of the Fibonacci sequence. or Create a function to return the n-th element of the Fibonacci sequence. The Fibonacci sequence is defined as ... are two instances of prompts for the task Fibonacci sequence. As such, each task can be expressed by many different (even infinite) prompts, including different wordings, and amount of contextual information about the task. HardEval aims to evaluate the difficulty of those tasks for code LLMs in order to generate targeted benchmarks to probe potential hard tasks of LLMs.



A general overview of HardEval is given in Figure <ref>. HardEval is operated in three steps: -.9pt1 First, transformations on the prompts are identified. Those transformations are any dimension that can lead to a variety of prompts for an LLM while keeping the essence of the task. In a way, such transformations can be seen as metamorphic transformations <cit.> as they should not alter the semantics of the task. In our study, we focus on two transformations: Rephrasing and Context Information. Those transformations are then used in conjunction with each task to have a Prompt LLM generate prompts following the transformations' specifications. -.9pt2 Secondly, the obtained prompts are fed into several Code LLMs to generate code. The aim of those LLMs is to generate code for each given prompt and give us an assessment of how LLMs fare when confronted with the different prompts. Multiple LLMs are used to encompass the diversity of architecture and training settings that could react differently to a particular prompt. From the generated code, several metrics are extracted that will help us to obtain a Difficulty Score for each task in the benchmark. This score will further help us to classify tasks into easy and hard tasks. -.9pt3 Finally, hard tasks are analyzed to generate further targeted new tasks based on those topics. In particular, such tasks can be clusterized into tasks of similar semantics (e.g., sequence tasks, and sorting tasks), i.e., topics, and fed into a Task LLM. This allows us to generate new tasks similar in semantics to the identified hard tasks (e.g., new sequence generation based tasks), effectively generating a sub-benchmark of targeted tasks which can bring a finer analysis of LLM shortcomings for code generation. In the following, we will describe each step in more detail.



 §.§ Generating Different Prompts based on Transformations


The first step involves choosing transformations to generate different prompts for the task. We want those transformations to introduce variabilities in our prompts to represent the many ways an LLM could be prompted for the same task. However, those transformations should not introduce changes that alter the semantics of the tasks. In our case, we use two types of such transformations, but any number could be used and combined. 

First, we use the Context Information transformation. Indeed, a straightforward way to act on the variability of a prompt is to simply disclose or withdraw certain information in the prompt, while not changing the task. EvoEval <cit.>, for instance, proposes also to act on the difficulty. However, they do so by adding/removing constraints in the prompt thus modifying the semantics of the task which is not desirable in our case. In our study, we instead propose to divide the information provided in the prompt into three levels. Those three levels of prompts are generated incrementally using the previous level as a starting point. We assume that prompts with more information should be more likely to lead to correct code for LLMs. The levels are defined such as: 


    
  * First level: contains a minimal amount of information that is the inputs/outputs of the task as well as a high-level description of what is intended. We could use the original prompt in benchmarks for each task as the first level, yet Siddiq et al. <cit.> showed that there can be inconsistencies across tasks' prompts in benchmarks. Thus, we prefer to generate a new prompt to have similar formatting across all tasks and levels to reduce biases. Moreover, this level is the closest one to the original prompt in benchmarks and contains the same level of information regarding the tasks.
    
  * Second level: further adds in the prompt a description of the targeted function in natural language using available oracle code. However, it does not include any direct reference to the function (variable names, helper functions used, ...) unless it was explicitly mentioned in the original prompt of the benchmark. 
    
  * Third level: complete the prompt with direct references to the function such as variable names, and helper functions used similarly to a pseudo-algorithm but in natural language.


We present an example of different prompt levels obtained for a task in Listing <ref>.

The second transformation used is Rephrasing which is paraphrasing a given prompt. Indeed, how the prompts or instructions are formulated has been shown to have an impact on LLM responses <cit.>. Thus, similarly to these studies, we try to see if different ways of rephrasing a given prompt for an LLM can affect how the LLM addresses the task. To rephrase our prompts, we will use a similar approach as done by Gonen et al. <cit.>, that is prompting an LLM to rephrase a given text. Other techniques could be used to assist in rephrasing such as Chains-Of-Thoughts <cit.> but assessing the effectiveness of different rephrasing techniques is out-of-scope of this study.

[language=python, caption=Example of different prompt levels for the task unique from HumanEval+, with the oracle code at the bottom., label=list:levels]
def unique(l: list):
    Level1: """ Write a function named 'unique' that returns a sorted list of unique elements from a given list. """
    Level2: """ Write a function named 'unique' which takes a list as input and aims to return a sorted list containing only the unique elements from the input list. The function should first convert the list to a set to remove any duplicates and then convert it back to a list which should be sorted before returning. """
    Level3: """ Write a function named 'unique' which takes an input list 'l' and returns a list containing only the unique elements from 'l', sorted in ascending order. The function should first convert "l" into a set, using 'set(l)', to remove any duplicates, then convert this set back to a list, and finally return this list after sorting it in ascending order using 'sorted()'. """
    return sorted(set(l))


Using those two transformations, we come up with general templates that can be used to generate prompts (see replication package <cit.>). Such general templates are then tailored for each task in the benchmark under consideration. Those templates are fed to a Prompt LLM (see Figure <ref>) to generate different prompts for a given task. The way prompts are generated is as follows: we start from the original docstring to have a common starting point for all prompts. Then, we apply the template to obtain different prompt levels with Context Information for a given task. Then, we apply the obtained prompts to the Rephrasing template. In our case, we first remove the examples from the original prompt, if any, before using our templates with the Prompt LLM. While removing examples can decrease the performance of LLM <cit.>, this allows for better control over the level of information we inject in our prompts. Indeed, examples can vary from simple basic examples to showcasing corner cases of the task. As such, this might artificially bias the output of LLM and the potential difficulty of the task because of the quality of the examples. 

At each step of our generation, that is after using each template on the prompts, we do a sanity check by manually checking them. This is accomplished by the first author, and cross-checked by the second author independently. We make sure that the Context Information levels are respected and the Rephrasing does not alter the semantics of the tasks. This is done in order not to have tasks that are artificially considered hard further down the line because of misleading prompts for the code LLM and not because of the task itself.



 §.§ Assessing difficulty of tasks


The second step is to evaluate the difficulty of each task. To do so, we leverage several Code LLM that will serve as our evaluators for code generation. For each prompt per task, each LLM is asked to generate code with varying seeds using the sampling mode of LLMs. The rationale is that most LLMs, especially non-open-source ones, are used with this non-greedy approach in order to generate code. Moreover, using the sampling method for the generation allows us to have a wider variety of generated code in order to evaluate the difficulty of each task. Indeed, we assume an easier task should not be as impacted by the stochasticity of the generation compared to a harder task. The obtained code is then processed to be executed. Sample codes were then executed against available test inputs in the benchmark.

Once we have obtained all code for all prompts of all tasks and executed them against tests, we collect different metrics. In our case, we are interested in two metrics: the Correctness and the Similarity to Correct Code of a code sample. Correctness simply quantifies if a given code sample is correct or not using available test cases. As such, it assesses the functional correctness of a code sample. Formally, the Correctness Cor for a code sample s is defined as:

    Cor_s = {[                      1 ∀ (i, o) ∈ T, s(i) = o;                      0                        ].

where T is the set of test cases symbolized by an input i and an output o. That is, the code sample passes all test cases, i.e., gets a value of 1. Intuitively, we expect a harder task to lead to a lower number of prompts leading to correct codes across rephrasing and different levels of context information.

Nonetheless, code samples on a task can be wrong to different degrees depending on the bugs affecting the code sample <cit.>. Indeed, a code sample could simply be wrong because of a missing corner case (e.g., forgetting an initial condition) or it could follow a completely false logic and not implement the desired tasks at all. One way to assess this aspect would be to quantify the proportion of test cases for which a code sample s is correct <cit.>. However, this can lead to a biased measure depending on the test set. For instance, if a test set has more test cases assessing one corner case rather than another one, some code samples might end up with a higher proportion of passed tests artificially. Instead, we propose to use classical similarity metrics to measure the distance between a code sample and the the most similar correct code sample. To do so, we will rely on the CodeBLEU <cit.> metric to measure the similarity, which measures the syntactical correctness of the code sample. We chose CodeBLEU as it is a widely used metric in code generation tasks <cit.>. Thus, the Similarity to Correct Code Sim for a code sample s is defined as:


    Sim_s = max_sc(s, sc), C_sc = 1

where sc are correct code samples according to Correctness. Note that, in practice, we have at least one correct code as a reference which is the oracle code of the task in its benchmark. If s is correct, the similarity is 1 (closest is itself). This allows us to assess, when a code sample is incorrect, how similar it is to a correct code. The Correctness and Similarity to Correct Code are then used to evaluate a code sample separately, or in a Composite manner by aggregating the scores.

Finally, we use those metrics to compute a score per code sample which in turn helps to assess the difficulty score of the task itself for a given LLM. We used different transformations in this study: Rephrasing and Context Information. So, for instance, we expect an LLM to be more likely to generate a correct code if given a prompt with more context information (of a higher level) than given a prompt of a lower level. To represent this, the scores of the code samples obtained for each prompt can be weighted depending on the transformation. In our case, as we apply the Context Information before the Rephrasing, the difficulty score R_i of the i^th Rephrasing is defined as:


    R_i = ∑_j s_j × score_s_j

where score_s_j is the score obtained for the j^th code sample using the i^th rephrasing, as we generate multiple code samples for a given prompt through sampling generation. s_j are weights to give more or less importance to certain code samples. In our case, all samples are of equal importance and so all s_j equals 1 over the number of samples. Then, we calculate the score L_i for a given level of prompts i:

    L_i = ∑_j r_j × R_j

where r_j are weights used to give more or less importance to a certain rephrasing. In our case, all rephrasings are of equal importance and so all r_j equals 1 over the number of rephrasings. Finally, the difficulty score of the i^th task in the benchmark for a given LLM is calculated as:

    Difficulty_i = 1 - (α× L_1 + β× L_2 + γ× L_3)

where α, β and γ are weights over the scores obtained for each level. An explanation of the choice of parameters is given in Section <ref>.

Finally, once all the Difficulty scores have been computed for all tasks and LLMs, the score of a given task is simply the average over all LLMs. One can then use the score to determine which tasks are hard and which are not. To do so, we will use a threshold of 0.5. The choice of the threshold is motivated by our definition of the score for each code sample: a score of 0.5 would mean, in the most extreme case and for the Composite score, that all generated code samples for a task are incorrect but are very similar to a correct one. In practice, a score higher than 0.5 for a task would be more reflective of a high number of incorrect codes with a lower similarity, pointing towards a hard task.



 §.§ Analyzing and Generating new Hard Tasks
 
Once the hard tasks have been identified, the third step is to generate new tasks based on the identified hard tasks. To do so, we propose to leverage a topic modeling approach with BERTopic <cit.>, as was used in similar studies when dealing with natural language data <cit.>. For a given benchmark, the topic model clusters all tasks (easy and hard) using the prompts of Level 1 before rephrasing. We use this specific prompt type to have similar formatting across tasks while being as close as possible to the original prompt as mentioned in Section <ref>. As the process will inherently label some tasks as "noise", noisy tasks are not considered in the analysis. Then, for each topic, the proportion of hard tasks among all tasks in the topic is calculated and the topics with the highest proportion of hard tasks are analyzed. Those topics constitute tasks that have a similar functionality/goal, for instance, tasks dealing with sequence generation. Finally, we aim to generate more hard tasks leveraging the tasks from those selected topics with a Task LLM to generate additional new tasks. To do so, we leverage the Self-Instruct approach <cit.> which allows generating new tasks based on a subset of sampled tasks. Nonetheless, just providing easy and hard tasks to the Task LLM and asking it to generate more hard tasks did not lead to meaningful results according to our observation. Indeed, the small number of tasks per topic makes it quite impossible to determine what constitutes a hard task. Instead, we use a combination approach, where we ask the Task LLM to augment a hard task by adding a constraint, or by adding a constraint from a hard task to an easy one. This ultimately reduces the diversity of tasks generated but increases the likelihood of the generated tasks being hard. The newly generated tasks constitute a smaller sub-benchmark of targeted tasks related to the topic that can be used to further assess LLMs corner cases for code generation tasks.



§ EXPERIMENTS






 §.§ Hyperparameters used

To have different prompts, we choose to generate 3 levels and 6 rephrasings per level. We chose those parameters as the trade-off between generating diverse enough prompts while limiting the amount to generate. For a benchmark of n tasks, this effectively means we have to generate 18 × n prompts. We further ask each Code LLM to generate 5 code samples per individual prompt resulting in 90 × n code samples for each Code LLM to generate. As an example, for a benchmark containing 200 tasks, this process yields 18,000 generated code for our evaluation per Code LLM. We make available all the prompts and code generated in our replication package <cit.>.



 §.§ Benchmark used


In our experiment, we chose two benchmarks: HumanEval+ <cit.> and ClassEval <cit.>. Those benchmarks were chosen as they cover different types of tasks and dependency levels. HumanEval+ is widely used as a standard benchmark in code generation <cit.>. It is an extended variant of HumanEval <cit.>, that contains more tests to assess the correctness of code samples. The benchmark consists of 164 handwritten Python programming problems with a function signature, docstring, body and several unit tests. HumanEval+ tasks are at the function level dependency as they do not require external libraries, besides Python standard library, or class/file contexts to be operated. For the original prompt, we use the docstring as given in HumanEval+. 

ClassEval is a benchmark of 100 handcrafted classes coded in Python. Each class contains a description and on average 4 methods each with a docstring. We consider a task to be one of the 400 methods of ClassEval. Contrary to HumanEval+, ClassEval tasks are more similar to practical code used in projects and rely on a class context to operate. ClassEval proposes three prompting strategies: holistic, incremental and compositional. In our case, we use the compositional strategy, that is, providing the class context to the LLM and asking it to complete one method of the class, only modifying the docstring of the method to complete through our crafted prompts. We do not make use of the other generation approaches as it reduces the control we have over the prompts and could introduce biases: the holistic approach, as it implies asking the LLM to generate the whole class from scratch, can impact the performance artificially. Similarly, the incremental approach requires the reuse of codes generated by the LLM for other methods when generating new ones, which could also alter the difficulty. As we need to generate a high number of code samples and prompts, we reduce the number of tasks to evaluate by doing a stratified sampling at the 95% confidence interval. For 400 methods, this gives us 197 methods to sample which we round up to 200. Those 200 methods are sampled equally across classes resulting in 2 methods per class available. Methods chosen are available in our replication package <cit.>.



 §.§ Models used


In our experiments, we use GPT-4 <cit.> as both the Prompt LLM (to generate the different prompts) and the Task LLM (to generate new hard tasks). We do so as GPT-4 is the state-of-the-art LLM and so should provide answers with the best quality which would require the minimum human supervision. For the Code LLM, we use 5 LLMs with different architecture: CodeLLama 7B <cit.>, MagiCoder 6.7B <cit.>, DeepSeekCode 7B <cit.>, CodeGemma 7B <cit.> and GPT-3.5 <cit.>. We have not used GPT-4 as Code LLM to avoid biases as it is used to generate the prompts fed to the Code LLM. For all models, we used their instructions-tuned versions so the models can handle the different prompts properly. This gives us a diverse array of architecture and training procedures. For all code generation in HardEval, as we are using the sampling generation, we fix the temperature hyperparameter to 0.8 as it is one of the temperature settings commonly used in other similar studies <cit.>. In all cases except GPT models (for which we use OpenAI API), we use the HuggingFace <cit.> implementation of the models to allow for replication. For verification, we compute Pass@1 <cit.> with the samples obtained across levels for a given task. Results are reported in Table <ref>. Ranking of the LLMs' performance follows the assessment obtained on the original benchmarks.





§ RESULTS




 §.§ Validation of the Context information transformation


We aim to assess whether the level of information in the prompts impacts the generation of the code by the LLM in such a way that is quantifiable by our metrics. While for rephrasing, the impact on the output has been analyzed in the literature <cit.>, the context information transformation as we formulate it in our work has not been so. Some studies <cit.> showed that prompt-added contents can range from no benefit to negative effect on the code generated. For instance, Ouyang et al. <cit.> showed that a longer description in the prompt could decrease performance. As such, we verify that our added context transformation has the intended functionality, i.e., the more context we provide, the better the generation.

As reported in Table <ref>, the Pass@1 rate increases with the Context Information. We now set out to show there is a correlation between the scores computed by our metric and the Context Information. To do so, we use the scores of the generated code for each rephrasing (see Equation <ref>) and compare it to the prompt level it comes from. To see the impact of both the Correctness and Similarity metric, we compute each metric individually as well as the composite of the two. As we introduce gradation in the injected information in the prompts through the Context information, one expects to observe a correlation between the level of information and the score obtained (i.e., the more information in the prompt, the higher the score). To calculate the correlation, we use the Spearman-ρ <cit.>, as we are interested in a monotonic increase between the two variables. We report the results for each Code LLM independently and for the averaged over all LLMs. Results are reported in Table <ref>.




In Table <ref>, for all cases, we obtain a significant positive correlation, indicating that increasing the amount of information in the prompt generally improves the LLM-generated code in terms of the score for all of our metrics. The Composite score (i.e., Correctness + Similarity) yields the best correlation overall on HumanEval+ while Similarity is slightly better on ClassEval. As we have a positive correlation, it means a higher score for a level 3 prompt than a level 1 prompt and so they account for our Context Information. As the Composite metric seems to be better on HumanEval+ and is tied up on ClassEval, we consider only this metric in the rest of this paper to calculate the difficulty score.

We note that the correlations are around 0.2 to 0.4, signifying a positive significant but weak to moderate relationship. We suppose the reason is the variability across tasks. Indeed, due to the particular difficulty of each individual task or rephrasing used, there might be some variability within the tasks that could impact our outcome. For instance, easier tasks are more likely to have a higher score for Level 1 prompts while more difficult tasks can end up with a low score for Level 3 prompts. This can impact the outcome of the Spearman-ρ when taking all tasks into account at once. To account for this, we calculate the Spearman-ρ for each task individually before averaging the obtained correlations. As averaging correlations directly might lead to biases, we follow the recommended approach <cit.> of converting the Spearman-ρ correlations using the Fisher Z-transformation, averaging the results and converting it back to a Spearman-ρ. Doing so, we obtain Spearman-ρ > 0.5 in all cases, even reaching 0.70 and 0.64 respectively for HumanEval+ and ClassEval when using the scores averaged overall LLMs. This results in moderate to strong relationships. As such, our Context Information indeed helps LLMs to generate correct code when we account for task variability, and so it is a valid means to calculate the difficulty score.




























 §.§ Validation of the parameters for the Difficulty Score


Without ground truth, we describe the rationale behind the choice of parameters α, β and γ of the Difficulty score. In the previous section (Section <ref>), we show that the lower the amount of information in the prompt, the less likely the model could produce a correct code. As such, we should consider that α≥β≥γ, that is, emphasizing more on lower information context. Otherwise, this could artificially bias the difficulty score towards assessing tasks as easy. Consider that this would be equivalent to having an exam where most questions would have tips to help the participants: the exam would easily be passed, but it would not necessarily be representative of the tasks in the exam. The issue then becomes which values to assign to parameters. To decide, we analyze the shape of the Difficulty score distributions obtained by varying the parameters. In Figure <ref>, we give the distributions of the difficulty score obtained on the ClassEval benchmark used during experiments, by varying α. The values of β and γ are adjusted so the sum equals 1.



As illustrated in Figure <ref>, emphasizing too much on α will give too much importance to low information context. In this case, tasks will tend to be declared harder (see Figure <ref>, orange), because LLMs will struggle more in a low information context to produce correct code fragments. However, on the contrary, the lower the α the more emphasis is put on higher information context (see Figure <ref>, blue), for which we provide more information to the LLMs than typical prompts in benchmark normally would. This instead makes more tasks be declared as easy. A good balance would be found between those two extremes: too much and too little reliance on the low information context. To do so, we assume that benchmarks are balanced. That is, we consider that LLMs should not consider that tasks are overwhelmingly easy or hard. Given that no LLMs reach very high or very low accuracy on the study benchmark (see Table <ref>), this looks like a sound assumption. As such, we can consider the shape of the distribution as an indication for the choice of the parameters α, β, γ. 

In our case, we will consider two indicators of the distribution: the standard deviation σ <cit.>, which measures the spread, and the skew g_1 <cit.>, which measures the asymmetry.











To achieve the mentioned balance, we aim to minimize the skew and the standard deviation through our parameters α, β, γ. This implies minimizing the L2 norm of the skew/standard deviation vectors after normalization. To simplify the optimization, we only vary α from 0.33 to 1.0 and adjust β and γ such as β represent 60% of 1 - α and γ represent 40% of 1 - α (i.e. we put more weights on level 2 than level 3 information context as level 3 has more information). We report the values of standard deviation and skew obtained for the different combinations in Figure <ref>.

The minimum is obtained for α = 0.6 when considering HumanEval+ and α = 0.65 for ClassEval. To have a common set of parameters to allow for comparison between benchmarks, we will set the weights at α = 0.625, β = 0.225, γ = 0.15. In practice, for these values, the skew and standard deviation values are close to the minimum for each benchmark.



 §.§ RQ1: 




The goal of this RQ is to show HardEval can reliably identify hard tasks in a benchmark. To do so, using greedy decoding (i.e. Temperature = 0.0), as done in the-state-of-the-art leaderboard <cit.> for each of 5 LLMs of our study, we evaluate each task using a single prompt to see if each LLM outputs a correct code using available tests. For the single prompt used, we chose the level 1 prompt obtained before rephrasing for each task, which is the closest to the original prompt of each benchmark. We did so for a similar reason as mentioned in Section <ref>, i.e. removing potential formatting issues and providing a fair comparison across tasks. For greedy decoding, we consider a task to be hard for an LLM if it does not manage to generate a correct output. We consider a task to be hard for all LLMs if a majority of LLMs fails to generate a correct code for that task. We use the opposite settings for easy tasks. At the same time, we collect tasks that are labelled as hard/easy by HardEval using our difficulty score (see Equation <ref>), which is when the score is above/below 0.5. Using those sets of tasks, we can compute the overlap between the sets, e.g., the tasks that are hard according to HardEval but not for the greedy decoding definition, for a given LLM or all LLMs. This allows us to see if the sets of tasks are similar across approaches or not. Then, we compute the accuracy of the greedy decoding over the tasks by splitting them into two sets of tasks, using the set of hard/easy tasks of HardEval. Similarly, we compute the median of the difficulty score of HardEval obtained on the tasks by splitting them into two sets, based on the set of hard/easy tasks from the greedy assessment, for each LLM. Intuitively, we should see a drop in performance of the accuracy of the greedy decoding on the tasks that are noted hard by HardEval and an increase in performance for the tasks that are labelled as easy by HardEval compared to the greedy accuracy of each LLM.

Computing the overlap between the sets of hard tasks for both approaches results in an overlap ratio ranging from 0.24 to 0.63 for HumanEval+ and from 0.56 to 0.62 for ClassEval depending on the LLM. Similarly, the overlap for the sets of easy tasks ranges from 0.54 to 0.78 for HumanEval+ and from 0.69 to 0.80 for ClassEval depending on the LLM. If we consider all LLMs, the overlap on the hard tasks is 0.46 for HumanEval+ and 0.62 for ClassEval, while for the easy tasks, it is 0.71 and 0.78 respectively. Thus the assessment of hard tasks between the two approaches is dissimilar while there are more similarities in the assessment of the easy tasks. The assessment seems also to be more similar for ClassEval compared to HumanEval+. Thus, while there exists an overlap between our proposed approach and the greedy traditional assessment, the overlap is not perfect. 

We provide the results regarding computing the metric of each approach on the hard/easy tasks given by the other one. Results are reported in Table <ref>. For each row of the table, we get, for an LLM, the results (median difficulty score or accuracy) on both HumanEval+ and ClassEval benchmarks using the subsets of hard/easy tasks per approach. For instance, for CodeLLama on HumanEval+, the median difficulty score on the easy tasks of CodeLLama for the greedy approach Greedy Easy, that is when the output for a task is correct, is 0.27. We also give in that case the first and third quartiles. For instance, still on HumanEval+, the first quartile for MagiCoder in Greedy Hard is 0.29 (i.e. at least 25% of the tasks have a difficulty score lower than 0.29). Similarly, for GPT3.5, the accuracy of the greedy approach on hard tasks according to HardEval is 8.33%. Moreover, we also report the accuracy of the greedy decoding over the whole benchmark for each LLM and averaged across LLMs.

Using Table <ref>, we make the following observations: 
 


Greedy decoding based majority assessment agrees with HardEval 
 


When using HardEval sets of easy/hard tasks, we see an increase/decrease in performances for both benchmark and all LLMs compared to the performance calculated over the whole benchmark (see Table <ref> Greedy Acc). This shows that the tasks identified as easy/hard by HardEval are indeed easy/hard, as the subset of hard tasks will result in lower accuracy and the subset of easy tasks in higher accuracy compared to the whole benchmark baseline. 
 


HardEval assessment on easy tasks relatively agrees with Greedy decoding 
 


We see that, for all LLMs, easy tasks for the greedy decoding (i.e., whether individual LLMs manage the tasks and, Overall LLMs, a majority of LLMs manages the task) obtain a median of the difficulty score < 0.5 showing that most of those tasks tend indeed to be easy according to HardEval. We note however that for HumanEval+, the third quartile is quite high for most LLMs (but below 0.5). As such, for HumanEval+, there is a higher variation in the easy tasks difficulty which is not noticeable when using the greedy decoding approach but that we can identify with HardEval. 
 


HardEval assessment on hard tasks disagrees with Greedy decoding 
 


For the set of hard tasks for the greedy decoding, we get a median difficulty score fairly close to 0.5, the only exception being for DeepSeek and MagiCoder on HumanEval+ (0.35, 0.38). However, the first quartile is below 0.5 (around 0.4) for most LLMs and so at least 25% of the hard tasks that are hard when using greedy decoding, i.e. pass/fail on a single prompt, are not hard according to HardEval. This however was to be expected, as we saw in Listing <ref>, using a single prompt, a task can easily be hard for one or more LLMs because of the formulation and/or underspecifications in the prompt. However, rephrasing the prompts or adding information can lead to easy tasks, something that a single prompt can not tell us.


[colback=blue!5,colframe=blue!40!black]
Findings 1: While using the correctness of the code over multiple LLMs on a single prompt (greedy decoding based) can be used to identify easy tasks, it is not reliable for identifying hard tasks. On the contrary, the difficulty score of HardEval can identify hard/easy tasks in a benchmark. 




 §.§ RQ2: 


Our goal in this RQ is to study the distribution of difficulty scores across tasks as well as the type of tasks that are difficult in the benchmarks under study. To do so, first, we use HardEval difficulty scores for each task calculated in RQ1, both for the Code LLM individually as well as all LLMs. By doing so, we can output the cumulative distribution of the difficulty score across tasks for analysis. Then, we use the BERTopic model as described in Section <ref>. The topic analysis returns 17 topics for HumanEval+ and 21 topics for ClassEval. The two first authors manually check the tasks in those topics to provide a general name to the topic that is as representative as possible of the tasks within the topic. We then calculate the proportion of hard tasks in each topic. 




The results for the cumulative distributions are given in Figure <ref> for both benchmarks. On both the benchmarks under study, 63% of the tasks for HumanEval+ and 65% of the tasks for ClassEval, Overall LLMs, obtain a difficulty score of less than 0.4 and only 21% of the tasks for HumanEval+ and 27% of the tasks for ClassEval have a difficulty score higher than 0.5 thus making them hard tasks with our definition. Thus, from a high-level point of view, both benchmarks seem to have a similar difficulty profile, which echoes the accuracy presented in Table <ref>. However, there are a number of differences which we can analyze thanks to HardEval: 1) ClassEval tends to have a high proportion of trivial tasks (< 0.1) with 37% of ClassEval's tasks while on HumanEval+ it is 16% and there are much more variations across LLMs, 2) CodeLLama and CodeGemma have the most difference in terms of performance between ClassEval and HumanEval+. For instance, 50% and 33% of the tasks of HumanEval+ are hard for CodeLLama and CodeGemma respectively compared to 33% and 28% on ClassEval. We suppose the difference comes from the fact that HumanEval+ has more algorithmic tasks while ClassEval tasks are more related to general development tasks, which could benefit those two LLMs, and 3) ClassEval has more difficult tasks overall. This echoes the subjective assessment that ClassEval is harder than HumanEval+, as it relies on a class context. Finally, while the accuracy over both benchmarks seems similar (see Table <ref> and Table <ref> Greedy Acc for the same Code LLMs), the distribution of difficulty of tasks between two benchmarks can be quite different. To verify that, we compute a k-samples (k = 2) Anderson-Darling test <cit.> to assess if the distributions between ClassEval and HumanEval+ are the same for a given LLM or overall LLMs. This leads to a significant p-value < 0.05, i.e. the distributions are indeed different, in all cases except when comparing for GPT.




We then proceed to extract topics using BerTopics as described in Section <ref> to group tasks into topics and analyze them. The output is given in Figure <ref>. One can make several observations from the obtained results. First, not all LLMs necessarily encounter the same difficulty depending on the tasks. For instance, GPT has more trouble on Topic 12 in HumanEval+ but less so on the Topic 11 of ClassEval compared to Deepseek, as highlighted by the higher proportion of hard tasks in this topic as well as the average score over the tasks. This indicates that each LLM has its strengths and weaknesses that originate from its architecture and the data it was trained on. While the prompts we use could affect this outcome, the fact that other LLMs have, on the same topic, opposite behavior instead suggests that the root cause is mainly due to training and/or models' architecture. Secondly, some topics that are easy/hard across LLMs highlight some general behaviors. For example, Topic 15 of HumanEval+ or Topic 13 of ClassEval have almost no hard tasks in all models and overall. On the other hand, Topic 16 of HumanEval+ and Topic 3 of ClassEval have a higher proportion of hard tasks and this is for all LLMs. ClassEval seems to have more topics with a higher proportion of hard tasks compared to HumanEval+ but, at the same time, a higher proportion of topics with no hard tasks. These predominantly hard or easy topics highlight limitations in the ability of the LLMs. We note nonetheless, based on the extracted topics, there does not seem to be a divide between topics with topics with all hard tasks and topics with all easy tasks. This further motivates us to use a combination of hard and easy tasks from a similar topic to generate new tasks: 1) as topics are rarely only hard or easy, combining both easy and hard tasks or augmenting hard tasks, for which we already have a result, increase the likelihood of re-obtaining a hard task, 2) by leveraging topics, we can easily make use of the thematic similarity between tasks to combine/augment them. This, for instance, prevents potential arbitrary transformation when combining tasks that are not directly related which can have an unpredicted impact on the LLMs.

[colback=blue!5,colframe=blue!40!black]
Findings 2: At a high-level, HardEval difficulty score profiles of both benchmarks are similar, which echoes the similar Accuracy obtained in the Greedy decoding. However, ClassEval exhibits more hard tasks than Humaneval+ (27% against 21%) and the difficulty profile between benchmarks on the LLMs are different, which highlights HardEval local assessment usefulness. Inside the benchmarks, while there does not seem to be a clear divide between hard/easy topics, some topics for both benchmarks are predominantly composed of hard tasks or easy tasks, which can show potential ability limitations of the LLMs.




 §.§ RQ3: 


[language=python, caption=Oracle code of a new task generated. We highlight in yellow the change compared to the original hard task used for generation. Changes are made to be minimal so the task is as similar as possible to an already established hard task., escapechar=!, label=list:new_task]
def solve_with_vowel(s):
    !def vowel_change(ch):!
        !return ch if ch not in "aeiouAEIOU" else chr(ord(ch) + 2)!
    ans, has_letter = "", False
    for ch in s:
        if ch.isalpha():
            has_letter = True
            ans += !vowel_change!(ch.swapcase())
        else:
            ans += ch
    return ans if has_letter else s[::-1]



Finally, this RQ aims to show that, using identified hard tasks, we can further generate new hard tasks that are targeted at specific problems. We do not aim to be systematic, but merely to show that, using identified hard tasks and topics, it is possible to generate new tasks. To do so, we use the result from the RQ2 (i.e., the topics) and select the first three topics in terms of the number of tasks with at least a third of the tasks being hard. We do so to have a sufficient number of LLM-easy and hard tasks for the Task LLM to be able to provide meaningful new tasks. For each topic, using the Task-LLM, we generated tasks, checking each of them manually for validity, until we got 5 different new tasks using the process described at the end of Section <ref>. Then, we manually write an Oracle code as well as unit tests for this task. When possible, we try to remain as close as possible to the oracle code of the base task used in the generation process. Finally, we assess the newly generated tasks to the same Step 1 and Step 2 of HardEval to measure the difficulty score. We thus obtain, for each of the picked topics, new tasks that are assessed. We present an example of Oracle code from a newly generated task using the Task LLM in Listing <ref>.




The results for the difficulty score obtained for the newly generated tasks are given in Table <ref>. Most tasks generated following our approach turn out to also be hard. We note that a few of them end up being below the 0.5. We see two possibilities for this: 1) either the base task was easy and we added a constraint from a hard task, which is not what makes the task hard in the first place, or not completely, 2) adding a small constraint to a hard task will modify the prompt and so, we might end up with some sets of prompts that are more favourable for the generation of codes for LLMs. However, we note that the difficulty score is fairly close to 0.5, illustrating that the task is far from trivial. The only task with a fairly low score is on Topic 9 for ClassEval. This task was obtained by adding a hard constraint from an easy task. Yet, as the constraint was not applicable perfectly, we added to relax it in the implementation, which might explain the low score. In particular, the new task-generation process can vary across topics and benchmarks. For instance, while generating a sequence with added constraints on HumanEval+ is really simple for the Task LLM, coming up with new methods for already established classes on ClassEval can be quite hard, especially when the methods to use as a base are very different from each other. In practice, this results in the large majority of the tasks being a slight modification of a precise hard task. However, while not perfect, simply reusing identified hard tasks can be a simple way to generate new tasks that are focused on some specific topics. 

[colback=blue!5,colframe=blue!40!black]
Findings 3: Using identified topics and hard tasks rated by HardEval, we can generate new tasks. Those new tasks are similarly hard, showing it is possible to generate new hard tasks using our approach.




§ DISCUSSION


With HardEval, we show that assessing the difficulty of a task from the point of view of LLMs is feasible as well as identifying hard tasks. In particular, hard tasks identified with HardEval are truly hard as they are more resilient to the potential variability of the prompt contrary to assessing using a single prompt. While relying on human judgment to define a task difficulty could be used, it is unlikely to work for LLMs as they process code and task instructions differently than humans <cit.>. Ouyang et al. <cit.> showed that the difficulty assessed by humans on code competition only weakly correlates with the test pass rate variance across code samples generated by ChatGPT. Thus, it is likely that human assessment or related metrics might not be accurate for estimating LLMs' task difficulty contrary to the difficulty score of HardEval. Our proposed approach, while applied in experiments to code generation, can be extended to other code-related tasks. Indeed, the transformations we used are easily transferable to other paradigms and new transformations could be defined.

Analyzing the tasks that were labeled as hard, we observe that LLMs tend to disregard certain parts of the instructions, which can lead to certain test cases failing thus making the code incorrect. That is why similarity helps (i.e., code is incorrect but very similar to a correct one) (Finding 1). This phenomenon happens generally as the number of constraints in the prompt (and in the function) increases or if the prompt becomes longer. This explains why coding-competition or algorithmic-like problems in both benchmarks, which require a lot of description to express the many operations and constraints, generally lead to a higher difficulty, as the LLM is overwhelmed with information. This echoes Ouyang et al. <cit.>, as they showed that description length impacts the test pass rate in ChatGPT. This however can not on its own explain why difficulty could be higher as Level 3 prompts tend to have a better score and are longer than lower-level prompts. We also note that corner cases generally have to be explicitly mentioned otherwise LLM might miss it. For instance, stating when checking a date that February can have 29 days, or when checking a condition in a string, the string can be empty, can improve generation. This type of information is generally omitted or implicit in prompts in benchmark (and so in our level 1 prompts) but is more likely to be explicitly mentioned in level 2 or 3 prompts. This also can explain why using the correctness of codes over one simple prompt is not enough to quantify the difficulty of a task (Finding 2). 




§ THREATS TO VALIDITY

Internal validity: One threat to the validity of our results comes from the way we measure the difficulty of tasks. To mitigate this threat, we combined two metrics to see if a programming task is difficult for LLMs to code using functional and syntactic correctness and showed it correlates with the Context Information level. Those metrics were used in studies investigating LLM capabilities for code generation. We used a linear combination of all levels having the minimum weight for Level 3 (with maximum contextual information) and maximum for Level 1 (with the least information) to put more emphasis on the Level 1 prompt score. While this might influence the decision on the hard tasks, it is more likely to have some easy tasks be considered as hard, as we showed that the metrics correlate positively with the level. Another threat might come from the generation of the prompts. The transformations used are motivated based on existing works. Moreover, while we have made use of state-of-the-art GPT-4 to generate them, we manually checked the prompts to make sure that the generative process did not alter semantics. Finally, another threat to validity is the fact we did not consider examples in the prompts. Not adding examples in prompts, as we mentioned, can decrease the performance of the models and so add artificially hard tasks. This was nonetheless necessary to make the comparison fair across tasks as not all examples contain the same information. However, beyond having some easy tasks be considered as hard, it does not affect the applicability of HardEval.

External validity: Our sets of tasks from the selected benchmarks could be a threat. The collection of tasks might not be representative of real-world programming tasks to assess LLMs' capabilities. Nonetheless, both HumanEval+ and ClassEval are used to evaluate Code LLM performance <cit.>. Moreover, the objective of the study was to set the scene for the assessment of task difficulty rather than dwell on the practicalities of the task. So, we believe using tasks from these two benchmarks provides a diverse set of tasks to assess the difficulty of code generation using LLMs.

Reliability validity: We detailed our methodology, and to enable other researchers to reproduce or expand upon our work, we make our code and data publicly available <cit.>. 




§ CONCLUSION

In this paper, we realized the first attempt to investigate the difficulty of code generation tasks for LLMs. To do so, we proposed a novel framework, HardEval, to assess the difficulty of code generation tasks for LLMs, identifying hard programming tasks, and creating new hard tasks. We evaluated HardEval on two well-known benchmarks for LLM-based code generation: HumanEval+ and ClassEval. A set of diverse prompts is crafted for each task including different levels of contextual information about the task and various rephrasings, yielding 18 different prompts per task. 5 state-of-the-art LLMs were then employed to solve coding tasks using the set of prompts. We proposed a difficulty score based on functional and syntactic correctness. Tasks with high average difficulty scores across all LLMs were identified as hard coding tasks. Based on thematic similarities, we clustered tasks from both benchmarks into respectively 17 and 21 topics to inspect their characteristics. For selected topics, we generated new tasks. New tasks were then assessed to see if they were really hard based on the proposed score, revealing their difficulty for LLMs to be coded. This new set of tasks can be used as a targeted benchmark to assess the effectiveness of LLMs for code generation tasks or to improve on them. We hope this research paves the way for further research on improving LLMs for code-related tasks.

In future works, we plan on expanding the evaluation of HardEval to other domains such as code completion or summarization. Moreover, we aim to improve the new task generation method so the process is more reliable as well as to be able to have better control of the difficulty level of the task generated. In particular, this would allow us to generate a benchmark with a desired difficulty profile.


This work was supported by: Fonds de Recherche du Québec (FRQ), the Canadian Institute for Advanced Research (CIFAR) as well as the DEEL project CRDPJ 537462-18 funded by the Natural Sciences and Engineering Research Council of Canada (NSERC) and the Consortium for Research and Innovation in Aerospace in Québec (CRIAQ), together with its industrial partners Thales Canada inc, Bell Textron Canada Limited, CAE inc and Bombardier inc.


ACM-Reference-Format




 Paper Title: Enhanced Uncertainty Estimation in Ultrasound Image Segmentation with MSU-Net.





Uncertainty Estimation with MSU-Net






R. Banerjee et al.




Carnegie Mellon University, Pittsburgh PA 15213, USA {rohinib,cgmorale,awd}@andrew.cmu.edu


Enhanced Uncertainty Estimation in Ultrasound Image Segmentation with MSU-Net
    Rohini Banerjee* Cecilia G. Morales* Artur Dubrawski
    August 5, 2024
=============================================================================



*These authors contributed equally to this workfootnote

§ ABSTRACT 
Efficient intravascular access in trauma and critical care significantly impacts patient outcomes. However, the availability of skilled medical personnel in austere environments is often limited. Autonomous robotic ultrasound systems can aid in needle insertion for medication delivery and support non-experts in such tasks. Despite advances in autonomous needle insertion, inaccuracies in vessel segmentation predictions pose risks. Understanding the uncertainty of predictive models in ultrasound imaging is crucial for assessing their reliability. We introduce MSU-Net, a novel multistage approach for training an ensemble of U-Nets to yield accurate ultrasound image segmentation maps. We demonstrate substantial improvements, 18.1% over a single Monte Carlo U-Net, enhancing uncertainty evaluations, model transparency, and trustworthiness. By highlighting areas of model certainty, MSU-Net can guide safe needle insertions, empowering non-experts to accomplish such tasks.








§ INTRODUCTION


Trauma, the leading cause of death among young individuals in the U.S. <cit.>, often results in blood loss, which requires rapid fluid resuscitation for vital organ oxygenation. In austere settings, accessing timely medical care can be challenging due to limited access, dangerous conditions, time constraints, or the absence of medical infrastructure, making expertise for optimal needle insertion sites critical. Autonomous robotic systems can assist in intravenous fluid administration when medical experts are unavailable, providing support in emergencies. These systems can also guide non-experts in accurate performance of phlebotomy tasks, empowering them to contribute effectively in dire medical situations.

Ultrasound imaging is widely used for locating vessels to enable fluid resuscitation due to its affordability, speed, safety, and portability, unlike CT or MRI imaging that are not portable and use ionizing radiation <cit.>. Despite advancements in autonomous needle insertion into blood vessels <cit.>, a critical challenge persists: inaccurate predictions can lead to severe consequences. Failing to anticipate vessel structure during intravenous cannulation may result in catastrophic outcomes for critically injured individuals; incorrectly predicting two adjacent vessels as a single one could lead to laceration of the vessel wall and cause hemorrhage upon needle insertion <cit.>. Both medical personnel and automated tools need to assess the certainty of their estimates of vessel location and structure to mitigate this risk <cit.>. 
As healthcare increasingly adopts automation, reliance on Artificial Intelligence (AI) models grows. Building a  model that communicates its uncertainty becomes essential to help users focus their attention and actions on where the model is confident, such as guiding needle insertion accurately within vessel segments. In this paper, we attack the problem of uncertainty estimation of models trained to guide needle insertion tasks. Our contributions include: 

(1) We introduce MSU-Net, a novel Multistaged Monte Carlo U-Net, demonstrating significant advancements in uncertainty quantification; (2) Our model achieves state-of-the-art performance in ultrasound image segmentation, marking the first known improvement in uncertainty estimation for this application.





§ RELATED WORKS



 §.§ Uncertainty Quantification in Deep Learning

Uncertainty quantification is vital in understanding the reliability of AI model predictions. Traditionally, the frequentist approach assumes a single point estimate of network weights, using estimated class likelihoods as confidence measures for predictions. However, studies have shown that these likelihoods often overestimate accuracy <cit.>, and the popular metric used to quantify confidence, expected calibration error, has been criticized for bias and inconsistency <cit.>. This limits utility, motivating the need for alternative approaches to accurately quantify model uncertainty rather than just confidence.

Predictive uncertainty comprises of aleatoric and epistemic components <cit.>. Aleatoric uncertainty accounts for inherent noise in observations, while epistemic uncertainty arises from limited training data and model parameter uncertainty. Recent advancements in Bayesian inference and Bayesian neural networks have provided robust frameworks to quantify both forms of uncertainty by estimating posterior distributions over model weights. Gal and Ghahramani <cit.> introduced Monte Carlo (MC) dropout for Bayesian inference in deep learning, leveraging dropout in convolutional layers for stochastic forward passes to approximate Bayesian variational inference. Bayesian approximation using MC dropout has been extensively applied: Kendall and Gal <cit.> developed Bayesian SegNet for scene understanding, while Dechesne and Lassalle <cit.> used it in U-Net for high-accuracy image segmentation. Seedat <cit.> proposed a human-in-the-loop system using model uncertainty. Yet, single-model architectures are now supplanted by model ensembles due to difficulties in capturing inherent variability.




 §.§ Multistage Neural Network Ensembles

Training multiple individual models through model ensembling offers a range of tools to encourage member model diversity and improve overall accuracy. Diversity can be achieved through several techniques, the most popular being bagging <cit.>, stacking <cit.>, and boosting <cit.>. Non-Bayesian deep ensembles improve predictive uncertainty estimation compared to single model architectures <cit.>, yet remain limited by naïve aggregation strategies such as simple or weighted averaging or majority vote <cit.>. Yang et al. <cit.> addresses drawbacks of traditional ensembles by training a secondary neural network to adaptively assign weights, inspired by stacked generalization. This approach offers flexibility in candidate selection and leverages non-linear modeling capabilities of neural networks. 


Lai et al. <cit.> designed a multistage reliability-based neural network ensemble learning approach for credit card evaluation based on a decorrelation maximization procedure to select diverse models. Their approach combines outputs using different aggregation strategies, achieving superior accuracy over single and hybrid models. Multistage ensembling with U-Nets also proves beneficial; Yin and Hu <cit.> designed Paw-Net, a multistage ensemble for semantic segmentation. Paw-Net demonstrates higher final Intersection-over-Union (IoU) scores by integrating outputs of multiple U-Nets specialized in different classes. 

Our proposed MSU-Net architecture combines the benefits of these appro­aches: candidate U-Nets are first overproduced and then selected based on decorrelation to compose an ensemble. The final stage combines outputs using a Monte Carlo U-Net approach to produce accurate segmentation maps.






§ METHODS



 §.§ Monte Carlo U-Net (MCU-Net)

We introduce stochasticity into our inference process by incorporating dropout layers into our chosen U-Net architecture, thereby enabling MC dropout, as in MCU-Net  <cit.>. MCU-Net is henceforth referred to as our baseline model. Indiscriminately placing dropout after each convolutional layer can lead to reduced fit and poor test performance <cit.>. We opt to situate dropout layers in the decoder section of our U-Net instead. Each decoder block consists of two sets of convolutional layers (3x3 convolution filter, batch normalization, ReLU activation) and an attention block. This enables us to approximate Bayesian inference by conducting T forward passes, or MC samples, of the U-Net during testing. Our empirical results show no significant improvement beyond T=30. Our model outputs both logits and logit variances <cit.> to use for segmentation and epistemic uncertainty maps,



with  [p_t, σ^2_t] =f^ω_t(x) representing the t-th forward pass of MCU-Net with learned weights ω_t. Given sigmoid activation, σ_SIG, the ensemble prediction is  aggregated by averaging individual model outputs, σ_SIG(p_t),  across all T samples. Similarly, the raw epistemic uncertainty map is obtained by averaging logit variances across all MC samples.




 §.§ Multistage U-Net (MSU-Net)

The MSU-Net architecture, illustrated in Fig <ref>, is structured in three main stages: (1) ensembling, (2) candidate selection, and (3) combining. In stage 1, we train multiple models referred to as candidates. Bagging is employed to induce diversity into our training data; given the training dataset of size n_TR, we repeatedly sample with replacement to generate M training subsets, each also of size n_TR. We then proceed to train a total of 15 bootstrapped models <cit.>. Hyperparameter optimization, specifically determining the number of epochs of training for each candidate model, is performed using the VS1 validation set, held out for this purpose. Early stopping is applied to each candidate when validation loss begins to increase, to mitigate model overfitting. 

In stage 2, candidates are selected to form a diverse and efficient ensemble, aiming to reduce computational costs by minimizing correlation. We apply the decorrelation maximization method from <cit.>, using the Brier score loss matrix computed from models f^ω_1_1, f^ω_2_2, ⋯, f^ω_M_M on a separate validation set VS2 within the specified region of interest, replacing binary error with Brier score for finer error measurement. The mean, variance, and covariance of the Brier score loss matrix is then computed in order to construct the correlation matrix R, where the ith row and jth column represents the degree of correlation between f^ω_i_i and f^ω_j_j. We write R in a block form for candidate model f^ω_i_i as


    R [  R_-i   r_i; r_i^T     1 ]


where R_-i is the principle submatrix of R resulting from deleting the ith row and ith column. Subsequently, we compute the plural-correlation coefficient <cit.>

    ρ_i^2 = r_i^T R^-1_-i r_i

by which candidate f^ω_i_i is kept if ρ_i^2 ≤θ for some threshold θ, else discarded, in order to build our final ensemble. Empirical trials testing θ thresholds indicate that K=3 performs almost equally as well as K=15, suggesting that minimal additional training will suffice for our new architecture.






















In stage 3, we train a final MCU-Net combiner taking ensemble member outputs as inputs and outputting segmentation maps. We compute the segmentation and epistemic uncertainty maps of MSU-Net analogous to MCU-Net.





§ EXPERIMENTS



Dataset and Training  
We used a ultrasound scanning system described by Morales et al. <cit.>  
to scan the CAE Blue Phantom anthropomorphic gel model simulating femoral vessels. Equipped with a 5MHz linear transducer, the system can scan up to 5cm in depth, producing 2D transverse ultrasound images. Expert clinicians annotated these images using the Computer Vision Annotating Tool (CVAT) <cit.>, followed by cropping and resizing to 256 × 256 pixels. The dataset was split into training (1392 images), validation (907 images), and testing (856 images). The validation set was further randomly split into two disjoint sets, VS1 and VS2.

For image segmentation, we employed U-Net architectures with a ResNet34 backbone in Pytorch, utilizing the Segmentation Models library <cit.> pretrained on ImageNet. To integrate MC Dropout, dropout layers with optimal rates of 0.4 and 0.5 were added after each ReLU activation  in the decoder <cit.>. Training used a batch size of 8, Adam optimizer with a learning rate of 0.0001, and early stopping based on validation loss stabilization. MSU-Net incorporated bagging for ensembling and the plural-correlation coefficient as a correlation metric.







































Distribution Divergence Estimation
Model quality hinges on uncertainty distribution: low for correct, high for incorrect predictions, enhancing calibration and user trust <cit.>. We utilize the Rényi divergence (RD) statistic as our metric for comparison. It is a generalization of Kullback-Leibler (KL) divergence that measures the dissimilarity between two probability distributions p and q. van Erven and Harremoës <cit.> relay an operational characterization of RD as the number of bits by which a mixture of two codes, p and q, can be compressed.  For some measurable set ℳ_0 that p and q lie in and order α∈ℝ∖{1 }, RD of distribution p from distribution q is defined as:


    R_α(p || q) = 1/α-1ln∫_ℳ_0 p^α(x) q^1-α(x)  dx


A nonparametric estimator of RD that is conditionally L_2-consistent using only k-nearest-neighbor statistics has been proposed to considerably reduce computational effort <cit.>. Let X_1:n_0=(X_1, ⋯, X_n_0) be an i.i.d. sample from a distribution with density p and Y_1:n_1=(Y_1, ⋯, Y_n_1)  an i.i.d. sample from a distribution with density q. We denote ρ_k(i) to be the k-th nearest neighbor of observation X_i in X_1:n_0 and v_k(i) the k-th nearest neighbor of X_i in Y_1:n_1. With B_k, α = Γ(k)^2/Γ(k-α+1) Γ(k+α-1), we can estimate the RD by:


    R_α(p || q) = 1/α-1ln( n_0^-1∑_i=1^n_0((n_0-1) ρ_k(i)/n_1 v_k(i))^1-α B_k,α)


We use RD to measure the ability of our method to distinguish correct predictions (p) from incorrect predictions (q). We conduct permutation testing to assess deterministicity and bootstrapping to obtain confidence intervals for the results. Nearest neighbor estimators are sensitive to perturbations in the underlying distribution, and hence their limited variance cannot be consistently estimated by a naïve Efron-type bootstrap <cit.>. 
Since this behavior may result in a non-negligible positive bias in bootstrap estimates, we instead apply a direct M-out-of-N (MooN) type bootstrap <cit.> for this metric, shown in Algorithm <ref>.









§ RESULTS AND DISCUSSION


We evaluate MSU-Net using quantitative and qualitative metrics. To alleviate heavy class imbalance, we choose to average metrics over a predefined region of interest (ROI), shown in  Fig. <ref>. Fig. <ref> displays model uncertainty distributions on the test set. We employ kernel density estimation with a Gaussian kernel and optimal bandwidth via Silverman’s rule of thumb to visualize difference in means between correct and incorrect prediction uncertainty distributions over 100,000 samples. To quantify separation between these distributions with RD, we select k=4 for k-nearest neighbors and α=0.85 for numerical stability and interpretability <cit.>. We perform B=1000 permutations and bootstrapped samples. For MooN-type bootstrap, γ=0.8 is selected to maintain the largest proportion of original data while achieving the closest coverage probability of 0.95 for 95% intervals. Our final results are displayed in Table <ref>.

 


Fig. <ref>a shows training behavior for both models. MSU-Net shows a more stable convergence and consistently outperforms MCU-Net during validation performed after each epoch. MSU-Net achieves 18.1% better mean IoU and a significant improvement in sensitivity and false negative rate scores at alpha level 0.05, while other metrics remain similar. We additionally utilize precision-recall curves, seen in  Fig. <ref>b, which are resilient to unbalanced classes since they only focus on positive class predictions. Performance results are displayed in Table <ref>.












Permutation tests for MCU-Net and MSU-Net (p ≤ 0.003 for both) indicate that the observed separation between correct and incorrect predictions is not random in both. Yet, at the 95% confidence level, we see no overlap between their 95% confidence intervals, showing that the ability of MSU-Net to distinguish correct from incorrect predictions is significantly better than that of MCU-Net.

Qualitative uncertainty maps visually validate our findings and capture local variations in model performance <cit.>. MCU-Net exhibits indiscriminately low uncertainty in Fig. <ref>, whereas MSU-Net provides more interpretable uncertainty values. MSU-Net shows clearer vessel tops with lower uncertainty and higher uncertanties at the bottom of vessels which is consistent with our assessment. 
MSU-Net yields a higher credibility in assessing its predictive reliability. 



MSU-Net generally outperforms MCU-Net. Although it performs marginally worse at specificity and false positive rate (FPR), the precision-recall curves from Fig. <ref>b show a higher average precision score for MSU-Net at 0.945 than MCU-Net at 0.875 compared to the baseline score of 0.09 for a `no-skill' classifier. As such, at the same level of recall, MSU-Net correctly classifies a greater proportion of pixels that are actually vessels than MCU-Net. Crucially, MSU-Net achieves a considerably lower false negative rate (FNR) than MCU-Net. In this context, failing to anticipate a true vessel can have more catastrophic consequences to a critically injured individual than incorrectly anticipating a vessel. MSU-Net improves credibility not only through higher quality results, but also through more accurate results while avoiding potentially disastrous deficiencies of predictive modeling in the context of medical image segmentation.




§ CONCLUSION

This paper introduces MSU-Net, a multistaged Monte Carlo U-Net for uncertain ultrasound image segmentations. It improves model transparency and trustworthiness compared to standard U-Net by enhancing uncertainty evaluation, despite minimal additional training. Our framework sets a benchmark for future studies, offering qualitative maps for intuitive model assessment. While our preliminary results are limited to phantom data and binary image segmentation, we aim to validate these findings on live animal and human data in future work.





  §.§.§ 
 Authors thank Nico Zevallos, Dr. Michael R. Pinsky, and Dr. Hernando Gomez for gathering experiment data. This work was partially supported by the U.S. Dept. of Defense contracts W81XWH-19-C0083 and W81XWH-19-C0101.



  §.§.§ 

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








 splncs04
 























 Paper Title: Alleviating Hallucination in Large Vision-Language Models with Active Retrieval Augmentation.








xiaoye@hust.edu.cn

  Huazhong University of Science and Technology
  China



chenqiyuan1012@gmail.com

  Zhejiang University
  China



weiw@hust.edu.cn

  Huazhong University of Science and Technology
  China



gasolsun36@gmail.com

  Xiamen University
  China





dongjf24@gmail.com

  Zhejiang Gongshang University
  China







§ ABSTRACT 
Despite the remarkable ability of large vision-language models (LVLMs) in image comprehension, these models frequently generate plausible yet factually incorrect responses, a phenomenon known as hallucination. 
Recently, in large language models (LLMs), augmenting LLMs by retrieving information from external knowledge resources has been proven as a promising solution to mitigate hallucinations. 
However, the retrieval augmentation in LVLM significantly lags behind the widespread applications of LVLM. 
Moreover, when transferred to augmenting LVLMs, sometimes the hallucination degree of the model is even exacerbated. 
Motivated by the research gap and counter-intuitive phenomenon, 
we introduce a novel framework, the Active Retrieval-Augmented large vision-language model (ARA), specifically designed to address hallucinations by incorporating three critical dimensions: 
(i) dissecting the retrieval targets based on the inherent hierarchical structures of images. (ii) pinpointing the most effective retrieval methods and filtering out the reliable retrieval results. (iii) timing the retrieval process to coincide with episodes of low certainty, while circumventing unnecessary retrieval during periods of high certainty. To assess the capability of our proposed ARA model in reducing hallucination, we employ three widely used LVLM models (LLaVA-1.5, Qwen-VL, and mPLUG-Owl2) across four benchmarks. 
Our empirical observations suggest that by utilizing fitting retrieval mechanisms and timing the retrieval judiciously, we can effectively mitigate the hallucination problem. 
We hope that this study can provide deeper insights into how to adapt the retrieval augmentation to LVLMs for reducing hallucinations with more effective retrieval and minimal retrieval occurrences.



†The first two authors contribute equally.

*Corresponding author.

































Alleviating Hallucination in Large Vision-Language Models with Active Retrieval Augmentation
    Jianfeng Dong
    August 5, 2024
============================================================================================





§ INTRODUCTION


Recently, Large Vision-Language Models (LVLMs)  <cit.> have gained widespread application across various scenarios due to their significant capacity to generate contextually appropriate text for visual inputs. 
Benefiting from the advancements in model architecture, training methods, and data diversity, LVLMs exhibit superior performance for various tasks such as visual question answering <cit.> and image captioning <cit.>.
Despite these advancements, LVLMs still suffer from a significant challenge termed “hallucination”, whereby the models produce semantically plausible but factually inaccurate text, misaligned with the ground-truth content of the associated image  <cit.>.
This problem damages the practical employment of LVLMs, particularly in fields that require accurate content generation, such as medical <cit.> and robotics <cit.>, potentially resulting in severe consequences. 
Hence, addressing the hallucination problem is paramount for bolstering the credibility of LVLMs across real-world scenarios.

To mitigate hallucinations in LVLMs, a series of attempts have been recently proposed and can be broadly classified into two categories. 
The first class of approach retrains the LVLMs with constructed hallucination-related datasets by
supervised fine-tuning (SFT)  <cit.>, or Reinforcement Learning from Human Feedback (RLHF) <cit.>.
Although effective, these methods introduce significantly additional training costs. 
The other solutions focus on designing more robust decoding strategies <cit.>. For example, VCD  <cit.> introduces visual contrastive decoding to contrast output distributions derived from original and distorted visual inputs. 
In this way, the LVLM reduces the over-reliance on statistical bias and unimodal prior.
Although these methods are training-free, they still suffer from the limitations of LVLMs’ static parametric capacity. 
Recently, in the realm of large language models (LMs), augmenting LLMs <cit.> by retrieving information from external knowledge resources has shown promise in reducing language hallucinations. Furthermore, this retrieval augmentation serves as a flexible way to extend beyond the model's inherent knowledge without the burden of extensive training costs.
In this paper, we aim to propose a novel framework to augment LVLMs with external knowledge by introducing an Active Retrieval-Augmented large vision-language model (ARA) for mitigating hallucinations. As shown in Figure <ref>, given an input image and query, the LVLM tends to produce hallucinated answers. Instead, our model can accurately identify the most relevant pairs and output the correct answer.



Concretely, our proposed ARA is grounded in three critical dimensions for mitigating hallucinations in LVLMs. 
Firstly, given the inherently hierarchical nature of images, simple full-image retrieval may result in noise and irrelevant outcomes, thus it is imperative to decompose the target object causing hallucination from the input image for more accurate retrieval. 
Secondly, in contrast to augmenting LLMs, LVLMs offer diverse retrieval modalities due to their multimodal inputs. 
In this way, discerning the most effective retrieval technique and securing trustworthy results is crucial for ensuring retrieving performance. 
Lastly, excessive retrieval activations may lead to undue time expenditure and unnecessary retrieval.
By initiating retrieval only during periods of low LVLM certainty and knowledge deficiency, one can circumvent unnecessary or improper information retrieval.


Based on the above analysis, given the input image and query, our proposed ARA first devises a coarse-to-fine retrieving paradigm. 
The target object is first extracted from the input query and subsequently positioned within a particular region of the input image. To build a robust retrieving system, we simultaneously engage in coarse-grained and fine-grained retrieval process which retrieves the full image as well as the specific regions pertaining to the target object. Recognizing the plurality of retrieval strategies due to the multimodal input, an extensive analysis of different methods is conducted to determine the most optimal retrieving approach. 
Subsequent to acquiring pertinent text and image pairs from external databases, a reranking strategy is employed to eliminate unreliable outcomes, thereby augmenting LVLMs with the appropriate external information.
Particularly, retrieval operations are circumvented when the LVLM exhibits high certainty, to prevent the inclusion of superfluous retrieval that may introduce redundant information. 
In this context, retrieval is activated using a difficulty metric that relies on the mutual information between multimodal inputs.

To assess the efficacy of ARA, empirical evaluations are performed using three prevalent LVLMs across four benchmarks related to hallucination challenges.
The promising results demonstrate that our framework can effectively mitigate hallucinations. 
To sum up, our contributions are summarized as follows:


    
  * Our research thoroughly investigates the three critical dimensions of augmenting LVLMs with external knowledge, including accurately dissecting the retrieval targets, pinpointing the most effective retrieval methods, and the sensible triggering of the retrieval process. 
    
    
  * We introduce an Active Retrieval-Augmented large vision-language model (ARA) for mitigating hallucination. 
    Our empirical evidence indicates that, by employing appropriate retrieval mechanisms and carefully timing the retrievals, we can significantly address the issue of hallucination while maintaining a moderate frequency of retrieval.

    
  * Through comprehensive experiments, we demonstrate the superior performance of our proposed ARA in alleviating hallucinations using three well-known LVLMs across four hallucination-related benchmarks. The results in our paper also suggest that with optimal retrieval settings, the potential of augmenting LVLMs can be more effectively harnessed.




§ RELATED WORK




 §.§ Large Visual-Language Models


Recently, large-scale models have attracted significant attention and gained widespread applications <cit.>. 
LVLMs significantly enhance the interaction between humans and AI in a more natural and intuitive
ways and demonstrate remarkable capabilities in understanding and generating multi-modal content.
With the aid of advanced Large Language Models like LLaMA <cit.> and Qwen <cit.>, a batch of LVLMs such as LLaVA-1.5 <cit.>, Qwen-VL <cit.>, and mPLUG-Owl2 <cit.> have emerged, which can comprehend and generate a wide array of content by utilizing information from distinct modalities like texts and images. 
The training of LVLMs is divided into two critical phases: feature alignment and instruction fine-tuning. The former aims to model visual and textual inputs coherently, while the latter strives to enhance the ability of LVLMs to comprehend user queries. 
Despite the advancements, these LVLMs continue to face significant challenges with hallucination issues. Thus, in this paper, we focus on solving hallucination problems to promote the use of LVLMs in practical scenarios.



 §.§ Hallucination in LVLMs


In LVLMs, hallucination refers to models that generate seemingly plausible outputs inclusive of objects,
attributes, or relations that do not correspond with the images <cit.>.
Approaches to address hallucinations in LVLMs have largely fallen into two camps. The first emphasizes model refinement through supervised fine-tuning or application of reinforcement learning techniques. For instance, LRV <cit.> seeks to diminish hallucinatory outputs by utilizing broad and varied supervised fine-tuning datasets. 
For methods based on reinforcement learning, LLaVA-RLHF <cit.> is the pioneer in applying Reinforcement Learning with Human Feedback (RLHF) to mitigate hallucination in LVLMs. 
This approach is further refined by RLHF-V <cit.>, which incorporates detailed corrective human feedback.
Considering the instability and training difficulty of RLHF, Zhao et al. <cit.> employ Direct Preference Optimization (DPO) and build a hallucination-aware dataset for alleviating hallucination. 
Although 
While these methods have yielded notable advancements, they also introduce significant training overheads and are prone to overfitting the training data. 
The second category consists of training-free strategies designed to circumvent hallucination without incurring additional training expenses.
VCD <cit.> endeavors to rectify the model's over-reliance on linguistic priors and statistical leanings by comparing the output distributions from both unaltered and visually perturbed inputs.
OPERA <cit.> mitigates overconfidence in model predictions by integrating a penalty term into the model logits during beam search decoding, further enhanced with a strategic rollback feature. 
Although these methods are effective, they still suffer from the limitations of LVLMs’ static parametric capacity. 
In this paper, different from the above two paradigms, we focus on augmenting LVLMs with external knowledge to mitigate hallucination in LVLMs.  



 §.§ Retrieval-Augmented Generation


In the realm of large language models (LLMs), Retrieval-Augmented Generation (RAG) has been widely used and shown promising results in mitigating hallucinations. RAG in LLMs retrieves relevant information from an external knowledge base before LLMs respond to a query <cit.>, thereby enabling them to collaboratively generate responses by leveraging the retrieved external non-parameterized knowledge alongside their internal parameterized knowledge. 
However, its application in LVLMs has not been thoroughly explored. Since LVLMs present a multimodal nature, findings pertinent to LLMs cannot be indiscriminately extrapolated to LVLMs. 
Predominantly, existing research on retrieval augmentation in multimodal tasks has been confined to image captioning <cit.> or image generation <cit.>, with a narrow focus that overlooks the broader implications of RAG in LVLMs, especially concerning its capacity for hallucination reduction.
In this paper, our research conducts an extensive evaluation of RAG's applicability in LVLMs to mitigate hallucination by tailoring retrieval strategies and optimizing active retrieval processes. 
It is worth noting that the active retrieving strategy has also been discussed in <cit.> in LLM. However, they propose a token-level confidence for triggering the retrieval, which is not applicable in LVLMs and we will investigate it in the ablation study.





§ METHOD


In this section, we first describe the generation procedure of the LVLMs, laying the groundwork for comprehending our approach. Subsequently, we will introduce our Active Retrieval-Augmentation large vision-language model (ARA) in detail. As shown in Figure <ref>, given the input image and query, the LVLM first analyzes them based on mutual information and determines whether to trigger the retrieval process. If the retrieval is necessary, our ARA will follow a coarse-to-fine paradigm to dissect the retrieval targets and conduct jointly coarse-grained and fine-grained retrieval. Following this, a reranking of the retrieval outcomes is performed for additional refinement. Finally, the LVLM harnesses this externally sourced knowledge to generate the final response.





 §.§ LVLM Input and Decoding

 
The input of LVLMs contains both image and text. Typically, the LVLMs use a vision encoder (e.g. CLIP) to extract visual tokens from the input images, and then map them into the input space of LLMs with a linear projection module. 
Subsequently, the mapped visual tokens are used as part of the LLM input, along with the text input. 
In this paper, we denote the visual tokens as 𝐱^v = {x_0,x_1,…,x_N-1}. Here N is the length of the visual tokens and it is a fixed number in most cases.  
Correspondingly, the input text is tokenized and we denote it as 𝐱^p = {x_0,x_1,…,x_M-1}.
The image and text tokens are concatenated as the final input sequence and we denote it as  {x_i}_t=0^T-1 that T=N+M.


During decoding, LVLM gets the probabilities for the next token prediction, formally:

    p(x_t|x_<t) = SoftMax[LVLM(h_t)]_x_t,    x_t∈𝒱,

where h_t denotes the output hidden states at time step t for the input image and query, x_<t to simplify the sequence {x_i}_i=0^t-1 and 𝒱 means the whole vocabulary set. For greedy decoding, the token with the highest probability will be chosen.




 §.§ Active Triggering Retrieval


As shown in Figure <ref>, given the input image and query, our ARA model first determines whether the input pair requires retrieval. In this way, we can circumvent unnecessary retrieval during periods of high certainty. 
In this section, we propose three active retrieval methods to explore the difficulty metrics that influence when to retrieve, including Confidence-aware Active Retrieval, 
Question-aware Active Retrieval, and Image-aware Active Retrieval. The first method relies on the confidence of the model output, and the latter two matrices depend on the mutual information between the inputs.
An ideal metric should effectively reduce the frequency of retrievals and can easily apply to different LVLMs.

Confidence-aware Active Retrieval. We feed the image and query into LVLMs to generate the answer. If the LVLM is confident about the answer, we accept it without retrieving additional information. Specifically, we actively trigger retrieval if any token of the answer has a confidence lower than the threshold θ. In this method, θ=0 indicates retrieval will never be used and θ=1 means that all inputs will use the retrieval. This strategy is inspired by <cit.> in LLM.

Query-aware Active Retrieval. To better evaluate the difficulty of queries from a view of language bias, we evaluate the degree of LVLMs answering relying on the input image. To this end, we define a difficulty metric using mutual information based on the difference below:


    M_ij = logP(a_ij | V_i, Q_i)/P(a_ij) - logP(a_ij | Q_i)/P(a_ij) = logP(a_ij | V_i, Q_i)/P(a_ij | Q_i).


where P(a_ij | Q_i) indicates the output probability with only query input to the LVLMs and P(a_ij | V_i, Q_i) means the output probability with both query and image as model input. The higher M_ij indicates more visual reasoning is needed to answer. M_ij<0 illustrates that the answer excessively depends on the input query. Thus, we define a threshold θ and trigger retrieval if M_ij<θ, where insufficient visual information is provided for answering the query. Incorporating retrieved information will enhance the confidence of LVLMs to answer.  

Image-aware Active Retrieval. Similar to the above query-aware active retrieval, we use the probability difference between the original and noisy image to evaluate the image information to answer the query. During the implementation, we use the method in <cit.> to add noise.


    D_ij = logP(a_ij | V_i, Q_i)/P(a_ij) - logP(a_ij | V'_i, Q_i)/P(a_ij) = logP(a_ij | V_i, Q_i)/P(a_ij | V'_i, Q_i).


where P(a_ij | V_i, Q_i) indicates the output probability with query and input image, P(a_ij | V'_i, Q_i) means the output probability with query and noisy image. Analogously, we define a boundary θ and trigger retrieval if D_ij<θ. 





 §.§ Coarse-to-Fine Hierarchical Retrieval


Given the inherently hierarchical nature of images, simple full-image retrieval may result in noise and irrelevant outcomes. As shown in Figure <ref>, using the full image for retrieval can not obtain the most relevant information for reasoning, namely “red shirt". Thus it is imperative to decompose the target object causing hallucination from the input image for more accurate retrieval. 
In the following, we will first describe the coarse-grained retrieval and then introduce the fine-grained retrieve. 

Coarse-grained Retrieve. 
Initially, based on the input image, we use the CLIP to extract the visual embedding. Similarly, we build the embeddings for the images and corresponding texts in the knowledge base. Then, we can retrieve the relevant images from knowledge bases with the similarity between images. Finally, the top-K image and text pairs will be fetched for the following reasoning. 
Benefiting from the multimodal encoding ability of CLIP, we can also retrieve the top-K pairs by matching the input image with texts in the database. We will compare different retrieving methods in the ablation study.


Fine-grained Retrieve. While full-image retrieval can incorporate valuable information, it may overlook fine-grained details, such as diminutive objects and intricate attribute data. Therefore, it is crucial to propose a fine-grained retrieval mechanism that concentrates on the distinct objects within an image. 
To conduct fine-grained retrieval, we initially extract key entities from the input query utilizing the large language model LLaMA2-7B <cit.>. This extraction is facilitated by in-context learning <cit.>, where demonstrations are presented to the LLM, which subsequently outputs the specific entities from the input query. 
Following this, Grounding Dino <cit.> is deployed to identify the objects in the image that correspond to the extracted entities. As illustrated in Figure <ref>, we crop the targeted object and then perform retrieval in a manner akin to the previously described coarse-grained retrieval.





 §.§ Reranking Retrieving Results


Throughout the above retrieving process, the similarity is computed by the CLIP embeddings between images. In this way, the visually similar but semantically different pairs may be recalled. 
To avoid noisy retrieving results which bring different semantic, after obtaining the retrieved pairs for both coarse-grained and fine-grained retrieval, we subsequently propose a reranking strategy designed to recalibrate the retrieved outcomes. 
Specifically, we adopt the LVLM to describe the input image and generate a detailed caption. For coarse-grained retrieval, the input full image is characterized, while describing the cropped image for the fine-grained retrieval.
In this way, we are able to derive the semantics of the input image.

Then, we extract the textual embeddings for both image captions and captions of retrieved images with the CLIP. Finally, we can rerank the retrieved pairs according to the semantic similarity between these captions. In this stage, we also explore other reranking methods, such as k-reciprocal encoding <cit.> or using an external visual entailment tool <cit.>. We will compare them in the ablation study. 






 §.§ Joint Coarse-grained and Fine-grained Decoding


After retrieving and ranking the pairs, we can conduct a joint coarse-grained and fine-grained decoding to yield the final responses. According to the different ways of using the retrieved results, we have two schemes including probability-level fusion and instance-level fusion. The former integrates the results of coarse-grained and fine-grained retrieval into a single prompt for inference, while the latter conducts separate inferences on the coarse-grained and fine-grained retrieval results, followed by a probabilistic-level weighting.

Probability-level Fusion. In this paradigm, we use the prompt for both coarse-grained retrieval and fine-grained retrieval: “Here are the image-caption pairs similar to the test image: ⟨Retrieval Pairs⟩. Based on these pairs and this image: ⟨Input Image⟩. Answer this question: ⟨Input query⟩". 
Here we denote the retrieval-augmented LVLMs as below:


    p^coarse(x_t|x_<t) = SoftMax[LVLM([R_I:h_t])]_x_t,    x_t∈𝒱,


where R_I indicates the hidden states of retrieved image-level pairs. For simplicity, we directly use [:] here. We use p^coarse to represent decoding with coarse-grained retrieval. Similarly, 


    p^fine(x_t|x_<t) = SoftMax[LVLM([R_O:h_t])]_x_t,    x_t∈𝒱,



where R_O indicates the retrieved fine-grained object-level pairs. We represent the decoding with object-level fine retrieval by p^fine(x_t|x_<t). For the probability-level fusion, we can combine them with a hyperparameter α in [0,1] to control the importance of two kinds of retrieval. 


    p^coarse&fine(x_t|x_<t) = α p^coarse(x_t|x_<t) + (1-α)p^fine(x_t|x_<t),



Instance-level Fusion. 
For this kind of fusion, we can integrate both coarse-grained retrieval and fine-grained retrieval with prompts: “Here are the image-caption pairs similar to the test image: ⟨Coarse-grained Retrieval Pairs⟩. Here are the image-caption pairs: ⟨Fine-grained Retrieval Pairs⟩ similar to the ⟨entity⟩ in the input image. Based on these pairs and this input image: ⟨Input Image⟩. Answer this question: ⟨Input query⟩." Here we extract the entity from the input query and fill it in the prompt. 
 

    p^coarse+fine(x_t|x_<t) = SoftMax[LVLM([R_I:R_O:h_t])]_x_t,    x_t∈𝒱,







§ EXPERIMENT


This section will elaborate on the evaluation of our proposed Active Retrieval-Augmented large vision-language model (ARA) across different LVLMs.



 §.§ Evaluation Metrics


POPE <cit.>, the Polling-based Object Probing Evaluation proposed a simplified method for evaluating object hallucinations. In this evaluation framework, the LVLM is tasked with determining the presence of specific objects in a given image, with a balanced distribution of queries for object presence and absence. The evaluation protocol encompasses three distinct sampling settings: random, popular, and adversarial, each employing unique strategies for constructing negative samples. In the random setting, objects not present in the image are chosen randomly. The popular setting selects missing objects from a high-frequency pool, whereas the adversarial setting prioritizes coexisting objects that could be easily mistaken for existing objects in the image. 
The POPE benchmark aggregates data from three diverse sources: MSCOCO <cit.>, A-OKVQA <cit.>, and GQA <cit.>. Within each dataset, every sampling setting comprises 500 images, with 6 questions posed for each image, resulting in a total of 27,000 question-answer pairs derived from the development sets of these datasets. The evaluation predominantly assesses metrics such as accuracy, precision, recall, and F1 score.

MME  <cit.> is a comprehensive benchmark designed specifically for evaluating LVLMs on multiple dimensions. It includes ten sub-tasks related to perception and four sub-tasks emphasizing cognition. According to previous studies <cit.>, we utilize the existence and count subsets for object-level hallucination evaluation, and the position and color subsets for attribute-level hallucination evaluation. Performance is quantified by accuracy and accuracy+, which is also the official evaluation method.


MMStar  <cit.>, foundational to vision-reliant multi-modal evaluation, consists of 1,500 challenge samples meticulously chosen by human experts. It is tailored to systematically benchmark six key capabilities alongside eighteen granular facets, aiming to rigorously assess the multi-modal proficiency of LVLMs using a selection of samples that have been carefully refined and equilibrated.


MMbench  <cit.> is collected from multiple sources, including public datasets and the Internet, and currently, contains 2974 multiple-choice questions, covering 20 ability dimensions. Similar to MME, we focus on hallucination-related ability. Thus, we select 4 ability subsets including object localization, attribute recognition, spatial relationship, and action recognition as our evaluation benchmarks. These subsets form a comprehensive assessment of hallucinations, addressing issues at the object-level, attribute-level, and relation-level comprehensively.





 §.§ Implementation Details
 

In the case of coarse-grained retrieval, we use the COCO dataset <cit.> as the retrieval database. For fine-grained retrieval, we use the VisualGenome <cit.> to build the fine-grained retrieval database as it provides a detailed caption for each grounding box in the image. CLIP-Large is used to extract image embeddings for retrieval.   
To decompose target objects from the input query, we extract key entities from the input query utilizing the large language model LLaMA2-7B <cit.>. 
Grounding Dino <cit.> is deployed to identify the objects in the image that correspond to the extracted entities. If the objects can not be located in the image, we only use coarse-grained retrieval for the following reasoning process.
In this paper, we implement our ARA model on three LVLMs, including LLaVA 1.5, Qwen-VL, and mPLUG-Owl2 to evaluate the effectiveness of our method. For LLaVA 1.5 and mPLUG-Owl2, we use the retrieved texts to augment, while we apply the retrieved image and text pairs to boost Qwen-VL. For these three LVLMs, we implement augmentation using 3, 4, and 5 retrieval components respectively. For active retrieval, we use the query-aware active retrieval. For the hyperparameter α presented in Equation 6, we apply a value of 0.8 for the POPE, MMstar, and the object-level subsets within the MME benchmark, as well as for OL, SR, and ACR in the MMbench. A value of 0.4 is utilized for the remaining evaluation subsets. For active triggering retrieval, we use query-aware active retrieval for its simplicity and effectiveness. The threshold of active triggering is determined by grid search. 

In our experiments, to have a strict comparison of our RAR and vanilla LVLMs, we consistently utilize greedy decoding instead of other probability sampling methods. With this strategy, the token with the highest probability from the post-softmax distribution is directly selected as the next token. 



 §.§ Experiment Results



Results on POPE. Experimental results on POPE under the random, popular, and adversarial settings are summarized in Table <ref>. This benchmark mainly focuses on the object-level hallucination. 
Specifically, the performances of our ARA consistently surpass the baseline results on all of the LVLMs. 
This suggests its pivotal role in augmenting LVLMs with retrieval, thereby reducing instances of object-level hallucination. 
In a more detailed model-specific analysis, RAR demonstrates varied effects across different LVLMs. For all three LVLMs, the F1 score elevation is predominantly driven by a recall boost (e.g., up to 10 points for mPLUG-Owl2), showcasing retrieving to external knowledge base can effectively help detect object presences.








Results on MME Hallucination Subset. 
The MME subset evaluations extend beyond POPE’s scope, encompassing both object-level and attribute-level hallucination. Results in Table <ref> show that our ARA leads to a uniform enhancement in addressing object-level hallucination for all models, except the count score of Qwen-VL. Meanwhile, our ARA demonstrates significant improvements in attribute-level Color scores, contributing to substantial overall performance gains. These improvements are attributable to our designed coarse-to-fine retrieval paradigms, which effectively focus on the target object.

Results on MMStar. In a more difficult benchmark MMStar which rigorously assesses the multi-modal proficiency, as depicted in Table <ref>, our method ARA secures a notable enhancement across all six subsets. For instance, LLaVA-1.5 realizes an average improvement of 8.8 points. Unlike POPE, which may include statistical bias, answering questions in MMStar necessitates reliance on image information. Hence, our retrieval augmentation method introducing external image and text knowledge significantly bolsters model performance.

Results on MMBench Hallucination Subset. 
This benchmark extends our hallucination evaluation to incorporate the relation level in addition to the above object and attribute levels, providing a more holistic assessment for reducing hallucinations. As evidenced in Table <ref>, benefiting from coarse-grained retrieval, our ARA manifests discernible enhancements in the object and relation-level subsets, such as OL, SR, and ACR. For instance, our ARA model surpasses the vanilla LLaVA-1.5 with 5.08 percent points on object localization (OL).  
Concurrently, the integration of fine-grained retrieval into our model facilitates superior performance in attribute-level subsets, namely attribute recognition (ATR). 












§ ABLATION STUDY


In this section, we undertake a series of ablation studies with the aim of identifying the critical elements integral to constructing a resilient and efficient retrieval-augmented large vision-language model.



 §.§ Ablation of Coarse-to-Fine Retrieval



In this section, in order to explore the effects of different factors on coarse-to-fine retrieval in a simple and effective way, we did not use the rerank method or active trigger strategy.


Which kind of retrieval is better? 
Due to the multimodal characteristics of the input, there are various retrieval methods at our disposal. As shown in Table <ref>, we attempt to use the input query and image to search for images and descriptions in the database. For this study, we only conduct coarse-grained retrieval and we use the COCO dataset as the database which contains a caption for each image. 
As a result, there are four retrieval methods available to us. 
As reflected in the results within this table, we adopt the input image as our primary retrieval modality for images in the database throughout our series of experiments.
Surprisingly, using the input query to search the database for image leads to substantially poorer outcomes, achieving only 51.6% accuracy. 
This lower performance may stem from the fact that, despite the presence of relevant keywords, the images retrieved this way tend to be surrounded by a significant degree of noise.

Additionally, we investigate the influence of different embedding models on the efficacy of image-to-image retrieval. 
In our comparison, we examine three models: CLIP-Base/Large <cit.> and DINOV2-Base <cit.>. Among them, we notice that CLIP-base performs better than DINO (86.33 vs 86.07) and this may be related to their different training methods. Moreover, when comparing CLIP models of various sizes, we observe that the larger-scale CLIP models contribute to more pronounced improvements in performance. Therefore, in the context of this study, CLIP-Large is selected as the image embedding model for retrieval purposes.







How to augment the LVLM? After retrieving, we can obtain retrieved images and paired captions. In this section, we explore how to augment the LVLM with this external information. For a fair comparison, in this study, we fix the number of retrieved pairs to 3 for both coarse-grained and fine-grained retrieval. As demonstrated in Table 6, when using coarse-grained retrieval, we observe that the image and caption pair perform worse than the single caption (86.93 vs 87.43). 
However, the opposite results are observed in fine-grained retrieval and the Fine (I+T) achieves 86.53 compared with 86.20 of Fine (T). These results imply that
retrieval augmentation is sensitive to specific configurations. In this experiment, we do not carefully compare with Coarse (I) or Fine (I) as we have found that the LVLMs we use in our experiments are still quite weak in multi-image reasoning.

In Section 3.5, we propose a joint coarse-grained and fine-grained decoding, designed to optimally utilize the two retrieval mechanisms. 
In Table <ref>, we also investigate the effectiveness of probabilistic-level fusion denoted as Coarse & Fine and instance-level fusion labeled as Coarse+Fine. 
As shown from the results, we observe that Coarse+Fine even performs than Coarse retrieval, signifying that simply integrating the coarse-grained and fine-grained retrieved results in a prompt fails to enhance performance, given the distinct granularity of the information retrieved through these two approaches. Instead, the probabilistic-level fusion Coarse & Fine achieves better accuracy than both two kinds of retrieval. 
Thus, in our paper, we use probabilistic-level fusion for the following experiments. It is worth noting that Coarse & Fine (T) performs better than Coarse & Fine (I+T) for LLaVA-1.5 but Coarse & Fine (I+T) achieves better for Qwen-VL (we will describe in Figure 3 later). This result further validates the necessity of carefully selecting retrieval configurations to obtain the most favorable experimental results.

How much external information do we need? 
With the above retrieval, we can obtain multiple image-caption pairs from the database. This research delves into the efficacy of varying retrieved pair configurations. 
Considering that different LVLMs may be sensitive to retrieval configurations, our study incorporates both (T) and (I+T) configurations for a thorough evaluation. The analysis is constrained to examining between one to five retrieved pairs due to the input token limitations inherent in contemporary LVLMs. 
As depicted in Figure <ref>, LLaVA (T) reaches peak accuracy with three retrieved texts, whereas mPLUG (T) begins to saturate when there are four retrieval texts.
Conversely, QwenVL (I+T) demonstrates improved outcomes with a greater number of retrieval pairs; however, it still lags considerably behind the performance achieved by the other two LVLM models. 












 §.§ The effectiveness of Reranking Method


To substantiate the efficacy and indispensability of reranking, we begin by providing a case study within the POPE benchmark. 
As depicted in Figure <ref>, in instances where an image portrays people in a pool, leveraging CLIP for image retrieval may inadvertently introduce irrelevant content due to its predominant emphasis on visual congruence.
For instance, it might retrieve an incongruent image captioned “fruit and a bowl of bread sit on the table." Therefore, an adept reranking methodology can proficiently filter and organize the retrieved images in a more relevant manner. In this particular scenario, our reranking strategy aligns the input image caption with those of the retrieved images. 
By assessing these similarity metrics, we are able to reorganize the pool of retrieved images into a more accurately ranked series, thereby enhancing the precision of the retrieval process.

In order to further quantify the effect of reranking and compare different reranking methods, we conduct experiments in Table <ref>. 
The result reveals that the omission of re-ranking precipitates a decline in accuracy, dropping from 87.17% to 86.93%
This suggests that re-ranking effectively mitigates noise, particularly when employing the top three re-ordered retrieval pairs. 
For comparative purposes, we executed two alternative methodologies. The k-Reciprocal approach <cit.>, a prevalent offline person re-identification method, re-ranks images using the k-Reciprocal nearest neighbor algorithm. While this method offers a swifter execution by circumventing the captioning process of input images, its efficacy slightly underperforms our strategy. Additionally, we examined the employment of the external tool VSR <cit.>, which determines the entailment between the caption of the input image and the retrieved images. As shown in Table 7, this method gains only a slight increase in accuracy (86.93 vs 86.97). This may be due to the visual entailment model not being sufficiently versatile.









 §.§ Ablation of Active Triggering Retrieval


As excessive retrieval activations may lead to undue time expenditure and unnecessary retrieval, in this paper, retrieval is activated using a difficulty metric. In this ablation study, we compare three kinds of active retrieval methods, including confidence-aware active retrieval, query-aware active retrieval, and image-aware active retrieval. As shown in Figure <ref>, we have drawn the proportion of queries that need to be retrieved to achieve peak performance. Firstly, it is quite intuitive, the fluctuations of confidence-aware active retrieval among different models are very significant. For mPLUG-Owl2, we even need to retrieve all the queries, which indicates that confidence-aware active retrieval is not a good indicator of whether to perform a retrieval. For both query-aware active retrieval and image-aware active retrieval, the LVLMs can better determine a trigger threshold. 
Considering that query-aware active retrieval is more concise and the process of adding noise is omitted, we use query-aware active retrieval in all our experiments. 
It is worth noting that the very slight decline in the model's performance after continuing to use retrieval beyond the trigger threshold is due to the fact that the model itself has a high level of confidence, but the retrieval has introduced some redundant noise information, which we cannot filter out through reranking.



 §.§ Comparison with State-of-the-Art Method



In this paper, we augment LVLMs with external knowledge to mitigate hallucinations. 
To further demonstrate the effectiveness of our method, we compare it with VCD <cit.> and the results are demonstrated in Table <ref>. VCD is a state-of-the-art training-free method that explores intra-model knowledge with visual contrastive decoding. From the results, we can notice that our method significantly surpasses VCD in all three settings. For instance, in popular settings, our method gains a 4.06 percent improvement. It is worth noting that VCD performs comparably with the vanilla LLaVA-1.5 when using greedy decoding. This phenomenon is also observed in <cit.>. Instead, our method significantly surpasses the vanilla LLaVA-1.5 model as presented in Tabe 1. This result further verifies the robust improvement of introducing external information when applying our ARA method.







 §.§ Text Generation Quality Evaluation


In this section, we evaluated the effectiveness of our method in improving text generation quality on the MS-COCO validation set[We use the official COCO evaluation package.]. As shown in Table 9, all models LLaVA, Qwen-VL, and mPLUG exhibited significant improvements after applying our method, with an average improvement of 7.7% across the three models. Notably, the improvement was most pronounced for Qwen-VL, with an average increase of approximately 14.9%.




 §.§ Qualitative Results


To qualitatively verify the effectiveness of our ARA method on downstream tasks, we presented five examples from the POPE MSCOCO dataset and MME benchmark. As shown in Figure <ref>, we present three image-caption pairs for both coarse-grained retrieval and fine-grained retrieval. In the above three examples of determining the existence of objects, fine-grained retrieval helps to increase the model's existence information about specific objects. For the last two examples, the coarse-grained retrieval contributed effectively to capturing the overall information of the images. Thus, with both kinds of retrieval, our RAR model can effectively increase the accuracy compared with vanilla LLaVA-1.5. 








§ CONCLUSION


In this paper, we propose a novel framework to augment LVLMs with external knowledge by introducing an
Active Retrieval-Augmented large vision-language model (ARA) for mitigating hallucinations. Specifically, our proposed ARA is grounded in three critical dimensions for mitigating hallucinations in LVLMs, including coarse-to-fine retrieval for dissecting the retrieval targets, pinpointing the most effective retrieval methods, and triggering the retrieval process properly. 
Our findings indicate that by utilizing fitting retrieval mechanisms and timing the retrieval judiciously, we can effectively mitigate the hallucination problem. Our extensive experiments across four benchmarks and three widely-used LVLM confirm ARA’s efficacy in reducing hallucination.

ACM-Reference-Format







