Paper Title: Load Balancing in Federated Learning.


 


 


Load Balancing in Federated Learning
    
   Alireza Javani, and Zhiying Wang
   
                     Center for Pervasive Communications and Computing

  					University of California, Irvine

                     {ajavani, zhiying}@uci.edu
                     

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



§ ABSTRACT 
Federated Learning (FL) is a decentralized machine learning framework that enables learning from data distributed across multiple remote devices, enhancing communication efficiency and data privacy. Due to limited communication resources, a scheduling policy is often applied to select a subset of devices for participation in each FL round. The scheduling process confronts significant challenges due to the need for fair workload distribution, efficient resource utilization, scalability in environments with numerous edge devices, and statistically heterogeneous data across devices. This paper proposes a load metric for scheduling policies based on the Age of Information and addresses the above challenges by minimizing the load metric variance across the clients. Furthermore, a decentralized Markov scheduling policy is presented, that ensures a balanced workload distribution while eliminating the management overhead irrespective of the network size due to independent client decision-making. We establish the optimal parameters of the Markov chain model and validate our approach through simulations. The results demonstrate that reducing the load metric variance not only promotes fairness and improves operational efficiency, but also enhances the convergence rate of the learning models. 








§ INTRODUCTION



Federated learning (FL), introduced by McMahan et al. <cit.>, emerged as a solution to the limitations of traditional machine learning models that require centralized data collection and processing. FL enables client devices to collaboratively train a global model while keeping all the training data localized, thus addressing privacy concerns. 





The iterative training process of FL involves clients updating local models with their data before sending these updates to a central server for aggregation. 
To reduce communication load, a client selection policy can be applied such that only a subset of the clients transmits to the central server in a round. For example, random selection <cit.> selects each client uniformly at random in each round. However, client selection may lead to an unbalanced communication load and uneven contribution to the global model among clients. 

In this paper, we propose the load metric X, defined as the number of rounds between subsequent selections of a client. Assuming the distribution of the load metric for all clients is identical, our goal is to minimize Var[X], the variance of X.


The load metric X is related to the concept of Age of Information (AoI) <cit.>, which is the time elapsed since the last transmission of a client. In fact, X is called the peak age <cit.>. While most previous approaches in studying AoI in networking problems involve minimizing the average age or the average peak age <cit.>, this paper focuses on reducing the variance of X given the probability of selecting each client. 



The benefits of minimizing the variance of X are as follows:

    
  * Predictable Update Intervals and Load Balancing: By reducing the variance of X, we stabilize peak age and ensure consistent intervals between client selections. This predictability implies a fair and balanced selection process and improves operational efficiency.
    
  * Improved Convergence and Accuracy: A lower variability in X is essential for maintaining a diverse and representative global model, preventing biases and enhancing overall model accuracy.




We introduce a decentralized policy based on a Markov chain. Here, the state of the Markov chain is determined by the age of the client, and the transition probabilities represent an age-dependent selection policy.  

We establish the optimal transition probabilities of the Markov model that minimize Var[X].
We demonstrate that the optimal Markov model is equivalent to selecting the oldest age. The Markov model is suitable for decentralized client selection, alleviating the management overhead of the central server. Furthermore, the Markov model can be potentially modified to allow for real-time adjustment of client selection probabilities, to accommodate the dynamic nature of client data quality, computation resources, and network configurations and conditions. 

Simulation results demonstrate that the application of the Markov model enhances the convergence rate of the learning processes, achieving target accuracies in fewer communication rounds compared to random selection, with improvements exceeding 9% and more than 20% in certain datasets. These findings confirm the efficacy of our approach in addressing issues related to client overloading and under-utilization while maintaining efficiency throughout the learning process.




















Traditional machine-learning approaches require centralized data training in data centers, resulting in centralized machine-learning algorithms for data analysis and inference. However, it is often undesirable and impractical for edge devices to transmit data to parameter servers due to privacy constraints in wireless networks and limited wireless communication resources. Federated learning can solve these problems. Introduced in <cit.>, Federated Learning is a technique designed to train scaled machine learning models, using on-device data in a privately preserved manner. It can enable devices to train machine learning models without data sharing and transmission and is required to fulfill the needs of data privacy and security concerns.

Compared to traditional distributed machine learning, Federated learning needs to address the following challenges <cit.>. The first challenge is systems heterogeneity. The storage, computational, and communication capabilities of the devices part of a federated network may differ significantly. Additionally, only a tiny portion of the devices may be active at once, and fast clients must wait for slow clients, referred to as stragglers, before model aggregation.
The second challenge is statistical heterogeneity, reflected by the non-independent and non-identical distributions of data (Non-IID) and imbalance, resulting in instability and slow convergence.
The third challenge is privacy concerns. Sharing other information, such as model updates, as part of the training process can reveal sensitive information to a third party or the central server.

In the context of wireless federated learning (FL), where a vast array of clients operate within dynamically changing wireless environments, maintaining stable network connections poses a significant challenge. This scenario is further complicated by the inherent non-independent and identically distributed (non-IID) nature of the data across different clients, a characteristic reflective of diverse real-world applications.

The proposed FL framework introduces an innovative approach by integrating the concept of Age of Information (AoI) to manage and prioritize client participation based on the freshness of the data they hold. At the outset, each client initializes its AoI value, which serves as a measure of the time elapsed since its last successful interaction with the federated learning server. This measure is instrumental in determining the selection of clients for each learning cycle, with a preference for those whose data may have become relatively outdated, ensuring that the updates contribute effectively to the global model.

Following the selection process, clients synchronize their local models with the global model from the federated learning server, establishing a common starting point for model training. The training process then unfolds locally on each client's dataset, enabling the model to learn from a wide array of data patterns and nuances, reflective of the non-IID data distribution.

Once training is complete, the clients upload their updated models to the server, where they are aggregated to refine the global model. This aggregation step may involve grouping clients into clusters based on various criteria to optimize the process. Subsequently, the updated global model is disseminated back to the clients, ensuring that all participants benefit from the collective learning outcomes and are equipped for subsequent iterations.

This iterative process, characterized by a cycle of client selection, local training, model aggregation, and global model dissemination, continues until the learning objectives are met or a predefined stopping condition is satisfied. By accommodating the challenges of network instability and data heterogeneity, this framework aims to enhance the efficacy and applicability of federated learning in distributed environments, leveraging the collective intelligence of a wide array of clients to achieve robust and scalable learning outcomes.

Federated learning systems frequently operate in dynamic environments where the quality of client data, network stability, and client availability can fluctuate unpredictably. Traditional methods such as selecting clients based on the oldest age typically fall short in these settings due to their static nature and inability to adapt to ongoing changes. To address these challenges, we adopt a Markov chain approach for client selection. This method empowers us to dynamically adjust client selection probabilities in real-time, facilitating responses to changes in client performance and system conditions, thereby optimizing the learning process for improved outcomes.


Related Work.
The advancements in improving the Federated Averaging (FedAvg) algorithm by McMahan et al. <cit.> in terms of partial participation include exploring system diversity and non-IID data distribution.

Importance sampling has been employed, focusing on gradient computations on the most informative data, thereby enhancing learning outcomes <cit.>. Complementing this, various client selection strategies have been developed. A loss-based sampling policy <cit.> targets clients with high loss values to quicken convergence, though it requires careful hyper-parameter tuning. Alternatively, a sample size-based sampling policy <cit.> opts for clients with extensive local datasets, which, while effective in accelerating convergence, can falter in non-IID scenarios.

Addressing further challenges such as data staleness and the efficient use of network resources, the Age of Information metric has been adopted <cit.>. By prioritizing updates based on their timeliness, AoI-centric approaches enhance both accuracy and convergence rates, proving particularly effective in dynamic FL environments. In <cit.>, authors use the AoI metric to determine device selection and resource allocation in federated learning over wireless networks. Their approach manages sub-channel assignments and power allocation. To the best of our knowledge, this paper is the first to address the load balancing issue based on Age of Information.


Paper organization. The paper is organized as follows. In Section <ref>, FedAvg algorithm is reviewed, and the load metric is proposed. Client selection policies are presented in Section <ref>, including random selection and decentralized decision-making using a Markov chain. The optimal Markov model is investigated to minimize the load metric variance. Simulation results are shown in Section <ref> to demonstrate the efficiency of the proposed client selection policy in terms of convergence rate and accuracy. Finally, conclusions and future directions are discussed in Section <ref>.


Notations. Random variables are denoted by capital letters. E[·] and Var[·] denote the expectation and variance of a random variable, respectively. 




 


§ PROBLEM SETTING







This section briefly reviews FedAvg algorithm <cit.>. The load metric X is defined and the optimization problem on Var[X] is formulated afterward.

FedAvg algorithm. Consider a group of n clients coordinated by a central server to collaboratively train a common neural model using federated learning (FL). Each participating client i, i ∈{1,2,…,n}, has a dataset D_i = {X_i, Y_i}, where X_i is the input vector and Y_i is the corresponding output vector. These datasets are used to parameterize local models directly through weights W_i. The goal of FL is to optimize the collective average of each client’s local loss function, formulated as:

    min_{W_i}_i=1^n∑_i=1^n|D_i|/∑_j|D_j| L(X_i, Y_i, W_i),

where n represents the number of clients participating in the training, |·| denote the cardinality of a set, and L denotes the loss function. In this paper, we assume |D_i| are all equal.

FedAvg algorithm comprises two steps that are repeated until the model converges or a set number of rounds is completed:

(i) Local training: In Round t, clients use their local data D_i to update the local model parameters to W_i^(t). This involves calculating the gradient of the loss function L, typically through batches of data to perform multiple gradient descent steps.

(ii) Global aggregation: all clients periodically send their updated parameters W_i^(t) to the central server. The central server aggregates these parameters using a specified algorithm, commonly by averaging: W_global^(t) = 1/k∑_i=1^k W_i^(t).  The aggregated global model parameters W_global^(t) are then sent back to all participating clients.
Each client sets its local model parameters to W_global^(t). 




















Client selection and load balancing.
Since communication usually occurs through a limited spectrum, only a subset k out of n total clients can update their parameters during each global aggregation round. This necessitates a strategy for client selection. 




We aim to achieve load balancing by equalizing the number of iterations between consecutive client selections, assuming each client is selected with equal probability k/n. To that end, we define the load metric X as the random variable for the number of rounds between subsequent selections of a client. Assume X follows the same distribution for all clients. We propose the following optimization problem: 

    min_𝒮Var[X], 
       s.t.  P(S_i^(t)=1)=k/n, i ∈{1,…,n}, t ∈{1,…,T}.

Here, 𝒮 represents the set of permissible client selection policies, and T is the total number of communication rounds. In addition, S_i^(t) takes value 1 if client i is selected during the t-th round, and takes value 0 otherwise.

Despite that all clients share the same selection probability and average workload over all iterations, as required in constraint (<ref>), Var[X] focuses on the workload dynamics even when a small number of iterations is considered.

Minimizing Var[X] creates predictable update intervals and promotes equitable load distribution among all clients. 
Moreover, the staleness of local updates, represented by a large X, may slow down the convergence of the training process; and frequent updates from any single client, or a small X, may disproportionately affect the model. 
Thus, (<ref>) can improve learning convergence and accuracy by ensuring consistent data freshness, which is verified in our simulation.  


Relation to Age of Information.
A relevant metric is AoI, defined as the number of rounds elapsed since the last time a client was selected.

The age A_i for Client i in Round t evolves as follows:

    A_i^ ( t+1  ) = ( A_i^ ( t  ) +1  )  ( 1-S_i^ ( t  ) ),  S_i^ ( t  ) ∈{ 0,1 },

where A_i^(0)=0. Therefore, in each round, a client's age increases by one if it is not selected and it resets to zero if it is selected.
The load metric X is equal to the peak age, or the age before a client is selected.









§ CLIENT SELECTION



This section investigates random selection and a decentralized policy based on a Markov chain, assesses their load metric variances, and explores how to optimize the Markov model to contribute to a balanced  federated learning environment.










    











 §.§ Random Selection



The random selection method <cit.> involves choosing k out of n clients uniformly at random in each round for federated training. This approach, modeled mathematically by a geometric distribution, assumes each client has an equal probability p = k/n of being selected per round. The selection dynamics are described by:

    P(X=x) = p(1-p)^x-1,

where X is the number of rounds until a client is selected and x ∈{1,2,…}. The mean and variance of X are:

    E[X] = 1/p = n/k, 
       Var[X] = 1-p/p^2 = n(n-k)/k^2.

























 §.§ Decentralized Decision-Making Based on Local State

This policy utilizes a decentralized approach where each client independently decides whether to send its data based on its current Age of Information, ranging from 0 to m, with m representing the maximum permissible age. 
We formulate a Markov chain whose states represent the age of the client.
The probability of sending data at State j is p_j, for j = 0, 1, …, m, and in this case the state transitions to State 0. With probability 1-p_j, a client does not send data at State j. Accordingly, the state transitions to State j+1 for j=0,1,…,m-1; and the state remains the same for j=m.

States, transitions, and transition probabilities are depicted in Figure <ref>.

For a fair comparison with the random selection method, we impose a condition that the average number of clients selected in each round at steady state is k out of n. 


This approach allows each client to autonomously manage its participation based on its current AoI without requiring communication or coordination with the central server or other clients. 



In the following, we find the optimal transition probabilities to minimize the variance of X assuming that the client is at the steady state.
We first develop the expression of Var[X], then find the optimal Markov model with maximum permissible age m=1 in Theorem <ref>, and at last establish the optimal Markov model with a general m in Theorem <ref>.

Steady-state probabilities for the Markov chain are as follows:

    π_0= ∑_i=0^mπ_i p_i=k/n
       π_i = π_i-1· (1 - p_i-1)    i ∈{1,2,...,m-1}


    π_m= π_m-1· (1 - p_m-1) +π_m· (1 - p_m) 
    
       ∑_i=0^mπ_i = 1.

By solving these equations we obtain:

    π_0   = k/n= 1/1+∑_i=0^m-2∏_j=0^i(1-p_j) + 1/p_m∏_j=0^m-1(1-p_j),

for i ∈{1,..,m-1},

    π_i   = ∏_j=0^i-1(1-p_j)/1+∑_i=0^m-2∏_j=0^i(1-p_j) + 1/p_m∏_j=0^m-1(1-p_j),



    π_m   = 1/p_m∏_j=0^m-1(1-p_j)/1+∑_i=0^m-2∏_j=0^i(1-p_j) + 1/p_m∏_j=0^m-1(1-p_j).


Our objective is to compute and minimize the variance of 
X, which represents the number of rounds required for a specific client to be returned to State 0. Noting that Var[X] = E[X^2] - E[X]^2, we will derive both terms in this expression. Let us define E_0=E[X], and E_i to be the average time it takes to go to State 0 from State i. Then we have the following set of equations:

    E_i    = 1+ (1-p_i)E_i+1,    i ∈{0,1,...,m-1},
    
        E_m    = 1+ (1-p_m)E_m.

Equation (<ref>) implies that 
E_m= 1/p_m.
By (<ref>) and solving the above equations, we can calculate E_0 as:

    E_0= 1+ ∑_i=0^m-2∏_j=0^i (1-p_j) + 1/p_m∏_i=0^m-1(1-p_i) = n/k,

and for k ∈{1...,m-1} we have:

    E_k = 1+ ∑_i=k^m-2∏_j=k^i (1-p_j) + 1/p_m∏_i=k^m-1(1-p_i).

Similarly, we can write the equations for E[X^2] as follows:

    E[X^2]    = 1+ (1-p_0)(2E_1+ E[X_1^2]),
    
        E[X_i^2]    = 1+ (1-p_i)(2E_i+1+ E[X_i+1^2]),  i ∈{1,...,m-1},
    
        E[X_m^2]    = 1+ (1-p_m)(2E_m+ E[X_m^2])=2-p_m/p_m^2.
 
By solving equations (<ref>), (<ref>), and (<ref>) and performing algebraic simplifications, we derive the expression for E[X^2] as follows:

    1+ ∑_i=0^m-2∏_j=0^i(1-p_j)+2∑_i=0^m-1E_i+1∏_j=0^i(1-p_j)
       +E[X_m^2]∏_i=0^m-1(1-p_i).

Accordingly, we can write the expression of Var[X]. 


If m=1, the variance of X is given by Var[X] = (1 + p_0 - p_1)(1 - p_0)/p_1^2. Considering this, the optimal values of p_0, p_1 depend on the relationship between k and n/2. Specifically:

    
  * If k ≤n/2, then the optimal values are [p_0^*, p_1^*] = [0, k/n-k], resulting in a variance of (n-k)(n-2k)/k^2.
     
  * If k ≥n/2, then the optimal values are [p_0^*, p_1^*] = [2k-n/k, 1], resulting in a variance of (n-k)(2k-n)/k^2.



We simplify the constraint (<ref>) to obtain:

    1+1-p_0/p_1=n/k.

By writing down equations for Var[X], we have:

    Var[X]   = 1+(1-p_0)2+p_1/p_1^2 - (n/k)^2
       = 1+ n/k-(n/k)^2+ 21-p_0/p_1^2
       = 1+ n/k-(n/k)^2+ 2 n/k-1/p_1,

where in (<ref>) and (<ref>), we use (<ref>).
In order to minimize Var[X], we should set p_1 to its maximum value within its feasible range given the constraint (<ref>). Since 0 ≤ p_0 ≤ 1, constraint (<ref>) implies 1/p_1≥n/k-1, or p_1 ≤k/n-k. When k ≤n/2, k/n-k≤ 1 and it becomes the maximum feasible value for p_1^*. Consequently, we can find p_0^*=0 from 1+1-p_0/p_1=n/k. If k ≥n/2, then the maximum feasible value for p_1^* becomes 1 and then we can find p_0^*= 2k-n/k. 


    1+np_0-k/n-k/(1-p_0)(k^2/(n-k)^2) = (n-k)n-2k+np_0/k^2(1-p_0)

If p_i^* =1 that implies that i+1 ≥n/k. We need to show that this is the threshold where after that p_i has to be 1 given the condition π_0 =k/n assuming m+1≥n/k. Once m+1 < n/k then p_m becomes a value less than 1.


In both scenarios of Theorem <ref>, the variance of the Markov model is less than n(n-k)/k^2, which is the variance achieved by random selection.



If m=2, the variance of X is given by:

    Var[X] = n/k(3 - n/k) - 2 + 2 n/k - 1 - p_0/p_2.

Optimal values for p_i are found to be:

    
  * If 1 ≤n/k≤ 2, then [p_0^*, p_1^*, p_2^*] = [2 - n/k, 1, 1],
    
  * If 2 ≤n/k≤ 3, then [p_0^*, p_1^*, p_2^*] = [0, 3 - n/k, 1],
    
  * If 3≤n/k,     then [p_0^*, p_1^*, p_2^*] = [0, 0, 1/n/k - 2].



Using the equations derived from the expectations and variances associated with each state:
    
    E_0=1+1-p_0 +1/p_2(1-p_0)(1-p_1)=n/k, 
       E_1= 1+1/p_2(1-p_1),
       E_2=1/p_2,
       Var[X]=n/k(3-n/k)-2 +2n/k-(1-p_0)-1/p_2.

Differentiate the variance function with respect to p_0 and p_2:



 §.§ Differentiation with respect to p_0


    ∂var/∂ p_0 = 2/p_2




 §.§ Differentiation with respect to p_2


    ∂var/∂ p_2 = -2(x + p_0 - 2)/p_2^2



The derivative 2/p_2 = 0 (which cannot be true as this would make p_2 undefined) does not provide useful information directly about p_0.

Setting the derivative equal to zero:

    -2(x + p_0 - 2)/p_2^2 = 0


    x + p_0 - 2 = 0  p_0 = 2 - x


From this analysis, we find that p_0^* = 2 - x and p_2 should be any nonzero value for the derivative condition. This conclusion assumes x (i.e., n/k) is known and fixed, which simplifies the selection of p_0.


Given the optimal setup, p_0 = 2 - x, where x = n/k must be chosen such that 1 ≤n/k≤ 2. Substituting p_0 = 2 - x into the variance function, we derive the following:


    Var[X] = x(3 - x) - 2 + 2(x - (1 - (2 - x)) - 1/p_2)


Simplifying the expression inside the brackets, we have:

    x - (1 - 2 + x) - 1 = x + 1 - 2


Therefore, the variance becomes:

    Var[X] = x(3 - x) - 2 + 2(x - 1)/p_2



To find the value of p_2 that minimizes this expression, consider the term 2(x - 1)/p_2. Since x - 1 will be non-negative if x ≥ 1 (which is within our established constraint), the term 2(x - 1)/p_2 will decrease as p_2 increases, because the denominator is getting larger.

To minimize the variance, you want p_2 as large as possible within the allowed range. Therefore, set p_2 = 1 to minimize the variance while satisfying the probability constraint that 0 < p_2 ≤ 1.

For the range where 1 ≤n/k≤ 2, assuming p_0^*=2-x and p_2=1, we can find p_1 from (<ref>) as follows:

    1+1-(2-x)+(1-(2-x))(1-p_1^*)   =x,
    
        p_1^*   =x-1/x-1.

Therefore, for 1 ≤n/k≤ 2, we have [p_0^*,p_1^*,p_2^*]=[2-n/k,1,1].

When 2 < n/k≤ 3, regardless of the values of p_i, the variance expressions is positive. Hence, in this regime, we can only consider and minimize the term n/k-2+p_0/p_2. This term is minimized when p_0=0 and p_2=1. We can find p_1 using (<ref>) as:
1+1+1-p_1=n/k, therefore p_1=3-n/k. Consequently, for 2 < n/k≤ 3 we have 
[p_0^*,p_1^*,p_2^*]=[0,3-n/k,1].

When n/k>3, from the expression in (<ref>), it is obvious that we want p_2 to be as big as possible and p_0 to be as low as possible. Consider n/k=3+x. From the condition in (<ref>), given that p_0 and p_1 are between 0 and 1 we have:

    1+1+1/p_2≥ 3+x,
    
        p_2 ≤1/x+1.

Therefore, assuming p_2=1/x+1, and p_0=0 gives us p_1=0 from (<ref>). Hence, the optimal p_i in this regime becomes [p_0^*,p_1^*,p_2^*]=[0,0,1/n/k-2].



Consider an arbitrary m ≥ 1.

    
  * When m ≤⌊n/k⌋-1, the optimal values of p_i for minimizing the variance are [p_0^*, p_1^*, …, p_m-1^*, p_m^*] = [0, 0, …, 0, 1/n/k - m], and the minimum variance is (n/k-m)(n/k-(m+1)).
    
    
  * When m ≥⌊n/k⌋, setting i = ⌊n/k⌋,
    
    the optimal values are [p_0^*, p_1^*, …, p_i-2^*, p_i-1^*, p_i^*, …, p_m^*] = [0, 0, …, 0, i+1 -n/k, 1, …, 1], and the minimum variance is 
    c(1-c), where c = n/k-⌊n/k⌋. 
    



    Based on (<ref>) and (<ref>), Var[X] can be calculated as:
    
    1+ ∑_i=0^m-2∏_j=0^i(1-p_j)+2∑_i=0^m-1E_i+1∏_j=0^i(1-p_j)
       +2-p_m/p_m^2∏_i=0^m-1(1-p_i) -E_0^2.

Using the condition that E_0=n/k, we can simplify this expression as:

    Var[X]   = n/k-(n/k)^2  +2∑_i=0^m-2E_i+1∏_j=0^i(1-p_j)
       +1/p_m∏_i=0^m-1(1-p_i) +2-p_m/p_m^2∏_i=0^m-1(1-p_i).

Based on (<ref>), we have:

    E_i+1∏_j=0^i(1-p_j) 
    
      =   ∏_j=0^i(1-p_j) +∑_t=i+1^m-2∏_j=0^t (1-p_j) + 1/p_m∏_t=0^m-1(1-p_t) 
    
      =   n/k-1 +∏_j=0^i (1-p_j) - ∑_t=0^i∏_j=0^t (1-p_j).

Therefore we can further simplify the expression in (<ref>) as:

    Var[X]   = n/k- (n/k)^2 +2/p_m^2∏_i=0^m-1(1-p_i)
       + 2 ∑_i=0^m-2(n/k-1 - ∑_t=0^i-1∏_j=0^t (1-p_j)) 
       =(2(m-1)+n/k)(1-n/k)+2/p_m^2∏_i=0^m-1(1-p_i)
       -2∑_i=0^m-2∑_t=0^i-1∏_j=0^t (1-p_j)
       =(2(m-1)+n/k)(1-n/k)
       +2/p_m(n/k-1-∑_i=0^m-2∏_j=0^i (1-p_j))
       -2∑_i=0^m-2∑_t=0^i-1∏_j=0^t (1-p_j).

From (<ref>), it is obvious that for Var[X] to be minimized, assuming other parameters are fixed, p_i for i ∈{0,1,...,m-1} should be minimum in its feasible range where the condition in (<ref>) is satisfied. Similarly, by condition (<ref>), the second term in (<ref>) is positive, and p_m should be maximum when other parameters are fixed. 

When m ≤⌊n/k⌋-1, we show that the above minimum and maximum values can be achieved at the same time. From (<ref>), we find the following upper bound for p_m and it happens when all p_i values are 0:

    n/k≤ 1+ m-1 + 1/p_m
        
        p_m ≤1/n/k-m.

Consequently, by choosing p_i^*=0 for i ∈{0,1,...,m-1} and p_m^*=1/n/k-m, we achieve the minimum of Var[X] while satisfying the constraint in (<ref>).





Next consider m ≥⌊n/k⌋.
By defining y_t  ≜∏_j=0^t (1-p_i) we can rewrite the expression of Var[X] and make further simplifications.  
 
Based on Equation (<ref>), we need to minimize:

    2/p_m(n/k-1-∑_i=0^m-2 y_i)-2∑_i=0^m-2∑_t=0^i-1y_t.

The constraint (<ref>) becomes

    1+y_0+y_1+...+1/q_my_m-1=n/k.

The constraint 0 ≤ p_i ≤ 1, for all i ∈{0,1,…,m-1}, becomes

    0≤ y_m-1≤ y_m-2≤ ... ≤ y_0 ≤ 1.


Increasing p_m can only increase y_i's due to (<ref>), and hence reduce the objective function (<ref>). Thus, we choose p_m^* to its maximum feasible value, p_m^*=1.




Now, minimizing (<ref>) becomes maximizing

    2∑_i=0^m-2 y_i +2∑_i=0^m-2∑_t=0^i-1y_t
       
        =2∑_i=0^m-2 ((m-1-i)y_i),











subject to 1+y_0+y_1+...+y_m-1=n/k, and (<ref>). 

This is a linear programming optimization problem, where we start with y_0 and set it to the maximum value possible, and then continue with the rest of y_i's. Let i = ⌊n/k⌋, we should set y_0,...y_i-2 to 1, y_i-1=n/k-i, and y_i,...,y_m-1 to 0. This translates to p_0^*,...,p_i-2^*=0, p_i-1^*=1-(n/k-i), and p_i^*,...,p_m-1^*=1. 

Finally, the minimum variance can be computed by plugging the optimal p_0^*,…,p_m^* into the expression of Var[X].






The optimal Markov model in Theorem <ref> resembles the oldest-age method, which selects k clients with the highest age in each iteration.  
The general Markov model, although not fully utilized in our current implementation, holds significant potential due to its flexibility. One can adaptively choose the transition probabilities in the Markov model based on real-time performance, system conditions, and client dropout probabilities.
For example, one can set p_i > 0 for all i∈{0,1,…,m}, and a client can be selected even if it does not have the oldest age, increasing the chance of getting an update from the client before it drops out of the network.




For fixed n/k, as we increase m, the optimal variance of X gets smaller, and the optimal Var[X] reaches the minimum for all m ≥⌊n/k⌋. For any m ≤⌊n/k⌋ -1, the optimal Var[X] is given by (n/k-m)(n/k-(m+1)), which is smaller than the variance of the random selection policy, n/k(n/k-1). Hence, for all choices of m, the optimal Markov model provides better load balancing than random selection.










Ensuring timely updates from multiple clients is crucial for maintaining the freshness of information, as quantified by the Age of Information (AoI). However, equally important is the principle of fairness, which dictates that all clients should have equitable opportunities to transmit their updates and thus contribute to the system's overall information currency. We postulate that minimizing the variance in the number of rounds 
X to return to state 0 in our Markov chain model serves a dual purpose: it directly impacts AoI by fostering more predictable update intervals, and it promotes fairness by equalizing update opportunities across clients.

The AoI, by its definition, is sensitive to the distribution of intervals between successful update transmissions. By minimizing Var[X], the variability in these intervals is reduced, leading to a more stable and predictable AoI. This stability is crucial for applications requiring consistent information freshness, as it ensures that the information does not become stale or overly variable in age. Therefore, minimizing 
Var[X] indirectly contributes to maintaining an optimal AoI across the system.

By promoting more uniform intervals between successive updates, each client in a federated learning system has a fair chance to contribute their model updates, which is vital for incorporating diverse data perspectives and preventing bias. This fairness fosters a collaborative environment where all clients are equally valued, enhancing the overall learning process.

Moreover, minimizing Var[X] and thereby reducing peak AoI ensures that the model at the central server is regularly updated with fresh information from clients. This is particularly advantageous in federated learning, where the relevance and timeliness of updates directly impact the model's accuracy and convergence speed. Frequent and predictable updates prevent the model from becoming stale and ensure that it evolves in a way that accurately reflects the current data distribution across the network.

In summary, the benefits of minimizing 
Var[X] extends to federated learning by ensuring equitable participation across clients and maintaining the freshness of updates, thereby improving the collaborative training of the model. This alignment between fairness, timely information exchange, and learning effectiveness positions minimizing 
Var[X] as a key strategy in optimizing federated learning systems for better performance and more democratic participation.











§ SIMULATION RESULTS


In this section, we present the model's convergence and accuracy in scenarios where only a partial subset of clients (15%) participate in each communication round. We compare random selection with the proposed decentralized client selection policy based on the optimal Markov model. The experiments are conducted on the MNIST, CIFAR-10, and CIFAR-100 datasets. In our simulations, we utilize a convolutional neural network (CNN) as described in <cit.> to classify samples from the datasets. We employ the stochastic gradient descent (SGD) algorithm as the local optimizer, standardizing the batch size at 50 during the local training phase. The local training epochs are set to 5 per round, with an initial learning rate of 0.1 and a learning rate decay of 0.998. In our tests, we configure the parameters as follows: the number of devices n=100, the number of selected devices in a round k=15, and for the Markov model the largest age m=10. 

As illustrated in Figure <ref>, for the CIFAR-10 dataset, our proposed Markov model for client selection achieves faster convergence. Specifically, it reaches the target accuracy of 80% in just 240 communication rounds, compared to 265 rounds required by the random selection method, representing a 9.4% improvement in convergence speed.
Additionally, as depicted in Figure <ref>, for the CIFAR-100 dataset, we observe that the Markov model reaches the target accuracy of 40% in just 500 communication rounds, compared to over 600 rounds required by the random selection method. This results in an improvement of more than 20% in convergence speed.





In Figure <ref>, we present the performance analysis for the MNIST dataset under both IID and non-IID client data distributions. In the IID setting, our model achieves the target accuracy of 97% in just 37 communication rounds, which is faster than the 41 rounds required by the random selection method, marking a 9.7% increase in convergence speed. For the non-IID scenario, where client data follows a Dirichlet distribution with a parameter of 0.6 <cit.>, our model reaches the same target accuracy in 56 rounds, compared to 64 rounds with random selection, reflecting an improvement of 12.5% in convergence speed.











§ CONCLUSION

In conclusion, our research suggests that the decentralized client selection policy based on the Markov chain can lead to load balancing and efficient training in federated learning.

Future investigations should focus on a dynamic Markov chain policy for federated learning, adapting to the ongoing changes in network and client activity. By dynamically adjusting transition probabilities, we could enhance control over the Age of Information and improve responsiveness, thereby optimizing the trade-off between equitable client participation and overall system performance. 





 Paper Title: High Performance Im2win and Direct Convolutions using Three Tensor Layouts on SIMD Architectures.








High Performance Im2win and Direct Convolutions using Three Tensor Layouts on SIMD Architectures
    

    Xiang Fu1
    Xinpeng Zhang1
    Jixiang Ma1
    Peng Zhao2
    Shuai Lu1
    Xu T. Liu340000-0003-3980-9803This work is not related to the author's position at Amazon.
    
    
    1Nanchang Hangkong University, China
    2Microsoft, China
    3University of Washington, USA
    4Amazon Web Services, USA
    

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



§ ABSTRACT 
Convolution is the core component within deep neural networks and it is computationally intensive and time consuming. Tensor data layouts significantly impact convolution operations in terms of memory access and computational efficiency. Yet, there is still a lack of comprehensive performance characterization on data layouts on SIMD architectures concerning convolution methods. This paper proposes three novel data layouts for im2win convolution: NHWC, CHWN, and CHWN8, and introduces a set of general optimization techniques for both direct and im2win convolutions. 
We compare the optimized im2win convolution with the direct convolution and PyTorch's im2col-based convolution across the aforementioned layouts on SIMD machines. The experiments demonstrated that the im2win convolution with the new NHWC layout achieved up to 355% performance speedup over NCHW layout. 
Our optimizations also significantly improve the performance of both im2win and direct convolutions. 
Our optimized im2win and direct convolutions achieved up to 95% and 94% of machine's theoretical peak performance, respectively.



direct convolution, im2win convolution, NHWC layout, CHWN layout, CHWN8 layout 




§ INTRODUCTION


Convolution is the essential component of deep neural networks for computer vision tasks such as feature exaction from large-scale image data<cit.>. It not only comprises 50%-90% of computational operations including convolutional, pooling, ReLU, and fully-connected layers<cit.>, but also consumes more than 90% of the total execution time of many popular neural networks<cit.>. Hence, optimizing convolution operations is crucial for enhancing the performance of neural networks.

Convolution methods can be classified into three main categories based on how they transform the input tensor: direct, im2col-based, and im2win. Direct convolution performs convolution operations directly on the tensor without changing its format<cit.>. This approach avoids extra memory consumption compared to im2col-based and im2win convolutions, but it suffers from nonconsecutive memory access.
Im2col-based convolution transforms the convolution into general matrix-matrix multiplications (GEMM)<cit.>, leveraging optimized Basic Linear Algebra Subprograms (BLAS)<cit.> for excellent performance. It has regular memory access but a significant extra memory footprint, which greatly limits its applicability on memory-constrained devices. 
Previously, we propose a memory-efficient convolution called image-to-window (im2win), which reorganizes the input tensor into a row of dot product windows and flattens the unique elements of these windows into a row that correspond to the convolution operation's receptive fields<cit.>. It provides sequential memory access and data reuse, and thus greatly reduces memory overhead. 

A tensor memory/data layout (referred as layout henceforth) refers to how the data of a tensor is physically arranged in memory. There are commonly three layouts for tensors: NCHW, NHWC and CHWN, where N is the batch size, C is the number of channels, H is the image height, and W is the image width. It significantly impacts convolution operations in terms of memory access, computational efficiency and compatibility with deep learning frameworks<cit.>.




Tensors are stored as one-dimensional arrays in memory, although logically represented as four-dimensional arrays. An input tensor and an im2win tensor in a NCHW layout are illustrated in <Ref>. Different colors represent different channels. In the NCHW layout, the elements in the width dimension are contiguous in memory by prioritizing the width dimension first, followed by height, channel and batch. 
Assuming a stride of 1 and a filter tensor in the NCHW layout with dimensions of 1x3x2x2, the elements in the solid-lined boxes are used to compute the first output element in the output tensor, while the elements in the dashed-lined boxes are for the next. 
Further details will be provided in <Ref>.

In general, the memory efficiency and performance implications of various tensor layouts with different convolution algorithms on single-input, multiple-data (SIMD) architectures have received limited attention.
The previous im2win works optimize the NCHW layout on CPU and GPU but have not tried the NHWC and CHWN layouts<cit.>.
Li et al. reveal the performance impact of the NCHW layout with the im2col-based convolution and the CHWN layout with the direct convolution in different CNN layers, and propose a fast multi-dimension layout transformation algorithm on GPU<cit.>. 
Li et al. optimize the cache utilization of the NCHW and NHWC layouts in DNN training along with DNN pruning on CPU<cit.>.
To date, no performance characterization have been conducted on NCHW, NHWC and CHWN layouts on SIMD architectures with the three convolution methods. 
 
To address aforementioned gaps, we propose three new layouts for the im2win convolution: NHWC, CHWN, and CHWN8,
and propose a set of optimization techniques based on the roofline model for both direct and im2win convolutions.
We compare the optimized im2win convolution with the optimized direct and PyTorch's im2col-based convolutions with the above layouts on SIMD machines.
Our experiments demonstrate that the new NHWC layout on the im2win convolution achieves 11% to 355% performance speedup compared with the NCHW layout. Our proposed optimizations have been empirically validated to enhance the performance of both im2win and direct convolutions. 
Our optimized im2win convolution and our optimized direct convolution achieve up to 95% and 94% of the theoretical peak performance of the machine, respectively.
As the main contributions, this work

1) proposes three novel layouts for im2win convolution.

2) comes up with a set of general optimization techniques that not only can be applied to im2win convolution and but also to direct convolution on different layouts.

3) compares the optimized im2win convolution with the im2col-based and the optimized direct convolutions using four tensor layouts on SIMD machines.





§ PRELIMINARY AND RELATED WORKS




 §.§ Notation

The input of a convolution operation includes an input tensor (ℐ), a filter tensor (ℱ), and an output tensor (𝒪). In these tensors,
N_i is the batch size, s is the stride size, C_i and C_o are the number of input and output channels (also known as feature maps), H_i/f/o and W_i/f/o denote the height and width of a feature map.



 §.§ Tensor Layouts: NCHW, NHWC, CHWN



In the NCHW layout, the input (ℐ), the filter (ℱ), and the output tensors (𝒪) are expressed as ℐ[N_i][C_i][H_i][W_i], ℱ[C_o][C_i][H_f][W_f], and 𝒪[N_i][C_o][H_o][W_o], respectively. The convolution is defined as:

    [ 𝒪_(i, j, m, n)=∑_j=1^C_i∑_m=1^H_f∑_n=1^W_f(ℐ_(i, j, m × s+u, m × s+v). 
    .×ℱ_(j, r, u, v)), ]


In the NHWC layout, the input (ℐ), the filter (ℱ), and the output tensors (𝒪) are expressed as ℐ[N_i][H_i][W_i][C_i], ℱ[C_o][H_f][W_f][C_i], and 𝒪[N_i][H_o][W_o][C_o], respectively. The convolution is defined as:

    [ 𝒪_(i, m, n, j)=∑_j=1^C_i∑_m=1^H_f∑_n=1^W_f(ℐ_(i, m × s+u, m × s+v, j). 
    .×ℱ_(j, u, v, r)), ]


In the CHWN layout, the input (ℐ), the filter (ℱ) and the output tensors (𝒪) are expressed as ℐ[C_i][H_i][W_i][N_i], ℱ[C_i][H_f][W_f][C_o], and 𝒪[C_o][H_o][W_o][N_i] respectively. The convolution is defined as:

    [ 𝒪_(j, m, n, i)=∑_j=1^C_i∑_m=1^H_f∑_n=1^W_f(ℐ_(j, m × s+u, m × s+v, i). 
    .×ℱ_(r, u, v, j)), ]

The definitions (1), (2) and (3) above are all subject to

    j=1,2,..,C_o, m=1,2,..,H_o, n=1,2,..,W_o, 
       i=1,2,..,N_i, u=1,2,..,H_f, v=1,2,..,W_f,
       r=1,2,..,C_i.





 §.§ Convolution Algorithms and Related Works



Direct convolution performs on the original ℐ and ℱ without any tensor transformation.
It has seven nested for loops and an AXPY operation in the innermost loop.
Based on the tensor layouts that direct convolution works with, the AXPY operation needs to read at different indices of ℱ and ℐ, and writes at different indices of 𝒪.

Direct convolution can compute with the original input tensors, so they usually adopt the NCHW layout of raw images. A study shows that, for convolution instances with large C, NHWC layout outperforms NCHW layout<cit.>. Grouping a certain dimension of the input tensors by a fixed size can also enhance the performance of direct convolution, such as NC32HW32 layout<cit.>.
Several works <cit.> have shown that the performance of direct convolution can be greatly improved by designing specific layouts based on the loop ordering of the algorithm on SIMD architecture. 

The im2col-based convolution transforms a convolution operation into a GEMM operation. ℐ[N_i][C_i][H_i][W_i] is processed in N_i batches, each batch contains data ℐ'[C_i][H_i][W_i] (that is, a single image). 
The im2col algorithm flattens the elements of each dot product window of  ℐ' and copies them into a single row of a 2D matrix<cit.>.
In addition to the conventional im2col data transformation algorithm, 
the MEC algorithm compresses the matrix layout, while still enabling the utilization of high-performance BLAS algorithms to perform convolution operations<cit.>.



§ HIGH-PERFORMANCE IM2WIN AND DIRECT CONVOLUTION USING THREE TENSOR LAYOUTS


In this section, we first review three layouts in the context of the direct convolution. Then we present three new layouts for the im2win convolution, following by how to determine the loop ordering based on the layouts and a set of optimizations for the im2win and direct convolutions on SIMD architectures.








 §.§ Motivations for Different Tensor Layouts



An example of direct convolution on an original input tensor in the NHWC layout is illustrated in <Ref>. 
The NHWC layout prioritizes the storage of elements in the last logical dimension—C_i—followed by that of W_i, H_i, and N_i. 
In the NHWC layout, the elements with the same N_i, H_i, and W_i but different C_i are contiguous in memory, which have unit stride access.
Assuming a stride of 1,
the red and green elements (outlined by solid lines, representing the convolutional windows to compute a single output element) in the input tensor are multiplied with the corresponding elements in the filter tensor, and the results are summed up (i.e., an AXPY) to obtain O_00. 
Next, the green and yellow elements (outlined by dashed lines) are used to calculate O_01. This process continues, with the convolutional window moving by the stride length in H_i or W_i, until all elements of the output tensor are computed.


Recall in <Ref>, the NCHW layout prioritizes the storage of elements in the last logical dimension—W_i—followed by that of H_i, C_i, and N_i.
Note that all C_i but not all W_i are used during the AXPY operation to obtain one output element. NCHW has non-unit stride access during the tensor convolution. However, the non-unit stride access may not be harmful, depending on the effects of caching and the access patterns used (determined by the loop ordering).

As shown in <Ref>, the CHWN layout stores elements by prioritizing N_i, followed by W_i, H_i, and C_i in memory. 
Previous research on GPU recommends to use N as the lowest dimension for coalesced memory access and data reuse in registers<cit.>. It has been observed that the performance is sensitive to the value of N.
The elements within the solid-lined boxes are used to compute the first eight output elements, while the elements within the dashed-lined box are used for the next eight output elements. 
This facilitates the use of vector registers for vectorization.
Assuming s=1, the convolutional window moves by one element in both the H_i and W_i dimensions until all output elements are computed.


 













 §.§ Im2win Tensor Transformation on Three Tensor Layouts



In this subsection, we present the im2win tensor transformation on three tensor layouts. The im2win tensor transformation process for the NWHC layout is shown as <Ref>. For other layouts, slight modifications need to be made on it. The im2win transformation flattens the elements of a convolutional window, storing them contiguously in memory and prioritizing them in the C_i dimension, as outlined in <Ref> from <Ref> to <Ref>. 
An example of transforming the original input tensor into an im2win tensor in the NHWC layout is shown in <Ref>. It is important to note that the green elements can be reused between two adjacent windows, which eliminates the need for repetitive storage in memory like the im2col transformation<cit.>.
In the output im2win tensor of <Ref>, the elements involved in each convolution operation are contiguous and more compact in memory, 
improving spatial locality,  cache and SIMD efficiency.

An example of the im2win convolution using the NHWC layout is shown in <Ref>. A naive im2win convolution method is shown as <Ref>. It is similar to the direct convolution, which involves seven nested loops and an AXPY in the innermost loop, 
except prior to that, the im2win convolution performs a tensor transformation upon ℐ to get Î, and ℱ in the NHWC layout is transformed into a NWHC layout to match Î.
All output elements can be computed through the outer four loops of <Ref> from <Ref> to <Ref>.

An im2win tensor in CHWN/CHWN8 layout is shown in <Ref>. The characteristics of C_i, H_i, and W_i in this layout are similar to NCHW, but the CHWN layout prioritizes storage in N_i. 
The CHWN layout's efficiency is constrained by 256-bit vector registers, which process only 8 outputs at once. When N_i > 8, this leads to low cache utilization, as the cache holds unnecessary data for these 8 output calculations.
Motivated by the approach of dividing C_i into blocks in the direct convolution <cit.>, we propose a new CHWN8 layout.
CHWN8 lays 8 N_i in the innermost layer and remaining N_i in the outermost layer, which takes full advantage of the vector registers without sacrificing the cache utilization.
N_i can be set to a multiple of 8 (with padding if necessary), 
CHWN8 layout may provide better data continuity than the NHWC layout when C_i is relatively small.







 §.§ Loop Reordering


In this subsection, we reorder the loops for the im2win and direct convolutions based on different tensor layouts.
Ideally, we arrange the inner loops to access data closer in memory to enjoy the unit stride access as long as possible.
In both direct and im2win convolutions, for NCHW, CHWN, and CHWN8 layouts, as they have the CHW memory access pattern, we use the width of the convolution window as the innermost loop, followed by the height, and the channel (W_f, H_f, and C_i). 
Conversely, for the NHWC layout, the innermost three loops iterate over the channel, the width, and the height (C_i, W_f, and H_f) of the convolution window. 
For the im2win convolution, since the im2win transformation flattens the elements of a convolutional window, we usually coalesce the height and the width layers (W_f, H_f) into a W_f × H_f.

Next we determine the order of the outer four loops. It not only determines the order in which the elements of the output tensor are produced but also influences the order in which the element in the input tensor are accessed. 
The layout mainly affects the inner three levels of the loop order, for the outer four levels the loop order applies to all data layouts.
Recall in <Ref>, the green elements of the input tensor are shared between the convolution windows of adjacent output tensor elements. Therefore, we consider to position the width (W_o) of the traversed output tensor in the fourth layer of loops at <Ref> in <Ref>.
Because the memory access of the input tensor is expensive in a convolution operation, we set the channel (C_o) as the third layer to reduce its access.
Since the batch of the output tensor corresponds to the batch of the input tensor, we place the batch (N_i) in the first layer of the loop. 
Finally we have the order of the outer four loops as NHCW in <Ref>. 



 §.§ Optimizations for the Im2win and Direct Convolutions




In this subsection, we use the Roofline Model<cit.> to determine how to optimize the im2win and direct convolutions on SIMD systems. 
We propose a set of optimizations for both the im2win and direct convolutions and apply the optimizations to them (with a slight modification on the direct convolution). The optimizations are classified into two categories: reducing the memory bottleneck and increasing the arithmetic intensity of the convolution kernel. The former includes hoist, memory alignment, register and cache blocking. The latter includes
loop unrolling, vectorization and FMA instructions, loop coalescing, and parallelization strategies.
We have three level parallelization: Non-Uniform Memory Access (NUMA) level, thread level using OpenMP and instruction level using SIMD.

We present the optimized im2win convolution as <Ref>. 
Since each element of the output tensor in  each batch (N_i) can be computed independently, there is abundant parallelism available<cit.> and this is NUMA-friendly.
Within each batch, the shared elements of adjacent windows can be maximized across the threads.
When the dimension we choose to parallelize is small and the number of CPU cores is large, this leads to a workload imbalance in many core architecture. The solution is to coalesce multiple dimensions/loops into one parallel loop to achieve better load balance. In practice, we find that coalescing two dimensions of the output tensor yields the best load balance. 
Hence, we apply parallel strategies and coalesce N_i and H_o in a parallel loop at <Ref> in <Ref>.

In <Ref>, we hoist three things: the indices of the elements in the 1D array of the tensor, the elements of the input tensor, and the entire filter tensor at <Ref>.
We use register and cache blocking<cit.> at <Ref> to reduce cache misses because the neighboring convolutional windows share duplicate elements of the input tensor, which W_o,b is the blocking size.
Next, we unroll the loop at <Ref> to take advantage of the spatial locality,
because the innermost loop accesses consecutive elements.
Since we are using FP32 and AVX2, we set the loop unrolling size N_vec=8. Finally, we vectorize the input tensor and filter tensor in units of eight to use the FMA instruction in AVX2<cit.> at <Ref>. 

A cache-line is 512 bits on contemporary x86_64 architectures, which is the minimum data quantity that can be fetched from memory to cache. Without memory alignment of the tensor data structure, the CPU needs to issue two memory requests to access an element, because the element may be in the middle of the cache-line. This hurts the performance greatly because not only the CPU has to wait for memory access but also more caches are used. This leads to higher cache miss and lower cache utilization. We store the elements of the tensor in a memory-aligned way using posix_memalign in C during the memory allocation.













§ EXPERIMENTS



 §.§ Experimental Setup

Architectures. We use a server with two Intel® Xeon® Gold 6330 CPUs and 251 GB RAM. Each CPU has 28 physical cores, running at 2.0 GHz, with 48 KB L1d cache, 32 KB L1i cache, 1.28MB L2 cache and 43 MB L3 cache.

Benchmarks We aim to cover the majority of the convolutional layers in commonly used DNNs. 
Hence for our experimental evaluation, we select an state-of-the-art DNN benchmark<cit.> shown in <Ref>, which includes twelve unique convolution layers, conv1-conv12.



Software We compare the direct, im2win, and im2col-based convolutions with four tensor layouts: NHWC, NCHW, CHWN, CHWN8. We intend to compare with the state-of-the-art implementation of direct convolution and the layout proposed in<cit.>, but their implementation is not open-sourced. 
We use the im2col-based convolution in PyTorch 2.1<cit.> with MKL<cit.>. Note that PyTorch only supports the NHWC and NCHW layouts.
Our code is compiled with GCC 9.5.0 compiler and  compilation flags. We use OpenMP 4.0 for parallelization with guided scheduling. 





 §.§ Performance of Different Convolution Algorithms


We run each algorithm 50 times on each benchmark with N_i=128 and report the best runtime. 
<Ref> shows the performance results in TFLOPS and <Ref> shows the memory usage of our high-performance direct convolution and high-performance im2win-based convolution, and the im2col-based convolution using MKL in PyTorch. In <Ref>, the left y-axis shows the performance in TFLOPS, and the right y-axis shows the performance of the machine peak.
Overall, our optimized im2win convolution achieves eight best TFLOPS out of twelve benchmarks; our optimized direct convolution achieves three best TFLOPS out of twelve benchmarks and the im2win convolution achieves close performance on these three; the im2col-based convolution in PyTorch achieves one out of twelve best TFLOPS on conv12. 
All twelve best TFLOPS are all yielded from the NHWC layout across these three methods.
Our proposed optimization techniques are proven effective on both im2win and direct convolutions. Our im2win convolution achieves 95% and 91% of the theoretical peak performance of the architecture on conv5 and conv6 respectively. Our direct convolution achieves 91% and 94% of the theoretical peak performance on conv5 and conv6 respectively. 

With the NHWC layout and performance normalization, 
excluding conv6 and conv12, our im2win convolution achieves between 1.1× and 4.6× performance speedup, and our direct convolution achieves between 1.1× and 3.8× performance speedup against the im2col-based convolution. All the best performance of the im2win convolution on twelve benchmarks is achieved using the NHWC layout on CPU. Our im2win convolution using the NHWC layout outperforms the NCHW layout by at least 11% and up to 355% across all benchmarks, showcasing significant efficiency gains in various scenarios.
For the direct convolution, nine out of twelve best TFLOPS are achieved using the NHWC layout on CPU, the rest three are achieved using the CHWN8 layout.

With the NCHW layout,
our im2win-based convolution achieves between 1.4× and 2.4× performance speedup on all benchmarks, and PyTorch's im2col-based convolution achieves between 1.1× and 7.5× performance speedup (excluding conv7) against the direct convolution. The direct convolution performs poorly on the NCHW layout. 

With the CHWN/CHWN8 layout, im2win_CHWN8 outperforms im2win_CHWN on all benchmarks from 3.7× to 16×; direct_CHWN8 outperforms 2.3× and 8× over direct_CHWN except in conv7. 
This shows our proposed CHWN8 layout overwhelmingly beats the CHWN layout on these two methods.
The direct convolution with the CHWN8 layout 
 performs better than the CHNW, NCHW and NHWC layouts when C_i is small (C_i=3 for conv1, conv2, conv3). Similar results have also been observed in the previous GPU work<cit.>.

In <Ref>, using the same tensor layout, the memory usage of each convolution is the same hence we only annotate one number per method. In all benchmarks, the direct convolution uses the least memory, while the im2col-based convolution consumes the most memory.
On average, the im2col-based convolution has 3.9× more memory usage than direct convolution, and im2win-based convolution has 1.5× more memory usage than direct convolution. The im2win-based convolution uses on average 39% of the memory of the im2col-based method, and in conv5, it uses only 24% of the memory of the im2col-based convolution.




§ CONCLUSION


We proposed three new layouts for the im2win convolution: NHWC, CHWN and CHWN8, and a set of general optimization techniques for both direct and im2win convolutions on SIMD architectures. We applied these optimizations on the im2win and the direct convolutions, and compared with PyTorch's im2col-based convolution on the above layouts along with the NCHW layout. 
Our experiments demonstrated that im2win convolution using the new NHWC layout achieved 11% to 355% performance speedup compared to NCHW layout. 
The proposed optimizations were proven to boost the performance of the im2win convolution and direct convolution. 
Our optimized im2win and direct convolutions achieved up to 95% and 94% of the theoretical peak performance of the machine, respectively.



§ ACKNOWLEDGMENT

This work was partially supported by the NSF OAC-2402542, and the National Natural Science Foundation of China (Grant No. 42261070).

IEEEtran








We leave some information out from the main body of this paper and have them as part of an appendix. Because we consider they are not essential to the main argument but may be useful for readers who want to dive deeper into the topic. 




 §.§ Peak Performance

We use the following formula to calculate the peak-gflops of the server:



    peak_flop/s = (#processors) × (#cores_per_processor) 
    × (clock_speed[1/s]) × (2 ×#FMA_units) ×vector_size[bits]/64



Based on <Ref>, the peak GFLOPS of the server that we use in our experimental evaluation is 3584 GFLOPS.



 §.§ Batch Size Scaling on Different layouts


We perform a strong scaling on the batch size from 32, 64, 128, 256 to 512 with different layouts using the direct convolution and the im2win convolution.
The performance results of the direct convolution with the CHWN, CHWN8, NCHW, and NHWC layouts are shown in <Ref>, <Ref>, <Ref>, and <Ref>, respectively.
The performance results of the im2win convolution with the CHWN, CHWN8, NCHW, and NHWC layouts are shown in <Ref>, <Ref>, <Ref>, and <Ref>, respectively.

From the figures, we can tell that the CHWN layout is most sensitive to the batch size among four layouts. The performance of the direct and im2win convolutions is the best on twelve benchmarks when the batch size is 32 (except for the direct convolution on conv12). Recall in <Ref>, the efficiency of the CHWN layout is constrained by the number of vector registers on the SIMD machines. When N_i>8, this leads to low cache utilization.

With our proposed CHWN8 layouts, the performance of the direct convolution and the im2win convolution exhibits similar patterns on twelve benchmarks, that is, when the channel sizes of the benchmarks are small (C_i=3 for conv1, conv2, conv3), the smaller the batch size is, the better performance the convolutions have on these benchmarks; when the channel sizes of the benchmarks (conv4-conv12) are large, the larger the batch size is, the better performance the convolution have on these benchmarks.

With the NWHC and NCHW layouts, both convolution methods show no obvious evidence that they are sensitive to the batch size across all benchmarks. The performance variance is contributed by the overall dimension of the input/filter tensors, and the loop coalescing which we coalesce the N_i and H_o dimensions into one parallel loop to achieve better load balance. 
























  


 Paper Title: MoFO: Momentum-Filtered Optimizer for Mitigating Forgetting in LLM Fine-Tuning.



  
 
 

 
    
       
            
       
       
       
      
         












  














 



 

 

















            
       
    
       
       
      






 








arrows
































































ℝ
ℤ
𝒜
ℬ
𝒟
ℰ
ℱ
𝒢
ℋ
𝒮
ℐ
ℒ
ℳ
𝒩
𝒫
𝒮
𝒯
𝒱
𝒲
𝒵
𝕊
ℕ
𝐏
TV
MSE

𝐰
𝐚
𝐙










































thmTheorem
lemLemma
corCorollary
propProposition
asmpAssumption
defnDefinition
oracleOracle
factFact
conjConjecture
remRemark
conditionCondition
exerciseExercise
messMessage



thmTheoremTheorems
lemLemmaLemmas
corCorollaryCorollaries
propPropositionPropositions
asmpAssumptionAssumptions
defnDefinitionDefinitions
oracleOracleOracles
factFactFacts
conjConjectureConjectures
remRemarkRemarks
exampleExampleExamples
conditionConditionConditions
exerciseExerciseExercises
algorithmAlgorithmAlgorithms
tableTableTables
figureFigureFigures
sectionSectionSections
subsectionSectionSections
appendixAppendixAppendices
messageMessageMessages






















@mathaccentsingle[3]
  0"0362#1^H
  2"03620pt#1^H
  0=2 #3#2
kern[1]#1@kerna

bar[2]@single#1@bar@#1#21@bar@#1#22
bar@[3]
  ##1##2
    @mathaccent
    #32 @nucleus@char @@style@nucleus_
    @@style@nucleus_
    @@
    @-@
    @ 3
    tempdima@
    tempdima-tempdima 10
    @-tempdima
    @>@ @0pt@kern0.6-@
    #31
      @kern-0.6@@nucleus@kern0.4@
      @0.4@kerna
      @kern#2
      @<@ @kern1@kern1 -@@kern-0.6@#1
    
  @depthne
  @bgroupempty @egroup@set@skewchar
  @ @everymath@group
  @set@skewchar@nested@a
  #31
    @nested@a111#1
  @till@marker##1
    @char@till@marker#1@char A@char
    @nested@a111@char
  















#1figure <ref>
#1Figure <ref>
#1#2figures <ref> and <ref>
#1#2#3#4figures <ref>, <ref>, <ref> and <ref>
#1section <ref>
#1Section <ref>
#1#2sections <ref> and <ref>
#1#2#3sections <ref>, <ref> and <ref>
(<ref>)1(<ref>)
#1Equation <ref>
#1<ref>
#1chapter <ref>
#1Chapter <ref>
#1#2chapters<ref>–<ref>
#1algorithm <ref>
#1Algorithm <ref>
#1#2algorithms <ref> and <ref>
#1#2Algorithms <ref> and <ref>
#1part <ref>
#1Part <ref>
#1#2parts <ref> and <ref>

#1⌈#1 ⌉
#1⌊#1 ⌋
1




ϵ



ϵ
θ
𝐚
𝐛
𝐜
𝐝
𝐞
𝐟
𝐠
𝐡
𝐢
𝐣
𝐤
𝐥
𝐦
𝐧
𝐨
𝐩
𝐪
𝐫
𝐬
𝐭
𝐮
𝐯
𝐰
𝐱
𝐲
𝐳

a
b
c
d
e
f
g
h
i
j
k
l
m
n
o
p
q
r
s
t
u
v
w
x
y
z

𝐀
𝐁
𝐂
𝐃
𝐄
𝐅
𝐆
𝐇
𝐈
𝐉
𝐊
𝐋
𝐌
𝐍
𝐎
𝐏
𝐐
𝐑
𝐒
𝐓
𝐔
𝐕
𝐖
𝐗
𝐘
𝐙

A
B
C
D
E
F
G
H
I
J
K
L
M
N
O
P
Q
R
S
T
U
V
W
X
Y
Z

0
1
μ
θ
a
b
c
d
e
f
g
h
i
j
k
l
m
n
o
p
q
r
s
t
u
v
w
x
y
z

α
β
ϵ
λ
ω
μ
ψ
σ
θ
a
b
c
d
e
f
g
h
i
j
k
l
m
n
o
p
q
r
s
t
u
v
w
x
y
z

A
B
C
D
E
F
G
H
I
J
K
L
M
N
O
P
Q
R
S
T
U
V
W
X
Y
Z
β
Φ
Λ
Σ

msl
boldbxn

A
B
C
D
E
F
G
H
I
J
K
L
M
N
O
P
Q
R
S
T
U
V
W
X
Y
Z


𝒜
ℬ
𝒞
𝒟
ℰ
ℱ
𝒢
ℋ
ℐ
𝒥
𝒦
ℒ
ℳ
𝒩
𝒪
𝒫
𝒬
ℛ
𝒮
𝒯
𝒰
𝒱
𝒲
𝒳
𝒴
𝒵

𝔸
𝔹
ℂ
𝔻
𝔽
𝔾
ℍ
𝕀
𝕁
𝕂
𝕃
𝕄
ℕ
𝕆
ℙ
ℚ
ℝ
𝕊
𝕋
𝕌
𝕍
𝕎
𝕏
𝕐
ℤ

Λ
A
B
C
D
E
F
G
H
I
J
K
L
M
N
O
P
Q
R
S
T
U
V
W
X
Y
Z
Σ


Λ
A
B
C
D
E
F
G
H
I
J
K
L
M
N
O
P
Q
R
S
T
U
V
W
X
Y
Z











 




















 



 Paper Title: GOProteinGNN: Leveraging Protein Knowledge Graphs for Protein Representation Learning.



printacmref=false 
 
plain 

















  Technion
 

kalifadan@cs.technion.ac.il




  Meta AI
 

urielsinger@meta.com




  Technion
 

kirar@cs.technion.ac.il























































































§ ABSTRACT 
Proteins play a vital role in biological processes and are indispensable for living organisms. Accurate representation of proteins is crucial, especially in drug development.
Recently, there has been a notable increase in interest in utilizing machine learning and deep learning techniques for unsupervised learning of protein representations. 
However, these approaches often focus solely on the amino acid sequence of proteins and lack factual knowledge about proteins and their interactions, thus limiting their performance.
In this study, we present GOProteinGNN, a novel architecture that enhances protein language models by integrating protein knowledge graph information during the creation of amino acid level representations. Our approach allows for the integration of information at both the individual amino acid level and the entire protein level, enabling a comprehensive and effective learning process through graph-based learning.
By doing so, we can capture complex relationships and dependencies between proteins and their functional annotations, resulting in more robust and contextually enriched protein representations.
Unlike previous fusion methods, GOProteinGNN uniquely learns the entire protein knowledge graph during training, which allows it to capture broader relational nuances and dependencies beyond mere triplets as done in previous work.
We perform a comprehensive evaluation on several downstream tasks demonstrating that GOProteinGNN consistently outperforms previous methods, showcasing its effectiveness and establishing it as a state-of-the-art solution for protein representation learning.


































<ccs2012>
   <concept>
       <concept_id>10010405.10010444.10010450</concept_id>
       <concept_desc>Applied computing Bioinformatics</concept_desc>
       <concept_significance>500</concept_significance>
       </concept>
   <concept>
       <concept_id>10010147.10010178</concept_id>
       <concept_desc>Computing methodologies Artificial intelligence</concept_desc>
       <concept_significance>500</concept_significance>
       </concept>
 </ccs2012>


[500]Applied computing Bioinformatics
[500]Computing methodologies Artificial intelligence






























GOProteinGNN: Leveraging Protein Knowledge Graphs for Protein Representation Learning
    Kira Radinsky
    August 5, 2024
=====================================================================================





§ INTRODUCTION


Proteins play a vital role in biological processes and are central to understanding the intricate mechanisms of living organisms. In the pursuit of unraveling their complexity, the challenge of accurately representing proteins has gained considerable attention. Accurate representation holds paramount significance, especially for drug development <cit.>.

Amino acids serve as the fundamental building blocks of proteins. In recent years, there has been a surge of interest in employing machine learning and deep learning techniques to create unsupervised amino acid representations that can be used to represent a protein <cit.>. Prevailing approaches have predominantly focused on unraveling amino acid relationships using text-based methodologies, such as ProtBert <cit.>. The amino acid representations are then averaged to create the protein representation.

Recently, numerous approaches have attempted to directly create a protein representation by encapsulating external information about the protein, rather than relying solely on the average of amino acid representations. 
Those methods, such as KeAP <cit.>, harness vast protein knowledge repositories. These knowledge graphs contain diverse information about interactions the protein participated in, the processes in which it is involved, etc. 
The factual knowledge is conveyed through text descriptions.
Figure <ref> provides an illustrative example of such a protein knowledge graph, showcasing the intricate connections between the protein, relations, and Gene Ontology (GO) terms. 
However, current state-of-the-art (SOTA) methods still grapple with the incorporation of this extensive information. Existing paradigms, at times, reduce the intricate protein knowledge graph to mere triplets (Protein, Relation, GO Term), neglecting the intricate tapestry of the knowledge graph <cit.>. These methods have overlooked the importance of capturing a holistic graph representation. For instance, although they may encode annotations, crucial information such as broader relational nuances and contextual dependencies within the graph may be lost unless the entire graph structure is considered. Moreover, these studies have predominantly concentrated on either amino acid level learning <cit.> or protein level learning <cit.>, with none addressing both simultaneously.

In this work, we present GOProteinGNN, a novel architecture that integrates amino acid level and protein level representations within a single framework. Our method addresses existing limitations by comprehensively utilizing the entire protein knowledge graph structure during pretraining, ensuring the capture of intricate relationships. Specifically, GOProteinGNN leverages a novel Graph Neural Networks (GNN) Knowledge Injection (GKI) procedure that leverages the [CLS] token representation of the protein language model (PLM) during the encoder stage. GOProteinGNN interprets amino acid sequences as coherent “sentences” that describe proteins, learning embeddings for individual amino acids in its initial encoder layers. The GKI layer then assigns embeddings to protein nodes in the protein knowledge graph and incorporates knowledge into these embeddings by leveraging neighboring embeddings through a Relational Graph Convolution Network (RGCN) <cit.>. 
This process alternates between amino acid level embedding and graph representation, enhancing the overall representation by allowing new insights to spread across the amino acid sequence and ensuring a thorough encoding that incorporates both sequential and graph-enhanced information.

Comparison to Prior Art:
Traditionally, most methods represented proteins solely based on their amino acid level interactions. Such methods <cit.> typically generate embeddings focused on the chemical characteristics of the proteins. For instance, these embeddings are capable of capturing the chemical interactions between protein pairs. In this work, by incorporating a knowledge graph into our approach, we are able to extend the encoding to include biological information, such as the specific biological processes in which the protein participates.
In terms comparable to natural language processing, our methodology introduces a novel technique for integrating external sentence knowledge into word-embedding learning. During the training phase, we directly infuse a knowledge graph representation into the amino acid embeddings, leveraging the [CLS] token representation for this novel fusion. This enhancement significantly improves the representations, even for novel proteins that are not present in the knowledge graph.
Other methods have concentrated on injecting knowledge graphs but have simplified the problem to only involve pairs of nodes <cit.>, often neglecting the broader structure of the knowledge graph. This simplification could adversely affect the quality of protein embeddings.
For example, as illustrated in Figure <ref>, the KeAP method <cit.>, which focuses solely on node pairs within the knowledge graph, might not adequately capture the complex interactions involved in a protein's role in diverse biological activities such as intracellular cAMP-activated cation channel activity, visual perception, and signal transduction pathways. This limitation stems from KeAP's failure to account for a holistic view of the protein's interactions within a biological context, potentially leading to a fragmented understanding of its functional significance, unlike our approach.

We perform an empirical evaluation over numerous protein tasks, spanning a broad spectrum of biological domains, comparing to SOTA methods and achieving SOTA performance across various benchmarks. We present ablation tests to better study the performance of the algorithm.

The contributions of this work are threefold:
(1) We present a novel algorithm for protein representation learning that allows the integration of knowledge graph information into PLMs. This enables a comprehensive and effective learning process through graph-based learning. 
(2) We introduce a novel GNN Knowledge Injection (GKI) layer that enables protein knowledge information to propagate through the amino acid sequence representations. In addition, we leveraged the [CLS] token representation, to integrate the GKI with the PLM in a novel manner.

(3) We present an extensive empirical evaluation of our work over numerous protein tasks establishing SOTA results. Also, we contribute our code to the community [GitHub repository: <https://github.com/kalifadan/GOProteinGNN>].



<https://www.zjukg.org/project/ProteinKG65>



§ RELATED WORK





 §.§ Proteins Representation Learning

Proteins are essential for living organisms and have vital roles in biological processes. Accurate representation of proteins is crucial, particularly in fields such as drug development. 
In recent years, there has been a surge of interest in applying machine learning and deep learning techniques to learn protein representations in an unsupervised manner. 
The popularity gain of Transformers <cit.> 
introduced many novel PLMs such as ESM <cit.>, MSA Transformer <cit.>, ProteinBERT <cit.>, ProteinLM <cit.> and ProtTrans <cit.>. 
These models were trained on data from UniRef <cit.> which contains hundreds of billions of amino acids.

A protein’s biological mechanism is determined by its three-dimensional (3D) native structure <cit.>. AlphaFold <cit.> is a recent model predicting the 3D structure of a protein.
A common SOTA method that resembles the AlphaFold architecture but tuned for protein representation is ESM <cit.>.
Unlike AlphaFold, which primarily focuses on 3D structure prediction and relies on multiple sequence alignments and template structures, ESM requires only a single input sequence (i.e., protein) and has demonstrated similar performance across various tasks.
Our comparative analysis focuses on ESM, emphasizing its utilization of single-sequence input and acknowledging its performance similarity to AlphaFold.



 §.§ Knowledge-Enhanced Language Models

Prior research has demonstrated the effectiveness of leveraging knowledge in language models to enhance their performance <cit.>. Specifically, KEPLER <cit.> endeavored to integrate factual knowledge through knowledge embedding optimization, while KG-BERT <cit.> and K-BERT <cit.> treated knowledge graph triples as textual sequences or transformed the original input sentence into a knowledge-rich sentence tree, respectively. In contrast, our algorithm differs significantly, as it considers graph topology and diverse node types present in protein knowledge graphs.

While PLMs traditionally learn protein representations from sequences, <cit.> proposed learning from biomedical text descriptions. 
Similarly, OntoProtein <cit.> integrated external knowledge graphs into protein pre-training, focusing on Gene Ontology (GO) term and protein entity alignment. KeAP <cit.> explored more granular knowledge graph interaction for improved protein decoding. OntoProtein's contrastive learning and masked modeling were enriched by KeAP's focus on the masked language modeling objective, achieving superior performance in various tasks.
Otter-Knowledge <cit.> extended knowledge graph integration beyond proteins, incorporating diverse entities into a comprehensive knowledge graph. 

The aforementioned models predominantly incorporate knowledge post-encoder, either via loss objectives or pre-training enrichment. This limitation hinders their ability to effectively incorporate external knowledge into protein embeddings during the encoder stage. 
Both OntoProtein <cit.> and KeAP <cit.>, together with other existing models, have predominantly focused on either amino acid level learning or protein level learning. Also, they reduce intricate protein knowledge graph to mere triplets, neglecting the intricate tapestry of the knowledge graph. GOProteinGNN sets itself apart by uniquely encompassing both levels of representation within a single framework thus learning the entire protein knowledge graph and allowing it to capture broader relational nuances and dependencies beyond mere triplets. 

We specifically target protein-related tasks, capturing critical biological contexts and interactions unique to proteins.



§ METHODS








 §.§ GOProteinGNN

We present the GOProteinGNN architecture (Fig. <ref>).
GOProteinGNN extends the capabilities of PLMs by incorporating knowledge graph information into the models, using our proposed GKI layer (See Section <ref>).
We leveraged the [CLS] token representation, to integrate between the proposed GKI and the PLM in a novel manner.
This fusion of language model-based representation learning with knowledge graph processing allows GOProteinGNN to generate enriched protein embeddings that capture essential biological contexts and interactions. The detailed algorithm is presented in Algorithm <ref>.







  §.§.§ Training Phase

Following the approach of KeAP <cit.>, we use a large pre-trained PLM (ProtBert <cit.>) as our initialization for the encoder layers' weights, leveraging its contextual understanding of proteins. The usage in ProtBert is a common practice <cit.> and allows us to show the significance of the injection method compared to the baselines in a fair manner. Alternatively, other advanced encoders can be used in the training algorithm. 

During the training stage, each amino acid sequence (i.e., protein) is passed through multiple encoder layers, yielding an intermediate representation of the amino acid sequence, denoted by z_p ∈ℝ^L_p × D, where L_p is the sequence length and D is the feature dimension (line 5 in the algorithm).

In language models, including a specialized classification token ([CLS]) at the beginning of the sequence is a common practice. The final hidden state of this token serves as a comprehensive sequence representation for various classification tasks <cit.>.
In our framework, we leverage the [CLS] token representation to integrate external knowledge into our representation framework and enhance the protein representation.

Within the context of a given encoder layer denoted by l, we introduce the GKI layer. The GKI operation (See Section <ref>) takes the input representation [CLS]_l along with the protein knowledge graph denoted by G (See Section <ref>), resulting in a refined representation [CLS]'_l (line 9): 
    [CLS]'_l = GKI([CLS]_l, G) ∈ℝ^D

Where the GKI utilizes the entirety of the knowledge graph during each injection process.

Subsequently, this refined [CLS] token representation assumes the role of the new [CLS] token in the ensuing encoder layer. 
Then, the model propagates and refines the representations of all other amino acids within the sequence, resulting in a new representation of the amino acid sequence (line 7).

This mechanism enables GOProteinGNN to achieve a holistic learning process, harmonizing amino acid level representations with protein level representations. The outcome is an integration of contextually enriched representations that intricately capture both sequential insights and the graph-enhanced knowledge pertaining to proteins. 

As the protein representation undergoes transformations across the encoder layers, our aim is to uniformly integrate knowledge a total of k times across the L encoder layers. Therefore, we add a GKI layer after every L ÷ k encoder layer (line 8).




  §.§.§ Inference Phase

To eliminate the dependency on a knowledge graph (KG) during inference, we introduce a weight-less injection mechanism. 
During the inference phase, GOProteinGNN omits the GKI components utilized solely for pre-training. 
In addition, we incorporate task-specific classification heads to enable predictions for these specific tasks.
This separation between pre-training and inference empowers GOProteinGNN to generalize effectively across various protein tasks, including those involving unseen proteins. 





 §.§ GNN Knowledge Injection (GKI)


GNNs <cit.> are powerful neural networks designed for graph-structured data <cit.>, making them suitable for learning effectively the protein knowledge graph, similar to common usages in other domains <cit.>.
By utilizing the protein knowledge graph denoted by G, we introduce a methodology called GNN Knowledge Injection, denoted by GKI. This method produces an updated protein representation that incorporates information from the knowledge graph into the input representation.




  §.§.§ Protein Knowledge Graph

GOProteinGNN is trained on a knowledge graph consisting of two different types of nodes: Proteins and GO-Terms, with Relations that act as edges. 
Proteins are represented as a sequence of amino acids, while the relation and the GO term are terms of factual knowledge described in natural language. The GO terms cover three aspects of biology: Molecular Function, Cellular Component and Biological Process.




  
Protein Knowledge Graph Definition

Let V_proteins be a set containing proteins (amino acid sequences).
The GKI leverages the accompanying descriptions of GO terms extracted from the protein knowledge graph. These descriptions are processed through a language encoder [We use PubMedBERT <cit.>] to generate comprehensive knowledge representations. 
To bridge the gap between text and protein, the GKI utilizes affine transformation to project those representations to the same space.
These representations in the protein knowledge graph are denoted by  V_terms.
We define the knowledge graph G = (V, E) where:
V = V_proteins∪ V_terms and E = { (u, v) | u ∈ V_proteins, v ∈ V_terms, (u, v) ∈ Relations }.





  §.§.§ GNN Method
 
A Protein Knowledge Graph consists of multiple types of relations between proteins and GO terms.
In response to the complex diversity of edge types in the graph, we select a GNN that accommodates the multifaceted edge information. Recognizing that conventional GCNs <cit.> lack the capacity to fully exploit this characteristic, we turn to the Relational Graph Convolution Network (RGCN) <cit.>, a specialized variant of GNN.
By doing so, we can effectively harness the rich information present in the knowledge graph to refine and enhance the protein embedding. This approach allows us to leverage the graph structure of the protein knowledge graph and to capture complex relationships and dependencies between proteins and their variable number of functional annotations, leading to more powerful and contextually enriched protein representations, a capability not attainable through other methods, such as embedding concatenation.
Moreover, other graph algorithms could also be applied. We compare different graph algorithms in Section <ref>.





  §.§.§ Knowledge Injection

Given a protein representation (i.e., the [CLS] representation) z_protein∈ℝ^D corresponding to node i in the knowledge graph G, set:

    h_i^(0) = z_protein


The given protein is connected to multiple nodes in the knowledge graph G, such as the GO terms nodes. To effectively utilize the textual information of these nodes, the GO term descriptions are processed through a language encoder (See Section <ref>) to generate comprehensive knowledge representations. To align these text-based representations with the protein embeddings, we project both types of representations into a unified space by applying an affine transformation.

The (l+1) RGCN layer performs the following update to the node representation z_protein = h_i^(l):

    h_i^(l+1) = σ( ∑_r ∈ R∑_j ∈ N_r(i)1/c_i,r W_r^(l) h_j^(l) + W_0^(l) h_i^(l))


where h_i^(l+1)∈ℝ^D is the new node representation, σ(·) is the activation function, N_r(i) is the set of neighbor nodes of node i connected by relation type r, R is the set of all relation types in the graph, h_j^(l)∈ℝ^D is the node representation of the neighbor node j at layer l, c_i,r is a normalization term for relation r specific to node i, and W_r^(l)∈ℝ^D × D and W_0^(l)∈ℝ^D × D are learnable weight matrices associated with relation r and the self-loop, respectively.

The new protein embedding after the L RGCN layers is given by:

    z^'_protein = h_i^(L)

where z^'_protein∈ℝ^D represents the updated protein embedding that captures essential biological contexts and interactions. 





 §.§ Masked Language Modeling Objective

For each input protein, 20% of the amino acids are randomly masked. Each masked amino acid a_(j) has an 80% chance of being masked for prediction, 10% chance of replacement with a random amino acid, and 10% chance of remaining unchanged. Suppose the number of masked
amino acids is M, the training objective ℒ_MLM is to minimize:


    ℒ_MLM = -∑_j=0^M - 1log P(a_(j) | z_p, Θ)


where z_p is the protein representation and Θ are the parameters of the GOProteinGNN model.







§ EMPIRICAL EVALUATION


In this section we outline our empirical setting. We performed analysis on both intrinsic (e.g., semantic similarity inference of proteins) and extrinsic tasks (e.g., protein-protein interactions), demonstrating the value of the novel representation we provide in real world bioinformatics tasks.




 §.§ Pre-Training Dataset

We employed the ProteinKG25 dataset <cit.> for pre-training, consistent with all knowledge-enhanced baseline methods.
ProteinKG25 provides a knowledge graph with aligned descriptions and protein sequences, respectively, to GO terms <cit.> and protein entities.
GO terms can be a molecule, a cellular component, or a biological process.
The GO terms and their relations to proteins are described using natural language. 
The dataset contains around 5 million triplets, including nearly 600k proteins, 50k attribute terms, and 31 relation terms, with an average of 8.64 relations per protein.
We address a data leakage concern by following the methodology proposed by KeAP <cit.>. After removing the leakage from the dataset, their algorithm is considered SOTA across all tasks. Given the robust performance demonstrated by KeAP in various downstream applications, outperforming other methods including OntoProtein <cit.>, we focus our comparative analysis solely on KeAP <cit.>.





 §.§ Tasks


We use Tasks Assessing Protein Embeddings (TAPE) <cit.> as the benchmark to evaluate protein representation learning. TAPE is a set of biologically relevant semi-supervised learning tasks spread across different domains of protein biology. 
All tasks were divided into specific training, validation, and test splits to ensure that each task tested a biologically relevant generalization that transferred to real-life scenarios.
Furthermore, we tested our representations on Semantic Similarity Inference and Binding Affinity Estimation tasks from PROBE <cit.> and on Protein-Protein Interactions (PPI) from GNN-PPI <cit.>.




  §.§.§ Contact Prediction

Contact prediction is a pairwise amino acid task, where each pair ⟨ x_i,x_j⟩ of input amino
acids from a sequence (protein) x is assigned to a label y_ij∈{0, 1}, where the label denotes whether the amino acids are in contact (< 8Å) <cit.>.
The data comes from ProteinNet <cit.>, and the evaluation is based on the precision metric.




  §.§.§ Remote Homology Detection

Remote Homology Detection is a sequence-level classification task. Each input sequence
(protein) x is assigned a label y ∈{1, ..., 1195} representing different possible protein folds. The data is from DeepSF <cit.> and the evaluation is based on average accuracy on a fold-level holdout set. 




  §.§.§ Protein-Protein Interaction Identification

Protein-protein interactions (PPI) refer to the physical contacts between two or more amino acid sequences, or two proteins, in our study. The goal is to predict the interaction type(s) of each protein pair, out of 7 possible interactions (i.e., multilabel classification task).
The evaluation is conducted on two datasets: SHS27K <cit.> and SHS148K <cit.>. Breadth-first search (BFS) and depth-first search (DFS) are used to generate test sets from the two datasets. F1 score is used as the evaluation metric.




  §.§.§ Semantic Similarity Inference

This analysis aims to evaluate the extent to which representation models capture information regarding biomolecular functional similarity, with a specific emphasis on the biological process (BP) category within the Gene Ontology (GO) dataset. We adopt the semantic similarity measures outlined by  and compute Spearman's rank correlation between representation vector similarities and the actual GO-based semantic similarities of corresponding protein pairs. Higher correlation values indicate greater success in representing functional similarity.





  §.§.§ Protein–Protein Binding Affinity Estimation

This task involves assessing the ability of a protein representation to predict alterations in binding affinity caused by protein mutations. Each protein pair is associated with a real value, making it a regression task that reflects binding affinity changes. We follow the approach of  and use the SKEMPI dataset <cit.> with an evaluation based on the mean squared error obtained by 10-fold cross-validation.





 §.§ Baselines

Incorporating SOTA models as our baselines, we used ProtBert <cit.> and ESM-1b <cit.> as benchmarks, following KeAP <cit.> for fair comparison. The former is based on a BERT architecture whereas the latter is based on the AlphaFold <cit.> architecture adapted for the tasks mentioned in Section <ref>.
These models are considered as the top-performing models in the absence of factual knowledge integration. 
Furthermore, we subjected our analysis to the scrutiny of the most recent and advanced knowledge-driven pre-training model, KeAP <cit.>.
We draw the reader's attention that we exclude the results of OntoProtein <cit.> in this paper, as it was presented by KeAP <cit.> that it is significantly inferior on a dataset without leakage.
For all baseline models, we generate representations during the training phase and subsequently assess their performance in various biological tasks. For each specific task, we followed the convention of selecting SOTA methods, just as prior research has done.




  §.§.§ Task Specific Baselines

Building upon the KeAP <cit.> baselines, we extend our evaluation to encompass additional SOTA models that are known to perform well on specific tasks within the protein domain.
For the contact and homology tasks we include variants of LSTM <cit.>, ResNet <cit.>, and Transformer <cit.>, as well as ProteinBert <cit.> for the homology.
Following , we examine DPPI <cit.>, DNNPPI <cit.>, PIPR <cit.>, and GNN-PPI <cit.> for the protein-protein interactions (PPI) task.
Also, we incorporate PIPR <cit.> for affinity estimation and
MSA Transformer <cit.> and ProtT5-XL <cit.> for semantic similarity inference.




 §.§ Implementation Details

During pre-training, GOProteinGNN is trained for 50k steps using a learning rate of 1e-4, weight decay of 0.02, and a batch size of 512 (proteins) over four GPUs (NVIDIA A40, 48GB Memory each). RGCN layers L are set to 1 and GNN Knowledge Injection (GKI) components (referred to as k) are set to 3.
For downstream tasks, we use hyperparameters from GNN-PPI <cit.> for PPI prediction and PROBE <cit.> for binding affinity and similarity inference. For TAPE <cit.>, we follow KeAP <cit.> and perform a grid search to optimize batch size and learning rate.
In terms of computational complexity, GOProteinGNN, like other GCNs, operates at approximately O(|V| * |E|). Several optimizations, including limiting the number of hops per protein, reduce the complexity to O(|V|), achieving performance akin to KeAP. During application, the complexity of both algorithms is O(1).
For full implementation details, refer to the provided code.




§ EMPIRICAL RESULTS

In this section, we present our empirical results. We validate the statistical significance of the results using a one-tailed paired Student’s t-test with a 95% confidence level. Significant results are marked with an asterisk (*) in the tables.




 §.§ Contact Prediction

In Table <ref>, we present GOProteinGNN's performance on the contact prediction task compared to six baseline methods. GOProteinGNN significantly outperformed all baseline methods across all categories and metrics.
In particular, GOProteinGNN showed superior performance compared to ProtBert <cit.> and ESM-1b <cit.>, which lack knowledge graph information.
This highlights the importance of incorporating knowledge during representation learning, especially for pairwise amino acid tasks (e.g., contact prediction task), where success depends on understanding the proteins' functions. 

Moreover, KeAP <cit.> incorporates a knowledge post-encoder, limiting its capacity to seamlessly integrate external knowledge into protein representations, as reflected in the performance gains of GOProteinGNN over KeAP. These findings highlight the significant impact of incorporating knowledge graph information in GOProteinGNN via the GKI layers, leading to superior performance in the contact prediction task.




 §.§ Remote Homology Detection

In Table <ref>, we present GOProteinGNN's performance compared to seven baseline methods on the remote homology detection task. 

Despite the challenging nature of the task, due to the large number of potential protein folds, GOProteinGNN significantly outperformed all methods, demonstrating its effectiveness in handling the complexities of remote homology detection. Specifically, GOProteinGNN achieved an advantage over KeAP <cit.>, further affirming its ability to capture essential biological contexts and interactions which is crucial for predicting the protein fold structure, as it processed the whole protein knowledge graph.

As remote homology detection is a sequence-level classification task, we can see that both KeAP <cit.> and ProtBert <cit.> achieved equal results, suggesting KeAP fails to leverage the factual knowledge in the protein level in this task. 

This task highlights the novelty of our algorithm which allows for the integration of information at both the individual amino acid level and the entire protein level.








 §.§ Protein-Protein Interaction Identification

Table <ref> presents the results of the protein-protein interaction (PPI) identification task. GOProteinGNN showed strong performance across the board, achieving the highest F1 scores across the majority of datasets and settings with statistical significance,
which can be beneficial for drug discovery or personalized medicine. 

In particular, for SHS27K, GOProteinGNN obtained SOTA results with an increase of approximately 1% using BFS and 4% using DFS compared to KeAP <cit.>. For the SHS148K dataset <cit.>, GOProteinGNN demonstrated superiority over all other models when using BFS, while maintaining its performance as the second-best performer when using DFS. The trend of declining improvement can be attributed to the increasing amount of fine-tuning data (from SHS27K to SHS148K) that reduces the impact of pre-training and the knowledge graph injection. 

In conclusion, the results highlight the strengths of our GOProteinGNN, which enable better biological understanding of the protein, leading to success in predicting interactions between proteins.








 §.§ Semantic Similarity Inference

Table <ref> presents the results of the semantic similarity inference task. In particular, our model outperformed all baselines significantly in capturing semantic similarities within the biological process (BP) category.

Biological processes play a crucial role in many drug development processes, making this category particularly important as it pertains directly to drug development. The success in this task underscores GOProteinGNN's efficacy in acquiring enriched representations that encompass the complex biological processes in which proteins are involved, due to the whole protein knowledge graph learning, surpassing the capabilities of alternative approaches. 

Furthermore, KeAP <cit.> showed inferior performance compared to ESM-1b <cit.>, which does not use factual knowledge. This result is surprising given that both GOProteinGNN and KeAP utilize the same factual knowledge repository, yet KeAP's method of knowledge injection failed to outperform the model that does not use knowledge at all.










 §.§ Protein–Protein Binding Affinity Estimation

Table <ref> presents the results of the binding affinity estimation task. Evidently, GOProteinGNN outperformed PIPR <cit.> and ProtBert <cit.> with a significant margin of improvement. 
Furthermore, it is noteworthy that GOProteinGNN exhibited a slight improvement over KeAP <cit.>.

The performances of ESM-1b <cit.> and GOProteinGNN are comparable, with no statistically significant differences observed between them. This similarity in performance may be attributed to the task's sensitivity to subtle structural changes, which ESM-1b captures effectively through its extensive pre-training. In this context, the external factual knowledge leveraged by GOProteinGNN does not provide an advantage.







§ ABLATION TESTS




 §.§ GNN Knowledge Injection (GKI) Ablation


The ablation test for the GKI component is presented in Table <ref>. The full model results highlight the effectiveness of GOProteinGNN's model components, where the model demonstrated superior performance across all metrics, underscoring its pivotal role.

Ignoring relations through the GKI component implies using GCN <cit.> instead of RGCN <cit.> and ignoring edge type during the knowledge graph learning. Without relation learning, we observe a slight precision decrease, emphasizing its importance.
Moreover, excluding the GKI component (i.e., without GNN) substantially degrades performance, reinforcing the significance of the GNN mechanism for effective knowledge integration. 

These experiments highlight the contribution of the GKI component and the GNN mechanism to GOProteinGNN's ability to capture essential biological contexts and interactions, leading to performance improvements.







 §.§ Number of GKI Components Ablation

We wish to experiment with the number of optimal GKI components. Each GKI component adds an additional knowledge graph injection to the model.
In the case of GOProteinGNN, the number of GKI components (referred to as k) is set to 3. In this ablation test, we reduce the number of GKI components to observe their impact.
As shown in Table <ref>, reducing knowledge injections in GOProteinGNN results in a slight performance decrease. This decrease highlights the significance and meaningfulness of the GKI component.






§ CONCLUSIONS


In this paper, we introduced GOProteinGNN, a novel architecture that enhances PLMs by integrating protein knowledge graph information during the creation of amino acid-level representations.
Specifically, GOProteinGNN leverages a novel GKI procedure that leverages the [CLS] token representation of the PLM during the encoder stage.
Our methodology enables the incorporation of information at both the individual amino acid and protein levels, fostering a comprehensive learning process through graph-based techniques.
Furthermore, GOProteinGNN uniquely learns the entire protein knowledge graph during training, allowing it to capture broader relational nuances and dependencies beyond mere triplets as done in previous work.

GOProteinGNN showcases its potential across diverse protein-related tasks, such as amino acid contact prediction and protein-protein interaction identification. We posit that the observed performance improvements arise from a deeper comprehension of proteins themselves and the associated factual knowledge. This understanding allows generalization for never-before-seen proteins, as demonstrated on downstream tasks.
The ability of GOProteinGNN to jointly model sequential data and knowledge graph information empowers it to handle complex biological scenarios effectively. This fusion of language model-based representation and knowledge graph processing makes GOProteinGNN a valuable tool for advancing research in the field of bioinformatics and drug development.

GOProteinGNN can enhance multiple real-world tasks, including protein structure prediction and drug-target interaction. By leveraging both the protein knowledge graph and sequential data modeling, GOProteinGNN effectively propagates information regarding proteins with biological connections. 
In drug development, it can facilitate virtual screening by prioritizing compounds for binding to specific protein targets, underscoring its potential to provide deeper insights into complex biological phenomena and expediting the drug discovery process.

For future directions, we envision the possibility of incorporating additional factual knowledge beyond proteins and extending our graph learning framework to more advanced methodologies. Additionally, exploring the applicability of our algorithm to general Natural Language Processing (NLP) tasks could be an intriguing avenue to explore.
















































































































































































































































































































































































































































































































































































ACM-Reference-Format




 Paper Title: "Patriarchy Hurts Men Too." Does Your Model Agree? A Discussion on Fairness Assumptions.





A Discussion on Fairness Assumptions

















University of Antwerp, Antwerp, Belgium {marco.favier, toon.calders}@uantwerpen.be




“Patriarchy Hurts Men Too.”
 Does Your Model Agree?
 A Discussion on Fairness Assumptions
    Marco Favier Toon Calders
    August 5, 2024
=========================================================================================



§ ABSTRACT 
The pipeline of a fair ML practitioner is generally divided into three phases: 1) Selecting a fairness measure. 2) Choosing a model that minimizes this measure. 3) Maximizing the model's performance on the data. In the context of group fairness, this approach often obscures implicit assumptions about how bias is introduced into the data. For instance, in binary classification, it is often assumed that the best model, with equal fairness, is the one with better performance. However, this belief already imposes specific properties on the process that introduced bias. More precisely, we are already assuming that the biasing process is a monotonic function of the fair scores, dependent solely on the sensitive attribute.

We formally prove this claim regarding several implicit fairness assumptions. This leads, in our view, to two possible conclusions: either the behavior of the biasing process is more complex than mere monotonicity, which means we need to identify and reject our implicit assumptions in order to develop models capable of tackling more complex situations; or the bias introduced in the data behaves predictably, implying that many of the developed models are superfluous.



§ INTRODUCTION

It is now well known that AI models may exhibit unfair behaviors against historically marginalized groups and minorities due to biases present in the data used to train these models. 
Famous cases like SyRI <cit.>, the toeslagenaffaire <cit.>, and the COMPAS <cit.> algorithm have proven beyond any doubt that the risks involved in using classifiers without concern for bias are real and have negatively affected, and still affect, communities that already suffer from unjust discrimination.

Since this acknowledgment, much effort has been spent in the last decade to study the phenomenon from a scientific and philosophical perspective, and many solutions have been proposed to improve the fairness of models.

Group fairness measures try to tackle this problem by evaluating metrics that describe how unfairly entire sensitive groups have been treated compared to each other.

For instance, the Demographic Parity Difference evaluates how statistically likely it is to be assigned a positive label for one group against another. 

Similarly, the Equal Opportunity Difference computes the difference between the false negative rates evaluated for different communities. Barocas et al. <cit.> provide a systematic overview of the most common fairness measures.

On top of these measures, many different models able to minimize (un)fairness measure have been developed.

In general, these models are evaluated by checking which one performs best according to fairness and/or accuracy, with the underlying assumption that, as long as the fairness measure is low, the preferable model is the one that performs better.

Nonetheless, understanding when or even why a model is better than another is not an easy question to answer; fairness will always remain undoubtedly a subjective concept, and any comparison between models or fairness measures is reasonable only within very specific contexts.

But also formally, the field of fairness lacks a theoretical and foundational background to better explain fair models and their behavior. It is often unclear how the choice of a specific model impacts the outcome and under what circumstances a given model should be used.

We believe that one of the reasons for this confusion comes from the fact that a number of these models seem to work under similar assumptions, trying to be as faithful as possible to the data. But, if we believe that our data is biased and thus not representative of the population, it's debatable if fidelity is the best approach to fairness.

This assumption is often implicitly made and not discussed, and it has important practical but also theoretical consequences on how we should approach fairness.


In this work, we explore a number of implicit assumptions often present in fairness measures and models.
In particular, we show how these assumptions all share the same underlying idea: the biasing process, that is, the process that has introduced bias into our data, is a monotonic function on each sensitive group. 

This means that if two individuals from the same sensitive group are ranked fairly, then the one with the highest score will always be ranked higher according to unfair probabilities as well. People are swapped in the ranking only if they belong to different sensitive groups.

But this has very strong implications and often does not reflect reality. For instance, consider a model screening CVs for a job position that unfairly privileges men. If a male candidate writes using the first person, while another uses the third person, the model may privilege the latter.
A different outcome only because it's easier to assess the gender from all the “he"s instead of the “I"s. Wouldn't it then be fair to improve the ranking of the first candidate as well?

Discrimination within the same sensitive group is a well-known phenomenon, so discussing whether a model should consider it or not is an important question that shouldn't be implicitly answered.

So, is your model aware that the patriarchy hurts men too?



§ RELATED WORK

We work under the so-called Fair World Framework <cit.>. Under this framework, it is assumed that an inaccessible fair world exists. This would be an ideal world where no bias or prejudice exists, untouched by the systemic issues our world suffers from. Data sampled from this world would be considered fair, and decisions based on it would be considered unbiased. We don't have access to this world; instead, we have access to data collected from our world, which is a biased version of the fair world.

This framework has been used to study the relationship between fairness and accuracy <cit.>, providing a useful tool to understand the objective of fairness in ML. It represents a paradigm shift from the classical fairness definition: a fair model is not one that satisfies some fairness constraint, but one that is accurate on the fair world.

We will adopt this framework to explore the space of theoretically optimal fair models for multi-label classification. For binary classification, Corbett-Davies et al. <cit.> and Menon et al. <cit.> have proven that the best fair model, where fairness means minimizing a fairness measure, can always be achieved by applying a different threshold for each sensitive group of the score function. Our work will explore how their results apply within the fair world framework.

On the other hand, Wick et al. <cit.> have shown that a fair model can be more accurate than an unfair one. They argue that when data are biased, fairness can improve accuracy on the originally unbiased data, contrary to the common belief that fairness and accuracy are in trade-off.

They claim that this phenomenon is often misunderstood because we do not have access to unbiased data, and we should not simply assume that the data at hand is the best data to evaluate a model.

More recently, Goethals et al.<cit.> have conducted experiments to see if fair models change the ranking within sensitive groups. Our framework provides a much-needed theoretical background for their results. Moreover, analyzing how a fair model can impact individuals within the same sensitive group has also been explored in the context of individual fairness <cit.>.



§ NOTATION AND PROBLEM SETTING

We will focus on a fairness multi-label classification problem. More formally, we consider a probability space (X,μ) called the feature space, where

    
  * X is the set of possible individuals, each described by a feature vector x∈ X.
    
  * μ is a probability measure used to sample individuals from  X.

In a classical group fairness fashion, we assume that the set X is partitioned into different sensitive groups. 

These groups generally represent communities that have historically been treated differently for a variety of unjust reasons. For instance, the sensitive groups can represent different genders or ethnicities. 
For convenience, we consider only two sensitive groups. More formally, a sensitive partition {,} of X is a binary partition of X into non-zero measurable sets. An element A∈{,} is called a sensitive group.

The reader should be aware that all our results generalize to any number of sensitive groups. Moreover, the symbols “” and “” used to identify the groups are completely arbitrary; they are not meant to represent genders or any other specific sensitive attribute.

Together with the space X, we also have a categorical label Y∈{y_1,…, y_n}  that is assigned to each individual. The main difference between a classical decision problem and our setting is how the label is sampled for each individual.

Instead of assuming the existence of a single conditional distribution, we consider two of them: p_u(y | x) and p_f(y | x). More specifically:

    
  * p_u(y | x):=(p_u(y_1| x),…,p_u(y_n| x)) is the observable distribution that we can infer from the data we have, but it is considered unfair towards some sensitive groups due to its biased information.
    
  * p_f(y | x):=(p_f(y_1| x),…,p_f(y_n| x)) is instead the fair distribution. Decisions based on this distribution are considered unbiased, but it's not possible to access it directly. This distribution represents what we would observe if there were no bias or prejudice in our society.

As conditional distributions, both p_f and p_u are measurable functions

    p_u,p_f X →Δ^n

where Δ^n is the n-dimensional simplex: 

    Δ^n:={𝐭∈ [0,1]^n t_1+…+t_n=1}


The relationship between the fair and unfair distributions is determined by a biasing process β that transforms fair probabilities into unfair ones. Similarly, we can describe a debiasing process β^-1 that does the opposite.
We can show that these processes always exist:

    Given probability distributions p_f and p_u, there exists a biasing and a debiasing process, β,β^-1Δ^n× X→Δ^n, respectively, such that
    
    p_u(y | x) = β(p_f(y | x), x) and  p_f(y | x) = β^-1(p_u(y | x), x)

    
        The proof is almost tautological: since β and β^-1 depend directly on X, then the use of the conditional distributions is superfluous. We can simply set 
        
    β(- , x):=p_u(y | x) =  and β^-1(-, x):=p_f(y | x)

        which proves the thesis.
    

Formally, using β^-1 to indicate the debiasing process is an abuse of notation, since the debiasing process is not technically the inverse of the biasing process. Nonetheless, we will use this notation to underline their mutual relationship.

The fact that the biasing (or debiasing) process depends on x as well as p_f (or p_u) may be confusing. This is done to underline that the biasing process transforms the fair probabilities into unfair ones. However, this requirement alone is not sufficient, as the transformation may also vary depending on the specific individual considered, necessitating the dependency on x.


Our objective will be to assign a label to each individual using a decision function. Formally, a decision d is a function 

    d X→Δ^n   x↦d(x)=(d_1(x),…,d_n(x))

The use of Δ^n as the codomain of a decision allows us to consider probabilistic decisions: for a data point x, d(x) represents a distribution on Y, from which we sample a predicted label Y∈{y_1,…, y_n}. Notice that this doesn't force us to consider only probabilistic decisions; in fact, we can also consider deterministic decisions by simply setting d_i(x)=1 for some i∈[n].

We can evaluate a decision based on how well it performs on each label. In particular, we define  P(Y = y_i, Y = y_j) as the following value:

    P(Y = y_i, Y = y_j):=∫_Xd_i(x)p(y_j | x)μ(x)

We will use P_u when the calculation is computed using p_u(y | x), and similarly P_f when fair probabilities are used. This allows us to define the following:

    Given two decisions d,d', we will say that d is Pareto non-inferior (or non-superior) to d', and write d≥d' if and only if for all i∈ [n]
    
    P(Y= y_i, Y=y_i)≥ P(Y' = y_i, Y=y_i)

where Y∼d(x) and Y'∼d'(x). We will use ≥_u or ≥_f to indicate according to which conditional probability p_f,p_u we have calculated P(Y= y_i, Y=y_i).

It's important to notice that if d≤d' according to some probability, d' is better in the sense that it can better represent the data. It doesn't mean that d' is more fair.

The Pareto order on decisions is not a total order, meaning two decisions may not be comparable. So it is possible that d≱d' and d'≱d. Moreover, there are cases where two different decisions d≠d' are equivalent, meaning d≥d' and d'≥d. 

We are obviously interested in finding the best fair decision possible, that is, a decision that is Pareto maximal according to ≤_f.


    A decision d on U⊆ X is said to be unbiased if it is maximal on U according to ≤_f and depends only on fair probabilities. Formally:
    
        
  * d is Pareto maximal for ≤_f; that is, if d'≥_f d then d≥_f d' also holds.
        
  * d depends only on fair probabilities; that is, there exists a function D such that d(x) = D(p_f(y | x)).
    

Notice that, since we do not have access to the fair distribution p_f, we are unable to determine if a decision is unbiased or not. This is why we need to use fairness measures as proxies to ensure fair decisions. The following example should clarify this point.

    Suppose every individual in X is described by two independent features: the number of languages they speak, ranging from one to three (s=1,2,3), and whether they are religious or not (r=0,1), which is a sensitive attribute.
    Moreover, all six possible combinations of attributes appear to be equally distributed. You want to predict whether a person will obtain a PhD in their lifetime (y={0,1}). The fair probabilities are the following:
    
    p_f(y=1| s,r=0)=
            
                1%   if  s=1
    
                2%   if  s=2
    
                4%   if  s=3
                  p_f(y=1| s,r=1)=
            
                1%   if  s=1
    
                3%   if  s=2
    
                3%   if  s=3

    These probabilities can be considered fair according to demographic parity: the chances of obtaining a PhD are the same for religious and non-religious individuals.
 
    Sadly, academia is known to be biased against religion. Consequently, in our unfair society, many religious people don't even apply for a PhD. As a result, the unfair probabilities are as follows:
    
    p_u(y=1| s,r=0)=
            
                1%   if  s=1
    
                2%   if  s=2
    
                4%   if  s=3
                  p_u(y=1| s,r=1)=
            
                0%   if  s=1
    
                1%   if  s=2
    
                2%   if  s=3

    Notice how the new probabilities do not satisfy demographic parity anymore. Consider now the decision d such that
    
    d_1(s,r)=
            
                1    if  p_u(y=1| s,r)≥ 2%
    
                0    otherwise

    Despite being purely based on the probabilities derived from the data, the decision is biased because d_1(2,1)=0 and d_1(2,0)=1, whereas p_f(y=1| s=2,r=1)> p_f(y=1| s=2,r =0).
    Instead, if we take the decision d' such that
    
    d'_1(s,r)=
            
                1    if  p_u(y=1| s,r)≥ 2% and  r=0
    
                1    if  p_u(y=1| s,r)≥ 1% and  r=1
    
                0    otherwise

    then d' is unbiased since it's equivalent to 
    
    d'_1(s,r)= 1  p_f(y=1| s,r)≥ 2%

    Notice how the second decision explicitly depends on the sensitive attribute r, which is necessary to contrast the bias present in the data.

Unbiased decisions are not always able to satisfy fairness constraints. For instance, in the previous example, it's not hard to find unbiased decisions that are still not fair according to demographic parity.

But it is true that if there are differences in the outcomes for the two communities, these can be fully explained using fair probabilities. This framework allows us to discuss fairness even in cases where the communities have different distributions. This is because the definition of fairness is now tied to the fair world and its probabilities. For instance, if we assume that p_f is identically distributed on  and , then unbiased decisions will also be fair.

This idea formalizes why fairness metrics and losses can be a useful tool for finding fair and unbiased decisions: within our framework, including a penalty term for unfairness in a loss function can be viewed as placing a prior on the decision space that favors unbiased decisions.

That being said, finding optimal decisions, especially in multi-label classification, is not an easy task. Fortunately, Weller's Theorem <cit.> provides a way to find such decisions when we are able to model the underlying probability distribution. To introduce it we need first the following definition:

	A sequence (ω^(k))_k∈ of point in Δ^n∩ (0,1)^n is said to be well-behaved if for all i,j∈ [n] the limit lim_k→∞ω_i^(k)/ω_j^(k)



	converges in [0,∞].


    A decision d X →Δ^n is Pareto maximal if and only if there exists a well behaved sequence (ω^(k))_k∈ such that almost everywhere on X, for all i,j∈ [n]
    
    d_i(x)>0 ⇒∃ k∈ s.t. p(y_i| x)/p(y_j|  x)≥ω_i^(k)/ω_j^(k)  ∀ k≥k

    Such sequence (ω^(k))_k∈ is said to be w-associated with d(x).





§ THE GROUP FAIRNESS ASSUMPTIONS


In this section, we will introduce a number of assumptions that fairness literature often implicitly uses. 

Before continuing, we want to make it clear that we don't think having assumptions is wrong, nor that these specific assumptions necessarily are. In fact, we believe that having assumptions is often unavoidable when developing a model. The famous quote by George Box, “All models are wrong, but some are useful,” makes it explicit that without assumptions on our data, even overly simplistic ones, we couldn't build a model. Even the entirety of the field of fair ML is based on the assumption that some sensitive attributes should not influence decisions.

It is important, though, to make these assumptions explicit, especially when they are used to justify the development of a model. In the same way, it is also important to realize the consequences of these assumptions. Fair ML is a clear example of the importance of doing so: when researchers began formalizing the assumption that a model should not be biased, they realized there were different ways to achieve this, to the point that different views on fairness can conflict with each other. Without stating the assumptions on fairness, we wouldn't have been able to prove the now-famous impossibility results <cit.>.

Moreover, we are aware that the assumptions in the following list, from a mathematical point of view, are very similar. However, we believe that even if two assumptions may be equivalent mathematically, their different rephrasing may evoke different opinions and ideas. We hope that an engaged reader will consider some of these assumptions more reasonable than others and be surprised to discover how they are interconnected.

Let d,d' A→Δ^n be decisions on a  sensitive group A. Then

    d<_u d' ⇒d<_f d'


The “No Harm, No Foul” assumption states that decisions can be unfair only when compared between different sensitive groups. In isolation, any decision regarding a single sensitive group is considered fair, which means that improving a decision for a single sensitive group can also be considered an improvement in fair probabilities. This resembles the point of view of earlier works on fairness, where fairness was considered simply as a constraint to satisfy <cit.>.

    Let d,d' A→Δ^n be decisions on a  sensitive group A. Then
    
    d<_f d' ⇒d<_u d'


The “Fairness as Optimization” assumption states that if, for a sensitive group, we theoretically manage to improve a decision according to fair probabilities, then we should at least be able to see the improvement for that group according to unfair probabilities as well. This means that an optimization step in fair probabilities can also be found in unfair probabilities. This, we believe, is the point of view of many in-processing methods <cit.>.

    Let d X→Δ^n be a decision. If  d|_A is not a maximal decision according to ≤_u for some sensitive group A, then d is not unbiased.

The “Representation Matters” assumption states that if we fail to correctly represent a sensitive group according to our data, then our decision cannot be considered fair. This is tied with the concept of “cherry-picking” in fairness literature <cit.>


    Let d X→Δ^n be an unbiased decision. There exist well-behaved sequences (ω_^(k))_k∈, (ω_^(k))_k∈ and associated decisions d_→Δ^n, d_→Δ^n on ≤_u such that d|_=d_ and d|_=d_.

The “Affirmative Action” assumption states that in order to build a fair decision, we can set different thresholds for different sensitive groups to counterbalance the bias present in the data.


    Let d X→Δ^n a maximal decision for ≤_u. Then d|_ and d|_ are ω-associated with two well-behaved 
    sequences (ω_^(k))_k∈ and (ω_'^(k))_k∈ according to ≤_f.

The “Double Standard” assumption states that an optimal unfair decision treats the two communities as if there were two different “standards” for being selected as a candidate by the decision. This and the previous assumption are the two most common assumptions for post-processing methods <cit.>. 


§ CONSEQUENCES OF THE FAIRNESS ASSUMPTIONS



 §.§ Preliminary Results

We start by showing how all the presented fairness assumptions imply either the “Affirmative Action” assumption or the “Double Standard” one. 

    
        Suppose by contradiction, that for a sensitive group A the decision d|_A cannot be associated with any well-behaved sequence according to ≤_u. Then, by Weller's Theorem, d|_A cannot be maximal on ≤_u. So there exists a decision d' A→Δ^n such that d|_A<_u d'. By the No Harm, No Foul assumption we have d|_A<_f d', then we can extend d' to a new decision on X so defined
        
    d'(x)=
            	d'(x)    if  x∈ A	
    d(x)    otherwise

        Clearly d<_f d', contradicting the maximality of d.
    



    
        As before, with p_f and p_u swapped.
    










    
        Representation Matters can be restated as for an unbiased d then d|_A is maximal according to ≤_u. By Weller's Theorem we have the thesis.
    






 §.§ The Main Consequence

Since we have proven that our previous assumptions had to satisfy either Affirmative Action or Double Standard, we can now see how these final two assumptions require specific properties on the biasing process. Namely, that the biasing process or the debiasing process needs to be monotonic.

	The Affirmative Action assumption holds if and only if the probabilities satisfy the following conditions:
    
    	
  *  for all i∈ [n] and for almost all x∈ X 
    p_u(y_i| x)=0⇒ p_f(y_i| x)=0

    	
  *  for all i,j∈ [n] and for almost all x,x'∈ A
        
    p_f(y_i| x)/p_f(y_j| x)≤p_f(y_i| x')/p_f(y_j| x')⇒p_u(y_i| x)/p_u(y_j| x)≤p_u(y_i| x')/p_u(y_j| x')

        for all sensitive groups A∈.
    
    
    	(⇒<ref>) Suppose by contradiction that there exist i and a non-zero set S such that
    	
    p_f(y_i| x)>0
        	 and 
        	p_u(y_i| x)=0  ∀ x∈ S

    	and let's consider an optimal decision d on the fair probabilities w-associated with (ω^(k))_k∈ defined as follows
    	
    ω_s^(k):=
        	
        		1/k    if  s = i
    
        		(k-1)/k(n-1)    otherwise

    	Clearly d_i(S)=1, let's now take a sensitive group A such that μ(A∩ S)>0. Such group must exists since the sensitive groups are a finite partition of X. Now we have p_u(y_i| x)=0 for all x∈ S∩ A, so d|_A it is not optimal according to Theorem <ref>.

    	
        (⇒<ref>) Let's suppose by way of contradiction that the implication is not true for some sensitive group A. Formally this means that there exists a non-zero subset S× S'⊆ A× A such that for all x∈ S, x'∈ S' we have 
        
    p_f(y_i| x)/p_f(y_j| x)≤p_f(y_i| x')/p_f(y_j| x')  
            and  p_u(y_i| x)/p_u(y_j| x)>p_u(y_i| x')/p_u(y_j| x')

        Notice that we do not have to check the case where we have p_u(y_i| x)/p_u(y_j| x)=0/0 since we would also have  p_f(y_i| x)/p_f(y_j| x)=0/0 by condition <ref>.
        Let's now take a threshold t such that 
        
    μ(p_f(y_i| x)/p_f(y_j| x)≤ t≤p_f(y_i| x')/p_f(y_j| x')| (x,x')∈ S× S')>0

        such threshold always exists since 
        
    p_f(y_i| x)/p_f(y_j| x)≤p_f(y_i| x')/p_f(y_j| x')  ∀ x∈ S, x'∈ S'

        implies
        
    sup_x∈ Sp_f(y_i| x)/p_f(y_j| x)≤inf_x'∈ S'p_f(y_i| x')/p_f(y_j| x')










		Consider now the following well-behaved sequence (ω^(k))_k∈ defined as follows
		
    ω_s^(k):=
    		
    			 t/k(t+1)    if  s = i
    
    			1/k(t+1)    if  s = j
    
    			(k-1)/k(n-2)    otherwise

		and let d be a decision w-associated with it.
		In particular we can take d such that if we consider the following sets
		
    S_j:={x∈ Sp_f(y_i| x)/p_f(y_j| x)≤ t}

		
    S'_i:={x'∈ S'p_f(y_i| x')/p_f(y_j| x')≥ t}

		we have that d_j(S_j)>0 and d_i(S'_i)>0.

		We claim that such decision cannot be optimal according to p_u: if it were it would exist a w-associated well-behaved sequence (ω'^(k))_k∈. In particular the value ω_i^(k)/ω_j^(k) converges to some value t'. Then for all x∈ S_j and x'∈ S'_i
		
    t'≤p_u(y_i| x',a)/p_u(y_j| x',a)
    		<
    		p_u(y_i| x,a)/p_u(y_j| x,a)≤
    		t'

		which is absurd.














        

        (⇐<ref>+<ref>) We first prove that β, the biasing function that transforms fair probabilities to unfair ones does not depend on x∈ X almost everywhere. This follow directly from the hypothesis since almost everywhere we have
        
    p_f(y_i| x)/p_f(y_j| x)=p_f(y_i| x')/p_f(y_j| x')⇒p_u(y_i| x)/p_u(y_j| x)=p_u(y_i| x')/p_u(y_j| x')

        for all i,j∈[n].

        Hence 
        
    p_f(y| x)=p_f(y| x')
            ⇒ p_u(y| x)=p_u(y| x')

        So people from the same sensitive group with the same fair probabilities are assigned the same unfair probabilities as well. So for all sensitive groups A, we can find a parameterized function β^A such that almost everywhere
        
    p_u(y| x)=β(p_f(y| x),x) = β^A(p_f(y| x))

        and for all t,t'∈Δ^n∩ (0,1)^n
        
    t_i/t_j≤t'_i/t'_j⇒β_i^A(t)/β_j^A(t)≤β_i^A(t')/β_j^A(t')

        Let us now consider a optimal decision d according to fair probabilities w-associated with a well-behaved sequence (ω^(k))_k∈. Furthermore, we can assume that for all i,j∈[n] the sequence (ω_i^(k)/ω_j^(k))_k∈ is monotonic. If that was not the case for some i,j, that is if (ω_i^(k)/ω_j^(k))_k∈ were not a monotonic sequence, we can consider a monotonic subsequence, which always exists since any real sequence has a monotonic subsequence. By repeating this process for all i,j∈[n] we can ensure the sequence's monotonicity for every couple of components.
        We claim that for all sensitive groups A, the decision d|_A is, according to p_u(y| x), w-associated with the well-behaved sequence (β^A(ω^(k)))_k∈ and hence maximal.

        First, let us show that the sequence is indeed well-behaved. Notice that because of condition <ref> we know that β^A(Δ^n∩ (0,1)^n)⊆Δ^n∩ (0,1)^n in particular (β^A(ω^(k)))_k∈ is indeed a sequence in Δ^n∩ (0,1)^n. Also for all i,j∈ [n], the sequence (β_i^A(ω^(k))/β_j^A(ω^(k)))_k∈ is monotonic and hence converges in [0,+∞]. The monotonicity follows from the fact that
        
    ω_i^(k)/ω_j^(k)≤ω_i^(k+1)/ω_j^(k+1)⇒β_i^A(ω^(k))/β_j^A(ω^(k))≤β_i^A(ω^(k+1))/β_j^A(ω^(k+1))

		Now we only need to prove that d|_A is w-associated with the sequence.
		We know that almost everywhere on A, for all i,j∈ [n]
		
    d_i(x)>0 ⇒∃ k∈ s.t. p_f(y_i| x)/p_f(y_j|  x)≥ω_i^(k)/ω_j^(k)  ∀ k≥k

		then using property <ref> for β^A we obtain for all k≥k
		
    p_u(y_i| x)/p_u(y_j|  x)=β^A(p_f(y_i| x))/β^A(p_f(y_j|  x))≥β_i^A(ω^(k))/β_j^A(ω^(k))

		Hence
		
    d_i(x)>0 ⇒∃ k∈ s.t. p_u(y_i| x)/p_u(y_j|  x)≥β_i^A(ω^(k))/β_j^A(ω^(k))  ∀ k≥k

     



    The Double Standard assumption holds if and only if the probabilities satisfy the following conditions:
    
    	
  * for all i∈ [n] and for almost all x∈ X 
    p_f(y_i| x)=0⇒ p_u(y_i| x)=0

    	
  * for all i,j∈ [n] and for almost all x,x'∈ A
    p_u(y_i| x)/p_u(y_j| x)≤p_u(y_i| x')/p_u(y_j| x')⇒p_f(y_i| x)/p_f(y_j| x)≤p_f(y_i| x')/p_f(y_j| x')

        for all sensitive groups A∈.
    
    
    	The proof is analogous to the one of Theorem <ref> with p_u and p_f swapped.

    	
    

In the binary case Y={0,1}, the requirements on the biasing (and debiasing) function β can be simplified as follows:

    Let Y={0,1} be a binary label. The Affirmative Action assumption holds if and only if, for almost all x,x'∈ A, it holds
    
    	
  * p_u(y=1| x)=0⇒ p_f(y=1| x)=0
    	
  * p_u(y=1| x)=1⇒ p_f(y=1| x)=1
    	
  * p_f(y=1| x)≤ p_f(y=1| x')⇒ p_u(y=1| x)≤ p_u(y=1| x')
    
    where A is a sensitive group. In particular, on each sensitive group, the biasing function β is non-decreasing function of p_f(y=1| x).
    
        It follows directly from Theorem <ref> and the fact, because the label is binary, we have p(y=0| x)+p(y=| x)=1.
    

And, similarly, for the Double Standard assumption:

    Let Y={0,1} be a binary label. The Double Standard assumption holds if and only if, for almost all x,x'∈ A, it holds
    
    	
  * p_f(y=1| x)=0⇒ p_u(y=1| x)=0
    	
  * p_f(y=1| x)=1⇒ p_u(y=1| x)=1
    	
  * p_u(y=1| x)≤ p_u(y=1| x')⇒ p_f(y=1| x)≤ p_f(y=1| x')
    
    where A is a sensitive group. In particular, on each sensitive group, the debiasing function β is non-decreasing function of p_u(y=1| x).

As previously stated, for a binary label, the assumptions we have shown are requiring that the biasing or the debiasing process is monotonic. When the label is not binary, the requirements are more complex, but the idea remains the same: the process must be relatively monotonic according to the ratio between labels. It's also noteworthy that, unless bias is strictly monotonic, the Affirmative Action and Double Standard assumptions are not exactly equivalent. 


§ DISCUSSION

In the previous section, we have proven how all the assumptions we have presented imply that the biasing process behaves as a monotonic function. Given this result, it would be interesting to see if fair models actually behave accordingly. Luckily, there is already some literature that has started to investigate this.
In their paper “Reranking individuals: The effect of fair classification within-groups”<cit.>, Goethals et al. have researched how fair machine learning models behave with respect to each sensitive group. In particular, they measure how often different fair models change the ranking of individuals within a sensitive group. They first adopt a fair-agnostic model to create a baseline ranking based on the score of each individual. Then, they apply different fair models to the same data and measure how often the ranking changes. To do so, they use the Kendall Tau coefficient, a measure of the similarity between two rankings, defined as

    number of concordant pairs-number of discordant pairs/number of pairs

Based on their results, they conclude that a number of models often inexplicably rerank individuals within the same group. In particular, they caution against the use of pre- and in-processing methods when no within-group bias is expected, since these methods are the ones that suffer the most from this issue.

Our theoretical framework is necessary to formally support their claims. They correctly point out that the reranking of individuals within a group is a sign that the model is not behaving as expected, and they intuitively connect this to the fact that this behavior can be considered acceptable only when bias affects even people within the same sensitive group.

Given our results, we can now formally claim that if the within-group reranking happens, it means that all the assumptions listed in Section <ref> are to be considered false.

For some models, this is worrisome: a number of them, in fact, seem to operate under the "Fairness as Optimization" assumption. This means that the reranking of individuals should not be an expected outcome of these models.

In other words, the models seem to be self-contradicting: they operate under the assumption that being as faithful as possible to the data is important, yet they fail to be optimal according to their own assumptions.

Another important discussion regards some of the claims from <cit.>, where they prove that the best model can be found by applying different thresholds on the probability distribution for each sensitive group. Our paper shows how "being best" in this context is equivalent to believing that the biasing process is monotonic. And since the process is monotonic, then indeed the best model can be found by applying thresholds, as a consequence of Weller's Theorem.

This conclusion is tautological; these papers claim that their models are “the best" because they follow their definition of what "best" is.


§ CONCLUSIONS

In this paper, we have shown how assumptions regarding fair models are tied to assumptions about the probabilities of the data. We believe that being transparent about these assumptions is crucial to understanding the behavior of a model.

As others have already claimed, in order to make a fair model, practitioners should start by understanding how bias influences the data.

One of the first questions we need to ask ourselves is whether bias is monotonic or not. This is important, especially when discrimination is possible even among people of the same sensitive group.
Two possible outcomes arise from asking how bias behaves:

    
  * Bias is monotonic. Then, searching for the best model is equivalent to searching for the best thresholds. However, it is also important to note that the fairest model is not necessarily the one that satisfies a fairness constraint. As proven in <cit.>, even when bias is monotonic, depending on the kind of bias, the fairest model can be very unintuitive.
    
  * Bias is not monotonic. If this is the case, we need to understand why it is not so and determine if we can identify the people who are more likely to be discriminated against. If bias is based on different demographic characteristics, it's likely that it affects people who correlate with stereotypes from other groups. This begs the question: should fairness be based on how predictable the sensitive attribute is?

The answers to these questions cannot be universal but need to be weighed based on the problem at hand. Often, educated assumptions must be made in order to build a fair model. Despite what some old proverbs may say, assuming is not wrong, as long as we are explicit about it.
splncs04



 Paper Title: Two Completely Parameter-Free Alternating Gradient Projection Algorithms for Nonconvex-(strongly) Concave Minimax Problems.









[1]This work is supported by National Natural Science Foundation of China under the grants 12071279. 

[2]Corresponding author.

My editor

Two Completely Parameter-Free Alternating Gradient Projection Algorithms for Nonconvex-(strongly) Concave Minimax Problems[1]
    Junnan Yang yangjunnan@shu.edu.cn 

	Department of Mathematics

	Shanghai University

	Shanghai 200444,People's Republic of China
	Huiling Zhang zhanghl1209@shu.edu.cn 

	Department of Mathematics

	Shanghai University

	Shanghai 200444,People's Republic of China
    Zi Xu[2] xuzi@shu.edu.cn 

	Department of Mathematics

	Shanghai University

	Shanghai 200444, People's Republic of China

	and

	Newtouch Center for Mathematics of Shanghai University

	Shanghai University

	Shanghai 200444, People's Republic of China
    August 5, 2024
============================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================



§ ABSTRACT 
Due to their importance in various emerging applications, efficient algorithms for solving minimax problems have recently received increasing attention. However, many existing algorithms require prior knowledge of the problem parameters in order to achieve optimal iteration complexity. In this paper, we propose a completely parameter-free alternating gradient projection (PF-AGP) algorithm to solve the smooth nonconvex-(strongly) concave minimax problems using a backtracking strategy, which does not require prior knowledge of parameters such as the Lipschtiz constant L or the strongly concave constant μ. The PF-AGP algorithm utilizes a parameter-free gradient projection step to alternately update the outer and inner variables in each iteration. We show that the total number of  gradient calls of the PF-AGP algorithm to obtain an ε-stationary point for nonconvex-strongly concave minimax problems is upper bounded by 𝒪( Lκ^3ε^-2) where κ is the condition number, while the total number of gradient calls to obtain an ε-stationary point for nonconvex-concave minimax problems is upper bounded by 𝒪( L^4ε^-4). As far as we know, this is the first completely parameter-free algorithm for solving nonconvex-strongly concave minimax problems, and it is also the completely parameter-free algorithm which achieves the best iteration complexity in single loop method for solving nonconvex-concave minimax problems. Numerical results validate the efficiency of the proposed PF-AGP algorithm.



  minimax optimization problem; backtracking strategies; parameter-free alternating gradient projection algorithm; iteration complexity; machine learning




§ INTRODUCTION



We consider the following minimax optimization problem:

    min_x ∈𝒳max_y ∈𝒴 f(x, y), P

where 𝒳⊆ℝ^n and 𝒴⊆ℝ^m are nonempty closed and bounded convex sets, and f: 𝒳×𝒴→ℝ is a smooth function. The minimax problem is a critical class of optimization problems, known for its extensive research history and broad applications in areas such as machine learning, signal processing, and various other fields in recent years. Many practical problems can be formulated as in (<ref>), including machine learning, such as generative adversarial networks (GANs)  <cit.>, reinforcement learning <cit.>, power control and transceiver design in signal processing <cit.>, distributed nonconvex optimization <cit.>, robust optimization <cit.>, statistical learning <cit.> and several other applications.

There are many research results on convex-concave minimax problems, i.e., f(x, y) is convex with respect to x and concave with respect to y. Nemirovski <cit.> introduced the mirror-prox  method, and Nestrov <cit.> proposed a dual extrapolation algorithm, both of which achieve an optimal iteration complexity of 𝒪(ε^-1) for solving convex-concave minimax problems. For more details, see <cit.>  and references therein.

Recently, various algorithms have been proposed for nonconvex-(strongly) concave minimax problems. For nonconvex-strongly concave minimax problems, several recent works <cit.> have proposed various algorithms that achieve a gradient complexity of 𝒪̃(κ_y^2 ε^-2) for stationary points of Φ(·) = max_y ∈𝒴 f(·, y) (when 𝒳 = ℝ^n), or for stationary points of f, where κ_y is the condition number for f(x, ·). <cit.> proposed an accelerated algorithm that improves the gradient complexity bound to 𝒪̃(√(κ_y)ε^-2).  Furthermore, Zhang et al. <cit.> proposed a generic acceleration framework which can improve the iteration complexity to 𝒪(√(κ_y)ε ^-2).  


For general nonconvex-concave minimax problems, there are two types of algorithms, i.e., multi-loop algorithms and single-loop algorithms. Various multi-loop algorithms have been proposed in <cit.>. The best known iteration complexity of multi-loop algorithms for solving nonconvex-concave minimax problems is 𝒪̃( ε ^-2.5), which was achieved by <cit.>.  For solving nonconvex-linear minimax problems, Pan et al. <cit.> proposed a new alternating gradient projection algorithm with the iteration complexity of 𝒪( ε ^-3). 
Various single-loop algorithms have also been proposed,  e.g., the gradient descent-ascent (GDA) algorithm <cit.>, the hybrid block successive approximation (HiBSA) algorithm <cit.>, the unified single-loop alternating gradient projection (AGP) algorithm <cit.>, and the smoothed GDA algorithm <cit.>. Both the AGP algorithm and the smoothed GDA algorithm achieve the best known iteration complexity, i.e., 𝒪( ε ^-4), among single-loop algorithms for solving nonconvex-concave minimax problems.



It should be noted that for most of the existing results mentioned above, achieving optimal complexity of the algorithm requires the assumption of knowing precise information about some parameters of the problem, such as the Lipschitz constant, or the strongly concave constant μ, etc. Accurately estimating these parameters is often challenging, and conservative estimates can significantly impact algorithm performance <cit.>. Therefore, designing parameter-free algorithms with complexity guarantees without relying on inputs of these parameters has attracted considerable attention recently. 








For nonconvex-strongly concave minimax problems, there are two main parameter-free algorithms: one based on AdaGrad stepsizes and the other based on backtracking. In the framework of AdaGrad stepsizes, Yang et al. <cit.> proposed NeAda, which can achieve convergence rates of 𝒪̃(ε^-2) and 𝒪̃(ε^-4) in deterministic and stochastic settings, respectively.  Li et al. <cit.> proposed TiAda, achieving convergence rates of 𝒪(ε^-2) and 𝒪(ε^-4) in deterministic and stochastic settings, respectively. In the framework of backtracking,  Xu et al. <cit.> proposed a stochastic gradient descent ascent (GDA)  method with backtracking (SGDA-B), achieving an upper bound on the total number of gradient calls of 𝒪(Lκ^2log(κ) ε^-2) with 𝒪(L κ^2 ε^-2) gradient calls required per backtracking step. Meanwhile, Zhang et al. <cit.> proposed a proximal alternating gradient descent ascent method (AGDA+), achieving an upper bound on the total number of gradient calls of  𝒪(Lκ^4ε^-2) with 𝒪(1) gradient calls required per backtracking step. Both methods do not require knowledge of the specific value of Lipschtiz constant L, but still require information about the strongly concave constant μ in the function. In the framework of backtracking, Xu et al. <cit.> proposed a stochastic GDA method with backtracking (SGDA-B) for nonconvex-concave minimax problems, achieving an upper bound on the total number of gradient calls of  𝒪̃(L^3ε^-4) with 𝒪(L^3ε^-4) gradient calls required per backtracking step.





 §.§ Contributions


In this paper, based on the framework of backtracking, we propose two completely parameter-free alternating gradient projection algorithms, i.e., the PF-AGP-NSC algorithm and the PF-AGP-NC algorithm, for solving nonconvex-(strongly) concave minimax problems respectively, which does not require prior knowledge of parameters such as the Lipschtiz constant L or the strongly concave constant μ.  

Moreover, we show that the total number of  gradient calls of the PF-AGP-NSC algorithm to obtain an ε-stationary point for nonconvex-strongly concave minimax problems is upper bounded by 𝒪( Lκ^3ε^-2) where κ is the condition number,  while the total number of gradient calls of the PF-AGP-NC algorithm to obtain an ε-stationary point for nonconvex-concave minimax problems is upper bounded by 𝒪( L^4ε^-4). As far as we know, the PF-AGP-NSC algorithm is the first completely parameter-free algorithm for solving nonconvex-strongly concave minimax problems, and the PF-AGP-NC algorithm is the first completely parameter-free algorithm which achieves the best iteration complexity in single loop algorithms for solving nonconvex-concave minimax problems.







 §.§ Organization

In Sect. <ref>, we propose a parameter-free alternating gradient projection (PF-AGP) algorithm for nonconvex-(strongly) concave minimax problems, and we then analyze the corresponding gradient complexity for two different settings.  We report some numerical results in Sect. <ref> and make some concluding remarks in the last section.


Notation.
For vectors, we use · to denote the l_2-norm. For a function f(x, y): ℝ^n ×ℝ^m →ℝ, we use ∇_x f(x, y) (or ∇_y f(x, y)) to denote the partial gradient of f with respect to the first variable (or the second variable) at point (x, y). Let 𝒫_𝒳 and 𝒫_𝒴 denote projections onto the sets 𝒳 and 𝒴. Finally, we use the notation 𝒪(·) to hide only absolute constants which do not depend on any problem parameter, and 𝒪̃(·) notation to hide only absolute constants and log factors.
A continuously differentiable function f(·) is called θ-strongly convex if there exists a constant θ > 0 such that for any x, y ∈𝒳,

    f(y) ≥ f(x) + ⟨∇ f(x), y - x ⟩ + θ/2y - x^2.

If -f satisfies (<ref>), f(·) is called θ-strongly concave. 
Denote the sign function as (x), i.e.,

    (x) =
     
    1    if  x > 0, 
    
    -1    if  x ≤ 0.







§ TWO COMPLETELY PARAMETER-FREE ALTERNATING GRADIENT PROJECTION ALGORITHMS

In this section, we propose two completely parameter-free alternating gradient projection (PF-AGP) algorithms for solving the nonconvex-strongly concave and the nonconvex-concave minimax problems (<ref>) respectively. Both algorithms do not require prior knowledge of Lipschitz constants, and the algorithm under nonconvex-strong concave settings also does not require prior knowledge of strong concave constants. The two proposed algorithms are based on the AGP algorithm <cit.>, which uses the gradient of a regularized version of the original function, i.e.,

    f_k(x,y) = f(x,y) - c_k/2y_k^2,

where c_k ≥ 0 is a regularization parameters. At the k-th iteration, AGP algorithm consists of the following two gradient projection steps for updating both x and y:

    x_k+1   = 𝒫_𝒳( x_k - 1β_k∇ _xf ( x_k,y_k ) ),
    
    	y_k+1   = 𝒫_𝒴( y_k+1γ_k∇ _y f( x_k,y_k ) - 1/γ_k c_k y_k ),

where  𝒫_𝒳 and 𝒫_𝒴 is the projection operator onto 𝒳 and  𝒴, respectively, and  β_k > 0, γ_k >0 are stepsize parameters. 


Although the AGP algorithm achieves the optimal iteration complexity among single loop algorithms for nonconvex-(strongly) concave minimax problems, it requires the knowledge of the Lipschtiz constant L and the strongly concave coefficient μ, or at least knowledge of the Lipschtiz constant L  to calculate the stepsizes β_k and γ_k under the nonconvex-strongly concave setting and the nonconvex-concave setting, respectively, which limits its applicability. To eliminate the need for the prior knowledge of the Lipschtiz constant L or the strongly concave coefficient μ, we propose two parameter-free AGP algorithms by employing backtracking strategies for solving the nonconvex-strongly concave and the nonconvex-concave minimax problems, respectively. Similar backtracking strategies have been used in
<cit.>. 




 §.§ Nonconvex-strongly Concave Setting

We first propose a completely parameter-free AGP algorithm for solving nonconvex-strongly concave minimax problems.
More specifically, in order to estimate the Lipschitz constants for ∇_x f(·,y), ∇_y f(·,y), ∇_y f(x,·) denoted as L_11, L_12 and L_22, respectively, and the strongly concave constant μ, at each iteration of the proposed algorithm, we aim to find a tuple (l_11^k,i, l_12^k,i, l_22^k,i, μ_k,i) by backtracking such that the following conditions for (x_k,i, y_k,i) are satisfied:

    C_1^k,i   = f(x_k,i, y_k) - f(x_k, y_k) - ⟨∇_x f(x_k, y_k), x_k,i - x_k ⟩ - l^k,i_11/2x_k,i - x_k^2 ≤ 0, C1
    
    	C_2^k,i   = ∇_y f(x_k,i, y_k) - ∇_y f(x_k, y_k) - l_12^k,ix_k,i - x_k≤ 0,C2
    
    	C_3^k,i   = ∇_y f(x_k,i, y_k,i) - ∇_y f(x_k,i, y_k) - l_22^k,iy_k,i - y_k≤ 0,C3
    
    	C_4^k,i   = ⟨∇_y f(x_k,i, y_k,i) - ∇_y f(x_k,i, y_k), y_k,i - y_k ⟩ + μ_k,iy_k,i - y_k^2 ≤ 0.C4

Otherwise, we enlarge l_11^k,i, l_12^k,i, l_22^k,i by a factor of 2 or reduce μ_k,i by half according to the sign of C_1^k,i, C_2^k,i, C_3^k,i and C_4^k,i, respectively. We then update β_k and γ_k by the new estimates, i.e., l_11^k,i, l_12^k,i, l_22^k,i, μ_k,i.
The proposed algorithm for solving nonconvex-strongly concave problems, denoted as PF-AGP-NSC,  is formally presented in Algorithm <ref>.



It is worth noting that the parameter-free NeAda <cit.> and TiAda <cit.> algorithms with adaptive step size are only used to solve the minimax problems when x is unconstrained, while the  proposed PF-AGP-NSC algorithm can solve the constrained case. Compared with the SGDA-B <cit.> algorithm and AGDA+ <cit.> algorithm using the backtracking framework under the nonconvex-strongly concave setting, AGDA+ and SGDA-B still require knowledge of the strongly concave constant mu, while the proposed PF-AGP-NSC algorithm is a completely parameter-free algorithm that does not require any prior knowledge of the parameters.





  §.§.§ Complexity analysis

In this subsection, we analyze the iteration complexity of Algorithm <ref> for solving nonconvex-strongly concave minimax optimization problems, i.e., f(x, y) is nonconvex with respect to x for any fixed y ∈𝒴, and μ-strongly concave with respect to y for any given x ∈𝒳. 

Before we prove the iteration complexity of the PF-AGP-NSC algorithm for solving (<ref>), we need to make the following assumption about the smoothness of f(x, y).

f(x, y) has Lipschitz continuous gradients, i.e., there exist positive scalars L_11, L_12, L_22 such that for any x, x̃∈𝒳, y, ỹ∈𝒴,

    ∇_x f(x, y) - ∇_x f(x̃, y)    ≤ L_11 x - x̃, 
    ∇_y f(x, y) - ∇_y f(x̃, y)    ≤ L_12 x - x̃,
    ∇_y f(x, ỹ) - ∇_y f(x, y)    ≤ L_22ỹ - y .
 

We denote  L = max{L_11, L_12, L_22}. To analyze the convergence of Algorithm <ref>, we define the stationarity gap as the termination criterion as follows.

At each iteration of Algorithm <ref>, the stationarity gap for problem (<ref>) with respect to f(x, y) is defined as:

    ∇ G(x_k, y_k) = [ β_k (x_k - 𝒫_𝒳(x_k - 1/β_k∇_x f(x_k, y_k))); γ_k (y_k - 𝒫_𝒴(y_k + 1/γ_k∇_y f(x_k, y_k))) ].

We denote ∇G_k= ∇G(x_k, y_k),
(∇ G_k)_x= β_k( x_k-𝒫_𝒳( x_k-1/β_k∇ _xf( x_k,y_k))),
and (∇ G_k)_y=γ_k( y_k-𝒫_𝒴( y_k+ 1/γ_k∇ _yf( x_k,y_k )) ).








At each iteration of Algorithm <ref>, the stationarity gap for problem (<ref>) with respect to f_k(x, y) is defined as:

    ∇G̃(x_k, y_k) = 
    [  β_k(x_k - P_x(x_k - 1/β_k∇_x f_k(x_k, y_k))); γ_k (y_k - P_y(y_k + 1/γ_k∇_y f_k(x_k, y_k))) ].
 

Under Assumption <ref>, we first establish the following lemma, which provides bounds on the changes in the function value when x_k is updated at each iteration of Algorithm <ref>. This proof is similar to that of Lemma 2.1 in <cit.>, for the sake of completeness, we give its proof.



 Suppose that Assumption <ref> holds. 
Let {(x_k, y_k)} be a sequence generated by Algorithm <ref>, then we have

    f(x_k+1, y_k) - f(x_k, y_k) ≤ - (β_k - l^k_11/2) x_k+1 - x_k^2 .


	
	

   



  By the optimality condition for (<ref>), we have
		
    ⟨∇ _xf( x_k,y_k ) +β_k ( x_k+1-x_k ) ,x_k-x_k+1⟩≥ 0.

		The backtracking strategy (<ref>) implies that
		
    f( x_k+1,y_k ) -f( x_k,y_k )≤⟨∇ _xf( x_k,y_k ) ,x_k+1-x_k ⟩  +l^k_112 x_k+1-x_k  ^2.

		By adding (<ref>) and (<ref>), we obtain
		
    f( x_k+1,y_k ) -f( x_k,y_k )    ≤ -( β_k -l^k_112)  x_k+1-x_k  ^2,

which completes the proof. 











Next, we provide an estimate of upper bound of the difference between f(x_k+1, y_k+1) and f(x_k, y_k). 

First, we need to make the following assumption on the parameter γ_k.

{γ_k} is a nonnegative monotonically increasing sequence. 


Suppose that Assumptions <ref> and <ref> hold. 
Let {(x_k, y_k)} be a sequence generated by Algorithm <ref>, then we have

    f(x_k+1, y_k+1) - f(x_k, y_k) 
         ≤ -(β_k - l_11^k/2 - l_12^k/2γ_k-1)x_k+1 - x_k^2+γ_k-1y_k+1 - y_k^2
            -(μ_k-1 - γ_k-1/2 - (l_22^k-1)^2/2γ_k-1)y_k - y_k-1^2  .
 
 


The optimality condition for y_k in (<ref>) implies that





    ⟨∇_y f(x_k, y_k-1) - γ_k-1 (y_k - y_k-1), y_k+1 - y_k ⟩≤ 0.

By the concavity of f(x,y) with respect to y, and combining (<ref>), we have 

    f(x_k+1, y_k+1) - f(x_k+1, y_k) 
    ≤   ⟨∇_y f(x_k+1, y_k), y_k+1 - y_k ⟩
    ≤   ⟨∇_y f(x_k+1, y_k) - ∇_y f(x_k, y_k-1), y_k+1 - y_k ⟩+ γ_k-1⟨ y_k-y_k-1, y_k+1 - y_k ⟩.

Denoting v_k+1 = (y_k+1 - y_k) - (y_k - y_k-1), we can write the first inner product term in the r.h.s. of (<ref>) as

    ⟨∇_y f(x_k+1, y_k) - ∇_y f(x_k, y_k-1), y_k+1 - y_k ⟩
    
    		=   ⟨∇_y f(x_k+1, y_k) - ∇_y f(x_k, y_k), y_k+1 - y_k ⟩ + ⟨∇_y f(x_k, y_k) - ∇_y f(x_k, y_k-1), v_k+1⟩
        + ⟨∇_y f(x_k, y_k) - ∇_y f(x_k, y_k-1), y_k - y_k-1⟩.

Next, we estimate the three terms on the right-hand side of (<ref>) respectively. By the backtracking strategy (<ref>), (<ref>) and the Cauchy-Schwarz inequality, we have

    ⟨∇_y f(x_k+1, y_k) - ∇_y f(x_k, y_k), y_k+1 - y_k ⟩≤l_12^k/2γ_kx_k+1 - x_k^2 + γ_k/2y_k+1 - y_k^2

and

    ⟨∇_y f(x_k, y_k) - ∇_y f(x_k, y_k-1), v_k+1⟩≤(l_22^k-1)^2/2 γ_k-1y_k - y_k-1^2 + γ_k-1/2v_k+1^2 .

By the backtracking strategy (<ref>) of μ_k-1, we obtain

    ⟨∇_y f(x_k, y_k) - ∇_y f(x_k, y_k-1), y_k - y_k-1⟩≤ -μ_k-1y_k - y_k-1^2.

Moreover, it can be easily checked that

    ⟨ y_k - y_k-1, y_k+1 - y_k ⟩ = 1/2y_k - y_k-1^2 + 1/2y_k+1 - y_k^2 - 1/2v_k+1^2.

Plugging (<ref>)-(<ref>) into (<ref>) and rearranging the terms, we conclude that

    f(x_k+1, y_k+1) - f(x_k+1, y_k)    ≤l_12^k/2γ_kx_k+1 - x_k^2 - (μ_k-1 - γ_k-1/2 - (l_22^k-1)^2/2γ_k-1) y_k - y_k-1^2 
          + (  γ_k/2 + γ_k-1/2) y_k+1 - y_k^2.

By Lemma <ref>, we have

    f(x_k+1, y_k) - f(x_k, y_k) ≤ - (β_k - l_11^k/2) x_k+1 - x_k^2.

The proof is then completed by combining (<ref>) with (<ref>).



We now establish an important recursion for Algorithm <ref>.

Suppose that Assumptions <ref> and <ref> hold. 
Let {(x_k, y_k)} be a sequence generated by Algorithm <ref>, and denoting

    f_k+1    = f(x_k+1, y_k+1),    S_k+1 = 2γ^2_k /μ_ky_k+1 - y_k^2, 
    
      F_k+1    = f_k+1 + S_k+1 - ( μ_k + 7γ_k/2 -  (l^k_22)^2/2γ_k - 2(l^k_22)^2/μ_k) y_k+1 - y_k^2-2γ_k^2/μ_k𝒟^2_y,

then we have

    F_k+1 - F_k
    	≤    -( 2μ_k γ_k - (l_22^k)^2/2γ_k + μ_k γ_k - 4(l_22^k)^2/2μ_k) y_k+1-y_k^2
          - ( β_k - l_11^k/2 - l_12^k/2γ_k-1- 2(l_12^k)^2γ_k-1/μ_k-1μ_k) x_k+1 - x_k^2,

where 𝒟_y = max{y_1 - y_2| y_1, y_2 ∈𝒴}. 



By the optimality condition for y_k in (<ref>) we obtain

    ⟨∇_y f(x_k+1, y_k) - γ_k (y_k+1 - y_k), y_k - y_k+1⟩≤ 0.

Combining (<ref>) and (<ref>), we have

    ⟨∇_y f(x_k, y_k-1) - ∇_y f(x_k+1, y_k)+γ_k( y_k+1 - y_k)-γ_k-1( y_k - y_k-1), y_k+1 - y_k⟩≤ 0,

which implies that

    γ_k y_k+1 - y_k^2 - γ_k-1⟨ y_k-y_k-1, y_k+1 - y_k ⟩≤⟨∇_y f(x_k+1, y_k) - ∇_y f(x_k, y_k-1), y_k+1 - y_k ⟩ .

Denoting v_k+1 = (y_k+1-y_k)-(y_k-y_k-1), and using the identity 1/2⟨ v_k+1, y_k+1 - y_k⟩ = 1/2y_k+1 - y_k^2 + 1/2v_k+1^2 - 1/2y_k - y_k-1^2, we conclude that

    γ_k y_k+1 - y_k^2 - γ_k-1( 1/2y_k+1 - y_k^2 + 1/2y_k - y_k-1^2 - 1/2v_k+1^2 )
         ≤⟨∇_y f(x_k+1, y_k) - ∇_y f(x_k, y_k-1), y_k+1 - y_k ⟩.

By rearranging the terms of the above inequality, we then have

    (γ_k - γ_k-1/2) y_k+1 - y_k^2 + γ_k-1/2v_k+1^2 - γ_k-1/2y_k - y_k-1^2 
    ≤   ⟨∇_y f(x_k+1, y_k) - ∇_y f(x_k, y_k-1), y_k+1 - y_k ⟩
    
        =   ⟨∇_y f(x_k+1, y_k) - ∇_y f(x_k, y_k), y_k+1 - y_k ⟩+⟨∇_y f(x_k, y_k) - ∇_y f(x_k, y_k-1), v_k+1⟩
        +⟨∇_y f(x_k, y_k) - ∇_y f(x_k, y_k-1), y_k - y_k-1⟩.

 By backtracking strategy (<ref>) and the Cauchy-Schwarz inequality, we can easily get that

    ⟨∇_y f(x_k+1, y_k) - ∇_y f(x_k, y_k), y_k+1 - y_k ⟩≤(l_12^k)^2/2μ_kx_k+1 - x_k^2 + μ_k/2y_k+1 - y_k^2.

Combining (<ref>), (<ref>), (<ref>) and (<ref>), we have

    γ_k/2y_k+1 - y_k^2 - γ_k-1/2y_k - y_k-1^2 
    	   ≤(l_12^k)^2/2μ_kx_k+1 - x_k^2 + (μ_k/2 + γ_k-1/2 - γ_k/2) y_k+1 - y_k^2 
          - (μ_k-1 - (l_22^k-1)^2/2γ_k-1) y_k - y_k-1^2.

By multiplying 4γ_k-1/μ_k-1 on both sides of (<ref>), using the definition of S_k+1 and Assumption <ref>, we obtain








    S_k+1 - S_k 
    		   ≤2(l_12^k)^2γ_k-1/μ_k-1μ_kx_k+1 - x_k^2 + 2γ_k-1y_k+1 - y_k^2  
          - (4γ_k-1 - 2(l_22^k-1)^2/μ_k-1) y_k - y_k-1^2  + (2γ_k^2/μ_k - 2γ^2_k-1/μ_k-1) 𝒟_y^2.

The proof is then completed by Lemma <ref> and the definition.


We are now ready to establish the iteration complexity for the PF-AGP-NSC  algorithm. In particular, letting ∇ G(x_k, y_k) be defined as in Definition <ref> and ε > 0 be a given target accuracy, we provide a bound on T(ε), the first iteration index to achieve an ε-stationary point, i.e., ∇ G(x_k, y_k)≤ε, which is equivalent to

    T(ε) = min{k |∇ G(x_k, y_k)≤ε}.

 
Suppose that Assumptions <ref> and <ref> hold. 
Let {(x_k, y_k)} be a sequence generated by Algorithm <ref>, β_k  =l_11^k + l_12^kμ_k-1/8 (l_22^k-1)^2  +  32(l_12^k)^2(l_22^k-1)^2/μ_kμ_k-1^2, γ_k = 8 (l_22^k)^2/μ_k,








then it holds that

    T(ε) ≤F_2 -F/d_1 ε^2 +1,

where  F_2 is defined in (<ref>), 
F = f - ( 224L_22^2/μ + μ_1 - 2 (l_22^1)^2/μ_1 - (l_22^1μ)^2/128L_22^2+2^13L_22^4/μ^3) 𝒟^2_y,








with f = min_(x, y) ∈ X × Y f(x, y),   𝒟_y = max{y_1 - y_2| y_1, y_2 ∈𝒴}, d_1 = min{β_min/2β_max^2+4L^2_12,15μ^3+2^7l_22^1μ/2^15L_22^4},  β_min=l^1_11 + l^1_12μ/64 (L_22)^2 + 32 (l_12^1)^2 (l^1_22)^2/μ_1^3 and  β_max = 2L_11+ μ_1 L_12/4(l^1_22)^2 +  4096(L_12)^2 (L_22)^2/μ^3.





 It follows immediately from Definition <ref> and (<ref>) that

    (∇ G_k)_x  = β_k  x_k+1 - x_k .

On the other hand, by  the triangle inequality, we obtain that

    (∇ G_k)_y    ≤γ_k  P_y (y_k + 1/γ_k∇_y f(x_k+1, y_k)) - P_y (y_k + 1/γ_k∇_y f(x_k, y_k)) 
          + γ_k  y_k+1 - y_k 
       ≤ L_12 x_k+1 - x_k  + γ_k  y_k+1 - y_k .

By combining (<ref>) and (<ref>), and using the Cauchy-Schwarz inequality, we obtain

    ∇ G_k ^2    ≤ (β_k^2 + 2 L_12^2)  x_k+1 - x_k ^2 + 2 γ_k^2  y_k+1 - y_k ^2.

Let d^(k)_1= min{2μ_kγ_k - (l^k_22)^2/2γ_k + μ_kγ_k-4 (l_22^k)^2/2μ_k/2 γ_k^2, β_k - l^k_11/2-  l^k_12/2γ_k-1 - 2 (l^k_12)^2 γ_k-1/μ_k-1μ_k/β_k^2 + 2 L_12^2}>0. Multiplying both sides of (<ref>) by d^(k)_1, and using (<ref>) in Lemma <ref>, we have

    d^(k)_1 ∇ G_k ^2    ≤ F_k - F_k+1.

By using the definitions of d_1, β_k and γ_k, we have
d_1<d^(k)_1, then we get

    d_1 ∇ G_k ^2 
    	≤ F_k - F_k+1.

Summing up the above inequalities from k = 2 to k = T(ε), we obtain

    ∑_k=2^T(ε) d_1 ∇ G_k^2    ≤ F_2 - F_T(ε)+1  .

Note that by the definition of F_k+1 in Lemma <ref>, we have

    F_T(ε)+1   = f_T(ε)+1 + S_T(ε)+1 -2γ_T(ε)^2/μ_T(ε)𝒟^2_y
          - ( 7/2γ_T(ε)+1 + μ_T(ε)+1 - 2 (l_22^T(ε)+1)^2/μ_T(ε)+1 - (l_22^T(ε)+1)^2/2γ_T(ε)+1) ||y_T(ε)+1 - y_T(ε)||^2 
       ≥f - ( 224L_22^2/μ + μ_1 - 2 (l_22^1)^2/μ_1 - (l_22^1μ)^2/128L_22^2+2^13L_22^4/μ^3) 𝒟^2_y 
       =F,

where the inequality follows from the definitions of f and 𝒟_y, and the fact that S_k ≥ 0 (∀ k ≥ 1). 
We then conclude from (<ref>) that

    ∑_k=2^T(ε) d_1 ∇ G_k^2    ≤F_2 - F_T(ε)+1≤F_2 - F  ,

which, in view of the definition of T(ε), implies that

    d_1 (T(ε)-1) ε^2 ≤ F_2 - F ,





or equivalently,

    T(ε) ≤F_2 -  F/d_1 ε^2+1.

We complete the proof.

 
At the k-th iteration of Algorithm <ref>, the number of backtracking step is upper bounded by log_2(2L_11/l_11^1)+log_2(2L_12/l_12^1)+log_2(2L_22/l_22^1)+log_2(2μ_1/μ)=𝒪(1). 
By Theorem <ref>, the total number of  gradient calls of Algorithm <ref> to obtain an ε-stationary point  that satisfies ∇ G(x_k,y_k)≤ε is upper bounded by 𝒪 ( Lκ^3ε^-2) under the nonconvex-strongly concave setting, where κ is the condition number.






 §.§ Nonconvex-concave Setting

In this subsection, we propose a completely parameter-free AGP algorithm for solving nonconvex-concave minimax problems. Similar to the idea of the PF-AGP-NSC Algorithm, we estimate the Lipschitz constants L_11, L_12 and L_22 iteratively, however we do not need to estimate the strongly concave constant μ. At each iteration of the proposed algorithm, we aim to find a tuple (l^k,i_11,l^k,i_12,l^k,i_22) by backtracking such that 
conditions (<ref>), (<ref>), (<ref>) and the following condition (<ref>) are satisfied at (x_k,i,y_k,i):

    C_5^k,i   = ⟨∇_y f_k,i(x_k,i, y_k,i) - ∇_y f_k,i(x_k,i, y_k), y_k,i - y_k⟩+ c^2_k,i+c_k,il^k,i_22/l^k,i_22 + 2c_k,iy_k,i - y_k^2
            +1/l^k,i_22 + 2c_k,i∇_y f_k,i(x_k,i, y_k,i) - ∇_y f_k,i(x_k,i, y_k)^2  ≤ 0. C5

The proposed algorithm, denoted as PF-AGP-NC, is formally presented in Algorithm <ref>.

Please note that the existing SGDA-B algorithm <cit.> for solving nonconvex-concave minimax problems requires the realization of the diameter of the feasible set 𝒴 in order to construct the key auxiliary problem, while the PF-AGP-NC algorithm we propose is a completely parameter-free algorithm that does not require any prior parameters.

 



  §.§.§ Complexity Analysis

In this subsection, we analyze the iteration complexity of Algorithm <ref> for solving (<ref>) under the general nonconvex-concave setting. 






The backtracking strategy (<ref>) is a key inequality that we will use to establish some important recursions for the PF-AGP-NC algorithm under the nonconvex-concave setting in the following two lemmas. This is also one of the main differences between proofs in the nonconvex- strongly concave setting and proofs in the nonconvex-concave setting.




Suppose that Assumptions <ref> and <ref> hold. Let {(x_k, y_k)} be a sequence generated by Algorithm <ref>. If ∀ k, β_k > l^k_11 and γ_k≥l_22^k + 2c_k/2, then

    f(x_k+1,y_k+1) - f(x_k,y_k) 
    ≤    - (   β_k/2 - (l_12^k)^2/2γ_k-1) x_k+1 - x_k^2 +γ_k-1y_k+1- y_k^2 + γ_k-1/2y_k - y_k-1^2
          + c_k-1/2( y_k+1^2 - y_k^2 )+ c_k/2y_k+1-y_k^2.






	
	


The optimality condition for y_k in (<ref>) implies that




    ⟨∇_y f_k-1 (x_k, y_k) - γ_k-1 (y_k - y_k-1), y_k - y_k+1⟩≤ 0.





	
	
	
	
	

The concavity of f_k(x_k+1, y) with respect to y together with (<ref>) then imply that

    f_k (x_k+1, y_k+1) - f_k (x_k+1, y_k)    ≤⟨∇_y f_k (x_k+1, y_k)- ∇_y f_k-1 (x_k, y_k), y_k+1-y_k ⟩
         + ⟨∇_y f_k-1 (x_k, y_k-1), y_k+1 - y_k ⟩
       = ⟨∇_y f_k (x_k+1, y_k) - ∇_y f_k-1 (x_k, y_k), y_k+1 - y_k ⟩
         + ⟨∇_y f_k-1 (x_k, y_k) - ∇_y f_k-1 (x_k, y_k-1), v_k+1⟩
         + ⟨∇_y f_k-1 (x_k, y_k) - ∇_y f_k-1 (x_k, y_k-1), y_k - y_k-1⟩
         + γ_k-1⟨ y_k - y_k-1, y_k+1 - y_k ⟩,

where v_k+1 = y_k+1 - y_k - (y_k - y_k-1). We now provide upper bounds for the inner product terms on the right hand side of (<ref>). Firstly, by the definition of f_k(x_k+1, y_k) and f_k-1(x_k, y_k), the backtracking strategy (<ref>)  and the Cauchy-Schwarz inequality, we have

    ⟨∇_y f_k (x_k+1, y_k) - ∇_y f_k-1 (x_k, y_k), y_k+1 - y_k ⟩
       = ⟨∇_y f(x_k+1, y_k) - ∇_y f (x_k, y_k), y_k+1 - y_k ⟩  - (c_k - c_k-1)⟨ y_k, y_k+1 - y_k ⟩
       ≤(l_12^k)^2/2γ_k-1x_k+1 - x_k^2 + γ_k-1/2y_k+1 - y_k^2 - c_k - c_k-1/2 (y_k+1^2 - y_k^2)
          + c_k - c_k-1/2y_k+1-y_k^2
       ≤(l_12^k)^2/2γ_k-1x_k+1 - x_k^2 + γ_k-1/2y_k+1 - y_k^2 - c_k - c_k-1/2 (y_k+1^2 - y_k^2) 
          + c_k/2y_k+1-y_k^2.

Secondly, by the Cauchy-Schwarz inequality, we get

    ⟨∇_y f_k-1 (x_k, y_k) - ∇_y f_k-1 (x_k, y_k-1), v_k+1⟩
    ≤   1/2 γ_k-1∇_y f_k-1 (x_k, y_k) - ∇_y f_k-1 (x_k, y_k-1) ^2 + γ_k-1/2 v_k+1^2.

Thirdly, by (<ref>) we have that

    ⟨∇_y f_k-1(x_k, y_k) - ∇_y f_k-1(x_k, y_k-1), y_k - y_k-1⟩
    ≤    -1/l_22^k-1 + 2c_k-1∇_y f_k-1(x_k, y_k) - ∇_y f_k-1(x_k, y_k-1) ^2
         - c_k-1l_22^k-1+c^2_k-1/l_22^k-1 + 2c_k-1 y_k - y_k-1^2
    ≤    - 1/l_22^k-1 + 2c_k-1∇_y f_k(x_k, y_k) - ∇_y f_k(x_k, y_k-1) ^2.

We also have the following equation that

    γ_k-1⟨ y_k - y_k-1, y_k+1 - y_k ⟩ = γ_k-1/2 y_k+1 - y_k ^2 + γ_k-1/2 y_k - y_k-1^2 - γ_k-1/2 v_k+1^2.

By γ_k≥l_22^k + 2c_k/2, plugging (<ref>)-(<ref>) into (<ref>), 
 and using the definition of f_k(x_k+1, y_k+1) and f_k(x_k+1, y_k),   we obtain

    f(x_k+1, y_k+1) - f(x_k+1, y_k) 
    ≤   (l_12^k)^2/2 γ_k-1 x_k+1 - x_k ^2 + γ_k-1 y_k+1 - y_k ^2 + γ_k-1/2 y_k - y_k-1^2
          + c_k-1/2( y_k+1^2 - y_k^2 )+ c_k/2y_k+1-y_k^2  .

By  β_k > l^k_11, we have

    f(x_k+1, y_k) - f(x_k, y_k) ≤ -β_k /2 x_k+1 - x_k ^2 .

By adding (<ref>) and (<ref>), we complete the proof.


Suppose that Assumptions <ref> and <ref> hold. Let {(x_k, y_k)} be a sequence generated by Algorithm <ref>. Denote

    S_k+1   = 8γ_k^2/c_k y_k+1 - y_k ^2 + 8γ_k( c_k+1/c_k-1 )  y_k+1^2,
    
    	F_k+1   = f(x_k+1, y_k+1) + S_k+1 - 15γ_k/2y_k+1 - y_k^2 - c_k/2y_k+1^2.

If 

    β_k > l^k_11,   γ_k+1/c_k+1 - γ_k/c_k≤1/5,   γ_k ≥l^k_22 + 2c_k/2 ,

then ∀ k ≥ 2,

    F_k+1 - F_k =    -( 9γ_k/10-c_k/2)y_k+1 - y_k^2 
       -(   + β_k/2 - (l^k_12)^2/2γ_k-1-16(l^k_12)^2γ_k-1/c_k-1^2)x_k+1 - x_k^2 
        +8( γ_kc_k+1/c_k - γ_k-1c_k/c_k-1) y_k+1^2+c_k-1-c_k/2y_k+1^2.



The optimality condition for y_k in (<ref>) implies that

    ⟨∇_y f_k(x_k+1, y_k)-γ_k(y_k+1-y_k), y_k - y_k+1⟩≤ 0,

By (<ref>) and (<ref>), we have  

    ⟨∇_y f_k-1(x_k, y_k-1) - ∇_y f_k(x_k+1, y_k) +   γ_k( y_k+1 - y_k)  - γ_k-1 (y_k - y_k-1), y_k+1 - y_k⟩≤ 0,

and

    γ_k y_k+1 - y_k^2 - γ_k-1⟨ y_k+1 - y_k, y_k - y_k-1⟩
    ≤   ⟨∇_y f_k(x_k+1, y_k) - ∇_y f_k-1(x_k, y_k-1), y_k+1 - y_k⟩,

 which can be rewritten as

    (γ_k - γ_k-1/2) y_k+1 - y_k^2 + γ_k-1/2v_k+1^2 - γ_k-1/2y_k - y_k-1^2
    ≤   ⟨∇_y f_k(x_k+1, y_k) - ∇_y f_k-1(x_k, y_k-1), y_k+1 - y_k ⟩
    
    =   ⟨∇_y f_k(x_k+1, y_k) - ∇_y f_k-1(x_k, y_k), y_k+1 - y_k ⟩+ ⟨∇_y f_k-1(x_k, y_k) - ∇_y f_k-1(x_k, y_k-1), v_k+1⟩
    
    	   + ⟨∇_y f_k(x_k-1, y_k-1) - ∇_y f_k-1(x_k, y_k-1), y_k - y_k-1⟩.

Using an argument similar to the proof of (<ref>)-(<ref>), by (<ref>) and the Cauchy-Schwarz inequality, we conclude from the above inequality that

    (γ_k - γ_k-1/2) y_k+1 - y_k^2 + γ_k-1/2v_k+1^2 - γ_k-1/2y_k - y_k-1^2
    ≤    (l_12^k)^2/2a_kx_k+1 - x_k^2 + a_k/2y_k+1 - y_k^2 - c_k - c_k-1/2 (y_k+1^2 - y_k^2)
        + c_k - c_k-1/2y_k+1 - y_k^2 + 1/2γ_k-1∇_y f_k-1(x_k, y_k) - ∇_y f_k-1(x_k, y_k-1)^2  
        +γ_k-1/2v_k+1- 1/l_22^k-1 + 2c_k-1∇_y f_k(x_k, y_k) - ∇_y f_k(x_k, y_k-1)^2  
         - c_k-1l_22^k-1+c^2_k-1/l_22^k-1  + 2c_k-1y_k - y_k-1^2,

for any a_k > 0. Obviously, we have -c_k-1l^k-1_22+c^2_k-1/2c_k-1 + l^k-1_22≤ -c_k-1l^k-1_22+c^2_k-1/2(l^k-1_22+c_k-1) = -c_k-1/2.
Combining γ_k ≥l^k_22 + 2c_k/2 and rearranging the terms in (<ref>), we obtain

    γ_k/2y_k+1 - y_k^2 + c_k - c_k-1/2y_k+1^2 
         ≤γ_k-1/2y_k - y_k-1^2 + ( α_k/2 + γ_k-1/2 - γ_k/2) y_k+1 - y_k^2  + c_k - c_k-1/2y_k^2 
           +  (l^k_12)^2 /2a_kx_k+1 - x_k^2 - c_k-1/2y_k - y_k-1^2.

By multiplying 16γ_k-1/c_k-1 on both sides of the above inequality, setting a_k = c_k/2, and using the definition of S_k+1, (<ref>) and {γ_k} is nodecreasing, we then obtain

    S_k+1 - S_k    ≤ 8( γ_kc_k+1/c_k - γ_k-1c_k/c_k-1) y_k+1^2 + 16 (l^k_12)^2 γ_k-1/c_k-1^2x_k+1 - x_k^2 
         + 28γ_k/5y_k+1 - y_k^2 - 8γ_k-1y_k - y_k-1^2.




	
	





	
	
	
	
	
	

Combining  (<ref>) and (<ref>), and using the definition of F_k+1,  we complete the proof.


We are now ready to establish the iteration complexity for the PF-AGP-NC algorithm to achieve an ε-stationary point under the general nonconvex–concave setting.

Suppose that Assumptions <ref> and <ref> hold. Let {(x_k, y_k)} be a sequence generated by Algorithm <ref>. If β_k = l_12^k/20l^k-1_22 + 1280(l_12^k)^2 √(k)/19^2l_12^k-1, γ_k =20l^k_22 ,  c_k = 19l^k_22/k^1/4,  then for any given ε > 0, k ≥ 2, we have

    T(ε) ≤max(  ( 5120 ( L_12)^2  d_2d_3/19^2l_22^1 ε^2+1)^2,  76^4L_22^4𝒟̂_y^4/ε^4),

where  d_2= F_2 - F +339L_22𝒟̂_y^2, d_3 = max{d̅_1, 20·19^2L_22^2/17(l_12^1)^2}, 𝒟̂_y = max{y  |  y ∈ Y} with F = f - 639L_22𝒟̂_y^2, f = min_(x,y) ∈ X × Y f(x, y), 
  d̅_1 =  32+ 1+1600L^2_22/128(l_12^1)^2.
 
















We can easily see that the relations in (<ref>) are satisfied by the settings of c_k, γ_k and β_k. Denote α_k = 16 (l_12^k)^2 γ_k-1/c_k-1^2. Then it follows from the selection of  β_k and α_k that

    β_k/2 - (l^k_12)^2/2γ_k-1-16(l^k_12)^2γ_k-1/c_k-1^2 = α_k .

This observation, in view of Lemma <ref>, c_k/2≤19γ_k/40k^1/4≤19γ_k/40 implies that

    α_k x_k+1 - x_k ^2 + 17γ_k/40 y_k+1 - y_k ^2 
    ≤    F_k -F_k+1 
    		+8( γ_kc_k+1/c_k - γ_k-1c_k/c_k-1) y_k+1^2 +c_k-1-c_k/2y_k+1^2.

We can easily check from the definition of f_k(x, y) that

    ∇ G_k  - ∇̃ G_k ≤ c_k  y_k .

By replacing f with f_k, similar to (<ref>) and (<ref>), we immediately obtain that

    (∇̃ G_k)_x  =   β_k     x_k+1 - x_k ,

and

    (∇̃ G_k)_y ≤γ_k  y_k+1 - y_k  + l^k_12 x_k+1 - x_k .

Combining (<ref>) and (<ref>), and using the Cauchy-Schwarz inequality, we have

    ∇̃ G_k ^2 ≤ (  β_k ^2 + 2(l^k_12)^2)  x_k+1 - x_k ^2 + 2γ_k^2 y_k+1 - y_k ^2.

 Since both α_k and β_k are in the same order when k becomes large enough, it then follows from the definition of d̅_1 that ∀ k ≥ 1,

    (  β_k )^2 + 2(l^k_12)^2/α_k^2   = 
    	(l^k_12/γ_k-1 + 64 (l_12^k)^2γ_k-1/c_k-1^2)^2 + 2(l^k_12)^2/α_k^2
    
    	
    	
    	
     	   ≤  32+ 1+1600L^2_22/128(l_12^1)^2 = d̅_1.

Combining the previous two inequalities in (<ref>) and (<ref>), we obtain

    ∇G̃_k^2 ≤d̅_1(α_k)^2 x_k+1 - x_k ^2 + 2 γ_k^2  y_k+1 - y_k ^2 .

Denote d_k^(2) = 1/max{d̅_1 α_k, 17/80γ_k}, c̃_k = 19/20 k^1/4, by multiplying d^(2)_k on the both sides of (<ref>), and using (<ref>), we have

    d_k^(2)∇G̃_k ^2 
    ≤    F_k - F_k+1  +8( γ_kc_k+1/c_k - γ_k-1c_k/c_k-1) y_k+1^2 
    	+c_k-1-c_k/2y_k+1^2
    ≤    F_k - F_k+1  +8( γ_kc_k+1/c_k - γ_k-1c_k/c_k-1) y_k+1^2 
    	+40L_22c̃_k-1-c̃_k/2y_k+1^2.

Denoting

    T̃(ε) = min{ k | ∇G̃(x_k, y_k) ≤ε/2, k ≥ 2 },

summing both sides of (<ref>) from k = 2 to k = T̃(ε), we then obtain

    ∑_k=2^T̃(ε) d^(2)_k ∇G̃_k ^2 
    	   ≤  F_2 - F_T̃(ε)  + 8 (  γ_T̃(ε)-1c_T̃(ε)/c_T̃(ε)-1-γ_1c_2/ c_1) 𝒟̂_y^2 + 40L_22( c̃_1/2 - c̃_T̃(ε)/2) 𝒟̂^2_y 
       ≤ F_2 -  F_T̃(ε) + 339L_22𝒟̂_y^2.

Note that by the definition of F_k+1 in Lemma <ref>, we have

    F_T̃(ε)   ≥f -  639L_22𝒟̂_y^2 = F,

where f = min_(x, y) ∈ X × Y f(x, y). We then conclude from (<ref>) that

    ∑_k=2^∞ d^(2)_k ∇G̃_k ^2 ≤ F_2 - F +339L_22𝒟̂_y^2= d_2.

We can see from the selection of  d_3 that d_3 ≥max{d̅_1, 80γ_k/17α_k}, which implies d^(2)_k ≥1/d_3α_k,  by multiplying d_3  on the both sides of (<ref>), and combining the definition of  d_2,  we have  ∑_k=2^∞1/α_k∇G̃_k ^2 ≤ d_2d_3,
which, by the definition of T̃(ε), implies that

    ε^2/4≤d_2d_3/∑_k=2^T̃(ε)1/α_k≤d_2d_3/∑_k=2^T̃(ε)1/α̅.

By using the fact ∑_k=2^T̃(ε)1/√(k)≥√(T̃(ε)) - 1  and (<ref>), we conclude ε^2/4≤1280·( L_12)^2  d_2d_3/19^2 l_22^1( √(T̃(ε)) -1 ) or equivalently,

    T̃(ε)  ≤( 5120 ( L_12)^2  d_2d_3/19^2l_22^1 ε^2+1)^2.

On the other hand, if  k ≥76^4L_22^4𝒟̂_y^4/ε^4, then  c_k = 19γ_k/20 k^1/4≤ε/2𝒟̂_̂ŷ. This inequality together with the definition of  𝒟̂_y, then imply that c_k  y_k ≤ε/2. Therefore, there exists a

    T(ε)    ≤max( T̃(ε), 76^4L_22^4𝒟̂_y^4/ε^4)
       ≤max(  ( 5120 ( L_12)^2  d_2d_3/19^2l_22^1 ε^2+1)^2,  76^4L_22^4𝒟̂_y^4/ε^4),

such that ∇ G_k ≤∇G̃_k  + c_k  y_k ≤ε/2  + ε/2 = ε. 

















At the k-th iteration of Algorithm <ref>, the number of backtracking step is upper bounded by  log_2(2L_11/l_11^1)+log_2(2L_12/l_12^1)+log_2(2L_22/l_22^1)=𝒪(1).
 According to Theorem <ref>, the total number of  gradient calls of Algorithm <ref> to obtain an ε-stationary point  that satisfies ∇ G(x_k,y_k)≤ε is 𝒪 ( L^4ε^-4) under the nonconvex-concave setting.











§ NUMERICAL RESULTS
 
In this section, we compare the numerical performance of the proposed PF-AGP-NC algorithm with the AGP algorithm.



 §.§ Dirac-GAN problem


The Dirac-GAN problem <cit.> can be formulated as the following nonconvex–concave minimax problem :


    min_x max_y L(x, y) =-log(1 + exp(-xy)) + log 2,

where (0, 0) is the unique stationary point.

Experimental setup.
We set the initial point of the experiment at (x, y) = (1, 1). The termination condition is defined as the norm of the gradient ∇ f(x_i, y_i) being less than or equal to 10^-5, where i represents the number of gradient calls.

For the AGP algorithm, let α_x and β_y be the stepsize of x and y, respectively. We set α_x = 0.8/√(k), β_y = 0.3 and c_k =0.5/k^1/4. For the PF-AGP-NC algorithm, we do not need to preset specific parameters.

In the experiments, we counted the number of gradient calls of the PF-AGP-NC algorithm. The results of alternating gradient projection in the k th iteration are represented by nodes in the graph and the number of backtracking step computations is represented by the connecting lines.

Results Figure <ref> shows the convergence curves of the AGP and PF-AGP-NC algorithms. It can be seen that although the PF-AGP-NC algorithm involves a double-layer structure, requiring the determination of points in the outer layer through computations in the inner layer, even when considering the number of iterations in the inner layer, the overall convergence speed is still faster than that of the AGP algorithm.





§ CONCLUSIONS AND DISCUSSION

In this paper, based on the framework of backtracking, we propose two completely parameter-free alternating gradient projection algorithms, i.e., the PF-AGP-NSC algorithm and the PF-AGP-NC algorithm, for solving nonconvex-(strongly) concave minimax problems respectively, which does not require prior knowledge of parameters such as the Lipschtiz constant L or the strongly concave constant μ.   Moreover, we show that the total number of  gradient calls of the PF-AGP-NSC algorithm to obtain an ε-stationary point for nonconvex-strongly concave minimax problems is upper bounded by 𝒪( Lκ^3ε^-2) where κ is the condition number,  while the total number of gradient calls of the PF-AGP-NC algorithm to obtain an ε-stationary point for nonconvex-concave minimax problems is upper bounded by 𝒪( L^4ε^-4). As far as we know, the PF-AGP-NSC algorithm and the PF-AGP-NC algorithm are the first completely parameter-free algorithms for solving nonconvex-strongly concave minimax problems and nonconvex-concave minimax problems respectively.  Numerical results demonstrate the efficiency of the two proposed algorithms.





















































































































































































































 Paper Title: Spatial Transformer Network YOLO Model for Agricultural Object Detection.





Spatial Transformer Network YOLO Model for Agricultural Object Detection


 



    Yash Vivek Zambre^1, Ekdev Rajkitkul^2, Akshatha Mohan^1, Joshua Peeples^1
 ^1Department of Electrical and Computer Engineering, Texas A&M University, College Station, TX, USA

	^2Department of Computer Science and Engineering, Texas A&M University, College Station, TX, USA

		yashzambre@tamu.edu, ekdev@tamu.edu, akshatha.mohan@tamu.edu, jpeeples@tamu.edu
    
=========================================================================================================================================================================================================================================================================================================================================================================



§ ABSTRACT 
Object detection plays a crucial role in the field of computer vision by autonomously identifying and locating objects of interest. The You Only Look Once (YOLO) model is an effective single-shot detector. However, YOLO faces challenges in cluttered or partially occluded scenes and can struggle with small, low-contrast objects. We propose a new method that integrates spatial transformer networks (STNs) into YOLO to improve performance. The proposed STN-YOLO aims to enhance the model's effectiveness by focusing on important areas of the image and improving the spatial invariance of the model before the detection process. Our proposed method improved object detection performance both qualitatively and quantitatively. We explore the impact of different localization networks within the STN module as well as the robustness of the model across different spatial transformations. We apply the STN-YOLO on benchmark datasets for Agricultural object detection as well as a new dataset from a state-of-the-art plant phenotyping greenhouse facility. Our code and dataset are publicly available [<https://github.com/Advanced-Vision-and-Learning-Lab/STN-YOLO>].




Spatial transformer network, Object detection, YOLO, plant phenotyping






§ INTRODUCTION

Plant phenotyping is critical for crop improvement <cit.>, yield optimization <cit.>, and sustainable practices <cit.>. Artificial intelligence (AI), particularly object detection algorithms, has transformed plant phenotyping, enhancing efficiency and performance <cit.>. The You Only Look Once (YOLO) has been used effectively in various Agricultural applications such as pest detection <cit.>, crop disease detection <cit.>, and crop harvesting <cit.>. Despite the vast number of use cases, YOLO has some limitations due to various spatial transformations <cit.>. Spatial transformer networks <cit.> are an approach to improve an artificial neural network's robustness to spatial transformations. We propose to integrate STNs with the YOLO model to incorporate spatial invariance. The STN applies a learnable affine transformations to the images that will help with object detection. The STN-YOLO model demonstrates spatial invariance and outperforms the baseline YOLO model on several Agricultural benchmark datasets.



Additionally, we present a new, high quality, annotated dataset of various plants to advance the study of agriculture plant detection and phenotyping. The new Plant Growth and Phenotyping (PGP) dataset differs from other available image datasets in the following aspects: (1) multi-spectral images captured across varying heights; (2) challenging images of multiple crops across varying illumination conditions; (3) precise annotations assisted with the use of Segment Anything model <cit.>, (4) large size and shape variations of plants. The features of the PGP dataset present new challenges for agricultural object detection. We perform extensive experiments on the proposed database and new model (STN-YOLO) for object detection. The key contributions of our work are the following:

    
  * Integration of STN within YOLO for improved object detection performance
    
  * New benchmark dataset for plant object detection.



§ RELATED WORK



 §.§ Object Detection Models

Object detection is a prominent challenge in computer vision, where researchers use deep learning to boost performance <cit.>. Two-stage detectors, such as Faster R-CNN <cit.>, involve distinct stages with a region proposal network (RPN) and region of interest (ROI) pooling for candidate bounding boxes <cit.>. In contrast, single-stage detectors such as single-shot detectors (SSD) predict bounding boxes directly by utilizing grid boxes and anchors, as described in <cit.>. Notably, YOLO has emerged as a prominent example of a single-stage detector. The advancements of CNNs architectures introduces innovations such as anchor boxes to enhance object detection <cit.>. Key advancements include R-CNN, integrating region proposals with CNN using support vector machine (SVM) classification and bounding-box regression <cit.>. However, R-CNN incurs substantial computational costs and information loss. Fast R-CNN extends this with ROI pooling and proposal refinement, improving test speed and algorithm precision <cit.>. Faster R-CNN further reduces computational requirements by introducing the RPN for regional proposals, selecting anchors based on specific criteria for outstanding recognition precision <cit.>. Ongoing advancements focus on speed optimization, including anchor-free detection <cit.>.



 §.§ Spatial Transformer Network (STN)


STN is a module with differentiable properties that performs learnable spatial transformation on input feature maps <cit.>. The transformation is dependent on the specific input, generating a singular output feature map. In the case of multi-channel inputs, an identical transformation is applied on each channel. The STN module is comprised of three components, as illustrated in the Figure <ref>. The first component is a localization network, which passes input feature maps through several hidden layers to generate the parameters for a learnable affine transformation that is applied to each feature map. Subsequently, the predicted transformation parameters θ are used to construct a sampling grid—a set of points indicating where the input feature maps should be sampled to generate the transformed output. Finally, the feature maps and the sampling grid serve as inputs to the sampler that generates the output map sampled from the input at the specified grid points.

Localization Networks The localization network (f_loc) takes the input feature maps 𝐗∈ℝ^H × W × C with width W, height H, and C channels. The features extracted by the localization network are then passed into a fully connected layer that outputs θ, the parameters of the transformation T_θ to be applied to each feature map: θ = f_loc(X). The size of θ can vary depending on the transformation type that is parameterized <cit.>. The localization network function f_loc() can take any form, such as a fully-connected network or a CNN such as ResNet18 <cit.>, but should include a final regression layer (i.e., fully connected layer) to produce the transformation parameters θ (e.g., affine transformation have six parameters).

Parameterized Sampling Grid
To perform a transformation on the input feature map, each resulting “pixel" is determined by applying a sampling kernel centered at a specific location in the input feature map. The term “pixel" in this context refers to an element of a general feature map, not necessarily an image. In a general context, the resulting “pixels" are positioned on a regular grid G = {G_i} of pixels G_i = (x^t_i, y^t_i), creating an output feature map V ∈ℝ^H' × W' × C. Here, H' and W' represent the height and width of the grid, respectively, and C is the number of channels, which is same in the input and output feature maps.

For clarity of exposition, assume that T_θ is a 2D affine transformation A_θ. In the affine case, the point-wise transformation is shown in Equation <ref>:

    [ x_i^s; y_i^s ]
    = T_θ(G_i) = A_θ[ x_i^t; y_i^t;     1 ]
    = 
    [ θ_11 θ_12 θ_13; θ_21 θ_22 θ_23 ][ x_i^t; y_i^t;     1 ]



In this scenario, (x_i^t, y_i^t) denotes the target coordinates of the regular grid in the output feature map, and (x_i^s, y_i^s) represents the source coordinates in the input feature map that define the sample points. The relationship between these coordinates is determined by the affine transformation matrix A_θ <cit.>. The STN is self-contained module that is versatile and can be seamlessly integrated at any point in an artificial neural network architecture. The inclusion of an STN at the beginning of a CNN enables the network to learn how to dynamically transform input images to account for spatial variations.



§ METHOD




 §.§ Plant Growth and Phenotyping Dataset



The Plant Growth and Phenotyping (PGP) dataset is a publicly available, high-resolution (512 × 512 pixels), multi-spectral collection of plant images captured from a state-of-the-art greenhouse. PGP consists of 1137 images from different crops: 442 corn, 246 cotton, and 449 rice plants. The dataset includes images taken from various heights of the same plants. 



The plant phenotyping facility consisted of a camera mounted on a robotic arm, capturing side and top views of the plants (only side view images are used in this work). The Multispectral Imaging System (MSIS-AGRI-1-A), integrating an MSC2-AGRI-1-A snapshot multispectral camera and a 4-channel LED illuminator, was used to collect images across four spectral bands: 580 nm (Green), 660 nm (Red), 730 nm (Red Edge), and 820 nm (Near Infrared). 

For this work, the Red, Red Edge, and Green bands were combined to create pseudo “RGB" images (as shown in Figure <ref>) compatible with pre-trained machine learning models. The individual bands were superimposed, and min-max normalization was applied to the final image. These plant images were captured between 2022 and 2024 in controlled greenhouse environments. Data augmentation techniques (e.g., rotation, cropping) can be applied using frameworks like Roboflow <cit.> to expand the dataset. The same framework is used for labeling, a two-step process involving auto-generation of initial labels using the Segment Anything Model (SAM) <cit.>, followed by manual refining of the auto-generated labels.




 §.§ STN-YOLO MODEL
    For the proposed STN-YOLO model, the STN module is first integrated in the beginning of the YOLO model to account for spatial transformations in the input images. After the images passes through STN module, the YOLO model is then used to perform the task of object detection. As stated in Section <ref>, the STN has can use any backbone for the localization networks. The localization network within the STN module comprises various convolutional layers, linear layers, and activation functions. 


Localization in STN model 
The STN models have the localization network which takes the input feature maps with width W and height H and C channels and outputs θ which are the transformation parameters. The localization network is selected by the user and can vary from a shallow or deeper network. The performances of different localization network are compared to evaluate if a shallow or deep, pre-trained model will lead to better object detection performance quantitatively (e.g., accuracy, precision, and recall) and computationally. The shallow network is comprised of one convolution layer with max pooling and ReLU activation function was used in contrast to the deep pretrained models. Following the convolution layer in the shallow network, an adaptive average pooling layer was used to aggregate the spatial resolution of the feature maps to a desired size (e.g., 28 × 28).

YOLO Backbone The output of the STN module is then given to the YOLO model where, P1, P2, P3, P4, and P5 typically refer to different levels or stages in the feature pyramid and these are the different convolutional layers as shown in the Figure <ref>. Output features are obtained by extracting features at different scales from the input image  to detect objects of various sizes and capture both fine-grained details and global context. P1 corresponds to the finest level with the highest resolution, while P5 is the coarsest level with lower resolution but capturing more context. The head in the YOLO model usually refers to the part of the network responsible for making predictions. The head takes features from multiple pyramid levels and produces predictions for object classes, bounding box coordinates, and other relevant information. 

Loss in YOLOv8
A two term loss function is used during training. The objective function consists of two components: classification loss (CLS), which penalizes errors in predicting object classes, and bounding box (BBOX) loss, which penalizes errors in predicting the coordinates of the bounding boxes around objects. Combining these losses helps train the model to simultaneously improve both classification and localization aspects of object detection. Class imbalance in single-class detection refers to the discrepancy between detection target objects and the background, typically considered the negative class. Addressing this imbalance is crucial for improving detection performance. Bounding box loss in YOLOv8 uses the complete intersection over union (CIOU) loss <cit.>, which considers box overlap and aspect ratio differences, enhancing regression performance. Distributed focal loss (DFL) <cit.> effectively tackles class discrepancy by dynamically adjusting the loss for each class during training to counteract biases caused by class imbalances.







§ EXPERIMENTAL RESULTS AND DISCUSSION




 §.§ Experimental Setup

 Three experimental runs of random initialization were used for each model and the YOLO backbones were pretrained on the COCO dataset. We used the AdamW optimizer with an initial learning rate of 0.002. The optimizer was configured with distinct parameter groups, each with specific weight decay values assigned to all parameters. The batch size for all experiments was set to 16. We used 100 epochs with early stopping (patience of 50) for each model. All the experiments were performed using the Ultralytics <cit.> framework and performed on an NVIDIA A100 GPU.

In our study, we used various benchmark datasets to assess the effectiveness of our proposed plant detection models.
The GlobalWheat2020 dataset contains 4,000 images, with 3,000 for training and 1,000 for testing, focusing on wheat detection.
The PlantDoc dataset comprises 2,569 images, with 2,330 for training and 239 for testing, addressing plant identification across varying environments.
The MelonFlower dataset includes 288 images, with 193 for training and 95 for testing, specifically targeting melon flower identification. These datasets collectively provide a robust evaluation framework for assessing model generalization across diverse plant species and environments. 





 §.§ Ablation Studies using PGP Dataset


Impact of Localization Networks
In Table <ref>, we outline the performance metrics of various localization networks, emphasizing the comparable quantitative results achieved by both the shallow localization network and deep pre-trained models.  Both the shallow and deep networks used a global average pooling layer after the convolutional layers. Based on these findings, we opted for the shallow network as the localization network in the STN-YOLO model. This decision stemmed from the observed effectiveness of the shallow network in extracting features suitable for learning affine transformations comparable to the capabilities of the pre-trained models (based on similar quantitative measures). The shallow network and the pretrained networks were finetuned  with the baseline YOLO model. The number of learnable parameters for the deeper networks, VGG16 and ResNet18 (approximately 82 million and 49 million respectively), were greater than the shallow network (approximately 5 million) for marginal improvements in performance hence the shallow network was chosen for the remaining experiments.




Impact of Spatial Information
To assess the impact of retaining more spatial information in the STN-YOLO model, modifications were applied to the last layer of the shallow localization network. Specifically, the adaptive average pooling layer, was evaluated for various spatial sizes of 1 × 1, 7 × 7, and 28 × 28 as shown in Table <ref>. This investigation aims to evaluate the model's sensitivity to variations in spatial information preservation and identify the optimal configuration for improved performance in object detection tasks. 
We wanted to look at no spatial information (1 × 1) and retaining some spatial information (7 × 7, 28 × 28). The results demonstrated that retaining the most amount of spatial information (28 × 28) resulted in the best object detection performance. If the feature maps have more spatial information, then there is more context that the model will be able to use for maximizing performance. However, as the spatial resolution is increased, so does the computational cost as the number of parameters will increase. 



Model Analysis using Explainable AI The proposed model is better on average at reducing the detection of false positives, as the precision of the STN-YOLO model is better than the baseline YOLO model as shown in Table <ref>. The STN-YOLO model also exhibits improved recall and mean average precision (mAP), which demonstrates a better object detection model as illustrated in Figure (<ref>-<ref>), where the baseline YOLO model fails to eliminate the bottom right vase area in comparison to the proposed STN-YOLO model. In addition to these metrics from Table <ref>, we also provide insights from Eigen class activation maps (EignenCAMs) <cit.> as seen in Figure (<ref>-<ref>). These EigenCAMs offer valuable visualization for model explainability, as they highlight significant aspects of the images that contribute to the model's predictions. The results support the hypothesis that incorporating the spatial invariance in the input image can result in better object detection both qualitatively and quantitatively as the STN-YOLO emphasizes more of the plant area.





Augmentations Testing One of the most important set of experiments performed was the testing  of models on the unseen augmented data as shown in Table <ref>. To do this, no data augmentation was added to the training dataset, but data augmentations were performed on the test dataset to evaluate the robustness of the STN-YOLO model across various transformations. The augmentations used on the test dataset were random cropping (15% zoom), shear (± 10^∘ horizontal and ± 10^∘ vertical), and rotation (between ± 10^∘) to simulate image conditions that may happen in the facility. We observed that the proposed STN-YOLO model gave better quantitative results in case of rotation and shear for PGP dataset compared to the baseline YOLO model. This result is intuitive as rotation and shear are affine transformations that can be learned by the STN-YOLO model. However, in the cases where cropping is used, the performance is comparable to the baseline YOLO method as cropping is not an affine transformation. Furthermore, for the combinations of the different augmentations (e.g., rotation with shear, rotation with crop), we observed that for most of the data augmentation cases (six out of eight), the STN-YOLO model performs better on average than the YOLO baseline for the precision values indicating that the STN-YOLO was effective at reducing the number of false positives by incorporating the added spatial invariance properties of the STNs.







 §.§ Benchmark Datasets Results


GlobalWheat2020 Dataset The GlobalWheat2020 dataset <cit.> is a comprehensive Agricultural dataset containing images of wheat heads sourced from various countries. From the results in Table <ref>, the  STN-YOLO model has a higher average value for most metrics (except accuracy), but the baseline YOLO model is comparable to the average value when you take the standard deviation in comparison. We can also qualitatively look at this observation from Figure <ref> where the number of detections by YOLO in the GlobalWheat2020 dataset are less as compared than the detections of STN-YOLO.
For the GlobalWheat2020 dataset, the images consisted of “zoomed"-in views of the wheatheads. In this case, the spatial invariance would not impact the object detection results and this is indicated by the slight improvements in the average object detection metrics (except for accuracy) shown in Table <ref>.

PlantDoc Dataset The PlantDoc dataset <cit.> is derived from the Plant Village dataset (PVD) <cit.>. The dataset was created using an automated system leveraging GoogleNet <cit.> and AlexNet <cit.>. Interesting observations can also be seen from the Table <ref> such as unlike the GlobalWheat2020 dataset, there is an increase in all the object detection metrics in the STN-YOLO model as compared to the baseline model. Particularly, the precision for STN-YOLO model is higher than the baseline model indicating the model reduced the number of false positives detected as noted for the PGP dataset. Qualitatively, we can see that in Figure <ref> where the YOLO detects the leaf correctly but wrongly classifies the leaf as tomato early blight leaf. However, the STN-YOLO detects and correctly classifies the leaf as Tomato Septoria. In the PlantDoc dataset, the emphasis remains on various plants and zoomed-in images. However, spatial invariance remains insignificant due to the absence of pre-translated or rotated input images. Despite this, spatial invariance could potentially offer some benefits with this dataset, particularly as certain images provide a more “zoomed-out" perspective of the plants compared to the GlobalWheat2020 dataset.

MelonFlower Dataset The MelonFlower dataset <cit.> is dataset available on Roboflow and the images were captured in a greenhouse (similar to the PGP dataset). This dataset was used to demonstrate the model's capability in detecting small objects (e.g., flowers) and different viewing geometries of the plants. As seen from Table <ref>, there is a significant difference in the precision values with the STN-YOLO model as compared to the baseline model. The precision metrics reveal a substantial improvement, showcasing the model's enhanced ability to accurately detect objects.  Visually, we can see this difference in the Figure <ref> where in the YOLO detections of the Melonflower dataset, the model detects the melon flower with an IOU of 0.50 and the STN-YOLO detections of the Melonflower dataset has a better improved detection with an IOU value of 0.60. This difference not only emphasizes the effectiveness of the STN-YOLO model but also underscores its potential for applications where precision is a critical performance indicator. Since the MelonFlower dataset had the most spatial variance among all datasets evaluated in this work, the STN-YOLO model exhibited notably improved performance, as evidenced by the precision values and other object detection metrics presented in Table <ref>.



§ CONCLUSION

This work focused on integrating STN with YOLO, creating the STN-YOLO model to address spatial invariance challenges. The model aimed to enhance plant image detection quality and can be used for downstream applications such as phenotypic feature extraction. Results demonstrated that STN improved model robustness and reduced number of false positives in the datasets as indicated by the higher precision scores. The STN-YOLO model improved performance on benchmark datasets, showcasing potential in handling real-world spatial transformations. Future work includes integrating STNs into other object detection models (e.g., future versions of YOLO <cit.>), developing new objective functions to improve the learning of the model with STNs, expanding our PGP dataset to include more images across multiple crops with different image conditions (e.g., illumination), and incorporating the near infrared channel. 



§ ACKNOWLEDGEMENT


This material is based upon work supported by Texas A&M AgriLife. Portions of this research were conducted with the advanced
computing resources provided by Texas A&M High Performance Research Computing.

IEEEtran




 Paper Title: Tabular Data Augmentation for Machine Learning: Progress and Prospects of Embracing Generative AI.




An Extended Version

cuilingxi.cs@zju.edu.cn

lihuan.cs@zju.edu.cn

chenk@zju.edu.cn

should@zju.edu.cn

cg@zju.edu.cn

  The State Key Laboratory of Blockchain and Data Security, Zhejiang University
  Hangzhou
  China
  310027







§ ABSTRACT 
Machine learning (ML) on tabular data is ubiquitous, yet obtaining abundant high-quality tabular data for model training remains a significant obstacle. 
Numerous works have focused on tabular data augmentation (TDA) to enhance the original table with additional data, thereby improving downstream ML tasks.
Recently, there has been a growing interest in leveraging the capabilities of generative AI for TDA. Therefore, we believe it is time to provide a comprehensive review of the progress and future prospects of TDA, with a particular emphasis on the trending generative AI.
Specifically, we present an architectural view of the TDA pipeline, comprising three main procedures: pre-augmentation, augmentation, and post-augmentation.
Pre-augmentation encompasses preparation tasks that facilitate subsequent TDA, including error handling, table annotation, table simplification, table representation, table indexing, table navigation, schema matching, and entity matching.
Augmentation systematically analyzes current TDA methods, categorized into retrieval-based methods, which retrieve external data, and generation-based methods, which generate synthetic data. We further subdivide these methods based on the granularity of the augmentation process at the row, column, cell, and table levels.
Post-augmentation focuses on the datasets, evaluation and optimization aspects of TDA.
We also summarize current trends and future directions for TDA, highlighting promising opportunities in the era of generative AI.
In addition, the accompanying papers and related resources are continuously updated and maintained in the GitHub repository at  to reflect ongoing advancements in the field.





Tabular Data Augmentation for Machine Learning: Progress and Prospects of Embracing Generative AI
    Gang Chen
    August 5, 2024
=================================================================================================





§ INTRODUCTION



Tabular data, such as relational tables, Web tables and CSV files, is among the most primitive and essential forms of data <cit.> in machine learning (ML), characterized by excellent structural properties, readability, and interpretability.
A testament to its significance, more than 65% of datasets available on the Google Dataset Search platform are tabular files <cit.>. This prevalence underscores its critical role across a myriad of fields, such as finance <cit.>, healthcare <cit.>, education <cit.>.
The growing availability of repositories containing structured or semi-structured data offers new opportunities for tabular data research and applications built upon it, particularly in the fields of ML and artificial intelligence (AI).

However, acquiring substantial amounts of high-quality tabular data for ML model training remains a persistent challenge <cit.>.
This is especially demanding because each individual table is modest in size and self-contained, making the overall data collection process resource-intensive and time-consuming.
According to the oft-cited <cit.> statistics, data scientists spend over 80% of their time on ML data preparation tasks, including data discovery and augmentation.
The complexity and uneven quality of massive tabular datasets from various domains further complicate the acquisition of high-quality tabular data <cit.>. 
Furthermore, in the era of large language models (LLMs), tabular data is one of the preferred data formats that LLMs consume, and existing high-quality tabular datasets may soon be exhausted <cit.>. 
Additionally, in the industrial sector where tabular data is most commonly used, the availability of data is often limited due to privacy concerns <cit.>.
All of these factors have led to significant efforts being devoted to developing techniques that support tabular data augmentation (TDA).
Through our extensive investigation, we have collected a total of 70 highly relevant studies from 2010 to 2024.
In this study, we define TDA as the process of augmenting the original dataset (table) to enhance the performance of downstream ML models.
An illustration of TDA and its role in ML scenarios is provided in Fig. <ref>. 
TDA techniques aim to enrich the ML tasks by incorporating additional data, either from external data sources[In this study, we also refer to external data sources as table pools as they can contain various table sources, including databases, Web tables, and so on.] or synthesized by generative methods[In addition to the recently much-discussed language models <cit.>, generative methods also include statistical approaches such as MICE <cit.>, deep generative models like diffusion models <cit.>, and so on.].

On the business side, the market for tabular data persists in its growth. For example, by year 2022, the US open data market () has amassed a total of 335,000 datasets, contributing to a staggering $3 trillion to the US economy <cit.>.
Moreover, the Augmented Analytics Market is expected to attain $35.6 billion expanding at a 22.70% CAGR (Compound Annual Growth Rate) report by Market Research Future[<https://www.marketresearchfuture.com/reports/augmented-analytics-market-7464>.].
Overall, TDA has garnered extensive attention from the research community and generated significant demand from the business sector.



However, the TDA task can be particularly challenging due to the unique characteristics of tabular data and the potential scale of table pools. 
Unlike homogeneous data such as images or text, tabular data is heterogeneous, typically containing both dense numerical and sparse categorical attributes <cit.>. 
Additionally, it has complicated structure, such as row and column permutation invariance and hierarchical organization, where cells belong to rows and rows belong to tables. 
Moreover, many TDA tasks involve large-scale table pools, sometimes encompassing millions of tables.
These tables often have inconsistent attribute naming and value formatting, and table pools  themselves are dynamic, changing over time.
Various TDA approaches have been proposed to address these challenges in diverse ways.
As mentioned and illustrated in Fig. <ref>, these methods can be broadly categorized into retrieval-based approaches, which involve retrieving data from table pools, and generation-based approaches, which involve generating new content.
With the rise of generative AI, many recent TDA works including both retrieval- and generation-based approaches (16 from 2023 to 2024 in our review) have embraced these models from different perspectives, and this trend continues to grow.
Generative AI, designed for generating new content based on patterns learned from training data, broadly includes various models such as pre-trained language models (e.g., BERT <cit.> and T5 <cit.>), large language models (e.g., ChatGPT <cit.> and LLaMA <cit.>), variational autoencoders <cit.>, generative adversarial networks <cit.>, and diffusion models <cit.>.
Such models have swept across various fields from computer vision (CV) <cit.> to natural language processing (NLP) <cit.>, and are now beginning to permeate the field of tabular data analysis.
In this context, we observe that there has yet to be a systematic review, synthesis, and categorization of existing TDA methods, let alone a discussion on the integration of current trending generative AI methods. Therefore, we believe it is both timely and essential to conduct such a survey to help readers and practitioners grasp the advancements in this critical field and pinpoint significant research opportunities in the era of generative AI.









Previous efforts, as summarized in Table <ref>, have focused on other research fields or only touched on aspects of TDA. 
For example, Zhang et al. <cit.> mention tabular data augmentation only briefly in their research on Web tables.
Another tutorial <cit.> focuses on data augmentation for data preparation and data integration, but does not specifically discuss its support for ML training; moreover, this tutorial does not place a particular emphasis on tabular data.
A more recent survey <cit.> on data augmentation in the era of large models mainly focuses on NLP and CV tasks, rather than tabular data.
Similarly, Wang et al. <cit.> cover data augmentation techniques across various data modalities (i.e., image, text, graph, table, and time-series), but TDA is only a small part of their broader scope. 
Their modality-independent taxonomy may not be the most suitable for an in-depth exploration of TDA specifically.
For example, one of the key methods covered in our survey, retrieval-based TDA, is not addressed in their work.
Other literature reviews on tabular data focus on table discovery <cit.>, table representation <cit.> or tabular data generation <cit.>, rather than augmentation.






Our survey stands out from previous reviews and tutorials by providing a comprehensive examination of TDA methods tailored for ML scenarios, with a special emphasis on the recent advancements in incorporating generative AI techniques.
We have meticulously selected 70 significant works from the fields of data management and artificial intelligence, dating back to 2010, to offer a diverse range of perspectives. Based on these works, we thoroughly investigate TDA-related research.
Specifically, we introduce taxonomies that categorize these works from both task and table content granularity perspectives. This allows us to clearly compare retrieval-based and generation-based methods, highlighting their respective strengths and weaknesses. Furthermore, we present a complete TDA pipeline that covers the entire process from preparation to augmentation to evaluation. Finally, we summarize future trends and highlight new opportunities in the TDA field, particularly in the context of generative AI.


The survey is structured as follows. Section <ref> provides preliminaries on TDA for ML and outlines an overall pipeline to characterize and classify approaches.
Section <ref> offers a comprehensive overview of eight classes of pre-augmentation techniques. 
Section <ref> delves into the details of specific TDA techniques at different levels, categorized into retrieval-based and generation-based methods.
Section <ref> explores post-augmentation techniques mainly for evaluation and optimization after TDA.
Section <ref> discusses challenges and future directions, with an emphasize on the trending technologies like generative AI. 
Section <ref> concludes the survey.



§ PRELIMINARIES



In this section, we will start by introducing the notation related to TDA and outlining the level-based taxonomy that defines the various levels of TDA methods (i.e., row, column, cell, and table) in Section <ref>. 
Subsequently, we will present the TDA pipeline and offer a taxonomy of methods from a task-oriented perspective in Section <ref>. 
In the following sections, tasks will serve as the primary basis for categorization, with levels providing a more granular categorization criterion.



 §.§ Notation in TDA and Level-based Taxonomy



First, we provide a formalization of tables, a prevalent data structure essential for the organization and presentation of data as follows.


A table T is an arrangement that organizes data into rows and columns, forming a grid of cells for systematic information representation. Each cell, denoted as T[i,j], is at the intersection of row i and column j, serving as the basic unit for data storage. 
The rows (T[i, :]) run horizontally and group data entries, while columns (T[:, j]) extend vertically, with each focusing on a specific data attribute.
Additionally, metadata, such as table captions, provides contextual textual information around the table.


An intuitive example of a table and its primary components — columns, rows, and cells — is depicted in Fig. <ref>.
Given a table T, we use T.ℛ and T.𝒜 to denote the set of its rows and columns (attributes), respectively. 
Notably, our analysis is restricted to tables that solely manage numerical and textual data structured in rows and columns.
This explicitly excludes tables that incorporate nested tables, lists, forms, images, or any other non-textual and non-numerical values within their cells.
We now provide the formal definition of tabular data augmentation as follows.


Given an original table T^O and a specific ML model f(Θ) parameterized by Θ, the task of Tabular Data Augmentation aims to expand T^O into an augmented table T^A that includes additional data values in its rows and/or columns.
The goal is for the ML model f(Θ), trained with T^A, to achieve superior model performance compared to the version trained with T^O.
Formally,

    T^A TDA(T^O, , [𝕋], [𝔾]), 
       s.t.    𝔼(f(Θ^A)) < 𝔼(f(Θ^O)),

where = {, , , } refers to the granularity at which the TDA operates;
𝕋, an optional input for retrieval-based TDA,
represents the pool of tables (simply called a table pool) for information enrichment use[The pool of tables, 𝕋, is often necessary for retrieval-based TDA methods. However, for generation-based TDA methods, the original table as input is usually sufficient. This is because generative models have typically been pre-trained on large amounts of external data, allowing them to retain and leverage a wealth of background information.];
𝔾, an optional input for generation-based TDA, implies the use of a specific generative method.
and 𝔼(f(Θ^A)) and 𝔼(f(Θ^O)) refer to the empirical errors of the ML models trained on the datasets T^A and T^O, respectively[Notably, the model training process can involve advanced feature engineering techniques such as coreset <cit.> and feature selection <cit.> to refine the training data from the input table. Here, we assume that these techniques will be identically applied despite the difference in the input tables.]. 



    
    Fig. <ref> depicts a typical TDA scenario:
    A data scientist aims to build a house price prediction model, but the available training data (the original table T^O containing locations and prices) is limited in features and records, and contains numerous missing or incorrect metadata and cell values (shown in gray). 
    A model trained on this low-quality data is likely to produce subpar results. 
    To address this issue, the data scientist needs to augment the original training set with more comprehensive data. This can be done a) by retrieving additional data from
    table pools for retrieval-based TDA, or b) by generating new data using existing generative methods for generation-based TDA. 
    The augmented data can include additional attributes, records, and/or cell values, reflected as enriched features and samples in the training set.
    The primary goal of this TDA process is to improve the overall quality and performance of the downstream house price prediction model.




Referring to Fig. <ref>, these TDA sub-tasks are instantiated based on the  that the TDA procedure is acting on. We formally define these tasks one by one as follows.


One significant challenge in ML is the scarcity and often uneven distribution of samples, such as in long-tail data.
To address this, row-level TDA involves adding additional rows to a given table T^O.
This procedure aims to obtain more samples for training, potentially increasing the variety of sample categories and altering the sample distribution to some extent.
For row-level TDA, only additional rows are considered, and the relationship between T^A and T^O satisfies:

    (T^O.𝒜 = T^A.𝒜) ∧ (T^O.ℛ⊂ T^A.ℛ).




    
    Continuing from the scenario in Example <ref>, the data scientist aims to expand the size and diversity of samples in the training set. 
    Fig. <ref> (a) demonstrates a generation-based TDA that only takes the original table T^O as input without external data. Generative methods typically learn the structure and the pattern of the original table T^O and then generate new synthetic record (shown in purple), resulting in the augmented table T^A_1.
    These missing values can be addressed during post-processing steps, such as filtering and imputation.




Adequate features are crucial for training high-performing ML models, but good features alone are not always enough. Therefore, column-level TDA involves extending the original table with additional columns, enriching the feature set.
Retrieval-based TDA at the column level often involves joining related tables. This procedure may necessitate feature engineering to remove less informative records and features, which can result in the augmented table having fewer rows than the original (|T^O.ℛ| > |T^A.ℛ|).
Generation-based column-level TDA, on the other hand, retains the number of rows (|T^O.ℛ| = |T^A.ℛ|).
All in all, it satisfies:

    (T^O.𝒜⊂ T^A.𝒜) ∧ (|T^O.ℛ| ≥ |T^A.ℛ|).




    
    Continuing from Example <ref>, the data scientist decides to expand the number of features in T^O. Still, we take the retrieval-based TDA as an example, as shown in Fig. <ref> (b). 
    This procedure involves computing column similarity. For example, if the column “Location” (the 2nd column in T^O) and the column “City” (the 1st column in table T^C_1) have largely similar values, T^O can be augmented with the column “DIS (distance)” (the 2nd column from T^C_1), resulting in the augmented table T^A_2. 



Empty table cells are common, and generating data representations with  values can lead to suboptimal results in downstream tasks.
Cell-level TDA involves filling these empty table cells to improve the quality of training data for ML tasks. 
For cell-level TDA, the relationship between T^A and T^O satisfies:

    (T^O.𝒜 = T^A.𝒜) ∧ (|T^O.ℛ| = |T^A.ℛ|) ∧ (∀ i,j: T^A[i,j] ≠).




    
    Continuing from Example <ref>, the data scientist decides to fill the empty cells using a generation-based TDA, as shown in Fig. <ref> (c). 
    Generative methods leverage the context (e.g., statistical distribution) from the original table T^O to fill in missing metadata and data values (shown in blue), resulting in the augmented table T^A_3.



Table-level TDA involves enriching the original table with both additional rows and columns. This aims to acquire more features and samples for ML purposes. For table-level TDA, the relationship between T^A and T^O satisfies:

    (T^O.𝒜⊂ T^A.𝒜) ∧ (|T^O.ℛ| < |T^A.ℛ|).




    Continuing from Example <ref>, the data scientist now seeks a more balanced approach that can augment the table along both the row and column dimensions. 
    Using a retrieval-based TDA, as shown in Fig. <ref> (d), the approach first retrieves the top-related tables from the table pool (e.g., T^C_1 and T^C_2). These tables are then integrated with T^O, resulting in the augmented table T^A_4, which includes additional rows (e.g., the 1st row from T^C_2) and columns (e.g., the 2nd column from T^C_1).


The aforementioned TDA methods of four different levels are summarized in Table <ref>. 
While the approaches differ, they all aim to enhance the quality of the original training dataset, thereby improving the performance of the resulting trained model.






 §.§ Pipeline of TDA and Task-based Taxonomy



In this section, we provide an overview of the key topics covered in the survey, structured around the TDA pipeline from a task-oriented perspective, as shown in Fig. <ref>.
We categorize tasks by levels for finer classification.
The pipeline highlights critical stages and procedures from the original training dataset T^O to the augmented training dataset T^A. We first overview the entire TDA pipeline, followed by a brief introduction to each pivotal procedure within it.


The TDA pipeline has two main stages:


  * Data-centric stage includes pre-augmentation and augmentation procedures to transform the original data into augmented data. Pre-augmentation involves preparation to enhance augmentation, while augmentation details existing TDA methods, divided into retrieval-based and generation-based methods.

  * Model-centric stage focuses on post-augmentation, primarily involving ML model training and evaluation to assess and optimize the augmented dataset. If the augmented dataset is not satisfactory, it is necessary to return to the data-centric stage for further augmentation.





  
(1) Pre-Augmentation

In the TDA pipeline, pre-augmentation encompasses preparation tasks to facilitate effective augmentation. For tabular data, issues like missing or incorrect cell values and unreliable metadata are common due to inappropriate data sharing and incompatible naming conventions <cit.>.
For the table pool, with tables reaching millions or more, data preparation before TDA is crucial. 
Pre-augmentation aims to a) improve the quality of both the original table and the tables from the pool (in the case of retrieval-based methods) and b) better organize the table pool for improved acceleration and scalability. 
This can involve a range of tasks, applicable to both single-table and multi-table settings, as listed in Table <ref>.
The table's right part presents their support for various different TDA tasks. The following observations can be drawn:


  * The tasks in single-table setting are applicable to all target TDA tasks, as all target TDA tasks require preprocessing of the original table. 

  * The multi-table setting tasks are specifically for handling tables from the pool and may not be suitable for generation-based TDA tasks. 

  * Entity matching, a pre-augmentation task that focuses on the relationships between rows, may not be much beneficial for the schema augmentation (), a TDA task at the column level. 

Section <ref> will provide a detailed introduction to relevant techniques for pre-augmentation tasks.






  
(2) Augmentation

Augmentation is the core procedure of the TDA pipeline,
enhancing the original table with more data to improve downstream ML tasks.
The techniques for TDA can be broadly divided into retrieval-based and generation-based methods. 
Retrieval-based methods are considered a data-driven TDA task based on the original table T^O (called query table), with table pools as additional input. 
The key to this type of methods lies in properly modeling the similarity between the query table and the tables from the pool. 
Generation-based methods can effectively leverage pre-existing knowledge acquired from pre-training to augment the input tables, without requiring additional data sources. 
Both methods can be further subdivided by the  of the augmentation, as listed in Table <ref>.
Section <ref> will cover the typical techniques for these TDA tasks.






  
(3) Post-Augmentation
Finally, the TDA pipeline includes a post-augmentation procedure that occurs primarily after ML model training. This involves evaluating and optimizing the augmented TDA results to enhance the performance of the downstream ML task.
We focus on three key aspects in post-augmentation: TDA datasets, evaluation policies, and optimization strategies.
We first elaborate on the commonly used datasets associated with TDA and these datasets' characteristics.
We then analyze evaluation policies from two perspective: original-table-based evaluation, which compares the augmented table with the original one (e.g, comparing their statistical distributions), and model-based evaluation, which compares the performance of ML models trained on augmented datasets versus baseline datasets.
We finally delve into the optimization strategies. The downstream ML model can iteratively optimize augmented results by retaining data that improves performance until a target accuracy is reached. More complex strategies, like reinforcement learning (RL) based frameworks, can also guide the data optimization process.
Section <ref> will detail post-augmentation techniques from these three aspects.




















§ TECHNIQUES IN PRE-AUGMENTATION



In this section, we review the techniques used in the pre-augmentation procedure, as introduced in Section <ref>.
As shown in Table <ref>, we have selected a collection of representative TDA works and summarized the pre-augmentation tasks they involve. 
Most of the selected works are oriented to TDA and have been published in well-known conferences or journals with high citation counts, reflecting their significance within the field. 
We have also included several target tasks other than TDA (see the rightmost part of Table <ref>), namely table search <cit.> and semantics detection <cit.>, as these often serve as intermediate steps in TDA.

Our task-oriented approach, illustrated in Table <ref>, examines four pre-augmentation tasks for the single-table setting (Sections <ref> to <ref>) and four for the multi-table setting (Sections <ref> to <ref>).
Pre-augmentation is essential for most TDA works, and the pre-augmentation tasks in Table <ref> are not mutually exclusive. A TDA work may involve one or more of these eight tasks. 
For example,  <cit.> (No.6 work in Table <ref>) employs multiple pre-augmentation tasks (table representation, table annotations, etc.) to complete its entire TDA process.

[!htbp]
Overview of pre-augmentation tasks. The table contains a non-exhaustive list of representative TDA-related works (arranged in chronological order) and the corresponding pre-augmentation tasks they involve. Detailed categorization of each task is provided in the respective sections.


!

No.     1c|Reference                               Pub. Year     16c|Pre-augmentation Methods                                                                                                                                                                                                                                                                                                                     6cTarget Tasks                                                                                                        
     1c|                                                                              Error Handling                              Table Annotation                                 Table Simplification                              Table Representation                                  Table Indexing                              Table Navigation                                    Schema Matching                                Entity Matching                              Augmentation                                      Table Search                                      Semantics Detection  
 1       Limaye et al. <cit.>    2010                               ∖                               Ontology                                         ∖                                     Content                                               ∖                               ∖                                       ∖                                  ∖                                                                                                               

2       Venetis et al. <cit.>    2011                              ∖                               Ontology                                         ∖                                     Content                                               ∖                               ∖                                       Textual                                        ∖                                                                                                               

3        <cit.>    2011             ∖                               ∖                                    ∖                 Content                           ∖                               ∖                                       ∖              ∖                                

4        <cit.>    2012            ∖                               ∖                                    ∖                 Content                           ∖                               ∖                                       ∖              ∖                                

5       Sarma et al. <cit.>     2012                               ∖                               ∖                                    ∖                                     Content+Metadata                                      ∖                               ∖                                       Textual+Metadata                               KB                                                                                                                                               

6        <cit.>    2012                                 ∖                               ∖                                    Summarization                                     Content+Metadata                                      Inverted index                              ∖                                       Textual+Metadata                                                                                                                                                                       

7       Ahmadov et al. <cit.>    2015           ∖                               ∖                ∖                 Content+Metadata                  Inverted index          ∖                   ∖              ∖                                

8        <cit.>    2015                              ∖                               ∖                                    ∖                                     Content                                               ∖                               ∖                                       ∖                                  DB                                                                                                                          

9       Kanter and Veeramachaneni <cit.>    2015           ∖                               ∖                                    ∖                 Content                           ∖                               ∖                                       ∖              ∖                                

10      Christophides et al. <cit.>    2015                              ∖                               ∖                                    ∖                                     Content                                               ∖                               ∖                                       ∖                                  KB                                                                                                                          

11       <cit.>    2016          ∖                               ∖                                    ∖                 Content                           ∖                               ∖                                       ∖              ∖                                

12       <cit.>    2016        ∖                               ∖                                    ∖                                     Content                                               LSH                                         ∖                                       ∖                                  ∖                                                                                                               

13       <cit.>    2017                              ∖                               ∖                                    ∖                                     Content                                               Inverted index                              ∖                                       ∖                                  KB+DB                                                                                                                                            

14       <cit.>     2018                              Implicit                                    Supervised-learning                              ∖                                     Content                                               LSH                                         ∖                                       Textual                                        ∖                                                                                                               

15       <cit.>    2018                                ∖                               ∖                                    ∖                                     Content                                               ∖                               Linkage graph                                       Numerical                                      ∖                                                                                                                                    

16       <cit.>    2018               ∖                               ∖                                    ∖                                     Content                                               ∖                               ∖                                       ∖              ∖                                

17       <cit.>     2018              ∖                               ∖                                    ∖                                     Content                                               ∖                               ∖                                       ∖              ∖                                

18       <cit.>    2019                               ∖                               ∖                                    ∖                                     Content                                               ∖                               ∖                                       ∖                                  KB+DB                                                                                                                                            

19       <cit.>    2019                              ∖                               ∖                                    ∖                                     Content                                               Inverted index                              ∖                                       ∖                                  ∖                                                                                                               

20       <cit.>    2019                                        ∖                               Supervised-learning                              ∖                                     Content                                               ∖                               ∖                                       ∖                                  ∖                                                                                                                                    

21       <cit.>    2019           ∖           ∖                Summarization                 Content+Metadata                  ∖                               ∖                                       ∖              KB+DB                                        

22       <cit.>    2019             ∖                               ∖                                    ∖                 Content                           ∖                               ∖                                       ∖              ∖                                

23       <cit.>   2019           ∖                               ∖                                    ∖                                     Content                                               ∖                               ∖                                       ∖              ∖                                

24       <cit.>    2019           ∖                               ∖                                    ∖                                     Content                                               ∖                               ∖                                       ∖              ∖                                

25      <cit.>    2019             ∖                               ∖                                    ∖                                     Content                                               ∖                               ∖                                       ∖              ∖                                

26       <cit.>    2019             ∖                               ∖                                    ∖                                     Content                                               ∖                               ∖                                       ∖              ∖                                

27       <cit.>    2020           ∖                               ∖                                    ∖                 Content                           ∖                               ∖                                       ∖              ∖                                

28       <cit.>    2020                               Implicit                                    Supervised-learning                              ∖                                     Content                                               LSH                                         ∖                                       Textual+Numerical                              ∖                                                                                                               

29       <cit.>     2020          Explicit                               ∖                                    Sampling                 Content                                        ∖           ∖                      ∖                                  ∖                                          

30       <cit.>    2020                              Explicit                                    ∖                                    ∖                                     Content+Context                                       ∖                               ∖                                       ∖                                  ∖                                                                                                                                    

31       <cit.>    2020                               Implicit                                    Supervised-learning                              Summarization                                     Content+Context                                       ∖                               ∖                                       ∖                                  ∖                                                                                                                                    

32      Nargesian et al. <cit.>    2020                              ∖                               ∖                                    ∖                                     Content                                               ∖                               Linkage graph                                       ∖                                  ∖                                                                                                                                    

33       <cit.>    2021              ∖                               ∖                Sampling                      Content+Metadata                  ∖           ∖                   ∖                                  ∖                                

34       <cit.>    2021                                  Implicit                                    ∖                                    ∖                                     Content                                               Inverted index                              ∖                                       Textual                                        ∖                                                                                                               

35      Liu et al. <cit.>    2021              ∖                               ∖                                    ∖                 Content                                        ∖           ∖                      ∖                                  ∖                                          

36       <cit.>    2021                                 ∖                               ∖                                    ∖                                     Content                                               ∖                               Hierarchical structure                              ∖                                  ∖                                                                                                                                    

37       <cit.>    2021              ∖                               ∖                                    ∖                                     Content                                               ∖                               ∖                                       ∖              ∖                                

38       <cit.>    2021               ∖                               ∖                                    ∖                                     Content                                               ∖                               ∖                                       ∖              ∖                                

39       <cit.>    2021             ∖           ∖                                    ∖                                     Content                                               ∖                               ∖                                       ∖              ∖                                

40       <cit.>    2022                                 Explicit                                    ∖                                    ∖                                     Content                                               ∖                               ∖                                       ∖                                  ∖                                                                                                                                    

41       <cit.>    2022                                  ∖                               ∖                                    ∖                                     Content                                               Inverted index                              ∖                                       ∖                                  ∖                                                                                                                                    

42       <cit.>     2022              ∖           PLMs                         ∖                 Content                           ∖           Clustering                      ∖                                  ∖                                

43       <cit.>   2022             ∖                               ∖                Sampling                      Content                                         ∖           ∖                                       ∖              ∖                                

44      Chai et al. <cit.>    2022                   ∖                               ∖                                    ∖                 Content                                        ∖           Clustering                      ∖                                  ∖                               

45      Santos et al. <cit.>    2022                                 ∖                               ∖                                    ∖                                     Content                                               ∖                               ∖                                       Numerical                                  ∖                                                                                                               

46       <cit.>     2022                                           Implicit                                    ∖                                    Summarization                                     Content+Context+Metadata                              ∖                               ∖                                       ∖                                  ∖                                                                                                                                    

47       <cit.>    2022                              Implicit                                    ∖                                    Summarization                                     Content+Context+Metadata                              ∖                               ∖                                       ∖                                  ∖                                                                                                               

48      Nargesian et al. <cit.>    2022                                 ∖                               ∖                                    ∖                                     Content                                               ∖                               Linkage graph                                       ∖                                  ∖                                                                                                                                    

49       <cit.>     2022                                 Implicit                                    PLMs                                             ∖                                     Content+Context                                       ∖                               ∖                                       ∖                                  ∖                                                                                                                                    

50       <cit.>    2022            ∖                               ∖                                    ∖                                     Content                                               ∖                               Linkage graph                                       ∖                                  ∖                                

51       <cit.>    2022            ∖                               ∖                                    ∖                                     Content                                               ∖                               ∖                                        ∖                                  ∖                                

52       <cit.>    2023                                 Implicit                                    PLMs                                             Summarization                                     Content+Context+Metadata                              ∖                               ∖                                       ∖                                  ∖                                                                                                                                    

53       <cit.>     2023                              Implicit                                    ∖                                    Sampling                                          Content                                               HNSW                                        ∖                                       Textual                                        ∖                                                                                                               

54       <cit.>     2023                               Implicit                                    Supervised-learning                              ∖                                     Content+Context                                       Inverted index                              ∖                                       Textual                                        ∖                                                                                                               

55       <cit.>     2023                              Implicit                                    ∖                                    Sampling                                          Content+Context                                       LSH+HNSW                                    ∖                                       Textual                                        ∖                                                                                                               

56       <cit.>    2023                             Implicit                                    ∖                                    Sampling                                          Content+Context                                       ∖                               ∖                                       ∖                                  ∖                                                                                                               

57       <cit.>    2023           Implicit                ∖                ∖                 Content+Context+Metadata          ∖           ∖                   ∖              ∖                                                                        

58       <cit.>     2023          Implicit                ∖                ∖                 Content+Context                   ∖           ∖                   ∖              ∖                               

59       <cit.>    2023          ∖           ∖                ∖                 Content+Context                   ∖           ∖                   ∖              ∖                                

60       <cit.>    2023          ∖                               ∖                                    ∖                                     Content                                               ∖                               ∖                                       ∖                                  ∖                                

61       <cit.>    2023          ∖                               ∖                                    ∖                                     Content                                               ∖                               ∖                                       ∖                                  ∖                                

62       <cit.>    2023          ∖                               ∖                                    ∖                                     Content                                               ∖                               ∖                                       ∖                                  ∖                                

63       <cit.>    2023          ∖                               ∖                                    ∖                                     Content                                               ∖                               ∖                                       ∖                                  ∖                                

64       <cit.>     2023          ∖                               ∖                                    ∖                                     Content                                               ∖                               ∖                                       ∖                                  ∖                                

65       <cit.>    2023          ∖                               ∖                                    ∖                                     Content                                               ∖                               ∖                                       ∖                                  ∖                                

66       <cit.>    2024          Implicit                ∖                                    ∖                 Content                           ∖           Linkage graph                   ∖                                  ∖                                

67       <cit.>    2024          Implicit                ∖                                    ∖                 Content                           ∖           ∖                   ∖                                  ∖                                

68       <cit.>     2024          ∖                               ∖                                    ∖                 Content                                               ∖           ∖                   ∖              ∖                                                    

69       <cit.>     2024          Implicit                ∖                                    ∖                 Content+Context                   ∖                               ∖                   ∖                                  ∖                                

70       <cit.>     2024          Implicit                ∖                                    ∖                 Content                  ∖                               ∖                   ∖                                  ∖                                
















 §.§ Error Handling



Error handling in pre-augmentation refers to the preprocessing of the dirty data in tables. 
Real-world tabular data often contain errors, such as mistakenly substituted proximal characters and unnecessary repetition of tokens in cell values.
Generally, there are three types of errors that considered in pre-augmentation, missing values <cit.>, misspellings <cit.>, and numerical outliers <cit.>.
When generating table representations based on token embedding with such errors, incorrect forecasts are inevitable <cit.>. Therefore, various methods have been developed to address tabular data errors, which can be categorized into explicit and implicit approaches.

Explicit error handling involves detecting and correcting errors directly. 
For example,  <cit.> handles missing data with a simple approach that uses uniform random sampling for categorical columns and the median value for numerical columns.
Cappuzzo et al. <cit.> handle missing values using classical database techniques like Skolemization <cit.>, which introduces a new constant or function symbol (called a Skolem function) to represent the unknown or missing value. 
Zhao and Fernandez <cit.> detect missing data through a voting mechanism. 
They observe that missing values are usually covered by multiple attributes while common values typically only appear under very few. 
Thus, they identify value nodes with votes exceeding a predefined threshold for different attributes as missing data.

Implicit error handling does not directly detect errors but instead enhances the model's robustness to them. 
For example, Hu et al. <cit.> propose a table noise generator in their automatic augmentation framework.
The generator introduces artificial table noise to the training data to enhance the model's ability to handle potential errors. 
They generate cell value noise in three ways: (1) replacing/inserting characters with proximal characters, (2) deleting/repeating characters, and (3) changing the numeral display format such as using scientific notation. 
Furthermore, multiple studies <cit.> adopt word embedding that can tolerate misspellings to some extent. 
For example, Dong et al. <cit.> support semantic joins by using word table representation learning to join cells with similar meanings, thus handling data with misspellings and discrepancies in formats or terminologies.


Real-world tables often contain various types of errors, making it essential for table augmentation models to be robust in combating such issues. 
Explicit methods require an established procedure of error detection and correction, adding additional steps and potentially leading to error propagation. 
On the other hand, implicit methods, while more streamlined, lack the interpretability of explicit methods. 
Moreover, both methods typically ignore addressing errors in numerical columns, such as out-of-bound values, which are often harder to detect and correct than textual errors.
It is worth mentioning a recent trend of exploring the use of pre-trained language models (PLMs) for table representation. These approaches <cit.> show certain resistance to spelling errors and offer a promising direction for handling noisy data. 





 §.§ Table Annotation



Table annotation involves inferring metadata information about a table, such as column types and the relationship between columns <cit.>. 
This task helps recover the semantic information within a table and is particularly useful for table augmentation by assessing the similarity between tables. 
Typically, metadata for tables is unreliable or incomplete due to inappropriate data sharing methods <cit.>. 
Even when metadata is available, tables from a wide range of sources can have incompatible metadata with different naming conventions and terminologies. 
Consequently, table annotation is crucial for retrieving syntactically and semantically relevant tables to augment the original table.
These approaches are divided into the three following categories.

Ontology-based table annotation. 
In earlier research, ontologies were frequently used for annotating table elements.
For example, as illustrated in Fig. <ref> (a), Limaye et al. <cit.> develop a classic table discovery system that utilizes an ontology for table annotation at multiple levels.
This includes cell-level annotation with ontology entities (e.g., labeling the cell “mechanical ape!” as a song name), column-level annotation with ontology types (e.g., identifying the column header “Title” as “Song”), and pairwise column annotation with ontology relationships (e.g., annotating the relationship between “Artist” and “Song” as “written-by”).
Similarly, Venetis et al. <cit.> annotate columns using class labels from a “isA” database for column labels and binary relationships automatically extracted from the Web for relationship labels. 





Supervised-learning-based table annotation. 
Recent studies have leveraged labeled data and supervised learning techniques for table annotation. 
Specifically, researchers <cit.> train ML models on tables labeled with a fixed set of 78 types (address, description, year, etc.) and then use the trained models to annotate unknown tables, as illustrated in Fig. <ref> (b). 
On the other hand, feature-based approaches <cit.> that compute several similarity signals (e.g., value overlap) to learn column representations have been used to detect semantic types.
 <cit.> further enhances semantic type detection at the column level by incorporating the topic of table (a.k.a. global context[In the context of learning from tabular data, global context refers to the values from the entire table, whereas local context pertains to values from neighboring table elements. For example, for a specific row in the original table, the rows immediately above and below it represent the local context, while the entire original table represents the global context.]) as a new similarity signal for column representation. 


PLM-based table annotation.
More recently, PLMs have been employed for table annotation due to their superior performance <cit.>.
For example,  <cit.> first leverages PLMs to encode columns, and then these column embeddings are used for clustering to determine the column IDs.
Suhara et al. <cit.> propose , a multi-task learning framework, as illustrated in Fig. <ref> (c). Specifically,  uses a single BERT model to complete both tasks of predicting the column types and the relationships. 
Additionally,  incorporates the entire table as input to capture the table context. 
Despite recent efforts, existing methods rely heavily on large-scale and high-quality labeled data. 
Miao and Wang <cit.> propose , a lightweight framework for column annotation.
 employs contrastive learning techniques to learn robust representations for tables while maintaining minimal overhead, leveraging a large-scale unlabeled table pool.


Real-world tables often have incomplete and incorrect metadata, necessitating table annotation to recover table semantics. 
Ontology-based table annotation can be limited by the ontology's coverage, particularly for domain-specific data, while supervised-learning table annotation depend heavily on large-scale, high-quality labeled data. 
Ontology-based methods typically offer high efficiency, whereas learning methods (including those based on supervised learning and those based on PLMs) achieve higher accuracy.
One potential direction for improvement involves integrating these methods and exploring the efficient use of PLMs for table annotation. Additionally, integrating LLMs and Retrieval-Augmented Generation (RAG) <cit.> with ontologies can lead to promising developments, particularly in domain-specific scenarios where leveraging external business knowledge is crucial.








 §.§ Table Simplification



Table simplification involves streamlining a table down to its essential elements, which can be addressed from both the content and semantic perspectives.
From the table content perspective, this procedure, known as table sampling, involves selecting portions of the table to retain as much information as possible. This is particularly useful for fitting data into limited token lengths for language models.
From the table semantics perspective, the procedure, referred to as table summarization, entails identifying the main topic or theme of the table to better understand its meaning.
Accurate summarization helps in comparing tables for similarity and ensures that any added rows and columns remain consistent with the table's original theme.
These two different perspectives are introduced as follows.

Table sampling selectively choose table content to preserve the original information as much as possible.
An early work <cit.> directly selects the top-K samples for each column as input.  <cit.> and  <cit.> adopts the stratified sampling method that divides the samples into several subsets and randomly selects samples from each subset.
<cit.> adopts a frequency-based approach that samples the most frequent cell values for each column.
Meanwhile,  <cit.> and  <cit.> sample rows based on an importance score derived from the cell-level TF-IDF score for each row.

Table summarization aims to extract semantic themes or topics within tables. 
For example, Zhang et al. <cit.> incorporate a topic-aware prediction module into their framework , responsible for summarizing table semantics. 
In particular, this module produces a topic vector from the values across the entire table, representing the global context of that table. 
's ablation experiments show that considering a table's global context improves table understanding and mitigates ambiguities.
For web tables, there are often pre-existing table metadata, such as table captions that summarize the table's contents <cit.>. 
In this case, various works <cit.> directly leverage these metadata by converting them into vectors and concatenating them with the table's representation.



Table simplification aims to extract the core information of a table either in content or semantic level. 
At the content level, table sampling selectively samples table content for constraints like limited token length for LLMs.
Table sampling can also improve efficiency and scalability when facing large-scale tables.
At present, most table sampling methods are statistical-based which are fast and resource-efficient, yet they fall short in terms of precision. Sampling based solely on the value distribution (such as sampling the most frequent values), can lead to an incomplete representation of the original table. Therefore, there is a need to investigate more efficient and accurate sampling algorithms, potentially leveraging learning-based approaches.
At the semantic level, table summarization serves both to extract the key semantic essence of the table and verify the coherence of any table augmentations. At present, as a preprocessing step, there are relatively few techniques that directly extract the primary topics or themes of tables. This is largely due to the challenge of achieving lightweight, efficient table summaries. Emerging LLMs may offer a promising solution in this regard. 







 §.§ Table Representation



Table representation involves converting the table elements such as rows, columns, and cells into a latent vector space.
This transformation prepares the table for robust use in subsequent TDA model. 
The past decade has witnessed the effectiveness of employing deep neural networks for table representation learning <cit.>, thereby enhancing TDA. 
The basic idea of table representation is to create vector representations of tables that preserve their syntax and semantics.
These vectors can then be compared using methods like cosine similarity to assess their relatedness. 
The effectiveness of these table representations depends on the information source they contain, which can be summarized as table content, table context, and metadata.
We briefly introduce each type of approach as follows.

Content-based table representation.
Typically, several works <cit.> derive features from table content, such as column distribution, column type, to create table representations.
Also, many studies <cit.> directly transform table content into vectors. For example, discrete columns can be represented as one-hot vectors <cit.>, analog bits vectors <cit.>, and embedding vectors <cit.>.
More recently, there are increasing studies <cit.> leveraging language models to encode tables. They first sequentialize table content, and then feed the sequential text into language models for encoding.


Context-based table representation. 
Recent research <cit.> underscores the importance of contextual information, such as column relationships <cit.>, alongside table content.
For instance, Khatiwada et al. <cit.> enhance the  model by incorporating semantic relationships between column pairs, refining its comprehension of table contexts and filtering out tables with similar columns but different contexts during the TDA process.
Hu et al. <cit.> utilize PLMs (e.g., BERT) to capture the contextual relationships between columns, aiding in the identification of relevant tables in the table pool for TDA.
 <cit.> employs a hybrid model to utilize signals from both the global context (values from the entire table) and the local context (predicted types of neighboring columns). 
Similarly, several studies <cit.> capture the global context of the table by feeding the entire table into column encoders, resulting in the encoded column vectors that contain distilled global table context.
Several works <cit.> leverage graphs to represent tables, with nodes representing cell values and edges indicating relationships between nodes (e.g., cell belonging to one specific column), aiming to better capture the table structure.
These approaches have demonstrated high prediction accuracy, highlighting the effectiveness of integrating table context into table representations.

Metadata-based table representation. 
In addition to using data values organized in rows and columns, several studies <cit.> also utilize metadata. 
Tables are often accompanied by secondary information such as the table captions and the containing webpage's title, forming part of the textual information <cit.>. 
Yakout et al. <cit.> consider element-wise and cross-element similarities, where elements encompass both the table content (e.g., tuples and column values) and metadata (e.g., table captions and URLs). 
 <cit.> improves upon  <cit.> by incorporating metadata such as headers, captions, and topics. 



Table representation forms the basis of tabular data processing and is not limited solely to TDA tasks. 
Consequently, numerous studies have concentrated on this domain. 
Early approaches that relied on table content were found inadequate. 
Current research typically leverages both the content and contextual information of tables to create richer table representations, yielding improved outcomes. 
The use of LLMs has further improved the extraction of semantic information from tables. 
However, a notable challenge remains in accurately capturing the structural information of tables, such as row/column permutations invariance.






 §.§ Table Indexing



Table indexing involves assigning a unique identifier (index value) to table elements, allowing for quick and efficient lookup and retrieval based on their index values. 
Many retrieval-based TDA methods use indexes to enhance efficiency and scalability, particularly when dealing with large-scale table pools with millions of tables. 
Researchers have utilized various types of indexes, such as the inverted index <cit.>, Locality Sensitive Hashing (LSH) index <cit.>, and graph index such as Hierarchical Navigable Small World (HNSW) <cit.>. 



Inverted index is a structure of posting lists that maps each distinct value to its containing structures, such as tables (see Fig. <ref> (a)), rows, columns, or cells. 
When searching for candidates, the inverted index requires reading all or a large subset of these posting lists, which can lead to memory management issues and long read times in large-scale table pools. 
To address this,  <cit.> incorporates a dictionary for quick access to the posting lists, avoiding the need to search the entire inverted index. 
 <cit.> extends the single-attribute inverted index by adding a super key, a fixed-size hash value that aggregates all possible key value combinations into a single entry. This approach enables efficient multi-attribute join discovery without the prohibitive storage requirements of a full multi-attribute index.

LSH index has been widely adopted in approximate nearest neighbor search in high-dimensional spaces <cit.>.
The basic idea is that similar vectors are more likely to be hashed into the same bucket <cit.> (see Fig. <ref> (b)), which significantly improves query time with minimal accuracy loss.
For instance,  <cit.> uses an extension of the LSH index to ensure that search time remains constant regardless of repository size. 
 <cit.> adjusts the LSH index using the simHash function to better estimate the similarity between column embedding vectors.

HNSW index is a graph-based approximate nearest neighbor search technique designed for high-dimensional data. 
HNSW involves a hierarchy of graphs, where each layer's graph is built on top of the previous layer, allowing for efficient search: a query can start in any layer and traverse the graph to find the nearest neighbors, as illustrated in Fig. <ref> (c). 
The HNSW index has recently been applied to accelerating the table retrieval process for TDA, reportedly achieving a 400× performance gain over the LSH index <cit.>. 


Table indexing substantially improves the efficiency and scalability of table retrieval for TDA. 
While inverted indexes provide high accuracy, they can encounter memory management issues and prohibitively long read times in large-scale table pools.
To address these limitations, recent works <cit.> have explored techniques like LSH and HNSW for table indexing, which can significantly improve efficiency and maintain accuracy. 
Furthermore, some study <cit.> has adopted a hybrid approach, combining both LSH and HNSW indexing techniques to utilize the strengths of each method. 
However, most existing approaches are designed for static table pools. In today's information age, table pools are continuously evolving. Adapting table indexing methods to handle dynamic, ever-changing data remains a major challenge that has yet to be fully addressed.





 §.§ Table Navigation



Table navigation involves establishing a navigational framework over a table pool. 
Essentially, it refers to organizing the table pool in a way that highlights connections between similar tables, such as through edges in a graph or by clustering them together.
With table navigation, the subsequent TDA can retrieve relevant data for augmentation more easily and efficiently.
Existing works typically employ cluster structures <cit.>, hierarchical structures <cit.>, or linkage graphs <cit.> to manage tables in table pools.



Cluster structure groups related tables in the table pool.
For example, Chai et al. <cit.> employ Multivariate Gaussian Mixture Model to cluster similar data points (rows), as illustrated in Fig. <ref> (a), aiding in the downstream data point selection and augmentation tasks.
To facilitate the subsequent column integration,  <cit.> clusters similar columns using hierarchical clustering by iteratively merging the closest clusters based on Euclidean distances between column embeddings. This method can also be considered as a hierarchical structure to some extent.


Hierarchical structure allows navigation from broader concepts to more specific ones, as demonstrated in Fig. <ref> (b).
For instance, Ouellette et al. <cit.> propose a hierarchical structure called  for organizing the tables in the pool. This model allows users to narrow down to potential datasets by navigating between groups of attributes sets connected by edges indicating subset relationships in a hierarchical manner.
 seamlessly integrates table search and table navigation, making the acquisition of related table more effective. 
 
Linkage graph uses edges to indicate the relevance between tables in the table pool, as shown in Fig. <ref> (c). 
For instance, Nargesian et al. <cit.> make use of Directed Acyclic Graphs (DAG) to structure the table pool.
In the DAG, nodes represent attribute sets, with the labels of a non-leaf node summarizing the content of the attribute sets in their corresponding subgraph.
Another approach  <cit.> constructs an enterprise knowledge graph, with nodes representing attributes within the table pool, edges indicating relationships between two nodes. 
Specifically,  also contains hyperedges connecting any number of nodes hierarchically related, such as attributes of the same table, or tables of the same table pool, thereby facilitating the exploration of relevant data items. 
More recently,  <cit.> constructs a similarity graph where columns are represented as nodes. These nodes are connected by different types of edges corresponding to various features, such as embedding similarity and distribution similarity.
GNNs are then used to propagate these similarity signals to facilitate join discovery.


Table navigation restructures the table pool, allowing retrieval-based TDA to locate and access similar tables more effectively.
This concept is fairly new and continues to evolve.
Clusters offer a relatively simple method for organizing table pools, indicating whether tables have similarity or belong to the same class; however, they lack the capability to convey more complex information, such as hierarchical relationships.
For both cluster and hierarchical structure, they may not fully capture the relationships within table structures, such as the “entity-property” relationship between columns (e.g., entity “person” has the property “gender”).
Both hierarchical structure and linkage graph face efficiency issues. 
Developing scalable and robust methods for navigating such vast table repositories remains a significant challenge in this field.







 §.§ Schema Matching



Schema matching involves evaluating the relatedness between two table columns. 
In this context, the set of column headers is typically referred to as the table schema <cit.>.
Schema matching methods are frequently adopted in retrieval-based TDA for identifying and fetching those related columns and tables to expand the features in ML models. 
Due to the varied data types within tables, schema matching methods are categorized into textual matching, numerical matching, and metadata matching.

Textual matching is the most commonly used schema matching technique because textual columns usually contain more information than numerical ones. 
Below is a concise overview of some common textual matching methods.


  * Value overlap <cit.>: If there is a significant overlap in the value sets of two columns, then the columns are considered related.

  * Semantic overlap <cit.>: When leveraging table annotations (see Section <ref>) to derive labels describing column semantics, two columns are considered a match if there is a substantial overlap between their respective labels.

  * Embedding similarity <cit.>: Related columns are identified by computing the similarity of their corresponding embeddings in vector space.


Numerical matching is concerned with numerical columns. 
These methods typically evaluate the value distribution <cit.> to derive insight from numerical data. 
For instance,  <cit.> utilizes the Kolmogorov-Smirnov statistic to decide whether two numerical columns come from the same domain distribution and thus can be matched. 
Santos et al. <cit.> propose the Quadrant Count Ratio (QCR) hashing scheme. This method divides the numerical values of two columns into four quadrants based on their sign (i.e., positive or negative). Only the points in the same quadrant would be assigned with the same hash value. Columns with a certain number of matching values are likely to be correlated.

Metadata matching is also commonly used, as tables are usually accompanied by secondary information such as column headers. For instance, several studies <cit.> consider two columns to be related when their column headers have a semantic overlap above a given threshold.


Schema matching is a crucial step in retrieval-based TDA, used to determine the similarity between table attributes (columns) and infer the overall similarity between tables. 
This field has been extensively studied, and different types of schema matching methods are often adopted simultaneously, yielding promising results.
However, there remain a wide range of research opportunities, particularly in the domain of numerical schema matching, such as handling different numerical display formats. 
Even current powerful LLMs cannot handle numbers well.
Moreover, the vast majority of existing work in this area has focused solely on the similarity between single columns. The similarity between combined or composite columns is rarely addressed.




 §.§ Entity Matching



Entity matching involves identifying connections between entities in different tables, facilitating the localization of relevant entities and tables for retrieval-based TDA. 
These methods are particularly relevant in the context of horizontal tabular tables, where entities are typically represented as rows and their attributes are in columns. 
Based on the data source to which the entities are matched, the methods are categorized into KB-referenced entity matching and DB-referenced entity matching.

KB-referenced entity matching maps table entities to their referenced entity in a knowledge base (KB) <cit.>, aiming to enhance the semantic understanding of tables. 
Essentially, if two table entities are related to the same KB entity, the likelihood of these two table entities being related is high.
For example, Sarma et al. <cit.> link table entities to KB entities to acquire weighted label sets for representing table entities. 
However, not all table elements can be mapped to predefined types and relationships in the referenced KB.
To address this,  <cit.> proposes an alternative way to weaken the strict mapping, using soft constraints based on graphical model to encode a higher preference for sets of related entities. 

DB-referenced entity matching determines whether entities and their corresponding properties from table pools (or databases) refer to the same real-world object as the entities in the original table.
Christophides et al. <cit.> outline a general framework for this task, comprising two main components: (1) similarity metrics, which compare entity descriptions and (2) blocking techniques, which group tables in the table pool that are approximately similar for enhanced efficiency of this process.
More recent efforts have begun to use iterative approaches, where previously discovered matching entities serve as input for computing similarities between further tables in the table pool <cit.>.

 
Hybrid entity matching combines both KB- and DB-referenced approaches to obtain a wider range of information sources. 
For example,  <cit.> identifies candidate entities from both sources. Entities in KB that share categories or types with the original table entities are considered good candidates. Similarly, entities in DB tables that contain the original table entities or have related captions to the original table are also considered candidates.  <cit.> enhances  by incorporating word embeddings for table entities. 
 <cit.> further improves  by carefully designing features that combine evidence from multiple sources (e.g., the similarity between table headings can be calculated using edit distance from DB or mapping probabilities from KB and DB).



Entity matching uncovers relationships between entities across different tables, facilitating the retrieval of relevant entities and tables.
For KB-referenced matching, knowledge bases may have limited coverage when applied to real-world table pools.
To address this, several works <cit.> have integrated KB- and DB-referenced approaches to broaden the range of information sources.
However, a common limitation in both KB- and DB-referenced entity matching is the assumption made in previous works that the leftmost column in a table contains the entity ID or name, which is not always the case. 
Indeed, entity matching has seen decreased use in the past few years compared to schema matching. 
An emerging field might be the combination of LLMs and RAG to replace traditional KB-based methods for entity matching. LLMs can better capture entity semantics, reducing the reliance on the leftmost column. Meanwhile, RAG techniques can effectively leverage and update the KB, counteracting the issue of limited scope.




 §.§ Scenarios for Pre-augmentation Techniques



This section empirically summarizes the specific pre-augmentation tasks used in different TDA scenarios, as depicted in Fig. <ref>. 
The four single-table-setting pre-augmentation tasks apply to both retrieval and generation-based TDA, while the multi-table-setting tasks are only suitable for retrieval-based methods, which require an external table pool to handle multi-table relationships.



Retrieval-based TDA operates on a table pool often possessing large-scale data, thus leading to issues such as inconsistent formatting and dynamic data. 
Common pre-augmentation methods in these scenarios include: 
(1) table simplification for reducing table information; (2) table indexing and table navigation for fast retrieval; (3) schema matching and entity matching for addressing inconsistent formatting; and (4) table indexing and table navigation for managing dynamic and large-scale table pools.

Generation-based TDA using a single original table may suffer from issues in data-scarce scenarios. In such scenarios, pre-augmentation techniques like table annotation (providing additional information or labels) and table summarization (extracting key information) are necessary.
Furthermore, in privacy-preserving scenarios, generating synthetic data often benefits from table sampling, which involves using only partial data.
Note that both table summarization and table sampling are part of table simplification, as discussed in Section <ref>.

Meanwhile, both retrieval- and generation-based TDA approaches face some common challenges, such as low-quality and imbalanced tables. 
In low-quality scenarios (e.g., tables with null or erroneous values), common preaugmentation techniques include error handling and table annotation (to annotate column when column names are missing). 
In imbalanced scenarios, table sampling within table simplification is a common scheme. 
Additionally, TDA often requires semantic awareness; for these semantic-aware scenarios, and commonly used techniques include table annotation, table summarization within table simplification, and table representation.







































        





























§ TECHNIQUES IN TABULAR DATA AUGMENTATION

 

This section will delve into the current state-of-the-art techniques for tabular data augmentation (TDA). 
As outlined in Table <ref>, we first classify TDA tasks into two primary categories: retrieval-based TDA (see Section <ref>) and generation-based TDA (see Section <ref>). Within these two categories, the approaches can be further divided into different levels: adding rows <cit.> or columns <cit.>, augmenting individual cells <cit.>, and extending the original table with both rows and columns <cit.>. Thus, for each category, we will discuss the corresponding table augmentation work at the , , , or  level. Fig. <ref> provides a concise overview of the TDA works discussed in this section, along with their detailed taxonomy.


  FARROW/.style=arrows=-Latex[length=1.25mm, width=1.mm], , 
  U/.style = circle, draw=melon!400, fill=melon, minimum width=1.4em, align=center, inner sep=0, outer sep=0,
  I/.style = circle, draw=tea_green!400, fill=tea_green, minimum width=1.4em, align=center, inner sep=0, outer sep=0,
  cate/.style = rectangle, draw, minimum width=8em, minimum height=2em, align=center, rounded corners=3, 
  cate2/.style = rectangle, draw, minimum width=33em, minimum height=2em, text width=33em, align=left, rounded corners=3,anchor=west,
  cate3/.style = rectangle, draw, minimum width=43em, text width=43em, minimum height=2em, align=left, rounded corners=3,
  encoder/.style = rectangle, fill=Madang!82, minimum width=10em, minimum height=3em, align=center, rounded corners=3,
















 §.§ Retrieval-based TDA



By retrieval-based, we refer to the process of enhancing the original table (query table T^O) with realistic data sourced from table pools 𝕋 = {T_i}.
One important difference between augmenting tabular data and augmenting other data modalities lies in the availability of existing data resources (such as databases and data warehouses), which provide opportunities for discovering fresh, realistic data.
In contrast, other data modalities (such as images) primarily rely on transforming the original data to generate new data that has not been seen before <cit.>.
Retrieval-based TDA tasks are further divided into Entity Augmentation () at the  level, Scheme Augmentation () at the  level, Cell Completion () at the  level, and Table Integration () at the  level. 
They will be introduced as follows.



  §.§.§ Entity Augmentation ()


Previous works <cit.> define retrieval-based TDA at the row-level as entity augmentation, as rows in tabular data generally correspond to specific entities. 
Entity augmentation extends a given table with more rows or row elements retrieved from table pools. 
Directly retrieving table rows without considering the context of the entire table is not feasible. Therefore, existing solutions typically first search for tables that are unionable with the input table and then select entities from these unionable tables to augment the original table. 
As we shall see, table search <cit.> is inherently involved in this process, serving as an intermediate step that feeds into tabular data augmentation <cit.>. 

To measure the unionability of two tables, existing solutions typically start by using table representation techniques (see Section <ref>) to convert tabular data into latent-space vectors. 
These vectors, at the higher level, are then used to compute the relatedness between the source and target tables, with the resulting scores used to rank the target tables.
For instance, Sarma et al. <cit.> search for the top-k related tables that are entity complements to the input table and then use these related tables to populate the input table. 
Most related works follow a similar approach as <cit.>, with distinctions mainly in the methodology employed to identify the top-k unionable tables. These methodologies can be further categorized into four domains: statistical, KB-based, graph-based, and PLM-based.


Statistical entity augmentation methods use statistical models to estimate the unionability between two tables. 
For example,  <cit.> measures the context-to-context and table-to-context similarity by calculating the cosine similarity of their TF-IDF vectors.
 <cit.> utilizes three statistical models to statistically test the value overlap, semantic overlap, and embedding overlap between attributes (columns), and then aggregate the results to derive the unionability between tables. 
 <cit.> extends  by incorporating additional statistical measures for numerical value distribution, which is based on Kolmogorov-Smirnov (KS) statistic. 


KB-based entity augmentation methods consider a knowledge base (KB) for identifying potential unionable tables, instead of solely relying on tables within a table corpus or a table pool. 
For example, Das Sarma et al. <cit.> represent table entities as weighted label sets from a knowledge base (WebIsA <cit.> or Freebase[<http://www.freebase.com>]) or from a table corpus, and take their dot product to compute the unionability between the input and candidate tables. 
 <cit.> incorporates the DBpedia[<https://www.dbpedia.org>] knowledge base to identify candidate entities. In their approach, they not only collect entities from similar tables, but also collect entities sharing the same types or categories from DBpedia with the input entities. They have proved that using related tables and using a KB are complementary when searching for candidate entities. 
Subsequently,  <cit.> further improves  for entity augmentation by incorporating Word2Vec to train table embeddings for entities.
More recently,  <cit.> leverages external knowledge bases, YAGO 4[<http://yago-knowledge.org>], to annotate both columns and the binary relationships between columns. To address the limited coverage of KB,  proposes a self-curated knowledge base on top of the table pool.
To be more specific,  hypothesizes that columns with common semantics have overlap values. Thus, it first assigns a unique synthesized label to each column in the table pool with a confidence score 1; if two columns have overlap values, their synthesized labels can be assigned to each other with a confidence score based on the relatedness that are computed using the two columns' values.

Graph-based entity augmentation methods convert tables to a graph to compare the unionability between tables. 
An illustrative example is  <cit.>, which computes topic-sensitive pagerank (TSP) over a weighted graph, where nodes represent tables and edges indicate direct pairwise match between tables. 
 <cit.> determines entity similarity from a table representation using compact tripartite graphs. Initially,  forms a heterogeneous graph based on the table, where cells, entities, and columns act as nodes, and the relationships derived from the table schema serve as edges.  then uses random walks to traverse the graph to construct sentences that can describe the table. The sentences are fed into word embedding algorithms like word2vec for entity representation and unionability computation.
More recently,  <cit.> models a table as a hypergraph and then encodes the hypergraph for table similarity prediction. In the hypergraph, nodes indicate table cells and three different types of hyperedges represent row, column, and the entire table, respectively.


PLM-based entity augmentation methods have been adapted for tabular data more recently, where TDA is one of the cases.
PLMs often serve as table encoders, enhancing the capture of table semantics and structure for subsequent similarity calculations.
Recently,  <cit.> leverages PLMs, i.e., RoBERTa, to encode columns.
It then derives a unionability score between two tables using column aggregation algorithms.
Recently, Hu et al. <cit.> employ PLMs (BERT and RoBERTa) to obtain contextualized representation of column pair relationships, capturing nuanced table contexts to identify unionable tables in the pools.


Entity augmentation can expand and diversify the samples for ML, making it a long-standing research focus.
Statistical entity augmentation methods are typically based on fixed distribution assumptions and hard to refine. Meanwhile, it may encounter issues of computational efficiency and scalability when dealing with large-scale datasets.
KB-based entity augmentation, on the other hand, suffers from limited KB coverage, often leading to low recall in practical scenarios. However, with the development of LLMs, the incorporation of RAG may be a future direction. RAG combines the dynamically retrieved information with the ability of generative models, reducing information loss and improving information relevance, thereby increasing recall rates.
Graph-based entity augmentation, which converts tables to graphs, can better grasp the structure relationship of tables, but it also requires an additional encoding mechanism for the constructed graph (e.g.,  <cit.> uses a transformer module to encode the hypergraphs). The entire process would be resource-intensive and time-consuming.
For PLM-based entity augmentation, most methods directly sequentialize tabular data ignoring the table structure. 
These limitations highlight the need for more advanced, structure-aware methods that can effectively leverage the strengths of language models while addressing the unique challenges posed by tabular data.
Furthermore, a significant portion of the existing research has focused primarily on table union search, without addressing the potential differences between the retrieved tables and the original data. 
Even when unionable tables are found, discrepancies may still need resolution before the tables can be combined.











  §.§.§ Schema Augmentation ()

We refer to retrieval-based TDA at the column level as schema augmentation, the process of extending the schema of the original table with additional columns from table pools. 
This process is somewhat similar to the join operation commonly executed in databases. 
To be specific, given a table pool 𝕋 containing m tables {T_i}^m_i=1, the task of schema augmentation  is to search 𝕋 and find the columns T_i[:, j] joinable to the query table T^O. 
As a result, the result table T^R =  (T^O, , 𝕋) can boost the downstream ML task.
To determine the joinability between two columns, the column values, column semantics, and table structure are the most considered three perspectives.
The corresponding subcategories are introduced below.



Value-based joins. 
Early approaches <cit.> for schema augmentation primarily focus on value-based join, where only exactly matching value can be joined <cit.> (see Fig. <ref> (a)). 
For example,  <cit.> formulates the join problem as an overlap set similarity search by treating columns as sets and matching values as intersections between sets.
 <cit.> further improves  by using intersection estimation to reduce the cost of set reads and the top-k set search.  starts by calculating the cost of reading posting lists and sets using statistical approximation techniques. It then uses an adaptive algorithm that switches between reading posting lists to gather candidates and reading sets to compute exact intersection sizes.

Semantic-based joins. 
Recent studies <cit.> have explored the concept of semantic joins, which join those columns with similar meanings (see Fig. <ref> (b)) instead of those ones having identical values. 
This approach is capable of handling misspellings and formatting variations, resulting in larger set of join results.
For instance,  <cit.> performs joins on soft keys without requiring exact match. However,  only consider one-hop semantic join (e.g., customer  order, customer  product), while  <cit.> and  <cit.> consider multi-hop semantic join (e.g., customer  order  product).
 <cit.> targets the case that columns are embedded into high-dimensional vectors and they are joined based on similarity predicates. 
 <cit.> enhances the capabilities of  <cit.> by employing PLMs (e.g., BERT) as the column vector encoder. Meanwhile,  can perform both value- and semantic-based joins for textual columns, breaking the drawback of previous methods <cit.> that could only handle one type of join and surpassing these works in performance. However,  can only handle textual columns with a relatively small cardinality.
 <cit.> also detects both value- and semantic-based join between columns.  incorporates Graph Neural Networks for similarity signals propagation to better capture column-pair similarity. Meanwhile,  adopts a self-supervised learning approach by generating positive and negative join examples from the table pool, eliminating the need for large amounts of labeled data.

Structure-based joins. Most recent works tend to incorporate the table structure information in the column embeddings, such as using graphs to capture the table structure as illustrated in Fig. <ref> (c).
For example,  <cit.> first constructs a hypergraph capturing the table structure. The hypergraph contains one type of nodes representing columns.  then derives column embeddings from the column nodes containing the table structure information. 
 <cit.> also uses a graph to capture the table structure, but the graph contains the information from the entire database with unique values in the table pool as nodes. 
Similarly, Bharadwaj et al. <cit.> also construct a database-level graph, with columns as nodes and the relationships between columns as edges.





Schema augmentation extends the original table with more features, thereby improving the downstream ML tasks. 
This area has evolved from traditional value-based joins, which depend on exact value matches, to more sophisticated semantic- and structure-based joins.
Semantic-based joins have seen a trend of leveraging generative AI (e.g., PLMs) as encoders to capture the inherent semantics and contextual relationships within and across table columns. 
However, due to token length limitations, these methods are often restricted to handling table columns with relatively small cardinalities. 
When dealing with columns that have high cardinalities, additional sampling or chunking steps are required, which can lead to potential information loss and reduced join performance.
In structure-based joins, most methods utilize graphs to capture the table structure. It has been observed that more complex graph structures can better represent the nuanced aspects of the table schema and relationships.
However, this increased complexity often requires more resources. Therefore, it is important to explore a balanced approach that finds a careful equilibrium between the level of structural detail captured and the computational resources needed. 




       





  §.§.§ Cell Completion ()


As with previous work <cit.>, we refer to retrieval-based TDA at the cell level as cell completion.
This task involves filling in empty cells within the input table by leveraging information extracted from table pools. 
Cell completion can be divided into several subtasks based on cell types: populating attribute names, adding additional entity IDs/names, and filling values for table cells <cit.>. 
The approaches per category are introduced as follows.

Attribute name completion. 
This task is to populate the input table with additional possible column headers or labels. As shown in Fig. <ref> (a), the original song-related table T^O is augmented with a new column header “Singer” retrieved from a related table T^C_2 in the table pool, involving new features for ML.
As a typical study, Yakout et al. <cit.> search for similar tables and then match column labels. Their proposed approach for matching column labels utilize additional information, including similarities that are based on context, attribute names, and column values, respectively. 
 <cit.> ranks a list of labels to be added as headings of new columns by using probabilistic models. The ranking is based on information obtained from a knowledge base and similar tables retrieved from a table corpus. 
A recent approach,  <cit.>, addresses this problem using a retrieval-augmented self-trained transformer model. Specifically,  first indexes and searches tables from the table pool using a bi-encoder retrieval model, and then identifies augmentations from retrieved tables using a reader transformer.

Entity ID/name completion. 
Entity ID/name completion refers to populate the key column (i.e., the entity's ID or name) for a row (an entity). For example in Fig. <ref> (b), this procedure augments T^O with a new entity name “smile like you mean it” (a song) retrieved from the table pool table T^C_1, thereby enriching the potential samples.
Zhang and Balog <cit.> populate rows with additional entities using a two-step approach: (1) candidate entity selection: search for table pool tables that contain the same or similar entities in the original table or search for KB entities with similar KB labels as the original table entities; and (2) entity ranking: leverage a probabilistic formulation based on Bayes' theorem to calculate the relatedness between entities. 
 <cit.> tackles the entity ID/name completion in a similar manner to its attribute name population, using a retrieval-augmented strategy. This strategy pretrains the retrieval-based model by randomly removing entities from the corpus and then reconstructing the removed entities, aiming to better capture semantics and structure of tables.

Cell value completion. 
In this task of finding values for empty data cells, models are designed to estimate a specific value to fill in based on the information from the table pool. Referring to Fig. <ref> (c), this procedure fills in the empty cell in T^O with a new value “songs about jane” retrieved from the table pool table T^C_1 to avoid null values in learning tasks.
A common approach is to retrieve tables from the table corpus and then extract values from those tables <cit.>. 
Zhang and Balog <cit.> build upon this approach such that they consider integrating supplementary information from a knowledge base. 
Ahmadov et al. <cit.> further improve these approaches <cit.> by introducing a hybrid method that combines a retrieval-based method with a
generation-based method (a value prediction model). The prediction model is implemented using a black box approach that can automatically choose the best ML model (e.g., kNN) and the corresponding parameters.
More recently, Deng et al. <cit.> propose the pre-trained  model, with cell value completion as one of its downstream tasks. 
Specifically, they pretrain a structure-aware transformer encoder with table pool tables to model the row-column structure by using a new Masked Entity Recovery (MER) objective. 
Likewise,  <cit.> addresses the cell value completion problem with a transformer model. 
However, it incorporates a reader-or-selection component, that reads the retrieved table and selects the most related ones. This component is based on an extractive approach, ensuring that the model's predictions are always based on existing data rather than speculative assumptions.


Missing values are common and typically have a negative impact in practice, making cell completion crucial. 
This field has been actively studied for decades, and multiple works <cit.> consider addressing various cell types simultaneously. 
For attribute name and entity ID/name completion, an internal problem is that additional cell completion is required after these two operations to obtain a complete table, which may lead to problem propagation. 
Another issue is that retrieval-based methods may sometimes be outperformed by simple statistical techniques like averaging, especially in domain-specific tasks.
Consequently, given the successful application of RAG in the NLP field, there is an opportunity to combine retrieval-based and generation-based methods <cit.>, as generative approaches can better capture the inherent structure and patterns within tables. 
Future advancements could involve integrating generative AI models, such as language models, into this hybrid approach. 




  §.§.§ Table Integration ()


We define retrieval-based TDA at the table level as table integration, the procedure of extending the original table with both rows and columns from related tables retrieved from the table pool via table search algorithms. 
There are two main methods for implementing table integration: one is compositional table integration, which combines retrieval-based TDA results at various levels (rows, columns, and/or cells), while the other is direct table integration, which directly enriches the original table with the content from the retrieved related tables.

Compositional table integration. 
An early work  <cit.> presents a holistic augmentation framework that can perform column and/or row augmentation simultaneously. 
To be more specific, they introduce three core operations: entity augmentation by attribute name, entity augmentation by example, and attribute discovery. 
 <cit.> targets two tasks: augmenting rows with additional entities and augmenting columns with new headers. 
More recently,  <cit.> integrates three TDA tasks: row population, column population, and cell completion.  functions as an end-to-end model that initially retrieves related tables and then extracts various table elements from these tables to perform different levels of TDA tasks. 

Direct table integration. 
More recently, the concept of direct table integration was proposed by <cit.>, which aims to find the right operators (e.g., join, nest, group, link, and twist) to integrate tabular data into a desired form. 
A recent work <cit.> has implemented this approach, proposing a model that extends the original table with both rows and columns simultaneously. They integrate two tables using Full Disjunction, which first connects two tables through an outer-join and then eliminates redundant rows. 
 <cit.> takes a different approach by integrating tables in the latent space. 
To be specific, they construct a data-lake-level graph, and the embedding of this graph is used to featurize the original table.


Table integration is a relatively new concept compared to other levels of retrieval-based TDA tasks, thus it still has a vast range of unexplored opportunities for further research and development. 
For compositional approaches that concatenate combining different levels of TDA, there is a risk of error propagation. 
The compounding of potential errors or biases from the individual TDA tasks can diminish the overall effectiveness and reliability of the integrated table.
For direct approaches, there are relatively few works so far, possibly due to the lack of integration benchmarks. 
Most existing works rely on strong assumptions and are only applicable to small table pools. 
Researchers and practitioners may need to develop more robust and scalable integration techniques, as well as establish comprehensive benchmarking platforms.




 §.§ Generation-based TDA



Generation-based TDA refers to the augmentation of tabular data through the generation of synthetic data. Unlike retrieval-based methods, generation-based methods do not require external data sources and are often built upon generative models. Generation-based TDA tasks can be further categorized into the following sub-tasks: Record Generation () at the  level, Feature Construction () at the  level, Cell Imputation () at the  level, and Table Synthesis () at the  level. These generation-based TDA tasks will be introduced and discussed in more detail in the following sections.



  §.§.§ Record Generation ()


Record generation aims at generating additional records (table rows) based on the original table and its associated information. 
In many practical scenarios, publicly released tables often contain only a small subset of the total available records due to various concerns such as legal constraints or privacy issues. 
They are referred to as sub-tables in this context.
As a result, ML models trained on these limited sub-tables may suffer from suboptimal performance. 
Therefore, it is crucial to generate more synthetic records from the released sub-table. 
Based on the distribution statuses of the original and result tables, the goals of record generation approaches differ: one approach focuses on preserving the original distribution, while the other aims to address imbalanced tabular data through oversampling.


Distribution-preserving record generation ensures that the generated records maintain the same distribution as the original table. 
For example, early works use statistical approaches (e.g., Bayesian networks <cit.> and Fourier decomposition <cit.>) to model the distribution of the original table and then generate synthetic records by sampling from the distribution.
Recently,  <cit.> leverages generative adversarial networks (GAN) to generate synthetic records that are statistically similar to distribution of the original table. Specifically, 's framework comprises three neural networks: a generator, a discriminator, and a classifier, which collectively enhance the semantic coherence of the synthetic records.
Similar to , several other works such as  <cit.>,  <cit.> and  <cit.> also employ GAN to generate distribution-preserving records.  introduces differential privacy guarantees for privacy concern;  <cit.> further maintains functional dependencies to capture the relationships between attributes; while  addresses the interpretation limitation of previous GAN-based methods and further consider explicit feature interactions.
Diffsion models have also been applied for record generation.
 <cit.> adopts a score-based generative model that uses the reverse diffusion process to generate records based on the score function aiming at enhancing sampling quality and diversity.
 <cit.> leverages two diffusion models to process continuous and discrete columns separately and then co-evolve the two diffusion models by transforming conditions from and to each other during training.
 <cit.> use diffusion models to generate tuples that not only cater to the original distribution but also meet specific conditions, such as satisfying a particular criterion for a given attribute. In particular,  first trains an unconditional generative model (diffuser module) to capture the overall distribution of the original tabular data. Then, it uses controllers to measure how well the synthetic data matches the condition and guide the diffuser module to generate data that better satisfies the condition.
 <cit.> further considers the relational structure of the original table when generating records. 
It first learns an approximate relational structure through the construction of a graph, which serves as the foundation of the generative modeling. Subsequently, it uses a variational autoencoders (VAE) architecture to gradually generate synthetic records similar to the original table through message passing over the constructed graph.
Most recently,  <cit.> explores LLMs for record generation. The LLM undergoes a two-stage fine-tuning process, one for capturing table formats, and the other for learning the feature distributions and dependencies.
 

Class-imbalance-aware record generation particularly focuses on oversampling imbalanced samples/instances (rows) within the tabular data, since class imbalance can significantly hinder the predictive performance of classification models.
For example,  <cit.> deals with the imbalanced tabular data by using a conditional generator that can condition on one of the discrete attributes.
It also uses a training-by-sampling strategy, which ensures even sampling from the discrete attributes.
Similarly,  <cit.> first estimates the underlying distribution of the original table using Wasserstein GAN, a derivate of GAN optimizing GAN's training process, and then uses the trained generator to synthesize additional samples of the minority class.
 <cit.> employs a sequence-to-sequence recurrent neural network (RNN) to generate minority class instances. The RNN is trained on the minority class instances to learn their data distribution, and then used to generate synthetic instances to augment the original dataset and balance the minority class.
 <cit.> proposes a transfer-based oversampling technique.  leverages a score-based generative model to transform majority class records to those “fake” minority class records.


Record generation has a broad range of applications in ML, such as increasing sample size, oversampling imbalanced classes, and preserving privacy. This field has evolved from using statistical models to incorporating deep neural network models.
Despite the progress made, a key challenge remains: how to effectively capture the complex attribute correlations present in the original tabular data. The synthetic records generated should not only reflect the overall statistical distribution of the source table but also accurately preserve the intricate relationships between different columns or features. 
Since the database community has extensively studied attribute correlations, integrating database techniques with newly proposed neural network models could be a promising direction. 
Additionally, tabular data often comes with a limited number of samples, making it challenging to train neural network models effectively and leading to a risk of overfitting.
Techniques such as employing regularization methods (e.g.,  <cit.>) and adopting self-supervised learning should be further explored to address this issue.





  §.§.§ Feature Construction ()


Feature construction refers to the task of generating additional columns, also known as features, to enhance the original table. 
Adequate features are essential for training ML models to achieve optimal performance, though good features alone are not always sufficient <cit.>. 
Consequently, numerous studies have focused on feature construction. Based on whether features are transformed into vectors and manipulated in hidden layers, feature construction methods can be categorized into two main groups: explicit and implicit.



Explicit feature construction generates new features by directly manipulating existing ones. For example, it can partition continuous columns into segments, as shown in Fig. <ref> (a), where the “Age” column is divided into [0, 18], [19, 50], and [51, 80].
For instance, Kanter and Veeramachaneni <cit.> use existing raw features and apply a variety of mathematical transformations (e.g., AVG, MIN, and MAX) to create a hierarchy of synthesized features that express more complex relationships in the data. 
Similarly,  <cit.> generates new features by applying predefined operators, including unary, binary and higher-order operators, to existing features. 
More recently,  <cit.> uses Foundation Models (FMs) to construct new features. It uses a prompt strategy to instruct the FMs to generate transformation functions for generating diverse and meaningful features. 

Implicit feature construction generates new features by indirectly manipulating existing feature vectors in hidden layers, as illustrated in Fig. <ref> (b). 
For example,  <cit.> starts by training an encoder to map features into vectors, and then uses an evaluator to optimize the feature vector along the gradient direction. 
Finally, a decoder generates optimal feature subsets.
Wang et al. <cit.> combine the explicit and implicit methods by encoding feature transformation operation sequences into embedding vectors. 
They modify  by replacing the original feature vectors with these operation sequence vectors to disclose a new feature space with discriminative patterns. 



ML is heavily based on high-quality features, which makes feature construction of great importance. 
For explicit methods, previous work typically relies on predefined operators that are hard to refine and often generate meaningless features. 
Recently,  <cit.> has begun to use foundation models for feature construction, indicating a potential future direction.
For implicit methods, manipulating existing feature vectors in hidden layers lacks interpretability. 
Integrating explicit methods, such as  <cit.> may be a potential solution.
Meanwhile, since implicit methods involve encoding features into vectors, using LLMs for feature encoding to better capture feature semantics is also a promising direction.





  §.§.§ Cell Imputation ()


Cell imputation is the procedure of generating assumed values to replace the unknown or missing values within the tabular data. 
Estimating such unseen values within the dataset is particularly challenging due to the high heterogeneity in data types and the large size of datasets, which often comprise millions of rows. 
Cell imputation methods can be divided into two main categories: statistical cell imputation and deep-learning-based cell imputation.

Statistical cell imputation identifies and computes missing cell values based on the statistical characteristics of the original table. 
For example,  <cit.> creates a statistical model for each variable (i.e., a column in the tabular data setting) with missing values, filling in those missing values iteratively until convergence is reached.
The statistical model can be a simple mean or median, or it can be a more complex statistical model like a regressor. 
 <cit.> trains a random forest on the observed parts of the dataset and then use this trained random forest to predict the missing values. 

Deep-learning-based cell imputation has recently utilized a deep generative network to model tabular data for cell imputation, aiming to better capture the distribution and correlations between attributes in the original table. These methods can be further categorized on the basis of the type of generative model they use:


  * VAE-based methods typically employ an encoder to capture the underlying structure of table and a decoder to generate the imputed value.
For example,  <cit.> handles missing data using the encoder to approximate the posterior of the latent variables given the observed data and the decoder to reconstruct the complete data. 
 <cit.> modifies the VAE architecture to handle missing data by marginalizing out the missing data. To be specific,  uses input dropout to make the encoder rely only on the observed data. Meanwhile,  modifies the VAE decoder to factorize the likelihood into separate components for observed and missing data.


  * GAN-based methods generally employ the generator to impute missing data and the discriminator to distinguish between real and imputed data. 
For instance,  <cit.> and  <cit.> use the generator to impute missing data by analyzing real data vectors, while the discriminator, aided by a hint vector that provides additional information about the missing data, differentiates between observed and imputed components.
However,  and  are theoretically only supported under the Missing Completely at Random (MCAR) mechanism, where the probability of missing data is uniform across all cells regardless of their values. In contrast, the Missing at Random (MAR) mechanism, where the likelihood of missing data depends on the observed values, presents a greater challenge in practice.
 <cit.> addresses this by modeling the conditional distribution of missing values based on the observed data for each pattern of missing data.


  * Diffusion-based methods use a reverse denoising process to learn the distribution of the data and perform the imputation. 
 <cit.> adapts the diffusion model to model the distribution of the missing parts given the observed parts by employing the reverse denoising process.





Cell imputation addresses missing values in tabular data, a common issue in real-world datasets, and has been a longstanding concern for the database community.
Statistical methods rely on analytically formulated techniques that often require human intervention and may not capture complex relationships between attributes. 
Recently, there has been a noticeable shift from these methods to deep-learning-based methods.
However, deep-learning-based methods, while powerful, come with the challenges of more complex optimization and often necessitate fully observed datasets for training.
A promising future direction is the adoption of self-supervised methods, which can leverage partially observed data without requiring complete datasets. Furthermore, some recent work, such as  <cit.>, utilizes an AutoML framework, presenting another promising direction for the field.




  §.§.§ Table Synthesis ()


We define generation-based TDA at the table level as table synthesis, which involves generating both rows and columns based on the original table.
At present, most research focuses on generating only rows, columns, or cells individually.
However, table synthesis as a whole could be a promising direction for future research.
Tables with small size of samples and limited features are common, such as web tables. In this case, a simple and feasible approach is to combine previous works that generate rows or columns separately, but additional steps bring increased resource consumption and error propagation.

Therefore, table synthesis that generates rows and columns simultaneously in one pass is necessary and worth exploring.
Currently, a transformer-based method  <cit.>,
which represents a combination of retrieval- and generation-based TDA.
 encodes tables into tokens to retain knowledge across the table pool, redistributes attention on these tokens to highlight important features, and ultimately synthesizes a new encoded table.
Prompt engineering is a potential direct table synthesis strategy, where crafting appropriate prompts could guide LLMs in generating tables.


Table synthesis aims to enhance the original table by generating both rows and columns, particularly useful when the training dataset has low-quality features and limited samples. 
This emerging field requires further exploration. Directly combining row and column generation methods can lead to error propagation. 
Therefore, more targeted solutions are necessary, such as leveraging large language models.




 §.§ Retrieval vs. Generation in TDA
 


In this section, we summarize and analyze the aforementioned TDA techniques from the perspective of comparing the retrieval-based approaches and generation-based approaches. 
First, using the proposed level-based taxonomy, we will distinguish from the task objectives and methodologies of these two approaches at different levels of granularity in Section <ref>. Then, we will provide a general overview of the pros and cons of retrieval-based approaches and generation-based approaches in Section <ref>, enabling researchers to choose the approach that best matches their tasks and requirements.



  §.§.§ Comparison at Different Levels



We analyze the differences between retrieval-based and generation-based methods at each level:


  * Row-level ( vs ): Retrieval-based methods () focus more on the similarity between the original table and candidate tables in the table pool. This method can introduce new and more diverse samples, but often overlooks the impact of retrieved records on the distribution of the original table. In contrast, generation-based methods () pay attention to the original table (e.g., statistical distribution and the relationships between columns) to ensure that the generated data maintain the logic of the original table. However, these synthetic samples tend to be similar to the original ones and cannot introduce new, previously unseen samples.

  * Column-level ( vs ): Retrieval-based methods () are akin to the database join operations, requiring the specification of the target column in the original table, which can introduce completely new related features. Generation-based methods (), similar to feature engineering, optimize existing features to generate new ones and then select the best among them.

  * Cell-level ( vs ): Retrieval-based methods () target regular empty cells and special cells such as column headers. They identify candidate values by searching for similar tables in the table pool or KB entities and then rank these candidates. This process is repeated for each null value and may not be friendly to often-empty regular cells.
In contrast, generation-based methods () focus on regular empty cells, synthesizing new values by analyzing the statistical distribution and structural information of the original table. This method can handle multiple null values in one pass and can handle special cells with simple modifications.

  * Table-level ( vs ): Retrieval-based methods () that integrate both rows and columns from external tables can easily introduce null values. While generation-based methods () do not produce null values, their credibility is uncertain due to a lack of interpretability and potential model hallucination.




  §.§.§ Summary of Pros and Cons


We summarize the pros and cons of retrieval-based and generation-based TDA, highlighting their common challenges, as illustrated in Table <ref>.

Retrieval-based TDA first retrieves related tables from table pools and then uses the retrieved tables for augmentation. The row-and-column structure and the data abundance make this method uniquely suited for tabular data.
This approach enriches the original table with real external data, improving interpretability and introducing new, related information.
Although studied for over a decade, recent advancements in deep learning, such as PLMs <cit.>, have revitalized the field.
However, challenges remain. The retrieval process requires preprocessing and indexing potentially millions of tables, entailing improvements in efficiency and scalability. 
Additionally, the lack of labeled data in large-scale table pools suggests self-supervised approaches as a future direction. Furthermore, retrieval-based methods often struggle with generalization.

Generation-based TDA generates synthetic data for TDA.
This approach does not require external data sources, eliminating the need for preprocessing and indexing numerous tables, thereby saving time and resources.
Synthetic data generation also offers privacy protection. However, this field is nascent and faces challenges. Current deep generative models are often over-parameterized, leading to overfitting, especially with small tables. Moreover, the process lacks interpretability and may lead to models' hallucination.

In addition to the challenges of the two methods mentioned above, TDA faces some common issues due to the nature of tabular data:
First, capturing the table semantics, such as differentiating the context of “apple” in a technology table versus a fruit table. 
Second, capturing the structural information, including row (column) invariance and relationships between columns (columns only depend on a subset of other columns). 
Additionally, incorporating generative AI models in a more effective and interpretable manner is a future direction for both retrieval- and generation-based TDA approaches. Given the respective pros and cons of retrieval- and generation-based methods, exploring their combination is also worthwhile.













































§ TECHNIQUES IN POST-AUGMENTATION


In this section, we mainly focus on three crucial aspects of post-augmentation: publicly available datasets used for TDA and its assessment (see Section <ref>), policies for evaluating the performance of the augmentation methods (see Section <ref>), and strategies for further optimizing the augmentation module (Section <ref>).



 §.§ TDA Datasets



This section delves into the classic datasets commonly used in TDA work, aiming to assist newcomers to the field. We only consider datasets that have been utilized in multiple TDA works, as these are likely to serve as strong benchmarks for the community. 
Key characteristics of each dataset are summarized in Table <ref>. 
These datasets can be broadly categorized into two main groups: retrieval-based TDA datasets and generation-based TDA datasets. The main difference between these categories lies in the input data requirements for the respective TDA approaches. 
Retrieval-based TDA necessitates the collection of additional external tables along with the original table, whereas generation-based TDA only requires the original table as input.
Nevertheless, training a generative model for TDA might require a substantial quantity of tabular data. Yet, from a post-augmentation viewpoint, the extensive data used for training is not taken into account here.





Retrieval-based TDA datasets usually consist of a table pool with hundreds to thousands of tables.
The earliest examples, Web_Manual and Wiki_Link, originate from the same study <cit.>. 
In the Web_Manual dataset, the researchers use Wikipedia tables as their queries and retrieve 371 Web tables to serve as the target corpus. 
These Web tables are then manually annotated with entities, types, and inter-column relationships. 
In contrast, the Wiki_Link dataset is designed for larger-scale use without extensive human annotation. 
It is created by selecting Wikipedia tables where at least 90% of the cell values were internally linked to entities in Wikipedia. 
While this automated approach leads to a larger dataset, the annotations are limited to only entity information, without the more detailed and accurate annotations found in Web_Manual.
Because of the trade-offs between dataset size and annotation quality, both Web_Manual and Wiki_Link are less frequently used in recent TDA research.
Nargesian et al. <cit.> focus on table union search, an immediate step before entity augmentation, and propose two synthesized datasets TUS Small and TUS Large. 
They identify high-quality base tables (with abundant rows and at least 5 textual columns) from Canadian and UK Open Data and then partition the base tables horizontally and vertically to obtain non-overlapping unionable tables to the base ones.
Similarly, Khatiwada et al. <cit.> use the same dataset synthesis technique as  to create SANTOS Small and SANTOS Large. 
These four datasets are commonly used for recent entity augmentation tasks.
Additionally, open data repositories like the WDC Web table corpus and WikiTables corpus are frequently used. They contain a vast number of tables from various domains, adaptable for different types of retrieval-based TDA tasks (e.g., , , and ).







Generation-based TDA datasets typically do not require a table pool or the collection and annotation of multiple similar tables, as only a single original table is necessary for generation-based TDA. 
Consequently, there are no datasets specifically designed for generation-based TDA tasks. Most generation-based TDA methods utilize multiple original tables from various fields, targeting different tasks such as binary classification, multi-class classification, and regression.
For example,  and  adopt the BTS dataset, which is a single table containing 1 million records of domestic air ticket sales. 
These two methods take a portion of the BTS dataset as input and generate augmented tables. 
The effectiveness of these augmented tables is then assessed through regression tests on the "ticket price" attribute in the original BTS dataset.
Common sources for generation-based TDA datasets are primarily from three public platforms: UCI, Kaggle, and OpenML. These platforms are popular partly due to their large scale and diverse range of tables.


It is natural to observe that the same dataset, when processed differently, can be applied to different TDA tasks.
For example,  <cit.> adapts the WikiTables corpus for entity augmentation, while  <cit.> uses WikiTables corpus for schema augmentation.
 filters out those tables that focus on entities, specifically those where the leftmost column contains unique entities, to create a dataset for the  task. Entities in related tables, such as those containing entities from the original table or having similar captions, are considered candidate entities.
On the other hand,  preprocesses datasets by stipulating the key columns as the ones with the most unique values in each table for subsequent joins.
Another observation is that the current generation-based TDA methods do not have standardized benchmarks. These methods often choose their own datasets, which makes it challenging to evaluate and compare their performance on a shared dataset, thereby hindering the assessment of their effectiveness and differences.

With the development of generative AI, an increasing number of TDA works are incorporating the use of pre-training. This has led to a growing need for the construction of robust pre-training tabular datasets.
An ideal pre-training dataset should possess three main attributes: high quality, large scale, and wide coverage.
Previous datasets often compromise between quality and scale, but combining human expertise with generative models might address this issue.
Wide coverage means curating diverse datasets across multiple domains, such as finance and healthcare, and various data types like CSV, JSON, and Markdown. 
This diversity is essential for pre-training large-scale generative models to handle the variety of tabular data found in real-world scenarios. However, privacy concerns might pose challenges in developing high-quality TDA datasets.
















 §.§ Evaluation Polices



We summarize the common evaluation policies for TDA works. These polices can be categorized into two main groups based on the involvement of models: original-table-based evaluation and model-based evaluation. 
Many studies <cit.> have utilized both evaluation methods simultaneously.


Original-table-based evaluation refers to evaluating the augmented table by comparing it to the groundtruth, which is typically the original table or its derivatives.
For example, several works <cit.> start from a base table and derive a subtable from it to serve as the original table T^O, while the entities or columns outside this subtable are used as the ground truth. This method can efficiently construct the ground truth, but the range of truth values is rather limited, and there may be data related to the base table in the table pool that are not present in the base table (e.g., the base table about IT companies in one country with a table in the table pool about another country).
Several works <cit.> calculate the cumulative distribution functions (CDFs) of the augmented and original tables to compare their statistic similarity. 
While this method is simple and effective, it only evaluates statistical distribution information and cannot capture more complex details, such as relationships between columns (e.g., the connection between "position" and "salary").

Model-based evaluation refers to feeding the augmented table alongside other baseline tables to a specific ML model and then evaluating the model's performance. 
These baseline tables generally include three primary datasets:
(1) None <cit.> refers to the original dataset without any augmentation. 
For example,  <cit.> feeds the augmented and original table to a specific classification model that performs a grid search over RandomForest, AdaBoost, and GradientBoosting, then comparing the corresponding classification results.
(2) Random <cit.> refers to the original dataset augmented with randomly selected candidates. For instance, a schema augmentation work  treats features in a table pool as independent entities and randomly selects a predefined number of features to augment the original table as the Random baseline.
(3) All <cit.> refers to the original dataset augmented with all possible candidates. For example,  involves a ALL baseline that joins the original table with as many tables as possible.
In general, None baseline is suitable for both retrieval- and generation-based TDA while the other two are suitable only for retrieval-based TDA.


The two evaluation methods, original-table-based and model-based, each come with their own set of pros and cons.
The original-table-based evaluation is straightforward and efficient but may lack precision and be limited in scope, especially for specific augmentation needs like enhancing the minority class.
The model-based evaluation offers greater accuracy but requires more time and resources.
Additionally, using models introduces extra layers of uncertainty, such as model characteristics and hyperparameters, making it difficult to control variables and ensure fairness.
This also complicates determining whether issues arise from the model itself or the data augmentation method.
To address these tradeoffs, several studies <cit.> adopt a hybrid approach, using both evaluation methods. This raises an interesting question: do these two evaluation methods yield conflicting results?
It may be worth exploring the development of alternative evaluation methods, such as rule-based TDA evaluation policies, which could potentially strike a balance between efficiency and accuracy better than current methods.









 §.§ Optimization Strategies



Optimization strategies aim to further refine the augmented results based on the performance of specific downstream ML models. These techniques can be categorized into two main types: iteration-based and reinforcement-learning-based.

Iteration-based optimization involves a simple and direct method of using feedback from the downstream ML model to determine whether adding a candidate augmentation enhances the performance of the task.
For instance, Chepurko et al. <cit.> devise the random injection feature selection (RIFS) algorithm, which compares model performance using candidate features against deliberately constructed random features as a baseline.
The objective is to identify and eliminate irrelevant features, finding a subset of features that contain signals relevant to the downstream ML task.
Their subsequent work  <cit.> is an automated system that searches and joins data with the input table end-to-end. 
Similarly,  <cit.> leverages the supervision signal from the downstream ML task to filter out unnecessary information.  uses a graph to capture information from the entire database, including both useful and potentially spurious relationships. 
During training, the downstream ML model will automatically focus on using the valuable information while ignoring or downweighting the non-useful parts. 
More recently,  <cit.> assesses the actual utility gain of candidate features by running ML model on the original and augmented tables; it then iteratively selects features based on their  utility gain and the feasibility of the join path.

Reinforcement-learning-based optimization employs reinforcement learning (RL) to explore the features that improve the performance of the ML model. 
For example, Liu et al. <cit.> propose an RL-based automatic data search system that retrieves fresh training data from table pools and interacts with the downstream ML model. 
In this RL-based framework, each training data point and its corresponding influence score, calculated by the “Environment”, serve as the “State”. Given the “State”, the “Agent” composed of a Search-Policy selects the optimal “Action” (a set of training data points retrieved from the table pool), and feeds them back into the “Environment”.
 <cit.> makes further improvement by not only exploring features that boost performance but also utilizing rarely selected ones to avoid local optima.
Chai et al. <cit.> further extend  by broadening the scope of table pools, such as enterprise data warehouses, online repositories, and data markets.


Optimization strategies have proven effective in improving downstream ML task performance, with a noticeable trend towards the use of RL.
However, they often come with computational overhead, which can be a constraint in time-sensitive or resource-limited situations and can impact scalability.
As such, a worthwhile research direction is to improve the iteration efficiency of these optimization strategies while maintaining their effectiveness. This could involve reducing the number of required iterations or exploring the use of lighter surrogate models in place of the original, specific ML models.




§ TRENDS AND OPPORTUNITIES



This section examines the current landscape and future prospects of TDA techniques. 
We first explore emerging trends shaping the field (Section <ref>), followed by a discussion of promising opportunities for further research and application (Section <ref>). By highlighting these aspects, we aim to inspire continued interest in this vital area of study.



 §.§ Major Trends in TDA Development



Based on our review of numerous research articles and our three-step pipeline, we have identified three significant trends in current TDA work:


  T1. Enhanced Table Representation. In the pre-augmentation phase, researchers are increasingly employing more advanced table representations to better capture both structural and semantic details within tables;

  T2. Confluence of Retrieval and Generation. In the augmentation phase, retrieval- and generation-based methods offer distinct advantages and disadvantages. A fruitful path forward may lie in the integration of these two methodologies;

  T3. Automated TDA. Looking at the entire TDA pipeline, a promising future direction is the creation of end-to-end automated TDA systems that can efficiently manage the entire TDA process.
We proceed to elaborate on each of these trends.


T1. Enhanced Table Representation.
The representation of tables is crucial for effective TDA. 
By accurately capturing content, semantic, and structural information in table representation, downstream augmentation processes can achieve better outcomes.
Accordingly, there is a clear trend in TDA towards more complex and sophisticated table representations.
Recent TDA research <cit.> includes diverse information such as table context and metadata, unlike earlier work <cit.> that focused solely on table content. 
Additionally, newer TDA approaches <cit.> utilize graph structures to capture relational information within tables more effectively, encoding these structures for improved table representation.
Notably, many of the latest TDA techniques, both retrieval-based methods <cit.> and generation-based methods, utilize <cit.> generative AI (e.g., PLMs) to generate robust table representations, leveraging their semantic understanding and generalization capabilities.
While generative AI models have shown promise in tabular data representation, there remains ample opportunity for further development, especially compared to advancements in NLP and CV.
In conclusion, the move towards more sophisticated and comprehensive table representations is a key trend driving progress in TDA research.

T2. Confluence of Retrieval and Generation. 
As discussed in Section <ref>, retrieval-based and generation-based TDA have their own pros and cons. 
Retrieval-based methods incorporate external data sources for better interpretability but struggle with efficiency and scalability as data increases. Generation-based methods do not use external data, thus lacking interpretability and potentially leading to model hallucination. 
Given the trade-offs between these two approaches, combining their strengths while mitigating their weaknesses should be the future direction for TDA research. This aligns with the prevalent use of RAG models in the field of NLP, where the benefits of both retrieval- and generation-based techniques are harnessed. By blending these complementary methodologies, we can work towards more robust, efficient, and interpretable TDA solutions that can meet the evolving needs of diverse applications and domains.




T3. Automated TDA.
A notable trend in TDA is the move towards automating the entire process, creating end-to-end platforms that integrate various operators.
For example, Chepurko et al. <cit.> propose , an end-to-end system that integrates multiple operators (e.g., imputation, hyperparameters optimization, feature selection, etc.) for automatic schema augmentation.
Similarly,  <cit.> automates cell imputation by proposing an AutoML framework that utilize search algorithms to automatically select candidate ML models.
Despite these advancements, current solutions often focus on single subtasks of TDA, leaving room for further exploration and development.
An area for improvement is to consider the TDA pipeline as a whole, integrating various pre-augmentation and augmentation operators, and intelligently selecting suitable operators based on user requirements. Additionally, refining the automated TDA pipeline could involve an iterative step to identify and eliminate unnecessary or even harmful steps, based on augmentation evaluation and data characteristics.
By advancing towards more holistic, end-to-end TDA automation, the field can unlock greater efficiency, scalability, and customization, ultimately enhancing the practical value and applicability of TDA techniques.







 §.§ Emerging Opportunities for TDA



As we navigate the era of big data and the trends towards generative AI and autoML, we observe that obtaining high-quality data from massive amounts of data to facilitate generate AI is imperative. TDA is an important sub task of data quality and there is still a long way to go.
Firstly, from a data perspective, the quantity and complexity of tabular data have significantly increased. Tables now typically have multiple patterns <cit.> and contain millions of samples <cit.>, posing new challenges to TDA work.
Secondly, from a model perspective, ML models, especially generative AI, have always faced interpretability issues and potential of privacy leakage.
Below, we will elaborate on these key opportunities for further advancements in this field.

O1. Multimodal TDA.
Current TDA works often assume that tables contain only textual and numerical values. However, the reality is that modern tables often encompass a broader range of modalities, such as images <cit.>. 
Handling these multimodal tables presents unique challenges that current table processing techniques may not adequately address. 
The representation and indexing of such heterogeneous tables, for instance, may require fundamentally different approaches compared to traditional text-based and numerical tables.
Additionally, user needs might be expressed in modalities other than the tabular data itself <cit.>.
For instance, a user might use natural language to request the oversampling of a minor class. 
Addressing these multimodal user inputs and aligning them with TDA operations is another key challenge to tackle.



O2. Efficiency and Scalability.
Another key challenge for TDA is the issue of efficiency and scalability. 
Retrieval-based TDA approaches often involve similarity comparisons at the table pool level, which can encompass millions of tables or more. 
Despite that, most existing methods <cit.> are exact algorithms with a worst-case time complexity that is linear in relation to the product of the query column size and the table repository size, raising concerns about their scalability.
This makes it a crucial research direction to explore.
Meanwhile, the table itself may contain tens of thousands of records[Large-scale tables do not necessarily indicate high quality or no need for augmentation (e.g., many tables often have many records but few features for ML model training). Similarly, candidate tables in the table pool may also be large-scale.].
Moreover, both retrieval- and generation-based TDA methods are embracing the large-scale generative AI models such as PLMs and LLMs. 
The computational and memory requirements of training, fine-tuning, and deploying these models can hinder the efficiency and scalability of TDA techniques in practice. 
Addressing the scalability challenges associated with large-scale generative AI models in TDA is another crucial area of research. 
Potential solutions may include the development of more efficient model architectures the exploration of model compression <cit.> and distillation techniques <cit.>.



O3. Domain-specific Tasks.
Another promising direction for TDA research is the exploration of domain-specific applications. 
Domain-specific data often exhibit strong professionalism and have some unique characteristics <cit.>.
For instance, medical data typically possesses specialized terminology and intricate data distributions that differ from more general tabular data. 
Developing TDA techniques tailored to these domain-specific characteristics could result in more effective and domain-friendly TDA strategies.
This might involve incorporating domain knowledge, such as expertise from professionals <cit.>, knowledge base <cit.> and knowledge graphs <cit.>, and utilizing specialized table representations <cit.>.
Combining retrieval-based and generation-based TDA holds promise for domain-specific tasks, allowing for extensive use of limited domain knowledge.
An interesting question arises as to whether the knowledge and techniques developed for TDA in one domain can be efficiently transferred to other domains. 
For example, can the insights and models derived from financial data augmentation be effectively applied to medical data, or vice versa?
The ability to successfully transfer TDA across domains would be highly valuable, as it could accelerate research and development efforts, enable the reuse of existing resources, and promote the exchange of ideas.


O4. Interpretability.
Existing TDA works face a notable challenge with respect to interpretability, particularly in generation-based TDA approaches.
Retrieval-based TDA, on the other hand, has demonstrated better interpretability, as it can be interpreted as leveraging the retrieved relevant table information to augment the original data.
One promising future direction for TDA research would be to explore the combination of retrieval- and generation-based TDA approaches.
This hybrid method could improve the interpretability of generation-based TDA while also harness the respective advantages of both approaches.
At the same time, the recent trend towards the incorporation of deep neural networks <cit.>, such as language models, into TDA workflows has introduced a new challenge regarding interpretability. 
The complex, black-box nature of these models can make it difficult to understand and explain the underlying rationale behind the augmented data.
Addressing this interpretability challenge is crucial, as it impacts the trust, transparency, and practical adoption of TDA methods, particularly in domains where explainability is of paramount importance, such as high-stakes decision-making <cit.> or regulated industries <cit.>.

O5. Privacy and Security.
Last but not least, privacy and security represent critical opportunities for TDA. 
For retrieval-based TDA approaches, they carry out similarity comparisons between the original table and the tables in the pool, which may lead to information leakage of the user's original table. In contrast, generation-based TDA has inherent advantages in privacy protection, and several works <cit.> have utilized record genration for privacy-preserving data publishing.
However, the recent trend towards the incorporation of LLMs into TDA workflows introduces a new set of privacy and security challenges. While LLMs have powerful generative capabilities, they also carry the risk of information leakage <cit.>, potentially exposing or memorizing sensitive data from their training corpora.
Addressing the offensive and defensive aspects of LLMs behaviors, including the development of robust privacy-preserving training techniques <cit.>, remains an active area of investigation.



























































§ CONCLUSION



This survey presents a thorough investigation of tabular data augmentation (TDA) for ML, with a particular emphasis on the recent advancements in leveraging prevalent generative AI models.
Our work meticulously outlines the essential steps involved in TDA by constructing an end-to-end pipeline encompassing three critical procedures: (1) pre-augmentation, where we summarize and analyze the commonly used preparation techniques for TDA; (2) augmentation, where we systematically compare current TDA techniques, including both retrieval-based and generation-based approaches; and (3) post-augmentation, where we delve into the evaluation and optimization processes following TDA. Additionally, we provide a comprehensive analysis of the pros and cons of current methodologies and outline future trends and opportunities for TDA.

The era of generative AI heralds a transformational phase for TDA. ML on tabular data is ubiquitous and demands a substantial amount of high-quality data — a requirement that generative AI can significantly enhance. Despite the distinct characteristics of tabular data, generative AI models have predominantly been applied to fields like computer vision and natural language processing, with their application to tabular data still in its nascent stages. This leaves a vast landscape for further innovation and advancements in TDA through the use of generative AI.
Our comprehensive survey aims to bridge this gap by providing a detailed, systematic overview of the current state of TDA and its potential future directions. We believe that this work can contribute to the community and serve as a valuable resource for researchers and practitioners alike.
Our endeavor will be presented as a continuously updated literature repository, maintained online at .



§ ACKNOWLEDGMENTS

The authors would like to thank Yifan Wu, Jianfeng Zhang, Ankai Hao, Zhaoyi Yuan, Jiahui Long, and Hongwei Yuan for their invaluable feedback, which significantly improved the manuscript.






ACM-Reference-Format




 Paper Title: Block-Operations: Using Modular Routing to Improve Compositional Generalization.




Transmission of quantum information through quantum fields in curved spacetimes
    Erickson Tjoa
    August 5, 2024
===============================================================================





§ ABSTRACT 


We explore the hypothesis that poor compositional generalization in neural networks is caused by difficulties with learning effective routing.
To solve this problem, we propose the concept of block-operations, which is based on splitting all activation tensors in the network into uniformly sized blocks and using an inductive bias to encourage modular routing and modification of these blocks.
Based on this concept we introduce the Multiplexer, a new architectural component that enhances the Feed Forward Neural Network (FNN).
We experimentally confirm that Multiplexers exhibit strong compositional generalization.
On both a synthetic and a realistic task our model was able to learn the underlying process behind the task, whereas both FNNs and Transformers were only able to learn heuristic approximations.
We propose as future work to use the principles of block-operations to improve other existing architectures.










§ INTRODUCTION




Typical artificial neural networks (NNs) perform poorly on tasks of systematic generalization <cit.>.
Numerous people have argued that this failure to generalize is caused by Neural Networks' lack of compositionality: The ability to break complex concepts down into their atomic elements and to reuse and recombine them appropriately when faced with a new task <cit.>.



In an analysis drawing inspiration from neuroscience and cognitive psychology <cit.> formalized this as the Symbol Binding Problem: The problem of representing the world in terms of symbol-like entities and to dynamically bind information that is distributed throughout the network to those symbols.
They highlighted that routing in regular neural networks tends to be static, becoming fixed at training time.











<cit.> investigated compositionality in neural networks empirically and found that



neural networks tend to learn new feature mappings instead of representation-preserving mappings suitable for routing, even in situations where a consistent representation would clearly be beneficial.








Based on these theoretical analyses and empirical observations, we develop the concept of block-operations in Section <ref>.
Block-operations are a new approach to tackle compositional generalization that is based on building a favorable inductive bias into the network:
We first split all network activation tensors into uniformly sized blocks.
We then modify network modules to implement Modular Representation-Preserving Mappings (MRPMs), a new method that encourages the network to route and modify blocks dynamically and independently of each other.
As we later show, this causes the blocks to act as objects that are used consistently throughout the network and that can be modularly recomposed by changing the order of blocks within a tensor.


















We then introduce the Multiplexer and SMFR modules (Stack of Multiplexers and Feedforward Neural Networks with gated residual connections) as a first implementation of block-operations in Section <ref>.
The SMFR implements MRPMs and replaces the basic FNN module.

We ran several experiments to verify the effectiveness of our architecture.
In Sections <ref> and <ref>, we demonstrate useful properties of our module on synthetic data and explore the effects model size and architecture on performance and reliability.
In Section <ref>, we show that our module is particularly effective at learning logical rules and variable assignments, such as those seen in algorithmic tasks, beating both the FNN and the Transformer.
We show that the SMFR learned the underlying atomic operations of the task, which resulted in perfect generalization, whereas the other models only learned statistical correlations.
In Section  <ref>, we test our module on realistic data: We split MNIST images into four chunks and permute those, creating one new task per permutation. The resulting dataset requires pattern recognition to solve just like MNIST, but can be solved more efficiently if the network understands how to undo the permutations.
We found that the SMFR took longer than the FNN or the Transformer to learn this task, but it kept improving in test accuracy and generalization even after training accuracy converged. At longer runtimes, it ended up outperforming the other architectures.
Inspecting the models' internals showed that the SMFR correctly learned to decompose the task into its two base tasks: Disentangling the permutation and solving MNIST.







Our contributions are:

    
  * We propose the concept of block-operations, a routing-based approach for making existing neural network architectures better at compositional generalization.
    
  * We apply this concept to create the SMFR module as a replacement for the FNN.
    
  * We empirically demonstrate that SMFRs have the properties we expected and that they outperform both FNNs and Transformers on tasks designed to test compositional generalization.
    
  * We note several remaining limitations and motivate further research on block-operations, which could be used to enhance existing architectures such as Transformers.





§ RELATED WORK




Residual Connections.
Residual connections allow a neural network to route data through the network unchanged, and many variants of this have been used to great success in different areas <cit.>.
The Neural Data Router uses a variant with a learnable interpolation weight and achieved very strong generalization <cit.>.
We use this variant for our own architecture.



Attention.
Attention mechanisms have become a ubiquitous tool in deep learning, useful in a variety of fields such as natural language processing, image processing and speech recognition <cit.>.

We note that the term Attention is not used consistently in the literature and is often conflated with the more specific Self-Attention that is used in Transformers.


Attention in Transformers is based on pairwise comparisons. This makes it difficult for them to learn routing mechanisms based on more than two features, such as conditional statements. We experimentally confirm this in Section <ref>.

The Multiplexer module we introduce in this paper uses a form of Attention, but not Self-Attention. 





Routing.
Many architectures for routing exist and have been summarized by surveys <cit.>.

For a comprehensive overview of related work, we recommend reading <cit.>, especially the paragraph on Instance Slots on page 13, which is the type of network most similar to our own.










Recurrent Independent Mechanisms use multiple largely independent cells that communicate through a bottleneck of attention and compete for access to input data <cit.>. The RIM uses multiple modules with their own independent training dynamics, whereas our approach is simpler and uses an inductive bias to allow subnetworks to more easily share data in a consistent format.


Capsule Networks can perform routing between a fixed number of capsules based on entity recognition <cit.>. They are designed for modeling hierarchical relationships, but have a fundamental design difference to our approach: Each capsule represents a specific type of object, while each block in block-operations is a placeholder variable that can come to hold any type of object.

Slot Attention <cit.> comes closer to block-operations as their slots are similar to our blocks: Each slot can hold a different object. This architecture is good at extracting object-centric representations, but does not contain a routing mechanism.

Transformers with Competitive Independent Mechanisms (TIM) split data into blocks, which they call positions, and operate on them independently <cit.>. However, they do not include a way to preserve mappings while routing them past a layer of mechanisms, to a subsequent layer of the network.
They hypothesize, but do not test, a way to access information from earlier layers, which would be similar to our blockwise routing. However, this is based on Self-Attention and would therefore suffer from the same issues as normal Transformers (see Section <ref>).





§ METHODOLOGY: BLOCK-OPERATIONS




In this section we discuss the Symbol Binding Problem and introduce the concept of block-operations as our approach to solve it.

<cit.> decomposed the binding problem into three different aspects:


Representation problem: Representing multiple objects separately, in a common format, yet without interference between them.


Segregation problem: The ability to form object representations from unstructured inputs.


Composition problem: The ability to dynamically relate and compose these object representations with each other.














Most approaches based on the instance slots described by <cit.> rely on weight sharing to encourage consistent representations.

Inspired by the problems identified by <cit.>, we take a different approach:

We treat routing as the core of the problem and develop an inductive bias to make learning it as simple as possible.
We hypothesize that if the different components of a neural network can easily learn to share objects with each other, then they will take the path of least resistance and converge towards using a common format even without explicit weight sharing.


However, we also have to take into account the problem of the superposition catastrophe, described by <cit.>: If object representations encoded in a tensor do not have distinguishable roles, it becomes ambiguous which attributes belong to which object.
For example, the concepts (red, apple, green, pear), when naively encoded in a single tensor, could either refer to "red apple, green pear" or to "red pear, green apple".






To deal with this problem, block-operations combine two ideas: We split all network activation tensors into uniformly sized blocks, and we require all modules in the network to implement Modular Representation-Preserving Mappings (MRPMs) to operate on these blocks.


Blocks.
Through the inductive biases provided by MRPMs, each block of a tensor should come to hold a separate object and the position of the block within the tensor should designate its role.
If one imagines a neural network module as a function, then the blocks in its input tensor would be a list of discrete function arguments.
This discretization mirrors the way humans design algorithms and naturally lends itself to modular decomposition, as demonstrated by our experiments in Section <ref>.


MRPMs.
It is easiest to think of MRPMs as an enhancement to an existing module that adds optional routing abilities to it.
To implement MRPMs, a module must be able to make any of its output blocks a copy of any of its input blocks, and it must be able to decide to do this dynamically, based on its inputs.
One way to implement this is to use attention mechanisms or dynamic residual connections to interpolate the module's normal outputs with the input blocks.

MRPMs allow different submodules of the network to easily share data with each other in a representation-preserving manner without hampering their ability to generate new data: The modules can simply use some of their output blocks for their own calculations while using other blocks to pass their inputs through.
This easy sharing of data makes it likely that blocks in different parts of the network that represent the same concepts will use the same representations, even if different subnetwork have only encountered different subsets of the concept. Our experiment in Section <ref> demonstrates this empirically.


Implementation.
In Section <ref> we create the SMFR module as a replacement for the FNN that satisfies the MRPM criteria. It can seamlessly combine two different ways of processing data: Learning new mappings just like a regular FNN, and conditionally routing blocks of data.
Figure <ref> illustrates the general idea without going into the details.

We refer to Section <ref> for a discussion on adapting other modules and architectures to block-operations.




Tackling the Binding Problem.
Block-operations address all three aspects of the binding problem:


Representation problem: Each block represents one object or feature. Since each feature tensor comprises several blocks, multiple objects can be represented in parallel. The ability to route through MRPMs leads to an inductive bias that encourages all blocks to use the same representation.


Segregation problem: Modules using MRPMs can dynamically decide when to route blocks and when to generate new blocks.


Composition problem: The position of a block within a feature tensor can be changed during routing and influences its semantic meaning: The four-block tensor [red, apple, green, pear] is different from [red, pear, green, apple]. This modularity allows block-operations to avoid the superposition catastrophe.




§ MODULES





We introduce the Multiplexer to route blocks and the FNNR (Feedforward Neural Network with gated residuals) to learn new feature mappings. We combine these two modules into the MFNNR (Multiplexer plus Feedforward Neural Network with gated residuals), which can do both. Finally, we stack multiple MFNNRs in sequence to form our final architecture, the SMFR.


Multiplexer. A Multiplexer (Figure <ref>, left) takes M input blocks and produces N output blocks, each of which is a weighted average of all M input blocks. The weights are generated by an FNN that receives a concatenation of all M blocks as input and outputs an M*N weight matrix, which is normalized by applying a softmax over the first dimension.

The Multiplexer can learn to transfer or linearly interpolate blocks. It makes these routing decisions dynamically, for all blocks at the same time, and takes the order of the blocks into account.



























FNNR. An FNNR (Figure <ref>, right) takes N input blocks and produces N output blocks. It uses an FNN to generate new blocks and then combines them with the input blocks using residual connections.

These residual connections use a learned gating weight instead of simple addition, similar to <cit.>.

An FNNR can learn to either let an input block pass through unchanged, or to replace it with a new block created by an internal FNN, or to create a linear interpolation of both. It can make this decision for each of the blocks independently, but can base each decision on all input blocks.









MFNNR. An MFNNR is composed of a Multiplexer followed by an FNNR module. The FNNR module uses the input blocks of the Multiplexer as its extra input.

The MFNNR can learn to rearrange blocks, or to generate new blocks through an FNN, or to interpolate between both. It can learn to do this conditionally on the input and with different rules for each output block.
An MFNNR can emulate an FNN or a data copying mechanism by setting its gating weights to extreme values.
It can learn arbitrary mappings like an FNN, but it also has an inductive bias to copy or interpolate blocks of data if doing so leads to a simpler solution.









SMFR. Multiple MFNNR modules can be stacked one after the other to form a more powerful architecture.

The SMFR is the architecture we use in our experiments, as an alternative to FNNs.




§ EXPERIMENTS





Compositional generalization has multiple different aspects, making it difficult to pick a benchmark for it.
A meta analysis comparing different benchmarks on compositional generalization found little concurrence between them <cit.>.
We therefore chose a combination of specific synthetic tasks, to test the theory, and practical tasks, to confirm our findings against realistic data and strong architectures. We used grid-search to generate different architectures to compare (see appendix).

The first two of our experiments are based on the experiments used by <cit.> to test a network's ability to specialize and to reuse subnetworks.
As the SMFR is intended as a replacement for the FNN, we compare SMFRs with FNNs with a similar numbers of parameters.
The experiment in Section <ref> measures negative interference, i.e. the extent to which altering the training distribution causes the network to forget previously gained knowledge <cit.>. It shows that SMFRs are less prone to negative interference than FNNs. The experiment in Section <ref> measures generalization through the reuse of subnetworks. It shows that SMFRs can sometimes learn to reuse subnetworks perfectly even in the absence of training data, a case of transfer learning.

The experiment in Section <ref> measures performance on an algorithmic task that requires conditional logic and variable assignments.
Here we also perform an ablation study by comparing it to a Transformer that operates on blocks, to demonstrate that existing Attention mechanisms can not efficiently emulate our architecture.
Our results show that SMFRs developed a more accurate model of the underlying atomic operations of the task.

The experiment in Section <ref> tests the SMFR's performance on realistic data. We introduce block-permuted MNIST (BPMNIST), a task that can be decomposed into two steps: Learning a subtask-dependent permutation, and solving MNIST.
We find that the SMFR took longer to train than the FNN or the Transformer, but its test accuracy kept increasing even after training accuracy converged, similar to the "grokking" observed by <cit.>.
Unlike the other architectures, the SMFR was able to learn the underlying compositional rule of BPMNIST, allowing it to achieve much better Out-of-Distribution (OOD) accuracy when trained for long enough.





 §.§ Addition/Multiplication Experiments





Task.
The addition/multiplication dataset is designed to test the resilience to negative interference of a neural network. The task is to either add or multiply two single-digit numbers (modulo 10).





We use two training stages. During the preparation stage the multiplication task is trained on limited data and the addition task is trained on all data.
In the limited dataset, one number uses only low digits (0 to 4) and the other number uses only high digits (5 to 9).
Once the network reaches a threshold of accuracy, we switch to the negative interference stage, in which the rule for training is inverted: The addition task is trained on limited data and the multiplication task is trained on all data.

We run the negative interference stage for a fixed number of additional training steps, then measure the preparation-data accuracy: The accuracy on the data that was only used for training in the preparation stage but not the negative interference stage.


Variant.
We note that both the addition and the multiplication task are commutative. Our SMFR architecture is good at solving commutative tasks because the softmax used by the Multiplexer is also commutative, giving it a useful inductive bias. This is an additional strength of our architecture over FNNs. We perform an ablation study to measure how much of the SMFR's strength comes from its block-routing abilities and how much from its ability to easily handle commutativity. To do so, we run variants of the experiment in which we use Straight Through Gumbel-Softmax instead of softmax for the Multiplexer, because Gumbel-Softmax makes a discrete pick among candidates and therefore does not help with commutative tasks in the same way softmax does <cit.>.




























Results.
Table <ref> shows the preparation-data accuracy of different models. Each line compares averages and standard errors of 90 trials of FNN and 63 different architectures of SMFR for each of softmax and Gumbel-Softmax. The Threshold refers to the accuracy at which we switch to the second stage of training. The values show the preparation-data accuracy 2000 steps after switching to the negative interference stage, except for the last column, which shows the ratio between the values of SMFR_Softmax and FNN.
We see that SMFR_Softmax performs best in all cases and the difference to FNNs is larger for smaller thresholds. In other words, SMFRs suffer less negative interference than FNNs, especially if the switch between the two training regimes happens earlier.



Routing and Commutativity.
SMFR_Softmax consistently performs better than SMFR_Gumbel, and both outperform FNNs. This shows that both the block-routing abilities of SMFRs and their ability to learn commutativity efficiently have positive effects.

Model Size.
The above analysis is based on an average over all model sizes for both SMFRs and FNNs. An ablation study showed that SMFRs perform better at smaller model sizes and FNNs at larger ones. The SMFRs were still better than the FNNs at larger model sizes, but the difference was less pronounced.
The only cases where FNNs slightly exceeded the performance of SMFRs were when the number of parameters was high and we also waited until convergence (threshold=1.0) before starting negative interference.
See the appendix for details on this and on architecture optimization.





 §.§ Double-Addition Experiments





Task.
The double-addition dataset is designed to test the ability of a neural network to reuse subnetworks for similar tasks. The network receives two pairs of numbers and has to return either the sum of the first pair or the sum of the second pair (modulo 10).



We want to measure how well the network learns that both tasks require the same logic. To do so, we use a biased training procedure: The distribution of the input numbers is fully uniform for the first pair of numbers, but it is restricted for the second pair of numbers.
This uses the same distributions that we also used in the addition/multiplication experiments: One number uses only low digits (0 to 4) and the other number uses only high digits (5 to 9).
The training data for the second task is therefore a strict subset of the training data for the first task.
During testing, we measure the OOD accuracy on the second pair of numbers.



Variants.
As before, we run two variants of the SMFRs: One using softmax and one using Straight Through Gumbel-Softmax. Using Gumbel resulted in worse performance overall but otherwise showed the same patterns.



Results.
FNNs of all model sizes consistently get an OOD accuracy of 0.0, which is actually worse than guessing (0.1). This is because f(a,b)=a+b becomes a bijective function if either a or b are frozen, mapping ten possible inputs to ten possible outputs. The set of possible outputs of the limited training inputs has no overlap with the set of possible outputs of the OOD inputs, which results in an OOD accuracy of 0.0.
In contrast, SMFRs reach an average of 0.205 OOD accuracy across all architectures we tested. The surprising part here was that some trials achieved 100% OOD accuracy, implying perfect reuse of subnetworks.



Architecture.
We investigated the effect of the architecture on performance. We found that model size correlated with performance, but bigger was not always better (see the appendix for details).
Across all of our experiments with SMFRs, 10.4% achieved 100% OOD accuracy (25 of 240 trials). Best results were achieved when using softmax instead of Gumbel, at a stack depth of exactly 1 and a width of 8 or higher: For these architectures 75% of trials achieved 100% OOD accuracy (9 of 12).

We hypothesize that higher depths make routing more complicated, causing the SMFR to resort back to using FNN logic.









 §.§ Algorithmic Experiments





Task.
The ALGO task is designed to emulate patterns of information processing that typically occur in real-life coding tasks. It tests how good the network is at understanding conditional logic and variable-assignment operations. The task uses five variables as the input and expects the same five variables as the output. The task proceeds in several iterations, modifying one variable per iteration. On each iteration, a formula of the following form is applied:

"Variables A, B, C, D should remain unaltered. Variable E should be assigned  A+1 if C > D and B+1 otherwise."

We use five different permutations of the variables A, B, C, D and E. 


We want to find out if the network learns the actual underlying algorithm, or only a statistical correlate. To do so, we use a special way to train and test the network: During training, we always perform exactly two applications of the rule.
As a result, we can test the network for generalization in two different ways: OOD_even measures the average accuracy when the network is run for 4, 6, or 8 iterations. This is a form of length generalization. In contrast, OOD_odd measures average accuracy for 1, 3, 5, 7, or 9 iterations. Since training happened on 2 iterations, this tests if the network understands that the data generating process can be decomposed into two applications of the same atomic rule.

Baselines.
In addition to the FNN we also compare to a Transformer. The Transformer operates on the blocks, just like the SMFR, treating the variables on each iteration as a sequence of inputs.




Results.
Table <ref> summarizes the results.
OOD_even is high for all architectures, so they all demonstrate good length generalization on this task.
The Transformer performed worse than either the FNN or the SMFR.
We believe this is because Transformers have an inductive bias that is unsuitable for this task: Applying the task's formula requires considering the values of all variables at the same time, but the Key/Query Attention mechanism in Transformers can only perform pairwise comparisons.



While performance on length generalization was quite similar between architectures, OOD_odd was low for both FNNs and Transformers, but the best SMFR architectures actually achieved perfect scores on it.
This indicates that the SMFR learned the actual underlying rule behind the task, while the other architectures only learned a heuristic.


Architecture.
Table <ref> only shows the five best-performing models of each type, but the ability of the SMFR to generalize perfectly extended beyond that.
For a stack depth between 1 and 5 (we tested 0 to 10) all SMFR models converged with 100% accuracy on all iterations regardless of the values of other hyperparameters. This supports our previous finding that the performance of our architecture can be improved by tuning the number of layers.
In contrast, both FNNs and Transformers generalized better to greater lengths the larger the model was, but reached their best performance on OOD_odd at smaller model sizes, which is a sign of overfitting.
We conjecture that larger FNNs and Transformers take the easy way out, so to say, by using their many parameters to just learn extensive statistical mappings. These generalize well to longer sequences but do not capture the actual formula that generates the data.
See the appendix for details.




Working with Noisy Input.
In the above experiments, the inputs were all formatted in a way that is easy for the SMFR to work with because the data is already split into blocks in the appropriate way: Each block represents one variable.
This raises the question: What happens if the input is noisy? Can the SMFR manage to reconstruct the correct format?
To test this we ran a second set of experiments with a modification: A random but fixed permutation is applied to the state before each iteration, permuting all neurons across all blocks.
This random permutation had no effect on FNNs, which still give the same performance since they do not care about the order of neurons in a layer.
SMFRs lost some OOD_odd accuracy but remained better than FNNs on average. Notably, in one of the 30 trials we ran the SMFR did achieve 100% OOD-accuracy again, even though the inputs were permuted.
This shows that SMFRs can learn to automatically arrange unstructured data into the block format they need. They can sometimes even do so perfectly, but not reliably. Investigating how to do this more reliably remains as future work.





 §.§ BPMNIST Experiments




Task.
The Block-Permuted MNIST task is designed to test the ability to generalize compositionally on realistic data. It is based on the popular MNIST task <cit.>. We take the MNIST images and split them horizontally into four equal-sized blocks, then permute these four blocks.
[Note that this task is different from the "Permuted MNIST" task used by <cit.>: We apply the permutation to parts of the image rather than to each individual pixel.]

The model has to learn to recognize the number for each of eight different permutations.

Some numbers are omitted from some permutations in the training set. Therefore, good performance on the test set requires compositional generalization.
See the appendix for details on this and the validation set.
Intuitively, there are two different ways to solve the task: Learn the mapping for each permutation separately, or learn to reverse the permutation first and then just learn MNIST.









We compared FNNs, Transformers and SMFRs.
We noticed that some of the SMFR models continued to strongly improve in test accuracy even after their training accuracy converged, which is similar to the "grokking" behavior observed by <cit.>.
To investigate this, we measured performance at two different time points: After 25k training steps, when the graphs showed that training accuracy had converged on all models, and again after 250k steps, to investigate if any models "grokked" the data.

We also wanted to see if this grokking behavior could be induced artificially. To this end we developed two additional variants of the SMFR.
SMFR_bias was equipped with an artificial bias for the first 300 training steps only, causing it to learn the ideal permutation quickly.
SMFR_no_context used a modified FNNR in which the original input blocks before the Multiplexer are not used as extra input. These inputs are a distraction, because an optimal solution does not need this data to solve this task.





Results.
Table <ref> shows the results.
The SMFR models performed worse than the FNNs and the Transformers at lower training times, but they kept improving after they achieved 100% training accuracy, and eventually surpassed the other models.
The variants SMFR_bias and SMFR_no_context both converged much more quickly and the former achieved the best performance of all models.

Their strong performance suggests that it could be well worth looking for a way to make SMFRs converge to the optimal routing more quickly. We leave this as future work.

Analysis.
For SMFRs and Transformers, we recorded key indicators of internal network behavior to try to understand how they learned to solve the problem.
The SMFRs that performed well tended to rearrange the four partial images in the very first MFNNR layer to get the same mapping for each permutation.

The longer training progressed, the more of the SMFRs converged to this behavior. 


Unlike the SMFRs, Transformers performed extremely well early on, but then barely improved with longer runtime. It appears that they simply learned to take an average over all input blocks, which is the same for all permutations, and then learned to detect MNIST digits from this mix.
This heuristic is easy to learn and worked surprisingly well, but did not leave room for further improvement.
The Transformer models ended up stuck in this local optimum.

Effects of Contained FNN Size.
The SMFR contains FNNs, but since much of the model's parameter count is used for routing operations, the contained FNN is necessarily smaller than the pure FNNs we used in our comparison.
We ran some additional tests to find out how important the size of these contained networks is relative to the ability to perform routing. To do so, we took the best SMFR architecture and increased the size of its contained FNNs to be equal to the standalone FNNs.
The resulting models achieved an accuracy of 0.948 after 25k training steps. This was better than before, but still slightly worse than the FNN.
It appears that until the network figures out how to route effectively, the overhead of doing so can sometimes do more harm than good.





§ DISCUSSION








Block-Operations in Other Architectures.


Adjusting other architectures to block-operations requires both the use of blocks and of MRPMs throughout the entire architecture.
Any module that does not support MRPMs will act as a barrier that prevents efficient communication. 





These modifications are non-trivial and we leave them as future work.

Additionally, since FNNs are so fundamental, it is possible that replacing all of them inside large, established architectures will have harmful side effects that require additional research to fix.




Nevertheless, since the SMFR is both useful in its current state and shows room for improvement through variants and architecture optimization, we believe that experimenting with different implementations of block-operations is a promising line of research to improve compositional generalization.


Transformers are a good candidate for enhancement.
Several previous papers have improved a Transformer's ability to generalize compositionally by providing ways to route input around them <cit.>.
The similar but more thorough modifications of MRPMs will likely bring similar benefits.








Since block-operations focus on routing representations we also expect that they will go well together with existing architectures that focus on generating object representations, such as Capsule Networks <cit.>, Slot Attention mechanisms <cit.>, or TIMs <cit.>.







Commutativity and Argument Selection.
Our module exhibits an inductive bias that helps it to learn tasks that rely on commutative functions, or on selecting arguments from a set of options.

Commutativity and argument selection are fundamental concepts in mathematics and programming respectively, so learning them more easily is likely to be helpful for these types of tasks.

Interpretability.
In the course of our experiments we noticed that SMFRs were easier to interpret than FNNs. The activations of the softmax and sigmoid neurons that control the routing correlate with the use of different subtasks.
This is analogous to how Transformers can be inspected by visually highlighting how much attention the model pays to different words <cit.>.

See the appendix for details.




§ LIMITATIONS




Architecture Optimization.

Our experiments showed that increasing an SMFR's depth and width does not always improve performance.

FNNs used to have the same issue, until <cit.> introduced residual connections.
We hope that a similar fix can be discovered for SMFRs as well. 



Computational Overhead.
SMFRs take more time per training step than FNNs of equal size because they perform a large number of operations on small matrices.
We have found empirically that they take more time than FNNs to converge on tasks that require heuristic reasoning but less on tasks with a lot of compositional logic, where their useful inductive bias outweighs the computational overhead.







Stability.
We have found that the SMFR architecture can sometimes get stuck in local optima because the gating weights and softmax values take on too extreme values, which kills the gradient. We fixed this by adding a regularization loss: Whenever the absolute of the weight used for a softmax or sigmoid exceeds a threshold value, we apply an MSE-loss to that weight, pushing it back to the threshold. See the appendix for details.












§ CONCLUSION





Neural Networks have difficulty achieving compositional generalization because it is difficult for them to route objects between subnetworks without incidentally modifying them.

We introduced the idea of block-operations, which is based on aggregating network activation neurons into larger semantic units and giving modules an inductive bias to route or modify blocks independently of each other.

Based on this idea, we developed the SMFR as a replacement for the FNN.

In our experiments the SMFR demonstrated superior generalization compared to other models and exhibited numerous useful properties, such as reusing subnetworks, understanding algorithmic logic, and being easy to inspect.
In several experiments it learned to decompose a task correctly where other architectures failed to do so.

This indicates that the SMFR is helpful for tackling the challenging problem of compositional generalization.

We noted several remaining limitations and proposed future directions of research.

Additionally, our findings suggest that the underlying idea of block-operations may be useful to other researchers in neural architecture design.





















§ APPENDIX






 §.§ Code




The supplementary material contains an executable python script that runs a trial of the algorithmic experiment in Section <ref>. It includes pytorch modules for the Multiplexer, the FNNR, the MFNNR, and the SMFR.
These modules are written with understandability in mind rather than performance.

An updated version with more details and more efficient modules will be provided in a github repository on acceptance of the paper, as it takes time to extract the relevant parts from our code base.















 §.§ Experiment Details, General



We used grid-search to generate different architectures. For FNNs we varied the number and size of intermediate layers. For SMFRs we varied the number of MFNNR modules (depth) and the number of blocks in each layer of neurons between the MFNNR modules (width). A depth of zero means that only a single MFNNR is used to map the input blocks to the output blocks directly. We also varied the size of the FNNs inside the MFNNR modules, but only to a lesser extent. For Transformers we varied the number of heads, the internal dimensions, and the depth of the encoder and decoder.

For all of our experiments, we used an Adam optimizer with a learning rate of 3e-4 and gradient clipping at 0.1. We used LeakyRelu with a negative slope of 0.01 as the activation function. The FNNs contained inside SMFRs had a width of 100 and variable depth, unless stated otherwise.

Each of the first three experiments used a set of digits as inputs, as well as a task indicator to differentiate between the subtasks of the experiment. The digits were one-hot encoded, with a block size of ten. The task-indicator tensor became its own block and was padded to the block size.
The experiment in Section <ref> used a block size of 196 instead, which is one fourth of the size of an MNIST image.

We used single-digit numbers in the addition/multiplication experiment and the double-addition experiment. We ran some exploratory tests with two digits as well and got similar results. We focused our analysis on single digit experiments because those converged faster, so we could run more trials.




  §.§.§ Experiment Details: Addition/Multiplication Experiments



We performed grid search over the following parameters:

    
  * The threshold value before switching training regimes: 0.7, 0.8, 0.9, 0.95, 1.0
    
  * For FNNs
    
        
  * Layer widths: 100, 200, 300
        
  * Layer depths: 2, 3, 4
    
    
  * For SMFRs
    
        
  * Stack widths: 1, 2, 3, 4, 5, 6, 7
        
  * Stack depths: 1, 2, 3
        
  * Contained FNN layer depths: 1, 2, 3
        
  * Attention type: Softmax, Straight Through Gumbel-Softmax
    


This resulted in model sizes in the range of roughly 13k to 280k parameters for both FNNs and SMFRs.



Table  <ref> shows the effect of model size on the preparation-data accuracy after catastrophic interference. HIGH refers to all models with at least 200,000 parameters, LOW refers to models with less than 50,000 parameters, and MID refers to the rest. The table shows that FNNs perform better at higher model sizes. It should be noted that the model sizes of FNNs and SMFRs were varied in different ways. The table indicates that growing SMFRs by stacking more MFNNR modules together is not as effective as increasing the FNNs inside the SMFRs.





Architecture Optimization.
Note that we varied the model sizes of FNNs and SMFRs in different ways, and this may explain some of our findings.
Table <ref> shows that SMFRs with more complex FNNs inside them performed better, especially at low thresholds. However this also had diminishing returns at 3 layers.
Our tests only included small variations on the contained FNNs because we focused on varying the number of blocks and MFNNR modules instead of the complexity of each FNN contained within them.
Additionally, we only altered the depths of the contained FNNs, not their width as well, as we did for the baseline FNNs that weren't part of an MFNNR. For FNNs, altering the width made a large difference.
Increasing the width of the FNNs inside the SMFRs as well would have led to much larger SMFR models, which is why we did not test this.
Therefore, the reason that pure FNNs performed slightly better than SMFRs at high model sizes and thresholds may be because we grew the complexity of SMFRs in a suboptimal way.
It remains future work to test if tuning the FNNs inside the SMFRs gives even better results than varying the number of blocks and MFNNR modules.











  §.§.§ Experiment Details: Double-Addition Experiments


We performed grid search over the following parameters:

    
  * For FNNs
    
        
  * Layer widths: 100, 200, 300
        
  * Layer depths: 2, 3, 4
    
    
  * For SMFRs
    
        
  * Stack widths: 1, 2, 3, 4, 5, 6, 7, 8, 9, 10
        
  * Stack depths: 0, 1, 2
        
  * Contained FNN layer depths: 1, 2
        
  * Attention type: Softmax, Straight Through Gumbel-Softmax
    


Module size comparisons between FNN and SMFR are largely irrelevant here because all FNNs failed regardless of size.




Table <ref> shows all 25 architectures with an OOD accuracy of 1.0, because a detailed breakdown over several dimensions is much more confusing to look at. This shows that model size correlates with performance. However, bigger models are not always better because the number of FNN layers was often 1 instead of 2, even though this parameter has a large impact on the model size.





Convergence Behavior.
Figure <ref> shows examples of the development of the OOD accuracy over time for different SMFR architectures. Each line represents one training run. The colors are meaningless and are just there to make it possible to tell the lines apart.
Most experiments with a perfect OOD accuracy converged to 1.0 almost immediately. In cases where this didn't happen, convergence took significantly longer, presumably because the model gets stuck in a local optimum and takes a while to escape.
It happened more often that OOD accuracy increased than that it dropped. Perhaps most importantly, the OOD accuracy never dropped after reaching 1.0.
In other words, the SMFR usually came to reuse modules more rather than less as the training progressed and remained stable once converged. This suggests that subnetwork reuse will also occur if SMFRs are used as components inside larger architectures that train for longer.







  §.§.§ Experiment Details: Algorithmic Experiments



We performed grid search over the following parameters:

    
  * For FNNs
    
        
  * Layer widths: 100, 200, 300
        
  * Layer depths: 2, 3, 4, 5, 6
    
    
  * For Transformers
    
        
  * Number of heads: 1, 4, 8
        
  * Internal dimensions: 400, 1200, 2000
        
  * Number of Encoder and Decoder Layers: 1, 2, 3, 4, 5, 6
    
    
  * For SMFRs
    
        
  * Stack widths: 6, 8, 10
        
  * Stack depths: 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10
        
  * Contained FNN layer depths: 1, 2, 3
    


This resulted in model sizes in the range of roughly 20k to 500k parameters for all three types of architectures.






Figure <ref> shows how the average accuracy of the different architectures changes with the number of iterations. Note that the training accuracy is the entry at iterations=2. Unlike the data reported in the main paper, which was based on the best models only, these graphs show averages over all models we ran, to provide a better overview.




Table <ref> shows the average OOD_even and OOD_odd accuracies for models of different sizes.
As pointed out in the main text, the SMFR has a sweet spot in performance based on the depth of the architecture.
In contrast, FNN and Transformer models become better at length generalization (OOD_even) as the model size increases, but do not become better at understanding the underlying atomic operation of the task (OOD_odd).






  §.§.§ Experiment Details: BPMNIST





We used eight different permutations for the task. In order to avoid unintentional bias, we selected the permutations in such a way that each block appears in each position an equal number of times when averaged over all permutations. We reported averages over the five best trials. The validation accuracy is the average over the first four permutations, the test accuracy is the average over the last four permutations. For each of the permutations used for testing we removed one of the ten numbers from the training set.

Because larger FNN models always performed better than smaller ones on this task, we did not use grid search as before. Instead, we varied the parameters mentioned in the previous experiment while keeping the total parameter count of each model roughly the same, around 4 million parameters. We also increased the standard size of the FNNs contained within an SMFR to have a depth of 4 and a width of 400, which was larger than before but still smaller than the pure FNN. The best SMFRs had a low stack depth, which also made them easier to inspect.

We inspected the internal network behavior of Transformers and SMFRs by graphing how the following indicators developed over time:

    
  * Attention sharpness: How much does each attention mechanism pay attention to only a single input vs. all of them? Good SMFR runs had high sharpness.
    
  * Difference between permutations: If we apply the inverse of the permutation used in a training batch to the data after the first attention layer, do we get similar distributions for each permutation? This measures roughly how much the network learns to undo the permutation. Good SMFR runs did learn to undo permutations.
    
  * Fairness of attention: Does the network pay equal amounts of attention to each of the four input blocks, or does it drop some of them? MNIST images carry most of their information content in the middle of the image and not the borders, so the outer two blocks were sometimes ignored by the model. The best-performing runs learned not to throw away this data.





 §.§ Regularization Loss



As mentioned in the Limitations, our model used to suffer from a stability issue that we have since fixed. The problem was that softmax and sigmoid weights sometimes kept growing even after becoming extreme enough that they had a gradient of zero. Our inspection showed that this was caused by the fact that the softmax and sigmoid values are based on computations that are also used for generating data in the MFNNRs. This means that they are subject to destructive interference. They effectively receive gradients from two sources: The softmax/sigmoid they are supposed to work on and incidental gradients from other parts of the model, but the former irrecoverably become zero.

We fixed this by adding a regularization loss: Whenever the absolute of the weight used for a softmax or sigmoid exceeds a threshold value, we apply an MSE-loss to that weight, pushing it back to the threshold. The target value of the MSE-loss is simply the tensor itself with each neuron clamped to a threshold value. We empirically found that a threshold of 20 worked well.

It should be noted that we introduced the regularization loss while working on a larger, more complex experiment that we will report in a subsequent paper. It is possible that this loss is not needed for many of the experiments we report in this paper. We have noticed that using it slows convergence but increases reliability. We have opted to always use it because reliability is more important to us than convergence speed while developing new architectures.






 §.§ Interpretability



To inspect an SMFR model, it is often enough to look at the weights that determine routing decisions: The softmax values in the Multiplexer, and the gating weight in the FNNR. It can be helpful to have a set of example inputs for which you expect different routing behavior and to compare the routing weights for these.



On the double-addition task, our inspection showed that
it is possible to measure the degree of modularity in our SMFR directly by investigating the softmax weights on example inputs. We manually inspected the values of the weights that the multiplexers and FNNR modules use in a trial that achieved good OOD performance and had a simple architecture, with depth 1 and width 2. The architecture had only 12 Multiplexer weights and 3 FNNR residual weights, which made it very easy to investigate how data is routed in the network for any given input.
Our inspection showed that the irrelevant inputs did correctly receive a weight of 0, while the ones that are important received a weight of 0.5. As it turns out, the network learned to use the correct pair of numbers and used interpolation in a single block to take advantage of symmetry. Instead of always moving block A1 to B1 and A2 to B2 as a human would intuitively do it, it learned to interpolate B1=0.5*A1+0.5*A2, which is actually more efficient.
We also found that the gating weights of several FNNR modules approached zero as training progressed. This implied that the network came to rely more and more on copy operations and less on FNN operations.
It was quite easy to inspect the network and learn this information, since we only needed to inspect the gating weights.
In larger networks used for practical tasks, the number of network weights may be too large for a manual analysis. However there should still be few enough of them for an automated analysis to yield useful results, especially since the gating weights have inherent meaning (does a block of data get used or not?), unlike normal neurons in a fully-connected neural network, which can mean anything.

On the algorithmic task, our inspection confirmed that the SMFRs can learn to pass blocks through without altering them where doing so is useful, and to alter only the ones that need to be altered. We were able to confirm this desired behavior by only looking at the routing weights. In an FNN, confirming such behavior would have required a much more complex correlation analysis of the neurons.

On the BPMNIST task, our inspection allowed us to determine which models understood the data generating process by decomposing the task into its two subtasks (fixing the permutation and solving MNIST). We reported on this in detail in Section <ref> and Section <ref> of the appendix.




 Paper Title: Tractable and Provably Efficient Distributional Reinforcement Learning with General Value Function Approximation.




Captured molecules could  make a Bose star visible
    I. B. Samsonov
    August 5, 2024
==================================================




§ ABSTRACT 

Distributional reinforcement learning improves performance by effectively capturing environmental stochasticity, but a comprehensive theoretical understanding of its effectiveness remains elusive.
In this paper, we present a regret analysis for distributional reinforcement learning with general value function approximation in a finite episodic Markov decision process setting. 
We first introduce a key notion of Bellman unbiasedness for a tractable and exactly learnable update via statistical functional dynamic programming.
Our theoretical results show that approximating the infinite-dimensional return distribution with a finite number of moment functionals is the only method to learn the statistical information unbiasedly, including nonlinear statistical functionals.
Second, we propose a provably efficient algorithm, , achieving a regret bound of Õ(d_E H^3/2√(K)) where H is the horizon, K is the number of episodes, and d_E is the eluder dimension of a function class.









§ INTRODUCTION



Distributional reinforcement learning (DistRL) <cit.> is an advanced approach to reinforcement learning (RL) that focuses on the entire probability distribution of returns rather than solely on the expected return.
By considering the full distribution of returns, distRL provides deeper insight into the uncertainty of each action, such as the mode or median. 
This framework enables us to make safer and more effective decisions that account for various risks <cit.>, particularly in complex real-world situations, such as robotic manipulation <cit.>,
neural response <cit.>, 
stratospheric balloon navigation <cit.>, 
algorithm discovery <cit.>, and several game benchmarks <cit.>.


Although distRL has a solid theoretical basis and demonstrates effective empirical performance, its efficiency and tractability are still not fully understood.

In practice, distributions contain an infinite amount of information, and we should resort to approximations using a finite number of parameters or statistical functionals, such as categorical <cit.> and quantile representations <cit.>.
However, not all statistical functionals can be learned through the Bellman operator, as the meaning of statistical functionals is not typically preserved after updates. 
For example, the median is not preserved under the Bellman updates, as the median of a mixture of two distributions does not necessarily equate to the mixture of their medians.

Hence, a distinct analysis is required to determine whether there exists a corresponding Bellman operator for each statistical functional that ensures commutativity.


<cit.> defined Bellman closedness as a property of statistical functionals that can be exactly learned by the existence of a corresponding Bellman operator.








At this point, we take a closer look at the distributional Bellman update and revisit what additional properties of statistical functionals beyond Bellman closedness are needed to construct tractable algorithms. 
In doing so, we identify and provide the following challenges that arise when statistical functionals are used to express the update process:
























    
  * Representing a mixture distribution with a finite, fixed number of parameters inevitably leads to approximation errors during the update process. 
    For example, when expressing the mixture of two distributions, each represented by N parameters, compressing 2N into N parameters in the mixture results in inevitable information loss.
    
    
    
    
    

    
  * Due to the unknown nature of the transition probability (·|s,a), the target distribution is estimated by sampling the next state s'. 
    For a tractable algorithm, it is crucial that the statistical functionals of the target distribution can be unbiasedly estimated using the statistical functionals from the sampled distribution.    



In this paper, we describe the statistical functional and sketch of the distribution, and introduce the additional approximation errors that can occur with the distRL method. 
We define the notion of Bellman unbiasedness, a property for representing the distribution tractably and prove that the moment functional is the only solution that satisfies this property, including nonlinear statistical functionals.
Based on this property, we propose a provably efficient statistical functional reinforcement learning algorithm with general value function approximation, denoted as . 
We will discuss the inherent issues of distributional Bellman completeness, a structural assumption previously defined in the literature <cit.>, and investigate the benefits of redesigning this concept using a sketch.






In summary, our main contributions are as follows:

    
  * Introduce a key property of Bellman unbiasedness for collections of statistical functionals (sketch) and show that the moment functionals are the unique structure that can be unbiasedly estimated and exactly learned from distributions, including nonlinear statistical functionals.

    
  * Revisit the existing concept of distributional Bellman Completeness through a statistical functional lens.
    We enhance the tractability by using moment functional, that the agent can only access a finite-dimensional embedding space.
    This removes the misspecification errors commonly encountered when discretizing infinite-dimensional distributions.

    
  * Propose a tractable and provably efficient distributional RL algorithm called , achieving a tightest regret upper bound Õ(d_E H^3/2√(K)) [We ignore poly-log terms in H,S,A,K in the Õ(·) notation.] in distributional framework.
    Our framework yields a better regret bound with a weaker structural assumption compared to prior results.















§ RELATED WORK














  
Distributional RL.
In classical RL, the Bellman equation, which is based on expected returns, has a closed-form expression.
However, it is questionable whether any statistical functionals of return distribution always have their corresponding closed-form expressions. 
<cit.> introduced the notion of Bellman Closedness for collections of statistical functionals that can be updated in a closed form via Bellman update. 
They showed that the only Bellman-closed statistical functionals in the discounted setting are the moments _Z ∼η[Z^k]. 
More recently, <cit.> proposed a general framework for distRL, where the agent plans to maximize their own utility functionals instead of expected return and formalized the property of Bellman Optimizability.
They proved that the only W_1-continuous and linear Bellman Optimizable statistical functionals are shown to be exponential utilities 1/λlog_Z ∼η[exp(λ Z)] in the undiscounted setting.

In practice, C51 <cit.> and QR-DQN <cit.> are notable distributional RL algorithms where the convergence guarantees of sampled-based algorithms are proved <cit.>.
<cit.> expanded the class of policies on arbitrary distortion risk measures by taking the based distribution non-uniformly and improve the sample efficiency from their implicit representation of the return distribution.
<cit.> highlighted the drawbacks of optimistic exploration in distRL, introducing a randomized exploration method that perturbs the return distribution when the agent selects their next action. 

























  
RL with General Value Function Approximation.
Regret bounds have been studied for a long time in RL, across various domains such as bandit <cit.>, tabular RL <cit.>, and linear function approximation <cit.>.
In recent years, deep RL has shown significant performance using deep neural networks as function approximators, and attempts have been made to analyze whether it is efficient in terms of general function approximation.
<cit.> established a provably efficient RL algorithm with general value function approximation based on the eluder dimension d_E <cit.> and achieves a regret upper bound of Õ(poly(d_E H)√(K)).
To circumvent the intractability from computing the upper confidence bound, <cit.> injected the stochasticity on the training data and get the optimistic value function instead of upper confidence bound, enhancing computationally efficiency.
Beyond risk-neutral setting, several prior works have shown regret bounds under risk-sensitive objectives (e.g., entropic risk <cit.>, CVaR <cit.>), which align with our approach in that they are built on a distribution framework.
<cit.> achieved the regret upper bound of Õ(exp(H) √(|𝒮|^2 |𝒜| H^2 K)) and the lower bound of Ω(exp(H) √(|𝒮||𝒜|HK)) in tabular setting.




  
DistRL with General Value Function Approximation.
Recently, only few efforts have aimed to bridge the gap between two fields.
<cit.> proposed a distributional RL algorithm, , which enjoys small-loss bound by using a log-likelihood objective.
Similarly, <cit.> provided a risk-sensitive reinforcement learning framework with static lipschitz risk measure.
While these studies analyze within a distributional framework, they do not address the intractability of implementation in infinite-dimensional space of distributions.
In contrast, our approach focuses on a statistical functional framework, providing a detailed comparison with other distRL methods as shown in Table <ref>.








§ PRELIMINARIES









  
Episodic MDP.

We consider a episodic Markov decision process which is defined as a ℳ = (𝒮, 𝒜, H, ℙ, r) characterized by state space 𝒮, action space 𝒜, horizon length H, transition kernels ℙ = {ℙ_h }_h ∈ [H], and reward r = { r_h }_h ∈ [H] at step h ∈ [H]. 







The agent interacts with the environment across K episodes. For each k ∈ [K] and h ∈ [H], ℍ_h^k = (s_1^1 ,a_1^1, …, s_H^1, a_H^1, …, s_h^k , a_h^k ) represents the history up to step h at episode k.
We assume the reward is bounded by [0,1] and the agent always transit to terminal state s_end at step H+1 with r_H+1=0.















  
Policy and Value Functions.
A (deterministic) policy π is a collection of H functions {π_h : 𝒮→𝒜}_h=1^H.
Given a policy π, a step h ∈ [H], and a state-action pair (s,a) ∈𝒮×𝒜, the Q and V-function are defined as




Q_h^π(s,a) (: 𝒮×𝒜→ℝ ) _π[∑_h'=h^H r_h'(s_h', a_h') | s_h = s, a_h =a ] and
V_h^π(s) (: 𝒮→ℝ ) _π[∑_h'=h^H r_h'(s_h', a_h') | s_h = s ].






  
Random Variables and Distributions. For a sample space Ω, we extend the definition of the Q-function into a random variable and its distribution,



    Z_h^π(s,a) (: 𝒮×𝒜×Ω→ℝ ) ∑_h'=h^H r_h'(s_h', a_h') | s_h = s, a_h =a, a_h' = π_h'(s_h'),
        
       η_h^π(s,a) (: 𝒮×𝒜→𝒫(ℝ) ) law(Z_h^π (s,a)).

Analogously, we extend the definition of V-function by introducing a bar notation.



    Z̅_h^π(s) (: 𝒮×Ω→ℝ ) ∑_h'=h^H r_h'(s_h', a_h') | s_h = s, a_h' = π_h'(s_h'),
        
       η̅_h^π(s) (: 𝒮→𝒫(ℝ) ) law(Z̅_h^π (s)).

Note that Z̅^π_h(s) = Z^π_h(s, π(s)) and η̅_h^π(s) =η_h^π(s,π(s)).
We use π^⋆ to denote an optimal policy 
and denote V_h^⋆(s) = V_h^π^⋆(s), Q_h^⋆(s,a) = Q_h^π^⋆ (s,a), η_h^⋆(s,a) = η_h^π^⋆(s,a), and η̅_h^⋆(s) = η̅_h^π^⋆(s).
For notational simplicity, we denote the expectation over transition,
[ℙ_h V_h+1^π](s,a) = _s' ∼ℙ_h(· | s,a)  V_h+1^π(s'),
[ℙ_h Z̅_h+1^π](s,a) = _s' ∼ℙ_h(· | s,a)Z̅_h+1^π(s'),
and [ℙ_h η̅_h+1^π](s,a) = _s' ∼ℙ_h(· | s,a)η̅_h+1^π(s'). 
[Note that _s' ∼_h(· |s,a)η̅^π_h+1(s') is a mixture distribution.]
For brevity, we refer to η̅^π simply as η̅.

In the episodic MDP, the agent aims to learn the optimal policy through a fixed number of interactions with the environment across a number of episodes.
At the beginning of each episode k(∈ [K]), the agent starts at the initial state s_1^k and choose a policy π^k.
In step h(∈ [H]), the agent observes s^k_h (∈𝒮), takes an action a^k_h (∈𝒜) ∼π^k_h(·| s^k_h), receives a reward r_h(s^k_h, a^k_h), and the environment transits to the next state s^k_h+1∼_h(· | s^k_h, a^k_h).
Finally, we measure the suboptimality of an agent by its regret, which is the accumulated difference between the ground truth optimal and the return received from the interaction.

The regret after K episodes is defined as
Reg(K) = ∑_k=1^K V_1^⋆(s_1^k) - V_1^π^k(s_1^k) .




    





























    























  
Distributional Bellman Optimality Equation.
Recall that η^⋆_h satisfies the distributional Bellman optimality equation:



















    η_h^⋆ (s,a) = (Τ_h η_h+1^⋆) (s,a)
           = _s' ∼_h(·|s,a), a' ∼π_h^⋆(· | s') [(ℬ_r_h)_#η_h+1^⋆ (s', a')]
        
       =  (ℬ_r_h)_# [_h η_h+1^⋆] (s, a)

where ℬ_r : ℝ→ℝ is defined by ℬ_r (x) = r + x, and
g_#η∈𝒫(ℝ) is the pushforward of the distribution η through g, i.e., g_#η (A) = η( g^-1(A) ) for any Borel set A ⊆ℝ.















  
Additional Notations.

For a given N, we denote an N-dimensional function class 
ℱ^N ℱ^(1)×…×ℱ^(N)⊆{ f = [f^(1), ⋯, f^(N)] : 𝒮×𝒜→ℝ^N }.






Given a dataset 𝒟 = { (s_t, a_t, [z_t^(1) , …, z_t^(N)])}_t=1^|𝒟|⊆𝒮×𝒜×ℝ^N,
a set of state-action pairs 𝒵 = { (s_t, a_t) }_t=1^|𝒵|⊆𝒮×𝒜
and
for a function f: 𝒮×𝒜→ℝ^N, we define the norm f^(n)_∞, f_∞, 1, f_𝒟, f_𝒵 as written in Appendix <ref>.












For a set of (vector-valued) functions ℱ^N ⊆{ f: 𝒮×𝒜→ℝ^N},
the width function of (s,a) is defined as
w^(n)(ℱ^N , s, a) 
    max_f,g ∈ℱ^N  | f^(n)(s,a) - g^(n)(s,a) |.








































§ STATISTICAL FUNCTIONALS IN DISTRIBUTIONAL RL







































 §.§ Bellman Unbiasedness


We revisit the definitions of statistical functionals and illustrate the property of Bellman closedness described in <cit.>.
Then we introduce another pivotal property, Bellman unbiasedness, required for learning in terms of statistical functionals, rather than approximated distribution.



    A statistical functional is a mapping from a probability distribution to a real value ψ : 𝒫(ℝ) →ℝ. 
    A sketch is a vector-valued function ψ_1:N : 𝒫(ℝ) →ℝ^N specified by an N-tuple where each component is a statistical functional, 
    
    ψ_1:N(·) = (ψ_1(·), ⋯, ψ_N(·)).


We denote the domain of sketch as 𝒫_ψ_1:N(ℝ) and its image as 
I_ψ_1:N = {ψ(η̅) : η̅∈𝒫_ψ_1:N(ℝ) }.
We further extend to state return distribution functions 

ψ_1:N(η̅) = (ψ_1:N(η̅(s)) : s ∈𝒮).











    A sketch ψ_1:N is Bellman closed if there exists an operator Τ_ψ_1:N : I_ψ_1:N^𝒮→ I_ψ_1:N^𝒮 such that
    
    ψ_1:N (Τη̅) = Τ_ψ_1:Nψ_1:N( η̅)   for all η̅∈𝒫(ℝ)^𝒮

    which is closed under the distributional Bellman operator Τ : 𝒫(ℝ)^𝒮→𝒫(ℝ)^𝒮.

















Bellman closedness is the property that a sketch are exactly learnable when updates are performed from the infinite-dimensional distribution space to the finite-dimensional embedding space. 
Notably, <cit.> showed that the only finite linear statistical functionals that are Bellman closed are given by the collections of statistical functionals where its linear span is equal to the set of exponential-polynomial functionals 

where ψ_0 is the constant functional equal to 1.



    
    <cit.> discussed only linear statistical functionals when defining a sketch, leaving open questions about the Bellman closedness of nonlinear statistical functionals such as variance, central moments, or quantiles.
    However, it is noteworthy that nonlinear statistical functionals, such as maximum or minimum, can also be Bellman closed. 
    
    In their proof, there might be an ambiguity regarding the assumption that the corresponding sketch Bellman operator for quantile is linear.
    
    
    In this paper, we show the non-existence of sketch Bellman operator for quantile and discuss in detail in Appendix <ref>.
    
 























































    Quantile functional cannot be Bellman closed under any additional sketch.




We investigate whether there exists another sketch that is both nonlinear and Bellman closed, and introduce additional conditions under which the Bellman closed property might be missing.
During the implementation phase, it is important to note that the agent does not have access to the transition kernel. 

Specifically, the agent can only access the empirical transition kernel  which is derived from a finite number of sample estimates. 
This limitation implies that the operator should be treated as an empirical operator 𝒯̂_ψ, rather than 𝒯_ψ. 
Therefore, we introduce a new notion of Bellman unbiasedness that allows us to unbiasedly estimate the expected distribution (ℬ_r)_#_s' ∼(· | s,a) [η̅(s')], which is a mixture by transitions, from the sample distribution (ℬ_r)_#η̅(s').


























































    
    
    A sketch ψ is Bellman unbiased if there exists an integer k and a corresponding vector-valued estimator ϕ_ψ = ϕ_ψ(ψ(x_1), ⋯, ψ(x_k)) : (I_ψ^𝒮)^k → I_ψ^𝒮 where the sketch of expected distribution (ℬ_r)_#_s' ∼(· | s,a)[η̅(s')] can be unbiasedly estimated by ϕ_ψ using the k samples from the sketch of the sample distribution (ℬ_r)_#η̅(s'), i.e.,
    
    _s'_1, ⋯ ,s'_k ∼(· | s,a)[ ϕ_ψ( ψ( (ℬ_r)_#η̅(s'_1)_𝒯̂_ψψ(η) ),  ⋯, ψ( (ℬ_r)_#η̅(s_k')  ) )]
                    = ψ( (ℬ_r)_#_s' ∼(· | s,a) [ η̅(s') ]).















Bellman unbiasedness is another natural definition, similar to Bellman closedness, which takes into account a finite number of samples for the transition.
As a simple example, median functional is not Bellman unbiased as the sample median is not an unbiased estimate of the median functional.
In general, there is no unbiased estimator for the median, implying that the median functional cannot be unbiasedly estimated within an embedding space represented by a finite number of statistical functionals.
On the other hand, variance functional is a nonlinear statistical functional but can be unbiasedly estimated.
Then, the following question naturally arises;



    “Which sketches are unbiasedly estimatable under the sketch-based Bellman update?”



The following lemma answers this question.
















    Let F_η̅ be a CDF of the probability distribution η̅∈𝒫_ψ(ℝ)^𝒮. 
    Then a sketch is Bellman unbiased if and only if the sketch is homogeneous over 𝒫_ψ(ℝ)^𝒮 of degree k, i.e.,
    there exists some vector-valued function h = h(x_1, ⋯, x_k): 𝒳^k →ℝ^N such that
    
    ψ(η̅) =  ∫⋯∫ h(x_1, ⋯, x_k) dF_η̅(x_1) ⋯ dF_η̅(x_k).












Lemma <ref> states that in statistical functional dynamic programming, the unbiasedly estimatable embedding of a distribution can only be structured in the form of functions that are homogeneous of finite degree.

(e.g., variance functional is a homogeneous of degree 2 <cit.>).
Taking this concept further and combining it with the results on Bellman closedness, we prove that even when including a nonlinear statistical functional, the only sketch that can be unbiasedly estimated in a finite-dimensional embedding space is the moment functional.










































































    The only finite statistical functionals that are both Bellman unbiased and closed are given by the collections of ψ_1, …,  ψ_N where its linear span {∑_n=0^N α_n ψ_n | α_n ∈ℝ ,∀ N} is equal to the set of exponential polynomial functionals {η→_Z ∼η[Z^l exp(λ Z)] |   l = 0, 1, …, L, λ∈ℝ}, where ψ_0 is the constant functional equal to 1.
    In discount setting, it is equal to the linear span of the set of moment functionals {η→_Z ∼η[Z^l] |  l = 0, 1, …, L } for some L ≤ N.  















    























Compared to <cit.>, our results extend beyond linear statistical functionals to include unbiasedly estimatable nonlinear statistical functionals, showing the uniqueness of the moment functional. 
Our theoretical results not only show that high-order central moments such as variance or skewness are exactly learnable, but also reveal that other nonlinear statistical functionals like median or quantiles inevitably involve approximation errors due to biased estimations.
In Section <ref>, we emphasize that unbiasedly estimable sketches play an important role in mitigating  misspecification errors, thereby enabling us to achieve sublinear regret. 










































































 §.§ Statistical Functional Bellman Completeness



























We consider distributional reinforcement learning with value function approximation.
For successful TD learning, reinforcement learning with function approximation commonly requires the assumption, Bellman Completeness, that after applying Bellman operator, the output lies in the function class ℱ <cit.>.
On the other hand, our approach receives a tuple of function class ℱ^N ⊆{ f : 𝒮×𝒜→ℝ^N } as input to represent N-moment of distribution.
Building on this, we assume that for any η̅ : 𝒮→𝒫([0,H]), the sketch of target function that results from applying sketch Bellman operator lies in the function class ℱ^N.







In the seminal works, <cit.> and <cit.> assumed that the function class 
ℋ⊆{η : 𝒮×𝒜→𝒫([0,H])} follows the distributional Bellman Completeness assumption, i.e., 
if η∈ℋ for all π, h ∈ [H], then Τ_h^πη∈ℋ.
However, since the distributional Bellman operator mixes distributions for the next state transition, it implies that a function class ℋ must be closed under mixture.
Due to this closedness under mixture, 

both previous studies assumed a discretized reward MDP, which requires the prior knowledge that outcomes of return distribution are additively closed.


In general, representing a distribution with a fixed number of parameters inevitably increases the required number of parameters due to the emergence of a mixture distribution.
Consequently, trying to represent it with a limited number of parameters leads to approximating the learning outcome, thereby causing model misspecification


which generally introduces an additional linear regret.
Precisely, a regret lower bound with model misspecification error ζ is known as Ω(ζ K) in a linear bandit setting <cit.>.
Therefore, assuming Bellman Completeness on the distribution space is invalid as it significantly constrains the MDP structure or causes linear regret.

To circumvent this issue, we revisit the assumption of distributional Bellman completeness through the statistical functional lens. 
We propose a new framework that matches a finite number of all statistical functionals to the target function, rather than the entire distribution itself.



[Statistical Functional Bellman Completeness]
    
    For any distribution η̅ : 𝒮→𝒫([0,H]) and h ∈ [H], there exists f_η̅∈ℱ^N which satisfies
    
    f_η̅(s,a) = ψ_1:N( (ℬ_r_h)_# [ℙ_h η̅](s,a) ) 
              ∀ (s,a) ∈𝒮×𝒜




[Model  Misspecification]
    There exists a set of function ℱ^N and a real number ζ>0, such that for any η̅ : 𝒮→𝒫([0,H]) and h ∈ [H], there exists f_η̅∈ℱ^N which satisfies
    
    max_(s,a) ∈𝒮×𝒜 f_η̅(s,a) - ψ_1:N( (ℬ_r_h)_# [ℙ_h η̅](s,a) ) ≤ζ  ∀ (s,a) ∈𝒮×𝒜.

    We call ζ the misspecification error.






















































§ ALGORITHMS

    
    
    

In this section, we propose  for distributional RL framework with general value function approximation. 
Leveraging the result from Theorem <ref>, we introduce a moment least square regression.

This allows us to capture a finite set of moment information from the distribution, which can be unbiasedly estimated, thereby leading to the truncated moment problem.
Unlike previous work <cit.> that estimates in infinite-dimensional distribution spaces, our method enables tractable distribution estimation in finite-dimensional embedding spaces without approximation error.







    




  
Overview.
At the beginning of episode k ∈ [K], we maintain all previous samples {(s^τ_h', a^τ_h', r^τ_h') }_(τ, h') ∈ [k-1] × [H] and initialize a sketch ψ_1:N( η̅^k_H+1(·) ) = 0^N.
For each step h = H, …, 1, we compute the normalized sample moments of target distribution { (ℬ_r_h'^τ)_#η̅^k_h+1(s_h'+1^τ) }_h' ∈ [H] with the help of binomial theorem,

    ψ_n( (ℬ_r_h'^τ)_#η̅_h(s_h'+1^τ) ) 
           = [(Z̅_h+1^k (s_h'+1^τ) + r_h'^τ)^n] / H^n-1
        
       = ∑_n'=0^n H^n'ψ_n'( η̅_h(s_h'+1^τ) ) (r_h'^τ)^n- n'/ H^n-1

and iteratively solve the N-moment least square regression

    f̃^k_h,η̅←min_f ∈ℱ∑_τ=1^k-1∑_h'=1^H( ∑_n=1^N f^(n)(s_h'^τ, a_h'^τ ) - ψ_n ( (ℬ_r_h'^τ)_#η̅_h+1^k(s_h'+1^τ) )  )^2

based on the dataset 𝒟^k_h which contains the sketch of temporal target ψ_1:N( (ℬ_r_h'^τ)_#η̅_h+1^k(s_h'+1^τ) ). 
Then we define Q^k_h(·, ·) =  min{ (f̃^k_h,η̅)^(1)(·, ·) + b^k_h(·, ·), H }  and choose the greedy policy π^k_h(·) with respect to Q^k_h.
We repeat the procedure until all the K episodes are completed.




    For an optimistic planning, we define the bonus function as the width function b^k_h(s,a)  w^k_h((ℱ^N)^k_h, s,a) where (ℱ^N)^k_h denotes a confidence region at step h, episode k.
    When ℱ is a linear function class, the width function can be evaluated by simply computing the maximal distance of weight vector.
    For a general function class ℱ, computing the width function requires to solve a set-constrained optimization problem, which is known as NP-hard <cit.>.
    However, a width function is computed simply for optimistic exploration, and approximation errors are known to have a small effect on regret <cit.>.
    



















































§ THEORETICAL ANALYSIS


In this section, we provide the theoretical guarantees for  under assumption <ref>.
We apply proof techniques from <cit.> and extend the result to a statistical functional lens. 
First, we generalize the concept of eluder dimension <cit.> to the vector-valued function, which has been widely used in RL literatures <cit.> to measure the complexity of learning with the function approximators.


    Let ϵ≥ 0 and 𝒵 = { (s_i, a_i) }_i=1^n ⊆𝒮×𝒜 be a sequence of state-action pairs.


    

    
  * A state-action pair (s,a) ∈𝒮×𝒜 is ϵ-dependent on 𝒵 with respect to ℱ^N if f - g_𝒵≤ϵ for any vector-valued function f, g ∈ℱ^N, then |f^(1)(s,a) - g^(1)(s,a)| ≤ϵ. 
    

    
  * An (s,a) is ϵ-independent on 𝒵 with respect to ℱ^N if (s,a) is not ϵ-dependent on 𝒵.

    
  * The ϵ-eluder dimension dim_E(ℱ^N ,ϵ) of a vector-valued function class ℱ^N is the length of the longest sequence of elements in 𝒮×𝒜 such that, for some ϵ' ≥ϵ, every element is ϵ'-independent on its predecessors.
    


    








We assume that the function class ℱ^N and state-action space 𝒮×𝒜 have bounded covering numbers.
[Covering number]
    
    For any ϵ >0, the following holds:
    
        
  * there exists an ϵ-cover 𝒞(ℱ^N, ϵ) ⊆ℱ^N with size |𝒞(ℱ^N, ϵ)| ≤𝒩(ℱ^N, ϵ), such that for any g ∈ℱ^N, there exists g' ∈𝒞(ℱ^N, ϵ) with g-g'_∞, 1≤ϵ.

        
  * there exists an ϵ-cover 𝒞(𝒮×𝒜, ϵ) with size | 𝒞(𝒮×𝒜 , ϵ) | ≤𝒩(𝒮×𝒜 , ϵ), such that for any (s,a) ∈𝒮×𝒜, there exists (s', a')∈𝒞(𝒮×𝒜, ϵ) with max_f ∈ℱ | f(s,a) - f(s',a') | ≤ϵ     
    
    
    

    
    



















The following two lemmas give confidence bounds on the sum of the l_2 norms of all normalized moments.






Consider a fixed k ∈ [K] and a fixed h ∈ [H]. Let 𝒵_h^k = { (s_h^τ , a_h^τ)}_τ∈ [k-1] and 

𝒟^k_h, η̅ = {( s_h^τ , a_h^τ,  ψ_1:N( (ℬ_r_h'^τ)_#η̅(s_h'+1^τ) )) }_τ∈ [k-1] for any η̅ : 𝒮→𝒫([0,H]).
Define f̃_h, η̅^k = min_f ∈ℱ^N f ^2_𝒟^k_h, η̅.
For any η̅ and δ∈ (0,1), there is an event ℰ(η̅, δ) such that conditioned on ℰ(η̅, δ), with probability at least 1-δ, for any η̅' : 𝒮→𝒫([0,H]) with 
ψ_1:N( η̅') - ψ_1:N(η̅) _∞,1≤ 1/T 

,we have

    f̃_h,η̅'(· , ·) - ψ_1:N( (ℬ_r(·, ·))_#  [ℙη̅'](· , ·) ) _𝒵_h^k   ≤c' (N^1/2 H √(log(1/δ) + log𝒩(ℱ^N, 1/T)  ))

for some constant c' >0.






We remark that using moment functional as a sketch is a unique structure that leads to a concentration results for the difference between the target and sampled sketches
ψ( (ℬ_r)_# [η̅]) - ψ( (ℬ_r)_#η̅), as the tranformation ϕ_ψ of moment functional is an identity.
Furthermore, using normalized moments minimizes the size of the confidence region from O(H^N) to O(√(N)). 
This adjustment is akin to scaling the optimization function in multi-objective optimization to treat each objective equally, which effectively prevents the model from favoring objectives with larger scales. 


















    
    












































    








    































  
























    





























    


















    Let (ℱ^N)_h^k = { f ∈ℱ^N |  f - f̃^k_h, η̅^2_𝒵^k_h≤β(ℱ^N, δ) }, where
    
    β(ℱ^N,δ) ≥ c' · NH^2 (log (T/δ) + log𝒩(ℱ^N , 1/T) )

    for some constant c'>0.
    Then with probability at least 1-δ/2, for all k,h ∈ [K] × [H], we have
    
    ψ_n ( (ℬ_r_h(·, ·))_#[_h η̅_h+1^k](·, ·) ) ∈ (ℱ^N)_h^k

































    




Lemma <ref> guarantees that the sequence of moments from the target distribution ψ_1:N( (ℬ_r_h(·, ·))_# [_h η̅_h+1^k](·, ·) ) lies in the confidence region (ℱ^N)^k_h with high probability.
Supported by the aforementioned lemma, we can further guarantee that all Q-functions are optimistically estimated with high probability and derive our final result.















































    



    


















    




































































































    Under Assumption <ref>, with probability at least 1- δ,  achieves a regret bound of
    
    Reg(K) ≤ 
        2Hdim_E(ℱ^N, 1/T) + 4H√(KH log(1/δ))

    
    
    
    
    


Under weaker structural assumptions, we show that  enjoys sublinear regret of order Õ(d_E H^3/2√(K)), which is √(H) better than the state-of-the-art distributional framework algorithm  <cit.>.

For the linear MDP setting, we have d_E = Õ(d) and thus the regret bound of  is Õ(√( d^2 H^3 K)) which matches a lower bound Ω(√(d^2 H^3 K)) 
<cit.>.
In our analysis, we highlight two main technical novelties which significantly reduces the complexity of regret;


    
  * We refine previous lemma of <cit.> and <cit.> to remove the dependency of β(ℱ^N, 1/δ) (See Appendix <ref>), ensuring that regret bound depends only on the defined function class, not on the number of moment extracted.

    
  * As shown in Table <ref>, we define the eluder dimension d_E in a finite-dimensional embedding space ℱ^N, while other methods rely on an infinite-dimensional distribution space ℋ⊆ℱ^∞.


































§ CONCLUSIONS


We describe the sources of approximation error inherent in distribution-based updates and introduce a pivotal concept of Bellman unbiasedness, which enables to exactly learn the information of distribution.
We also present a provably efficient distRL algorithm, , with general value function approximation.
Notably, our algorithm achieves a near-optimal regret bound of Õ(d_E H^3/2√(K)), matching the tightest upper bound achieved by non-distributional framework <cit.>.




One interesting future direction would be to reformulate the definition of regret as discrepencies in moments rather than the expected return, and to show the sample-efficiency of distRL.
We hope that our work sheds some light on future research in analyzing the provable efficiency of distRL.




*
plainnat



























*relemmaLemma
*retheoremTheorem






§ APPENDIX









§ NOTATION











§ RELATED WORK AND DISCUSSION





 §.§ Technical Clarifications on Linearity Assumption in Existing Results




  
Bellman Closedness and Linearity.

<cit.> proved that quantile functional is not Bellman closed by providing a specific counterexample.

However, their discussion based on counterexamples can be generalized as it assumes that the sketch Bellman operator for the quantile functional needs to be linear.

They consider an discounted MDP with initial state s_0 with single action a, which transits to one of two terminal states s_1, s_2 with equal probability.
Letting no reward at state s_0, ([0,1]) at state s_1, and ([1/K , 1+ 1/K]) at state s_2, the return distribution at state s_0 is computed as mixture 1/2([0,γ]) + 1/2([γ/K , γ+ γ/K]).
Then the 1/2K-quantile at state s_0 is γ/K.
They proposed a counterexample where each quantile distribution of state s_1, s_2 is represented as 1/K∑_k=1^K δ_2k-1/K and 1/K∑_k=1^K δ_2k+1/K respectively, the 1/2K-quantile of state s_0 is 
ψ_q_2K( 1/2K∑_k=1^K δ_γ(2k-1)/K + δ_γ(2k+1)/K) = 3γ/2K.
However, this example does not consider that the mixture of quantiles is not a quantile of the mixture distribution 
(i.e., ψ_q(λη_1 + (1-λ) η_2) ≠λψ_q(η_1) + (1-λ) ψ_q(η_2)), due to the nonlinearity of the quantile functional. 
Therefore, this does not present a valid counterexample to prove that quantile functionals are not Bellman closed.





  
Bellman Optimizability and Linearity.

<cit.> proposed the notion of Bellman optimizable statistical functional which redefine the Bellman update by planning with respect to statistical functionals rather than expected returns.
They proved that W_1-continuous Bellman Optimizable statistical functionals are characterized by exponential utilities 1/λlog_Z ∼η [exp(λ Z)].
However, their proof requires some technical clarification regarding the assumption that such statistical functionals are linear.


To illustrate, they define a statistical functional ψ_f and consider two probability distributions η_1 = 1/2(δ_0 + δ_h) and η_2 = δ_ϕ(h) where ϕ(h) = f^-1(1/2(f(0) + f(h)) ).
Using the translation property, they lead ψ_f(η_1) = ψ_f(η_2) to
1/2(f(x)+f(x+h)) = f( x+ϕ(h)) for all x ∈ℝ.
However, this equality ψ_f (1/2 (δ_x + δ_x+h) ) = 1/2 (f(x) + f(x+h)) holds only if ψ_f is linear, which is not necessarily a valid assumption for all statistical functionals. 




 §.§ Existence of Nonlinear Bellman Closed Sketch.


The previous two examples may not have considered the possibility that the sketch Bellman operator might not necessarily be linear.
However, some statistical functionals are Bellman-closed even if they are nonlinear, so it is open question whether there is a nonlinear sketch Bellman operator that makes the quantile functional Bellman-closed. 
In this section, we present examples of maximum and minimum functionals that are Bellman-closed, despite being nonlinear.

In a nutshell, consider the maximum of return distribution at state s_1, s_2 is γ, γ + γ/K respectively.
Beyond linearity, the maximum of return distribution at state s_0 can be computed by taking the maximum of these values; 

    max( max(η̅(s_1)), max(η̅(s_2))) = max( γ, γ + γ/K) = γ + γ/K

which produces the desired result.
This implies the existence of a nonlinear sketch that is Bellman closed.
More precisely, by defining max_s' ∼(· | s,a) and min_s' ∼(· | s,a) as the maximum and minimum of the sampled sketch ψ( (ℬ_r)_#η̅(s') ) with the distribution (· | s,a), we can derive the sketch Bellman operator for maximum and minimum functionals as follows;


    
  * 𝒯_ψ_max( ψ_max(η̅(s)) )
    
    = max_s' ∼(· | s,a)( ψ_max((ℬ_r)_#η̅(s')) )
    
    
  * 𝒯_ψ_min( ψ_min(η̅(s)) )
    
    = min_s' ∼(· | s,a)( ψ_min((ℬ_r)_#η̅(s')) ).    










 §.§ Non-existence of sketch Bellman operator for quantile functional


In this section, we prove that quantile functional cannot be Bellman closed under any additional sketch.
First we introduce the definition of mixture-consistent, which is the property that the sketch of a mixture can be computed using only the sketch of the distribution of each component.


    A sketch ψ is mixture-consistent if for any λ∈ [0,1] and any distributions η_1, η_2 ∈𝒫_ψ(ℝ), there exists a corresponding function h_ψ such that
    
    ψ( λη_1 + (1-λ) η_2)
            = h_ψ( ψ(η_1), ψ(η_2), λ).




Next, we will provide some examples of determining whether a sketch is mixture-consistent or not.



  
Example 1.
    Every moment or exponential polynomial functional is mixture-consistent.
    
    For any n ∈ [N] and λ∈ℂ,
    
    _Z ∼λη_1 + (1-λ) η_2[Z^n exp(λ Z)]
            = λ_Z ∼η_1[Z^n exp(λ Z)] + (1-λ) _Z ∼η_2[Z^n exp(λ Z)].

    





  
Example 2.
    Variance functional is not mixture-consistent.
    
    Let λ = 1/2 and Z,Y be the random variables where Z ∼1/2δ_0 + 1/2δ_2 and Y ∼1/2δ_k + 1/2δ_k+2.
    Then, Var(Z) = Var(Y)=1.
    While RHS is constant for any k,
    LHS is not a constant for any k, i.e.,
    
    Var_X ∼1/2 (1/2δ_0 + 1/2δ_2) + 1/2 (1/2δ_k + 1/2δ_k+2) (X)
            
            
            = 1/4(k^2 + 5).

    
    
While variance functional is not mixture consistent by itself, it can be mixture consistent with another statistical functional, the mean.




  
Example 3.
    Variance functional is mixture-consistent under mean functional.
    
        Notice that mean functional is mixture-consistent. We need to show that variance functional is mixture-consistent under mean functional.

    
    Var_Z ∼λη_1 + (1-λ) η_2[Z]
            
        = _Z ∼λη_1 + (1-λ) η_2[Z^2] - (_Z ∼λη_1 + (1-λ) η_2[Z])^2 
            
        = λ_Z ∼η_1[Z^2] + (1-λ) _Z ∼η_2[Z^2] - (λ_Z ∼η_1[Z] + (1-λ) _Z ∼η_2[Z])^2
            
        = λ (Var_Z ∼η_1[Z] + (_Z∼η_1[Z])^2 ) + (1-λ) (Var_Z ∼η_2[Z] + (_Z∼η_2[Z])^2) 
            
            - (λ_Z ∼η_1[Z] + (1-λ) _Z ∼η_2[Z])^2.

    

This means that to determine whether it is mixture-consistent or not, we should check it on a per-sketch basis, rather than on a per-statistical functional basis.



  
Example 4.
    Maximum and minimum functional are both mixture-consistent.
    
    
    max_Z ∼λη_1 + (1-λ) η_2[Z] = max(max_Z ∼η_1[Z] , max_Z ∼η_2[Z])

    and
    
    min_Z ∼λη_1 + (1-λ) η_2[Z] = min(min_Z ∼η_1[Z] , min_Z ∼η_2[Z])

    

Since maximum and minimum functionals are mixture consistent, we can construct a nonlinear sketch bellman operator like the one in section <ref>.
This is possible because there is a nonlinear function h_ψ that ensures the sketch is closed under mixture.

Before demonstrating that a quantile sketch cannot be mixture consistent under any additional sketch, we will first illustrate with the example of a median functional that is not mixture consistent.



  
Example 5.
    Median sketch is not mixture-consistent.
    
    
    
    
    
    
    
    
    
    
    Let λ = 1/2 and Z,Y be the random variables where Z ∼ 0.2 δ_0 + 0.8 δ_1 and Y ∼ 0.6 δ_0 + 0.4 δ_k for some 0<k<1. 
    Then ψ_med(Z) = 1 and ψ_med(Y) = 0.
    However,
    
    med_X = Z+Y/2[X] = ψ_med(0.4 δ_0 + 0.2 δ_k + 0.4 δ_1)
            = k

    which is dependent in k.
    






    Quantile sketch cannot be mixture-consistent, under any additional sketch.    




    For a given integer N>0 and a quantile level α∈ [0,1], 
    let λ = 1/2 and a random variable Y ∼ p_y_0δ_0 + p_y_1δ_y_1 + ⋯ + p_y_Nδ_y_N (0 < y_1 < ⋯ < y_N < 1) where p_y_0 >α so that ψ_q_α(Y) = 0.
    Consider another random variable Z ∼ p_z_0δ_0 + p_z_1δ_1 where p_z_0 < α so that ψ_q_α(Z) = 1.
    Then the α-quantile of the mixture X = Y+Z/2 is 
    
    ψ_q_α[X] = y_n  where  n = min{ n ≤ N |  1/2∑_n'=0^n p_y_n' + 1/2p_z_0 > α}.

    Letting p_z_0 = 2 α - ∑_n'=0^n p_y_n', we can manipulate ψ_q_α[X] to be any value of y_n.
    Hence, ψ(q_α)[X] is a function of all possible outcomes of Y.

    If there exists a finite number of statistical functionals which make quantile sketch mixture-consistent, then such sketch would uniquely determine the distribution for any N.
    This results in a contradiction that infinite-dimensional distribution space can be represented by a finite number of statistical functional.







    If a sketch ψ is Bellman closed, then it is mixture-consistent.


    
        
        
        
        
        
        
        
        Consider an MDP where initial state s_0 has no reward and transits to two state s_1, s_2 with probability λ, 1-λ and reward distribution η̅_1, η̅_2.
        Since ψ is Bellman closed, ψ(η̅(s_0)) is a function of ψ(η̅(s_1)) and ψ(η̅(s_2)), (i.e., ψ(η̅(s_0)) = g_ψ(ψ(η̅(s_1)), ψ(η̅(s_2))) for some g_ψ).        
        
        
        
        
        
        
        
        Since ψ(η̅(s_0)) = ψ(λη̅(s_1) + (1-λ) η̅(s_2)), it implies that ψ is mixture-consistent.
    

Combining the results of Lemma <ref> and Lemma <ref>, we prove that a quantile sketch cannot be Bellman closed, no matter what additional sketches are provided.







§ PROOF




[<ref>]
    Quantile functional cannot be Bellman closed under any additional sketch.



    See Lemma <ref> and Lemma <ref>.




[<ref>]
    Let F_η̅ be a CDF of the probability distribution η̅∈𝒫(ℝ)^𝒮. 
    Then a sketch is Bellman unbiased if and only if the sketch is a homogeneous of degree k, i.e.,
    there exists some vector-valued function h = h(x_1, ⋯, x_k): 𝒳^k →ℝ^N such that
    
    ψ(η̅) =  ∫⋯∫ h(x_1, ⋯, x_k) dF_η̅(x_1) ⋯ dF_η̅(x_k).





    (⇒) 
    Consider an two-stage MDP with a single action a, and an initial state s_0 which transits to one of terminal state {s_1, ⋯, s_K } with transition kernel (· | s_0 ,a). Assume that the reward r(s_0) = 0. 
    Then η̅(s_0) = ∑_k=1^K (s_k) δ_r(s_k). 
    Note that s'_1 , ⋯, s'_k are independent and identically distributed random variable in distribution (·| s,a).

    
    _s'∼(· |s_0,a)[ 
            ϕ_ψ( ψ( (ℬ_r)_#η̅(s'_1) ), ⋯, ψ( (ℬ_r)_#η̅(s'_k))
            ) ]
            = ψ_1:N( (ℬ_r)_#_s' ∼(· | s_0,a) [ η̅(s') ])
            
       ⟹_s' ∼(· | s_0,a)[ ϕ_ψ( ψ( δ_r(s'_1)) , ⋯ , ψ( δ_r(s'_k))  ) ]
            = ψ( _s' ∼(· | s_0,a) [ δ_r(s') ])
            
       ⟹_s' ∼(· | s_0,a)[ ϕ_ψ(
            g(s'_1) , ⋯ , g(s'_k) 
            ) ]
            = ψ( η̅(s_0) )
            
       ⟹∫⋯∫ h(s'_1, ⋯ , s'_k) dF_η̅(s'_1) ⋯ dF_η̅(s'_k)
            = ψ( η̅(s_0) ).


    (⇐)
    
    ψ((ℬ_r)_#_s' ∼(· | s,a)[η̅(s')] )
        
        = ∫⋯∫ h(x_1, ⋯, x_k)dF_(ℬ_r)_#_s' ∼(· | s,a)[η̅(s')] (x_1) ,⋯, dF_(ℬ_r)_#_s' ∼(· | s,a)[η̅(s')](x_k)
        = ∫⋯∫ h(x_1 + r, ⋯, x_k + r) d(_s' ∼(·| s,a)F_η̅(s')(x_1) ) , ⋯, d(_s' ∼(·| s,a)F_η̅(s')(x_k) )
        
        = _s' ∼(· | s,a)[ ∫⋯∫ h(x_1 +r, ⋯, x_k+r) dF_η̅(s')(x_1) ⋯ dF_η̅(s')(x_k) ]
        
        = _s' ∼(·|s,a)[ ψ( (ℬ_r)_#[η̅(s')]
        ) ]


























































[<ref>]

    The only finite statistical functionals that are Bellman unbiased and closed are given by the collections of ψ_1, …,  ψ_N where its linear span {∑_n=0^N α_n ψ_n | α_n ∈ℝ ,∀ N} is equal to the set of exponential polynomial functionals {η→_Z ∼η[Z^l exp(λ Z)] |   l = 0, 1, …, L, λ∈ℝ}, where ψ_0 is the constant functional equal to 1.
    In discount setting, it is equal to the linear span of the set of moment functionals {η→_Z ∼η[Z^l] |  l = 0, 1, …, L } for some L ≤ N.  







    Our proof is mainly based on the proof techniques of <cit.> and we describe in an extended form.
    Since their proof also considers the discounted setting, we will define ℬ_r,γ(x) = r + γ x for discount factor γ∈ [0,1).
    By assumption of Bellman closedness,
    ψ_n( (ℬ_r,γ)_#η̅(s') ) will be written as g(r, γ, ψ_1:N(η̅(s')) for some g.
    By assumption of Bellman unbiasedness and Lemma <ref>,
    both ψ_1:N(η̅(s')) and ψ_n( (ℬ_r, γ)_#η̅(s') ) are affine as functions of the distribution η̅(s'),
    
    ψ_1:N( αη̅_1(s') + (1-α) η̅_2(s') )
            
        = 𝔼_Z_i ∼αη̅_1(s') + (1-α) η̅_2(s') [h_1:N(Z̅_1, ⋯, Z̅_k)]
            
        = α𝔼_Z̅_i ∼η̅_1(s') [h_1:N(Z̅_1, ⋯, Z̅_k)] + (1-α )𝔼_Z̅_i ∼η̅_2(s') [h_1:N(Z̅_1, ⋯, Z̅_k)]
            
        = αψ_1:N(η̅_1(s')) + (1-α) ψ_1:N(η̅_2(s'))

    and
    
    ψ_n( (ℬ_r,γ)_# (αη̅_1(s') + (1-α) η̅_2(s')) )
            
        = 𝔼_Z_i ∼αη̅_1(s') + (1-α) η̅_2(s') [h_n( r+ γZ̅_1, ⋯, r+ γZ̅_k)]
            
        = α𝔼_Z̅_i ∼η̅_1(s') [h_n(r+ γZ̅_1, ⋯, r+ γZ̅_k)] + (1-α )𝔼_Z̅_i ∼η̅_2(s') [h_n(r+ γZ̅_1, ⋯, r+ γZ̅_k)]
            
        = αψ_n((ℬ_r,γ)_#η̅_1(s') ) + (1-α) ψ_n( (ℬ_r,γ)_#η̅_2(s') )

    Therefore, g(r, γ , ·) is also affine on the convex codomain of ψ_1:N.
    
    Thus, we have
    
    _Z̅_i ∼η̅[ϕ_ψ_n(r+ γZ̅_1, ⋯, r+ γZ̅_k)] 
        = a_0(r, γ) + ∑_n'=1^N a_n'(r, γ) _Z̅_i ∼η̅[ϕ_ψ_n'(Z̅_1, ⋯, Z̅_k)]

    for some function a_0:N : ℝ× [0,1] →ℝ.
    By taking η̅(s') = δ_x, we obtain
    
    ϕ_ψ_n(r+ γ x, ⋯, r+ γ x) = a_0(r, γ) + ∑_n'=1^N a_n'(r, γ ) ϕ_ψ_n'(x, ⋯, x).

    According to <cit.>, for any translation invariant finite-dimensional space is spanned by a set of function of the form
    
    { x ↦ x^l exp(λ_j x) | j ∈ [J], 0 ≤ l ≤ L }

    for some finite subset {λ_1, ⋯, λ_J } of ℂ.
    Hence, each function x ↦ϕ_ψ_n(x, ⋯, x) is expressed as linear combination of exponential polynomial functions.
    In addition, the linear combination of ϕ_ψ_n should be closed under composition with for any discount factor γ∈ [0,1], all λ_j should be zero. 
    Hence, the linear combination of ϕ_ψ_1, ⋯, ϕ_ψ_N must be equal to the span of { x ↦ x^l | 0 ≤ l ≤ L } for some L ∈ℕ.
















[<ref>]
Consider a fixed k ∈ [K] and a fixed h ∈ [H]. Let 𝒵_h^k = { (s_h^τ , a_h^τ)}_τ∈ [k-1] and 

𝒟^k_h, η̅ = {( s_h^τ , a_h^τ,  ψ_1:N( (ℬ_r_h'^τ)_#η̅(s_h'+1^τ) )) }_τ∈ [k-1] for any η̅ : 𝒮→𝒫([0,H]).
Define f̃_h, η̅^k = min_f ∈ℱ^N f ^2_𝒟^k_h, η̅.
For any η̅ and δ∈ (0,1), there is an event ℰ(η̅, δ) such that conditioned on ℰ(η̅, δ), with probability at least 1-δ, for any η̅' : 𝒮→𝒫([0,H]) with 
ψ_1:N( η̅') - ψ_1:N(η̅) _∞,1≤ 1/T or ∑_n=1^N ψ_n (η̅')- ψ_n(η̅) _∞≤ 1/T, we have

    f̃_h,η̅'(· , ·) - ψ_1:N( (ℬ_r(·, ·))_#  [ℙη̅'](· , ·) ) _𝒵_h^k   ≤c' (N^1/2 H √(log(1/δ) + log𝒩(ℱ^N, 1/T)  ))

for some constant c' >0.





    Define the sketch of target f_η̅ : 𝒮×𝒜→ℝ^N,
    
    f_η̅(· , ·) ψ_1:N( (ℬ_r(·, ·))_#  [ℙη̅](· , ·) )

    for all i ∈ [N].

    
    
    
    
    
    
    
    
    
    For any f ∈ℱ,
    
    f ^2_𝒟^k_h,η̅' -  f_η̅'^2_𝒟^k_h,η̅'
       = ∑_n=1^N ∑_s_h^τ,a_h^τ∈𝒵^k_h,η̅'( f^(n)(s_h^τ, a_h^τ) - ψ_n ( (ℬ_r_h^τ)_#η̅'(s_h+1^τ) ) )^2 - 
            ( f_η̅'^(n)(s_h^τ, a_h^τ) - ψ_n ( (ℬ_r_h^τ )_#η̅'(s_h+1^τ) ) )^2
            
       = ∑_n=1^N ∑_ s_h^τ,a_h^τ∈𝒵^k_h,η̅' 
            ( f^(n)( s_h^τ,a_h^τ ) - f_η̅'^(n)( s_h^τ,a_h^τ ) )^2
            
            + 2 ( f^(n)( s_h^τ,a_h^τ ) - f_η̅'^(n)( s_h^τ,a_h^τ ) ) 
            ( f^(n)_η̅'( s_h^τ,a_h^τ ) - ψ_n ( (ℬ_r_h^τ)_#η̅'(s_h+1^τ) ) )
            
       ≥f - f_η̅'^2_𝒵_h^k- 4  ∑_n=1^N   f_η̅^(n) - f_η̅'^(n)_∞ (H+1)|𝒵_h^k|
            + 
            ∑_n=1^N ∑_ s_h^τ,a_h^τ∈𝒵^k_h,η̅'[ 2 ( f^(n)( s_h^τ,a_h^τ ) - f_η̅^(n)( s_h^τ,a_h^τ ) ) 
            ( f^(n)_η̅( s_h^τ,a_h^τ ) - ψ_n ( (ℬ_r_h^τ)_#η̅(s_h+1^τ) ) )_χ_h^τ(f^(n))]
            
       ≥f - f_η̅'^2_𝒵_h^k
            - 4 N(H+1) 
            - | ∑_n=1^N ∑_s_h^τ, a_h^τ∈𝒵^k_h, η̅'χ^τ_h(f^(n)) |.


    For the first inequality, we change the second term  from η̅' to η̅ which are the ϵ-covers.
    Notice that AC -BC' ≥ -|AC -BC'| ≥ -|(A-B)C| -|(A-B)C'| ≥ -2|A-B| |max(C,C')|.
    
    ( f^(n)( s_h^τ,a_h^τ ) - f_η̅'^(n)( s_h^τ,a_h^τ ) ) 
            ( f^(n)_η̅'( s_h^τ,a_h^τ ) - ψ_n ( (ℬ_r_h^τ)_#η̅'(s_h+1^τ) ) )
            
            - ( f^(n)( s_h^τ,a_h^τ ) - f_η̅^(n)( s_h^τ,a_h^τ ) ) 
            ( f^(n)_η̅( s_h^τ,a_h^τ ) - ψ_n ( (ℬ_r_h^τ)_#η̅(s_h+1^τ) ) )
            
       ≥ - 2  f_η̅'^(n)( s_h^τ,a_h^τ ) - f_η̅^(n)( s_h^τ,a_h^τ ) 
           ×max( 
            | f^(n)_η̅'( s_h^τ,a_h^τ ) - ψ_n ( (ℬ_r_h^τ)_#η̅'(s_h+1^τ) ) |,
            | f^(n)_η̅( s_h^τ,a_h^τ ) - ψ_n ( (ℬ_r_h^τ)_#η̅(s_h+1^τ) ) |
            )
            
       ≥ -2   f_η̅'^(n)( s_h^τ,a_h^τ ) - f_η̅^(n)( s_h^τ,a_h^τ )  (H+1)

    
    For the second inequality, consider η̅' : 𝒮→𝒫([0,H]) with ∑_n=1^N ψ_n( η̅' ) - ψ_n ( η̅ ) _∞≤ 1/T. 
    We have
    
    f^(n)_η̅ - f^(n)_η̅'_∞   = max_s,a| ∑_n'=1^n H^n' [ ψ_n' ( [η̅](s,a) ) - ψ_n' ([η̅'](s,a) ) ]r^n-n'/H^n-1|
            
       ≤∑_n'=1^n max_s'| ψ_n'( η̅(s') ) -ψ_n' ( η̅'(s') ) |
            
       ≤ 1/T.

    
    Defining 𝔽_h^k as the filtration induced by the sequence
{ (s_h'^τ ,a_h'^τ)}_τ, h' ∈ [k-1] × [H]∪{(s_1^k , a_1^k), (s_2^k, a_2^k), …, (s_h^k , a_h^k) },
    notice that
    
    [ ∑_n=1^N χ_h^τ(f^(n)) | 𝔽_h^τ]
            
       =∑_n=1^N 2 (f^(n)( s_h^τ,a_h^τ ) - f_η̅^(n)( s_h^τ,a_h^τ )) (f^(n)_η̅(s_h^τ, a_h^τ) - [ ψ_n ((ℬ_r_h^τ)_#η̅(s_h+1^τ) ) | 𝔽_h^τ])
            
       =∑_n=1^N 2  (f^(n)( s_h^τ,a_h^τ ) - f_η̅^(n)( s_h^τ,a_h^τ )) (f^(n)_η̅(s_h^τ, a_h^τ) - _s_h+1^τ∼_h(· | s_h^τ, a_h^τ )[ ψ_n ((ℬ_r_h^τ)_#η̅(s_h+1^τ) ) ])
            
       =∑_n=1^N 2 (f^(n)( s_h^τ,a_h^τ ) - f_η̅^(n)( s_h^τ,a_h^τ )) (f^(n)_η̅(s_h^τ, a_h^τ) -  ψ_n ((ℬ_r_h^τ)_#_s_h+1^τ∼_h(· | s_h^τ, a_h^τ )[ η̅(s_h+1^τ)] ) )
       =0

    and
    
    | ∑_n=1^N χ_h^τ(f^(n)) |
                = | ∑_n=1^N 2 (f^(n)( s_h^τ,a_h^τ ) - f_η̅^(n)( s_h^τ,a_h^τ )) (f_η̅^(n)(s_h^τ, a_h^τ) - ψ_n ( (ℬ_r_h^τ)_#η̅(s_h+1^τ) )  )|
            
            
       ≤max_n ∈ [N]{ 2 (f_η̅^(n)(s_h^τ, a_h^τ) - ψ_n ( (ℬ_r_h^τ)_#η̅(s_h+1^τ) )  ) }∑_n=1^N | f^(n)( s_h^τ,a_h^τ ) - f_η̅^(n)( s_h^τ,a_h^τ)  |
            
       ≤ 2(H+1) ∑_n=1^N | f^(n)( s_h^τ,a_h^τ ) - f_η̅^(n)( s_h^τ,a_h^τ)  |



    
    
    
    
    
    
    
  
    In third equality, we emphasize that only Bellman unbiased sketch can derive the martingale difference sequence which induce the concentration result.
    Since every moment functional is commutable with mixing operation, the transformation ϕ_ψ_n in Definition <ref> is identity for all n ∈ [N].
    Hence, we choose the sketch as moment which already knows ϕ_ψ.
    
    By Azuma-Hoeffding inequality,
    
    [ | ∑_(τ, h) ∈ [k-1] × [H] ∑_n=1^N χ_h^τ (f^(n)) | ≥ϵ]
               ≤ 2 exp( -ϵ^2/2(2(H+1))^2∑_(τ, h) ∈ [k-1] × [H]( ∑_n=1^N |f^(n) - f_η̅^(n)|)^2 )
            
       ≤ 2 exp( -ϵ^2/2(2(H+1))^2∑_(τ, h) ∈ [k-1] × [H]( N ∑_n=1^N |f^(n) - f_η̅^(n)|^2) )
            
        = 2 exp( -ϵ^2/ 2N(2(H+1))^2 f-f_η̅_𝒵_h^k^2 )

    where the second inequality follows from the Cauchy-Schwartz inequality.

    We set
    
    ϵ = √(8N(H+1)^2  f- f_η̅^2_𝒵_h^klog( 𝒩(ℱ^N, 1/T) /δ) )


    With union bound for all f ∈𝒞(ℱ^N , 1/T), with probability at least 1-δ,
    
    | ∑_(τ, h) ∈ [k-1] × [H] ∑_n=1^N χ_h^τ (f^(n)) | 
            ≤ c' N^1/2(H+1)  f- f_η̅_𝒵_h^k√(log( 𝒩(ℱ^N, 1/T) /δ))

    for some constant c'>0.

    
    For all f ∈ℱ^N, there exists g ∈𝒞(ℱ^N , 1/T), such that  f-g_∞,1≤ 1/T or ∑_n=1^N f^(n)-g^(n)_∞≤ 1/T for all n ∈ [N],
    
    | ∑_(τ, h) ∈ [k-1] × [H] ∑_n=1^N χ_h^τ (f^(n)) | 
               ≤| ∑_(τ, h) ∈ [k-1] × [H] ∑_n=1^N χ_h^τ (g^(n)) | + 2(H+1) |𝒵^k_h| ∑_n=1^N 1/T
       ≤ c' N^1/2 (H+1) g-f_η̅_𝒵^k_h√(log( 𝒩(ℱ^N,1/T)/δ)) + 2 N (H+1)
            
       ≤ c' N^1/2 (H+1)(f-f_η̅_𝒵^k_h +1) √(log( 𝒩(ℱ^N,1/T )/δ)) + 2 N (H+1)
            
       ≤ c' N^1/2 (H+1)(f-f_η̅'_𝒵^k_h +2) √(log( 𝒩(ℱ^N,1/T )/δ)) + 2 N (H+1)

    where the third inequality follows from,
    
    f- g _𝒵_h^k^2 
               ≤∑_n=1^N ∑_(τ, h) ∈ [k-1] × [H] |f^(n)(s_h^τ, a_h^τ) -g^(n)(s_h^τ, a_h^τ)|^2
            
       ≤ N T (1/T)^2
            
       ≤ 1.




    Recall that f̃^k_h,η' = min_f ∈ℱf^2_𝒟^k_h,η'.
    We have f̃^k_h,η'^2_𝒟^k_h,η' - f_η̅'^2_𝒟^k_h,η'≤ 0, which implies,
    
    0    ≥f̃^k_h,η̅'^2_𝒟^k_h,η̅' -  f_η̅'^2_𝒟^k_h,η̅'
       = f̃^k_h,η̅' - f_η̅'^2_𝒵^k_h
            + 2 ∑_n=1^N ∑_(τ, h) ∈ [k-1] × [H][( ( f̃^k_h, η̅')^(n)(s_h^τ, a_h^τ) - f^(n)_η̅'(s_h^τ, a_h^τ) )( ( f_η̅'^(n)(s_h^τ, a_h^τ) - ψ_n ((ℬ_r_h^τ)_#η̅'(s^τ_h+1) ) )]
            
       ≥f̃^k_h,η̅' - f_η̅'^2_𝒵^k_h 
            - c' N^1/2 (H+1) (f̂^k_h,η̅' - f_η̅'_𝒵^k_h +2) √(log(2/δ) + log𝒩(ℱ^N, 1/T)) - 6N(H+1).


    Recall that if x^2 -2ax -b ≤ 0 holds for constant a,b>0, then x ≤ a+ √(a^2 + b)≤ c' · a for some constant c'>0.
    
    Hence,
    
    f̃^k_h, η' - f_η̅'_𝒵^k_h≤ c' (N^1/2 H √(log (1/δ) + log𝒩(ℱ^N, 1/T)) )

    for some constant c' >0.








[<ref>]
    Let (ℱ^N)_h^k = { f ∈ℱ^N |  f - f̃^k_h, η̅^2_𝒵^k_h≤β(ℱ^N, δ) }, where
    
    β(ℱ^N,δ) ≥ c' · NH^2 (log (T/δ) + log𝒩(ℱ^N , 1/T) )

    for some constant c'>0.
    Then with probability at least 1- δ/2, for all k,h ∈ [K] × [H], we have
    
    ψ_n ( (ℬ_r_h(·, ·))_#[_h η̅_h+1^k](·, ·) ) ∈ (ℱ^N)_h^k




    For all (k,h) ∈ [K] × [H],
    
    
    
    
    
    
    𝐒{(min{ f^(1)(· , ·) 
        + b^k_h+1(·, ·) 
        , H }) |  f ∈𝒞(ℱ^N , 1/T) }∪{0 }    n=1 
    {(min{ f^(n)(· , ·) 
        , H }) | f ∈𝒞(ℱ^N , 1/T) }∪{0 }    2 ≤ n ≤ N

    
    
    
    is a (1/T)-cover of ψ_1:N (η^k_h+1(·, ·)) where
    
    ψ_1:N ( η^k_h+1(·, ·) ) 
        = min{( f^k_h+1)^(1)(· , ·) 
        + b^k_h+1(·, ·) 
        , H }    n=1  and  h<H 
    min{( f^k_h+1)^(n)(· , ·) 
        , H }    2 ≤ n ≤ N  and  h<H 
    0^N     h=H
        
        ,
 
    
    i.e., there exists ψ_1:N (η) ∈𝐒 such that ψ_1:N(η) - ψ_1:N(η^k_h+1) _∞,1≤ 1/T.
    This implies
    
    𝐒̅{ψ_1:N( η(·, max_a ∈𝒜ψ_1 (η(·, a)) ) ) |ψ_1:N( η ) ∈𝐒}

    is a (1/T)-cover of ψ_1:N( η̅^k_h+1) with log(|𝐒̅|) ≤log𝒩(ℱ^N , 1/T).

    For each ψ_1:N( η̅) ∈𝐒̅, let ℰ(η̅ , δ/2|𝐒̅|T) be the event defined in Lemma <ref>.
    By union bound for all ψ_1:N (η̅) ∈𝐒̅, we have Pr[⋂_ψ_1:N(η̅) ∈𝐒̅ℰ(η̅, δ / 2|𝐒̅|T)] ≥ 1- δ/2T.

    Let ψ_1:N(η̅) ∈𝐒̅ such that ψ_1:N( η̅ ) - ψ_1:N( η̅_h+1^k ) _∞,1≤ 1/T.
    Conditioned on ⋂_s_N(η̅) ∈𝐒̅ℰ(η̅, δ / 2| 𝐒̅|T) and by Lemma <ref>, we have
    
    f̃^k_h,η̅ (·,·) - ψ_1:N( (ℬ_r_h(·,·))_# [_h η̅_h+1^k](·,·)) ^2_𝒵^k_h≤ c' ( N H^2 (log (T/δ) + log𝒩(ℱ^N , 1/T) ) )

    for some constant c' >0.
    
    By union bound for all (k,h) ∈ [K] × [H], we have ψ_1:N((ℬ_r_h(·,·))_#[_h η̅_h+1^k](·, ·)) ∈ (ℱ^N)_h^k with probability 1-δ/2.











    Let Q^k_h(s,a) min{ H, f̃^k_h(s,a) + b^k_h(s,a) } for some bonus function b^k_h(s,a) for all (s,a) ∈𝒮×𝒜. 
    If b^k_h(s,a) ≥ w^(1)((ℱ^N)^k_h ,s,a) ,
    then with probability at least 1-δ/2,  
    
    Q^*_h (s, a ) ≤ Q^k_h (s, a) and  V^*_h (s) ≤ V^k_h (s)

    for all (k,h) ∈ [K]×[H], for all (s,a) ∈𝒮×𝒜.










 
    We use induction on h from h=H to 1 to prove the statement.
    Let ℰ be the event that for (k,h) ∈ [K] × [H], ψ_1:N( (ℬ_r_h (· , ·))_# [_h η̅^k_h+1](·,·) ) ∈ (ℱ^N)^k_h.
    By Lemma <ref>, Pr|ℰ| ≥ 1-δ/2.
    In the rest of the proof, we condition on ℰ.
    
    When h=H+1, the desired inequality holds as Q^*_H+1(s,a) = V^*_H+1(s) = Q^k_H+1(s,a) = V^k_H+1(s)=0.
    Now, assume that Q^*_h+1(s,a) ≤ Q^k_h+1(s,a) and V^*_h+1(s) ≤ V^k_h+1(s) for some h ∈ [H].
    Then, for all (s,a) ∈𝒮×𝒜,
    
    Q^*_h(s,a)     = min{ H, r_h(s,a) + [_h V^*_h+1](s,a) }
       ≤min{ H, r_h(s,a) + [_h V^k_h+1](s,a) }
       ≤min{ H, f̃^k_h(s,a) + w^(1)(ℱ^k_h , s, a) }
        = min{ H, Q^k_h(s,a) - b^k_h(s,a) + w^(1)(ℱ^k_h , s, a) }
       ≤ Q^k_h(s,a)








    With probability at least 1-δ/4, we have
    
    
    Reg(K) ≤∑_k=1^K ∑_h=1^H (2 b^k_h(s^k_h, a^k_h) + ξ^k_h) ,

    where ξ^k_h = [_h (V^k_h+1 - V^π^k_h+1)](s^k_h, a^k_h) -(V^k_h+1(s^k_h+1) - V^π^k_h+1(s^k_h+1) ) is a martingale difference sequence with respect to the filtration 𝔽^k_h induced by the history ℍ^k_h.












    We condition on the above event ℰ in the rest of the proof.
    For all (k,h) ∈ [K] × [H], we have 
    
    f̃^k_h,η̅ (·,·) -ψ_1:N( (ℬ_r_h(·,·))_# [_h η̅_h+1^k](·,·)) ^2_𝒵^k_h≤β(ℱ^N , δ).

    Recall that (ℱ^N)_h^k = { f ∈ℱ^N | f - f̃^k_h,η̅^2_𝒵^k_h≤β(ℱ^N, δ) } is the confidence region.
    Since ψ_1:N( (ℬ_r_h(·,·))_# [_h η̅_h+1^k](·,·) ) ∈ (ℱ^N)^k_h, then by the definition of width function w^(1)(ℱ^k_h , s,a), for (k,h) ∈ [K] × [H], we have
    
    w^(1)(ℱ^k_h , s,a)
               ≥| ψ_1((ℬ_r_h(s,a))_#[_h η̅_h+1^k](s,a) ) - (f̃^k_h,η̅)^(1)(s,a)|
            
        = | r_h(s,a) + [_h V^k_h+1](s,a) - (f̃^k_h,η̅)^(1)(s,a) |.

    
    
    
    
    
    
    Recall that Q^*_h (·, · ) ≤ Q^k_h (·, ·).
    
    Reg(K)    = ∑_k=1^K V_1^⋆(s_1^k) - V_1^π^k(s_1^k)
            
       ≤∑_k=1^K V_1^k(s_1^k) - V_1^π^k(s_1^k)
            
        = ∑_k=1^K Q_1^k(s_1^k , a_1^k) - Q_1^π^k(s_1^k, a_1^k)
            
        = ∑_k=1^K Q_1^k(s_1^k , a_1^k) - (r_1(s_1^k, a_1^k) 
            + [_1 V_2^k](s_1^k, a_1^k)) + (r_1(s_1^k, a_1^k) + [_1 V_2^k](s_1^k, a_1^k))
            
           - Q_1^π^k(s_1^k, a_1^k)
            
       ≤∑_k=1^K w^(1)((ℱ^N)^k_1, s_1^k, a_1^k)  + b^k_1(s_1^k, a_1^k) + [_1 (V_2^k - V_2^π^k)](s_1^k, a_1^k)
            
       ≤∑_k=1^K w^(1)((ℱ^N)^k_1, s_1^k, a_1^k)  + b^k_1(s_1^k, a_1^k) + (V_2^k(s_2^k) - V_2^π^k(s_2^k) ) + ξ_1^k 
            
           ⋮
       ≤∑_k=1^K ∑_h=1^H ( w^(1)((ℱ^N)^k_h, s_h^k, a_h^k)  + b^k_h(s_h^k, a_h^k) + ξ_h^k )
            
       ≤∑_k=1^K ∑_h=1^H (2 b^k_h(s_h^k, a_h^k) + ξ_h^k )






It remains to bound ∑_k=1^K ∑_h=1^H b^k_h(s^k_h, a^k_h), for which we will exploit fact that ℱ^N has bounded eluder dimension.









    If b^k_h(s,a) ≥ w^(1)((ℱ^N)^k_h, s,a ) for all (s,a) ∈𝒮×𝒜 and k ∈ [K] where 
    
    (ℱ^N)^k_h = { f ∈ℱ^N |  f-f̃^k_h, η̅^2_𝒵^k_h≤β(ℱ^N, δ)},
 then
    
    
    
    
    
    ∑_k=1^K ∑_h=1^H 1{b^k_h( s^k_h,a^k_h) > ϵ}≤(4β(ℱ^N, δ)/ϵ^2 +1) dim_E(ℱ^N ,ϵ)

    for some constant c>0. 



























We first want to show that for any sequence {(s_1, a_1), … ,(s_κ, a_κ) }⊆𝒮×𝒜,
there exists j ∈ [κ] such that (s_j,a_j) is ϵ-dependent on at least L = ⌈ (κ-1)/dim_E(ℱ^N, ϵ) ⌉ disjoint subsequences in {(s_1, a_1), …, (s_j-1, a_j-1) } with respect to ℱ^N. 
We demonstrate this by using the following procedure.
Start with L disjoint subsequences of { (s_1, a_1), …, (s_j-1, a_j-1) }, ℬ_1, ℬ_2, …, ℬ_L, which are initially empty.
For each j, if (s_j , a_j) is ϵ-dependent on every ℬ_1, …, ℬ_L, we achieve our goal so we stop the process.
Else, we choose i ∈ [L] such that (s_j, a_j) is ϵ-independent on ℬ_i and update ℬ_i ←ℬ_i ∪{(s_j,a_j)}, j ← j+1.
Since every element of ℬ_i is ϵ-independent on its predecessors, |ℬ_i| cannot get bigger than dim_E(ℱ^N, ϵ) at any point in this process.
Therefore, the process stops at most step j = L dim_E(ℱ^N, ϵ) + 1 ≤κ.

Now we want to show that if for some j ∈ [κ] such that 

b^k_h(s_j , a_j) > ϵ,
then (s_j , a_j) is ϵ-dependent on at most 4β(ℱ^N , δ)/ϵ^2 disjoint subsequences in {(s_1, a_1), …, (s_j-1, a_j-1) } with respect to ℱ^N.
If 

b^k_h(s_j , a_j) > ϵ
and (s_j, a_j) is ϵ-dependent on a subsequence of {(s'_1, a'_1), …, (s'_l, a'_l)}⊆{(s_1, a_1), …, (s_κ, a_κ)}, 
it implies that there exists f,g ∈ℱ^N with f - f̃^k_h,η̅^2_𝒵^k_h≤β(ℱ^N,δ) and g - f̃^k_h,η̅^2_𝒵^k_h≤β(ℱ^N,δ) such that


f^(1)( s'_t , a'_t ) - g^(1)( s'_t , a'_t )  ≥ϵ.
By triangle inequality, f - g ^2_𝒵^k_h≤ 4 β(ℱ^N,δ).
On the other hand, if (s_j, a_j) is ϵ-dependent on L disjoint subsequences in {(s_1, a_1), …, (s_κ, a_κ)}, then

    4β(ℱ^N , δ ) 
        
        ≥ f-g ^2_𝒵^k≥f^(1) - g^(1)^2_𝒵^k≥ L ϵ^2

resulting in L ≤ 4 β(ℱ^N, δ)/ϵ^2.
Therefore, we have (κ/ dim_E(ℱ^N, ϵ))-1 ≤ 4β(ℱ^N, δ)/ϵ^2 which results in

    κ≤( 4β(ℱ, δ)/ϵ^2+1) dim_E(ℱ^N, ϵ)






























































    
    If b^k_h(s,a) ≥ w^(1)((ℱ^N)^k_h, s,a ) for all (s,a) ∈𝒮×𝒜 and k ∈ [K], then
    
    ∑_k=1^K ∑_h=1^H b^k_h(s^k_h, a^k_h ) 
        ≤ H dim_E(ℱ^N, 1/T).











    We first sort the sequence { b^k_h(s^k_h, a^k_h ) }_(k,h) ∈ [K]×[H] in a decreasing order and denote it by {e_1 , …, e_T } (e_1 ≥ e_2 ≥…≥ e_T).
    
    
    
    
    
    By Lemma <ref>, for any constant M >0 and e_t ≥ 1/√(M) T, we have
    
    t ≤( 4β(ℱ^N,δ)/M e_t^2 +1)dim_E(ℱ^N, √(M) e_t)
            ≤( 4β(ℱ^N,δ)/M e_t^2 +1)dim_E(ℱ^N, 1/T)

    which implies
    
    e_t ≤( t/dim_E(ℱ^N , 1/T)  -1 )^-1/2√(4 β(ℱ^N, δ)/M),

    for t ≥dim_E(ℱ^N ,1/T).
    Since we have e_t ≤ H,
    
    ∑_t=1^T e_t 
                = ∑_t=1^T e_t 1{ e_t < 1/√(M)T } 
            + ∑_t=1^T e_t 1{ e_t ≥ 1/√(M)T, t < dim_E(ℱ^N ,1/T) }
          + ∑_t=1^T e_t 1{ e_t ≥ 1/√(M)T , t ≥dim_E(ℱ^N ,1/T) }
       ≤1/√(M) + H dim_E(ℱ^N ,1/T) + ∑_dim_E(ℱ^N ,1/T) ≤ t ≤ T( t/dim_E(ℱ^N , 1/T)  -1 )^-1/2√(4 β(ℱ^N, δ)/M)
       ≤1/√(M) + H dim_E(ℱ^N ,1/T) + 2 ( T/dim_E(ℱ^N ,1/T) -1)^1/2dim_E(ℱ^N ,1/T) √(4 β(ℱ^N, δ)/M)
        =  1/√(M) + H dim_E(ℱ^N ,1/T) +  √(16 ·dim_E(ℱ^N ,1/T) · T ·β(ℱ^N , δ)/M).

    Taking M →∞,
    
    ∑_k=1^K ∑_h=1^H b^k_h(s^k_h, a^k_h ) 
        ≤ H dim_E(ℱ^N, 1/T).

    
























[<ref>]

    Under Assumption <ref>, with probability at least 1- δ,  achieves a regret bound of
    
    Reg(K) ≤ 
        2Hdim_E(ℱ^N, 1/T) + 4H√(KH log(2/δ)).

    
    
    



Recall that ξ_h^k = [_h(V^k_h+1 - V^π^k_h+1)](s^k_h, a^k_h) - ( V^k_h+1(s^k_h+1) - V^π^k_h+1(s^k_h+1)) is a martingale difference sequence where [ξ^k_h | 𝔽^k_h] = 0 and | ξ^k_h| ≤ 2H.
By Azuma-Hoeffding's inequality, with probability at least 1-δ/2,

    ∑_k=1^K∑_h=1^H ξ^k_h ≤ 4H√(KH log(2/δ)).


Conditioning on the above event and Lemma <ref>, we have

    
    Reg(K)    ≤ 2 ∑_k=1^K ∑_h=1^H b^k_h(s^k_h, a^k_h) + ∑_k=1^K ∑_h=1^H ξ^k_h 
            
       ≤ 2Hdim_E(ℱ^N, 1/T) + 4H√(KH log(2/δ))

    






























 




























































































































    

























    


















































































