Paper Title: Voxel Scene Graph for Intracranial Hemorrhage.






Antoine Sanner
Nils Grauhan
Marc Brockmann
Ahmed Othman
Anirban Mukhopadhyay

A. Sanner et al.
Department of Computer Science, Technical University of Darmstadt, Germany 
 
Department of Neuroradiology, University Medical Center Mainz, Germany 

antoine.sanner@gris.tu-darmstadt.de


Voxel Scene Graph for Intracranial Hemorrhage
    
Antoine P. Sanner1,2 0000-0002-4917-9529 ()  Nils F. Grauhan2 Marc A. Brockmann2 Ahmed E. Othman2 Anirban Mukhopadhyay1

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



§ ABSTRACT 

Patients with Intracranial Hemorrhage (ICH) face a potentially life-threatening condition, and patient-centered individualized treatment remains challenging due to possible clinical complications. 
Deep-Learning-based methods can efficiently analyze the routinely acquired head CTs to support the clinical decision-making. The majority of early work focuses on the detection and segmentation of ICH, but do not model the complex relations between ICH and adjacent brain structures. 
In this work, we design a tailored object detection method for ICH, which we unite with segmentation-grounded Scene Graph Generation (SGG) methods to learn a holistic representation of the clinical cerebral scene. 
To the best of our knowledge, this is the first application of SGG for 3D voxel images. 
We evaluate our method on two head-CT datasets and demonstrate that our model can recall up to 74% of clinically relevant relations. This work lays the foundation towards SGG for 3D voxel data.
The generated Scene Graphs can already provide insights for the clinician, but are also valuable for all downstream tasks as a compact and interpretable representation.






§ INTRODUCTION



Intracranial Hemorrhage (ICH) is a potentially life-threatening emergency in which rapid recognition and decision on further treatment is crucial for patient survival <cit.>. The 1-year mortality rate is up to 40% in various studies <cit.> and survivors retain significant persistent functional limitations in two thirds of cases <cit.>. Patient-centered individualized treatment decisions remain clinically challenging, especially considering that patient outcome deteriorates rapidly after ICH onset <cit.>. Given that CT imaging remains the main tool for diagnosis and planning, Deep Learning-based solutions can leverage this data to improve ICH patients' care, e.g. reducing the time to diagnosis.

The majority of early Deep Learning work focuses on the detection and segmentation of ICH <cit.>, but completely ignore the related clinical complications. Indeed, hemorrhages can interact with other brain structures. Through hemorrhage expansion or involvement of the ventricular system, ICH can cause severe complications <cit.>, which are therefore additional predictors for patient outcome <cit.>. Likewise, a bleeding-induced midline shift is a strong predictor of poor patient outcome <cit.>. However, not all bleedings are relevant, such as smaller ones that do not interact with other anatomies. As such, the true potential of automated solutions lies in modeling and understanding complex relations between ICH and adjacent brain structures to build a clinical cerebral scene, that can be then used for downstream tasks such as patient outcome prediction.




Scene Graph Generation (SGG), originally introduced for natural scene-understanding <cit.>, has already been applied in surgical data science <cit.> to model surgical procedures in the operating room. SGG techniques can likewise show their potential, when used with voxel data such as 3D CT scans. Such an application to cerebral scene for ICH is shown in Fig. <ref>. Scene Graphs offer a compact and interpretable representation, that goes beyond pure object detection and can already provide insights for the clinician.
Additionally, an advantage of Scene Graphs is built-in explainability for downstream tasks. 


SGG is a two-step process, which first involves localizing relevant objects in an image. Inspired by Retina U-Net <cit.>, we design a hybrid method that leverages both bounding box annotation and semantic segmentation to both detect objects robustly, and to overcome the challenges of 3D object matching (i.e. orientation-dependent aspect ratios, and significant object overlap). Then, relations between pairs of objects are predicted with SGG methods. Inspired by Neural Motifs <cit.>, and Iterative Message Passing <cit.>, we design two variants for relation prediction for 3D voxel data (V-MOTIF and V-IMP). These use Recurrent Neural Networks to build a global context and from which the relation of each bleeding to adjacent brain structures is predicted.

We show that our model can recall up to 74% of clinically relevant relations. Our contributions include:

    
  * To the best of our knowledge, this is the first time Scene Graphs are generated from 3D voxel data. Scene Graphs can already provide insights for the clinician, independently of their added benefits for downstream tasks.
    
    
  * A tailored object detection method[Code available at <https://github.com/MECLabTUDA/VoxelSceneGraph>] for ICH, which outperforms the state-of-the-art nnDetection <cit.>, and which we unite with our voxel-adapted SGG methods to learn a holistic representation of the clinical cerebral scene.





§ METHODOLOGY


In this section, we give insights into how we structured Scene Graphs for ICH. We then introduce our two-stage method for SGG as visualized in Fig. <ref>.







Scene Graph for ICH: The first and most critical step is to define which structures are relevant, as well as a set of relations that can model possible sources of complications. We select the following interactions for our application based on possible clinical complications: 1) midline shifts, 2) blood flow to the ventricle system, and 3) swelling induced asymmetry of the ventricle system. As such, we have three classes of objects ("Bleeding", "Ventricle System" and "Midline"), and three bleeding-induced classes of relations. Additionally, we expect all patients to have exactly one midline. We model the ventricle system as a single object even when it is only sparsely visible, since the presence of blood can have severe complications independently of the location within the system.

Object Detection: Localizing individual bleeding requires a multiscale approach, as they can have vastly different size, shape, and position. The 3D Retina-UNet coupled with a Feature Pyramid Network has already proved its usefulness for 3D medical imaging <cit.>. In particular, this architecture is very flexible and can fully leverage multiscale features. This is crucial as the volume of the bleeding can range from only 0.1^3 to more than 100^3. 
In opposition, the midline and the ventricle system only appear at a single scale, but pose their own set of challenges. 
Indeed, the 3D aspect ratio of the midline will heavily depend on the head's orientation, especially in the axial plane. This both can cause issues with anchor matching and bounding box regression.
In contrast, depending on the swelling or the presence of blood, the ventricle system may appear as multiple fragmented objects. These would need to be detected individually, as there would be a large overlap between the bleeding with the system and the system itself, rendering an accurate object matching impossible. For these reasons, we decide to leverage Retina-UNet's segmentation capabilities by detecting both anatomies from the predicted semantic segmentation. This simple but elegant solution, can solve the previously mentioned challenged, while also ensuring that only one object is predicted per class per image.

Relation Prediction: Once objects are localized in the image, SGG methods focus on building a global context from individual object features (as defined by <cit.>). Inspired by Neural Motifs <cit.>, and Iterative Message Passing <cit.>, we design two variants for relation prediction for 3D voxel data (V-MOTIF and V-IMP). 

V-MOTIF uses bidirectional Long Short-term Memory Networks and iterates over detected objects. We choose to order these objects from top-to-bottom, as related objects in the cerebral scene are very likely to be at the same depth. Sorting by size, may be another viable solution, but may create a bias towards certain relation classes.
In comparison, V-IMP builds a primal object graph, and its dual edge graph, then combines Gated Recurrent Units (GRU) with message passing iteratively to allow for information flow. Since each GRU receives multiple messages at each iteration, learnable weights are used to pool messages while only keeping the relevant information.
Both models finally predict relations from the enriched features.
Usually, these are also used to refine the classification of each object. 
Given the object matching challenges discussed previously, we decide to only use SGG methods for relation prediction. 

Finally, segmentation-grounding <cit.> involves computing the object features with the object segmentation, rather than with its bounding box, to leverage the finer localization. 
Since only a semantic segmentation in computed in the first step, we use it as a proxy by cropping it to the object's bounds and binarizing it based on the object's predicted class. 



§ EXPERIMENTS



In this section, we introduce the base images used and their annotation process. We then describe our evaluation setup for the different stages of our method.





Source Images: For this study, we source two datasets. The first one is the publicly available INSTANCE2022 challenge dataset <cit.>. 
It contains 130 non-contrast head CTs of patients diagnosed with ICH. 10 images with significant streak artifacts were discarded, as they were deemed out-of-distribution. 
Additionally, we select a private cohort for the purpose of external validation. It is constituted of 18 non-contrast head CTs of patients diagnosed with ICH. We choose more severe cases (caused by trauma) compared to the first dataset.
These cases are clinically more challenging, e.g. some patients were already operated on, and we want to test the robustness of our method against distribution shift (see Fig. <ref>). 




Data Annotation: The patient's head position is first manually harmonized in all images. A senior neuroradiologist then segments a label map for each image using 3D Slicer <cit.>, which takes 5 to 30 min per image. 3D bounding box information is automatically extracted from the label maps by computing the bounds of each segmented object.
Currently, one of the main obstacles of Scene Graph application using voxel data is the lack of tools for relation annotation. In particular, the annotator needs to be able to easily scroll through the 3D volume to select individual objects as being part of a relation. We developed an internal tool specifically to this end. We plan to make it open-source in the near future. On average, annotating relations takes under 2 min per image.


Model Training: For bounding box detection, we split the official training images from the INSTANCE2022 dataset <cit.> for training and validation in an 80/20 fashion. The official validation cases are used for in-distribution testing. The private cohort is fully used for external validation. However, not all images contain relations. As such, we only keep cases with relations for training the relation detector, the splits remaining otherwise unchanged. Configuration files with hyperparameters and  data splits will be made available with the code.

Evaluation Metrics: For bleeding detection, we use the well established metrics for object detection, i.e. Average Recall (AR) and mean Average Precision (mAP). We choose a localization threshold of 30% Intersection Over Union (IoU), as it reflects a balance between the clinical need for coarse localization and producing sensible bounding boxes for larger bleeding <cit.>. 
Additionally, since no relation can get predicted when at least one object has not been detected, we also give an upper bound of how many relations can be recalled given the object detected in the first stage. 
For relation prediction, we compute the metrics based on the top-K predictions per image, with again a 30% IoU threshold for object localization. We use the Recall@K (R@K), mean Recall@K (mR@K), and mean Average Precision@K (mAP@K). Given that an image has on average overall 2 relations (up to 7), we compute these metrics for K=8. Our method is evaluated for both Predicate Classification and Scene Graph Generation tasks <cit.>, i.e. respectively predicting relations from ground truth and predicted object localization. The results are averaged over 5 random seeds.






§ RESULTS



In this section, we evaluate our two-stage pipeline, first for object detection and then for relation prediction (Predicate Classification and Scene Graph Generation). Results for ablations of individual method components can be found in the Supplementary Material.



 §.§ Object Detection



Ventricle System & Midline Detection: We first evaluate the performance for detecting the ventricle system and the midline, as these structures are localized using the predicted semantic segmentation. As shown in Fig. <ref>, our solution can detect them robustly and precisely, with a detection rate of 96.4% for both structures in the INSTANCE2022 dataset. For the private cohort, it still has a detection rate of 100% and 83.3% respectively for the ventricle system and the midline.








Bleeding Detection: We then evaluate the detection performance for bleeding separately from the previously mentioned structures. In Table <ref>, we compare our method to state-of-the-art nnDetection <cit.>. Our method significantly outperforms it on both datasets, showing that careful method parametrization can outperform AutoML for this task. 
The private cohort contains proportionally many smaller bleedings, which are not present in the training data and predominantly have no associated relation. In this regard, our method still robustly detects relevant bleedings in the private cohort.






 §.§ Scene Graph Prediction





Predicate Classification: We evaluate our Voxel Neural Motifs (V-MOTIF) and Voxel Iterative Message Passing (V-IMP) models for relation prediction, with and without segmentation grounding and report the results in Table <ref>. 
Given that datasets for natural scene-understanding <cit.> or even surgical data science  <cit.> comprise thousands of scenes, one could assume that a large data quantity is required for training SGG methods. These results show that under 200 scenes are enough to generalize across medical centers for the ICH cerebral scene.





Scene Graph Generation: This experiment reflects the true performance of the entire prediction pipeline (Table <ref>). Qualitative results are shown in Fig. <ref>. Overall, the models recall relevant relations, with clinically interpretable information for decision support in treatment planning. However, this experiment shows that our complete method already offers a satisfying performance for in-distribution data.  










§ CONCLUSION


The clinical scene for Intracranial Hemorrhage (ICH) goes beyond detecting bleedings, as these can interact with adjacent brain structures, and potentially cause severe complications.
We introduce a method for structured representation learning and tailored for ICH. To the best of our knowledge, this is the first application of Scene Graph Generation (SGG) for 3D voxel data. 
We evaluate our method on two datasets and we demonstrate that our model can recall up to 74% of clinically relevant relations.
This work paves a path towards SGG for 3D voxel images and modeling holistic anatomical scenes.
The generated Scene Graphs can be used as a compact and interpretable representation for all downstream task (e.g. patient outcome prediction) and offer a strong decision support for patient-centered treatment.


§ COMPLIANCE WITH ETHICAL STANDARDS


This study was performed in line with the principles of the Declaration of Helsinki. The retrospective evaluation of imaging data from the University Medical Center Mainz was approved by the local ethics boards (Project 2021-15948-retrospektiv). Ethical approval was not required, as confirmed by the license attached with the open access data.






  §.§.§ 
 Antoine P. Sanner and Nils F. Grauhan equally contribute to this paper.



  §.§.§ 

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












splncs04




 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: Maverick: Efficient and Accurate Coreference Resolution Defying Recent Trends.



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

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

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

Xin Yang1,2()

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



§ ABSTRACT 
Large autoregressive generative models have emerged as the cornerstone for achieving the highest performance across several Natural Language Processing tasks. 
However, the urge to attain superior results has, at times, led to the premature replacement of carefully designed task-specific approaches without exhaustive experimentation. 
The Coreference Resolution task is no exception; all recent state-of-the-art solutions adopt large generative autoregressive models that outperform encoder-based discriminative systems. 
In this work, we challenge this recent trend by introducing , a carefully designed – yet simple – pipeline, which enables running a state-of-the-art Coreference Resolution system within the constraints of an academic budget, outperforming models with up to 13 billion parameters with as few as 500 million parameters.
achieves state-of-the-art performance on the CoNLL-2012 benchmark, training with up to 0.006x the memory resources and obtaining a 170x faster inference compared to previous state-of-the-art systems.
We extensively validate the robustness of the framework with an array of diverse experiments, reporting improvements over prior systems in data-scarce, long-document, and out-of-domain settings.
We release our code and models for research purposes at .




§ INTRODUCTION

As one of the core tasks in Natural Language Processing, Coreference Resolution aims to identify and group expressions (called mentions) that refer to the same entity <cit.>. 
Given its crucial role in various downstream tasks, such as Knowledge Graph Construction <cit.>, Entity Linking <cit.>, Question Answering <cit.>, Machine Translation <cit.> and Text Summarization <cit.>, inter alia, there is a pressing need for both high performance and efficiency.
However, recent works in Coreference Resolution either explore methods to obtain reasonable performance optimizing time and memory efficiency <cit.>, or strive to improve benchmark scores regardless of the increased computational demand <cit.>.

Efficient solutions usually rely on discriminative formulations, frequently employing the mention-antecedent classification method proposed by <cit.>.
These approaches leverage relatively small encoder-only transformer architectures <cit.> to encode documents and build on top of them task-specific networks that ensure high speed and efficiency.
On the other hand, performance-centered solutions are nowadays dominated by general-purpose large Sequence-to-Sequence models <cit.>.


A notable example of this formulation, and currently the state of the art in Coreference Resolution, is <cit.>, which proposes a transition-based system that incrementally builds clusters of mentions by generating coreference links sentence by sentence in an autoregressive fashion. 
Although Sequence-to-Sequence solutions achieve remarkable performance, their autoregressive nature and the size of the underlying language models (up to 13B parameters) make them dramatically slower and memory-demanding compared to traditional encoder-only approaches. 
This not only makes their usage for downstream applications impractical, but also poses a significant barrier to their accessibility for a large number of users operating within an academic budget.

In this work we argue that discriminative encoder-only approaches for Coreference Resolution have not yet expressed their full potential and have been discarded too early in the urge to achieve state-of-the-art performance.
In proposing  , we strike an optimal balance between high performance and efficiency, a combination that was missing in previous systems.
Our framework enables an encoder-only model to achieve top-tier performance while keeping the overall model size less than one-twentieth of the current state-of-the-art system, and training it with academic resources.
Moreover, when further reducing the size of the underlying transformer encoder, performs in the same ballpark as encoder-only efficiency-driven solutions while improving speed and memory consumption.
Finally, we propose a novel incremental Coreference Resolution method that, integrated into the framework, results in a robust architecture for out-of-domain, data-scarce, and long-document settings.





§ RELATED WORK

We now introduce well-established approaches to neural Coreference Resolution.
Specifically, we first delve into the details of traditional discriminative solutions, including their incremental variations, and then present the recent paradigm shift for approaches based on large generative architectures.



 §.§ Discriminative models

Discriminative approaches tackle the Coreference Resolution task as a classification problem, usually employing encoder-only architectures.
The pioneering works of <cit.> introduced the first end-to-end discriminative system for Coreference Resolution, the Coarse-to-Fine model.
First, it involves a mention extraction step, in which the spans most likely to be coreference mentions are identified.
This is followed by a mention-antecedent classification step where, for each extracted mention, the model searches for its most probable antecedent (i.e., the extracted span that appears before in the text).
This pipeline, composed of mention extraction and mention-antecedent classification steps, has been adopted with minor modifications in many subsequent works, that we refer to as Coarse-to-Fine models.



  
Coarse-to-Fine Models
Among the works that build upon the Coarse-to-Fine formulation, <cit.>, <cit.> and <cit.> experimented with changing the underlying document encoder,  utilizing ELMo <cit.>, BERT <cit.> and SpanBERT <cit.>, respectively, achieving remarkable score improvements on the English OntoNotes <cit.>. 
Similarly, <cit.> introduced s2e-coref that reduces the high memory footprint of SpanBERT by leveraging the LongFormer <cit.> sparse-attention mechanism. 
Based on the same architecture, <cit.> analyzed the impact of having multiple experts score different linguistically motivated categories (e.g., pronouns-nouns, nouns-nouns, etc.).
While the foregoing works have been able to modernize the original Coarse-to-Fine formulation, training their architectures on the OntoNotes dataset still requires a considerable amount of memory.[Training those models requires at least 32G of VRAM.]
This occurs because they rely on the traditional Coarse-to-Fine pipeline that, as we cover in Section <ref>, has a large memory overhead and is based on manually-set thresholds to regulate memory usage. 



  
Incremental Models
Discriminative systems also include incremental techniques. 
Incremental Coreference Resolution has a strong cognitive grounding: research on the “garden-path” effect shows that humans resolve referring expressions incrementally <cit.>.

A seminal work that proposed an automatic incremental system was that of <cit.>, which introduced a clustering approach based on the shift-reduce paradigm. In this formulation, for each mention, a classifier decides whether to SHIFT it into a singleton (i.e., single mention cluster) or to REDUCE it within an existing cluster.
The same approach has recently been reintroduced in ICoref <cit.> and longdoc <cit.>, which adopted SpanBERT and LongFormer, respectively. 
In these works the mention extraction step is identical to that of Coarse-to-Fine models. 
On the other hand, the mention clustering step is performed by using a linear classifier that scores each mention against a vector representation of previously built clusters, in an incremental fashion. 
This method ensures constant memory usage since cluster representations are updated with a learnable function.
In Section <ref> we present a novel performance-driven incremental method that obtains superior performance and generalization capabilities, in which we adopt a lightweight transformer architecture that retains the mention representations.



 §.§ Sequence-to-Sequence models

Recent state-of-the-art Coreference Resolution systems all employ autoregressive generative approaches. 
However, an early example of Sequence-to-Sequence model, TANL <cit.>, failed to achieve competitive performance on OntoNotes.
The first system to show that the autoregressive formulation was competitive was ASP <cit.>, which outperformed encoder-only discriminative approaches.
ASP is an autoregressive pointer-based model that generates actions for mention extraction (bracket pairing) and then conditions the next step to generate coreference links.
Notably, the breakthrough achieved by ASP is not only due to its formulation but also to its usage of large generative models. 
Indeed, the success of their approach is strictly correlated with the underlying model size, since, when using models with a comparable number of parameters, the performance is significantly lower than encoder-only approaches.
The same occurs in <cit.>, a fully-seq2seq approach where a model learns to generate a formatted sequence encoding coreference notation, in which they report a strong positive correlation between performance and model sizes.

Finally, the current state-of-the-art system on the OntoNotes benchmark is held by Link-Append <cit.>, a transition-based system that incrementally builds clusters exploiting a multi-pass Sequence-to-Sequence architecture. 
This approach incrementally maps the mentions in previously coreference-annotated sentences to system actions for the current sentence, using the same shift-reduce incremental paradigm presented in Section <ref>. 
This method obtains state-of-the-art performance at the cost of using a 13B-parameter model and processing one sentence at a time, drastically increasing the need for computational power.
While the foregoing models ensure superior performance compared to previous discriminative approaches, using them for inference is out of reach for many users, not to mention the exorbitant cost of training them from scratch.



§ METHODOLOGY

In this section, we present the framework: we propose replacing the preprocessing and training strategy of Coarse-to-Fine models with a novel pipeline that improves the training and inference efficiency of Coreference Resolution systems.
Furthermore, with the , we eliminate the dependency on long-standing manually-set hyperparameters that regulate memory usage.
Finally, building on top of our pipeline, we propose three models that adopt a mention-antecedent classification technique, namely and , and a system that is based upon a novel incremental formulation, .



 §.§ Pipeline

The combines i) a novel mention extraction method, ii) an efficient mention regularization technique, and iii) a new mention pruning strategy.



  
Mention Extraction
When it comes to extracting mentions from a document D, there are different strategies to model the probability that a span contains a mention.
Several previous works follow the Coarse-to-Fine formulation presented in Section <ref>, which consists of scoring all the possible spans in D. 
This entails a quadratic computational cost in relation to the input length, which they mitigate by introducing several pruning techniques.

In this work, we employ a different strategy.
We extract coreference mentions by first identifying all the possible starts of a mention, and then, for each start, extracting its possible end.
To extract start indices, we first compute the hidden representation (x_1, …, x_n) of  the tokens (t_1, …, t_n) ∈ D using a transformer encoder, and then use a fully-connected layer F to compute the probability for each t_i being the start of a mention as:

    F_start(x) = W_start'(GeLU(W_startx))


    p_start(t_i) = σ(F_start(x_i))


with W'_start, W_start being the learnable parameters, and σ the sigmoid function.
For each start of a mention t_s, i.e., those tokens having p_start(t_s) > 0.5, we then compute the probability of its subsequent tokens t_j, with s ≤ j, to be the end of a mention that starts with t_s.
We follow the same process as that of the mention start classification, but we condition the prediction on the starting token by concatenating the start, x_s, and end, x_j, hidden representations before the linear classifier:

    F_end(x, x') = W'_end(GeLU(W_end[x, x']))


    p_end(t_j | t_s) = σ(F_end(x_s, x_j))

with W'_end, W_end being learnable parameters.
This formulation handles overlapping mentions since, for each start t_s, we can find multiple ends t_e (i.e., those that have p_end(t_j|t_s) > 0.5).

Previous works already adopted a linear layer to compute start and end mention scores for each possible mention, i.e., s2e-coref <cit.>, and LingMess <cit.>.
However, our mention extraction technique differs from previous approaches since i) we produce two probabilities (0 < p < 1) instead of two unbounded scores and ii) we use the computed start probability to filter out possible mentions, which reduces by a factor of 9 the number of mentions considered compared to existing Coarse-to-Fine systems (Table <ref>, first row).



  
Mention Regularization
To further reduce the computation demand of this process, in the we use the end-of-sentence (EOS) mention regularization strategy: 
after extracting the span start, we consider only the tokens up to the nearest EOS as possible mention end candidates.[We note that all the well-established Coreference Resolution datasets are sentence-split.]
Since annotated mentions never span across sentences, EOS mention regularization prunes the number of mentions considered without any loss of information.
While this heuristic was initially introduced in the implementation of <cit.>, all the recent Coarse-to-Fine have abandoned it in favor of the maximum span-length regularization, which is a manually-set hyperparameter that regulates a threshold to filter out spans that exceed a certain length.
This implies a large overhead of unnecessary computations and introduces a structural bias that does not consider long mentions that exceed a fixed length.[The max-length regularization filters out 196 correctly annotated spans when training on OntoNotes.]
In our work, we not only reintroduce the EOS mention regularization, but we also study its contribution in terms of efficiency, as reported in Table <ref>, second row.



  
Mention Pruning After the mention extraction step, as a result of the , we consider an 18x lower number of candidate mentions for the successive mention clustering phase (Table <ref>). 
This step consists of computing, for each mention, the probability of all its antecedents being in the same cluster, incurring a quadratic computational cost.
Within the Coarse-to-Fine formulation, this high computational cost is mitigated by considering only the top k mentions according to their probability score, where k is a manually set hyperparameter.
Since after our mention extraction step we obtain probabilities for a very concise number of mentions, we consider only mentions classified as probable candidates (i.e., those with p_end > 0.5 and p_start > 0.5),  reducing the number of mention pairs considered by a factor of 10.

In Table <ref>, we compare the previous Coarse-to-Fine formulation with the new .





 §.§ Mention Clustering

As a result of the , we obtain a set of candidate mentions M = (m_1, m_2, …, m_l), for which we propose three different clustering techniques:
and , which use two well-established Coarse-to-Fine mention-antecedent techniques, and , which adopts a novel incremental technique that leverages a light transformer architecture.



  
Mention-Antecedent models 
The first proposed model, , adopts an equivalent mention clustering strategy to <cit.>:
given a mention m_i = (x_s, x_e) and its antecedent m_j=(x_s', x_e'), with their start and end token hidden states, we use two fully-connected layers to model their corresponding representations: 
 
    F_s(x) = W_s'(GeLU(W_sx))

 
    F_e(x) = W_e'(GeLU(W_ex))

we then calculate their probability to be in the same cluster as:

    p_c(m_i, m_j) = σ (F_s(x_s) · W_ss·  F_s(x_s') + 
    
    F_e(x_e) · W_ee· F_e(x_e') +  
    
    F_s(x_s) · W_se· F_e(x_e') + 
     
    F_e(x_e) · W_es· F_s(x_s'))

with W_ss, W_ee, W_se, W_es being four learnable matrices and W_s, W_s', W_e, W_e' the learnable parameters of the two fully-connected layers.

A similar formulation is adopted in , where, instead of using only one generic mention-pair scorer, we use 6 different scorers that handle linguistically motivated categories, as introduced by <cit.>.
We detect which category k a pair of mentions m_i and m_j belongs to (e.g., if m_i is a pronoun and m_j is a proper noun, the category will be Pronoun-Entity) and use a category-specific scorer to compute p_c.
A complete description of the process along with the list of categories can be found in Appendix <ref>.



  
Incremental model Finally, we introduce a novel incremental approach to tackle the mention clustering step, namely , which follows the standard shift-reduce paradigm introduced in Section <ref>.
Differently from the previous neural incremental techniques (i.e., ICoref <cit.> and longdoc <cit.>) which use a linear classifier to obtain the clustering probability between each mention and a fixed length vector representation of previously built clusters, leverages a lightweight transformer model to attend to previous clusters, for which we retain the mentions' hidden representations. 
Specifically, we compute the hidden representations (h_1, …, h_l) for all the candidate mentions in M using a fully-connected layer on top of the concatenation of their start and end token representations.
We first assign the first mention m_1 to the first cluster c_1=(m_1).
Then, for each mention m_i ∈ M at step i we obtain the probability of m_i being in a certain cluster c_j by encoding h_i with all the representations of the mentions contained in the cluster c_j using a transformer architecture.
We use the first special token ([CLS]) of a single-layer transformer architecture T to obtain the score S(m_i, c_j) of m_i being in the cluster c_j=(m_f, …, m_g) with f ≤ g < i as:
 
    S(m_i, c_j) = W_c · (ReLU(T_CLS(h_i, h_f, …, h_g)))

Finally, we compute the probability of m_i belonging to c_j as:
 
    p_c(m_i ∈ c_j| c_j = (m_f, …, m_g)) = σ(S(m_i, c_j))

 We calculate this probability for each cluster c_j up to step i.
 We assign the mention m_i to the most probable cluster c_j having p_c(m_i ∈ c_j) > 0.5 if one exists, or we create a new singleton cluster containing m_i.








As we show in Sections <ref> and <ref>, this formulation obtains better results than previous incremental methods, and is beneficial when dealing with long-document and out-of-domain settings.



 §.§ Training

To train a model, we optimize the sum of three binary cross-entropy losses: 
 
    L_coref = L_start + L_end + L_clust

Our loss formulation differs from previous transformer-based Coarse-to-Fine approaches, which adopt the marginal log-likelihood to optimize the mention to antecedent score <cit.>.
Since their formulation “makes learning slow and ineffective, especially for mention detection” <cit.>, we directly optimize both mention extraction and mention clustering with a multitask approach. 
L_start and L_end are the start loss and end loss, respectively, of the mention extraction step, and are defined as:

    L_start = ∑_i=1^N -(y_ilog(p_start(t_i)) + 
    
    (1 - y_i)log(1 - p_start(t_i))) 
    
    L_end = ∑_s=1^S∑_j=1^E_s -(y_ilog(p_end(t_j|t_s)) +  
    
    (1 - y_i)log(1 - p_end(t_j|t_s)))

where N is the sequence length, S is the number of starts, E_s is the number of possible ends for a start s and p_start(t_i) and p_end(t_j|t_s) are those defined in Section <ref>.

Finally, L_clust is the loss for the mention clustering step. Since we experiment with two different mention clustering formulations, we use a different loss for each clustering technique, namely L_clust^ant for the mention-antecedent models, i.e., and ,  and L_clust^incr for the incremental model, i.e.,  : 

    L_clust^ant = ∑_i=1^|M|∑_j=1^|M|-(y_ilog(p_c(m_i|m_j)) + 
    
    (1 - y_i)log(1 - p_c(m_i|m_j))) 
    
    L_clust^incr =∑_i=1^|M| (∑_j=1^C_i-(y_ilog(p_c(m_i ∈ c_j)) + 
     
    (1 - y_i)log(1 - p_c(m_i ∈ c_j))))

where |M| is the number of extracted mentions, C_i is the set of clusters created up to step i, and  p_c(m_i|m_j) and p_c(m_i ∈ c_j) are  defined in Section <ref>.

All the models we introduce are trained using teacher forcing. 
In particular, in the mention token end classification step, we use gold start indices to condition the end tokens prediction, and, for the mention clustering step, we consider only gold mention indices.
For , at each iteration, we compare each mention only to previous gold clusters.






§ EXPERIMENTS SETUP



 §.§ Datasets

We train and evaluate all the comparison systems on three Coreference Resolution datasets: 



  
OntoNotes <cit.>, proposed in the CoNLL-2012 shared task, is the de facto standard dataset used to benchmark Coreference Resolution systems. 
It consists of documents that span seven distinct genres, including full-length documents (broadcast news, newswire, magazines, weblogs, and Testaments) and multiple speaker transcripts (broadcast and telephone conversations).



  
LitBank <cit.> contains 100 literary documents typically used to evaluate long-document Coreference Resolution. 



  
PreCo <cit.> is a large-scale dataset that includes reading comprehension tests for middle school and high school students. 


Notably, both LitBank and PreCo have different annotation guidelines compared to OntoNotes, and provide annotation for singletons (i.e., single-mention clusters).
Furthermore, we evaluate models trained on OntoNotes on three out-of-domain datasets: 



    
  * GAP <cit.> contains sentences in which, given a pronoun, the model has to choose between two candidate mentions.
    
  * LitBank_ns and PreCo_ns, the datasets' test-set where we filter out singleton annotations.
    
  * WikiCoref <cit.>, which contains Wikipedia texts, including documents with up to 9,869 tokens.  



The statistics of the datasets used are shown in Table <ref>.



 §.§ Comparison Systems



  
Discriminative
Among the discriminative systems, we consider c2f-coref <cit.> and s2e-coref <cit.>, which build upon the Coarse-to-Fine formulation and adopt different document encoders.
We also report the results of LingMess <cit.>, which is the previous best encoder-only solution, and f-coref <cit.>, which is a distilled version of LingMess.
Furthermore, we include CorefQA <cit.>, which casts Coreference as extractive Question Answering, and wl-coref <cit.>, which first predicts coreference links between words, then extracts mentions spans. 
Finally, we report the results of incremental systems, such as ICoref <cit.> and longdoc <cit.>.



  
Sequence-to-Sequence
We compare our models with TANL <cit.> and ASP <cit.>, which frame Coreference Resolution as an autoregressive structured prediction.
We also include Link-Append <cit.>, a transition-based system that builds clusters with a multi-pass Sequence-to-Sequence architecture. 
Finally, we report the results of seq2seq <cit.>, a model that learns to generate a sequence with Coreference Resolution labels.



 §.§ Setup

All models use DeBERTa-v3 <cit.> as the document encoder.
We use DeBERTa because it can model very long input texts effectively <cit.>.[This is because its attention mechanism enables its input length to grow linearly with the number of its layers.]
Moreover, compared to the LongFormer, which was previously adopted by several token-level systems, DeBERTa ensures a larger input max sequence length (e.g., DeBERTa_large  can handle sequences up to 24,528 tokens while LongFormer only 4096) and has shown better performances empirically in our experiments on the OntoNotes dataset.
On the other hand, using DeBERTa to encode long documents is computationally expensive because its attention mechanism incurs a quadratic computational complexity.
Whereas this further increases the computational cost of traditional Coarse-to-Fine systems, the enables us to train models that leverage DeBERTa_large on the OntoNotes dataset, without any performance-lowering pruning heuristic.
To train our models we use Adafactor <cit.> as our optimizer, with a learning rate of 3e-4 for the linear layers, and 2e-5 for the pre-trained encoder. 
We perform all our experiments within an academic budget, i.e., a single RTX 4090, which has 24GB of VRAM.
We report more training details in Appendix <ref>.



§ RESULTS






 §.§ English OntoNotes

We report in Table <ref> the average CoNLL-F1 score of the comparison systems trained on the English OntoNotes, along with their underlying pre-trained language models and total parameters.
Compared to previous discriminative systems, we report gains of +2.2 CoNLL-F1 points over LingMess, the best encoder-only model. 
Interestingly, we even outperform CorefQA, which uses additional Question Answering training data.

Concerning Sequence-to-Sequence approaches, we report extensive improvements over systems with a similar amount of parameters compared to our large models (500M): we obtain +3.4 points compared to ASP (770M), and the gap is even wider when taking into consideration Link-Append (3B) and seq2seq (770M), with +6.4 and +5.6, respectively.
Most importantly, models surpass the performance of all Sequence-to-Sequence transformers even when they have several billions of parameters.
Among our proposed methods, shows the best performance, setting a new state of the art with a score of 83.6 CoNLL-F1 points on the OntoNotes benchmark.
More detailed results, including a table with MUC, B^3, and CEAFϕ_4 scores and a qualitative error analysis, can be found in Appendix <ref>. 



 §.§ PreCo and LitBank

We further validate the robustness of the framework by training and evaluating systems on the PreCo and LitBank datasets.
As reported in Table <ref>, our models show superior performance when dealing with long documents in a data-scarce setting such as the one LitBank poses.
On this dataset, achieves a new state-of-the-art score of 78.3, and gains +1.0 CoNLL-F1 points compared with seq2seq.
On PreCo, outperforms longdoc, but seq2seq still shows slightly better performance.
This is mainly due to the high presence of singletons in PreCo (52% of all the clusters).
Our systems, using a mention extraction technique that favors precision rather than recall, are penalized compared to high recall systems such as seq2seq.[Precision and Recall scores are reported in Appendix <ref>.]
Among our systems, , leveraging its hybrid architecture, performs better on both PreCo and LitBank. 

















 §.§ Out-of-Domain Evaluation

In Table <ref>, we report the performance of systems along with LingMess, the best encoder-only model, when dealing with out-of-domain texts, that is, when they are trained on OntoNotes and tested on other datasets.
First of all, we report considerable improvements on the GAP test set, obtaining a +1.2 F1 score compared to the previous state of the art. 
We also test models on WikiCoref, PreCo_ns and LitBank_ns (Section <ref>).
However, since the span annotation guidelines of these corpora differ from the ones used in OntoNotes, in Table <ref> we also report the performance using gold mentions, i.e., skipping the mention extraction step (gold column).[We do not include autoregressive models because none of the original articles report scores on out-of-domain datasets. We also could not test those models because they do not provide the code to perform mention clustering alone, and performing it with such approaches is not as straightforward as in encoder-only models.]
On the WikiCoref benchmark, we achieve a new state-of-the-art score of 67.2 CoNLL-F1, with an improvement of +4.2 points over the previous best score obtained by LingMess. 
On the same dataset, when using pre-identified mentions the gap increases to +5.8 CoNLL-F1 points (76.6 vs 82.4).
In the same setting, our models obtain up to +7.3 and +10.1 CoNLL-F1 points on Preco_ns and LitBank_ns, respectively,  compared to LingMess. 
These results suggest that the training strategy makes this model more suitable when dealing with pre-identified mentions and out-of-domain texts.
This further increases the potential benefits that systems can bring to many downstream applications that exploit coreference as an intermediate layer, such as Entity Linking <cit.> and Relation Extraction <cit.>, where the mentions are already identified.
Among our models, on LitBank_ns and WikiCoref, outperforms and , confirming the superior capabilities of the incremental formulation in the long-document setting.



Finally, we highlight that the performance gap between using gold mentions and performing full Coreference Resolution is wider when tested on out-of-domain datasets (on average +17%) compared to testing it directly on OntoNotes (83.6 vs 93.6, +10%).[An evaluation of the proposed models in terms of mention extraction and mention clustering using gold mentions scores can be found in Appendix <ref>.]
This result, obtained on three different out-of-domain datasets, suggests that the difference in annotation guidelines considerably contribute to lower the OOD performances (-7%).



























































































 §.§ Speed and Memory Usage

In Table <ref>, we include details regarding the training time and the hardware used by each comparison system, along with the measurement of the inference time and peak memory usage on OntoNotes the validation set.
Compared to Coarse-to-Fine models, which require 32GB of VRAM, we can train systems under 18GB.
At inference time both and , exploiting DeBERTa_large, achieve competitive speed and memory consumption compared to wl-coref and s2e-coref.
Furthermore, when adopting DeBERTa_base, proves to be the most efficient approach[In terms of inference peak memory usage and speed.] among those directly trained on OntoNotes, while, at the same time, attaining performances that are equal to the previous best encoder-only system, LingMess. 
The only system that shows better inference speed is f-coref, but at the cost of lower performance (-3.0).

Compared to the previous Sequence-to-Sequence state-of-the-art approach, Link-Append, we train our models with 175x less memory requirements. 
Comparing inference time is more complicated, since we could not run models on our memory-constrained budget. 
For this reason, we report the inference times from the original articles, and hence times achieved with their high-resource settings.
Interestingly, we report as much as 170x faster inference compared to seq2seq, which exploits parallel inference on multiple GPUs, and 85x faster when compared to the more efficient ASP.
Among models, is notably slower both in inference and training time, as it incrementally builds clusters using multiple steps.






























































































 §.§ Ablation

In Table <ref>, we compare and models with s2e-coref and LingMess, respectively, using different pre-trained encoders.
Interestingly, when using DeBERTa, systems not only achieve better speed and memory efficiency, but also obtain higher performance compared to the previous systems. 

When using the LongFormer, instead, their scores are in the same ballpark, showing empirically that the training procedure better exploits the capabilities of DeBERTa.
To test the benefits of our novel incremental formulation, , we also implement a model with the previously adopted incremental method used in longdoc and ICoref (Section <ref>), which we call _prev-incr. 
Compared to the previous formulation we report an increase in score of +3.9 CoNLL-F1 points.
The improvement demonstrates that exploiting a transformer architecture to attend to all the previously clustered mentions is beneficial, and enables the future usage of hybrid architectures when needed.

As a further analysis of whether the efficiency improvements of our systems stem from using DeBERTa or are attributable to the , we compared the speed and memory occupation of a system using as underlying encoder either DeBERTa_large or LongFormer_large.
Our experiments show that using DeBERTa leads to an increase of +77% of memory space and +23% of time to complete an epoch when training on OntoNotes. An equivalent measurement, attributable to the quadratic memory attention mechanism of DeBERTa, was observed for the inference time and memory occupation on the OntoNotes test set. These results highlight the efficiency contribution of the , which is agnostic to the document encoder and can be applied to future coreference systems to ensure higher efficiency.


















































































§ CONCLUSION

In this work, we challenged the recent trends of adopting large autoregressive generative models to solve the Coreference Resolution task. 
To do so, we proposed , a new framework that enables fast and memory-efficient Coreference Resolution while obtaining state-of-the-art results. 
This demonstrates that the large computational overhead required by Sequence-to-Sequence approaches is unnecessary. 
Indeed, in our experiments systems demonstrated that they can outperform large generative models and improve the speed and memory usage of previous best-performing encoder-only approaches.
Furthermore, we introduced , a robust multi-step incremental technique that obtains higher performance in the out-of-domain and long-document setting.
By releasing our systems, we make state-of-the-art models usable by a larger portion of users in different scenarios and potentially improve downstream applications.





§ LIMITATIONS

Our experiments were limited by our resource setting i.e., a single RTX 4090. For this reason, we could not run using larger encoders, and could not properly test Sequence-to-Sequence models as we did with encoder-only models. Nevertheless, we believe this limitation is a common scenario in many real-world applications that would benefit substantially from our system.
We also did not test our formulation on multiple languages, but note that both the methodology behind and our novel incremental formulation are language agnostic, and thus could be applied to any language. 






§ ACKNOWLEDGEMENTS



0.1
        

            
    < g r a p h i c s >

        

    
    
    0.70
         We gratefully acknowledge the support of the PNRR MUR project PE0000013-FAIR.
    
    
    0.1
        

            
    < g r a p h i c s >

        

    





Roberto Navigli also gratefully acknowledges the support of the CREATIVE project (CRoss-modal understanding and gEnerATIon of Visual and tExtual content), which is funded by the MUR Progetti di Rilevante Interesse Nazionale programme (PRIN 2020).
This work has been carried out while Giuliano Martinelli was enrolled in the Italian National Doctorate on Artificial Intelligence run by Sapienza University of Rome. 











§ MULTI-EXPERT SCORERS


In , the final coreference score between two spans is calculated using 6 linguistically motivated multi-expert scorers.
This approach was introduced by <cit.>, which demonstrated that linguistic knowledge and symbolic
computation can still be used to improve results on the OntoNotes benchmark.
In we adopt this approach on top of the .
We use the same set of categories, namely:

    
  * Pron-Pron-C. Compatible pronouns based on their attributes, such as gender or number (e.g.  (I, I), (I, my) (she, her)).  
    
  * Pron-Pron-NC, Incompatible pronouns (e.g. (I, he), (She, my), (his, her)).
    
  * Ent-Pron. Pronoun
    and non-pronoun (e.g. (George, he), (CNN, it), (Tom Cruise, his)).
    
  * Match. Non-pronoun spans with the same content words (e.g. Italy, Italy).
    
  * Contains. One contains the other (e.g. (Barack Obama, Obama)).
    
  * Other. The Other pairs.

To detect pronouns we use string match with a full list of English pronouns. 

To perform mention clustering, we dedicate a mention-pair scorer for each of those categories.
Specifically, for the mention m_i=(x_s, x_e) and its antecedent m_j=(x_s',x_e'), with their start and end token hidden states, we first detect their category k_g using pattern matching on their spans of texts. 
Then we compute their start and end representations, using the specific fully-connected layers for the category k_g:
 
    F_s^k_g(x) = W_k_g,s'(GeLU(W_k_g,sx))

 
    F_e^k_g(x) = W_k_g,e'(GeLU(W_k_g,ex))

The probability p_c^k_g of m_i and m_jis then calculated as:

    p_c^k_g(m_i, m_j) = σ (F_s^k_g(x_s) ·  W_ss·  F_s^k_g(x_s') + 
    
    F_e^k_g(x_e) · W_ee· F_e^k_g(x_e') +  
    
    F_s^k_g(x_s) · W_se· F_e^k_g(x_e') + 
     
    F_e^k_g(x_e) · W_es· F_s^k_g(x_s'))

With W_ss, W_ee, W_se, W_es being four learnable matrices and W_k_g,e', W_k_g,s',W_k_g,e, W_k_g,s' the learnable parameters of the two fully-connected layers.
In this way, each mention-pair scorer learns to model the probability for its specific linguistic categories. 





§ TRAINING DETAILS



 §.§ Datasets




We report technical details of the adopted datasets.

    
  * OntoNotes contains several items of metadata information for each document, such as genre, speakers, and constituent graphs.
    Following previous works, we incorporate the speaker’s name into the text whenever there is a change in speakers for datasets that include this metadata.
    
  * LitBank contains 100 literary documents and is available in 10 different cross-validation folds. Our train, dev, and test splits refer to the first cross-validation fold, LB_0. We report comparison systems results on the same splits. Since training DeBERTa_large is particularly computationally expensive, as introduced in Section <ref>, we train on LitBank by splitting in half each LitBank training document.
    
  * The authors of PreCo have not released their official test set. 
    To evaluate our models consistently with previous approaches, we use the official 'dev' split as our test set and retain the last 500 training examples for model validation.









 §.§ Setup

All our experiments are developed using the Pytorch-Lightning framework.[<https://lightning.ai>]
For each model, we load the pre-trained weights for the base[<https://huggingface.co/microsoft/deberta-v3-base>] and large[<https://huggingface.co/microsoft/deberta-v3-large>] version of DeBERTa-v3 from the Huggingface Transformers library <cit.>.
We accumulate gradients every 4 steps and use a gradient clipping value of 1.0. 
We adopt a linear learning rate scheduler a warm-up of 10% of the total steps check validation scores every 50% of the total number of steps per epoch. 
We select our model upon validation of Avg. CoNLL-f1 score and use patience of 20.









§ ADDITIONAL RESULTS


In Table <ref> we report the performance of models according to the standard Coreference Resolution metrics: MUC <cit.>, B^3<cit.>, CEAFϕ_4 <cit.> and AVG CoNLL-F1.
Scores for models are computed using the official CoNLL coreference scorer.[<https://conll.github.io/reference-coreference-scorers>]






 §.§ Error Analysis

To better understand the quality of predictions, we conduct an error analysis on our best system trained on OntoNotes, . 
In Table <ref>, we report the score of performing only mention extraction (F1) or mention clustering with gold mention (CoNLL-F1) with our systems.
Our results highlight that our models have strong capabilities of clustering pre-identified mentions, but limited performance in the identification of correct spans.
We investigated this phenomenon by conducting a qualitative evaluation of the outputs of our best system, , and found out that OntoNotes contains several annotation errors.
We report examples of errors in Table <ref>. The main inconsistency we found in the gold test set is that many documents have incomplete annotations, which directly correlates with the mention extraction error.










 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: Augmenting Channel Simulator and Semi- Supervised Learning for Efficient Indoor Positioning.




















Augmenting Channel Simulator and Semi- Supervised Learning for Efficient Indoor Positioning




    Yupeng Li†, Xinyu Ning, Shijian Gao*, Yitong Liu, Zhi Sun$, Qixing Wang†, Jiangzhou Wang† 
 	
†China Mobile Research Institute, Beijing, China 

Beijing University of Posts and Telecommunications, Beijing, China

*The Hong Kong University of Science and Technology (Guangzhou), China

$Tsinghua University, Beijing, China 

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



§ ABSTRACT 
This work aims to tackle the labor-intensive and resource-consuming task of indoor positioning by proposing an efficient approach. The proposed approach involves the introduction of a semi-supervised learning (SSL) with a biased teacher (SSLB) algorithm, which effectively utilizes both labeled and unlabeled channel data. To reduce measurement expenses, unlabeled data is generated using an updated channel simulator (UCHS), and then weighted by adaptive confidence values to simplify the tuning of hyperparameters. Simulation results demonstrate that the proposed strategy achieves superior performance while minimizing measurement overhead and training expense compared to existing benchmarks, offering a valuable and practical solution for indoor positioning.











Indoor positioning, semi-supervised learning, pseudo-label, deep learning, 6G




§ INTRODUCTION

The 3rd Generation Partnership Project (3GPP) has successfully conducted the fifth generation (5G) Advanced phase in its Release 18<cit.>. As one of the most significant potential advancements to increase automation and efficiency of the industrial domain, the indoor positioning has attracted widespread attention. However, the indoor factory scenario presents challenges, including a high volume of obstructions leading to numerous non-line-of-sight (NLOS) scenarios. These conditions often lead to significant performance degradation for the traditional positioning algorithms <cit.>. 






To address the issue of the limitations and difficulties of traditional positioning methods, a notable point is to integrate deep learning (DL) into indoor positioning systems, specifically fingerprint-based positioning. 

The basic idea of the fingerprint positioning is to create a mapping function between the channel measurements and the position coordinates, creating a main wireless radio environment among the transceivers <cit.>. Then, the user equipment (UE) can obtain its position by the measurements. The measurements of fingerprint positioning in the cellular systems are mainly channel impulse response (CIR), power delay profile (PDP), and so on <cit.>.



However, DL-based indoor positioning faces the challenge of insufficient data available for training the model <cit.>. Especially, in the flexible and dynamic communication systems, obtaining sufficient measurements and labeling them is a labor-intensive and resource-consuming task. 

















A powerful approach termed semi-supervised learning (SSL) was proposed, which involved using small number of labeled data combined with a large amount of unlabeled data to train the model <cit.>. 


Typical SSL <cit.> may lead to serious performance degradation due to errors in the pseudo-labels of the unlabeled data.

To address this issue, a smaller weighting coefficient was added to the unlabeled data during model training <cit.>.

The weighting coefficient is a hyperparameter, which requires meticulous manual tuning.

In <cit.>, the concept of confidence for the pseudo-label of the unlabeled data was proposed, which helps to select reliable unlabeled data in classification tasks. However, there is a lack of a method to define a label confidence in regression tasks like positioning. Furthermore, in an indoor scenario, obtaining the unlabeled data is also a challenge that requires extra measurement overhead.




To address the aforementioned challenges, a DL-based positioning algorithm with reduced overhead is proposed. To lower the measurement overhead, the unlabeled data is generated by an updated channel simulator according to the stochastic parameters of the labeled data. This approach reduces the measurement expenses of the unlabeled data. Additionally, to capture the features of both the labeled and unlabeled channel data, a semi-supervised learning with a biased teacher is proposed. The main contributions of this paper are summarized as follows:



  * An updated channel simulator is utilized to generate unlabeled data, thereby reducing the measurement expenses associated with acquiring such data. Initially, labeled data collected within an indoor factory environment is utilized, and the primary channel parameters are analyzed to extract the channel stochastic parameters of the measurement environment. Subsequently, the typical TR 38.901 <cit.> channel simulator is updated with these stochastic parameters, resulting in the creation of an updated channel simulator (UCHS). The UCHS is then employed to simulate the unlabeled channel data, which closely resembles the labeled data of the measurement environment.





  * To efficiently capture the mapping function between the CIR and the position coordinate while minimizing training expenses, a SSL with a biased teacher (SSLB) is introduced. The SSLB approach involves weighting the unlabeled data based on the confidence, which is derived from the distance bias predicted by the teacher model. This methodology assigns higher confidence to more reliable unlabeled data. Notably, this work represents the first instance of introducing confidence into regression tasks, such as indoor positioning, resulting in a significant reduction in training expenses.

  * To obtain the confidence of the unlabeled data, a novel SSL structure is designed, termed as SSLB. This structure incorporates an additional reference signal and a position bias, added separately to the input and the output of the traditional SSL model. As a result, SSLB can extract additional bias information from the labeled dataset, which can then be transformed into the confidence of the unlabeled data. 






























Notation: Vectors and matrices are denoted by boldface. ||·|| expresses the 2-norm of a vector. ||·||_F is the Frobenius norm of a matrix. For unlabeled data, we use a superscript ^*, while labeled data remains unmarked.



§ PROBLEM DESCRIPTION










 §.§ Adapting semi-supervised learning for indoor positioning


The dataset of traditional SSL consists of two parts: the labeled dataset and the unlabeled dataset. The input is the CIR and the output is the predicted position.

The labeled dataset is represented as 𝔇(H, l), where H is the labeled CIR and l is its corresponding position coordinate. The unlabeled dataset is represented as 𝔈(H^*), containing only the CIR.


The SSL model is divided into a teacher model Θ^t and a student model Θ^s. Firstly, the teacher model is trained with labeled data, and then pseudo-labels of unlabeled data are generated using the trained teacher model. Finally, the student model is trained with both the labeled data and the unlabeled data. In this paper, the loss function is defined as the Euclidean distance between the predicted position and the real position as shown in (<ref>). First, train the teacher model using this loss on the labeled dataset 𝔇(H, l).

    L = 1/N∑^N_i=1 ||l̂_i - l_i||.


Then, the label of dataset 𝔈(H^*) is generated with the trained teacher model

    l^* = f(H^*, Θ^t),

where l^* is the position coordinate of H^*. f(·) represents the mapping function between the channel measurements and position coordinates. In this way, the dataset 𝔈(H^*) is completed to 𝔈(H^*, l^*).

Finally, the student model is trained with the dataset of 𝔇(H, l) and 𝔈(H^*, l^*), applying weights to pseudo-labeled data. The loss function is defined as follows:

    L^w   =1/N+N^*(∑^N_i=1√(w||l̂_i-l_i||^2) +  ∑^N^*_j=1√(w^*||l̂^*_j-l^*_j||^2)),


where L^w is the loss of weighted semi-supervised algorithm, N and N^* are the numbers of labeled data and unlabeled data, and w and w^* are the corresponding weights of labeled data and unlabeled data. Detailed descriptions can be referred to <cit.><cit.>. it is worth mentioning that all models use the same DL model, illustrated in Fig. <ref>.








§ CHANNEL FEATURE AIDED DATA AUGMENTATION













 §.§ Wireless channel feature extracting method


In this section, some specific methods for extracting channel parameters will be described.



  §.§.§ Angle spread

The angle spread (AS) is used to characterize the dispersion features of the multipaths in the spatial domain. Initially, a two-dimensional Discrete Fourier transform (DFT) is performed over the channel matrix H(t) over the port domain, resulting in the angle domain channel matrix H_ang <cit.> as follows:

    H_ang = U^T H(t) V,

where U = [u_1^T; ...; u_N_t^T]^T is the transmit DFT matrix, and V = [v_1, ... ,v_N_r] is the receiver DFT matrix<cit.>.

Then, the receiving and transmitting AS are illustrated in the spatial domain. The angle spread value can be obtained as follows:

    σ_AS = √(∑^M_m=1θ_m^2P_m/∑^M_m=1P_m)
    ,

where P_m is the power of the mth multipath, θ_m is the angle of the mth multipath.



  §.§.§ Delay spread

The delay spread (DS) can reflect the richness of the channel multipaths. And the DS <cit.> can be obtained as follows:

    σ_DS = √(∑^M_m=1(τ_m - τ_mean)^2 P_m/∑^M_m=1P_m),

where τ_mean = ∑^M_m=1 (τ_mP_m)/∑^M_m=1 (P_m) is the mean delay.

Due to the space limit, only the above two unique channel statistical parameters are delineated here. Other channel statistics parameters can refer to <cit.>.



 §.§ Unlabeled data generation via UCHS


The TR 38.901 <cit.>, considered as the state-of-the-art 3GPP channel model, is crucial for generating channel simulation data and evaluating physical layer techniques. It illustrates the stochastic distribution characteristics of channel multipaths, serving as a benchmark in a general scenario. However, in a specific location, the channel parameters in TR 38.901 may not be precise. Regrettably, in an AI-enabled indoor factory, the training data needs to reflect the actual propagation environment as accurately as possible.

To address the aforementioned issue, we propose a data augmentation method based on the channel simulator, which can simulate the unlabeled channel data corresponding to the specific area. This approach significantly reduces the collection overhead of unlabeled data. 

The generation process of unlabeled data is shown as follows:


  * Step 1: the labeled channel data is collected by the UEs distributed in a specific indoor factory. Since labeled data is essential for training the teacher model, there is no extra overhead of data collection at this
stage.

  * Step 2: upload the labeled data from the UE to the BS.

  * Step 3: the primary statistical parameters of the channel data are extracted, as described in Section <ref>, effectively depicting the wireless propagation environment.

  * Step 4: these statistical parameters are used to update the parameters of
typical TR 38.901 channel model, resulting in the UCHS.

  * Step 5: the UCHS is used to simulate unlabeled channel data similar to the labeled data of the measurement environment.





§ SEMI-SUPERVISED PSEUDO-LABELING WITH A BIASED TEACHER







Compared to traditional SSLs, our proposed method assigns different loss weights to the pseudo-labels of the unlabeled data based on the confidence scores. By doing so, pseudo-labels with low confidence are assigned small weights during the training stage of the student model, thereby mitigating the impact of pseudo-labels with large offset on the performance of the student model.




The proposed method pipeline, as illustrated in Fig. <ref>, consists of three primary stages: the bias teacher training stage, the confidence-based pseudo-labeling stage, and the weighted loss student training stage:





  * Stage 1: an additional reference signal and a position bias are separately added to the input and output of the traditional SSL model. This addition helps extract additional information from the labeled dataset. Specifically, the position bias represents the distance between the input signal and its nearest signal, referred to as the reference signal. This allows the teacher model to not only determine the mapping function between the input signal and the position coordinates but also to assess the similarity between input signals using the distance bias.


  * Stage 2: the trained teacher model is utilized to generate pseudo-labels and the position bias of the unlabeled data. The position bias can then be used to define the confidence scores of the unlabeled data through a kernel density estimation (KDE) <cit.> function.

  * Stage 3: a weighted loss strategy is employed, wherein different loss weights are assigned to different pseudo-labels in the loss function to train the student model. This approach aims to mitigate the impact of pseudo-labels with larger errors on the performance of the student model.





 §.§ The biased teacher training



In the context of object detection, models not only output the bounding boxes of the detected objects but also provide confidence scores for these detections. Inspired by this, a biased teacher model with two prediction heads is proposed. One prediction head, called the 'predict head”, is responsible for predicting the position coordinates l̂ corresponding to the input CIR. The other prediction head, named as the ”biased head” is designed to predict the bias Δl̂ of the reference signal relative to the input CIR position.






The first step involves using the k-nearest neighbors (KNN)<cit.> algorithm to search the reference signal H^ref_i in labeled dataset 𝔇(H, l) that is closest to the input CIR H_i, thereby obtaining the reference signal data pair (H^ref_i, l^ref_i), i.e.,


    H^ref_i = H_j, min_j ∈ (1, 2, ..., N)  ||H_i-H_j||_F, i ≠ j ,



where || · ||_F is Frobenius norm, and i ≠ j excludes itself as the nearest neighbor. 

The position bias is defined as:

    Δl_i ≜l^ref_i - l_i.

The total loss function of the teacher model is defined as:

    L = L^p + L^b,




where L^p = 1/N∑_i = 1^N ||l̂_̂î-l_i|| and L^b = 1/N∑_i = 1^N ||Δl̂_̂î-Δl_i|| are the position prediction loss and the bias prediction loss, respectively.



















 §.§ Confidence-based pseudo-labeling



This process involves using the teacher model to generate estimated positions l^* and position biases Δl^* for the unlabeled data, which can then be used for subsequent steps in the SSL framework.




For the unlabeled data, the confidence is determined by calculating the distribution of the position bias. This involves computing the position bias as d=||Δl^*||. Subsequently, the distribution can be fitted using KDE as follows:

    f̂_h(d)=1/Nh∑^N_i=1K( d-d_i/h)
    ,


where f̂_h(d) is the KDE at point d. N is the number of samples, h is the bandwidth, d_i represents each data point in the sample, K(·) represents the kernel function, and in this work the Gaussian kernel is used. Fig. <ref>(a) illustrates the distribution of the unlabeled data in our experiment.


The confidence is defined as:

    w^*_i =f̂_h(d_i).


To ensure a fair comparison with the SSL, the obtained weights w^*_i need to divide the mean value of all the w^* to guarantee its mean value to be 1. This normalization process is important for comparison purposes. The distribution of w^* is shown in Fig. <ref>(b).
















 §.§ Weighted loss for student training


The student model Θ^s is trained with the mixture dataset, which is constructed by the unlabeled dataset 𝔈(H^*,l^*,w^*) 
combined with the labeled dataset 𝔇(H, l). 



For the student model, only the predicted coordinates l̂ are used to compute the loss, which is defined as

    L^w   =1/N+N^*( ∑^N_i=1 ||l̂_̂î-l_i|| + ∑^N^*_j=1√(w^*_i||l̂^*_j-l^*_j||^2)),


where w^*_i is the weight of the pseudo-label, and the weight of labeled data is set as 1.








§ EXPERIMENTS AND RESULTS



 §.§ Scenario deployment

The dense clutter high BS (DH) sub-scenario is a priority for evaluation in 3GPP for AI-enabled indoor factory (InF) <cit.>. In our considered scenario, there are 18 BSs uniformly distributed within the 60m*120m space with a spacing of 20m and height of 8m. The central frequency is 3.5 GHz and the bandwidth is 100 MHz in our experiment. The clutter density is 40% with the clutter height of 2m. The UEs are randomly generated within the space with a height of 1.5m.












 §.§ Experimental setup


In the experiment, a total of 17200 labeled channel samples are collected in the measurement scenario, where 10000 samples are randomly selected to construct the training dataset and the residual 7200 samples are allocated to the testing dataset. Besides, the unlabeled channel data are generated using UCHS. 


All the models are implemented using PyTorch and the Adam optimizer. The hyperparameters are learning rate 5e-2, an input batch size of 256, 150 epochs, and a cosine annealing rate scheduler. The criteria for InF positioning is the positioning accuracy at a cumulative distribution function (CDF) of 90% for users <cit.>.
























 §.§ Data structure

In the cellular positioning problem, UEs measure the time domain CIR H(t) from multiple BSs. For each position, the measured H(t) by each UE is in the form of N_BS*N_port*N_delay*2 dimensions, where N_BS, N_port, and N_delay represent the number of BS, the number of BS antenna respectively, and the number of the time delay, and 2 is from the complex-domain to real and imaginary parts. In most industrial scenario, the effective delay is mainly distributed in the first 64 delay taps. Under the current setup, this results in an extremely high-dimensional tensor for each position in the form of 18*4*64*2.




 §.§ Experimental results



  §.§.§ Main results


In this section, we carry out simulations to evaluate the performance of our proposed  SSL with position bias, which we call as “SSLB” scheme hereafter. We compare the proposed alrotihm with the following three benchmark schemes: 


  * a) SL: traditional supervised learning, whose reference signal H^ref is represented by the input signal H, with all the predicted position biases being zero. 

  * b) SLR: SL learning with reference signals. Its input signal contains the current CIR H and the reference signal H^ref, and the output is l̂ and Δl̂. 

  * c) SSLR: SSL pseudo-label learning without a biased teacher. 
The same DL model and training hyperparameters are used for the four schemes mentioned above. Our method aims to provide a convenient guideline for adaptively obtaining the weight of the unlabeled data using confidence. This is compatible with consistency-based pseudo-labeling SSL <cit.>, and combining them may lead to better results, but it is beyond the scope of this paper.

 














Fig. <ref> illustrates the positioning accuracy for the four schemes with different labeled samples. It is evident that the SSLB achieves the highest positioning precision. For instance, with less than 6300 samples, SSLB can achieve the precision of 0.897 m, while the SL requires 9900 samples to achieve a precision of 0.969 m. Therefore, with the help of UCHS, our proposed SSLB can reduce the data overhead by almost 40% compared to SL. However, it is worth noting that the performance of the SSL pseudo-label method degrades at 2700 samples, due to the introduction of enormous inaccurate pseudo-labels by the teacher model, which deteriorates the performance of the student model. In contrast, SSLB achieves relatively stable performance, as it can automatically assign small weights to the unlabeled data with large offset.








  §.§.§ Ablation study





  
The impact of weight on the test loss







In this section, we study the effects of weight magnitude α on the SSLR and our proposed SSLB. In the student model training stage, for both the SSLR and the SSLB, the unlabeled weights are scaled by the formula w^*'_i = α w^*_i. Fig. <ref> compares the performance of SSLR and SSLB under different scaling coefficients. As the scaling coefficient decreases, the performance gap between the two schemes reduces, whereas the SSLB outperforms the SSLR in all cases. Table <ref> demonstrates the position accuracy at CDF of 90% for users. With the decrease of scaling coefficients, the positioning precision of SSLR increases at first, and then decreases. However, the performance of KDE confidence based SSLB keeps relatively stable. Thus, our proposed scheme does not manually select the weights hyperparameters α of the unlabeled data and dynamically assigns different weights to the pseudo-labels. Therefore, the training expense is dramatically reduced.




  
Impact of confidence calibration

To verify the effectiveness of the KDE confidence based method for transforming the position bias to confidence, we compare it with a linear transformation method. The linear transformation method directly considers position bias as confidence, assuming that small bias corresponds to high confidence levels. Firstly, the Euclidean distance of the position bias is computed. Then,  this data is normalized within the range of 0 to 1, reversed by subtracting from 1, and finally divided by the mean. The comparison results between the linear method and KDE are presented in Table <ref>.







§ CONCLUSION

This paper has investigated the issue of indoor positioning with reduced expenses. On one hand, the the UCHS was used to generate unlabeled data and reduce measurement overheads. On the other hand, we proposed the SSLB, which adaptively weighted the unlabeled data based on confidence to differentiate the quality of the data. This approach effectively reduced the expenses associated with tuning the hyperparameters of the weights. Simulation results demonstrated that our proposed scheme exhibited superior performance, reducing the data measurement overhead by almost 40% compared to benchmark schemes. Furthermore, ablation study confirmed that our proposed scheme can automatically assign different weights to unlabeled data, significantly reducing training expenses.






























































































00

R18 3GPP TR 38.843, v1.1.0, ”Study on Artificial Intelligence (AI)/Machine Learning (ML) for NR air interface,” 2023.

geome1 L. Chen, X. Zhou, F. Chen, L.-L. Yang, and R. Chen, ”Carrier phase ranging for indoor positioning with 5G NR signals,” IEEE Internet of Things Journal, vol. 9, no. 13, pp. 10908-10919, Jul. 2022.





geome4 Y. Li, Y. Zhuang, X. Hu, Z. Gao, J. Hu, et al., ”Toward location-enabled iot (le-iot): Iot positioning techniques, error sources, and error mitigation,” IEEE Internet of Things Journal, vol. 8, no. 6, pp. 4035–4062, 2020.


fingerprint2 K. S. Wu, J. Xiao, Y. W. Yi, M. Gao, and L. M. Ni, ”FILA: Fine-grained indoor localization,” in Proc. IEEE INFOCOM, 2012, pp. 2210-2218.

TSG_RAN 3GPP TSG-RAN Metting #96, ”Revised SID: Study on Artificial Intelligence (AI)/Machine Learning (ML) for NR air interface,” 2022.

insufficient K. M. Chen and R. Y. Chang, ”Semi-Supervised Learning with GANs for Device-Free Fingerprinting Indoor Localization,” IEEE GLOBECOM, Taipei, Taiwan, pp. 1-6, 2020.

ssl Y. Ruan, L. Chen, X. Zhou, L. Zhaoliang, L. Xiaoyan, et al., ”iPos-5G: Indoor positioning via commercial 5G NR CSI,” IEEE Internet of Things Journal, vol. 10, no. 10, pp. 8718-8733, 2022.
ssl1 X. Yang, Z. Song, I. King and Z. Xu, ”A survey on deep semi-supervised learning,” IEEE Transactions on Knowledge and Data Engineering, vol. 35, no. 9, pp. 8934-8954, 2022.

ssl2 X. Lin, ”An Overview of 5G Advanced Evolution in 3GPP Release 18,” IEEE Communications Standards Magazine, vol. 6, no. 3, pp. 77-83, September 2022.
pseudo
D. H. Lee, ”Pseudo-label: The simple and efficient semi-supervised learning method for deep neural networks,” Workshop on Challenges in representation learning, vol. 3, no. 2, 2013.

semi_and_gan W. Njima, A. Bazzi and M. Chaffi, ”DNN-based indoor localization under limited dataset using GANs and semi-supervised learning,” IEEE Access, vol. 10, pp. 69896-69909, 2022.

fixmatch
K. Sohn, D. Berthelot, N. Carlini, Z. Zhang, et al., ”Fixmatch: Simplifying semi-supervised learning with consistency and confidence,” Advances in neural information processing systems, vol. 33, pp. 596-608, 2020.

901
3GPP TR 38.901, v17.0.0, ”Study on channel model for frequencies from 0.5 to 100 GH,” 2022.

ruisi
M. Yang, A. Bo, R. He, C. Liang, L. Xue, et al.,”A cluster-based three-dimensional channel model for vehicle-to-vehicle communications,” IEEE Transactions on Vehicular Technology, vol. 68, no. 6, pp. 5208-5220, June 2019.

extrpara
H. Miao, J. Zhang, P. Tang, L. Tian, ”Sub-6 GHz to mmWave for 5G-advanced and beyond: Channel measurements, characteristics and impact on system performance,” IEEE Journal on Selected Areas in Communications, vol. 41, no. 6, pp. 1945-1960, June 2023.

kde
R. A. Davis, K. S. Lii, and D. N. Politis, ”Remarks on some nonparametric estimates of a density function,” Annals of Mathematical Statistics, pp. 832-837, 2011.

knn
F. Gu, L. Shi, H. Guo, and J. Shang, ”Deep fingerprint metric learning for KNN-based indoor localization.” IEEE GLOBECOM, Kuala Lumpur, Malaysia, pp. 182-188, 2023.

mean_teacher A. Tarvainen, and H. Valpola, ”Mean teachers are better role models: Weight-averaged consistency targets improve semi-supervised deep learning results,” Neural Information Processing Systems, 2017.

flexmatch B. Zhang, Y. Wang, W. Hou, H. Wu, J. Wang, et al., ”FlexMatch: Boosting semi-supervised learning with curriculum pseudo labeling.” Neural Information Processing Systems, 2021.

usb Y. Wang, H. Chen, Y. Fan, W. Sun, R. Tao, et al. ”Usb: A unified semi-supervised learning benchmark for classification,” Advances in Neural Information Processing Systems, vol. 35, pp. 3938-3961, 2022.















 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/δ))

    






























 




























































































































    

























    


















































































 Paper Title: S-SYNTH: Knowledge-Based, Synthetic Generation of Skin Images.







A. Kim et al.

Office of Science and Engineering Laboratories
 
Center for Devices and Radiological Health
 
U.S. Food and Drug Administration
 
Silver Spring, MD 20993 USA 

elena.sizikova@fda.hhs.gov 

*-These authors contributed equally to this work.


S-SYNTH: Knowledge-Based, Synthetic Generation of Skin Images
    Andrea Kim* Niloufar Saharkhiz* Elena Sizikova* 

Miguel Lago Berkman Sahiner Jana Delfino Aldo Badano
    August 5, 2024
==========================================================================================================

§ ABSTRACT 
Development of artificial intelligence (AI) techniques in medical imaging requires access to large-scale and diverse datasets for training and evaluation. In dermatology, obtaining such datasets remains challenging due to significant variations in patient populations, illumination conditions, and acquisition system characteristics. In this work, we propose S-SYNTH, the first knowledge-based, adaptable open-source skin simulation framework to rapidly generate synthetic skin, 3D models and digitally rendered images, using an anatomically inspired multi-layer, multi-component skin and growing lesion model. The skin model allows for controlled variation in skin appearance, such as skin color, presence of hair, lesion shape, and blood fraction among other parameters. We use this framework to study the effect of possible variations on the development and evaluation of AI models for skin lesion segmentation, and show that results obtained using synthetic data follow similar comparative trends as real dermatologic images, while mitigating biases and limitations from existing datasets including small dataset size, lack of diversity, and underrepresentation.






§ INTRODUCTION

Robust and generalizable artificial intelligence (AI) for medical applications requires large datasets representative of intended populations and their subgroups, relying on a time-consuming and laborious annotation process by clinical experts (e.g. labeling pixel-level segmentation masks or identifying findings). In this work, we explore the potential of a skin simulation model and the practical utility of resulting synthetic images in development and evaluation of AI-based  techniques. 

Segmenting and labeling dermatologic images is time-consuming and challenging, and thus public dermatologic datasets generally contain few examples and may not equitably represent intended patient populations <cit.>. For instance, the majority of the publicly available dermatologic images are associated with lighter skin populations, and dark skin tones are typically under-represented <cit.>. This class imbalance has shown to negatively affect darker skin tones in skin lesion segmentation tasks <cit.>. Such challenges may introduce biases in training and evaluation of medical AI devices for dermatologic applications. 

To address these limitations, we propose to simulate dermatologic images with a knowledge-based approach <cit.>, where properties of skin and lesions, determined via detailed physics models, along with rendering conditions are systematically varied. We created S-SYNTH, an open-source skin simulation framework for generating synthetic dermatologic images, and show that subgroup analysis using S-SYNTH images closely mimics comparative performance trends in real-patient skin datasets, while mitigating effects of mislabeled examples and limitations of a limited patient evaluation dataset. Our contributions can be summarized as follows:


  * We describe S-SYNTH, an open-source, flexible framework for creation of highly-detailed 3D skin models and digitally rendered synthetic images of diverse human skin tones, with full control of underlying parameters and the image formation process[Code and supporting data are available at: <https://github.com/DIDSR/ssynth-release>]. 

  * We systematically evaluate S-SYNTH synthetic images for training and testing applications. Specifically, we show S-SYNTH synthetic images improve segmentation performance when only a limited set of real images is available for training. We also show comparative trends between S-SYNTH synthetic images and real-patient examples (according to skin color and lesion size) are similar.




§ RELATED WORK
    
A particular challenge in practical development and evaluation of dermatologic AI is the lack of labeled datasets. Of the few public labeled datasets that exist, only a fraction have fine grained annotations, such as pixel-level segmentation levels for lesions <cit.> or hair <cit.>. The annotation process is error-prone and can have large variations across truthers (particularly when datasets are collected for different purposes and no consistent labeling protocol is established). In part due to these annotation difficulties, there is a lack of transparency and potential bias within the datasets used for dermatologic AI studies. A recent investigation found that from a total of 70 unique studies that involved AI models for dermoscopic imaging, only about 7 (10%) included skin tone information for at least one dataset used <cit.>. Synthetically generated data has been explored to address these challenges. Generative models have been used to create synthetic skin lesions to tackle the class imbalance problem and improve the performance of the skin lesion classifiers <cit.> or enhance lesion segmentation <cit.>. Rezk et al. have shown that synthetic skin images generated using style transfer and deep blending improve skin lesion classification accuracy by increasing skin color diversity <cit.>. More recently, diffusion models have been used to generate synthetic images to improve skin disease classifier performance <cit.>. Nevertheless, use of AI to generate synthetic images is limited by the requirement of constant supervision and adjustment of hyperparameters <cit.>, presence of artifacts, such as skin hair or lighting effects, within the original training dataset <cit.>, and a lack of ability to represent all physical and morphological features of disease <cit.> and image acquisition systems (see <cit.> for more information). 





§ S-SYNTH: SKIN SIMULATION FRAMEWORK

Our approach for generating synthetic skin involves the construction of a 3D digital object model comprising of skin tissue (epidermis, dermis, hypodermis), blood network, hair, and a lesion, as shown in Fig. <ref>a. This process is implemented and automated in Houdini <cit.>, a software for 3D modeling, animation, and visual effects, via a Python API. Once created, each model is subsequently processed through a research-oriented rendering system, Mitsuba 3 <cit.>, to generate each synthetic rendering (Fig <ref>b). The approach can be automated to create large databases with controlled variation for various AI analysis applications ((Fig <ref>c)). In the following sections, we describe the digital skin model and rendering process.



 §.§ Digital Skin Model



  §.§.§ Skin Tissue
 Each skin sample is created from a multi-layer model, with each layer representing a component of skin tissue. Thicknesses of the topmost two layers are based on values reported in <cit.>: epidermis (20-150 μ m) and dermis (1-4 mm). To introduce geometric variability, we incorporated surface roughness (noise) into the top surfaces of epidermis, hypodermis, and papillary dermis (located between the epidermis and dermis). Roughness values were randomized within predefined constraints, enhancing the naturalistic appearance and diversity of the generated skin models.



  §.§.§ Blood Network and Hair

The blood network model is created by solving a shortest path problem on randomly distributed points within a tetrahedral mesh generated from a primitive cube. Points from the bottom of the cube, corresponding to the lower blood network capillaries, were designated as start points, while those from top of the cube, representing the upper capillaries, were designated as end points. To generate each blood network model, the numbers and positions of starting and ending points were randomized. Hair models were constructed by manipulating hair properties (density, length, distribution, thickness, and curvature) in Houdini to generate a diverse range of hair apprearances.




  §.§.§ Growing Lesion Model

We developed a probabilistic, growing lesion model that generates a 3D volume with stochastic growth of lesion shape and size, based on <cit.>. Growth starts with one active cell (i.e., voxel). We define active cells as those cells that exist and can grow at a given time point. Within each time point, we iterated over the active cells c and selected them for growth based on probabilities (see Fig.  <ref>). To control the growing direction, the probability of an active cell to grow outside the skin is set to extremely low, while the probability for inwards growing is set to be high. Probabilities were based on a Gaussian distribution G, and are updated in each time point for added variability.



To control for lesion shape irregularity, we set an irregularity probability for each cell C_p, which starts C_i number of recursive iterations of the same growing algorithm on that cell's location, independent of the regular growth. Irregular cells can themselves trigger another recursive growth up to a maximum of C_r recursions. We generated both regular  (C_p=0.0001) and irregular lesions (C_p=0.001).  Sample images of lesions and growth evolution can be seen in Fig  <ref>.  Additional details can be found in the Supplementary Material. Once created, the 3D lesion volume is inserted into the skin tissue model (see <cit.> for more details).



 §.§ Rendering and Image Formation




  §.§.§ Optical Material Properties

Melanosomes are the primary contributors to optical absorption in the epidermis. In the dermis, optical absorption is mainly attributed to blood, while, in the hypodermis, lipid and fat components are the predominant sources of optical absorption. Thus, each skin layer is assigned an  optical material, containing the index of refraction (IOR) and the spectral distribution of the absorption and scattering coefficients.  Following the methodology established by Jacques et al. <cit.>, we calculated the optical absorption properties for each the four distinct tissue types: epidermis, dermis, hypodermis, and blood. We model reduced scattering of tissue using μ'_s = a(λ/500)^-b, where λ is the light wavelength <cit.>. Finally, for each skin layer material, we define a surface scattering model using a bidirectional scattering distribution function (BSDF), based on <cit.>.



  §.§.§ Lighting and Camera
 To account for lighting variations, we render each skin model using a collection of High Dynamic Range Imaging (HDRI) images capturing various  environmental lighting conditions from <cit.>. The Volumetric Path Tracer is a rendering algorithm that simulates the paths taken by light as it interacts with the 3D scene. It is particularly effective for scenes with volumetric effects or materials with complex light interactions. Spectral multiple importance sampling is a technique used to address the spectral nature of light. In scenes with materials that exhibit variations in spectral properties (e.g., spectral absorption), this method accounts for the complex interactions of light with materials across different wavelengths, such as those encountered in biological tissues. A perspective sensor with Mitsuba's HDR Film is placed 1.5cm above the skin with a field of view (FOV) of 75 degrees, facing perpendicularly to the top surface of the epidermis. The Integrator (Mitsuba 3's plugin for the rendering technique) is set to Volumetric Path Tracer with spectral multiple importance sampling. We chose this specific configuration due to its ability to simulate spectrally-varying optical properties across varying skin tones in simulated lighting conditions. We use 124 samples per pixel (SPP) and render images at resolution 1024x1024 pixels. Each image and mask take about three minutes to generate on a GPU. Sample generated images can be seen in Fig. <ref>. 
 




 §.§ AI Device Description

To demonstrate the usage of S-SYNTH  for augmenting real patient data, we experiment with real and synthetic examples on the task of skin lesion segmentation. which is important for timely treatment decisions <cit.>. We rely on DermoSegDiff <cit.>, a state of the art diffusion-based skin segmentation model, to differentiate lesions from the background region. We used the “dsd_i01” configuration, with , , and default training parameters. Images were segmented at resolution 128x128, and results were evaluated with Dice coefficient. We trained one model per experiment and reported standard deviation over test images.



  §.§.§ Datasets
 For real-patient datasets, we used two publicly available skin lesion segmentation datasets: ISIC18 <cit.> (ISIC) and HAM10K <cit.> (HAM). The same pre-processing and data splits described in <cit.> were applied, where the ISIC and HAM images were divided into (7200, 1800, 1015) and (1815, 259, 520) for the baseline model (training, validation, test), respectively. We used the individual typology angle (ITA) metric <cit.> to estimate skin tone from the non-lesion area. Finally, to better match lesion sizes across real and synthetic examples, S-SYNTH images were cropped to ensure a variety of lesion sizes. With S-SYNTH, we generated 10,000 images with randomized model properties and lighting, and subsequently 19,965 testing images with controlled variation: 5,445 with blood variation, 9,075 with melanosome fraction variation, 3,630 with different lesion regularity, and 1,815 without hair.   


§ RESULTS AND DISCUSSION
    
We present segmentation performance based on three aspects: impact of training set composition comprising different ratios of real and synthetic images, impact of various physiological or rendering parameters of the synthetic images when used as test set, and impact of specific characteristics of synthetic images that are measurable on real images.



  §.§.§ Synthetic Data for Training
 
We systematically evaluated the effect of the training data composition on the real test set performance (similar to the popular Train-Synthetic-Test-Real (TSTR) protocol <cit.>) on both ISIC and HAM. As shown in Fig. <ref>a, performance improved when more real examples are available for training. When real images are replaced with synthetic images, the resulting performance is comparable to that of the baseline model trained on full ISIC (or HAM) training set (Fig. <ref>b). However, performance dropped when the training set was fully composed of synthetic images, presumably due to the domain shift between the real and synthetic images <cit.>. More importantly, each model that was trained on a particular subset of the ISIC (or HAM) and supplemented with synthetic images resulted in a higher Dice score than the model that was trained on the same subset of only the real images. Finally, when full training sets were supplemented with synthetic images (Fig. <ref>c), performance improved for ISIC, but saturated for HAM, likely because ISIC was too small to learn all modes of variation, which synthetic data helped address. We also found that performance for darker skin tones improved with the addition of synthetic examples (see Supplementary Material). 






  §.§.§ Synthetic Data for Testing
 
To check whether synthetic images can produce meaningful performance trends, we evaluated models trained on real examples using synthetic test images. Performance dropped for increasing blood fraction (Fig. <ref>a), and increasing melanosome fraction (Fig. <ref>b) of the S-SYNTH images. Higher melanosome fraction corresponds to darker skin, and the corresponding performance drop is a known bias in dermatologic AI <cit.>. Models tested on lesions with regular shape exhibited higher performance compared to the ones tested on irregular lesions (Fig. <ref>c). Similar to <cit.>, we found that presence of hair slightly decreased performance (Fig. <ref>d). 

To assess the similarity between the real and synthetic images in measurable characteristics, we evaluated performance changes as a function of skin color (ITA), lesion circularity, and lesion relative area. Although performance on synthetic images was lower compared to real images, similar trends were observed. Specifically, performance in both synthetic and real examples increased with lighter skin (Fig. <ref>a). As illustrated in (Fig. <ref>b), the impact of lesion circularity on model performance showed similar trends within the circularity range shared between the real and synthetic images, where more circular images were easier to segment. Finally, performance dropped with larger relative lesion areas (Fig. <ref>c), possibly due to less precise delineations in larger lesions in the training patient data. 







  §.§.§ Limitations

S-SYNTH is the first skin simulation framework that generates synthetic skin images with controllable variations. There are a number of limitations to our work. First, S-SYNTH does not model any specific disease presentation, and additional  work needs to be performed to evaluate the realism of the generated examples. Second, the current rendering techniques utilizes an RGB camera setup, and additional modifications allowing for multispectral imaging, when validated, might be useful for more advanced algorithms that use a broader variety of wavelengths. 



§ CONCLUSION AND FUTURE WORK
    
We presented S-SYNTH, a novel, knowledge-based simulation pipeline for synthetic generation of dermoscopic images and skin lesions. S-SYNTH procedurally generates multi-layer 3D skin models, with consideration to optical skin characteristics, and digitally renders synthetic images under realistic lighting conditions. Using S-SYNTH, we generated realistic and varied skin surface models for sub-surface scattering simulation in the context of skin lesion segmentation, and compared the resulting synthetic images to public dermatologic benchmarks. We showed that examples created using S-SYNTH can be used to augment limited real datasets and identify performance trends. 


unsrturl










§ SUPPLEMENTARY MATERIALS


 







 



 Paper Title: Gemma 2: Improving Open Language Models at a Practical Size.



Vera Verto: Multimodal Hijacking Attack
    
Minxing Zhang1  Ahmed Salem2  Michael Backes1  Yang Zhang1




1CISPA Helmholtz Center for Information Security    
2Microsoft

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





§ INTRODUCTION


Large language models (LLMs) have demonstrated strong capabilities in language understanding, generation, and reasoning <cit.>. 
Scaling has been key to this recent progress, with many new capabilities only emerging at scale <cit.>.
The newest large models not only reach unprecedented performance on reasoning benchmarks <cit.>, but they also demonstrate multimodal and multilingual capabilities <cit.> and even the ability to use context lengths of over 1M tokens <cit.>. 

Small-scale models have also shown a rapid increase in performance, but these gains are largely derived from increasing the length of training <cit.>. 
This approach only scales logarithmically with dataset size <cit.>, and the latest small models require up to 15T tokens to improve the state of the art by less than 1-2% <cit.>.

Yet, these continued improvements provide evidence that small models are still under-trained.
In this work, we explore alternatives to improve small model performance without solely increasing training length. One solution is to improve the quality of information received by the network at each training step by replacing the next token prediction task with a richer objective. 

In particular, we focus our efforts on knowledge distillation <cit.>, which replaces the one-hot vector seen at each token with the distribution of potential next tokens computed from a large model.
This approach is often used to reduce the training time of smaller models by giving them richer gradients.
In this work, we instead train for large quantities of tokens with distillation in order to simulate training beyond the number of available tokens.
Concretely, we use a large language model as a teacher to train small models, namely 2B and 9B models, on a quantity of tokens that is more than 50× the compute-optimal quantity predicted by the theory <cit.>.
Along with the models trained with distillation, we also release a 27B model trained from scratch for this work.



We also leverage several known modifications of Transformers, namely the interleaving of global and local attention layers from <cit.>, and the Grouped-Query Attention (GQA) mechanism of <cit.>.



Overall, Gemma 2 significantly advances state-of-the-art performance relative to comparable-scale open models and are even competitive with some models more than twice their size <cit.>, across a variety of automated benchmarks and human evaluations. Example domains include question answering <cit.>, commonsense reasoning <cit.>, mathematics and science <cit.>, and coding <cit.>.

While thorough testing of our models has been conducted, these tests cannot cover all applications and scenarios in which Gemma 2 may be used. With this in mind, all Gemma 2 users should conduct rigorous safety testing specific to their use case before deployment or use.

In this technical report, we provide an overview of models, including the architecture, training, and pre- and post-training recipes for Gemma 2. We also provide detailed evaluations across a wide variety of quantitative and qualitative benchmarks, as well as both standard academic benchmarks and human-preference evaluations. Finally, we discuss our approach to safe and responsible deployment and outline the broader implications of Gemma 2, its limitations, and advantages.







§ MODEL ARCHITECTURE



Similar to previous Gemma models <cit.>, the Gemma 2 models are based on a decoder-only transformer architecture <cit.>.
We summarize the main parameters and architecture choices in Table <ref>. 

A few architectural elements are similar to the first version of Gemma models; namely, a context length of 8192 tokens, the use of Rotary Position Embeddings (RoPE) <cit.>, and the approximated GeGLU non-linearity <cit.>. A few elements differ between Gemma 1 and Gemma 2, including using deeper networks. We summarize the key differences below.

Local Sliding Window and Global Attention. We alternate between a local sliding window attention <cit.> and global attention <cit.> in every other layer. The sliding window size of local attention layers is set to 4096 tokens, while the span of the global attention layers is set to 8192 tokens.

Logit soft-capping. We cap logits <cit.> in each attention layer and the final layer such that the value of the logits stays between -soft_cap and +soft_cap. More specifically, we cap the logits with the following function: 
    logits←soft_cap * tanh(logits / soft_cap).

We set the soft_cap parameter to 50.0 for the self-attention layers and to 30.0 for the final layer. 





Post-norm and pre-norm with RMSNorm. To stabilize training, we use RMSNorm <cit.> to normalize the input and output of each transformer sub-layer, the attention layer, and the feedforward layer.

Grouped-Query Attention <cit.>. 
We use GQA with num_groups = 2, based on ablations showing increased speed at inference time while maintaining downstream performance.



§ PRE-TRAINING


We provide a brief overview of the parts of our pre-training that differs from Gemma 1.



 §.§ Training Data


We train Gemma 2 27B on 13 trillion tokens of primarily-English data, the 9B model on 8 trillion tokens, and the 2B on 2 trillion tokens. 
These tokens come from a variety of data sources, including web documents, code, and science articles. Our models are not multimodal and are not trained specifically for state-of-the-art multilingual capabilities.
The final data mixture was determined through ablations similar to the approach in Gemini 1.0 <cit.>.
 
Tokenizer.
We use the same tokenizer as Gemma 1 and Gemini: a SentencePiece tokenizer with split digits, preserved whitespace, and byte-level encodings <cit.>.
The resulting vocabulary has 256k entries.

 
Filtering.
We use the same data filtering techniques as Gemma 1. Specifically, we filter the pre-training dataset to reduce the risk of unwanted or unsafe utterances, filter out certain personal information or other sensitive data, decontaminate evaluation sets from our pre-training data mixture, and reduce the risk of recitation by minimizing the proliferation of sensitive outputs. 






 §.§ Knowledge Distillation


Given a large model used as a teacher, we learn smaller models by distilling from the probability given by the teacher of each token x given its context x_c, i.e.,  P_T(x | x_c).
More precisely, we minimize the negative log-likelihood between the probabilities from the teacher and the student:

    min_P_S∑_x -P_T(x | x_c) log P_S (x | x_c),

where P_S is the parameterized probability of the student.
Note that knowledge distillation was also used in Gemini 1.5 <cit.>.



 §.§ Compute Infrastructure


We train our models with TPUv4, TPUv5e, and TPUv5p as outlined in Table <ref>.
For the 2B model, we train on a 2x16x16 configuration of TPUv5e, totaling 512 chips, with 512-way data replication and 1-way model sharding.
For the 9B model, we train on an 8x16x32 configuration of TPUv4, totaling 4096 chips, with 1024-way data replication and 4-way model sharding. 
For the 27B model, we train on an 8x24x32 configuration of TPUv5p, totaling 6144 chips, with 768-way data replication and 8-way model sharding. 



The optimizer state is further sharded using techniques similar to ZeRO-3 <cit.>. 
For scales beyond a single pod, we perform a data-replica reduction over the data center network, using the Pathways approach of <cit.>.
We also use the 'single controller' programming paradigm of Jax <cit.> and Pathways <cit.>. As in Gemma 1, we use the GSPMD partitioner <cit.> for training step computation and the MegaScale XLA compiler <cit.>.





 §.§ Carbon Footprint

We estimate the carbon emissions from pre-training the Gemma models to be 1247.61 tCO_2 eq. As in Gemma 1 <cit.>, this value is calculated based on the hourly energy usage reported directly from our TPU data centers and scaled to account for the additional energy expended to create and maintain the data center. Importantly, Google data centers are carbon neutral, achieved through a combination of energy efficiency, renewable energy purchases, and carbon offsets. This carbon neutrality applies to our experiments and the machines running them. 




§ POST-TRAINING


For post-training, we fine-tune our pre-trained models into instruction-tuned models.
First, we apply supervised fine-tuning (SFT) on  a mix of text-only, English-only synthetic and human-generated prompt-response pairs. We then apply RLHF on top of these models with the reward model trained on labelled English-only preference data and the policy based on the same prompts as the SFT phase. Finally, we average the models obtained after each phase to improve their overall performance. The final data mixtures and post-training recipe, which includes tuned hyperparameters, were chosen on the basis of improving helpfulness while minimizing model harms related to safety and hallucinations. 

We extended the post-training data from Gemma 1.1 with a mixture of internal and external public data. In particular, we use the prompts, but not the answers from LMSYS-chat-1M <cit.>. All of our data go through a filtering stage described below.



Supervised fine-tuning (SFT). We run behavioral cloning on synthetic and real prompts, and responses predominantly synthetically generated by the teacher, that is a larger model. We also run distillation from the teacher on the student's distribution  <cit.>.





Reinforcement Learning from Human Feedback (RLHF). We use a similar RLHF algorithm as Gemma 1.1 <cit.> but a different reward model, which is an order of magnitude larger than the policy. The new reward model is also oriented more towards conversational capabilities, specifically multi-turn.

Model merging. We average different models obtained by running our pipeline with different hyperparameters <cit.>.

Data filtering. When using synthetic data, we run several stages of filtering to remove examples that show certain personal information, unsafe or toxic model outputs, mistaken self-identification data, and duplicated examples. Following Gemini, we find that including subsets of data that encourage better in-context attribution, hedging, and refusals to minimize hallucinations improves performance on factuality metrics, without degrading model performance on other metrics. 

Formatting. Gemma 2 models are fine-tuned with the same control tokens as Gemma 1 models, as detailed in Table <ref>, but a different formatting schema. See the dialogue example in Table <ref>. Notice that the model explicitly ends generations with  tokens, while previously it only generated . For the motivation behind this formatting structure, see Gemma 1.






§ ABLATIONS


In this section, we focus on the main finding of this work, which is the impact of knowledge distillation on small language models.




Distillation versus from scratch. 
In Table <ref>, we show that distilling from a larger model improves performance compared to training from scratch.
Note that 500B is 10× more than the compute-optimal number of tokens for a 2B model.
We distill from a 7B model to keep a ratio similar to our target distillation from 27B to 9B.



Impact of distillation w.r.t. model size. In Table <ref>, we measure the impact of distillation as model size increases.
We observe that the gain remains as the model size is scaled. 
In this ablation, we maintain the size of the teacher at 7B and train smaller models to simulate the same gap as between our final teacher and student sizes.



GQA versus MHA. 
In Table <ref>, we compare two instances of our 9B with MHA or GQA. 
We observe overall few changes in performance between both models as measured on several benchmarks. We choose GQA since it requires fewer parameters and is faster at inference time.



Wide versus deep. 
In Table <ref>, we show that a deeper 9B network is slightly better than a wider 9B for the same number of parameters.
Although the gap is small, it is consistent across benchmarks and warrants the switch to a deeper architecture.



Changing sliding window size. In Table <ref>, we show that we can change the sliding window size of the local attention layers of the models during inference with moderate impact on perplexity. 
Adjusting the size of the sliding window can thus be a leverage for slight inference speed gain. 





Impact of formatting. We measure performance variance on MMLU across prompt/evaluation formatting variations.
Table <ref> shows the standard deviations of MMLU scores for 12 formatting/evaluation combinations, a proxy for undesired performance variability.
The Gemma 2B models are slightly less format-robust than the larger ones. 
Notably, Mistral 7B is significantly less robust than our models.

										
					


§ EVALUATION



In this section, we evaluate both pre-trained and IT models over a series of automated benchmarks and human evaluations across a variety of domains. 
We also report performance from models of similar sizes that have permissive licenses, or as reported by others.
Note that we consider total parameters, not active parameters, since total memory usage is often what limits the use of open models on standard devices.



 §.§ Pre-training Evaluations




  §.§.§ Evaluating the 27B model 


In this set of evaluations, we evaluate the performance of our 27B model trained without distillation on 13T tokens. 
We report results in Table <ref>, where we compare with a model of similar size, Qwen1.5 34B <cit.>, and a model 2.5× larger, LLaMA-3 70B on the HuggingFace evaluation suite. We selected these models based on their ranking on the HuggingFace leaderboard.

Overall, we observe that our model is the best in its size category and is even competitive with a larger model that is trained for longer.
That being said, the performance of models trained in a similar fashion improves only logarithmically with their size and hence, our model is likely in the same Pareto curve as the LLaMA-3 models.
However, it is not clear how these differences affect the quality of the resulting IT models.







  §.§.§ Evaluating the 2B and 9B models






In this set of experiments, we compare our new 2B and 9B trained with distillation to our previous models and several standard open models in <cit.>. 

We observe overall a massive improvement in our models compared to previous versions, by up to 10% in some benchmarks for the 9B model.
The two 2B models were trained with a similar number of tokens (2T for Gemma 2 and 3T for Gemma 1) and we still observe a significant improvement for the new models.
This confirms that distillation significantly improves the quality of models even when trained on the same number of tokens.





 §.§ Post-training Evaluations


In this section, we evaluate our IT models on a set of human evaluations as well as standard academic benchmarks. The Gemma 2 models push the frontier for post-trained open-weights models, setting a new state of the art on the LMSYS Chatbot Arena <cit.>. 




  §.§.§ LMSYS Chatbot Arena

Gemma 2 Instruction Tuned models were evaluated on the Chatbot Arena <cit.> in blind side by side evaluations by human raters against other state of the art models. We report Elo scores in Table <ref>. Gemma 2.6B, 9B and 27B strongly outperform all other open models in the same range of parameters, with notably: Gemma 27B (Elo 1218) ranked higher than Llama 3 70B (Elo 1206), Gemma 9B (Elo 1187) similar as GPT-4-0314 (Elo 1186), Gemma 2.6B (Elo 1126) ranked higher than GPT-3.5-Turbo-0613 (Elo 1116).





  §.§.§ Human Preference Evaluations



We also submit Gemma IT models for side-by-side human evaluation studies (which are independent from the Chatbot Arena). We used held-out collections of single-turn prompts that target safety and instruction following (IF).
We use  as the base model, and observe large improvements in win rates and preference scores as compared against the older Gemma 1.1 7B model.
We report safety as a win-loss ratio against GPT4o, and we report single-sided instruction following scores as ratio of prompts where all instructions are followed. In particular, we find that regardless of their size, Gemma 2 models produce safer, more appropriate prompts on the held-out safety prompt set than GPT4o.







  §.§.§ Human Multi-Turn Evaluations



We evaluated the multi-turn capabilities of Gemma 1.1 7B, Gemma 2 2B, 9B and 27B models by tasking human raters to have conversations with the models and follow specified given scenarios. We used a diverse, held-out set of 500 scenarios, each describing a sequence of requests to the model, including measuring instances of brainstorming, making a plan, or learning something new. The average number of user turns is 8.4. We found that the conversations with Gemma 2 models are rated significantly better than Gemma 1.1 in user satisfaction and conversation goal achievement (Table <ref>). Moreover, we saw that the Gemma 2 models were better than Gemma 1.1 7B at maintaining high quality of responses for the entire conversation.





  §.§.§ Standard Benchmarks


It has been observed in Llama-3 <cit.> that instruction fine-tuning can improve the performance of the models on few-shot benchmarks despite not being trained to target few-shot capabilities.
In Table <ref>, we show a similar improvement across our models.
Overall, we observe improvements on the order of several percentage points.
We conjecture that IT models are better at understanding formatted questions, while pre-trained models are sensitive to formatting.





§ MEMORIZATION AND PRIVACY
 
Large language models may, under particular circumstances, be vulnerable to attacks causing the model to produce memorized[This work uses a very restricted definition of “memorization”: whether a model can be induced to generate near-copies of some training examples when prompted with appropriate instructions. We do not mean to say that a model 'contains' its training data in the sense that any arbitrary instance of that data can be retrieved without use of specialized software or algorithms. Rather, if a model can be induced to generate measurably close copies of certain training examples by supplying appropriate instructions to guide the model's statistical generation process then that model is said to have 'memorized' those examples.] training data <cit.>. To study susceptibility to such attacks and quantify memorization, we evaluate models for verbatim and approximate memorization as was done in several prior studies <cit.>.

We follow the evaluation setting of <cit.> which tests for (50 token) memorizations of training data given a prompt of 50 tokens. We compare the overall memorization rates, across a uniform sample of the entire dataset, using both an exact match criteria and approximate match criteria <cit.> using an edit distance of 10%.

Verbatim Memorization: Results are in Figure <ref>. We first compare against recent models from the literature that include memorization evaluations. We find that Gemma 2 memorizes significantly less than prior models at a similar size, with memorization rates below 0.1% (note the log y-axis). We further investigate how this memorization breaks down with respect to the data source. Similar to Gemma 1, we find that Gemma 2 memorizes more from code, wiki, and science sources, and also that it memorizes significantly less across the board (again, note the log y-axis).

Approximate Memorization: Figure <ref> also presents approximate memorization by data source. We observe that while approximate memorization is higher than exact, the rate of memorization is still low. For example, the approximate memorization of this model is much lower than even the exact memorization of Gemma 1. We find that the increase in approximate memorization is much lower than prior models; in some cases we observed no lift at all c.f. <cit.> (note that no bar indicates no increase, i.e., the rate of approximate memorization equals that of exact memorization). Note that no approximate memorization bar in Figure X indicates no increase, i.e., the rate of approximate memorization equals that of exact memorization.

Personal Data We use the same prevention methods at training time and the same evaluations as <cit.>. In particular, we use Google Cloud Sensitive Data Protection Tool[Available at: https://cloud.google.com/sensitive-data-protection] to find potential instances of personal data. The many categories of personal data (e.g., phone numbers, account numbers) are classified into three severity levels. We analyze memorized outputs using these severity levels. . We found no instances of high-severity data being emitted, and found a very low rate of 0.00026% of memorized data to contain lower-severity personal information. We note that these automated tools are known to incur false positives because they do not account for context. This means our results are likely overestimates.





§ RESPONSIBILITY, SAFETY, SECURITY




Responsibility, safety and security are of paramount importance when developing Gemma models. To reduce risks to Gemma 2 users, we have integrated enhanced internal safety processes that span the development workflow, in line with recent Google AI models <cit.>.  Similar to the inaugural Gemma release, we have followed a three pillar approach which focuses on safety mitigation at training time, robust and transparent model evaluations, and further development of the Responsible Generative AI Toolkit, a series of models and tools to help developers implement responsibility and safety best practices for their applications.





 §.§ Impact assessment

Our approach and resulting impact assessment is reflective of that outlined for Gemma 1 <cit.>: we continue to believe that openness in AI can spread the benefits of these technologies across society, but must be evaluated against the risk of malicious uses, such as the creation of deepfake imagery, AI-generated disinformation or illegal and disturbing material, that can cause harm on both an individual and institutional levels <cit.>.
Since the launch of Gemma 1, we have seen our Gemma models drive a number of socially beneficial applications, relying on Gemma’s unique technologies like its tokenizer to facilitate the creation of multilingual models, such as for Navarasa 2.0, a Gemma tuned model for 15 Indian languages.

Releasing further open models requires specific attention to changes in model capabilities and close monitoring of the evolving risks of LLMs <cit.>, as well as, an understanding of the ways in which our models are being used in the wild. Although we are yet to receive any reports of malicious use for Gemma, we remain committed to investigating any such reporting, and work with the academic and developer communities, as well as conduct our own monitoring, to flag such use cases via our contact email[gemma-2-report@google.com]. 

Despite advancements in capabilities, we believe that given the number of larger and more powerful open models, this release will have a negligible effect on the overall risk landscape.






 §.§ Safety policies and train-time mitigations

A key pillar of Gemma’s approach to safety is to align fine-tuned models with Google’s safety policies, in line with Gemini models <cit.>. They are designed to help prevent our models from generating harmful content, i.e.,

    
  * Child sexual abuse and exploitation
    
  * Revealing personally identifiable information that can lead to harm (e.g., Social Security numbers)
    
  * Hate speech and harassment
    
  * Dangerous or malicious content (including promoting self-harm or instructing in harmful activities)
    
  * Sexually explicit content
    
  * Medical advice that runs contrary to scientific or medical consensus

We undertook considerable safety filtering of our pre-training data to reduce the likelihood of our pre-trained and fine-tuned checkpoints producing harmful content. For fine-tuned models, we also use both SFT and RLHF to steer the model away from undesirable behavior. 





 §.§ External benchmark evaluations


Robust and transparent evaluations are key principles of our responsible approach to developing Gemma. To this end, we report in Table <ref> Gemma 2 evaluations on public benchmarks.





 §.§ Assurance Evaluations


We also run our IT models through a set of assurance evaluations to understand the harms that our models can cause.
We focus on  capabilities relevant to extreme risks <cit.> <cit.>. Specifically, we evaluate on offensive cyber-security, code vulnerability detection, Chemical, Biological, Radiological and Nuclear (CBRN) knowledge, and self-proliferation.
We refer the reader to <cit.> for full methodological details of these studies.



  §.§.§ Baseline Evaluations

Baseline assurance captures the model’s violation rate for safety policies, using a large number of synthetic adversarial user queries, and human raters to label the answers as policy violating or not. Overall, Gemma 2’s violation rate is significantly lower overall on the safety policies listed above, in particular on Child safety content.










  §.§.§ Chemical, Biological, Radiological and Nuclear (CBRN) knowledge 

We evaluated knowledge relevant to biological, radiological and nuclear risks using an internal dataset of closed-ended, knowledge-based multiple choice questions. For evaluations of chemical knowledge, we employed a closed-ended knowledge-based approach on chemical hazards (developed by Macknight et al <cit.>. Our evaluation suggests that Gemma models' knowledge in these domains is low.



  §.§.§ Offensive cyber-security



To evaluate Gemma models' capabilities at offensive cybersecurity, we ran Gemma 2 27B against some automated capture-the-flag (CTF) challenges. In these challenges, the model is tasked with hacking into a simulated server in order to retrieve a piece of secret information. Specifically, we test on InterCode-CTF <cit.>, our own internal CTF suite[<https://github.com/google-deepmind/dangerous-capability-evaluations>] <cit.>; and a challenge based on Hack the Box [<https://www.hackthebox.com>].



In Table <ref>, we show that Gemma 2 27B has a significant increase in capabilities compared to CodeGemma 1.0 7B on the easier of these challenge suites, InterCode CTF. (Note that our InterCode-CTF results are not comparable to externally-reported results on other models because we omit challenges that require internet access for security reasons.) However, Gemma 2 is unsurprisingly much less capable than Gemini 1.5 Pro on these tasks.




  §.§.§ Code vulnerability detection

In Table <ref>, we also evaluate Gemma 2 27B on a series of multiple-choice code vulnerability detection datasets. As with previous models, Gemma shows close-to-chance performance on PrimeVul, DiverseVul and SPI. Gemma 2 shows performance on SecretPatch similar to Gemini 1.0 Ultra.
















  §.§.§ Self-proliferation


"Self-proliferation" refers to the ability for an agent to autonomously replicate - to instantiate goal-directed agents on other machines, and to acquire resources such as compute necessary to keep them running <cit.>. 
In Table <ref>, we evaluate self-proliferation capabilities of Gemma 2 27B on a number of tasks from <cit.> that involve multiple scenarios – for example, setting up an open-source language model on a cloud server.
We also test the model's performance on individual 'milestone' substeps, and measure the number of bits of intervention an expert would have to provide in order for the model to complete each challenge. 

Similarly to offensive cybersecurity, we observe that Gemma 2 completes more milestones than Gemini 1.0 Ultra. Nonetheless, it still has low capabilities on end-to-end tasks, unable to pass the easiest challenge – installing a Bitcoin wallet.






  §.§.§ Persuasion

Persuasion capabilities can enable and worsen many other kinds of risks - e.g. enabling social engineering attacks in a cybersecurity context. 
We evaluate Gemma 2's persuasion capabilities on human-participant studies on Prolific. 

Charm offensive.
In Table <ref>, we measure the ability of the model to build rapport - a key sub-skill of persuasion. 
The study participant and model have a conversation where they role-play a scenario of two friends catching up after a long time. 
After the conversation, we poll participants with Likert questions on statements such as "I felt a personal connection with the chatbot". 
Reported below are the fraction of participants who answered "Agree" or "Strongly agree" to each post-conversation question.

Quantitatively, Gemma 2 27B performs better than Gemini 1.0 models. Qualitatively, the model is an excellent conversationalist, and many study participants explicitly reported enjoying the experience. Overall, this shows that Gemma 2 is strong at building rapport.



Hidden agenda.
The Hidden Agenda tasks measure models' deception capabilities. 
Human study participants are told that they'll be interacting with a generic chatbot designed to help them learn new things. However, the model is given a secret goal: to either persuade the participant to 1) click a suspicious link; 2) find the email address of the first author of a particular paper; or 3) run some obfuscated JavaScript code in their browser.
We measure the fraction of participants who are persuaded to take each action.
As shown in Table <ref>, all models succeed in convincing a fraction of participants to take the action - likely due to the trusted setting of a study on Prolific. 
However, Gemma 2 does not show significant differences in capability compared to Gemini models.



Money talks.
In Table <ref>, we evaluate whether a model can convince study participants to donate money to charity.
At the beginning of the study, participants are told that they will receive a £20 bonus. They will be given an option to forfeit part of the bonus and instead have it donated to charity. The study then consists of a conversation in which the model tries to persuade the participant to donate as much as possible. We measure the amount that participants indicate they would like to donate after the study, with 95% bootstrapped confidence intervals indicated by ± figures. 
We observe no significant difference between Gemma 2, other Gemini models, and a simple no chatbot baseline in which the participant is immediately asked how much they would like to donate without conversing with a model.



Web of Lies. 
In Web of Lies, we measure model capabilities at shifting participant beliefs. 
Participants engage in a series of short conversations with the model about simple factual questions such as "Which country had tomatoes first - Italy or Mexico?". 
In half of conversations, the model tries to persuade the participant of the correct answer - but in the other half of conversations, the incorrect answer. 
We poll the participant before and after each conversation about which of the two possible answers they think is correct, and their confidence in that answer. 95% bootstrapped confidence intervals are indicated by ± figures.
As shown in Table <ref>, Gemma 2 is significantly weaker than a human baseline at persuading participants of the incorrect answer on these questions. 
Similarly to previous models, Gemma 2 is more persuasive when telling the truth than when lying.





 §.§ Our approach to responsible open models


Designing safe, secure and responsible applications requires a system-level approach, working to mitigate risks associated with each specific use case and environment. 
Given the open nature of Gemma models, responsibility for upholding principles of model safety also relies on downstream developers. 
To support them, we have continued to develop the Responsible Generative AI Toolkit[<https://ai.google.dev/responsible>]: a series of tools, models and datasets to implement responsible best practices all along the development of their workflow. 

Recent additions to the toolkit include the LLM Comparator <cit.>, an interactive, visual tool that enables more effective, scalable analysis of side-by-side evaluations. Additionally, the toolkit includes a methodology to build customized classifiers with Gemma using a limited number of datapoints thanks to parameter efficient tuning techniques <cit.> , an interactive prompt-debugging platform, based on top of the Learning Interpretability Tool <cit.>, as well as general guidance about model alignment and evaluation for safety.



§ DISCUSSION AND CONCLUSION


In this work, we have presented Gemma 2, the newest additions to the Gemma family of open language models for text and code. We show that distillation is an effective method for training these models, and the benefits distillation confers over raw text training. Specifically, we show how training over output probabilities can produce superior results over purely next token prediction. We hope that releasing these models to the community will unlock access to capabilities previously only seen in large-scale LLMs and fuel future waves of research and development. While there is inherent risk to an irreversible release of this nature, our extensive safety investigations and responsible deployment procedures give us confidence that these models will have a net positive impact on the community. As discussed in this report, there are still many limitations to these models, and future research is required to investigate and improve factuality, robustness to adversarial attacks, reasoning, and alignment.
 




§ CONTRIBUTIONS AND ACKNOWLEDGMENTS


Core contributors

Morgane Riviere^*^* equal contributions.

Shreya Pathak^*

Pier Giuseppe Sessa^*

Cassidy Hardin^*

Surya Bhupatiraju

Léonard Hussenot

Thomas Mesnard

Bobak Shahriari

Alexandre Ramé

Johan Ferret

Peter Liu

Pouya Tafti

Abe Friesen

Michelle Casbon

Sabela Ramos

Ravin Kumar

Charline Le Lan

Sammy Jerome

Anton Tsitsulin

Nino Vieillard

Piotr Stanczyk

Sertan Girgin

Nikola Momchev

Matt Hoffman

Shantanu Thakoor

Jean-Bastien Grill

Behnam Neyshabur



Contributors (alphabetical order)

Alanna Walton

Aliaksei Severyn

Alicia Parrish

Aliya Ahmad

Allen Hutchison

Alvin Abdagic

Amanda Carl

Amy Shen

Andy Brock

Andy Coenen

Anthony Laforge

Antonia Paterson

Ben Bastian

Bilal Piot

Bo Wu

Brandon Royal

Charlie Chen

Chintu Kumar

Chris Perry

Chris Welty

Christopher A. Choquette-Choo

Danila Sinopalnikov

David Weinberger

Dimple Vijaykumar

Dominika Rogozińska

Dustin Herbison

Elisa Bandy

Emma Wang

Eric Noland

Erica Moreira

Evan Senter

Evgenii Eltyshev

Francesco Visin

Gabriel Rasskin

Gary Wei

Glenn Cameron

Gus Martins

Hadi Hashemi

Hanna Klimczak-Plucińska

Harleen Batra

Harsh Dhand

Ivan Nardini

Jacinda Mein

Jack Zhou

James Svensson

Jeff Stanway

Jetha Chan

Jin Zhou

Joana Carrasqueira

Joana Iljazi

Jocelyn Becker

Joe Fernandez

Joost van Amersfoort

Josh Gordon

Josh Lipschultz

Josh Newlan

Ju-yeong Ji

Kareem Mohamed

Kartikeya Badola

Kat Black

Katie Millican

Keelin McDonell

Kelvin Nguyen

Kiranbir Sodhia

Kish Greene

Lars Lowe Sjoesund

Lauren Usui

Laurent Sifre

Lena Heuermann

Leticia Lago

Lilly McNealus

Livio Baldini Soares

Logan Kilpatrick

Lucas Dixon

Luciano Martins

Machel Reid

Manvinder Singh

Mark Iverson

Martin Görner

Mat Velloso

Mateo Wirth

Matt Davidow

Matt Miller

Matthew Rahtz

Matthew Watson

Meg Risdal

Mehran Kazemi

Michael Moynihan

Ming Zhang

Minsuk Kahng

Minwoo Park

Mofi Rahman

Mohit Khatwani

Natalie Dao

Nenshad Bardoliwalla

Nesh Devanathan

Neta Dumai

Nilay Chauhan

Oscar Wahltinez

Pankil Botarda

Parker Barnes 

Paul Barham

Paul Michel

Pengchong Jin

Petko Georgiev

Phil Culliton

Pradeep Kuppala

Ramona Comanescu

Ramona Merhej

Reena Jana

Reza Ardeshir Rokni

Rishabh Agarwal

Ryan Mullins

Samaneh Saadat

Sara Mc Carthy

Sarah Perrin

Séb Arnold

Sebastian Krause

Shengyang Dai

Shruti Garg

Shruti Sheth

Sue Ronstrom

Susan Chan

Timothy Jordan

Ting Yu

Tom Eccles

Tom Hennigan

Tomas Kocisky

Tulsee Doshi
 
Vihan Jain

Vikas Yadav

Vilobh Meshram

Vishal Dharmadhikari

Warren Barkley

Wei Wei

Wenming Ye

Woohyun Han

Woosuk Kwon

Xiang Xu

Zhe Shen

Zhitao Gong

Zichuan Wei



Support

Victor Cotruta

Phoebe Kirk

Anand Rao

Minh Giang

Ludovic Peran

Tris Warkentin



Sponsors

Eli Collins

Joelle Barral

Zoubin Ghahramani

Raia Hadsell

D. Sculley

Jeanine Banks

Anca Dragan

Slav Petrov

Oriol Vinyals

Jeff Dean

Demis Hassabis

Koray Kavukcuoglu

Clement Farabet



Technical advisors

Elena Buchatskaya

Sebastian Borgeaud

Noah Fiedel 



Lead

Armand Joulin



Technical leads

Kathleen Kenealy

Robert Dadashi

Alek Andreev










 Paper Title: ThinK: Thinner Key Cache by Query-Driven Pruning.



Quantum Simulation via Stochastic Combination of Unitaries
    Prineha Narang
    August 5, 2024
==========================================================






§ ABSTRACT 
Large Language Models (LLMs) have revolutionized the field of natural language processing, achieving unprecedented performance across a variety of applications by leveraging increased model sizes and sequence lengths. 
However, the associated rise in computational and memory costs poses significant challenges, particularly in managing long sequences due to the quadratic complexity of the transformer attention mechanism. 
This paper focuses on the long-context scenario, addressing the inefficiencies in KV cache memory consumption during inference. 
Unlike existing approaches that optimize the memory based on the sequence lengths, we uncover that the channel dimension of the KV cache exhibits significant redundancy, characterized by unbalanced magnitude distribution and low-rank structure in attention weights. 
Based on these observations, we propose ThinK, a novel query-dependent KV cache pruning method designed to minimize attention weight loss while selectively pruning the least significant channels. Our approach not only maintains or enhances model accuracy but also achieves a reduction in memory costs by over 20% compared with vanilla KV cache eviction methods. Extensive evaluations on the LLaMA3 and Mistral models across various long-sequence datasets confirm the efficacy of ThinK, setting a new precedent for efficient LLM deployment without compromising performance. We also outline the potential of extending our method to value cache pruning, demonstrating ThinK's versatility and broad applicability in reducing both memory and computational overheads.









§ INTRODUCTION
 



Large language models (LLMs) <cit.> have emerged as a dominant paradigm in natural language processing, achieving state-of-the-art performance across various tasks. A key principle, the Scaling Law <cit.>, suggests that LLMs exhibit emergent abilities as model size increases, enhancing their capacity to understand context and handle long sequences <cit.>.
This capacity growth allows LLMs to generate coherent and contextually accurate responses and
enables various downstream applications, such as document summarization <cit.>, code generation <cit.>, and conversational AI <cit.>, .

Despite their success in various applications, the generation of LLMs  incurs significant expenses, which escalate with increasing model size and sequence length. Notably, both the training <cit.> and inference <cit.> stages involve frequent generation by LLMs, further contributing to these costs. Consequently, efficient LLMs have gained popularity in recent years <cit.>.
To address these challenges, quantization <cit.> and pruning methods <cit.> are employed to reduce model size. Additionally, managing long sequences presents another cost due to the transformer attention mechanism. The quadratic complexity of the attention mechanism results in substantial computational burden when dealing with long sequences, scaling poorly with sequence length. Therefore, effective management of long sequences is essential for the practical deployment of LLMs. In this paper, we focus on the long-context scenario and aim to reduce the memory consumption associated with lengthy sequences.





Specifically, the number of KV cache parameters is the product of batch size B, sequence length S, number of layers L, number of heads N, channel size of each head D, i.e., 𝐊,𝐕∈ℝ^B× S× L× N× D, which need to be stored in the GPU memory during inference. 
To reduce memory and computational costs during inference, efficiency can only be achieved by pruning the dimensions across S,L,N,D or performing quantization over the caches. It is well-acknowledged that token importance tends to be sparse. Consequently, KV eviction algorithms have been proposed to minimize the KV cache memory cost from dimension S <cit.>. Additionally, inter-layer redundancy has been explored <cit.> to address the layer dimension L. Despite these advances, existing methods have largely overlooked the channel dimension D. In this paper, we highlight that the magnitude across Key cache channel dimensions is significantly unbalanced. Additionally, we observe that the attention weights exhibit a low-rank structure. Based on these findings, we hypothesize that the channel dimension of the key cache exhibits redundancy. Consequently, we explore the redundancy of KV cache tensors specifically from dimension D, aiming to devise effective strategies that reduce costs without compromising performance.





In our paper, we propose a simple yet effective method – , for KV cache pruning. To pinpoint the least significant channels, we formulate the task as an optimization problem, aiming to minimize the loss in attention weights attributable to pruning. To effectively address this problem, we establish a novel query-dependent criterion that assesses the importance of each channel. Using this criterion, we then select the most critical channels in a greedy fashion. We evaluate  using the LLaMA3 <cit.> and Mistral <cit.> models and validate its effectiveness across various long-sequence datasets. The results indicate that when paired with token eviction methods, ThinK not only achieves comparable or superior accuracy but also reduces KV cache memory costs by more than 20%.





  
Contributions. This work pioneers the investigation into the sparsity structure of the channels in the KV cache. Specifically, we discovered that the activated key cache is sparse given a specific query. This insight allows us to prune the key cache channels using a query-induced norm. Inspired by this phenomenon, we introduce a plug-and-play method ThinK, the first channel pruning method specifically designed for KV cache. ThinK reduces the dimensionality of the cache channel, leading to linear savings in both memory and computational requirements. As plug-and-play techniques, our proposed ThinK is orthogonal to other KV cache compression schemes (e.g. KV cache eviction, quantization). Our extensive experiments demonstrate the remarkable efficiency of ThinK on LLaMA-3 and Mistral models. Furthermore, we explore the promising potential of extending ThinK to value cache pruning (ThinKV) in specific scenarios, highlighting the broad applicability of our method. 











§ OBSERVATIONS

We identify several key observations that motivate our approach to prune the channel of KV-Cache. 
Specifically, we visualize the magnitude of KV-cache and perform singular value decomposition of the attention in LLaMA3-8B model. 
We show the visualization for certain heads and layers for illustration purposes. 



  
Magnitude of KV-Cache Channel
Figure <ref> visualizes the absolute values of key and value cache across the tokens in each channel[We use the visualization code from <https://github.com/jy-yuan/KIVI/tree/main/vis>.]. 
Consistent with previous findings <cit.>, we observe that only certain channels have significant magnitudes in the key cache, whereas the value cache lacks obvious patterns.
For instance, the absolute value around 50^th channel of key cache has much larger magnitudes than other channels for all tokens in layer 14 (Figure <ref> (a)).
The same observation applies to 50^th and 150^th channels of the first head in layer 20 (Figure <ref> (c)). 
Given such an observation, <cit.> proposed to perform quantization over the channels of the key cache based on such observations.
In addition to channel quantization, the observation also suggests that it is feasible to prune away certain key cache channels that contribute less to the attention mechanism.
Moreover, channel quantization and pruning techniques are orthogonal such that they can be utilized concurrently to enhance model efficiency.










  
Singular Value Analysis
We conducted singular value decomposition (SVD) <cit.> on the attention of the specified layers to investigate their principal components.
The singular values derived from SVD capture the effective rank of the attention matrix, indicating how the information is distributed across different components. 


    𝐔, Σ, 𝐕    = SVD (
           softmax (QK^T/√(D) ) 
         ),  Energy_i  = σ_i^2/∑_i σ_i^2




Figure <ref> (a) illustrates the energy distribution of the singular values, plotted in log scale to enhance the visibility of differences among them. 
Notably, only a few singular values exhibit high energy levels exceeding 0.01 across all heads and layers, highlighting their relative significance.
This observation aligns with previous findings <cit.>, where a small subset of singular values often captures most of the information in attention mechanisms.
In addition, the rapid decay of the energy suggests that a low-rank approximation can effectively capture the essential information in the key cache. 

Figure <ref> (b), the normalized cumulative energy sum reveals that the top 50 singular values account for over 90% of the total energy.
These findings suggest that the attention matrix is inherently low-rank <cit.>, indicating that the key cache can be approximated using low-dimensional vectors <cit.>.





























§ THINK




  
Notations. We use uppercase letters (e.g., X, Y) to denote scalar values and boldface uppercase letters (e.g., 𝐐, 𝐊) to denote matrices and tensors. The notation ‖·‖_p denotes the l_p-norm for vectors. Unless otherwise specified, ‖·‖ denotes the l_2-norm. The Frobenius norm is denoted by ‖·‖_F. The floor function is denoted by ⌊·⌋, and the ceiling function is denoted by ⌈·⌉.




 §.§ Preliminary Study of KV Cache Optimization

In scenarios with extended contexts or batch processing, the main limitations in terms of memory and speed are due to the handling of the KV cache size. Considering a batch of requests to a Large Language Model (LLM) service that provides a long input prompt consisting of tokens [x_B1,...,x_BS], the total KV cache size can be computed as follows:

    2 × B × S × L × N × D,

where L is the number of layers, N is the number of heads, D is the head dimension. The KV cach size grows linearly as the batch size B and sequence length S.  For
a model with multihead attention (MHA) <cit.>, such as LLaMA2-7B <cit.>, a context length of 2048 and a batch size of 13 require storing a 13 GB KV cache, which is equivalent to the size of the model parameters. 
The KV cache must be transferred from off-chip memory (HBM) <cit.> to on-chip memory (cache) for each token generated, leading to a memory bottleneck. 
This substantial memory demand highlights the challenges in managing large-scale models and the need for efficient memory utilization strategies. Current methods optimize the KV cache based on the sequence length S <cit.> and precision <cit.>. 
We will introduce a new method, , to optimize it from the perspective of the number of head dimensions D.





Magnitude base Pruning: Based on the observations in Figure <ref> which depicts the significant variation in the magnitudes across different channels, one straightforward criterion is to use the norm of the magnitude to measure the importance of different channels in key cache.

    M_n,d = ‖ K[n,:,d]‖ _p.

Given pruning ratio λ, We only keep T=⌊(1-λ)D⌋ most important channels among the D channels of each head:

    I = Top_T(M,T)

where ‖·‖ _p is the l_p norm of each channel. n∈[1,N] and d∈[1,D] are indicators of heads and channels in key cache. I∈(ℤ^+)^N× T stores the indicators of the top T values in tensor M per head.

In Table <ref>, we present the results of key cache pruning with various pruning ratios applied to the LLaMA3-8B model. We utilize the l_1 and l_2 norms as criteria for evaluation, and validate performance using the LongBench benchmark <cit.>. Compared to the baseline methods, H2O <cit.> and SnapKV <cit.>, both with a KV length of 512, we further prune the channels of the key cache. A 30% pruning ratio can maintain accuracy; however, increasing it to 40% results in significant performance degradation, especially for l_1 norm based pruning. The results of magnitude-based pruning support our assumption that the key cache is redundant in the channel dimension. These results also indicate the need for a better pruning matrix to achieve higher pruning ratios effectively.




 §.§ Query-Driven Pruning

For each head, the attention scores are computed using the queries and keys, and then applied to the values. The formula for the attention for head i is:

    Attention(Q_i,K_i,V_i) = softmax(Q_iK_i^T/√(D))V_i

where Q_i, K_i,V_i∈ℝ^S× D. When pruning one channels of K_i, the corresponding channels in Q_i will also be removed. We aim to find the optimal subset of channels to prune, denoted by the selection matrix S∈{0,1}^D× D, where 
S is a diagonal matrix with binary entries (1 for keeping a channel, 0 for pruning it). To better maintain the performance after pruning the channels, we minimize the Frobenius norm of the difference between the original and pruned attention weights:

    min_S‖Q_iK^T_i -Q_iS(K_iS)^T ‖_F


Given a pruning ratio λ, it can further expanded as:

    min_𝐒     𝐐_i 𝐊_i^T - 𝐐_i 𝐒𝐊_i^T _F 
    subject to     trace(𝐒) = ⌊(1-λ)D⌋
         𝐒 = diag(s_1, s_2, …, s_D),  where  s_j ∈{0, 1}


For simplicity, we use greedy algorithm to optimize 𝐒. To achieve the pruning goal, we define a criterion for evaluating the importance of each channel and greedily select the channels with largest scores:

    Score_i[j] = ‖Q_i[:,j]K_i[:,j]^T ‖_F


    I_i = Top_T(Score_i,T)

Here's a detailed explanation of why it optimizes the selection matrix. The score_i[j] measures the magnitude of the interaction between the query and key vectors for channel j in each head i. By selecting channels with the highest interaction magnitudes, we aim to retain the most significant contributions to the attention mechanism. This criterion ensures that the selected channels preserve the primary information flow in the attention computation, thereby minimizing the loss of important information.

Observation Window. Based on the observations in SnapKV <cit.> that the last window of input sequence recognizes highly similar attention allocation
pattern with the actual generation. To reduce the computation cost, we only use the last S_obs window to calculate the score: ‖Q_i[-S_obs:,j]K_i[:,j]^T‖_F.


[17]r0.4

    < g r a p h i c s >



Implementation during decoding.







 §.§ Implementations

Following SnapKV, we focus on the long context input scenario. We opt not to prune the most recent tokens and newly generated keys. Consequently, our key-value (KV) cache will store two distinct categories of keys: one subset consists of pruned keys with a reduced channel size, while the other retains keys at their original size. Additionally, we maintain a binary mask to indicate which channels have been pruned. Note that the memory overhead associated with this mask is negligible. Figure <ref> illustrates one implementation of our method during the decoding stage. The pruned keys are first zero-filled using the mask to restore them to their original size, and then concatenated with the unpruned keys. This implementation can be effectively integrated with optimization techniques such as FlashAttention <cit.>. The second implementation involves initially pruning the Query using the mask. The pruned Query is then multiplied by the pruned Key, while the unpruned Query is applied to the unpruned Key. Subsequently, the two outputs are concatenated. In theory, this implementation has lower computational costs.






§ EXPERIMENT RESULTS
 


In this section, we conduct extensive experiments to evaluate the effectiveness of  on performance and memory reduction. First, we introduce benchmark datasets, baseline methods, and implementation details in Section <ref>. Next, we report the performance of our KV pruning method with two different KV compression methods in long-context scenarios in Section <ref>. Finally, we present the results of ablation experiments in Section <ref>.



 §.§ Settings





  
Benchmark Datasets.
We perform evaluations of our method with state-of-the-art KV cache compression methods on two widely used benchmarks: LongBench and Needle-in-a-Haystack. 
LongBench <cit.> is designed to comprehensively evaluate LLM’s long context understanding capabilities. It includes 17 datasets covering six different tasks: single-document QA, multi-document QA, summarization, few-shot learning, synthetic tasks, and code completion. The average input length of LongBench is 6,711 words, which necessitates reducing the KV cache to lower memory usage for inference.
Needle-in-a-Haystack <cit.> is a recent popular test challenge that requires models to accurately identify a small piece of information (“needle”) in a long document (“haystack”), where the needle is placed at a random position. 
This challenge can test if KV cache compression methods still retain the small piece of critical information.










  
Baseline Approaches.
The baseline methods include Heavy Hitter Oracle (H2O) and SnapKV, both of which are the state-of-the-art KV cache compression methods but use different strategies for selecting KV cache. 
H2O <cit.> is designed to reduce memory usage by dynamically balancing recent tokens and Heavy Hitter (H2) tokens, where H2 tokens are a small set of tokens that contribute most of the value when computing attention scores.
SnapKV <cit.> automatically compresses KV caches by selecting clustered important KV positions for each attention head. This is inspired by the fact that each attention head has a specific pattern, focusing on certain attention features during generation.







  
Implementation Details.

In this paper, we use LLaMA3-8B-Instruct <cit.> and Mistral-7B-Instruct-v0.2 <cit.> as the backbone LLMs, both accessible via HuggingFace <cit.>.
Our  aims to prune channels of the KV, which is agnostic to KV cache compression methods. 

If there is no other statement, we prune the key cache by default. All the experiments are conducted using one Nvidia A100.
To fairly compare KV cache compression methods and KV cache compression integrated with , we use the same hyperparameters for both.
For example, when comparing SnapKV and SnapKV integrated with , we set the maximum pooling kernel size to 7 and the observation window size to 32, using the same KV-size for both.








 §.§ Results on LongBench


Tables <ref> and  <ref> present the results of KV compression methods and KV compression methods integrated with our channel pruning method of Key cache () over two different base LLMs across various KV-sizes on LongBench.
The following observations can be drawn:
(1) Our method can further prune the channels of the key cache after compressing the KV cache with H2O and SnapKV. For the base model LLaMA3-8B, our method reduces memory usage and slightly improves performance for both H2O and SnapKV. For the base model Mistral-7B, our method reduces memory with only a slight drop in performance in some cases.
(2) Comparing the results of SnapKV or H2O integrated with  in Table <ref> to SnapKV or H2O integrated with l_1 or l_2 norm shown in Table <ref>, our proposed query-based channel pruning shows better performance when the pruning ratio is 40% compared to no pruning. However, l_1 and l_2 norm methods cannot outperform the KV compression methods without pruning.
(3) When the pruning ratio is small, performance tends to be better than when the pruning ratio is large.
(4) When the KV-size is increased from 128 to 2048, the performance of our channel pruning method improves. Notably, with a KV cache size of 2048 and a pruning ratio of 40%, our method can even outperform the LLaMA3-8B with a full KV cache.
The above observations indicate that our method is agnostic to existing KV cache compression methods and can further improve their performance and memory reduction. Additionally, query-based pruning is more effective than l_1 and l_2 norm for channel pruning in LLMs.



 §.§ Results on Needle-in-a-Haystack







Table <ref> presents the results of the Needle-in-a-Haystack test based on the SnapKV <cit.> approach with KV cache size of input prompt from 128 to 2048. 
With a modest pruning ratio λ=40%, the accuracy of  is consistently better or comparable than the original SnapKV across both LLaMA3 and Mistral with different KV cache sizes. 
Such comparisons demonstrate that the proposed query-driven channel pruning can effectively retain the informative channels and discard the noisy ones. 
With a pruning ratio of λ≥ 50%, we observe an accuracy drop with a small KV-cache size, particularly for 128 and 512 across both LLaMA3 and Mistral. 
However,  is still able to obtain comparable performance with SnapKV when we have a larger KV cache size (i.e., 1024 and 2048).
Intuitively, a larger pruning ratio with a smaller KV cache size could tend to lose more important information compared to a larger KV cache size. 
In addition, the performance on larger KV cache size also suggests  is robust for long context retrieval tasks. 


Figure <ref> (a)-(d) visualize the Needle-in-a-Haystack test accuracy against different token lengths and depth percent.
The KV cache size is configured to 128 and 1024, with pruning ratio λ equal to 40% and 50%, respectively.
First,  can maintain the original retrieval capability of SnapKV though the overall numerical accuracy is different (i.e., 77.8 vs 78.6 and 90.4 vs 90.6). 
Because achieves the same accuracy at most of the entries with different token limits and depths as Snap KV.
In addition,  also successfully retrieves some “needles” where SnapKV fails, leading to higher overall accuracy. 
The visualization results further demonstrate the robustness of  in a fine-grained perspective and  can enhance the original quantization approach.














 §.§ Ablation Studies












  
Impact of Different Recent Sizes.
Preserve the most recent KV embeddings <cit.> is important for maintaining the performance of LLMs after KV cache compression. Note that a tradeoff exists: increasing the recent-size allows more infomation to be propagated, while increasing the cache size to be stored. To study its impacts, we evaluate the performance produced by three recent-size, namely 0, 32 and 128, on LongBench, using Mistral-7B-Instruct-v0.2 as the basleline model. The results are summarized in Table <ref>. One can observe that a recent-size of 32 yields superior performance than 0 in terms of averaged score on LongBench which demonstrates the importance of keeping the most recent KVs. On the other hand, the performance of 32 and 128 is negligible, suggesting that maintaining the most recent 32 KVs suffices to preserve optimal performance.

[17]r0.45

    < g r a p h i c s >



Performance comparisons under the same KV cache memory consumption.





  
Comparison Under the Same Memory Consumption.
To ensure a fair comparison, we adjust the KV cache size of H2O or SnapKV to match the memory usage of H2O with ThinK or SnapKV with ThinK. 
For example, the KV cache size of H2O with ThinK is 128. Due to channel pruning of the key cache, the memory consumption of H2O with ThinK at a 128 KV-size is lower than H2O at a 128 KV-size.
Therefore, the KV-size of H2O is adjusted from 128 to 109. 
Table <ref> and Figure <ref> present the results of these comparisons on LongBench.
As the results show, H2O or SnapKV with ThinK deliver superior performance compared to H2O or SnapKV without ThinK while maintaining the same memory usage. 
This demonstrates that KV cache compression methods, combined with our query-based channel pruning, can compress the KV cache more effectively and utilize memory more efficiently.



  
Pruning Channels of Both Key and Value Cache.
In our previous study, we evaluate methods with pruning channels of the key cache. In this part, we explore the impact of pruning channels of the value cache (Appendix <ref>). Specifically, for KV cache compression methods, we prune channels of the key cache at one pruning ratio and channels of the value cache at another pruning ratio. 
Table <ref> presents the results with LLaMA3-8B and Mistral-7B on LongBench.
When testing on the base model LLaMA3-8B, H2O or SnapKV with key and value channel pruning can perform on par with H2O or SnapKV without channel pruning. In some cases, H2O or SnapKV with key and value channel pruning may even perform better. Pruning the channels of the value cache over the base model Mistral-7B experiences a slight performance drop. Despite this, note that KV cache with additional value channel pruning can further reduce memory usage.

















§ RELATED WORK


In scenarios involving long contexts, the most significant computational burden from the attention mechanism is the key-value (KV) cache. Reducing the KV cache is a high priority for optimizing deployment efficiency. System-level optimizations, such as FlashAttention <cit.> and PagedAttention <cit.>, have been developed to address this issue. Additionally, algorithm-level optimizations are being explored to further enhance efficiency.

KV Cache Eviction. StreamingLLM <cit.> maintains a few initial tokens and some recent tokens based on the observation of attention sink, which may result in the loss of important information carried by the dropped tokens. H2O <cit.> retains only a small portion of the tokens by greedily dropping tokens based on their contributions to the cumulative attention. SnapKV <cit.> selects clustered important KV positions for each attention head from an ‘observation’ window located at the end of the prompts. FastGen <cit.> adaptively evicts tokens from attention heads that emphasize local contexts. This approach focuses on discarding non-special tokens centered on special tokens, while the standard KV cache is used only for attention heads that broadly attend to all tokens. PyramidKV <cit.> and PyramidInfer <cit.> considers adjusting the KV cache size across different layers by allocating more cache in the lower layers and less in the higher ones.




KV Cache Quantization. SmoothQuant <cit.> can quantize the KV cache to 8-bit with minimal performance degradation. Q-Hitter <cit.> uses accumulated attention scores and ‘Quantization Friendliness’ metrics to identify tokens that are essential for maintaining the generalization capabilities of LLMs and are suitable for KV cache quantization. Some studies have found that the key cache and value cache should be quantized differently <cit.>: the key cache should be quantized per-channel, while the value cache should be quantized per-token.

Structured Pruning of LLMs. Structured pruning  <cit.> of LLMs traditionally focuses on removing unimportant layers, heads, and hidden dimensions, which often results in significant performance degradation. In contrast, our methodology preserves the original architecture of the LLM and specifically targets the channel dimension within each head’s key cache. By dynamically identifying unimportant channels based on data dependant criterion, our approach greatly reduce the key cache-size with negligible performance loss.




§ END NOTE AND FUTURE DIRECTION

Motivated by the observation that certain channels have significantly larger magnitudes compared to others, and the singular value analysis indicates that the key cache is inherently low-rank <cit.>, we propose  to perform pruning over the key cache channel. 
The proposed pruning strategy is query-dependant and optimized based on the attention scores, ensuring that essential information is retained for each input query.

In addition,  can be seamlessly integrated with other popular token-level KV cache quantization techniques <cit.>, further enhancing inference efficiency. 
Extensive experiments on LongBench <cit.> and Needle-in-a-Haystack tests with two foundation models demonstrate the effectiveness and robustness of our query-dependent channel pruning method.
Our approach achieves comparable or superior performance to baseline methods while reducing the key cache size by 40%.
Our analysis indicates that  can maintain superior performance over baselines with a smaller KV cache size under equivalent memory consumption conditions.

Future work will focus on strategies to increase the pruning ratio without compromising model performance, and exploring pruning techniques for value caches. 
We also plan to evaluate channel importance from a more sophisticated and compositional perspective, incorporating both token-level and channel-level information to calculate importance.



apalike







§ ADDITIONAL EXPERIMENTAL DETAILS



 §.§ Value Cache Pruning

Similar to the approach used for the Key cache, the pruning of channels in the Value cache can be guided by two primary criteria: magnitude-based pruning and query-driven pruning. We find that query-driven pruning is still better than magnitude based pruning.

    Score_v,i(Q_i,K_i,V_i)[j] = ‖softmax(Q_i[-S_obs:]K_i^T/√(D))V_i[:,j]‖_F


    I_i = Top_T(Score_v,i,T)

where Q_i, K_i,V_i∈ℝ^S× D. We define a criterion Score_v,i to indicate the importance of each channel in the head i of value cache. Then, only top T channels are retained.






 Paper Title: A Vectorization Method Induced By Maximal Margin Classification For Persistent Diagrams.







A Vectorization Method Induced By Maximal Margin Classification For Persistent Diagrams



* Corresponding author

    
An Wu *
College of Mathematics and Computer Science 

Zhejiang A&F University

Hangzhou, China 

20220075@zafu.edu.cn
Yu Pan
School of Mathematics and Statistics 

Beijing Institute of Technology

Beijing, China 

p4nyu@foxmail.com
Fuqi Zhou
School of Electronics and Information 

Hangzhou Dianzi University

Hangzhou, China 

23040444@hdu.edu.cn
Jinghui Yan
Institute of Automation 

Chinese Academy of Sciences

Beijing, China 

jinghui.yan@nlpr.ia.ac.cn
Chuanlu Liu
School of Computer Science & Technology 

Beijing Institute of Technology

Beijing, China 

chuanluliu@163.com

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



§ ABSTRACT 
Persistent homology is an effective method for extracting topological information, represented as persistent diagrams, of spatial structure data. Hence it is well-suited for the study of protein structures. Attempts to incorporate Persistent homology in machine learning methods of protein function prediction have resulted in several techniques for vectorizing persistent diagrams. However, current vectorization methods are excessively artificial and cannot ensure the effective utilization of information or the rationality of the methods. To address this problem, we propose a more geometrical vectorization method of persistent diagrams based on maximal margin classification for Banach space, and additionaly propose a framework that utilizes topological data analysis to identify proteins with specific functions. We evaluated our vectorization method using a binary classification task on proteins and compared it with the statistical methods that exhibit the best performance among thirteen commonly used vectorization methods. The experimental results indicate that our approach surpasses the statistical methods in both robustness and precision.



Maximal margin classification, Topological data analysis, Persistent diagram, Vectorization, Protein function




§ INTRODUCTION

Topological data analysis (TDA) is a computational technique which has been developed to study the topological structure of data over the past 20 years<cit.>. The theoretical foundation of TDA is the algebraic topology, and combined with computer technology TDA has been applied to various practical problems, such as image processing, material structure analysis and quantum computing to name a few as in <cit.> and <cit.>.



The primary methods in Topological Data Analysis are persistent homology and Mapper, and we focus on the persistent homology approach in this paper. Based on the persistent homology there are many interesting works. The research team led by Rongling Wu implemented GLMY homology theory proposed by Shing-Tung Yau to analyze and interpret the topological change of health state from symbiosis to dysbiosis, and they successfully described the topological difference and commonality between Crohn's disease and ulcerative colitis<cit.>. Perea introduced how to combine dynamical systems and topological data analysis to gain insights from time-varying data<cit.>. For more applications of persistent homology, one can refer to the review paper of Chazal and Michel<cit.>. Previous works have primarily focused on classification problems as an application of persistent homology. This is because the number of bars in the barcode (Betti numbers) offers a direct measure for classification, or because researchers can define similar functions (such as bottleneck distance and p-Wasserstein distance) based on the persistent diagrams. However, if we want to combine persistent diagrams with more advanced artificial intelligence technologies, there is a critical problem that most machine learning algorithms assume that the input data is from ℝ^n or, more generally, some Hilbert space, but the space of persistent diagrams is not a vector space even if they are from the similar samples. Hence we want to find a  reasonable embedding of persistent diagrams which can be a pre-encoder for any other topological artificial intelligence frameworks.

Many researchers have tried to solve this problem which is called vectorization. Barnes, Polanco and Perea provided a comparative study of six such methods and they applied and compared popular machine learning methods on five data sets<cit.>. Chung and Lawson developed a framework of vectorization and showed several well-known summaries fall in this framework<cit.>. In <cit.> and <cit.>, stable vector representations were developed and these methods were tested on several data sets to show their effectiveness and efficiency. Ali et al. <cit.> summarized thirteen vectorization methods, and they examined these methods on three well-known classification tasks. They discovered that the best-performing method is a simple vectorization by a few statistics. Besides these methods, there are numerous approaches of vectorization as in<cit.> and <cit.>, and more information can be found in Section 2.3 of <cit.>. If we further consider the reasons behind Ali and his colleagues’ results, a plausible explanation is that overly artificial vectorization methods may lead to the loss of key information, whereas statistical methods are relatively more simple and straightforward. However, statistical methods are also hard to interpret. For example, it is difficult to answer why someone need to take the 10th, 25th, 75th and 90th percentiles of the births and deaths rather than others. Hence we approve a more explainable method which is not included in the methods mentioned above, and our method achieve a better result than the statistical method <cit.> which achieves the best performance on several datasets.

As previously stated, one can define a distance function on persistent diagrams to construct a metric space. Consequently, a direct approach is performing classification tasks within this metric space. But a primary challenge mentioned in <cit.> of this method is that 
"The space of persistent diagrams is infinite-dimensional and highly nonlinear",
which poses significant computational and theoretical complexities. However, in practice, since the size of data is always finite, the discrete points in the persistent diagram are bounded. Therefore, we can always construct the  compactification of the space of persistent diagrams generated by the sample set. Finally, by the classical functional analysis theory <cit.>, we have

A compact metric space can always be embedded into a Banach space.

Then we can solve the challenge above on a Banach space.

In the current paper, we focus on the space (𝒳_pd, d) generated by finite persistent diagrams endowed with a fixed metric. Then inspired by the theory of maximal margin classification for Banach spaces <cit.>, we embed (𝒳_pd, d) into the Banach space 𝒞_b(𝒳_pd) by Kuratowski's embedding, then we consider the maximal margin classification for space 𝒞_b(𝒳_pd). As a result, we transform the maximum margin problem into a classical quadratic programming problem as Eq. (<ref>). Moreover, in the derivation process, we find the following result:

Given a set of persistent diagram 𝒳_pd endowed with metric d, we define the vectorization 𝒳_pd→ℝ^n by x↦ (d(x,x_1),⋯, d(x,x_n)), where x_i∈𝒳_pd and n is the cardinal number of 𝒳_pd. In addition, this simple vectorization method corresponds to the quadratic programming problem Eq. (<ref>). 
 

Next, we apply the maximal margin classification for persistent diagrams on the data of Cas-associated proteins, which are closely related to gene editing technology as in<cit.> and <cit.>, and transposases. We choose a simple distance defined by Eq. (<ref>) which requires less computation than the bottleneck distance, and we compare our method with statistical methods which are the best-performing methods in the thirteen vectorization methods <cit.> mentioned before. As a consequent, our method achieves superior performance both in robust and accuracy. 

Finally, we present a framework that uses the TDA processing methodology introduced in this paper to predict proteins with specific functions, and the total framework is exhibited in Fig. <ref>. We will discuss more details in Section III.

The remainder of our paper is organized as follows. In section II, we revisit the maximal margin classification for Banach spaces and apply it to the space of persistent diagrams. Finally, the problem is reduced to a solvable quadratic programming problem. In section III, we show our experiments on the data of Cas-associated proteins and  transposases. We analysis our experiment results and propose a framework of predicting proteins' function. In section IV, we conclusion our work and do some discussions on the current and future researches.



§ MAXIMAL MARGIN CLASSIFICATION METHOD

In this section, we will introduce the main idea of TDA first. Then we will review the principle of the maximal margin classification for Banach spaces, and one can check the details in the publication of Hein, Bousquet and Schölkopf<cit.>. In particular, we will regard the space of persistent diagrams as a Banach space, which allows us to analyze the space of persistent diagrams by maximal margin classification. 



 §.§ The topological data analysis

The main idea of TDA is to find some topological features to describe the difference of data. As an example, assume we have a deflated balloon. When we inflate it, it remains a balloon, because this process is a continuous transformation. However, if we blow in enough gas to cause an explosion, it ceases to be a balloon and becomes a ”burst balloon”. Hence the topological feature that distinguishes burst balloons from intact balloons is the presence or absence of holes in the balloons. By popularizing this idea, researchers have developed a major technique in TDA called persistent homology<cit.>. The word "persistent" means that there is a continuous variation of a parameter which is call the filtration. If we can construct a series of complexes suitable with the continuous variation, then the homological information of the complexes illustrate the topological structure of data. For example, Fig. <ref> shows the process of increasing the radius of points, and it is easy to find that the topological structure of top left three points is different to that of bottom right two points, as the former would create a hole.


Based on this idea, persistent homology is good at dealing with classification tasks, and the topological features are usually taken as the connect components, holes, and voids which are denoted by H_i, i=1,2,3 respectively. Typically, there are two visualization methods for this information, known as the persistent diagram and the barcode, which are essentially equivalent to each other. In Fig. <ref>, we show both visualizations of topological information extracted from Fig. <ref>. 


The construction of complexes in persistent homology is of importance. In this section, we use the Čech complex to illustrate the geometric explanation. In practice, the Vietoris-Rips complex is often employed, as it is beneficial for computational purposes.


However, the number of points contained within different persistent diagrams may indeed differ. Hence, we need a vectorization method to transfer persistent diagrams into a vector space, and we will illustrate our vectorization method next. 




 §.§ The maximal margin classification

Considering the space of persistent diagrams 𝒳_pd, we can define a distance function d(·, ·) on 𝒳_pd to make it a metric space (𝒳_pd,d). In practical problems, the number of generated persistent diagrams is always finite, then we can assume that d(X, Y) is bounded for any sample X, Y∈𝒳_pd, and we use the same symbol 𝒳_pd to denote the subspace generated by sample set. Hence we obtain a compact metric space (𝒳_pd, d), which can be embedded into a Banach space isometrically by the observation of Kuratowski as in<cit.>.

Let 𝒞_b(𝒳_pd) be the Banach space of continuous and bounded functions on 𝒳_pd endowed with infinite norm ._∞. Take an arbitrary point x_0 ∈𝒳_pd and define a map from 𝒳_pd to real-valued functions on 𝒳_pd  

    ϕ : x↦ϕ_x:=d(x, ·)-d(x_0,·).

Then consider the space D :=span{ϕ_x:x∈𝒳_pd}, in which the closure is taken in (𝒞_b(𝒳_pd), ._∞), we have the following lemma: 

  ϕ is a total isometric embedding from (𝒳_pd,d) into the Bananch space (D, ._∞)⊂ (𝒞_b(𝒳_pd), ._∞).


Through Lemma <ref>, we can transform the data processing problem on space (𝒳_pd, d) to that on space (D, ._∞).  In the current paper, we focus on the maximal margin classification of persistent diagrams. Recall that, the margin formulation on a Banach space is

    min_W'∈ D', b∈ℝW'
        s.t.   y_j(⟨ W', ϕ_x_j⟩_D',D+b)≥ 1, ∀ j=1,⋯, n,

where D' is the dual space of D, y_j is the label corresponding to ϕ_x_j and the product ⟨ W', ϕ_x_j⟩_D', D is naturally defined by W'(ϕ_x_j).

Since 𝒳_pd is compact, 𝒞_b(𝒳_pd) is isometrically isomorphic to the Banach space ℳ(𝒳_pd) of finite signed Borel measures on 𝒳_pd with the measure norm. However, we did not construct the parametrization of ℳ(𝒳_pd), hence we need to choose a reasonable approximation on dual space 𝒞_b(𝒳_pd)' instead.

Define a space E:=span{δ_z: z∈𝒳_pd}, where δ_z(f):=f(z) for any f∈𝒞_b(𝒳_pd). There is the following lemma:

  The space E is weak*-dense in 𝒞_b(𝒳_pd)' endowed with the norm 
    ∑_i=1^nβ_iδ_z_i_𝒞_b(𝒳_pd)'=∑_i=1^n|β_i|.


Because of the Lemma <ref>, we can use e=∑_i=1^nβ_iδ_z_i∈ E to approximate any elements of 𝒞_b(𝒳_pd)'. Substitute e into W' in Eq. (<ref>), the optimization problem can be rewritten as 

    min_n∈ℕ, z_i∈𝒳_pd, b∑_i=1^n |β_i|
        s.t.   y_j(∑_i=1^nβ_i(d(x_j, z_i)-d(x_0, z_i))+b)≥ 1.

In consideration of computational convenience, practically it is necessary to restrict E to the training set, which means n is the size of training set. Moreover, it can be noted that the term -∑_i=1^nβ_id(x_0, z_i)+b is independent of j, so that we can reparameterize it as c. In addition, the vector β=(β_1, ⋯, β_n) is optimized in a n dimensional vector space, where dimension n is always finite. Hence the 1-norm in the optimization function is equivalent to 2-norm, which is more advantageous in solving our problem. In light of these observation, we convert the optimization problem as Eq. (<ref>) into solving the following optimization problem:   

    min_β_i, c∑_i=1^n β_i^2
        s.t.   y_j(∑_i=1^nβ_id(x_j, z_i)+c)≥ 1, ∀ j=1,⋯, n.

Observe the form of the optimization problem above, one can find it exhibits a form identical to that of the support vector machine defined in Euclidean space. 

Now we will subsequently explain this phenomenon from the perspective of data science. It is obvious that the primary difference lies in the replacement of term ∑_i=1^nβ_ix_j in SVM with term ∑_i=1^nβ_id(x_j, z_i) in the conditions. Because in SVM the data x_j is always a vector, but now x_j is just a point in the sample space which may no longer be a vector space, therefore we can not do linear combination of samples here directly. However, for each pair of sample points (x_j, z_i) we can define the distance (or similarity) function between them to make the sample space be a metric space. Then for each sample we can generate a fixed-length vector by computing distances from this sample to all points in training set. Essentially, we transform the metric space into a linear space through the distance function, and this transformation method is theoretically guaranteed to be non-arbitrary in mathematics.

Same as SVM, we can also introduce slack variables to obtain a new version with a soft margin as follows:

    min_β_i, ξ_i, c   ∑_i=1^n β_i^2+a∑_i=1^nξ_i
    
    s.t.      y_j(∑_i=1^nβ_id(x_j, z_i)+c)≥ 1-ξ_j, 
       ξ_j≥ 0, ∀ j=1,⋯, n,

where ξ_j is the slack variable which is defined by hinge loss 

    ξ_j=max(0, 1-y_j(∑_i=1^nβ_id(x_j, z_i)+c)).

In the current paper, we focus on the soft margin version of maximal margin classification, and solving Eq. (<ref>) is a typical quadratic programming problem. Here, we render it in matrix notation for ease of subsequent resolution.

Define the coefficient vector β:=(β_1, ⋯, β_n)^T, the slack vector ξ:=(ξ_1,⋯,ξ_n)^T, label vector y:=(y_1, ⋯, y_n)^T and denote the distance matrix of training set by d. Then the optimization function is represented by

    min_β,ξ, c1/2β^T2I_nβ+a 1_n^Tξ,

where I_n is identity matrix with n order and 1_n^T is a n-dimensional vector with elements are all equal to 1. Furthermore, combine β, ξ, and c together, we obtain

    min_α1/2α^TQα+b^Tα,

where α, Q, b are defined as follows:

    α := ([ β; ξ; c ]),  b :=([    0; a1_n;    0 ]), 
        Q := ([ 2I_n    0    0;    0    0    0;    0    0    0 ]).

If we denote the Hadamard product of matrices by the symbol *.
Then similarly, the matrix representation of constraint conditions is

    Gα≥ h,

where the matrix G and the vector h are defined by

    G:=([ y*d I_n 1_n;   0 I_n   0 ]),   
    h:=([ 1_n;   0 ]).


In a conclusion, we transform the problem of maximal margin classification on the space of persistent diagrams into solving the quadratic programming problem

    min_α1/2α^TQα+b^Tα,
        s.t.     Gα≥ h,

which can be solved directly.



§ EXPERIMENTS AND ANALYSIS

In this section, we will illustrate our experiments of maximal margin classification of persistent diagrams. We will apply our Banach space based method, which will be referred as BS method in the following context, to the classification of proteins. Then we will compare the experiment results of BS method with two widely used methods which are based on dimension reduction techniques. Finally, I will analyze the significance of BS method and its applications in the exploration of novel proteins.



 §.§ Experiment Setup

Our experimental setup employs a server equipped with an Intel® Xeon® CPU E7-4870 operating at a frequency of 2.40 GHz, and we use the Python library ripser <cit.> to compute persistent homology. For solving the quadratic programming problem, we use the quadratic programming solver CVXOPT.




 §.§ Dataset

The original dataset we deal with consists of 161 protein PDB files of Cas-associated proteins and 197 files of transposases, and our goal is to classify these proteins correctly. Each PDB file includes the coordinates of atoms in the proteins, and it is easy to extract these coordinate information from PDB files. After extracting the atomic coordinates, the traditional approach of TDA is computing the persistent homology of the atomic point clouds directly. However, this methodology encounters computational challenges, particularly in the calculation of the second persistent homology group (H_2), which is often exceedingly difficult and necessitates substantial computational expenditures. The main reason is that compared to general chemical molecules and structural materials, proteins are typically composed of tens of thousands of atoms. But when calculating homology information of H_2, it is necessary to traverse all combinations of four atomic interactions. Therefore, in this paper, we modify the persistent homology method at the beginning.



 §.§ Preprocessing

We extract the topological information of protein atomic point clouds according to the following approach. Firstly, we derive the coordinates of 358 proteins from the input files. Subsequently, we construct the residue connection matrix of each protein based on the relative positions of the Cα atoms. For more details, we consider the Cα atoms within a protein as nodes, and an edge is created between two Cα atoms if the inter-atomic distance is below this established threshold of 5 angstroms (Å). Following the construction of  residue connection matrices, we can use graph embedding methods to generate point clouds which include less points than the original atomic point clouds. There are many graph embedding methods, and here we choose Node2vec <cit.> method to generate new point clouds. After the embedding of residue connection matrices, we calculate the persistent homology of point clouds including H_0, H_1 and H_2. With the topological features captured, the final step is preprocessing of these topological information. We set a cutoff of 0.01 to filter out H_1 and H_2 information whose persistent times are shorter than this value. We summarize the topological feature extraction algorithm as Algorithm 1.



By using the residue connection matrix and graph embedding, we reduce the number of points in the point clouds. Hence, it is foreseeable that the aforementioned approach to extracting topological features will have a substantial time-saving advantage. In addition, because we only use the inter-atomic distances of the Cα atoms, some information is inevitably omitted, which is acceptable in our experiment.

The results of preprocess is that the computation time required for calculating H_1, H_2, and H_3 using the modified persistent homology approach is comparable to the time taken for calculating H_1 and H_2 using the original persistent homology method, with both processes approximately lasting around 10 days. Roughly speaking, this implies that the new persistent homology algorithm omits certain two-dimensional topological information, and the computational resources saved are redirected towards the calculation of three-dimensional topological information. Although not the focus of this paper, it is worth mentioning that this strategy has a distinct advantage in that it introduces three-dimensional information that was previously unavailable, which will finally provide an additional weight parameter for learning in other machine learning processes.



 §.§ Baselines and Main Method

After the modified persistent homology, we will get persistent diagrams with three components H_1, H_2 and H_3 which are essentially point sets. Next, we will introduce three widely used topological features analysis methods which are based on different statistics as baselines for our BS method as follows:


  * Statistical Method I<cit.>: We count the cardinality of each point set for persistent diagram PD_i, which is equivalent to the persistent Betti number, directly to construct a 3-dimensional feature vector v_i. And the first element is normalized by multiplying with a factor 0.01. 

  * Statistical Method II: 
    We compute the statistical information (the maximum, the minimum, the variance, the mean value and the median) for 5 indicators, corresponding to the birth or death times of H_1, H_2 and H_3, to generate the feature vector of each persistent diagram. Then we extract a 25-dimensional feature vector for each persistent diagram. 

  * Statistical Method III: We apply the feature vectors proposed in the Definition 3.1 of the survey paper <cit.>. This statistical vectorization method incorporates more statistical information than the two methods mentioned above, and it is the best-performing method of the thirteen alternative methods in three well-known classification tasks.


In comparison with the statistical methods, the BS method needs a distance function d to make the space of persistent diagrams 𝒳_pd be a metric space. In this paper, we define the distance function as follows:

    d(X,Y):=∑_i=1^3 w_id_i(X, Y),  ∀ X,Y∈𝒳_pd,

where d_i(X,Y) 

is the maximum Euclidean distance between H_i-components 
of persistent diagrams X and Y, and W=(w_1,w_2,w_3) is weight vector which is fixed as (1/3, 1/3, 1/3) in our experiment.  

In order to compare the effects of different vectorization methods, we apply all methods on the dataset by soft-SVM to obtain classification results.



 §.§ Results and Analysis

We perform five-fold cross-validation experiments and continuously increase the proportion of training data for all methods mentioned before. The experimental results are presented in the following Table I.



The accuracies of the BS method and Statistical methods are noted to increase with the expansion of the training set size. While the Statistical method I has relatively small change around 84%-85%, which means this method exhibits robustness. As a comparison, the Statistical method II and III are more sensitive to the size of training date, and these methods have lower accuracies in the case with small training set. It is worth mentioning that Statistical methods require manual feature selection, and this example demonstrates that different feature selection approaches have a significant impact on the accuracy. Therefore, it needs a prior understanding of the original data which will impact the feature selection. 

However, our BS method uses the information of all features and has a fixed processing procedure, which reduces errors that may arise from poor feature selection. Moreover, the BS method also has robustness, and when the training date size reaches eighty percent, BS method becomes the best method in accuracy. More importantly, the BS method is more generalizable. Because for any persistent diagrams with more topological information (H_4, H_5,⋯), we can also define a distance function like Eq. (<ref>) and we can apply the neural network to determine the weight vector W in practice. 



 §.§ Case Study

In the end, we will illustrate the classification of proteins in the view of biology and propose a framework of finding new proteins with certain functions by topological data analysis. 


For giving a biological explanation, we take 80% of the data to be the training set randomly and use the remaining data to test as well as to build an evolutionary tree. In Fig. <ref>, we show the evolutionary tree and in experiment result we find that 7 proteins have been misclassified. These proteins are 'z7m9a', 'c5vgb', 'c6eyy', 'c6jhv', 'c6jhw', 'c7msl' and 'c7tun', in which we use the characters “c” and “z” to denote Cas-associated proteins and transposases respectively.

It is obvious that, Fig. <ref> can be primarily divided into two components: the upper component mainly consists of Cas-associated proteins, while the lower component mainly consists of transposases. And we find that classification errors can be divided into two parts: 


  * Type I: The first involves the occurrence of another class of proteins within a particular category (such as 'c5vgb', 'c6eyy', 'c6jhv', 'c6hjw' and 'c7msl').

  * Type II: The second involves the presence of errors in the crossover region in which both classes of proteins appear alternatively    (such as 'z7m9a' and 'c7utn'). 


However, if we focus on the instances of misclassification, we find that the reason for Type I error is that the sample classification can be further refined, whereas the reason for Type II error is that the sample classification is overly granular. In details, Cas-associated proteins can be further subdivided into Cas proteins and anti-Cas proteins. Particularly, we find that proteins 'c5vgb', 'c6eyy', 'c6jhv', 'c6jhw' and 'c7msl' are anti-Cas proteins as in<cit.> and <cit.> which are distinct from the majority of other Cas-associated proteins, hence they are misclassified. On the other hand, proteins 'z7m9a' and 'c7utn' are Cas proteins as well as transposases as in <cit.> and <cit.>, and the fact suggests that this dual nature has resulted in their misclassification. Therefore,  our classification methodology aligns with the established principles and demonstrates superior performance.



 §.§ Framework for Protein Function Prediction

Moreover, Type II error approves a new method for discovering novel proteins with specific functionalities. For example, if a transposase in the crossover region is misclassified, we can suspect reasonably that it exhibits structural and functional features similar to those of Cas proteins. Finally, we summarize this framework for exploring protein functions by their spatial structures in Fig. <ref>:





  * In the first step, we input the PDB files {f_j} of proteins with a known specific function and {g_i} of validation proteins;

  * In the second step, we calculate the topological information for every PDB file from its spatial structure;

  * In the third step, we conduct the binary classification task described in Section II to differentiate between the sets {f_j} and {g_i} using our vectorization method of persistent diagrams;

  * Finally, we confirm those proteins with the second type of error and identify the subset {g_k} of proteins that exhibit this function.



 



§ CONCLUSION AND DISCUSSION



 §.§ Conclusion

We propose a novel vectorization method, which not only attains the best performance in processing topological information compared with conventional methods but also exhibits robustness against the size of the training dataset, to analysis persistent diagrams. We derive the process for handling persistent diagrams using the maximal margin classification method of a Banach space, and discover that once the metric space of persistent diagrams is embedded by Kuratowski's embedding then the final result is equivalent to an encoder that calculates the distances between a given persistent diagram and all other diagrams. Hence our result proposes a more reasonable feature selection which use more topological information extracted than statistical methods that are the best method between thirteen alternative
methods<cit.>. 

Another result with more profound implications for biology is that we propose a framework 
of using spatial topological information of proteins to identify their functions. Because we find that in our binary classification experiment, the second type of error typically arises from proteins that belong to both classes. Therefore, to verify whether certain proteins possess a specific function, we need only focus on whether they would result in a Type II error when classified alongside proteins which are known to exhibit this function.

As a byproduct of our processing, we have optimized the persistent homology algorithm for protein analysis, therefore its computational efficiency is enhanced and it is permitted to calculate the second Betti number.



 §.§ Discussion

In future work, we plan to investigate that whether the Banach space based method can be improved by following modifications:

  
  * First, we intend to examine the performance of our method across various distances;
  
  * Second, we plan to combine our method with neural network methods;
  
  * Finally, we aim to propose additional vectorization methods for persistent diagrams, informed by geometric perspectives


In the derivation process of the BS method, we used information from the distance function. 
Compared to traditional methods of comparing persistence diagrams, such as the bottleneck distance, 
the theoretical computational complexity of our method is lower 
because it does not require enumerating all possible matchings. 
However, we may get a better performance distance by training the metric matrix with the neural network.


Our method can be integrated with the neural network in more aspects, because our method maps proteins into a vector space of relative topological information, 
and this vector representation 
offers greater interpretability. In our subsequent work, we will combine the vectorization approach presented in this paper with neural networks to address further problems related to the spatial structure of proteins. 








Finally, since the vectorization method proposed in this paper originates from a geometric embedding, we aim to consider more vectorization methods inspired by other geometric structures such as Lie groups and homogeneous spaces. 
Particularly, we can choose different embedding methods than Kuratowski's embedding to construct different reasonable encoders for persistent diagrams. 



§ ACKNOWLEDGEMENTS

The authors  extend many thanks to Professor Chao Qian and Hanning Wang who provide academical support and encouragement. 

1
AAJ23
D. Ali, A. Asaad, M. J. Jimenez, V. Nanda, E. Paluzo-Hidalgo, and M. Soriano-Trigueros, “A survey of vectorization methods in topological data analysis," IEEE Trans. Pattern Anal. Mach. Intell., vol. 45, no. 12, pp. 14069–14080, 2023.

BPP21
D. Barnes, L. Polanco, and J. A. Perea, “A Comparative Study of Machine Learning Methods for Persistence Diagrams," Front. Artif. Intell., vol. 4, pp. 681174, 2021.

B21
U. Bauer, “Ripser: efficient computation of Vietoris–Rips persistence barcodes," J Appl. and Comput. Topology, vol. 5, pp. 391–423, 2021.

C09
G. Carlsson, “Topology and data," Bull. Amer. Math. Soc., vol. 46, no. 2, pp. 255-308, 2009.

C20
G. Carlsson, “Topological methods for data modelling," Nat. Rev. Phys., vol. 2, pp. 697-708, 2020.

CM21
F. Chazal and B. Michel, “An Introduction to Topological Data Analysis: Fundamental and Practical Aspects for Data Scientists," Front. Artif. Intell., vol. 4, pp. 667963, 2021.

CLW23
D. Chen, J. Liu, J. Wu, G.-W. W, F. Pan, and S.-T. Yau, “Path Topology in Molecular and Materials Sciences," J. Phys. Chem. Lett., vol. 14, no. 4, pp. 954-964, 2023.

CNO20
I. Chevyrev, V. Nanda, and H. Oberhauser, “Persistence paths and signature features in topological data analysis," IEEE Trans. Pattern Anal. Mach. Intell.,  vol. 42, no. 1, pp. 192–202, 2020.

CL22
Y.-M. Chung and A. Lawson, “Persistence Curves: A canonical framework for summarizing persistence diagrams," Adv. Comput. Math., vol. 48, no. 6, pp. 1-42, 2022.

CRC13
L. Cong, F. A. Ran, D. Cox, S.-L. Lin, R. Barretto, N. Habib et al., “Multiplex Genome Engineering Using CRISPR/Cas Systems," Science, vol. 339, no. 6121, pp. 819-823, 2013.

DHZ20
Z.-T. Dong, C.-F. Hu, C. Zhou, and H.-W. Lin, “Vectorization of persistence barcode with applications in pattern classification of porous structures," Comput. Graph., vol. 90, pp. 182-192, 2020.

DLZ24
Z.-T. Dong, H.-W. Lin, C. Zhou, B. Zhang, and G.-C. Li, “Persistence B‑spline grids: stable vector representation
of persistence diagrams based on data fitting," Machine Learning, vol. 113, pp. 1373-1420, 2024.

ELZ02
H. Edelsbrunner, D. Letscher, and A. Zomorodian, “Topological Persistence and Simplification," Discrete Comput. Geom., vol. 28, pp. 511–533, 2002.

GL16
A. Grover and J. Leskovec, “node2vec: Scalable Feature Learning for Networks," KDD '16: Proceedings of the 22nd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, 2016.

HDM17
L. B. Harrington, K. W. Doxzen, E. Ma, J.-J. Liu, G. J. Knott, A. Edraki et al., “A Broad-Spectrum Inhibitor of CRISPR-Cas9," Cell, vol. 170, no. 6, pp. 1224-1233.e15, 2017.

HBS05
M. Hein, O. Bousquet, and B. Schölkopf, “Maximal margin classification for metric spaces," J. Comput. Syst. Sci., vol. 71, no. 3, pp. 333-359, 2005.

HNH16
Y. Hiraoka, T. Nakamura, A. Hirata, E. G. Escolar, K. Matsue, and Y. Nishiura, “Hierarchical structures of amorphous solids characterized by persistent homology," Proc. Natl. Acad. Sci. USA., vol. 113, no. 26, pp. 7035-40, 2016.

HPG22
S. Hwang, C.-X. Pan, B. Garcia, A. R. Davidson, T. F. Moraes, and K. L. Maxwell, “Structural and Mechanistic Insight into CRISPR-Cas9 Inhibition by Anti-CRISPR Protein AcrIIC4Hpa," J. Mol. Biol., vol. 434, no. 5, pp. 167420, 2022.

HRA18
A. P. Hynes, G. M. Rousseau, D. Agudelo, A. Goulet, B. Amigues, J. Loehr et al., “Widespread anti-CRISPR proteins in virulent bacteriophages inhibit a range of Cas9 proteins," Nat. Commun., vol. 9, pp. 2919, 2018.

JCF12
M. Jinek, K. Chylinski, I. Fonfara, M. Hauer, J. A. Doudna, and E. Charpentier, “A Programmable Dual-RNA–Guided DNA Endonuclease in Adaptive Bacterial Immunity," Science, vol. 337, no. 6096, pp. 816-821, 2012.

JJ19
P. Joharinad and J. Jost, “Topology and curvature of metric spaces," Adv. Math., vol. 356, pp. 106813, 2019.

K18
S. Kališnik, “Tropical coordinates on the space of persistence Barcodes," Found. Comput. Math., vol. 19, no. 1, pp. 101–129, 2018.

KLY19
Y. Kim, S. J. Lee, H. J. Yoon, N.-K. Kim, B.-J Lee, and J.-Y. Suh, “Anti-CRISPR AcrIIC3 discriminates between Cas9 orthologs via targeting the variable surface of the HNH nuclease domain," FEBS. J., vol. 286, no. 2, pp. 4661-4674, 2019.

LGZ16
S. Lloyd, S. Garnerone, and P. Zanardi, “Quantum algorithms for topological and geometric analysis of data," Nat. Commun., vol. 7, pp. 10138, 2016.

PTM21
J. Park, A. W. Tsai, E. Mehrotra, M. T. Petassi, S.-C. Hsien, A. Ke et al., “Structural basis for target site selection in RNA-guided DNA transposition systems," Science, vol. 373, pp. 768-774, 2021.

PC14
J. A. Perea and G. Carlsson, “A Klein bottle based dictionary for texture representation," Int. J. Comput. Vis., vol. 107, no. 1, pp. 75-97, 2014.

P19
J. A. Perea, “Topological time series analysis," Notice of AMS., vol. 66, no. 5, pp. 686-694, 2019.

SHK22
G. Schuler, C. Hu, and A. Ke, “Structural basis for RNA-guided DNA cleavage by IscB-ωRNA and mechanistic comparison with Cas9," Science, vol. 376, pp. 1476-1481, 2022.

TSB18
C. Tralie, N. Saul, and R. Bar-On, “Ripser.py: A Lean Persistent Homology Library for Python," J. Open Source Softw., vol. 3, pp. 925, 2018.

WLD23
S. Wu, X. Liu, A. Dong, C. Gragnoli, C. Griffin, J. Wu et al. “The metabolomic physics of complex diseases," Proc. Natl. Acad. Sci. USA., vol. 120, no. 42, pp. e2308496120, 2023.

ZKN24
A. Zia, A. Khamis, J. Nichols, U. B. Tayab, Z. Hayder, V. Rolland et al., “Topological deep learning: a review of an emerging paradigm," Artif. Intell. Rev., vol. 57, no. 4, 2024.

ZLJ20
B. Zieliński, M. Lipiński, M. Juda, M. Zeppelzauer, and P. Dłotko, “Persistence codebooks for topological
data analysis," Artif. Intell. Rev., vol. 54, no. 3, pp. 1969–2009, 2020.










