Paper Title: Small Object Few-shot Segmentation for Vision-based Industrial Inspection.





Enhancement of the molecular electron chirality by electronic excitation

    Masato Senami
    August 5, 2024
==========================================================================




§ ABSTRACT 
Vision-based industrial inspection (VII) aims to locate defects quickly and accurately. Supervised learning under a close-set setting and industrial anomaly detection, as two common paradigms in VII, face different problems in practical applications. The former is that various and sufficient defects are difficult to obtain, while the latter is that specific defects cannot be located. To solve these problems, in this paper, we focus on the few-shot semantic segmentation (FSS) method, which can locate unseen defects conditioned on a few annotations without retraining. Compared to common objects in natural images, the defects in VII are small. This brings two problems to current FSS methods: 172 distortion of target semantics and 173 many false positives for backgrounds. To alleviate these problems, we propose a small object few-shot segmentation (SOFS) model. The key idea for alleviating 172 is to avoid the resizing of the original image and correctly indicate the intensity of target semantics. SOFS achieves this idea via the non-resizing procedure and the prototype intensity downsampling of support annotations. To alleviate 173, we design an abnormal prior map in SOFS to guide the model to reduce false positives and propose a mixed normal Dice loss to preferentially prevent the model from predicting false positives. SOFS can achieve FSS and few-shot anomaly detection determined by support masks. Diverse experiments substantiate the superior performance of SOFS. Code is available at https://github.com/zhangzilongc/SOFShttps://github.com/zhangzilongc/SOFS.




§ INTRODUCTION



Vision-based industrial inspection (VII) automatically identifies defects in industrial processes and ensures product quality. Compared to object recognition in natural images, defect recognition in VII is fine-grained. As shown in Fig.<ref>, fine-grained defects appear as small objects, the area proportions of defects are mostly less than 0.3%, and the smallest is even 0.015%. 

The common paradigm in VII is supervised learning under a close-set setting <cit.>. This paradigm requires researchers to collect a large number of various defects in specific scenarios, and then design specific models for fitting the distribution of small defects in training. Since industrial processes are generally optimized to minimize the production of defective products <cit.>, collecting various and sufficient defects for training is often intractable, which makes it difficult for such methods to be widely applied. To tackle this problem, industrial anomaly detection (IAD) <cit.> has been widely adopted recently. IAD only requires normal samples to detect defects, while there are a large number of normal samples. Such a characteristic is conducive to the application of IAD in real scenarios. However, due to limitations of the problem formulation of IAD, it regards various defects as the same class, , abnormal class, which prevents it from locating specific defects. In real industrial processes, different defects have significantly different impacts on products, , scratches on the printed circuit board affect the appearance, while broken circuits affect the use. Ablation of aircraft engine blades affects engine efficiency, and cracks may cause engine fire <cit.>. Thus, locating specific defects has significant practical value. 



To break through the above limitations, this paper focuses on the few-shot semantic segmentation (FSS) method. FSS divides the input into the query and support sets following the episode paradigm <cit.>. It segments the query targets conditioned on the semantic clues from a few support annotations and can quickly adapt to unseen classes without retraining. These properties make FSS an ideal approach for VII. The achievement of previous FSS methods <cit.> is built upon the accurate extraction of semantic features of the support set, determined by the feature intensity in the downsampling of support annotations and the information in the resizing image. As shown in Fig. <ref>, applying these methods to segment small defects brings two problems: failure to segment and many false positives for backgrounds. The former is due to the distortion of target semantics, caused by the information loss during the resizing of the original image and the feature intensity bias in the downsampling of support annotations. The latter is caused by the missing defect, , the first and third rows in Fig. <ref>, and the train-set overfitting.



In this paper, to alleviate the above problems, we propose a small object few-shot segmentation (SOFS) model for VII. SOFS alleviates the distortion of small objects by 1) avoiding the resizing and 2) correctly indicating the intensity of target semantics. Specifically, 1) SOFS uses a non-resizing procedure to crop small objects in training and adopt the sliding window mechanism in the test. SOFS ensures that the pixel area of small objects encoded by the model is consistent with that in the original image. 2) We propose the prototype intensity downsampling of support annotations to correct the bias caused by the common interpolation. This downsampling is calculated by all pixels of corresponding regions of target semantics, providing a better estimate of target semantics.

To reduce false positives caused by the missing defect and the train-set overfitting, 1) We design an abnormal prior map in SOFS to guide the model to highlight the semantic while ignoring the normal background. This design also enables SOFS to have both FSS and few-shot anomaly detection (FAD) abilities. 2) We propose a mixed normal Dice loss to impose a large penalty when the model predicts false positives, preferentially preventing the model from predicting false positives.













Contributions: 1) We recognize FSS as an ideal solution for addressing the difficulty of obtaining diverse defects and the inability of anomaly detection to identify specific defects in VII. 2) We propose SOFS to alleviate the distortion of target semantics and many false positives. The key idea of the former is to avoid the resizing of the original image and correctly indicate the intensity of target semantics. The latter is to guide the model to focus more on backgrounds. SOFS can adapt to unseen classes without retraining, achieving FSS and FAD. To the best of our knowledge, SOFS is the first model to realize both FSS and FAD in VII. 3) FSS experiments on VISION V1 show that SOFS achieves state-of-the-art results, achieving 12.2% mean mIoU improvements compared to SegGPT. FAD experiments also demonstrate its competitive performance.



§ RELATED WORK





 §.§ Few-shot Semantic Segmentation


Few-shot semantic segmentation <cit.> predicts dense masks for novel classes with only a few annotations. According to the different fusion granularity of support features and query features, previous methods can be divided into image-level feature fusion <cit.> or pixel-level feature fusion <cit.>. The image-level feature fusion methods extract the semantic clue of support sets by a masked average pooling of a global image and then fuse this prototype with the query features. Since a pooling prototype feature cannot cover all regions of an object, the pixel-level feature fusion methods are proposed to mine the correspondence between the query pixel-level features and the support semantic-related pixel-level features, where the residual connection in the cross attention plays the role of fusing query and support features.



Previous methods mostly focus on the segmentation of PASCAL-5^i <cit.> and COCO-20^i <cit.>, where the segmented objects are mostly salient. In comparison, the segmented objects in VII are generally small. In VISION V1 <cit.>, 84% of defect types have area proportions less than 0.3%, and 38% of defect types have pixel areas less than 900. These characteristics either cause previous methods to predict many false positives or fail to segment.





 §.§ Small Object Recognition


The term “small object” refers to objects that occupy a small fraction of the input image <cit.>, , in the widely used MS COCO dataset <cit.>, it defines objects whose bounding box is 32× 32 pixels or less in a typical 480 × 640 image, while other datasets define the objects that occupy 10% of the image. The main challenges for small object recognition include information loss, low tolerance for bounding box perturbation, etc <cit.>. Information loss refers to the fact that the feature information of small objects is almost wiped out during the downsampling of the feature extractor, it has the greatest impact on performance. To alleviate this issue, there are mainly three kinds of methods. The first is to maintain the high-resolution feature map while retaining the fast process <cit.>. The second is a super-resolution-based method <cit.>, which restores the distorted structures of small objects. The third is to process small objects by multi-scale learning and hierarchical feature fusion <cit.>.

Previous methods mainly focus on small object recognition in a close-set setting. In this paper, we study this in an open-set setting, the target is to learn the meta-class-agnostic feature.







 §.§ Few-shot Anomaly Detection (Abnormal Segmentation)


Recently, few-shot anomaly detection (FAD) <cit.> has been proposed to simultaneously solve the problem of difficulty in collecting various and sufficient defects in VII and the problem of requiring retraining for different objects. FAD can quickly adapt to unknown classes by a few normal samples, , support images. FAD methods detect anomalies based on the similarity of pixel-level features between the query image and the support images. Furthermore, recent works <cit.> add textual descriptions to supplement the missing details of the image modality.

Anomaly detection regards various defects as the same class, , abnormal class, which prevents it from identifying specific defects. In VII, different defects have significantly different impacts on products, which makes anomaly detection unsuitable. In this paper, our proposed method can simultaneously achieve FSS and FAD and is applicable in more industrial scenarios.



§ METHODS





 §.§ Problem Formulation


Few-shot semantic segmentation trains model to capture the meta-knowledge of the training set, adapting to novel objects with only a few annotated support images. In definition, the model is trained on the training set D_train and tested on the test set D_test. Suppose the category sets in D_train and D_test are C_train and C_test respectively. In the few-shot setting, C_train∩ C_test=∅. Episodes are applied to D_train and D_test, each episode consists of a query set Q={(I^q, M^q)} and a support set S={(I^s_i, M^s_i)}^K_i=1 with the same class, where I^q, I^s_i∈ℝ^H̅×W̅× 3 represent the RGB images and M^q, M^s_i∈{0, 1 }^H̅×W̅ denote their binary masks, K denotes the number of support images. In training, both M^q and M^s_i are used, while only M^s_i are accessible in testing. During testing, the model requires no optimization for the novel classes. In FAD, M^s_i∈{0 }^H̅×W̅.





 §.§ Small Object Few-shot Segmentation



  §.§.§ Overview


The proposed small object few-shot segmentation (SOFS) model consists of three major steps: 1) small object enhancement designs, 2) feature augmenter, 3) feature fusion and meta prediction, as illustrated in Fig. <ref>, where the small object enhancement designs alleviate the distortion of small objects. In training, SOFS crops the semantic area on the support images and query images (Section <ref>). The cropped query/support images are then encoded by a feature augmenter to generate the semantic-related features (Section <ref>). Afterward, the encoded query and support features are fused via a non-learnable feature fusion. Finally, the fused feature interacts with the prototype features to produce the query mask (Section <ref>). In testing, SOFS uses a sliding window mechanism to produce multiple query images. To reduce false positives caused by the train-set overfitting, we elaborate on a mixed normal Dice loss (Section <ref>).





  §.§.§ Small Object Enhancement Designs





Non-resizing Procedure. The core idea is to ensure that the pixel area of small objects encoded by the model is consistent with that in the original image. As shown in Fig. <ref>, the non-resizing procedure randomly crops the small object on the original image in training and uses the sliding window mechanism to process all regions of the query image in the test. The reason why the non-resizing procedure can be used is that the segmentation of small defects in VII has a strong correlation with the surrounding areas, but a weak correlation with distant areas, which eliminates the need for the model to establish long-range dependencies. In addition to the above benefits, the non-resizing procedure can produce normal samples (the non-defective areas are all normal areas). We show in Section <ref> how these normal samples alleviate model overfitting. 

In the following part, we refer to query set and support set after the non-resizing procedure as Q={(I^q, M^q)} and S={(I^s_(1), M^s_(1))}, where we show the case K=1 without loss of generality.



Prototype Intensity Downsampling. To extract the prototype feature on the support image, we need to downsample the support mask M^s to ensure that it is consistent with the size of the support feature map. As shown in the patch-wise original mask of Fig. <ref>, every patch represents a region that is highly correlated with corresponding features on the feature map, since the segmented object is small, some corresponding features of the feature map may be weak, occupying only a small part of the original segmentation area. However, if we use the common bilinear/bicubic interpolation to downsample the original mask to indicate the feature intensity, since the bilinear/bicubic interpolation only uses 4/16 points to get the result, the result may be overestimated or underestimated, as shown in Fig. <ref>, which leads to the distortion of target semantics. To alleviate this issue, we propose a prototype intensity downsampling to replace the common bilinear/bicubic interpolation in the downsampling of M^s. Specifically, we employ a l× l convolution layer to process M^s: 


    M̂^s=Conv(M^s)/l× l ,


where M̂^s∈ℝ^H× W denotes the downsampling of the support mask, l denotes the multiple of downsampling of feature map, the parameters in the convolution kernel are all 1, the stride of convolution is l. M̂^s is calculated by all pixels of corresponding regions of features. It provides a better estimate of the intensity of prototype features, avoiding the mismatch of semantic clues.





  §.§.§ Feature Augmenter



Firstly, we encode I^s and I^q separately, resulting in feature maps F^s, F^q∈ℝ^H× W× C. Subsequently, we use a prior generation module to generate the prototype feature p∈ℝ^C, a semantic prior map M^q_s∈ℝ^H× W and an abnormal prior map M^q_a∈ℝ^H× W. p is extracted by a MaskAvgPool(F^s,M̂^s), where MaskAvgPool(·) denotes the masked average pooling (detailed descriptions are provided in Appendix). M^q_s and M^q_a are formulated as follows:










    M^q_s=Ψ(Φ(F^q) Φ(F^s)^T⊙φ(M̂^s)), 
    
    M^q_a=1-Ψ(Φ(F^q) Φ(F^s)^T⊙ (1-φ(M̂^s))),


where Φ(·): ℝ^H× W× C↦ℝ^HW× C refers to the reshape function and L_2 normalized on feature dimension, φ(·): ℝ^H× W↦ℝ^1× HW↦ℝ^HW× HW refers to the reshape and repeat along the row, Ψ(·): ℝ^HW× HW↦ℝ^HW× 1↦ℝ^H× W refers to taking the maximum value along the column and the reshape. M^q_s denotes the maximum similarity between each pixel-level query feature and support semantic features. The larger the value on M^q_s, the more relevant it is to the semantics in the support set. 1-M^q_a denotes the maximum similarity between each pixel-level query feature and support normal features. The larger the value on M^q_a, the smaller the similarity with normal features in the support set, and the more abnormal the corresponding region is.



The motivation for adding M^q_a is that we hope that SOFS can reduce false positives for backgrounds in VII. Specifically, in VII, there is a unique type of defect: the missing defect, where the target area misses some components, , the defect shown in Fig. <ref> is a missing defect, we can only observe the image background. If we only use semantic matching for the missing defect, it may highlight many false backgrounds. M^q_a matches every pixel-level query feature with the normal support features, if there is a missing defect, it can be highlighted and the normal background can not be. In addition, M^q_a enables SOFS to have FAD ability, we can input the normal support image.

After obtaining the above features, we concat F^s with p, concat F^q with p, M^q_a, and M^q_s, then use a 1× 1 convolution to produce F̅^s, F̅^q∈ℝ^H× W× C_1. Subsequently, we use the self-attention <cit.> to further improve F̅^s, F̅^q.



  §.§.§ Feature Fusion and Meta Prediction



F̅^s, F̅^q are subsequently fused by a non-learnable feature fusion. The matching mechanism follows <cit.> to replace the dot produce with the cosine similarities, which is formulated as follows:


    F̅^qs=softmax^'(Φ(F̅^q) Φ(F̅^s)^T/τ)Φ^'(F̅^s⊙ϑ(M̂^s))+Φ^'(F̅^q),


where Φ^'(·): ℝ^H× W× C↦ℝ^HW× C refers to the reshape function and non L_2 normalized, τ controls the distribution shape, softmax^'(·) refers to the normalization along the row, , reverse softmax <cit.>, ϑ(·): ℝ^H× W↦ℝ^H× W× C refers to first expanding the new dimension and then replicating along the expanded dimension. Eq. (<ref>) can be regarded as a type of cross-attention, where the learnable parameters are discarded. We think that the recognition of small objects does not need lots of parameters, more parameters may cause the risk of overfitting the category-specific information. After the feature fusion, we extract the prototype feature p̅∈ℝ^C_1 by a MaskAvgPool(F̅^s, M̂^s). Then the prediction result of SOFS model is formulated as follows:


    M̂^q_pred=𝕀(sum(M^s)>0) Sigmoid(ϕ(F̅^qsp̅))+(1-𝕀(sum(M^s)>0))M^q_a,


where 𝕀(·) refers to a indicator function, Sigmoid(·) is a sigmoid function, ϕ(·): ℝ^HW× 1↦ℝ^H× W refers to the reshape function. Compared with the learnable prediction classifier (a multi-layer perceptron) used in <cit.>, our meta prediction can be learned dynamically, it is encouraged to learn a class-agnostic meta feature. M̂^q_pred is a combination of a semantics prediction and an abnormal prior map, if the support images are all normal samples, SOFS produces a result of few-shot abnormal segmentation. Otherwise, SOFS produces a result of few-shot semantic segmentation.



  §.§.§ Loss Function



In our experiments, we find that training SOFS by Dice loss <cit.> and the cross entropy loss suffers from a severe overfitting problem, producing many false positives. To alleviate this issue, we propose a mixed normal Dice loss as follows:


    DiceLoss_normal(X, Y)= t(1-2 |X∩Y|/|X| + |Y |) + β(1-t)( 1- 1/η|X | + 1),


where t=𝕀(sum(Y)>0), X, Y refers to the prediction and the ground truth, η is a large hyperparameter, , 1e^4, β is a hyperparameter to control the weight, we set β=1 in the experiment. The normal sample term (the second term of Eq. (<ref>)) can be regarded as a regularization term of the original Dice loss. At the beginning of training, the model almost outputs 0 for all inputs. As the training evolves, once the model has a false prediction for the normal sample, the backpropagation of this term produces a larger gradient compared to defective samples, which preferentially prevents the model from predicting false positives. This loss facilitates that the model does not produce false positives while predicting the defects as much as possible. In our implementation, the normal query samples are from the non-defective areas in the non-resizing procedure, the probability of sampling normal query samples is α.

Our final loss function is a linear combination of the cross entropy loss and the mixed normal Dice loss, where the weight of the cross entropy loss is γ.



§ EXPERIMENTS





 §.§ Implementation Details


Datasets. FSS: We conduct experiments on VISION V1 <cit.> and Defect Spectrum <cit.>. VISION V1 <cit.> released in CVPR 2023 Vision-based InduStrial InspectiON workshop [https://vision-based-industrial-inspection.github.io/cvpr-2023/https://vision-based-industrial-inspection.github.io/cvpr-2023/], is one of the largest and most diverse manufacturing datasets in terms of both image number and available annotations. VISION V1 consists of 14 objects, spanning 44 defect types. 84% of defect types have area proportions less than 0.3%, and 38% of defect types have pixel areas less than 900.  Defect Spectrum <cit.> relabeled the data from MVTec AD <cit.>, DAGM2007 <cit.>, part of VISION V1 <cit.> and Cotton-Fabric <cit.>, including 97 defect types. Defect Spectrum is the dataset with the largest number of defect types so far. 63% of defect types have area proportions less than 2%. For the above datasets, cross-validation is conducted by dividing all objects into 3 folds, we ensure that the distributions of area proportions of defects in each fold are roughly the same. In the experiments, two folds serve as training data, while the remaining one is used for testing. To ensure performance stability and fairness for comparison, the query/support pairs sampling by all methods are the same. FAD:We use two benchmark FAD datasets, , MVTec AD <cit.> and VisA <cit.>, including 27 objects in total. More details of datasets and experiments are shown in Appendix.

Metrics. FSS: We adopt the mean intersection over union (mIoU) as the evaluation metric. We denote mIoU=1/CΣ^C_i=1IoU_i, where C is the number of defect classes in each fold, and IoU_i indicates intersection-over-union for class i. FAD: For classification, we report Area Under the Receiver Operating Characteristic (AUROC). For segmentation, we report pixel-wise AUROC and pixel-wise F_1-max.

Training and Test Details. FSS: We use ResNet-50 <cit.> and DINO v2 Base <cit.> as the encoder to verify the effectiveness of the proposed method and reimplementation methods on different backbones. The training input sizes of ResNet-50 and DINO v2 Base are set to 512× 512 and 518× 518 respectively. l in the prototype intensity downsampling are set to 8 and 14 for ResNet-50 and DINO v2 Base respectively. The probability α is set to 0.3, η is set to 1e^5, γ is set to 0.01, τ is set to 0.1. In the non-resizing procedure, the stride of sliding window and the input size in test are the same as the input size in training. All reimplementation methods and SOFS use the same training augmentations. FAD: We use DINO v2 Base <cit.> as the encoder. The input size is set to 518× 518. We only test SOFS. More details of SOFS and reimplementation methods are shown in Appendix.







 §.§ Comparison with State-of-the-Art Methods


FSS: In Table <ref>, we report the comparison of SOFS with other state-of-the-art FSS approaches on VISION V1. Compared with other methods, SOFS can gain passable improvement on ResNet-50. However, due to the insufficient discriminability of pre-trained features for small defects, the performance is extremely low. As we choose DINO v2 Base <cit.> as the encoder, the performance has a significant improvement and outperforms the others by a large margin. Despite using a model with fewer parameters, SOFS significantly outperforms SegGPT, achieving 11.4% (1-shot) and 12.2% (5-shot) of mean mIoU improvements. Notably, when applying the non-resizing procedure and the mixed normal Dice loss to comparison methods, the performance improves consistently, this illustrates the generality of the proposed method. In addition, we show qualitative results of SOFS in Fig. <ref>, <ref>, <ref>. Although there are different defect types in one image, SOFS can segment the specific small defects. Table <ref> shows the FSS performance on Defect Spectrum. Although this dataset contains larger defects than VISION V1, SOFS still achieves the best in most cases. The limitation of SOFS is that if the support images increase (K increases), the performance improvement for large defects is relatively weak (VISION V1 F2, DS-Spectrum F1 and F2). We think this is due to the underfitting to large defects. FAD: Table <ref> also shows the comparison of SOFS with other state-of-the-art few-shot anomaly detection approaches. Although SOFS only receives the image information, it can achieve a competitive performance with other multi-modal models. This shows that SOFS has strong capabilities for both FSS and FAD. 

In addition, we show the performance of FSS under domain shifts. We only test the model trained on VISION V1 <cit.> for Defect Spectrum <cit.> and AeBAD <cit.>, we verify the performance under different domain shifts. To avoid conflicts between test classes and training classes, we exclude VISION V1 from our test. Test defect types are 81. The results are shown in Table <ref>. Although these datasets contain larger defects than VISION V1 and different domain shifts, SOFS still achieves the best performance. More quantitative and qualitative results are shown in Appendix.







 §.§ Ablation Study


We report the ablation studies in this section to investigate the effectiveness of each component. All ablation experiments are conducted under VISION V1 1-shot setting with DINO v2 Base backbone.

Small object enhancement designs ablation. Table <ref> shows results regarding the effectiveness of small object enhancement designs. The baseline (line 1) uses the bilinear interpolation for M̂^s and the standard resizing procedure, its performance is extremely low. Then we add the prototype feature intensity map (PFIM), resulting in improved results. Since many defects account for small area proportions in F0 and F1, the performances of F0 and F1 are still low. Due to “Cylinder” with large defects in F2, the improvement for F2 is significant. Subsequently, we only add the non-resizing procedure (NR). Compared with line 2, this brings a greater improvement. However, the improvement in F1 is still not enough. This is mainly because there are many defects with small pixel areas in F1. NR cannot deal with this kind of problem. Finally, we use PFIM and NR to simultaneously improve the performance for defects with small pixel areas and small area proportions, resulting in a significant improvement for F1.

Other components ablation. Table <ref> shows results regarding the effectiveness of other components in SOFS. The baseline (line 1) is built upon the matching mechanism in <cit.> (a type of cross attention), the learnable prediction classifier (a multi-layer perceptron) used in <cit.> and trained by Dice loss <cit.>. We first replace Dice loss with the proposed mixed normal Dice loss (MNDL), resulting in a large improvement for F0 and F1. However, since the training on F2 does not include large defects (“Cylinder” with large defects is in F2), the model with too many parameters causes overfitting to small defects. Then we replace the matching mechanism in <cit.> with a non-learnable feature fusion, the performances for F0, F1, and F2 improve. Subsequently, we add an abnormal prior map (APM). As shown in Fig. <ref>, APM can reduce the false positive regions of missing defects, resulting in further improvements. Finally, we replace the learnable prediction classifier with a meta prediction, bringing 1.5% and 1.1% mIoU improvement for F0 and F1. But this also brings a slight degradation for F2. We find that this is due to insufficient prediction for large defects.





η and α in the mixed normal Dice loss. Fig. <ref> shows the results for different η and α. η: If η is small (0-10^-2), it does not penalize false positives enough, resulting in a relatively weak performance, especially for Fold 2. If η is set to 10^-1 or 10^0, the performance of Fold 0 and Fold 1 improve significantly. However, the performance of Fold 2 is 0. We found that after the training started, the model quickly fell into a local optimum with an output of all 0. This is mainly because the training for Fold 2 does not contain the large defects (Cylinder). In other words, if the training samples are all small defects, then some η will cause the training to fall into a useless local optimal solution. If η is set to a large number, this phenomenon will disappear. α: Different α have a significant impact on the results. If α is small, the number of normal query samples per iteration is small (possibly none), resulting in no penalty for false positives. If α is large, it causes the underfitting for defective samples.

r0.5
    

    
    < g r a p h i c s >

    

    The results as the test input size evolves in the sliding window mechanism.
    
    



Input size in test. Fig. <ref> shows the results as the test input size evolves in the sliding window mechanism (the stride is the same as the input size). We can find that different test input sizes have a significant impact on the segmentation of small objects, especially for the defects with small area proportions (Fold 0). The smaller the input size, the higher the latency of the model. We should choose the test input size based on actual requirements.



§ CONCLUSION



In this paper, we propose SOFS to tackle the problems existing in supervised learning under a close-set setting and industrial anomaly detection in VII. SOFS significantly improves the small object segmentation by avoiding resizing and a prototype feature intensity downsampling. Further, we add an abnormal prior map in SOFS to guide the model to reduce false positives for backgrounds, while enabling SOFS to realize both FSS and FAD. In addition, we propose a mixed normal Dice loss, non-learnable feature fusion, and meta prediction to alleviate the overfitting problem. The experiments for different defects under an open-set setting demonstrate the superior performance of SOFS. We hope that SOFS can provide some insights for the community of VII.

Limitations and Future Works 1) As shown in Fig. <ref>, the differences among defects in the same object are slight, which leads to some false positives. We think that test-time training/adaptation <cit.> may be a promising method to tackle this problem. 2) As shown in Fig. <ref> (7th row) and Fig. <ref> (3rd row), SOFS produces false negatives for defects with extreme aspect ratio. This may originate from the receptive field of the encoder. 3) If SOFS is used to process high-resolution images, due to the sliding window mechanism in the test, the latency of SOFS is high. We will explore some designs to reduce latency and segment defects with extreme aspect ratios in the future. 4) SOFS is allowed to adopt a mixture of normal support images and defective support images for training and inference, which is a more common scenario in reality. We will explore this direction in the future.





abbrvnat










§ APPENDIX




 §.§ Supplemental Details of SOFS


Masked average pooling. The masked average pooling is formulated as follows:


    MaskAvgPool(F,M)= sum^'(F⊙ϑ(M))/sum(M) ,


where F∈ℝ^H× W× C, M∈ℝ^H× W, ϑ(·): ℝ^H× W↦ℝ^H× W× C refers to first expanding the new dimension and then replicating along the expanded dimension, sum^'(·): ℝ^H× W× C↦ℝ^C refers to the addition of spatial features, sum(·): ℝ^H× W↦ℝ refers to the addition along spatial dimensions.  

Extension to K-shot setting. In extension to K-shot (K > 1) setting, K support images with their annoted masks S={(I^s_i, M^s_i)}^K_i=1 and the query set Q={(I^q, M^q)} are given. M̂^s_i, p_i, M^q_i_s and M^q_i_a are firstly calculated, where p_i is extracted by a MaskAvgPool(F_i^s,M̂^s), M^q_i_s denote a semantic prior map calculated by F_i^s and F^q. Then we stack all M^q_i_s and M^q_i_a along the number dimension of support set respectively. And we calculate the mean of maps concatenated toghter to obtain M^q_s and M^q_a. In the feature fusion, we calculate the cosine similarity between the query feature and all support features. We concat all F_i^s along the spatial dimension. The remaining steps are the same as 1-shot.

K-shot implementation. Due to the non-resizing procedure, we can sample multiple support images under a 1-shot scenario. These support images are from the random crops containing defects. In 1-shot experiments, we sample 4 random crops in one support image. In 5-shot experiments, we sample 1 random crop in one support image. We use this implementation for all methods.

Few-shot Anomaly Detection. We use the maximum value on M̂^q_pred to indicate the abnormal degree of the image.



 §.§ Supplemental Implementation Details


SOFS. SOFS is built upon the Pytorch <cit.> framework. All models are trained and tested on 8 NVIDIA GeForce RTX 2080 Ti GPUs. All reimplementation methods and SOFS use the same training augmentations, including rotation, blur, and flip. In the non-resizing procedure, we use the random crop, where the crop image includes the defect. We only use the non-resizing procedure if the area proportion of defect in the support set is less than 1%, otherwise, we use the common resizing. In addition, following the previous works, SOFS adopts features from different layers of backbones to further improve the useful information. 

SOFS is trained in an episode fashion for 50 epochs for VISION V1, the batch size is 4 for every GPU. During training, AdamW optimizer is adopted, and the learning rate is set to 1e^-5, the weight decay is 0.01, and the “poly” strategy is used to adjust the learning rate. During testing, predictions are resized back to the original sizes of the input images, keeping the groundtruth labels intact.

In the few-shot anomaly detection, SOFS adopts the common resizing procedure.

Reimplementation of Comparison Methods We compare SOFS with the recent methods: PFENet <cit.>, HDMNet <cit.>, and SegGPT <cit.>. We reproduce them following the official codes and adjust the hyperparameters to maximize performance. For PFENet and HDMNet, we train the models by cross entropy loss and Dice loss <cit.>. Due to the small objects, we adjust the downsampling modules in PFENet and HDMNet, we do not downsample the feature map. The input sizes of models are the same as SOFS. In addition, we do not mask the support features in the feature extraction of the freezing encoder like previous methods, masking leads to poor performance. We only test the model for SegGPT, and the input size is set to 448× 448, which is the same as the original setting. We do not observe the improvement if the input size is set to 512× 512.

We use ResNet-50 <cit.> and DINO v2 Base <cit.> as the encoder to extract features with freezing parameters to verify the effectiveness of the proposed method and reimplementation methods on different backbones. For ResNet-50, we use the multi-scale features from layer 2 and layer 3, and the feature processing method in <cit.> is adopted to upsample the high-level features. We do not use PSPNet <cit.> to train the base learner, which is used in the previous methods <cit.>. We find that this causes a severe overfitting problem to fit the category-specific distribution. Using the original pretrained model can alleviate this. The feature dimension C is set to 64 and C_1 is set to 64. For DINO v2 Base, we use the features from layers 5, 6, 7, 8, 9, 10. The feature dimension C is set to 256 and C_1 is set to 256. To ensure fairness, the features from the different layers of backbones for SOFS and reimplementation methods are the same. Since the backbone cannot be replaced, SegGPT is verified by ViT-Large <cit.>.





In the experiments on VISION V1 and Defect Spectrum, we train the models with 3 runs for PFENet, HDMNet, and SOFS, and then test each model with 5 runs, for a total of 15 runs. For SegGPT, we only test it with 15 runs. In the experiments of domain shifts (see Section <ref>), we test models with 10 runs. To ensure performance stability and fairness for comparison, the query/support pairs sampling by all methods are the same (the non-resizing procedure adopts the crops of the image).



 §.§ Supplemental Details for Dataset


VISION V1. VISION V1 <cit.> released in CVPR 2023 Vision-based InduStrial InspectiON workshop, is one of the largest and most diverse manufacturing datasets in terms of both image number and available annotations. VISION V1 consists of 14 objects sourced from Roboflow and <cit.>, 10k high-quality defect segmentation annotations on 4k images, spanning 44 defect types. Area proportions and pixel areas (number of pixels) for different defects are shown in Fig. <ref>. 84% of defect types have area proportions less than 0.3%, and 38% of defect types have pixel areas less than 900. These characteristics bring great challenges to defect segmentation. In the few-shot semantic segmentation, we divide all objects into 3 folds, we ensure that the distributions of area proportions of defects in each fold are roughly the same. The test objects for each fold are shown in Table <ref>.

Defect Spectrum. Defect Spectrum <cit.> relabels the data from MVTec AD <cit.>, DAGM2007 <cit.>, part of VISION V1 <cit.> and Cotton-Fabric <cit.>, ensuring that every image is accompanied by precise and diverse category annotations. It is the dataset with the largest number of defect categories so far, including in total 97 defect types. Area proportions and pixel areas (number of pixels) for different defects are shown in Fig. <ref>. 63% of defect types have area proportions less than 2%. In the few-shot semantic segmentation, we divide all objects into 3 folds, we ensure that the distributions of area proportions of defects in each fold are roughly the same. The test objects for each fold are shown in Table <ref>. Compared with the defects in VISION V1, the defects in Defect Spectrum are larger. The objects in this dataset are more under experimental conditions, , placing a single object in a noise-free background, and there is a certain gap between the defects that need to be detected in real manufacturing. Thus, we choose VISION V1 as the main experimental dataset. 





AeBAD. AeBAD <cit.> is used to detect defects in aero-engine blades. The defects of blades are under different domains, including different illuminations, backgrounds, and views. The defects contain 4 types, involving some minor defects. This dataset is used to verify the robustness of different models for different domains. In <ref>, we only test the model trained on VISION V1 <cit.> for Defect Spectrum <cit.> and AeBAD <cit.>, we verify the performance under domain shifts.



 §.§ Supplemental Experiments



The results (mean IoU) of different defect types of every object on VISION V1 are shown in Table <ref> and <ref>. SOFS achieves the best in most cases.









 §.§ Supplemental Ablation Study


The outputs of the model with and without the abnormal prior map are shown in Fig. <ref>. We can find that APM can reduce the false positive regions of missing defects.



 §.§ More Qualitative Results


We show the visualizations for different defects of a query image conditioned on different support images in Fig. <ref> and Fig. <ref>. This shows the segmentation results of SOFS for the specified defect. We show the qualitative results of every defect type on VISION V1 in Fig. <ref>, <ref>, <ref>, <ref>, <ref>, <ref>, <ref>.
























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



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



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







§ INTRODUCTION













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

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

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

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




§ RELATED WORK



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

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



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

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





§ CREW: DESIGN AND COMPONENTS






 §.§ Platform Vision and Design Philosophy




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

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

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

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

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



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

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





 §.§ Environments



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

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






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



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





 §.§ Human and Agent Role Assignment



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

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

[13]r0.25


  

  
    < g r a p h i c s >

  

  

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


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




 §.§ Multiplayer and Parallel Sessions



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







 §.§ Human and Agent Data Collection



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

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

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






 §.§ Designing Algorithms



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

[16]r0.4


  

  
    < g r a p h i c s >

  

  

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






 §.§ Modular Pipeline Design for Quantifying Human Characteristics



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








§ BENCHMARKING STUDY




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







 §.§ Experiment Setup
 


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

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


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

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

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

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




 §.§ Results



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













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


[17]r0.5


  

  
    < g r a p h i c s >

  

  

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

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




§ CONCLUSION, LIMITATION, AND FUTURE WORK



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

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



§ ACKNOWLEDGEMENTS

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













































































































    

















§ ACCESSING CREW

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








§ PLATFORM DETAILS

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




§ ADDITIONAL RESULTS




 §.§ Feedback Visualization

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





 §.§ Full Cognitive Test Analysis

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









§ COMPUTATIONAL RESOURCES


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





§ BENCHMARKING EXPERIMENT DETAILS



 §.§ Implementation details


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


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


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



 §.§ Hyperparameters

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






 §.§ Human subject experiment details


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



  §.§.§ Instruction Video Script


 General Introduction:  
 
 

Cognitive Test Introduction: 



Eye Alignment Instruction: 



Reflex Instruction: 



Theory of Behavior Instruction: 



Mention Rotation and Mental Fitting Instruction: 



Spatial Mapping Instruction: 



Human Guiding AI Introduction: 



Bowling Instruction: 



Find Treasure Instruction: 



1v1 Hide-and-Seek Instruction: 






  §.§.§ Compensation

We pay each human subject $20 for participation. 



§ AUTHOR STATEMENT

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




































 Paper Title: Tree-of-Traversals: A Zero-Shot Reasoning Algorithm for Augmenting Black-box Language Models with Knowledge Graphs.


An LLM-based Readability Measurement for Unit Tests' Context-aware Inputs
    Zhichao Zhou,
        Yutian Tang,
        Yun Lin,
        and Jingzhu He
Z. Zhou and J. He are with School of Information Science and Technology, ShanghaiTech University, China
Y. Lin is with the Shanghai Jiao Tong University, China
Y. Tang is with University of Glasgow, United Kingdom
J. He (hejzh1@shanghaitech.edu.cn) is the corresponding author.
    August 5, 2024
====================================================================================================================================================================================================================================================================================================================================================================

§ ABSTRACT 

Knowledge graphs (KGs) complement Large Language Models (LLMs) by providing reliable, structured, domain-specific, and up-to-date external knowledge. However, KGs and LLMs are often developed separately and must be integrated after training. 
We introduce Tree-of-Traversals, a novel zero-shot reasoning algorithm that enables augmentation of black-box LLMs with one or more KGs. The algorithm equips a LLM with actions for interfacing a KG and enables the LLM to perform tree search over possible thoughts and actions to find high confidence reasoning paths. We evaluate on two popular benchmark datasets. Our results show that significantly improves performance on question answering and KG question answering tasks. Code is available at <https://github.com/amazon-science/tree-of-traversals>
 






















































§ INTRODUCTION

Large language models (LLMs) are used for a range of knowledge-intensive tasks such as information retrieval <cit.>, summarization <cit.>, and question answering <cit.>. Trained on large amounts of textual data, these models learn a wide breadth of information. However, LLMs suffer from several limitations – they produce hallucinated information <cit.>, lack deep domain-specific knowledge <cit.>, and have a static knowledge cutoff when training ends.

Knowledge graphs (KGs) naturally complement LLM weaknesses. KGs contain up-to-date information covering general <cit.> and/or domain-specific topics <cit.> in highly structured and interpretable format. Augmenting an LLM's ability to reason and respond in natural language with an external KG's up-to-date knowledge presents a path toward a  reliable and factual LLM. 

The rise of powerful LLMs with new capabilities has renewed interest in combining LLMs with KGS. Numerous survey and position papers have recently emphasized their combined potential <cit.>. 
Existing works augmented LLMs with KGs in multiple ways, such as integration into pretraining <cit.>, fine-tuning <cit.>, or later adaptation with subsequently trained components <cit.>. All of these carry some limitations. In particular, training or fine-tuning a large scale LLMs is computationally expensive. 
In some cases, model weights are unavailable publicly.
Finally, the largest KGs require their own servers and cannot be integrated in-memory with LLMs. Additionally, previous works do not consider the case of augmenting with multiple KGs. 


An algorithm that allows the augmentation of a powerful black-box LLM with any number of internal or external KGs without training from scratch or fine-tuning the model is valuable. Such a zero-shot algorithm would enable several innovative use cases such as (1) customers using a black-box LLM API in conjunction with an internal domain-specific KG, (2) integration of personalized KGs into an LLM without the risks associated with training a model with such personal data, (3) integration of deep domain knowledge via an array of API accessible KGs (e.g., IMDb[https://aws.amazon.com/marketplace/seller-profile?id=0af153a3-339f-48c2-8b42-3b9fa26d3367IMDb API], MusicBrainz[https://musicbrainz.org/doc/MusicBrainz_APIMusicBrainz API]).

We introduce Tree-of-Traversals, a novel algorithm that addresses the above issues by allowing the augmentation of any powerful LLM with arbitrary number of KGs in a zero-shot fashion. It requires no training, functions with black-box access to the LLM, and works with any API accessible KG. Our contributions are:

    
  * Tree of traversals: A novel zero-shot algorithm for augmenting any powerful LLM with arbitrary number of KGs and enabling advanced KG reasoning using tree search. 
    
  * Evaluation of on two question answering tasks: 2WikiMultiHop and QALD-10 and comparison with baselines.
    
  * Development of a new dataset to test combined reasoning over a general and a domain-specific KG, and evaluation of on this dataset.


We conduct detailed experiments on three models of varied sizes hosted on Amazon Bedrock and present detailed ablation studies.



§ TREE-OF-TRAVERSALS

The algorithm maintains a local KG subgraph that is repeatedly expanded until it contains all the information required by an LLM to answer the given query. At the start, a local KG subgraph is initialized to contain the entities present in the original query. It is then expanded using a tree search algorithm to choose actions and thoughts generated by an LLM in order to obtain relevant knowledge from the KG using a KG interface. The algorithm halts when the LLM is able to answer the original query using the local KG subgraph as context. consists of three major components. 
(1) A knowledge graph interface implemented to interact with one or more required KGs.
(2) An action state machine (ASM) which is a finite state machine that defines the feasible space of actions, states, and prompt templates when the LLM interacts with a KG to expand the local KG subgraph.
(3) A tree search algorithm which defines the overall LLM search trajectory such as best first search, backtrack upon making a mistake, and termination condition when an answer is found. 




 §.§ Knowledge Graph Interface


The knowledge graph interface allows to interact with one or multiple KGs. Let 𝒦=(ℰ, ℛ, 𝒯) be a single KG. ℰ is the set of entities in which each entity consists of an identifier, a label, and an optional description (e.g., Q35332, `Christopher Nolan', `British-American filmmaker'). ℛ is the set of relation types, each consisting of an identifier, a label, and an optional inverse label (P57, `director', `is director of'). 𝒯 is the set of edges or facts in the KG in which each edge is of the form (s, r, o) where s, o ∈ℰ and r∈ℛ, e.g., (`Inception', `director', `Christopher Nolan'). can be used for 𝒦 as long as the following interfaces are implemented. 

    
  * (q)→ℰ_0: It takes as input a query q, extracts entities from q and returns the linked entities ℰ_0 ⊂ℰ where ℰ_0 are the entities from 𝒦 that are referenced in q. 
    
  * (ℰ_selected)→ℛ_options:
It takes a set of entities, ℰ_selected, and returns the relation types, ℛ_options⊂ℛ, that ℰ_selected have in 𝒦: {r|(s,r,o)∈𝒯,s∈ℰ_selected}
    
  * (ℰ_selected, r)
→ 
𝒯_added, ℰ_added: It takes a set of entities and a selected relation type, and returns all edges with relation type r for source entities in ℰ_selected:
{(s,r,o)∈𝒯|s∈ℰ_selected, r=r, o∈ℰ}. It also returns the new entities, ℰ_added, that are entities reached with 𝒯_added.


This interface is implemented with SPARQL queries when available; otherwise, we use the graph API that is available for the KG. For multiple KGs, each interface is implemented separately.



 §.§ Action State Machine (ASM)



One of the challenges with developing a zero-shot LLM algorithm that works with arbitrary KGs is that the LLM does not know what relations are available in the graph or what relations are valid for a given entity. Few-shot or in-context learning approaches can only cover a handful of the possible relation types (e.g., Wikidata has over 11,000 relation types) <cit.>. 



To overcome these issues we break the task of expanding a local KG subgraph into multiple subtasks. We use a finite state machine with the following actions: , 
, 
, 
, and 
, and states: default, selecting-entities, selecting-relation, and done as shown in Figure <ref>. This is named as Action State Machine (ASM) throughout the paper.

From the default state, can either , , or choose to . After chooses to , it first is prompted to  about which it needs more information (e.g., `Inception' in Figure <ref>). It is then prompted to  from a list of candidate relations provided by the KG interface's  method. After selecting a relation (e.g., `cast member' in Figure <ref>), all edges containing one of the selected entities as the source and the selected relation as the relation are added to the local KG subgraph. is then able to , , or  again.






  
Prompt Templates. Each state in the ASM other than the `done' state is associated with a unique prompt template. Prompt templates are filled with information from the local KG subgraph and the KG interface before presenting them to the LLM. A customized prompt for each state allows us to present precise and relevant information along with specific instructions for each state to the LLM simplifying the LLM's task. For instance, the prompt for `selecting-entities' presents the available options of entities to choose from that are in the local KG subgraph (see Figure <ref>). Prompt templates for all the states are in the Appendix <ref>-<ref>.

Local KG subgraph is represented using a a token efficient YAML format that minimizes repetition for multiple edges on the same entity.





  
Invoking KG Interfaces. Outside of initialization, there are two times in which the ASM needs to invoke the KG interface: (1) When constructing the selecting-relation prompt, the algorithm calls (ℰ_selected) to gather available choices of relations. (2) When executing the transition from selecting-relation to default after having selected a relation type r, the algorithm calls (ℰ_selected, r) so that it can add the new edges and entities to the local subgraph.



 §.§ Tree Search Algorithm

Our approach draws inspiration from the Tree-of-Thoughts approach <cit.> in which an LLM is given increased reasoning power by allowing generation of multiple thoughts at a time and building a search tree over the resulting reasoning chains. We extend this approach to augment LLMs with KGs through allowing the generation of actions in addition to thoughts, and building a search tree over them. Challenges arise because Tree-of-Thoughts was not designed to incorporate actions and was not designed for knowledge intensive question answering tasks. As a result, we introduce some modifications: incorporation of actions via the ASM, a slightly different search procedure and stopping condition to better handle QA, and a different sampling procedure for improved diversity when doing constrained sampling. 


Algorithm <ref> presents the tree search algorithm. Given a query q it begins with initialization using (q) as described in Section <ref>.  After initialization, it searches by (1) choosing the unexplored node to expand based on the node's value assigned by the value function (Best First→ Depth First as tie-breaker), (2) sampling k actions from the LLM (k is branching factor) using the prompt associated with the selected node’s state in the ASM, (3) for each of the sampled action, applying the transition function, and (4) evaluating the value of the resulting nodes with the LLM value function. The search stops when an answer is found with a node value exceeding the threshold τ. To bound the search space we add two hyper-parameters: max depth of the tree search, after which the algorithm is forced to transition to the done state (i.e., answer the query), and max expansions after which the model halts exploration and returns "Cannot find answer". 

    




















  
Value Function Guidance.
computes the value of a node to determine its utility. This value is created by the LLM using evaluation prompts (step  in Algorithm <ref>). The value can be between 0 to 1 where 1 indicates highest utility. We use two types of evaluation prompts: one for intermediate states and one for answer states. The prompts include the original query, the local KG subgraph, the trajectory of previous actions, followed by instructions for evaluating the node (see Appendix <ref> and <ref>). These values are then used to guide the exploration of the action space. Specifically,  returns the unexplored node with the highest value (Best First). If there are nodes with the same value, Depth First Search is used.



  
Chain-of-Traversals. In some cases we can find the answer to a query with a single sequence of thoughts and actions using the ASM and KG interface. This is equivalent to Tree-of-Traversals with a branching factor of k=1. We refer to this as Chain-of-Traversals. While useful for comparison, experiments show the benefit of considering multiple branches.



 §.§ with Multiple KGs

Augmenting an LLM with more than one KG mainly involves building KG interfaces for each of the added KGs. There are a few other changes to the algorithm. (1) The entities extracted in (q) are matched with each KG interface. (2) When presenting the options of relations during selecting-relation, (ℰ_selected) is called on each KG interface. (3) When adding a new entity to the local KG subgraph we call an entity linking function for the other KG interfaces. We allow for the entity linking function to be a separate function in the KG interface or to fallback on (o) where o is the text label of the entity that has just been added. This makes allows the use of explicit links between common KGs if available while still functioning without. 



§ EXPERIMENTS



We evaluate using three different models available on Amazon Bedrock[<https://aws.amazon.com/bedrock>]: Claude-Instant (), Llama2 70b (), and Llama2 13b (). AWS Bedrock provides on-demand access through a single API to various Foundation Models, including both open source and black box ones. This is precisely a use case is designed for.















       













 §.§ Tasks and Datasets.
 We first evaluate the algorithm on two common tasks used for evaluating an LLM's knowledge: 2WikiMultiHop and QALD-10. To allow testing on complex questions requiring knowledge from multiple KGs we create a new dataset that requires knowledge from multiple KGs. 


2WikiMultiHop Dataset <cit.> was constructed by extracting multi-hop templates from HotPotQA <cit.>, combining templates to get complex reasoning questions, generating candidate questions with Wikidata, and then confirming that the entity mentions for each edge appear in the Wikipedia paragraphs. Answers to these questions can be derived from both Wikipedia and Wikidata <cit.>. We subsample 500 questions from the test set following the approach used in ReAct <cit.>, including the sampling seed value of 233.


QALD-10 Dataset <cit.>  is a multi-lingual Knowledge Graph Question Answering (KGQA) dataset with 395 questions created by humans, translated into other languages, and then constructed as a SPARQL query over Wikidata[<https://github.com/KGQA/QALD-10>]. These questions are more varied in terms of reasoning structure than the questions in 2WikiMultiHop (e.g. requiring multiple answers or aggregations). We used the questions in English language. 








        































MusicBrainz-x-Wikidata Dataset One novel use-case of is synthesizing and reasoning over multiple knowledge graph sources. There is no existing dataset that requires synthesizing information from multiple KGs to answer individual questions. Therefore, to test the reasoning ability using multiple KGs, we create a new dataset, MusicBrainz-x-Wikidata, containing 109 questions that require reasoning with information from both MusicBrainz and Wikidata. Unlike Wikidata which contains general knowledge, MusicBrainz is a deep domain-specific database on the music industry. The vast majority of this information is unlikely to be known by large language models. We construct the MusicBrainz-x-Wikidata dataset with human annotators who were provided with instructions to find reasoning paths between these two KGs, the question types to focus on, and some example questions. The curated questions were checked for ambiguity and sensibility. By design, each question requires extracting information from both knowledge graphs in order to answer it successfully. In addition to the reasoning types in 2WikiMultiHop, this dataset  contains questions that involve aggregations, comparisons of aggregations, qualifications, and complex combinations of these. An example can be seen in Figure <ref>. Detail on instructions, question types, and examples are in Appendix <ref>. 



 §.§ Metric.
 We use Exact Match Included (EM-in) as the evaluation metric. EM-in is 1 if the ground truth answer appears with an exact match anywhere in the answer and 0 otherwise. This accounts for the propensity of an LLM to output answers in a sentence with varying syntax. This is a common metric but is often referred to interchangeably with Exact Match (EM) <cit.>. When there are multiple answers, we compute average EM-in over all labels. 



 §.§ Comparison Baselines.
 We experiment against three related approaches that can work with any black-box LLMs: (1) Chain-of-Thought (CoT) prompting <cit.> (2) ReAct <cit.> which iterates between generating thoughts and generating actions for searching and retrieving from a text knowledge base like Wikipedia, and (3) Forward Looking Active Retrieval (FLARe) <cit.> which iterates between generating thoughts and then retrieving from a knowledge base to correct inaccuracies. 



 §.§ Implementation details.
 For all models, we use a sampling temperature of 0.0 for the LLM when multiple samples are not required and a temperature of 1.0 when diverse samples are required. We test our approach in two settings: (1) with a branching factor of k=1 termed as Chain-of-Traversals, and (2) with a branching factor of k=3 termed as Tree-of-Traversals. In both setups, we set maximum depth to 7 which means that upon reaching the default action state beyond depth 7, the only available action is to answer the question. For Tree-of-Traversals, we set the maximum total expansions to 20. The answer threshold τ is set to 0.8 which corresponds to high confidence answers supported by the KG according the  prompt (Appendix <ref>).



For 2WikiMultiHop and QALD-10, we use Wikidata as the knowledge graph <cit.>. 
For MusicBrainz-x-Wikidata, we use MusicBrainz in addition to Wikidata. 

We implement the KG interface for Wikidata using Wikidata SPARQL queries[<https://query.wikidata.org/>], and implement the KG interface for MusicBrainz KG using the MusicBrainz API[<https://musicbrainz.org/doc/MusicBrainz_API>].



§ RESULTS AND DISCUSSION





Table <ref> presents the results of our experiments on 2WikiMultiHop and QALD-10. For all the models, outperforms the baselines on 2WikiMultiHop, setting state-of-the-art results for these tasks in the zero-shot setting. We hypothesize that much of this gain is due to Tree-of-Traversals' access to the knowledge base via proposed KG interface and its thought-action procedure guided by the ASM. This is evident, as even Chain-of-Traversals, which does not perform a tree traversal (including multiple thoughts/actions, backtracking, and node value computation), significantly outperforms ReAct's knowledge grounding: 8.7% higher accuracy than ReAct→ CoT on 2WikiMultiHop when averaged over all models. Compared to Chain-of-Traversals, further improves performance. It gives on average a 12.8% absolute accuracy increase for 2WikiMultiHop and 4.3% absolute improvement for QALD-10. We note that the better performing the model is, the more it stands to gain from as noted by the difference between Llama-70b and Llama-13b. 



On MusicBrainz-x-Wikidata (Table <ref>), which contains challenging reasoning questions requiring access to two KGs, there is an average of 37.4% relative improvement from over Chain-of-Traversals as shown in Table <ref>. Both Chain-of-traversals and outperform Chain-of-Thoughts on this dataset. We only compare with Chain-of-Thoughts as other retrieval methods have low coverage over the MusicBrainz knowledge base. 





 §.§ Effect of the Value Function.
 relies on signal from the value function to pick the final tree trajectory. If the values were arbitrary, then Tree-of-Traversals would not do any better than Chain-of-Traversals. Figure <ref> shows that there is a meaningful signal from the value function for all models. There is an average performance difference of 31.0% between the accuracy for answers valued at 1.0 vs 0.0. This represents an average relative improvement of 83.2% when selecting answers valued at 1.0 over those valued at 0.0. See Appendix <ref> for individual profiles of each model's value function.



 §.§ Effect of Backtracking.
 To determine the effect of backtracking in Tree-of-Traversals, we ask the counter-factual of how would the model perform if it could not backtrack (Figure <ref>). We limit the analysis to 2WikiMultiHop questions in which Tree-of-Traversals generates answers on a subtree, and the model eventually backtracks on that subtree (i.e., the cases where we have a counter-factual). As a result, these questions are generally more challenging than the overall distribution of questions. In these cases, we compare the result if the highest-valued answer was taken from the first searched subtree compared to the ultimate answer after backtracking. We find that the ability to backtrack gives a significant accuracy increase ranging from 4.1% to 12.3%.  





 §.§ Performance on MusicBrainz-x-Wikidata.
 For MusicBrainz-x-Wikidata dataset, we only compare against Chain-of-Thought since the implementations for other baseline algorithms do not have access to similar music-specific knowledge bases. Despite this, we observed some interesting results. Chain-of-Thought does far worse on MusicBrainz-x-Wikidata than on the more general datasets based exclusively on Wikipedia/Wikidata (10.1%-13.8% vs. 25.2%-47.4%). This could be because LLMs are trained on vast amount of general knowledge such as that found in Wikipedia, but they are not trained with as much data from specific domains such as music. Besides the presence of KG interface and ASM with Tree-of-Traversals, this could be an additional reason why Tree-of-Traversals does 2.2 times as well as Chain-of-Thought on MusicBrainz-x-Wikidata. This demonstrates the importance of augmenting LLMs with domain-specific KGs and/or multiple KGs which the proposed Tree-of-Traversals is capable of.



 §.§ Analysis of Baselines.
 ReAct underperforms Chain-of-Thought in most cases. This phenomenon was seen in the original ReAct paper which noted that their approach significantly reduced hallucinations despite lower accuracy <cit.>. Therefore, falling back on Chain-of-Thought results in an accuracy improvements for for Llama2-70b and Claude-Instant. We note that  vastly underperforms Llama2-70b on ReAct. The primary reason for this is  refuses to "search" people and apologizes after performing invalid actions. Despite this, ReAct→CoT still improves over CoT on both 2WikiMultiHop and QALD-10. FLARe improves model performance on 2WikiMultiHop but not on QALD-10. This may be due to better overlap between the text knowledge base for 2WikiMultiHop. 



§ RELATED WORKS




  
Knowledge Base Question Answering There is a large history of work that has looked into answering questions using a knowledge graph <cit.>. The approaches can broadly be broken into those that attempt to parse the question into a logical form <cit.>, and information retrieval approaches <cit.>



  
Knowledge Enhancement. Recently, researchers have looked into enhancing pretrained language models and LLMs using knowledge graphs. Many works have focused on incorporating knowledge graph data into LLMs through training, often with architectural changes to the model <cit.>. Some of these have found success using specialized graph encoder layers <cit.>. Some have sought to mix this process with pretraining of the language model <cit.>. The limitations of including KGs during LLM training are: the models are incapable of incorporating KG updates without retraining, are not able to change KG source (e.g., a domain-specific one), and may have low explainability resulting from learning knowledge in model weights rather than explicit retrieval. In addition, these methods add complexity to the training process, increase cost, and do not work with LLM without access to model weights. As a result, scaling these methods has often been seen as too risky. 

Other methods treat a KG as a more complete source of truth and use the LLM to generate structured queries for the KG. <cit.> and <cit.> teach the LLM to generate a logical KG query which then returns matching entities from the KG. These methods share similarities with the semantic parsing based approaches mentioned earlier. However, by returning a logical query for the KG, these methods lose the reasoning and commonsense abilities of the LLM. 

Some approaches have looked at injecting KG or knowledge base data into LLM prompts. Many methods do a single round of retrieval based on the query <cit.> using a retrieval mechanism, such as dense passage retrieval <cit.>. These methods cannot answer more complex multi-hop questions as the initial retrieval is unlikely to contain the secondary or tertiary information that will ultimately be required. Later methods have made the retrieval process iterative. E.g. FLARe <cit.> uses prediction of the upcoming sentence to retrieve relevant documents and regenerates until the sentence contains high-confidence tokens. In <cit.> multi-round QA format is used for multi-round retrieval. ReAct <cit.> does multiple rounds of thoughts and actions to query a text-based knowledge base. 
<cit.> repeatedly interfaces with a KG. In other parallel works, <cit.> explore methods for building KG context for LLMs using path and neighborhood based search methods. The above methods do not incorporate tree search or forms of advanced reasoning beyond the LLM's innate ability. They cannot explore multiple reasoning paths or solutions and cannot backtrack if making a mistake. This results in lesser capabilities. Additionally, none of the above methods explore reasoning over multiple KGs. 



  
Multiple Knowledge Base QA. Some works studied QA on questions derived from differing domains. These works learn to pick between different domain-specific QA models, each trained on a single knowledge base <cit.>. Besides requiring training, such systems cannot answer questions requiring the synthesis of information from different domains (MusicBrainz-x-Wikidata). 



§ CONCLUSION


Tree-of-Traversals is a powerful algorithm that enables LLMs to utilize KG reasoning with zero examples, no schema dependency, and no training. We demonstrate its efficacy experimentally on multiple LLMs and datasets. We hope continues to be developed and see value in studying integration with personalized user KGs as well as with other domain-specific datasets. 



§ LIMITATIONS


Tree-of-Traversals is slower than simpler retrieval alternatives. Improving the value function would reduce the number of LLM and KG API calls by avoiding incorrect paths along the tree to explore. Other engineering solutions, such as hosting graph servers, could be implemented to accelerate LLM and KG access. In terms of the token cost, the KG text representation is token efficient compared to many retrieval methods using unstructured knowledge bases as the latter tend to maximize usage of the LLM context window. However, compared to non-retrieval baselines, there is significant increase in token usage cost. 

The types of KG questions that can be answered are limited by the context window, search depth, and LLM's reasoning ability. For instance, a very large aggregation (`How many mountains have elevation above 3500 meters?') would require more entities than what fits in the LLM context. Future work could look at adding other actions to the ASM, such as aggregations, that could distill the local KG into more relevant information. 

EM-in is also an imperfect metric. False negatives can arise from discrepancies in date and number formatting, text formatting, and aliasing. False positives can arise in cases in which the model does not definitively answer but includes the correct answer in its response.



§ ETHICAL IMPACT


We do not anticipate to introduce new areas of risk but it may have unstudied effects on existing risks of LLMs or KGs. We highlight the following areas. (i) gives new capabilities to LLMs after training. While positive in terms of accuracy, we have not evaluated its effect on wider safety metrics. (ii) Our evaluation is limited to only English versions of KGs and datasets. should be evaluated in other languages to ensure consistent and fair experience. (iii) We have not performed analysis under misleading or deceptive knowledge graphs. Using a publicly modifiable knowledge graph does come with the risk that information could be deceptively changed.

In terms of positive ethical impact, making this research public democratizes access to knowledge-augmented LLMs as this method does not require a large training investment or owning the language model to build and customize. 



acl_natbib






§ MUSICBRAINZ X WIKI





Additional information on the creation of the MusicBrainz x Wiki dataset.



 §.§ Creation steps


The dataset was created by annotators in a semi-automatically supported manner. We developed an automated tool to support the following process: 

    
  * Find linking entities that are present in both MusicBrainz and Wikidata of various types (artist, label, place, event, etc.)
    
  * Find related entities that could be used in a question to perfectly identify the entities above (e.g., The place owned by the Detroit Tigers → Tiger Stadium)
    
  * Find related entities that could be used in a question to ambiguously identify a linking entity, so that a qualifier could be added to disambiguate the linking entity.
    
  * Count the number of edges in Wikidata or MusicBrainz for each relationship type an entity has.


We use these tools to create an initial set of multi-hop questions falling into specific reasoning categories. We then have human annotators check each question for sensibility (grammatically correct and can be definitively understood) and ambiguity (one right answer in the KGs). If possible, the question is reworded or fixed to not contain ambiguities. We are then left with a final set of 109 questions. 



 §.§ Composition of Questions


The composition of questions for Musicbrainz x Wiki can be found in table <ref>.



 §.§ Annotators


The annotators consisted of approximately 10 research scientists and engineers, all English speaking and based in the United States.



§ IMPLEMENTATION DETAILS




 §.§ Diversity Oversampling.

One problem with Tree-of-Thoughts is that when there are constrained options to choose from, the LLM becomes repetitive and fails to produce diverse outputs <cit.>. Their solution used a "propose prompt" with additional instructions for proposing multiple distinct thoughts. However, this adds significant complexity to the prompt for purposes that are only tangential to the task being completed. A simpler approach we employ is to generate diverse actions by oversampling. Specifically, if the branching factor is k, we sample 2*k possible actions from the LLM and then extract the first k unique ones to use. We only employ this for selecting-entities and selecting-relation as the action choices are limited and diversity is required. This change does not noticeably affect the computation cost of as multiple actions are sampled for the same prompt and the average prompt input size (100s-1000s of tokens) is generally much larger than the average generation length for these actions (<20 tokens). Latency remains similar as multiple actions are sampled in parallel. With this, we consistently generate diverse options when selecting-entities or selecting-relation. 



 §.§ Hyperparameters


does not have a huge number of hyperparamters. Most of the hyperparameters were chosen analytically, though some were chosen based on preliminary experiments. We share some information on choosing them.

τ was chosen analytically in conjunction with the prompt for  on the done state (Appendix <ref>). Chosen to be 0.8 so that only stops when it is confident that the answer is supported by the KG. 

The temperature was chosen to be 1.0 to promote diversity. If implementing with other models which can have temperature >1.0, we recommend hyperparameter tuning or using a sample prompt of each type and analytically choosing the temperature such that diverse but sensible actions are output. 

The branching factor k=3 was chosen based on small scale experimentation. Larger k makes the model spend longer searching, taking longer time and increasing expenses. Smaller k did not generate as much diversity and options. It is worth noting that even with k=3 and diversity sampling, sometimes the actual number of unique outputs is less than 3. 

The max depth was chosen analytically to be minimal while enabling answering of most questions. Some questions in QALD-10 and MusicBrainz-x-Wikidata would require a greater depth to answer. 

The maximum expansions cutoff was chosen purely to ensure the model does not hang on a single question too long. The actual value is somewhat arbitrary and has minimal impact on accuracy.



 §.§ KG Interface


We present some additional details of the KG Interface. 

The  function is a three step process. First a LLM call and prompt is used to extract named entities. Second, for each extracted entity candidates are searched for using the KG API. Finally, another LLM call is used to match the extracted entity to the best candidate. 

 and  are implemented exclusively using the respective KG API. For Wikidata, these are each one or two SPARQL query (forward and reverse edges). For MusicBrainz, they each require multiple API calls as the endpoints are separated for different entity types. E.g. there is a separate endpoint for Artists and Recordings. 




§ VALUE FUNCTION PROFILES



We profile each value model's value function to assess their characteristics. Figure <ref> shows the individual answer value vs accuracy profiles for each model on 2WikiMultiHop. All models have a positive correlation between the answer value and the answer's accuracy. However, as is to be expected, there are differences with Llama2-13b having the lowest correlation and Claude-Instant having the highest correlation. 

Figure <ref> only look at the value function response for final answer states as those are the only ones with a directly associated accuracy score. While accurately evaluating answer states is more important, we also want to see some useful signal from the values for intermediary states. Figure <ref> shows the average value for nodes on the search path to the answer, separated by whether the answer scores as correct or not. We would expect actions that ultimately lead to a correct answer to have higher values than those that lead to incorrect ones. We clearly see this with the correct distribution being further right than the incorrect distribution. This is least pronounced for Llama2-13b which is to be expected. We also note that the distributions shape is different between the Llama2 models and Claude-Instant. 

These results indicate that larger, stronger models, will likely see more improvement in the value function, and thus, more utility from . 







§ PROMPT ENGINEERING



Some prompt re-engineering was required when switching between models for both our methods and ReAct. This engineering was limited primarily to formatting of responses. For example, Llama2 would start responses with "So the next action should be" while we required the action to start with "THINK", "EXPAND_KG" or "ANSWER". Instead we incorporated “So the next action should be:” into the prompt so that the output began with the chosen action. For  we added “You have selected the following entities to expand:”. For  we added “I suggest selecting property:”. These issues could also be solved by changing the parsing of the presented prompts. 



§ FALLBACK TO CHAIN-OF-THOUGHT (→COT)


We fallback to Chain-of-Thought for the ReAct baseline when either no answer is provided after depth 7 or when one of the following phrases appears in the answer: 'determine', 'unable', 'cannot', 'unknown', 'unsure', or 'not possible'. While these certainly could appear in intentional answers, we find that they occur consistently and exclusively in non-answer style responses. 




§ PROMPTS


The following pages contains the prompts used in Tree-of-Traversals. The dynamic components that change based on the query, KG state, action history are shown in color.



 §.§ Prompt for default action state



        
Original Query: 
    |Who is Bob Dylan's maternal grandmother?|
    
Knowledge Graph Entities:
    |Q392: Bob Dylan - American singer-songwriter|
    |Q62519478: Beatrice Stone|
Knowledge Graph Edges:
    |Bob Dylan:|
        |mother:|
            |Beatrice Stone|

Previous Actions: 
|EXPAND KG: I should search for the mother of Bob Dylan|
|SELECT ENTITIES: Q392|
|SELECT RELATION: P25 - has mother|

You are a superintelligent AI equipped with the ability to search a knowledge graph for definitive, up-to-date answers. Your task is to interface with the knowledge graph in order to answer the above query. You will be able to expand the knowledge graph until you have found the answer. Think in detail before acting or answering.

Available actions:
'THINK' - Generate relevant thoughts to solving the problem. This could include recalling well known facts from memory.
    e.g.,
        THINK: I should search for the movies directed by...
        THINK: I know that Biden is American, therefore...
        THINK: I see that John Cena is in both The Independent and Vacation Friends, therefore...
'EXPAND_KG' - Search for edges of entities in the knowledge graph using an external API. This is a useful function for getting to the correct answer.
    e.g., EXPAND_KG: I should search for the country of origin of Jonathon Taylor
'ANSWER' - Generate the final answer once the problem is solved. Just state the best answer, do not output a full sentence.
    e.g., 
        ANSWER: No
        ANSWER: Harry Potter
        ANSWER: [Harry Potter, Ron Weasley, Hermione Granger]






 §.§ Prompt for selecting-entities action state




        
Original Query: 
    |Who is Bob Dylan's maternal grandmother?|
    
Knowledge Graph Entities:
    |Q392: Bob Dylan - American singer-songwriter|
    |Q62519478: Beatrice Stone|
Knowledge Graph Edges:
    |Bob Dylan:|
        |mother:|
            |Beatrice Stone|

Previous Actions: 
|EXPAND KG: I should search for the mother of Bob Dylan|
|SELECT ENTITIES: Q392|
|SELECT RELATION: P25 - has mother|
|EXPAND KG: I should search for the mother of Beatrice Stone|

Current task
EXPAND_KG takes two parameters. The first is the entity or entity group to get more information about. Select which entity or entities from the KG to expand.
Provide the QIDs. Options include |[Q392, Q62519478]|
SELECT ENTITIES:






 §.§ Prompt for selecting-relation action state




        
Original Query: 
    |Who is Bob Dylan's maternal grandmother?|
    
Knowledge Graph Entities:
    |Q392: Bob Dylan - American singer-songwriter|
    |Q62519478: Beatrice Stone|
Knowledge Graph Edges:
    |Bob Dylan:|
        |mother:|
            |Beatrice Stone|

Previous Actions: 
|EXPAND KG: I should search for the mother of Bob Dylan|
|SELECT ENTITIES: Q392|
|SELECT RELATION: P25 - has mother|
|EXPAND KG: I should search for the mother of Beatrice Stone|
|SELECT ENTITIES: Q62519478|

Your current task is to select the property (PID) to expand along for the selected entities.
The selected entities are: |[Q62519478]|
The options of properties to choose from are:
|P31 - has instance of|
|P21 - has sex of gender|
|P27 - has country of citizenship|
|P735 - has given name|
|...|

Select exactly one property (PID e.g., P10) from those listed above
SELECT PROPERTY:






 §.§ General Prompt for 




        
Original Query: 
    |Who is Bob Dylan's maternal grandmother?|
    
Knowledge Graph Entities:
    |Q392: Bob Dylan - American singer-songwriter|
    |Q62519478: Beatrice Stone|
Knowledge Graph Edges:
    |Bob Dylan:|
        |mother:|
            |Beatrice Stone|

Previous Actions: 
|EXPAND KG: I should search for the mother of Bob Dylan|
|SELECT ENTITIES: Q392|
|SELECT RELATION: P25 - has mother|
|EXPAND KG: I should search for the mother of Beatrice Stone|
|SELECT ENTITIES: Q62519478|

Your current task is to evaluate the above knowledge graph and action history.
Based on the original query, the current knowledge graph, and the action history, give the likelihood that the model will correctly answer the question. 
If the most recent action provided information towards the goal and followed the preceding thought, give a high score. 
If the last action was unhelpful, give a low score. 

The output should be a number between 0 and 1 with one decimal. Do not output anything else. 

RATING [0.0-1.0]:






 §.§ Prompt for  on answer (done)




        
Original Query: 
    |Who is Bob Dylan's maternal grandmother?|
    
Knowledge Graph Entities:
    |Q392: Bob Dylan - American singer-songwriter|
    |Q62519478: Beatrice Stone|
    |Q62519478: Florence Sara Stone|
Knowledge Graph Edges:
    |Bob Dylan:|
        |mother:|
            |Beatrice Stone|
    |Beatrice Stone:|
        |mother:|
            |Florence Sara Stone|

Provided answer: |The answer is Florence Sara Stone.|

Your task is to score the correctness of the provided answer based on the original query, and the knowledge graph.
Give a pessimistic score from 0.0 to 1.0 on how likely the answer is to be correct. 
0.0 if definitely wrong
0.0 if unable to answer based on the knowledge graph
0.5 if unsure
0.7 for probably correct but not confirmed in knowledge graph
1.0 for definitely correct and confirmed in knowledge graph.

Give reasoning to get to the correct answer. Then provide a score.
e.g., 
Reasoning...
So the score for the provided answer should be...





 Paper Title: DriveArena: A Closed-loop Generative Simulation Platform for Autonomous Driving.























-1.0ex
    < g r a p h i c s >
 : A Closed-loop Generative Simulation Platform for Autonomous Driving
    Xuemeng Yang^1, *   Licheng Wen^1, *    Yukai Ma^2, 1, *   Jianbiao Mei^2, 1, *   Xin Li^3,5, *, † 

 
 
  Tiantian Wei^1, 4, *   Wenjie Lei^2, †   Daocheng Fu^1   Pinlong Cai^1   Min Dou^1

 

  Botian Shi^1,   Liang He^5   Yong Liu^2,   Yu Qiao^1 




^1 Shanghai Artificial Intelligence Laboratory   ^2 Zhejiang University   ^3 Shanghai Jiao Tong University 




^4 Technical University of Munich   ^5 East China Normal University




Project Page: <https://pjlab-adg.github.io/DriveArena/>

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

 
 
 
* Equal contribution,    Corresponding author 

† Work performed during internships at Shanghai AI Laboratory



§ ABSTRACT 
This paper presented , the first high-fidelity closed-loop simulation system designed for driving agents navigating in real scenarios.  features a flexible, modular architecture, allowing for the seamless interchange of its core components: , a traffic simulator capable of generating realistic traffic flow on any worldwide street map, and , a high-fidelity conditional generative model with infinite autoregression. This powerful synergy empowers any driving agent capable of processing real-world images to navigate in 's simulated environment. The agent perceives its surroundings through images generated by  and output trajectories. These trajectories are fed into , achieving realistic interactions with other vehicles and producing a new scene layout. Finally, the latest scene layout is relayed back into , perpetuating the simulation cycle. This iterative process fosters closed-loop exploration within a highly realistic environment, providing a valuable platform for developing and evaluating driving agents across diverse and challenging scenarios.  signifies a substantial leap forward in leveraging generative image data for the driving simulation platform, opening insights for closed-loop autonomous driving.

Code will be available soon on GitHub: <https://github.com/PJLab-ADG/DriveArena>







§ INTRODUCTION



Autonomous driving (AD) algorithms have advanced rapidly in recent decades <cit.>, progressing from modular pipelines <cit.> to end-to-end models <cit.> and knowledge-driven methods <cit.>. Despite demonstrating outstanding performance across various benchmarks, significant challenges persist in evaluating these algorithms on replayed open-loop datasets, obscuring their real-world efficacy.
Public datasets <cit.>, while offering realistic driving data with authentic sensor inputs and traffic behavior, are inherently biased towards simple straight-ahead scenarios. In such cases, an agent can achieve seemingly good performance by merely maintaining its current state, complicating the assessment of actual driving capabilities in complex situations. 
Furthermore, the agent's current decision does not affect execution or subsequent decisions in the open-loop evaluation, which prevents it from reflecting cumulative errors in real-world driving scenarios.
Additionally, the static nature of recorded datasets, where other vehicles cannot react to the ego vehicle's behavior, further hinders the evaluation of AD algorithms in dynamic, real-world conditions.

As illustrated in Figure <ref>, we analyze existing AD methods and platforms, revealing that most of them are inadequate for a high-fidelity closed-loop simulation.
Ideally, as an aspect of embodied intelligence, agents should be evaluated in a closed-loop environment, where other agents react to the actions of the ego vehicle, and the ego vehicle receives changed sensor input accordingly. However, existing simulation environments either cannot simulate sensor inputs <cit.> or have a significant domain gap with the real world <cit.>, making it difficult to seamlessly integrate algorithms into the real world, thus posing a huge challenge for closed-loop evaluation. 
We believe that the simulator should not only closely reflect the visual and physical aspects of the real world, but also promote the continuous learning and evolution of the model within an exploratory closed-loop system for adapting to diverse complex driving scenarios. To achieve this goal, it is imperative to establish a high-fidelity simulator that complies with physical laws and supports interactive functionalities. 

Therefore, we present , a pioneering closed-loop simulator based on conditional generative models for training and testing driving agents. 
Specifically,  offers a flexible platform that can be integrated with any camera-input driving agent. It adopts a modular design and naturally supports iterative upgrades of each module.  consists of a  that manages traffic flow and a  based on auto-regressive generation.
 can generate realistic interactive traffic flow on any road network worldwide, while  is a high-fidelity conditional generative model with infinite autoregression. The driving agent should make corresponding driving actions based on the images generated by , and feed them back to  to update the status of vehicles in the environment. The new scene layout will be returned to  for a new round of simulation. This iterative process realizes the dynamic interaction between the driving agent and the simulation environment. The specific contributions are as follows:

 
    
  * High-fidelity Closed-loop Simulation: We propose the first high-fidelity closed-loop simulator for autonomous driving,  , which can provide realistic surround images and integrate seamlessly with existing vision-based driving agents. It can closely reflect the visual and physical properties of the real world, enabling agents to continuously learn and evolve in a closed-loop manner and adapt to various complex driving scenarios.
    
    
  * Controllability and Scalability: 
    Our  can dynamically control the movement of all vehicles in the scenarios and feed the road and vehicle layouts into , which utilizes a conditional diffusion framework to generate realistic images in a stable and controllable manner. Additionally,  supports simulation using road networks from any city worldwide, enabling the creation of diverse driving scenario images with varying styles. 
    
    
  * Modularized Design: The Driving Agent,  and  communicate via network interfaces, enabling a highly flexible and modular framework. This architecture allows each component to be replaced with different methods without requiring specific implementations. Functioning as an arena for these players,  facilitates comprehensive testing and improvement of both vision-based autonomous driving algorithms and driving scene generative models.






§  FRAMEWORK


As illustrated in Figure <ref>, the framework of our proposed  comprises two key components: a  functioning as the backend physical engine and a  serving as the real-world image renderer. Unlike conventional approaches,  does not rely on pre-built digital assets or reconstructed 3D road models. Instead, the  adapts to road networks of any city in OpenStreetMap (OSM) format <cit.>, which can be directly downloaded from the Internet. This flexibility enables closed-loop traffic simulations on diverse urban layouts.

The  receives ego trajectories output by the autonomous driving agent and manages the movement of all background vehicles. Unlike world model approaches <cit.> that rely on diffusion models for both image generation and vehicle movement prediction, our  utilizes explicit traffic flow generation algorithms <cit.>. This approach enables the generation of a wider range of uncommon and potentially unsafe traffic scenarios, while also facilitating real-time collision detection between vehicles.

 generates realistic camera images that precisely correspond to the 's output. It also allows for user-defined prompts to control various elements of the generated images, such as street view style, time of day, and weather conditions, enhancing the diversity of the generated scenes. Specifically, it employs a diffusion-based model that utilizes the current map and vehicle layouts as control conditions to produce surround-view images. These images serve as input for end-to-end driving agents. Given 's closed-loop architecture, the diffusion model is required to maintain both cross-view and temporal consistency in the generated images.

The generated multi-view images of the current frame are fed into the end-to-end autonomous driving agents, which can output the ego vehicle's movement. The planned ego trajectory is subsequently sent to  for the next simulation step. The simulation concludes when the ego vehicle either successfully completes the entire route, crashes, or deviates from the road. Upon completion,  performs a comprehensive evaluation process to assess the driving agent's capabilities.

It is noteworthy that  employs a distributed modular design. The , , and AD agent communicate via network using standardized interfaces. Consequently,  does not mandate specific implementations for the  or the AD agent. Our framework aims to function as an “arena” for these “players”, facilitating comprehensive testing and improvement of both end-to-end autonomous driving algorithms and realistic driving scene generative models.



§ METHODOLOGY



Following the  framework outlined above, we have implemented a preliminary version of . In this section, we elaborate on the implementation of each module: , , and AD agent, while describing necessary details that were not previously mentioned. At the end of this section, we present both the open-loop and closed-loop evaluation metrics for AD agents in .



 §.§ 

Most existing realistic driving simulators <cit.> rely on limited layouts from public datasets, lacking diversity for dynamic environments. To address these challenges, we utilize LimSim <cit.> as the underlying  to simulate dynamic traffic scenarios and generate road and vehicle layouts for subsequent environment generation. LimSim also provides a user-friendly front-end GUI, which directly displays the BEV map and results from  and the driving agent.

Our  enables interactive simulations of multiple vehicles in traffic flow, including comprehensive vehicle planning and control. We adopt a hierarchical multi-vehicle decision-making and planning framework, which jointly makes decisions for all vehicles within the flow and reacts promptly to the dynamic environment through a high-frequency planning module <cit.>. The framework also incorporates a cooperation factor and trajectory weight set, introducing diversity to autonomous vehicles in traffic at both social and individual levels.

Furthermore, our dynamic simulator supports various custom HD maps of any city from OpenStreetMap, facilitating the construction of diverse road graphs for convenient simulation. 
The  controls the movement of all background vehicles. For the ego vehicle, we provide two distinct simulation modes: open-loop and closed-loop. In closed-loop mode, the driving agent performs planning for the ego vehicle, and  uses the agent-outputted trajectory to control the ego vehicle accordingly. In open-loop mode, the trajectory generated by the driving agent is not actually used to control the ego vehicle; instead,  maintains control in a closed-loop manner. The details of these two modes are further elaborated in Section <ref>.



 §.§ 



Unlike recent autonomous driving generation methods <cit.> that use Neural Radiance Fields (NeRF) and 3D Gaussian Splatting for environment reconstruction from logged video, we design a diffusion-based . It utilizes control conditions of the map and vehicle layouts from the  to generate geometrically and contextually accurate driving scenarios. Our framework shares several advantages: (1) Better controllability. The generated scenes can be controlled by scene layouts from , textual prompts, and reference images to capture different weather conditions, lighting, and scene styles. (2) Better scalability. Our framework can adapt to various road structures without the need to model the scene in advance. In theory, we support the generation of driving scenes for any city in the world by leveraging layouts from OpenStreetMap.

We illustrate our diffusion-based  in Figure <ref>. Built upon the stable diffusion pipeline <cit.>,  utilizes an effective condition encoding module that accepts a variety of conditional inputs including map and vehicle layouts, text descriptions, camera parameters, ego poses, and reference images to generate realistic surround-view images. Considering the importance of ensuring synthesis scene consistency across different views and time spans for driving agents, we integrate a cross-view attention module, inspired by <cit.>, to maintain coherence across different views. Additionally, we adopt an image auto-regressive generation paradigm to enforce temporal consistency. This approach enables  to not only maximally maintain the temporal consistency of the generated videos, but also generate videos of arbitrary length in an infinite stream, which provides great support for autonomous driving simulation.

Condition encoding. Previous work <cit.> applied BEV layout as conditional input to control the output of the diffusion model, which increased the difficulty of the network in learning to generate geometrically and contextually accurate driving scenes. In this work, we present a new condition encoding module to introduce more guidance information, which helps the diffusion module generate high-fidelity surround images. Specifically, in addition to encoding camera poses for each view, text descriptions, 3D object bounding boxes, and BEV map layouts using a condition encoder similar to <cit.>, we also explicitly project the map and object layouts onto each camera view to generate layout canvases for more accurate lane and vehicle generation guidance. 
Specifically, the text embedding e_text is obtained by encoding the text descriptions with the CLIP text encoder <cit.>. The parameters 𝐏 = {𝐊∈ℝ^3 × 3, 𝐑∈ℝ^3 × 3, 𝐓∈ℝ^3 × 1} of each camera and the 8 vertices of the 3D bounding boxes are encoded to e_cam and e_box by Fourier embedding <cit.>, where 𝐊, 𝐑, 𝐓 represent camera intrinsic, rotations and translations respectively. The 2D BEV map grid uses the same encoding method as in  <cit.> to get embedding e_map. Then, each category of the HD maps and the 3D boxes is projected onto the image plane to obtain the map canvas and box canvas, respectively. These canvases are concatenated to create the layout canvas. The final feature e_layout can be obtained by encoding the layout canvas by the conditional encoding network <cit.>.

Moreover, we introduce a reference condition to provide appearance and temporal consistency guidance. During training, we randomly extract a frame from the past 𝐋 frames as a reference frame and use the pre-trained CLIP model <cit.> to extract reference features e_ref from the multi-view images. The encoded reference features imply semantic context and are integrated into the conditional encoder through the cross-attention module. In order to make the diffusion model aware of the motion changes of the ego-vehicle, we also encode the ego-pose relative to the reference frame into the conditional encoder to capture the motion change trend of the background. The relative pose embedding e_rel is encoded by Fourier embedding.
By incorporating the above control conditions, we can effectively control the generation of surround images.

Auto-regressive generation. To facilitate online inference and streaming video generation while maintaining temporal coherence, we have developed an auto-regressive generation pipeline. Specifically, during the inference phase, the previously generated images and the corresponding relative ego pose are used as reference conditions. This approach guides the diffusion model to generate current surround images with enhanced consistency, ensuring a smoother transition and coherence with the previously generated frames.

What we designed in this paper is just a simple implementation of the . We also verify that extending the auto-regressive generation to a multi-frame version (using multiple past frames as references and outputting multi-frame images) and adding additional temporal modules can improve temporal consistency.



 §.§ Driving Agent

Recent works <cit.> have demonstrated the challenges in justifying the planning behavior of driving agents through open-loop evaluation on public datasets <cit.>, primarily due to the simplistic nature of driving scenarios presented. While some studies <cit.> have conducted closed-loop evaluations using simulators like CARLA <cit.>, discrepancies such as appearance and scene diversity persist between these simulations and the dynamic real world. To bridge this gap, our  provides a realistic simulation platform with the corresponding interfaces for camera-based driving agents <cit.> to perform more comprehensive evaluations, including both open-loop and closed-loop testing. Moreover, by changing the input conditions, such as the road and vehicle layouts,  could generate corner cases and facilitate these driving agents' evaluation on out-of-distribution scenarios. Without loss of generality, we select a representative end-to-end driving agent, namely UniAD <cit.>, to conduct both open-loop and closed-loop testing in our . UniAD utilizes surround images to predict motion trajectories for the ego vehicle and other agent vehicles, which can be seamlessly integrated with the API of our dynamic simulator for evaluation. Furthermore, the perceptual outputs, such as 3D detection and map segmentation, contribute to enhancing the validation of realism in our environment generation.





 §.§ Ego Control Modes and Evaluation Metrics



 inherently supports “closed-loop” simulation mode of driving agents. That is, the system adopts the trajectory output by the agent at each timestep, updates the ego vehicle's state based on this trajectory, and simulates the actions of background vehicles. Subsequently, it generates multi-view images for the next timestep, thus maintaining a continuous feedback closed-loop.
Additionally, recognizing that some AD agents may be unable to perform long-term closed-loop simulation during the development process,  also supports the “open-loop” simulation mode. In this mode, the  will take over the control of the ego vehicle, while the trajectory output by the AD agent is recorded for subsequent evaluation.

In both open-loop and closed-loop modes, it is crucial to comprehensively evaluate AD agent performance from a results-oriented perspective. Drawing inspiration from NAVSIM <cit.> and the CARLA Autonomous Driving Leaderboard <cit.>,  adopts two evaluation metrics: PDM Score (PDMS) and Arena Driving Score (ADS).


PDMS, initially proposed by NAVSIM <cit.>, evaluates the trajectory output at each timestep. We adhere to the original definition of PDMS, which aggregates the following sub-scores:

    PDMS_t=   (∏_m ∈{NC, DAC} score _m)_penalties×
       (∑_w ∈{EP, TTC, C}weight_w  × score _w/∑_w ∈{EP, TTC, C}weight_w)_weighted average .

where the penalties include the drive with no collisions (NC) with road users and drivable area compliance (DAC), as well as the weighted average including ego progress (EP), time-to-collision (TTC), and comfort (C). 
We implement minor modifications tailored to : in score _NC, we do not differentiate “at-fault" collisions, and for score _EP, we utilize the 's Ego path planner as the reference trajectory instead of the Predictive Driver Model.
At the end of the simulation, the final PDM Score is averaged across all simulation frames.

    PDMS = Σ_t=0^TPDMS_t/T∈ [0,1]


For open-loop simulations, PDMS serves directly as the evaluation metric for AD agents. However, for driving agents operating under the “closed-loop" simulation mode, we employ a more comprehensive metric called Arena Driving Score (ADS), which combines the trajectory PDMS with route completion:

    ADS = R_c×PDMS

where R_c∈ [0,1] represents route completion, defined as the percentage of the route distance completed by an agent. Given that “closed-loop" simulations terminate upon agent collision with other road users or deviation from the road, ADS provides a suitable metric for differentiating agents' driving safety and consistency.




§ EXPERIMENTS

 



 §.§  Setups

Dataset. 
For , we use the nuScenes <cit.> dataset for training. Following the official configuration, we employ 700 scenes for training and 150 for validation. We focus on four road categories (lane boundary, lane divider, pedestrian crossing, and drivable area) and ten object categories. The nuScenes dataset contains data collected from four different cities, covering various light and weather conditions, including daytime, night, sunny, cloudy, and rainy scenarios, enabling  to conditionally imitate diverse appearances. We additionally annotated each scene using GPT-4V, providing detailed scene descriptions that include elements like time, weather, street style, road structure, and appearance. These descriptions serve as text prompt conditions.






Model Setup. 
The model is initialized with the pre-trained Stable Diffusion v1.5 <cit.>, with only the newly added parameters being trained. For various conditions, except for the encoding of reference images and text prompts, the encoders for other conditions are randomly initialized and trained from scratch. These conditions are then integrated into the UNet using a randomly initialized ControlNet <cit.> to control the denoising process.

Training and Inference. 
To utilize the reference images and achieve temporal correlation, we employ ASAP <cit.> to generate 12Hz interpolated annotations and crop them into image clips of length 𝐋=7. During training, we use the last frame of each clip as the current frame, select any frame from the clip as the reference frame, and calculate the relative pose between them to model the motion trend of the background. Accordingly, the surround images corresponding to the reference frame are input to the network as reference images. During inference, the generated result of the previous frame is used as the current reference images, enabling unlimited length generation. 
The experiment is conducted on 8 NVIDIA A100 (80GB) GPUs with a batch size of 4×8 and 200K iterations of training. The AdamW optimizer is used with a learning rate of 1e-4.
The network follows the same image resolution (224×400) as MagicDrive, and when input to the driving agent, it will be upsampled to the original image size of nuScenes (900×1600) through a super-resolution algorithm <cit.>.



 §.§  Setups


Operating Frequencies.
In our experiments, the  operates at a frequency of 10Hz, while the control frequency is set to 2Hz. This configuration results in the  sending the current layout to  every 0.5 simulation seconds, requesting surround images. These images are then forwarded to the driving agent, which predicts and plans the subsequent trajectory for the ego vehicle. The , , and driving agent communicate via HTTP protocol, enabling deployment across different servers.

Simulation Modes.
As detailed in Section <ref>, we implement two simulation modes. In the open-loop mode, all vehicles, including the ego vehicle, are controlled by  itself. The driving agent can predict the ego vehicle's trajectory, but its trajectory is not actually executed. In the closed-loop mode, the ego vehicle is controlled by the driving agent, and the simulation terminates if it crashes with other vehicles or leaves the road.


Supported Maps.
Currently,  supports four different maps, which are: , , , and . The first two maps closely resemble the corresponding areas in the nuScenes dataset, while the last one replicates the road network of the Town05 map in the CARLA simulator. Notably,  can download road network data for any area directly from OpenStreetMap and perform simulations, enabling  to simulate the road network of almost any city worldwide. OpenstreetMap also accepts customized operations, and users can draw the desired road network structure for simulation testing.




 §.§  Fidelity Validation

To assess the sim-to-real gap between our generated images and the original nuScenes images, we employ UniAD <cit.> as an evaluator. We generate videos for 150 scenes based on the original layout provided by the nuScenes validation set with 2Hz. For comparative analysis, we set MagicDrive as the baseline method and use its official codes and checkpoints for inference. Subsequently, UniAD is performed on these images to compute various metrics, including 3d object detection, BEV map segmentation, and planning. The results are summarized in Table <ref>. It shows that all our indicators are higher than the baseline method, and a few indicators even surpass the performance on the original nuScenes. Furthermore, it demonstrates our model's superior capability to accurately respond to control signals and strictly adhere to input conditions. These findings establish a solid foundation for using our generator as a reliable simulator.









 §.§ Visualization

Controllability.
In this section, we will comprehensively demonstrate the controllability of the model from various dimensions, including the control of lighting and weather, the fit of object boxes and maps, change of street style, and consistency over long periods of time.

We demonstrate the impact of the reference image on the generated image, as shown in Figure <ref>. We randomly select one frame of images from the nuScenes dataset as reference images and choose three scenes from OpenStreetMap and Carla. We perform inference on them with  respectively. It can be seen that the source and style of the road network are very different from the scope of the original nuScenes dataset. The pictures show that the generated vehicles and road networks conform closely to control conditions, demonstrating strong control capabilities. The style and weather of the generated pictures can also be consistent with the reference images. In other words, besides maintaining image generation continuity through reference images, we can also regulate image style accordingly.

Figure <ref> presents images generated using different text prompts and reference images on the same road network. Each set of images portrays the surrounding scenery at intervals of 8.5 seconds and 24 seconds respectively, with the layout projected on the image. The images clearly 
illustrate that the road structure and vehicles strictly adhere to the given control conditions while maintaining excellent consistency in the surround view. In addition, the four sets of images exhibit significant differences in weather and lighting and can maintain their own styles during the continuous iteration process.


Scalability.
The  can accept any map downloaded from OpenStreetMap and seamlessly connect to the Carla road network. Combined with Dreamer's excellent following capability, the entire framework demonstrates robust scalability. The specific results are shown in Figure <ref>. We used both MagicDrive and our  to generate realistic images on the same Carla road network, with the corresponding lane lines projected onto the images. The road style in Carla differs significantly from that of nuScenes. It is rare to encounter roads with such large curvature and such wide roads in nuScenes. Consequently, the performance of MagicDrive, which is based on the nuScenes BEV map, is slightly inferior in these conditions. As indicated by the yellow arrow, MagicDrive struggles to generate curved roads and fit wide roads accurately. , however, can produce reasonable pictures that follow the road structure.


We also demonstrate additional cases using data from the nuPlan dataset to validate the scalability. The nuPlan data originates from cities different from nuScenes and features varying camera numbers and parameters. We select 6 cameras with a similar layout to the nuScenes dataset, and nuPlan's camera parameters are employed to project object boxes and lane lines onto corresponding images as control conditions. As shown in Figure <ref>,  fully trained on nuScenes adeptly adheres to these conditions, generating coherent images when deployed in new cities and even with novel camera configurations.


















 §.§ Open-loop and Closed-loop Experiments


In this section, we adopt the prevailing end-to-end autonomous driving method UniAD <cit.> as the driving agent to test both the open-loop and closed-loop performance within the  framework. We utilized UniAD's open-source code and pre-trained weights without additional training. UniAD operates at 2Hz, outputting a trajectory of 6 path points over the next 3 seconds.  further interpolates this to a 10Hz trajectory.

Open-loop Evaluation.
We first assess UniAD's performance in 's open-loop mode.
UniAD is evaluated on three scenarios: 1) the original nuScenes image sequences; 2) -generated nuScenes image sequences, where the vehicles' trajectory remains identical to nuScenes ground truth, but surround images are replaced with -generated ones; and 3) 's own simulation sequences (i.e., 's open-loop mode).
Our evaluation metrics consist of the PDM Scores and its sub-scores, as detailed in Section <ref>. Additionally, we evaluate trajectories driven by human drivers in nuScenes as the human driver performance. Detailed results are presented in Table <ref>.

The results reveal that while UniAD performs optimally on the original nuScenes sequence with a PDMS metric of 0.91, the -generated sequence surprisingly achieves a PDMS of 0.902, representing a metric drop of less than 1%. We attribute this to both the high fidelity of our -generated images and UniAD's strong dependence on ego states, as corroborated by <cit.>.

In 's open-loop mode, Figure <ref> illustrates two sequences, demonstrating that UniAD's prediction of the road network and vehicle tracking are fundamentally accurate. However, in terms of metrics, UniAD's performance in such scenarios with unseen road and traffic flow is significantly degraded, with an average PDM Score of only 0.636. The output trajectories by UniAD exhibit a substantial increase in collision rates and instances of driving outside the drivable area.
The open-loop experimental results underscore the critical importance of closed-loop experiments and tests for autonomous driving methods.

Closed-loop Evaluation.
We further evaluated UniAD's performance in 's closed-loop mode. In this mode, the trajectory outputted by UniAD is directly used for ego vehicle control, and the evaluation metrics include PDM Score (PDMS), Route Completion (RC), and Arena Drive Score (ADS). Our closed-loop experiment was conducted on four pre-set paths, with two paths selected in Boston and two in Singapore. The simulation time to complete each trajectory was approximately 120 seconds. Detailed results are presented in Table <ref>.

The results indicate that the PDMS of UniAD-generated trajectories in closed-loop mode (0.667) is comparable to that of the open-loop mode. However, the Route Completions (RC) are consistently low, averaging only about 13.7% of the total route length. Specifically, UniAD performs better on straightaways but largely fails to navigate the first turning intersection in the route. Figure <ref> illustrates two failure cases where UniAD lacked sufficient trajectory correction. Despite a roughly correct prediction of the road structure, it ultimately mounted the central green belt or failed to complete a right turn successfully. The average Arena Driving Score for UniAD is 0.086.
It should be noted that these are preliminary results based on testing only 4 routes. We plan to expand the number of routes for a more comprehensive evaluation and explore the combined effect of 's timing consistency and the driver agent's performance on the final ADS.




§ RELATED WORKS





 §.§ Data Acquisition for Autonomous driving


The characteristics of the automated driving dataset can be categorized into two aspects: appearance fidelity and interactivity. First, in terms of appearance fidelity, NGSIM <cit.> and CitySim <cit.> provide only realistic driving trajectories that can provide safe and reliable driving planning guidance. On top of that, some datasets developed based on the Carla simulator, such as DriveLM-CARLA <cit.> and BenchDrive <cit.>, provide simulated sensor data. Taking it a step further, the Waymo <cit.> and nuScenes <cit.> datasets capture real-world sensor recordings and the driving behavior of human drivers. The datasets were produced in a complex process and with a limited amount of data. To add variety to the scenarios, MagicDrive <cit.> and DriveDreamer <cit.> provide editable scenario generation. So far, we have been able to obtain diverse and rich data for training. However, the above data can only be used for open-loop evaluation, i.e., current decisions do not affect future data distributions, which differs significantly from real driving. Works <cit.> that also have fidelity differences, improve the interactivity of the data, they usually use auto-regressive generation methods to realize the interaction, the generation process implies the model's understanding of the world, and usually can not be too much human intervention. Some simulators <cit.> make things more controllable by decoupling part of the mechanics of how the world works. Common examples include simulators <cit.> that provide realistic traffic flow, and simulators <cit.> that drive vehicles in game engines, and reconstructive simulations represented by <cit.> that provide the appearance of reality.




 §.§ Diffusion-based Generative Models

Recent advancements in generative models have seen diffusion models play a pivotal role in image and video generation <cit.>. Moreover, recent works have expanded the scope by integrating additional control signals beyond traditional text prompts <cit.>. For instance, ControlNet <cit.> incorporates a trainable version of the SD encoder for control signals, while studies such as Uni-ControlNet <cit.> and UniControl <cit.> have emphasized the fusion of multi-modal inputs into a unified control condition using input-level adapter structures. 
Our approach aims to study the generation of continuous and controllable sequence frames, thereby bridging the gap between simulation environments and reality, and establishing the required foundation for closed-loop learning of autonomous driving agents.




 §.§ Evolution of Autonomous Driving Generation
 

World Models <cit.> utilize diffusion models to generate future driving scenes based on historical information, these methods often lack the ability to control the scenarios through layout, are difficult to generate continuous and stable videos and lack the approximation of physical laws.
TrackDiffusion focused on generating videos based on 2D object layouts <cit.>.
BEVGen <cit.> pioneered the generation of synthetic multi-view images based on the BEV layout, laying the foundation for a controllable generation of autonomous driving scenarios. BEVControl <cit.> extended this approach by a height elevation process, enabling image generation aligned with surrounding projection layouts.
Further advancements includes MagicDrive <cit.>, DriveDreamer <cit.>, Panacea <cit.> and DrivingDiffusion <cit.>, which generate panoramic controllable videos through various 3D controls and encoding strategies.
However, their primary focus lies in augmenting training data to enhance algorithm performance, rather than serving as simulators for modeling dynamic environmental interactions.



 §.§ Simulator-Driven Scenario Generation

Autonomous vehicle development is significantly enhanced by driving simulators, which provide controlled environments for realistic simulation. Prominent research efforts have concentrated on generating virtual imagery and annotations, with some studies expanding to incorporate environmental variations and construct safety-critical scenarios for training based on real-world data logs. Nevertheless, these simulated images frequently fall short of achieving true realism, as evidenced by previous works <cit.>. 
While SimGen <cit.> made a breakthrough as the first work to generate diverse driving scenarios following conditions from a simulated environment, it mainly focused on the quality of the generated content with only front-view images, neglecting the exploration of closed-loop systems. Our research aims to bridge this gap by developing a system that can not only generate realistic scenarios but also allow agents to interact with them in a closed-loop manner.




 §.§ Closed-Loop Driving in Simulation

End-to-end vehicle control algorithms <cit.>, are typically trained and evaluated on open-loop datasets <cit.>. However, these algorithms lack the capability to generalize directly to simulators for closed-loop evaluation, which hinders the demonstration of their true performance potential. Recent studies have increasingly recognized the significance of closed-loop evaluation, as exemplified by <cit.>. Moreover, simulation environments offer a wealth of training data, a stark contrast to models trained on datasets that are constrained by data distribution <cit.>. A significant challenge arises due to the discrepancy between the simulated scene's appearance and real-world conditions, complicating the generalization of models trained on simulation data to actual scenarios. This creates a paradox: the desire to utilize simulation data for its diversity and editability, while also seeking data that closely mirrors reality. Our approach effectively addresses this issue by enhancing the realism of the simulator for certain closed-loop learning methods <cit.>.




§ CONCLUSIONS AND FUTURE WORKS


This paper introduces a novel closed-loop simulation platform named  for vision-based driving agents.  integrates a  that generates human-like traffic flow and a high-fidelity generative  with infinite generation. This combination allows realistic interaction and continuous feedback between the driving agent and the simulation environment. The system provides a valuable platform for developing and testing autonomous driving agents in a variety of scenarios, marking a substantial leap in driving simulation technology.


 is designed with a modular architecture, allowing for easy replacement of each module. This paper presents an initial implementation of . As the first high-fidelity closed-loop simulator, we still have a few limitations for future improvement:

1) Data Diversity: The current generative model is trained solely on the nuScenes dataset, which limits the diversity and emergence capabilities. We plan to expand training to include more varied datasets to enhance the model's robustness and versatility.

2) Temporal Consistency: While we can generate continuous videos with an autoregression strategy, maintaining motion trends and temporal consistency between frames remains challenging. Future work will explore multi-frame autoregressive networks and more scalable architectures <cit.> to address these issues.

3) Runtime Efficiency: Like many generative models,  requires significant runtime. Investigating faster sampling methods <cit.> and model quantization may alleviate these problems.

4) Expanded Agent Testing: We plan to incorporate a broader range of driving agents within , facilitating the continuous learning and evolution of knowledge-driven driving agents in the closed-loop environment <cit.>.

5) A Real Arena:  can not only evaluate the performance of different driving agents, but also act as a testing ground for AD generative models. By using the same driving agent as a referee, it can fairly assess the sim-to-real gap of different generative models. This approach even provides a more credible and convincing evaluation compared to traditional metrics like FID and FVD.

We recognize that practical application may still be a way off, but the potential and promise shown by this work are evident. We hope this research will advance closed-loop exploration in highly realistic environments and offer a valuable platform for developing and assessing driving agents across a range of challenging scenarios. We encourage the community to collaborate in advancing this field. The era of open loops is transitioning, and autonomous driving evaluation and learning are set to enter a new era of closed-loop systems.



§ ACKNOWLEDGMENTS

The research was supported by Shanghai Artificial Intelligence Laboratory, the National Key R&D Program of China (Grant No. 2022ZD0160104) and the Science and Technology Commission of Shanghai Municipality (Grant Nos. 22DZ1100102 and 23YF1462900).




ieeetr





 Paper Title: Formal Ethical Obligations in Reinforcement Learning Agents: Verification and Policy Updates.



Multipartite Entanglement for Multi-node Quantum Networks
    G. Araneda
    August 5, 2024
=========================================================



§ ABSTRACT 


When designing agents for operation in uncertain environments, designers need tools to automatically reason about what agents ought to do, how that conflicts with what is actually happening, and how a policy might be modified to remove the conflict.
These obligations include ethical and social obligations, permissions and prohibitions, which constrain how the agent achieves its mission and executes its policy.
We propose a new deontic logic, Expected Act Utilitarian deontic logic, for enabling this reasoning at design time: for specifying and verifying the agent's strategic obligations, then modifying its policy from a reference policy to meet those obligations.
Unlike approaches that work at the reward level, working at the logical level increases the transparency of the trade-offs.
We introduce two algorithms: one for model-checking whether an RL agent has the right strategic obligations, and one for modifying a reference decision policy to make it meet obligations expressed in our logic.
We illustrate our algorithms on DAC-MDPs which accurately abstract neural decision policies, and on toy gridworld environments.




§ INTRODUCTION: STRATEGIC OBLIGATIONS IN THE FACE OF UNCERTAINTY



In the rapidly evolving domain of reinforcement learning (RL), agents are trained to autonomously perform tasks within complex and often unpredictable environments. 
While significant advancements have been made in improving the performance and adaptability of these agents, a crucial dimension remains under-explored: how these agents align with societal and ethical expectations.
Given the increasing ubiquity of RL agents in real-world scenarios — from finance <cit.> to health care <cit.> and transportation <cit.> — it is essential to establish mechanisms that ensure their actions do not merely optimize for the mission's objective but also adhere to a broader spectrum of ethical norms and societal standards. For brevity, we refer to these ethical and social standards as obligations.
Without such normative constraints, the agent's behavior is likely to be badly surprising, and ultimately unsafe when we think of the reactions of human agents in the environment.



  
Running example. 



Figure <ref> shows a problem where a drone must carry biohazardous material across a city to a hospital. The RL problem is to maximize the drone's expected utility, and the rewards are assigned such that the hospital rewards the drone with 10 points, but every other space gives the drone a penalty of -1. The penalty is designed to reflect the operating cost of the drone, and the opportunity cost of not being at its goal. The reward reflects the utility of the drone arriving and remaining at the hospital. This reward encodes the drone's mission. Solving this RL problem gives the policy that leads to the shortest route to the hospital.

However, we also want the drone to avoid a children's playground, where accidental contamination is especially problematic. This indicates a moral conflict between delivering the material to hospital patients quickly, and avoiding contamination to third parties. Instead of tweaking the reward or designing a secondary reward, we assign the drone an explicit, legible, logical obligation to avoid the playground with a high probability. The requirement is probabilistic to account for the uncertainty in the environment, which could make a non-probabilistic obligation unachievable in every case.

This explicit obligation, and the algorithms we develop to handle it, avoid the pitfalls of reward manipulation, especially as the number of moral dilemmas increases, and the reward balancing potentially becomes more and more arbitrary.


For this task we need ways to formalize the right obligations, and tools to guarantee that the agent has adopted a policy that maximizes its utility only subject to meeting these obligations.
Traditional specification languages, like Linear Temporal Logic (LTL), are inadequate at drawing a distinction between what is the case (how the agent behaves) and what should be the case (how it should behave to meet its obligations). 
Instead, a deontic logic is needed for this distinction <cit.>.
We adopt the logic of Expected Act Utilitarianism (EAU), first introduced in <cit.>.
EAU allows specification and automatic reasoning about the obligations of utility maximizing systems, including the models that underlie reinforcement learning. But it has the shortcoming of grounding obligations in the agent's optimal action at a given moment, without regard for future actions, which conflicts with the way a policy is computed as the maximizer of long-term utility.
To remedy this, we introduce a strategic modality to EAU which grounds obligations in the agent's entire policy, not just instantaneous action.

Designers also need algorithms that can both verify if a given agent's policy meets the specified obligations, and guide modifications to the agent's policy if it falls short.
This paper introduces two algorithms for these challenges: the first efficiently model-checks an RL agent's policy against the strategic obligations formalized in EAU, and the second employs policy gradient ascent to refine the agent's policy until it aligns with the given obligations.

This paper's contributions are:

    
  * An extension of EAU to strategic obligations, which formalize obligations that must be met by the optimal (utility-maximizing) policy over an infinite horizon, not just the current time step (Section <ref>).
    
  * An algorithm for model-checking whether a Markov Decision Process (MDP), which models action under uncertainty, has given strategic obligations formalized in (the extended)  EAU. The model-checker can handle MDPs with tens of thousands of states (Section <ref>).
    
  * An algorithm for modifying a utility-maximizing policy to also satisfy a deontic obligation (Section <ref>).
    
  * An extension of the above policy search algorithm for the case where rewards are not known a priori (Section <ref>).
    
  * Experimental evidence for the effectiveness of the model checking and policy search algorithms (Section <ref>).










§ RELATED WORK


There are several works that propose approaches for developing ethical RL agents.
In <cit.>, the authors argue that RL provides a sufficient foundation to ground ethical decision making, and <cit.> explores methods for implementing ethics into RL systems.
Works such as <cit.> and <cit.> propose the use of inverse reinforcement learning to teach norms to agents.
However <cit.> argues for a hybrid approach, and <cit.> offers a purely logical approach.
In this paper, we aim to maintain the precision and formality of a logical approach while extending it to be compatible with common RL techniques.



Our algorithm to modify an agent's policy resembles work in policy gradient methods <cit.>.
Policy gradient methods are robust and versatile mechanisms within the realm of on-policy reinforcement learning and have been used in safe reinforcement learning <cit.>. These techniques are designed to optimize cumulative reward by directly manipulating policy functions, affording ease of implementation and compatibility with function approximations. The foundational work of REINFORCE <cit.> marked an early instance where policy gradients were computed using Monte Carlo returns. 
However, our approach relies entirely on the dynamics of the system, and is concerned both with cumulative reward and conformance to an obligation.




This reliance on system dynamics makes our problem similar to the problem addressed in <cit.>.
There, the authors produce an MDP that encodes  logical constraints (in a non-deontic logic) and can be solved with dynamic programming.
However, their solution maximizes the probability of satisfying the logical constraint while ours seeks to maximize expected utility subject to satisfying a logical constraint.



Our goal of maximizing expected utility subject to the satisfaction of an obligation is analogous to the constrained MDP (CMDP) problem <cit.>.
The CMDP problem is to find the policy that maximizes expected utility subject to a constraint on expected cost (or secondary reward).

We seek to constrain the MDP solution with a non-discounted probability of reaching a state instead of with a discounted cost.

The non-discounted nature of the probability of reaching a state makes our problem distinct from the CMDP problem.
Further, the model checkers we employ do not consider discounted rewards.
Thus, our problem is not solved by solutions to the CMDP problem or by solutions to a classical model checking problem.






§ TECHNICAL PRELIMINARIES













We introduce the Expected Act Utilitarian deontic logic for formalizing obligations of agents acting in stochastic environments.
We then draw a correspondence between the semantic frames of the logic, MDPs, and Bellman optimality.
And we discuss the policy gradients that are central to our policy search algorithms.






















 §.§ Expected Act Utilitarianism





Expected Act Utilitarianism, or EAU <cit.>, uses PCTL <cit.> to describe states of affairs in the world, and adds modalities to speak of action and obligation.

Letting α be an agent from a finite set of agents ,  and  be Boolean conjunction and negation, and  a PCTL formula, the syntax of EAU is defined by the following grammar,

    A  |  A | A A |    |  αA


Intuitively, PCTL formula  describes a state of affairs, such as P_≥ 0.9 g: the probability of predicate g eventually holding is at least 0.9. See <cit.> for details of PCTL.
Formula  says that α “sees to it that”, or ensures, that A is true, while αA says that α ought to ensure that A is true.
For example, the formula αP_≥ 0.75[ playground] specifies the obligation to avoid the playground in at least 75% of possible executions.
The formula αP_< 0.01[ checkpoint] gives the obligation to eventually reach the checkpoint in less than 1% of possible executions. 

Figure <ref> illustrates all definitions in this section.
Formally, EAU formulas are interpreted over a branching time model . 
It is made of the following components: 

    
  * A tree Tree, whose vertices are called moments, m, representing decision points of the agent. Two moments are related by a directed edge mm' if m'  follows m. This is a partial order on moments, written as m<m'.
    
  * the root is moment `0', which precedes all other moments.
    
  * AP is a set of atomic propositions, and v: V(Tree)→ 2^AP is a labeling function that assigns atomic propositions to each moment m.
    
  * a history h is a linearly ordered, possibly infinite, set of moments — a branch of the Tree. We write H_m for the set of histories that start at moment m
    
  * K: an action available to an agent at a moment m. By identifying K with the subset of histories in H_m that are still possible after taking the action, we can consider that K⊆ H_m. We write  for the set of actions confronting α at m, and (h) to refer to those actions that contain history h.
    
  * (m' | m): the probability of agent α moving from m to m', assuming that the agent takes some action K that leads to m' (formally,  K ⊆ H_m and K ∩ H_m'≠∅).
    
  * : h →(h): a function that assigns a real value - a utility - to a history.


The branching time model can represent the roll-out of an MDP, where moments are state visits, and probabilities are derived from transition probabilities.

An EAU formula holds (or not) at an index of evaluation in the model, which is a moment/history pair m/h.
The satisfaction relation is written , m/h  A, where it is always the case that h ∈ H_m.
The proposition defined by the EAU statement A at moment m is the set of histories, starting at m, in which the statement holds:

    |A|_m^{h ∈ H_m , m/h  A}

When it is unambiguous, we drop  from the notation.

For convenience, we write the probability with which an agent can execute a particular history h from moment m as (h|m).
We can determine this value by taking the product of the probabilities (m' | m) along history h — assuming that the agent always takes the action K that history h is a part of, and dropping K from the notation.

The quality of an action, Q(K), is defined as:

    Q(K) = ∑_m' ∈ M'_K(m' | m) max_K' ∈αm' Q(K')

where M'_K is the set of moments that follow m by taking action K.

An agent's set of optimal actions, then, can be defined as the action(s) with the best quality at the moment:

    {K ∈   ∄K' ∈   
    s.t. Q(K)< Q(K')}


Now we can provide the formal semantics for two important concepts in EAU: agency and obligation.
In EAU, agency is modeled by the `Chellas sees to it' operator cstit <cit.>. 
The agent α “sees to it that” the statement A holds at index m/h if and only if it takes an action K such that A holds in all histories  in K  - i.e., K guarantees A is true <cit.>.
We write this as ,m/h  [α cstit:A].


We model obligation with the expected Ought.
The agent α “ought to see to it that” A holds at index m/h if and only if A is guaranteed by every E-Optimal action.

In other words, an agent's obligations at m are defined by the states of affairs A guaranteed by taking its optimal actions at this moment.
To denote that an obligation is satisfied we write ,m/h αA.

	
    With α an agent and A an obligation in a model , 
    
    ,m/h αA    iff  K ⊆ |A|_m^
          for all  K ∈



Note that an obligation at moment m is a function of the optimal actions at m alone. No consideration is given to what is optimal in the future or what actions are taken in the future. By contrast, an optimal policy (that maximizes expected utility) is typically followed forever. To reason about an agent's obligations under the consideration that it will always act in accordance with its optimal policy, we introduce a new modality to the logic in Section <ref>.

From Markov processes to stit models. 
In the context of Markov processes, a stit model can be thought of as a roll-out of the process.
A moment m is assigned to each visit of each state in the roll-out, and the root moment is the roll-out's first visit to the starting state.
Each action a in the MDP has a corresponding action K for each moment that action can be taken, and the probability (m'|m) is the transition probability T(s, a, s') for the states s, s' associated with the moments m, m', and the action a associated with action K.
A history h is a particular realization of an execution, and its  would be its discounted sum of rewards.
Thus the Q-function derived for the MDP is isomorphic to the Q(K) value for an action in the Tree, and so, as shown in <cit.>, an agent's optimal action K ∈ in the stit model is its Bellman optimal action in the MDP <cit.>.





 §.§ Property Gradients for Parametric Markov Chains


In the second part of this paper we will propose an algorithm for updating an optimal policy so that it satisfies the content of an obligation - that is, the PCTL formula P_≥ρϕ that shows up in an obligation operator αP_≥ρϕ.
To do this we leverage recent work <cit.>, in which gradients with respect to transition probabilities are computed. 
Specifically, given a parametric Markov Chain (MC) whose transition probabilities are parameterized, and some function f: S→ℝ of the states which is obtained as the solution of a linear program, <cit.> show how the gradient of f with respect to the parameters can be computed efficiently.
In practice, the parameters on the transition probabilities represent a policy.
The function f represents either the reward function (when we want to maximize expected utility), or the probability of satisfying the formula ϕ (when we want to maximize the likelihood of satisfaction).





§ MODEL-CHECKING STRATEGIC OBLIGATIONS


We aim to extend EAU to describe what an agent should do given that it follows its optimal policy for all time. Following <cit.>, we call these strategic obligations.
This makes EAU more applicable to the RL paradigm, where the optimal policy is typically followed forever. 
We therefore extend EAU to speak of strategic obligations, then present a model-checking algorithm to handle strategic obligations.



 §.§ Expected Strategic Oughts


While EAU gives us the capability to specify and reason about an agent's obligations in stochastic environments, the obligations it defines are determined only by the agent's optimal action in the moment of evaluation.
When an agent follows a strategy however, it is natural to inquire about the implications that strategy has on its behavior (beyond its immediate impact).
This is especially pertinent when verifying RL systems, as they are expected to follow a learned policy.

To reason about an agent's obligations under a strategy, we introduce a strategic obligation in the manner of Horty's strategic ought <cit.>, though a broader treatment of strategic modalities can be found in <cit.>.


A strategy (or policy, or schedule) π, is a mapping from moments to actions that determines what action an agent will take when it finds itself in a given state.
A strategic obligation, then, is an agent's obligation to the state of affairs brought about by an agent following its optimal strategy.
We begin with the strategic stit modality: , which says that agent α has some strategy π that, if followed, ensures that A is the case.
To say that an agent has a strategic obligation we write αAπ.

We define a strategy π as a mapping from moments m in the stit tree to a subset π(m) of the actions available at m.
The set of histories realizable by π starting at m is then 

    H_m,π = {h ∈ H_m | h ∈π(m')  ∀ m' ∈ Tree  s.t.  m'≥ m}

Thus h∈ H_m,π iff h is a possible evolution of the system if the agent follows π.



    
    In a stit model ,
    , m/h
    iff there exists a policy π such that
    h ∈ H_m,π and  H_m,π⊆ |A|_m^.


To say that an agent has a strategic obligation we must return to the question of optimality.
We say that a strategy is optimal if all actions in it are optimal actions.
That is, π is α's optimal strategy if and only if K ∈ for all K ∈π, and for all m.
Then we can say that an agent α has an obligation to strategically ensure that A holds if and only if all optimal strategies ensure A.

    
    In a stit model ,
    , m/h αAst
    iff H_m,π⊆ |A|_m^ for all optimal π



To illustrate these operators, we turn to Figure <ref>.
Let us define π_1 as {m:K_1; m':K_3}, π_2 as {m:K_2; m':K_3}, and π_3 as {m:K_1, K_2; m':K_3, K_4}.
We can say that m/h_1 αAπ as π_1(m) = {K_1}, and π_1(m') = {K_3}, so H_m,π_1 = {K_1}∩{K_3} = {h_1} and |A|_m = {h1} — that is, there is a strategy (π_1) that guarantees A.
The optimal action at m is K_2, and the optimal action at m' is K_3, so we can say that π_2 is the optimal strategy, and that m/h_4 αBπ2.
Finally, we note that since H_m,π_3 = {h_1, h_2, h_3, h_4, h_5}, we can't say that π_3 guarantees any atomic proposition.

If (h_1) were 10 instead of 7, then π_1 would be the optimal strategy. In this case, m/h_1 αAπ_1. In EAU without strategic obligations, however, we cannot assume that the agent will continue to act optimally in the future. Thus m/h_1 αA.



 §.§ Model Checker of Strategic Oughts

To make the strategic EAU modalities a practical tool for the verification of RL agents, we developed and implemented a model checking algorithm for strategic obligations.


Our model checking algorithm takes as inputs an MDP , and an obligation αϕπ where ϕ is a formula in PCTL.
The algorithm then determines if it is the case that αϕπ
By definition <ref>, this model checking problem can be performed in two sequential steps: first, find the optimal strategy π^*, and second, determine if all histories consistent with π^* satisfy ϕ .


Our model checker assumes that the optimal policy π^* is unique.
To enforce this assumption we simply employ a tie-breaking rule, forcing the selection of only one optimal policy if multiple are available.
This is done for simplicity; if the existence of multiple optimal policies is critical then the algorithm could be modified to return true only if all optimal policies satisfy the obligation.

This algorithm is shown in Algorithm <ref>, and experimental results of its performance are given in Section <ref>.








§ POLICY UPDATE TO SATISFY OBLIGATIONS





In this part of the paper, we move from model-checking an MDP against a strategic obligation, to policy search with strategic obligation constraints. 
Namely, the design team is given an obligation  as part of design requirements. 
The team comes up with a reward structure for the agent. 
Ideally, the reward function induces an optimal policy π^* that also satisfies the obligation, but that is not guaranteed since the reward might be balancing several requirements, and reward shaping is notoriously difficult. 
But the ethical obligations aren non-negotiable. 
We therefore ask: how can we modify the reward optimal π^* to obtain a policy π' such that the controlled system satisfies P_≥ρφ while maintaining a high expected reward?

We do this by leveraging gradient computation for parametric MCs <cit.>, which was described in Section <ref>. 
In our case, the parametric MC is the underlying MDP controlled by π^*.
The parameters of the MC are the probabilities of taking a given action in a given state, i.e. π(a|s). 
The function f to be optimized is the probability of the parametric MC satisfying φ: this probability is computed as the solution of a linear program <cit.>. 
The gradient of f relative to the policy parameters is computed by solving a system E of linear equations <cit.>.
We call this the probability gradient ∇_π f.
By doing gradient ascent on f relative to the action probabilities, we can increase the probability of satisfaction, at the cost of veering away from the initial, reward-optimal policy π^*.
This process does not change the agent's reward function — allowing us to continuously evaluate our modified policy on the reward function given by the design team.
Note that the probability bound can also be ≤ρ with trivial changes to the algorithm.
Note also that we can move along the top k largest gradient entries, but potentially decrease search effectiveness. 
The impacts of this are explored in Section <ref>.



  
Approaches to policy updates.
Given the probability gradients needed to improve the policy's probability to satisfy φ, there are many heuristics that can be taken to update the policy while maintaining a high reward.






∙ Line Search: One approach is to do a line search over the line connecting π^*, the reward-optimal policy, and π^φ, the policy that satisfies P_≥ρφ (and which can be obtained by, e.g., STORM <cit.>). This line is depicted in Figure <ref>. This update simplifies the search space, and is doable using a classical projected gradient.

∙ Average Gradient: We can also use the policy gradient ∇_π V for expected utility to guide our policy updates (the utility gradient). By following the average of this gradient and probability gradient ∇_π f, we can allow a simple trade-off between these two objectives.


∙ Alternating Gradient: Finally, we can use the probability threshold ρ to check when we should be following the probability gradient ∇_π f vs. the utility gradient ∇_π V. If the current policy's probability does not meet the threshold ρ, then it is updated by following the probability gradient. If the threshold is met, then the utility gradient is followed. This allows the policy to improve its utility only after it has sufficiently improved its probability of satisfaction.

In the next section we explore these three heuristics.




















§ EXPERIMENTS


To demonstrate the performance of our algorithms we report on the results of experiments on constrained policy search and model checking.
An implementation of our model checking algorithm and our policy update algorithm is included with code to run the following experiments at: https://github.com/sabotagelab/formal-ethical-obligations.



 §.§ Illustrative Example




We first apply our methods to the “windy-drone” system depicted in Figure <ref>.
The “windy-drone” system represents a drone delivering a heart for transplant while battling high winds. The effects of the wind are represented by stochastic transitions in the system. The drone's objective (as represented by its reward function) is to reach the hospital as quickly as possible. However, in such high winds, we want to prevent the drone from flying bio-hazardous material over locations such as playgrounds.

The algorithm is given the stochastic policy that maximizes the agent's expected reward, and aims to modify the policy to satisfy αP_≥ρ[  playground ]π.
This obligation represents the requirement for the agent to avoid the playground with a probability greater than ρ.
The dynamics of this obligation are interesting as fulfilling it pushes the agent away from its optimal policy — encouraging it to take a path that is almost twice as long in the best case.
In this system, ρ can not be larger than 0.998, and for the following experiments we set ρ = 0.75.
In practice, this probability threshold should be based on a risk analysis, and should reflect the degree of risk that stakeholders are willing to accept.



  
Line Search experiment.

Our first experiment was on the performance of policies, on the line in policy space, between the the maximally satisfying policy π^φ and the reward optimal policy π^*.
In this case we simply interpolate between the two policies — taking 100 steps between π^φ and π^*.

In Figure <ref> we show the expected utility and probability of satisfaction of the policies on the line between π^φ (update 0) and π^* (update 100).
Each update i on the x-axis describes the evaluation of a policy π(i) defined by:


    (1-i/100)π^φ + (i/100) π^*

As i increases, the probability of satisfaction monotonically decreases.
The expected utility, however, initially decreases, but, around update 40, increases again as π_i approaches π^*.
This non-monotonicity means that using a simple hill-climbing algorithm to find the maximum expected utility among policies on the line between two given policies will not necessarily return a policy with a global maximum expected utility.
We also note that the first update in this experiment to violate the probability threshold ρ in our obligation of 0.75 is around i=50.
The policy that maximizes expected utility that is found before π_50 is π_0, which is π^φ.


This figure lays the baseline for our other experiments in this section.





  
Average Gradient experiment.

Our second experiment with this problem was on the trade-off between expected utility and probability of deontic satisfaction.
The algorithm does gradient ascent along (∇_π V +∇_π f)/2 - the average of the two gradients.
This method is depicted in Figure <ref>, and the results are shown in Figure <ref>.
This allows the probability f of satisfaction to rise to almost 1.0, while pulling the expected utility for the agent down below 110 from its maximum of 142.
This suggests that there is room to increase utility again at the cost of some lowering of the probability of satisfaction (while maintaining f ≥ρ). 
Further, this approach gives no guarantee that f will be raised above ρ. 
To address these two issues, we perform an experiment using the alternating gradient approach.












  
Alternating gradient experiment.
Our third experiment uses the alternating gradient method depicted in Figure <ref> where the policy is updated using the gradient that improves its satisfaction probability when below a probability threshold, and otherwise uses the utility gradient.
The performance of this method is shown in Figure <ref> for a threshold ρ=0.75.
Here, the probability of satisfaction comes to oscillate around the threshold, allowing the expected utility of the policy to rise to almost 120 — an increase of almost 10% over the line search and average gradient methods while maintaining satisfaction.
This method ensures that the satisfaction probability f is near the set threshold, or is otherwise as high as possible. Only when f is above the threshold does this method improve the agent's expected utility. The alternating gradient method is also less computationally demanding than the average gradient method, as each iteration only requires the calculation of one gradient.





































  
Hyperparameter choice.
These experiments used a learning rate of 1.0, and all 64 of the system's derivatives (4 for each of the 16 non-absorbing states).
However, choosing a learning rate, or number of derivatives is not straightforward.
To demonstrate, we performed a small grid search over these hyperparameters.
As seen in Figures <ref> and <ref>, the search for satisfying policies is robust to most choices of learning rate and number of partial derivatives.
It appears that only a few of the system's derivatives are required to get a good result. This is encouraging, as larger MDPs will require more time to solve, but may be solved without having the calculate a gradient value for every decision variable.
Using all of the derivatives, however, yields the highest final probability of satisfying the given specification. And, for small learning rates, it seems to lead to the best expected utility of those policies that satisfy the specification. At higher learning rates, using 48 of the 64 available gradients leads to a policy that satisfies the requirement, and performs slightly better than a policy found using all of the gradients.







 §.§ Policy Optimization With Exploration




We also sought to test if our method could perform well when rewards were not known a priori.
To do this, we guessed the rewards at each state as an arbitrary value, and updated those values as the agent explored its environment.
This gave us access to an approximation of the utility gradient ∇_π V which we can use to increase the expected reward of our policy.


However, if we want to ensure safety while exploring, then we need a way to prevent the agent from taking unsafe actions.
To this end we implement a PCTL shield <cit.> that prevents the agent from taking any action that would violate the content of a given obligation.

With the shield in place we can allow the agent to explore while we update its policy towards satisfaction of a given obligation.
Then, once we have found a safe policy, we can follow the approximate utility gradient that is based on our observations so far.
This is how we implement our alternating gradient method to allow for exploration.

To test if this method is effective across a broader range of environments we randomly generated 24 12-by-12 gridworlds.
Each had 10 impassible cells, 10 “pits” that would assign a reward of -10 and ends a run, 10 “coins” that would assign a reward of +5, and a goal state that would assign a reward of +10 and ends a run.
These cells would be randomly placed on the grid, and the agent would always start at the bottom-left of the grid.
The agent is initiated with a random policy, and may explore the environment ϵ-greedily for 100 steps.
The agent's memory of reward values updates after every run, and the updated knowledge is used when calculating the current utility gradient.
The agent's policy is updated after every run as well using a learning rate of 0.01.

The obligation given to the agent is 
    αP_ > 0.75 [   coin ]π

That is, the agent should ensure that, with probability greater than 0.75, it should never enter a “coin” state.
This obligation was chosen because it interferes with the agent's ability to maximize its utility — forcing π^* and π^φ to be different policies.


The performance of this method is shown in Figures <ref> and <ref>.
These show that as time goes on, the agent approaches the threshold of probability satisfaction (0.75), at which point the expected utility levels out.
We would expect the utility to decrease as satisfaction probability increases since the obligation prohibits the agent from collecting rewards from the “coin” states very often.





























We also show that, for certain initializations of the rewards, shielding alone is insufficient to prevent the agent from learning an unsafe policy — thereby requiring the shield to be used at execution time as well.
Figure <ref> shows that the satisfaction probability decreases, even though the agent never takes an action that would violate its obligations.
When shielded and following only the approximate utility gradient for its policy updates, the agent does manage to find an average reward near zero, but its probability of satisfying the obligation if the shield is removed is reduced below 0.2 on average.
By learning a safe policy using the alternating gradients method instead, we could remove the shield at runtime - saving on computational requirements.





 §.§ Large Model-Checking Experiment: Cartpole


This experiment illustrates the execution of our model-checker on a large MDP  - specifically, on an MDP modeling the cartpole system.
The cartpole system consists of a pole attached by an un-actuated joint to a cart, which moves along a frictionless track. The system is controlled by applying a force of either positive or negative magnitude to the cart, with the objective of keeping the pole balanced upright without the cart running off the track. A tabular approximation of such a system can be defined by a DAC-MDP compiled on a set of trajectories collected by a random policy in the original MDP <cit.>. The size of the DAC-MDP can be adjusted based on the data size, fan-out, and fan-in size of the compiled MDP. Fan-out is controlled by the number of candidate actions for each state, and fan-in by the number of neighbors used to compile the DAC-MDP.

The MDP we retrieved had 50,000 states, 15 actions at each state, and 5 transitions per action, for over 3,000,000 transitions.
We formulated 20 PCTL formulas φ_1,…,φ_20 to check as both strategic stit statements αφ_kπ, and as strategic obligations αφ_kπ.
We labeled each state with the quintile of the angle of the pole in that state aq0 … aq4, and with the quintile of the x-position of the cart xq0 … xq4.
The time it took to check a subset of these formulas is given in Table <ref>.
We found that the time to model check strategic stit statements were consistently completed between 47 and 60 seconds.
All times were measured on a system with 16 GiB of RAM and an Intel i7-2620M CPU at 2.7 GHz.
Strategic obligations, with times near 25 seconds, were checked much more quickly; likely thanks to the induced Markov chain's smaller size - weighing in at fewer than 200,000 transitions.





 §.§ Obligation Implication


To demonstrate our policy optimization procedure on a more complicated form of obligation we show how to find a policy that satisfies an obligation with an implication within it.

In the “windy-drone” system, we might consider allowing the drone to fly over a playground if it visits the checkpoint — where its cargo is checked for safety and security.
We can model the obligation of the drone to fly over the playground (and thus take a shorter path to the hospital) if it visits the checkpoint as αP_≥β[  playground]  P_≥γ[  checkpoint]π.
This is equivalent to saying that the drone should avoid the playground, or it should visit the checkpoint.
Thus, we can synthesize a policy that avoids the playground, and another policy that visits the checkpoint, and take the better performing of the two as the optimal satisfying policy.

In our experiments, we set β to 0.75, and γ to 0.9.
As shown in Figure <ref>, a policy that avoids the playground results in an expected utility just below 120, while we found that visiting the checkpoint results in an expected utility of less than 40. Thus we take the former policy and complete the search.[The poor performance of the policy that visits the checkpoint is due to the Markov property of the MDP. If the policy could change after it visited the checkpoint to allow the drone to take the faster northern route, then visiting the checkpoint would be the better-performing policy.]



 §.§ Contrary-to-Duty Obligation

As mentioned, a key strength of a deontic logic is the ability to distinguish between what ought to be the case (the obligation) and what actually is the case, and to reason over the divergence between these two.

A contrary-to-duty (CTD) obligation is an obligation that enters into force in case a primary obligation (the duty) is violated: e.g. if the agent ought to ensure that the medicine cabinet is full (the primary obligation), but it isn't (the violation), then the agent ought to ensure that next an order is placed (the CTD). 
This has the following general structure: If α ought to ensure A but A does not happen, then α ought to ensure B next. (Other structures are possible, notably using conditional obligations)

    αAπ A α Bπ

We model-check such a formula on the “windy-drone” system depicted in Figure <ref>.


By setting a high reward on the checkpoint state, the agent will have the duty α P_≥ 0.7 [checkpoint]π.
However, in the case that the agent slips north towards the playground (the violation), it will inherit the new CTD obligation to return to the start state so that it might make a second attempt at the checkpoint: α P_≥ 0.6 [ start ]π.
In the CTD structure we have:

    α P_≥ 0.7 [checkpoint]π north 
    α P_≥ 0.6 [ start ]π

Verifying that our agent has this obligation allows us to determine how it is expected to behave when it enters a less-than-ideal state.

We checked this contrary-to-duty obligations by forcing the agent to move north, and then verifying the truth of the CTD obligation from that state.
More generally, we check CTD obligations by looking at the successor states that would indicate a violation of the primary obligation and then verify the truth of the CTD obligations from those states.





§ CONCLUSIONS


Our modalities for strategic agency and strategic obligation give us the expressive power to reason about a large class of reinforcement learning agents.
We presented an algorithm for model-checking that an MDP, equipped with a policy, has the right obligations captured in deontic logic.
We introduced a new kind of constrained MDP problem where reward maximization is constrained by an obligation.
We also provided a way to modify a policy so that it meets such an obligation — without needing to toy with reward functions.

We hope that these algorithms will aid system designers in specifying normative constraints, checking their systems against those constraints, and refining their systems to meet their constraints.

In future work we're interested in extending our policy update algorithm beyond safety properties to support a larger class of PCTL. We're also interested in managing the trade-offs between following the utility gradient and the probability gradient, perhaps by using a linear weighting function, or different learning rates for each. Further, we'd like to try constraining the policy search with a maximum divergence from the initial policy.


































































































 Paper Title: Contrastive Factor Analysis.




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




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





§ INTRODUCTION

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

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


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

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










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

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


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




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



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


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


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





§ PRELIMINARY ON CONTRASTIVE LEARNING



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


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

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

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









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


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

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





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

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



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


 



§ CONTRASTIVE (NON-NEGATIVE) FACTOR ANALYSIS

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


 §.§ Bayesian generative model



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



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

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

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

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

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


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









 §.§ Reparameterizable latent variable distributions






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

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



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

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










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


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




 §.§ Variational inference


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

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

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

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

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



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



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

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

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

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



§ RELATED WORK

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

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

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


§ EXPERIMENT



 §.§ Feature learning



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




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






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



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





 §.§ Uncertainty evaluation

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


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


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





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

 


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

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










































































































    


















 


 §.§ Disentanglement representation learning

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


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



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






§ CONCLUSION

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




unsrtnat










































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







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


    Venkat Margapuri
Department of Computing Sciences 

Villanova University

Villanova, PA, USA 


Prapti Thapaliya
Department of Computing Sciences 

Villanova University

Villanova, PA, USA 


Trevor Rife
Pee Dee Research and Education Center  

Clemson University

Florence, SC, USA 



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



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



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




§ INTRODUCTION

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



§ RELATED WORK

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



§ MATERIALS AND METHODS

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



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



 §.§ Extraction of Region of Interest (ROI)

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



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


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

    


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




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


    






 §.§ Estimation of Leaf Angle

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




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

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





§ IMPACTING FACTORS




 §.§ Impact of Image Sharpness

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



 §.§ Multiple Leaf and Stem Detections

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





§ RESULTS


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



§ FUTURE WORK AND CONCLUSION

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

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

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

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

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

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

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

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

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

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

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

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


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

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

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

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

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

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

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

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

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

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

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

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

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





 Paper Title: Emotion-driven Piano Music Generation via Two-stage Disentanglement and Functional Representation.



[
    [
    
=====


§ ABSTRACT 

Managing the emotional aspect remains a challenge in automatic music generation.
Prior works aim to learn various emotions at once, leading to inadequate modeling. 

This paper explores the disentanglement of emotions in piano performance generation through a two-stage framework. 

The first stage focuses on valence modeling of lead sheet, and the second stage addresses arousal modeling by introducing performance-level attributes. 

To further capture features that shape  valence, an aspect less explored by previous approaches, we introduce a novel functional representation of symbolic music.

This representation aims to capture the emotional impact of major-minor tonality, as well as the interactions among notes, chords, and key signatures.

Objective and subjective experiments validate the effectiveness of our framework in both emotional valence and arousal modeling. 

We further leverage our framework in a novel application of emotional controls, showing a broad potential in emotion-driven music generation.




§ INTRODUCTION


With the recent advancements in symbolic music generation <cit.>, there has been a growing interest in controlling high-level musical features throughout the generation process.

Among these features, emotion-driven music generation <cit.> aims to generate music that conveys specific emotions, representing a crucial aspect for music appreciation and analysis. 

The downstream applications of such models have also been explored, such as music therapy for healthcare and educational purposes <cit.> and soundtrack generation for videos and movies <cit.>.


Emotion could be represented in two dimensions from the literature <cit.>: valence and arousal.

Valence refers to the positiveness of an emotion and arousal refers to energy or activation <cit.>.

These two dimensions can be further divided into four quadrants (4Q), namely high valence high arousal (Q1),  low valence high arousal (Q2), low valence low arousal (Q3), and high valence low arousal (Q4). 

In this paper, we focus on the emotion-driven piano performance generation of these four quadrants. 

Throughout prior works, we observe crucial challenges from the perspectives of both model design and musical inductive bias.


First, previous emotion-driven piano performance generation models <cit.> attempt to learn emotion quadrants and expressions in an end-to-end paradigm. 

In terms of model design, this approach poses training difficulty on the generation model, leading to the instability in achieving results of desired emotions. 

For example, many existing works <cit.> could effectively control the arousal levels of music, while their performance of valence modeling, especially in generating low valence (i.e., negative) music, is still poor. 

In terms of music, the creation process of music typically involves multiple stages, such as the lead sheet composition for melodies and chord progressions, and performance generation for textures and expressiveness. 

Consistently, emotion can be evoked through a combination of musical elements (e.g., melody, chord, texture). 

For example, major/minor chords have been found to seize different valence trends in psychological studies <cit.> and performance-level attributes like articulation, tempo, and velocity are more related to arousal <cit.>. 

It is worth to explore the potential relation between the disentanglement of the generation process and the emotion expression. 










Second, previous emotion-driven generation models have received limited attention regarding the influence of tonality on emotion modeling.

It has been widely shown that major-minor tonality in composition is highly related to valence perception <cit.>. 

For example, as depicted in Figure <ref>, the histogram of musical keys derived from the emotion-labeled music dataset EMOPIA <cit.> supports the distribution skews to major keys for high valence clips and opposite trend for low valence ones. 

Furthermore, different tonalities may reveal similar patterns in the relative relationships between melodies and chords, while the distribution of melodies, chords, and tonalities can exhibit distinct shapes across different emotions.

Current representations of symbolic music, such as REMI <cit.> and CP-Word <cit.>, do not explicitly incorporate such interactions nor address its connection to emotion adequately.


Therefore, it is necessary to consider a functional format of symbolic music representation considering the relationships between notes, chords and key signatures to better model the tonality in the emotion-driven music generation process.

In this paper, we contribute to combat above challenges:


    
  * We employ a two-stage Transformer-based model on emotion-driven piano performance generation. The first stage focuses on valence modeling via lead sheet composition, while the second stage addresses arousal modeling by introducing performance-level attributes.
    
    
  * We propose a novel functional representation for symbolic music, encoding both melody and chords with Roman numerals relative to musical keys, to consider the interactions among notes, chords and tonalities <cit.>.
    
    
  * Experiments demonstrate the effectiveness of our framework and representation on emotion modeling. Additionally, our method enables new capabilities to control the arousal levels of generation under the same lead sheet, leading to more flexible emotion controls.


As a minor contribution, we also refine key signature labels and extract lead sheet annotations for the EMOPIA dataset <cit.> to ensure the correct training of the two-stage framework. 

We share the data, open source our code[<https://github.com/Yuer867/EMO-Disentanger>] and present generation samples in the demo page.[<https://emo-disentanger.github.io/>]

















































§ RELATED WORK




 §.§ Emotion-driven Piano Performance Generation
 

Prior works apply emotion conditions on deep-learning models to guide the generation of piano performance <cit.>, or develop searching methods to generate music of desired emotions <cit.>.

Musical elements via feature disentanglement <cit.> or supervised clustering <cit.> can further be regarded as a bridge between emotion labels and performances for generation.

In contrast, our framework employs a two-stage generation approach to reduce the complexities of one-stage generation, fostering a more nature process of music creation as well as a better incorporation between emotion labels and generation results.





 §.§ Tonality, Functional Harmony, and Emotion


Musical keys and functional harmony have been explored in the field of roman numeral analysis <cit.>. 

The analysis of how modes and tonalities relate to mid-level perceptual features (e.g., dissonance, tonal stability, minorness) and affect the emotional perception of music pieces has also been discovered <cit.>.

While some music generation works attempted to combine key information into data representation <cit.>, loss function <cit.> and text conditions <cit.>, none of them explore the relation between musical keys and emotional perception. 

In this paper, we leverage both functional harmony knowledge and class-octave based pitch representation <cit.> to design a new data representation, incorporating the relationships between notes, chords and keys for emotion-driven music generation.



§ METHOD

In this section, we will first introduce the functional representation of symbolic music as the main generation unit.

Then we introduce the two-stage model as the main component of the emotion disentanglement and generation. 



 §.§ Functional Representation


Figure <ref> illustrates our proposed functional representation.

Its design is initially based on REMI <cit.>, a widely used event-based representation for symbolic music. 

We incorporate different note and chord events assisting to better learn the joint information of emotion and key signature.





  §.§.§ Emotion and Key Events


We follow CTRL <cit.> to set up the condition within the autoregressive generation process in Transformer architecture.

To denote distinct emotions and affect overall properties, we begin the event sequence with  event to indicate the emotion label of music clips. 

The  event is appended after  to provide the musical key property, with the total of 24 keys (12 tonic notes with two modes in EMOPIA <cit.>). 



  §.§.§ Bar, Sub-Beat, Tempo and EOS Events


Similar to REMI, 

a  event denotes the new start of a bar; 

a  event denotes one of 16 possible discrete beat locations within a bar;

a  event denotes local tempo changes every four beats;

and an  event denotes the end of sequence. 



  §.§.§ Chord Events


A musical chord name typically consists of root note and chord quality. 

For example,  represents the chord  with root  and major quality.


Such symbols describe correct note information in chord within the tonality, but they overlook the variations in chord functions of the same chord across different tonalities.

For example, while  serves the tonic function in F major scale, it serves the subdominant function in C major scale. 

Moreover, the chord progression follows these functional harmony rules to establish tonality and convey musical emotion <cit.>.

To introduce chord functions in the emotion modeling, we adopt Roman numerals from Roman Numeral Analysis <cit.> to notate chord roots in Figure <ref>.

Given the  event, root notes in the absolute pitch are directly converted into Roman numerals based on their scale degrees relative to the key (i.e., relative pitch). 

For roots outside the scale, we employ a direct conversion for , , ,  and  appearing in major keys, but randomly assign  and , which only appear in minor keys, as one of their neighboring degrees during the encoding and decoding process. 

This design ensures the notation to be key-independent and make every conversion of notes reasonable to the music theory. 

The notations of chord qualities remain unchanged, and the chord event  appears every four beats.





  §.§.§ Note-related Events
 

A note is denoted by ,  and  events, where  event indicates the onset of pitches from A0 to C8. 

Inspired by <cit.>, we decompose  into  and  events according to the note octave and degree in the certain key scale. 

The conversion rule from  to   is the same as that of chord roots in Figure <ref>.

For example, pitch D#4 is decomposed into  and  in c minor scale, but  in D# major scale.

Such degree-octave pitch representation narrows the difference between melodies, thus improves the learning of connections between emotions, chords, and melodies, as demonstrated in Figure <ref>.




 §.§ Two-stage Emotion Disentanglement


We use the idea of Compose & Embellish <cit.> to generate music in two stages: lead sheet first, and then piano performance. While Compose & Embellish is emotion-agnostic, we extend it so that the lead sheet model involves valence modeling and the performance model  arousal modeling. 











  §.§.§ Valence Modeling



The top left section of Figure <ref> denotes the first stage, where only emotion events  and  are considered as conditions.

The former includes music pieces of Q1 and Q4 (high valence) and the latter includes those of Q2 and Q3 (low valence).

The lead sheet model first predicts a key event k conditioned on the given emotion event e, and then generates the lead sheet sequence M={m_1, ⋯ m_T} of length T, as melody and chord progression, conditioned on previous tokens step-by-step:

    p(k, M|e) 
        = p(k|e)∏_t=1^T p(m_t | e, k, M_<t)  ,

where p(k|e) and p(m_t | e, k, M_<t) are jointly learned through the Transformer-based generation model <cit.>. 

Performance-related events  and  are removed in the first stage (i.e., lead sheet generation), as we mainly focus on the contributions of key, pitch and chord for valence perception.




  §.§.§ Arousal Modeling

The top right section of Figure <ref> denotes the second stage.

Given the lead sheet M, the performance model generates performance X conditioned on the true emotion label (Q1 to Q4). 

As the valence aspect has already been modeling in the first stage, this stage focuses on the generation of musical textures for the lead sheet, and more importantly, on how to perform it through variations of tempo, velocity, articulation, and other performance-level attributes that largely influence perceived arousal <cit.>.

During the training and inference phases, with the positions of  events, M and X are further segmented into {M_1, ⋯, M_b} and {X_1, ⋯, X_b}, where b is the number of bars. 

The segmented sequences are “interleaved” in the form of {⋯, M_i, , X_i ⋯} with additional  events to distinguish M and X tracks.

In that, the target performance bar X_i is appended to its corresponding conditions M_i, as mapping each lead sheet segment to its corresponding performance segment <cit.>.

With the emotion condition and key event from lead sheet as prefix tokens, the performance model is summarized as

    p(X|e, k, M) = ∏_i=1^b p(X_i|e, k, M_≤ i, X_<i)  .






  §.§.§ Training Objectives

Lead sheet and performance models are trained separately by both optimizing the negative log-likelihood loss of the sequence.

Since existing emotion-labeled music datasets are not large, we leverage large-scale music datasets without emotion annotations to pretrain both models for better music understanding.

During pretraining, the emotion event is marked as .

We then finetune two models on the emotion-labeled dataset (detail in Section <ref>)
to learn composition and performance styles specific to different emotion contexts.





  §.§.§ Two-stage Inference
 
The left bottom section of Figure <ref> denotes the inference process of both models.

In the first stage, the lead sheet model predicts the key event and generates the lead sheet sequence step-by-step given  or  event, creating a musical motif for the specific valence preference.

Even though our framework has the capability to generate any-key music of specific emotions, we observe that some generation results, such as a high-valence and high-arousal song with a minor key scale, may go beyond the current definition of emotion in <cit.>, where the valence naturally has a strong correlation to the major-minor tonality (Figure <ref>).

Therefore, we limit major keys to  and minor keys to  during the inference stage. 

We acknowledge that this can be overly simplifying.

Since this paper focuses mainly on the valence-arousal disentanglement during the generation process, we leave this exploration of generating any emotion within any key as an advanced topic for future research. 

In the second stage, the performance model generates piano performance with desired valence and arousal combination given the lead sheet from the first stage. 

For example, to generate a music piece of Q3, a “Negative” lead sheet and a “Q3” emotion event are selected as conditions.

Additionally, this two-stage framework enables the flexibility to generate different arousal levels of piano performance under the same lead sheet, delivering some scenarios when the music need to shift quickly to complement the scenes in movies or daily videos (detail in Section <ref> and the demo page). 





§ EXPERIMENTS




 §.§ Datasets and Preprocessing

As presented in Table <ref>, we collect different datasets for pretraining and finetuning phases as mentioned in Section  <ref>. 

For pretraining the lead sheet model, we use the HookTheory dataset <cit.>, where we choose 18,206 lead sheets with high-quality and human-transcribed melody, chord and key annotations in 4/4 time signature. 



We simplify 249 chord quality classes into 11 types[Major, minor, augment, diminish, suspend2, suspend4, major7, minor7, dominant7, diminish7, half-diminish7] as the same set in the other datasets below.

For pretraining the performance model, we use the Pop1k7 dataset <cit.>, consisting of 1747 transcribed pop piano performances.

Since Pop1k7 does not contain lead sheet annotations, we refer <cit.> to extract melodies using the skyline algorithm <cit.>, recognize chords using the chorder library <cit.>, and detect key signatures using <cit.> in MIDI Toolbox<cit.>.

For finetuning the models with emotion conditions, we use the EMOPIA dataset <cit.>, consisting of 1,071 music clips with human-annotated emotion labels. 

Similar to Pop1k7, we obtain the lead sheets of EMOPIA by extracting melodies using the algorithm in <cit.> and recognizing chords using the algorithm in<cit.>.


Empirically, we obverse that specifically in the EMOPIA dataset, melodies and chords extracted by these alternative algorithms are more correct compared to the skyline algorithm and the chorder library. 


Additionally, we found the key signature labels in EMOPIA are not fully correct since they are also obtained by the detection algorithm with error rates.


Since the valence modeling is strongly related to the musical keys and modes, we manually correct the key annotations of 367 clips in EMOPIA to ensure a high quality of lead sheets.

All datasets are randomly divided into respective training and validation sets at the ratio of 9:1. 

As a result in our functional representation, the vocabulary size of events is 215 for lead sheet and 324 for piano performance.








 §.§ Model Settings

The lead sheet model is a 12-layer Transformer Decoder <cit.> with 8 heads, 512 hidden dimensions and relative positional encoding <cit.>. 

The performance model is similar to the lead sheet model except with Performer attention <cit.>.

The total parameter sizes are 41 million and 38 million respectively.



Both models are trained with the batch size of 4, the maximum sequence length of 512 (lead sheet model) or 3072 (performance model), and the Adam optimizer with β=(0.99,0.9).

We adopt a 200-step warm-up to achieve the maximum learning rate of 1e-4 for pretraining and 1e-5 for finetune.

All models are implemented by PyTorch and trained on one NVIDIA Tesla V100 GPU. 

The lead sheet model took around 180,000 steps to converge and the performance model took around 200,000 steps.


Nucleus sampling <cit.> is employed in the inference phase.

We referred <cit.> to choose the sampling hyperparameters τ=1.2, p=0.97 for the lead sheet model and τ=1.1, p=0.99 for the performance model.



 §.§ Baseline and Ablations

We consider the emotion-driven piano performance generation model in EMOPIA <cit.> as our baseline, which generates music in an end-to-end paradigm instead of two stages. 

To ensure the fair comparison of generation performance, we trained the baseline model under ths same datasets in both pretraining and finetune phases, and replaced the original CP-Word representation with REMI as the former usually yields better generation performance and more comparable to our proposed functional representation.  


Two other related works <cit.> are not included in comparison due to the main reason that we focus more on the evaluation of the two-stage framework in valence and arousal modeling; and the partial reason that they are not open-source or releasing the reproducible model weights.

We conduct a comprehensive ablation study to evaluate if each proposed design benefits the emotion modeling of music generation. 

Specially, these designs include: 1) the two-stage generation, 2) the functional representation, and 3) the dataset pretraining.


In the following sections, models are denoted as <representation(stage)>. For example, REMI(one) denotes the one-stage generation model with REMI representation as the baseline, and REMI(two) denotes the two-stage generation as one variant.






 §.§ Objective Evaluation and Results

Even though previous studies <cit.> employ metrics, such as Pitch Range (PR) and Number of Pitch Classes (NPC), to evaluate the generation performance, they do not provide any evidences on the superiority of  melody development, chord progression, and texture arrangement of music. 

Therefore, a model with more similar PR and NPC values to those of the target dataset does not necessarily promise a better generation quality than others.

Instead of using such metrics, we wish to evaluate the consistency between the input musical conditions and the generation results. 

We introduce key consistency to assess if a model can generate music pieces that adhere to the desired input key signatures, which is highly correlated to the lead sheet development and valence modeling. 

Specifically, key consistency measures the match between the key condition  and the actual key detected in the generation using the algorithm <cit.> with an 81% accuracy rate. 

We compare REMI and our functional representation to determine if the functional representation can improve the key consistency via more close and interactive designs on key, melody, and chord. 

Since this metric requires the key as conditions, we add the  event in REMI after  (as REMI+key in Table <ref>) when training the model. 

The non-pretrained versions are also included for comparison.

Each model generates 200 lead sheets (100 high and 100 low valence) and 400 performance samples (100 per emotion quadrant) for evaluation.

From Table <ref>, the functional representation outperforms REMI (i.e., REMI+key) across all components and even achieves compatible accuracy to real data over the lead sheet component.

This demonstrates the effectiveness of the functional representation, by representing notes and chord roots relative to key events for key modeling.

In contrast, REMI struggles with associating chord events with keys due to the ignorance of chord labels serving different functions in different key scales. 

Moreover, pretraining process introduces musical priors to enhance the learning of key relationships with other musical elements, improving key consistency for both representations.







 §.§ Subjective Evaluation and Results

We leverage an online listening test to assess the emotion modeling ability of models. 

The test was conducted to collect user responses on three parts: 1) valence modeling, 2) arousal modeling, and 3) 4Q emotion modeling. 

During this test, the quality of the generated music has also been assessed implicitly as it is a prerequisite to the emotion expression in the music.

22 participants were engaging in this test, 5 with less than 2 years of musical training, 8 with 2-5 years, 3 with 5-10 years, and 6 with more than 10 years.




  §.§.§ Valence Modeling


In this part, each participant listened to 16 generated tracks of piano performance from four models [four tracks (two high valence and two low valence) per model]: 1) Real data; 2) REMI (one); 3) REMI (two); 4) Functional (two).


For each track, participants rated its positiveness from –2 (low valence) to 2 (high valence) with the step size 1.

The left of Figure <ref> (`(a)') presents the mean opinion scores for the valence-oriented test, where the Functional (two) model significantly outperforms both REMI (two) and REMI (one) models.

The REMI (two) model shows a slight improvement over REMI (one) due to its two-stage design.

Our proposed Functional (two) model even marginally exceeds real data in low valence scores, which could be due to the potential subjective biases in the negative emotion as discussed in <cit.>. 

And our model achieves both great performance in high valence and low valence results, demonstrating a good balance in valence modeling.





  §.§.§ Arousal Modeling


In the second part, the functional (two) and REMI (two) models are chosen to compare their arousal modeling performance.

Specifically, we wish to explore whether they can generate piano performance with either high or low arousal under the same lead sheet based on the given conditions (Q1 and Q4 for positive lead sheets, Q2 and Q3 for negative ones). 

Two pairs of generated tracks are randomly drawn for each model and each valence level, where every pair includes two tracks of different arousal conditions.

For each track, participants rated its arousal level from –2 (low arousal) to 2 (high arousal) with the step size 1. 

The right of Figure <ref>  (`(b)') presents the results. 

The Functional (two) model surpasses REMI (two) by an average of 0.5 point, highlighting its superior ability to differentiate between the musical features of the two arousal levels through performance. 

Additionally, it is rare for Functional (two) to be incorrectly identified as high arousal under low arousal conditions (Figure <ref> (b-2)).



  §.§.§ 4Q Emotion Judgement


In the last part, participates needed to choose the best option from four options (4Q) for each track, with 8 tracks in total for the two models the last section (4 tracks per model and 1 track per emotion).

Figure <ref> presents the confusion matrices of two models.

The Functional (two) model achieves the higher overall accuracy than that of REMI (two) (71.5% vs. 31.0%).

When examining each emotion category, Functional(two) demonstrates superior performance in Q3 and Q4 than Q1 and Q2. 

Furthermore, music pieces generated from it with high valence conditions are misidentified almost based on their arousal levels; for instance, pieces intended for Q1 are almost mistaken for Q4 and vice versa. 

In contrast, for REMI(two), the misclassifications are across all categories, demonstrating its limitations in modeling the four emotion classes although through two-stage generation.



All above evaluations support that the combination of two-stage framework and functional representation is effective in controlling the emotion of the music it generates to a certain extent.















§ CONCLUSION AND FUTURE WORK

In this paper, we first explore emotion disentanglement through a two-stage Transformer-based framework for emotion-driven piano performance generation. Then we propose a novel functional representation for symbolic music to capture the interactions among musical keys, modes, chords, and melodies in relation to the emotion contexts. 

An objective metric is designed to qualify the key modeling of the proposed method, and subjective evaluations further confirm its ability to convey desired emotional perception. 

In the future, we wish to focus on enhancing the flexibility of emotional music generation across all musical keys and investigating new applications fostered by our framework, such as the controls of valence and arousal attributes under similar music motifs.











§ ACKNOWLEDGMENT


The work is supported by a grant from the National Science and Technology Council of Taiwan (NSTC 112-2222-E-002-005-MY2). This work also benefits from the review comments of an unpublished work <cit.>, which is an early version of the present study. While the concept of functional representation has been mentioned in that previous work <cit.>, its focus was on emotion-driven melody harmonization and the effectiveness of the functional representation in that context was unclear, as the melody was pre-given and fixed. In the current work, melody and chords are generated together, allowing more flexibility for emotion control with the proposed functional representation. The source code of that previous work can be found at the GitHub repo: <https://github.com/Yuer867/EMO_Harmonizer>.





























 Paper Title: WAS: Dataset and Methods for Artistic Text Segmentation.





 



WAS: Dataset and Methods for Artistic Text Segmentation






X. Xie et al.




Huazhong University of Science and Technology, China
{xdxie,yzli12,yangliu1213,xbai}@hust.edu.cn Adobe, USA 

{zzhang,zhawang}@adobe.com, wxiongur@gmail.com

WAS: Dataset and Methods for Artistic Text Segmentation
    Xudong Xie1 Yuzhe Li1 Yang Liu1 Zhifei Zhang2 Zhaowen Wang2 Wei Xiong2 Xiang Bai1()
    August 5, 2024
=======================================================================================



Corresponding authorfootnote


§ ABSTRACT 
  Accurate text segmentation results are crucial for text-related generative tasks, such as text image generation, text editing, text removal, and text style transfer. Recently, some scene text segmentation methods have made significant progress in segmenting regular text. However, these methods perform poorly in scenarios containing artistic text. Therefore, this paper focuses on the more challenging task of artistic text segmentation and constructs a real artistic text segmentation dataset. One challenge of the task is that the local stroke shapes of artistic text are changeable with diversity and complexity. We propose a decoder with the layer-wise momentum query to prevent the model from ignoring stroke regions of special shapes. Another challenge is the complexity of the global topological structure. We further design a skeleton-assisted head to guide the model to focus on the global structure. Additionally, to enhance the generalization performance of the text segmentation model, we propose a strategy for training data synthesis, based on the large multi-modal model and the diffusion model. Experimental results show that our proposed method and synthetic dataset can significantly enhance the performance of artistic text segmentation and achieve state-of-the-art results on other public datasets. The datasets and codes are available at: https://github.com/xdxie/WAS_WordArt-Segmentationhttps://github.com/xdxie/WAS_WordArt-Segmentation.
  





§ INTRODUCTION


Text segmentation is dedicated to finely segmenting the strokes of text from complex scene images, discriminating whether each pixel belongs to the text foreground or the background. Accurate text segmentation results are the foundation for text-related generative tasks. For instance, tasks such as text image generation<cit.>, text style transfer <cit.>, and text removal <cit.> can produce excellent and practical generative outcomes based on text masks. However, although existing models have achieved outstanding performance in regular text segmentation tasks, they are difficult to accurately segment artistic text. Artistic text features complex appearances and shapes, making it difficult to distinguish from background patterns <cit.>. Therefore, artistic text segmentation is a challenging task that has not yet been studied by the academic community.

The first problem encountered in implementing this new task is the lack of datasets. The existing text segmentation datasets ICDAR13 FST <cit.>, COCO_TS<cit.>, MLT_S <cit.>, and Total-Text <cit.> suffer from the problems of low annotation quality and insufficient quantity. Moreover, these data are scene images with regular text. Although TextSeg <cit.> provides high-quality annotated data and some artistic text images, the number of these images is still insufficient to train a high-performance artistic text segmentation model with strong generalizability. Therefore, we construct a real WordArt Segmentation dataset called WAS-R, consisting of 7100 artistic text images with word-level annotations of quadrilateral boxes, masks, and transcriptions. Additionally, to further enhance the accuracy and generalization ability of the text segmentation model, we also propose a synthetic dataset called WAS-S. Our designed synthetic pipeline utilizes the popular large multi-modal model and diffusion model to achieve realism, accuracy, and diversity in the generated images.

Besides, artistic text segmentation presents two unique challenges compared to general object segmentation and regular text segmentation. (1) The strokes of artistic text have flexible and changeable local shapes, such as slender tails or twisted ligatures. (2) The global topological structure of the artistic text is very complex, with many holes and intricate connections within the text. In contrast, the local stroke shapes and the global structure of regular text are almost invariant, and the topological structure of general objects is very simple. Therefore, the task we propose has clear academic value and practical significance. 

There are currently few specialized models for text segmentation. Recent studies either require the aid of text detection modules <cit.> or the assistance of character-level recognizers <cit.>.
Moreover, these methods have not been specifically designed for artistic text. In view of this, we propose a WordArt segmentation model WASNet. To address the first challenge, we propose a Transformer decoder with the layer-wise momentum query. The input of the self-attention module is the momentum superposition of the masked queries from the current layer and the previous layers. This operation ensures that the model does not quickly fit a restricted regular mask area when updating attention, thereby ignoring some special-shaped stroke regions, which are precisely what the earlier layers are capable of capturing. To address the second challenge, we propose a skeleton-assisted head that enables the model to output both mask predictions and skeleton predictions simultaneously, guiding the model to capture the global topological structure. It enhances the decoder's ability to perceive the overall structure of the text.

We conduct extensive experiments to verify the effectiveness of the proposed method and the synthetic dataset on the task of artistic text segmentation. We also verified the generalizability on other public datasets <cit.>.
The results achieved state-of-the-art (SOTA) performance. More importantly, the model trained on the WAS dataset can be directly tested on other datasets without the need for fine-tuning and still achieve competitive results. This opens up a new experimental paradigm for the task of text segmentation.

In summary, our contributions are four-fold:


  (1) We present a new challenging task: artistic text segmentation, and construct a real dataset to benchmark the performance of various models.   

  (2) We design a training data synthesis strategy and generate a synthetic dataset consisting of 100k image-mask pairs.

  (3) We introduce the layer-wise momentum query to handle the changeable local strokes and skeleton-assisted head to capture the complex global structure.

  (4) We achieve new SOTA results in the tasks of artistic text segmentation and scene text segmentation, and simplify the experimental paradigm for text segmentation.









§ RELATED WORK






 §.§ Text Segmentation Method


Some early text segmentation methods relied on thresholding <cit.>, low-level features <cit.>, or Markov Random Fields (MRF) <cit.> to segment foreground text from the background, but these methods could only achieve limited success in document processing. 
With the continuous advancement of deep learning, the corresponding text segmentation methods have shown great potential in complex scenes <cit.>. For instance, SMANet <cit.> employs the encoder-decoder architecture of PSPNet <cit.> and achieves a multi-scale attention module for text segmentation. PGTSNet <cit.> employs a pre-trained detector to ground out text regions before segmentation, further enhancing the accuracy of segmentation. TexRNet <cit.> incorporates character recognition and attention-based similarity checking to aid the model in segmenting text. Building on these methodologies, Yu  <cit.> developed a model featuring a lightweight detection head and a Text-Focused Module, elevating text segmentation performance in complex scenes to a new level. Nevertheless, recent high-performance text segmentation methods either utilize extra bounding box annotations and rely on text detection, or employ character-level supervision. Also, these methods lack specialized design for artistic text.




 §.§ Text Segmentation Dataset


The construction of text segmentation datasets has not received enough attention in academia. ICDAR13 <cit.> and Total-Text <cit.> provide high-quality, pixel-level annotations for text segmentation, but their quantities are very limited with only 462 and 1,555 images, respectively. To address the issue of insufficient quantity, researchers have proposed a dataset COCO-TS <cit.> (14,690 images) based on COCO-Text <cit.> for text segmentation. Similarly, MLT_S <cit.> (6,896 images) is also a large-scale text segmentation dataset based on ICDAR MLT <cit.>. Both of these datasets use automatic annotation strategies, resulting in low-quality dataset annotations. In view of these problems, Xu  <cit.> introduced a larger-scale and high-quality text segmentation dataset TextSeg (4,024 images), which includes character-level and word-level annotations of masks, bounding boxes, and transcriptions. Moreover, different from the datasets mentioned above, BTS <cit.> is a bilingual text segmentation dataset. These real datasets are all derived from natural images but lack a dedicated dataset for segmenting artistic text. Although TextSeg contains some artistic text images, the quantity is insufficient to train a robust model for artistic text segmentation.




 §.§ Segmentation Dataset Generation

The construction of synthetic segmentation datasets plays an important role in studying visual perception problems. DatasetGAN <cit.> and BigDatasetGAN <cit.> only use a small number of manually labeled samples for each category to train the decoder and generate a large amount of new data. 
Diffumask <cit.> extends the text-driven image synthesis to semantic mask generation in Stable Diffusion <cit.> to create a high-resolution and class-discriminative pixel-wise mask. 
Dataset diffusion <cit.> leverages the pre-trained diffusion model and text prompts to generate segmentation maps corresponding to synthetic images. 
DatasetDM <cit.> is a generic dataset generation model that decode the latent code of the diffusion model as accurate perception annotations. 
MosaicFusion <cit.> is a diffusion-based data augmentation method that does not require training and does not rely on any label supervision, especially for rare and new categories. 
Then SegGen <cit.> integrates Text2Mask and Mask2Img synthesis to generate training data, improving the performance of state-of-the-art segmentation models in various segmentation tasks. 
These advanced methods often fail to align artistic text with their masks in images. In this paper, we avoid having the model generate both images and mask annotations. Instead, we pre-render the text masks and use ControlNet <cit.> to generate mask-conditioned images.








§ DATASET

As artistic text in the real world is incredibly diverse, we propose two new datasets: WAS-R composed of real-world text images, and WAS-S composed of synthetic text images. These multi-purpose artistic text datasets aim to bridge the gap between artistic text segmentation and real-world applications, accommodating the rapid advances in text vision research.



 §.§ WAS-R Image Collection

The WAS-R dataset is composed of 7,100 images sourced from a variety of contexts, including posters, cards, covers, logos, goods, road signs, billboards, digital designs, and handwritten text. Among these, 4,100 images serve as the training dataset, while the remaining 3,000 images constitute the test dataset. The artistic text can be categorized into two major types according to the way capturing images. A type of artistic text image is taken by cameras from various scenes, such as signboards. The other type is directly exported from design software, such as poster files.
During data collection, we specifically balances these two types to create a diverse dataset for research and development.



 §.§ WAS-R Image Annotation

The WAS-R dataset stands out due to its comprehensive annotations, surpassing existing datasets. Specifically, WAS-R provides minimum quadrilateral detection boxes with distinct segmentation mask labels for each word. It also provides text transcription for each word mask. Moreover, we annotates the word effects such as shadow, glow, 3D, which play a crucial role in distinguishing artistic text from conventional scene text and significantly impacts text segmentation. Fig. <ref> shows examples of collected images and their annotations in WAS-R. 









 §.§ WAS-S Synthetic Dataset Construction

Fig. <ref> shows the pipeline of generating synthetic text images. The core idea is that we build a text image generation model which can generate aligned text images from text masks and input prompts. To this end, we construct the training pipeline as illustrated in Fig. <ref> (a). Specifically, we first generate diverse and informative captions from the text images in the training set of WAS-R to obtain training triplets <caption, mask, image>. Following that, we train a ControlNet <cit.> with these triplets for generating diverse images that are pixel-wisely aligned with the input text mask. During inference, as shown in Fig. <ref> (b), we first construct diverse text masks using Mask Render, then use GPT-4 to extend the texts in the mask into a scene description caption. The constructed text mask and caption are send to our trained ControlNet to generate the synthetic text image. We describe the modeling details below. 





Training Pipeline. During training, we construct the  dataset for training ControlNet <cit.> based on the training set of WAS-R. To this end, we obtain the image captions of existing training samples in WAS-R using the advanced multi-modal large language model called Monkey <cit.>. Formally, let I_t represent the real image in the training set of WAS-R, and C_t denote the prompt generated from I_t, we have C_t = Monkey(I_t).

Having obtained the <caption, text mask, text image> training triplets, we train a ControlNet that maps the input prompt and text mask to a text image. Our goal is that the outline of the artistic text in the generated image should be well aligned with the input text mask. Moreover, the contents and styles of the generated images should be diverse enough as guided by the input prompts. 



Inference Pipeline.  During inference, as illustrated in Fig. <ref>(b), we first generate synthetic text masks denoted as M_syn using our proposed Mask Render technique. Specifically, for each mask, we randomly select 1-7 phrases from the 20 newsgroups dataset <cit.> based on the word distribution of each image in the real dataset WAS-R. These phrases are consist of 1-5 consecutive words. Additionally, we apply a random rotation in the range of -30^∘≤ϕ≤ 30^∘ to each phrase. The size of each phrase is limited to match the general width of the entire image, and we position them randomly within the image boundaries. Besides, we use 250 artistic fonts.
Finally, an affine transformation is applied to each phrase to introduce skewness and distortion. 

We use GPT-4 <cit.> to generate the prompts corresponding to the synthetic text masks. Specifically, we ask GPT-4 to mimic the style of captions we generate from the training set in WAS-R, and synthesize new prompts. Next, we incorperate the text information in the synthetic mask into the generated prompt to obtain the final prompt. Formally, we have: C_syn = GPT4(C_t, C_M), where C_syn is the caption we generate and C_M is the text in the synthetic mask. Fig. <ref> shows examples of our synthetic prompts. 

Following the construction of synthetic text mask and synthetic prompts, we use the trained ControlNet to generate the final text image. We have I_syn = ControlNet_θ(C_syn, M_syn), where I_syn is the image we generate, θ denotes the trainable parameters in ControlNet. Fig. <ref> shows examples of the final synthesized <text mask, prompt, image> triplet. More details can refer to the supplementary materials.










§ METHODOLOGY

In this section, we introduce our artistic text segmentation model WASNet. We first present the overall architecture, followed by detailed descriptions of the local and global designs.





 §.§ Overall Architecture

The overall framework of WASNet is shown in <ref>. We take an excellent semantic segmentation model Mask2Former <cit.> as the meta-architecture. It is a mask classification architecture that directly predicts multiple binary masks and corresponding category labels, instead of performing per-pixel classification.
We add a skeleton-assisted head and improve the Transformer decoder with a mechanism of layer-wise momentum query. The backbone extracts low-resolution features from an image. The pixel decoder upsamples the image features and generates a feature
pyramid. The multi-scale features are fed into the Transformer decoder, with each resolution corresponding to each layer's input of the decoder. Besides, each layer of the Transformer decoder also receives the mask prediction and query generated from the previous layer as input. Finally, the mask head and the skeleton head generate binary mask and skeleton predictions respectively, by decoding the per-pixel embeddings from the pixel decoder and object queries from the Transformer decoder. The ground truth of the skeleton is obtained by thinning the binary mask labels through a skeleton extraction algorithm <cit.>.





 §.§ Transformer Decoder with Layer-wise Momentum Query


Artistic text segmentation faces the challenge of the local stroke shapes being flexible and changeable. Due to designers using hundreds of different artistic fonts and applying various text effects, the local strokes of the same character can differ significantly. This results in some slender strokes spanning across other areas, as well as twisted ligatures leading to complex text edges.  In contrast, normal scene text typically utilizes regular printed fonts without special designs, and the stroke shapes are almost invariant. Therefore, it is necessary for the decoder to pay attention to these special local strokes.

First, we use the masked attention mechanism <cit.>, constraining cross-attention
to within the local text mask region for each query, instead of attending to the full feature map. This mechanism can be expressed as:

    𝐌𝐀_l=softmax(𝐌_l + 𝐐_l𝐈_l^T) 𝐈_l,

where l is the layer index, and 𝐐_l is the input queries. 𝐈_l ∈ℝ^H_l W_l × C is the image feature input to the l-th layer, which comes from the feature pyramid of the pixel decoder. H_l and W_l indicate the spatial resolution of the image feature and C is the feature dimension. 𝐌_l is transformed from the binary mask output of the previous layer, with the value of the text region being 0 and the value of the non-text region being -∞ <cit.>.
𝐌𝐀_l is the output of the masked attention module. We have omitted the residual connection and normalization here. 

Furthermore, since the masks predicted by each layer are different, the previous layers yield coarse masks that may include special-shaped stroke regions. However, the subsequent layers are inclined to predict more precise regions of regular strokes, overlooking those local special regions. Therefore, in order to prevent the model's attention from being quickly confined to regular regions, we design a mechanism of Layer-wise Momentum Query (LMQ). The momentum superposition of the masked queries from the current and previous layers is input to the self-attention module before the module gathers contextual information.  <ref> illustrates this mechanism.

    𝐌𝐐_l+1=α𝐌𝐐_l+(1-α)𝐌𝐀_l,

where α∈[0,1) is a momentum coefficient. 𝐌𝐐 is the momentum query that is input to the self-attention module. We ultimately use this decoder in WASNet with layer-wise momentum query.



 §.§ Skeleton-Assisted Head

Different from regular text and general objects, the global topological structure of artistic text is very complex, and there are many holes and intricate connections inside. This presents new challenges for the segmentation task. The model needs to capture the global structure of the text object rather than just a region. Inspired by DeepSkeleton <cit.> and DeepFlux <cit.>, 
we found that the skeleton is an effective representation to describe the shape and topology of text because it can extract the central axis of the object. Therefore, we use skeletons to assist text segmentation.

As shown in <ref>, we add a skeleton-assisted head to WASNet, enabling the model to simultaneously predict the mask and the skeleton, thus endowing it with the capability to perceive the global topological structure. Since the binary mask is a finely annotated label for semantic segmentation, the ground truth for the skeleton can be obtained by processing the mask with the classic Zhang-Suen <cit.> skeleton extraction algorithm. The algorithm progressively removes pixels that satisfy certain template structural conditions through an iterative process, until no more pixels meeting the conditions are deleted.

We use the binary cross-entropy loss and the dice loss <cit.> for our skeleton loss and mask loss:

    ℒ_skeleton=ℒ_mask= λ_ceℒ_ce + λ_diceℒ_dice.

We set λ_ce=λ_dice=5. The final loss is the combination of skeleton loss, mask loss, and classification loss:

    ℒ_final=ℒ_skeleton+ℒ_mask+ λ_clsℒ_cls,

where λ_cls=2 for
predictions matched with labels and 0.1 for predictions that have not been matched with any labels.

During the inference phase, it is unnecessary to output the predictions of the skeleton. Therefore, we follow the post-processing method in <cit.> to obtain the final output of text semantic segmentation.







§ EXPERIMENTS




 §.§ Implementation Details

Our experiments are mainly based on the MMSegmentation <cit.> toolbox. The overall hyperparameter configuration is the same as <cit.>. The pixel decoder is a multi-scale deformable attention Transformer <cit.> with 6 layers. The Transformer decoder consists of 9 layers, each with an auxiliary loss.
We use the AdamW <cit.> optimizer and the poly <cit.> learning rate schedule with an initial learning rate of 10^-4 and a weight decay of 0.05. The data augmentation strategies include random scale jittering, random color jittering, random cropping as well as random horizontal flipping. We use a crop size of 512×512 and a batch size of 16. The models are trained with 8 RTX4090 GPUs. 
If the model is only trained on the real dataset, we set the number of iterations to 100k. If the model needs to be pre-trained on the synthetic dataset WAS-S, we first pre-train the model for 50k iterations and then fine-tune it on the real dataset for 50k iterations.
For the momentum coefficient α in <ref>, we set α=0.8 by default.
Following the previous text segmentation methods <cit.>, we use foreground (text) Intersection-over-Union (fgIoU) as the major metric and F-score measurement on foreground pixels as the auxiliary metric. 






 §.§ Results of Artistic Text Segmentation



To verify the superiority of our method in the task of artistic text segmentation, we trained several representative models on our WAS-R dataset, including six semantic segmentation models and two text segmentation models. We use the officially released code for TexRNet <cit.>, DDP <cit.>, and TextFormer <cit.>, and the code reproduced by MMSegmentation <cit.> for other models. For a fair comparison, we did not apply the character-level glyph discriminator for TexRNet.

The experimental results in <ref> indicate that our WASNet outperforms all of these advanced models. Moreover, when we train the baseline models and WASNet with the synthetic dataset WAS-S, their performance can be further improved. Our final results have achieved a significant SOTA performance.




 §.§ Results of Scene Text Segmentation

To further verify the generalizability of WASNet, we also conducted experiments on three publicly available scene text segmentation datasets <cit.>, as shown in <ref>. We can draw the same conclusion as in <ref> regarding the effectiveness of WASNet and our synthetic dataset. It is worth mentioning that character-level annotations were used to train TexRNet on TextSeg. Extra bounding box labels were used to train the text detection module of TextFormer on all three datasets. However, we only use binary mask labels of the full images. Despite this, we still achieved competitive or state-of-the-art results. Due to the highly inaccurate annotation quality of COCO_TS <cit.> and the fact that Total-Text <cit.> contains only 300 test images, the conclusions drawn from the evaluation results of the models on these two datasets may be inconsistent.








Furthermore, we directly evaluate the performance of WASNet on the three datasets using the model trained on the synthetic and real WAS datasets, as shown in the last row of <ref>. Note that the results in this row have not been fine-tuned on specific datasets but are still competitive. Therefore, to simplify the experimental paradigm and evaluation process of text segmentation models, we encourage researchers to train on WAS and test directly on other datasets.



[t]0.45
captypetable



Ablation study on our proposed modules and datasets.


2*Methods     2cWAS 

    fgIoU     F-score 



Baseline <cit.>     80.21     0.890



+ LMQ     80.95     0.898 

+ Skeleton      82.11     0.901

+ WAS-S     84.18     0.913 





[t]0.48
captypetable


Ablation study on datasets.


2*Dataset     2cWAS 

    fgIoU     F-score 



WAS-S     84.18     0.913 



5w images    83.25     0.906 

20w images    84.01     0.912

BLIP2 <cit.>      83.07     0.906

1000 fonts     82.35     0.902 














 §.§ Ablation Study

In this section, we conduct the ablation study on the artistic text segmentation dataset WAS.
We first validate the effectiveness of our proposed modules and the synthetic dataset. As shown in <ref>, as we gradually apply the design of LMQ and the skeleton to the baseline, the performance of WASNet is incrementally improved. Pre-training WASNet on WAS-S can continue to enhance the performance of artistic text segmentation.



Therefore, the synthetic dataset is an important contribution of this paper. We conduct ablation experiments regarding some synthetic details of the dataset in <ref>. It is crucial to control the number of synthesized mask-image pairs. Less data will weaken performance, but more data will lead to a performance plateau.
The reasons will be analyzed in <ref>.
We also use other large multi-modal model BLIP2 <cit.> to generate the image captions, but the performance is limited. This is because the overall performance of BLIP2 is inferior to Monkey <cit.> we used. Besides, we applied more fonts to generate masks, but the performance actually decreased. The dataset of 1000 fonts includes a large number of regular fonts, which reduces the learning difficulty of the dataset.

Furthermore, we explored the impact of different momentum coefficient values α on the performance of WASNet in <ref> and found that the approximate optimal value is 0.8. A coefficient that is too large can cause the model to be overly influenced by the coarse predictions from earlier layers. A coefficient that is too small diminishes the positive effect of momentum queries.































 §.§ Further Analysis

To further verify the effectiveness of our proposed WASNet, we visualize the inference outputs of our baseline model Mask2Former <cit.> and WASNet in <ref>. According to <ref> (a), it is evident that WASNet can capture special-shaped stroke regions such as slender tails or twisted ligatures. This is attributed to our  Transformer decoder with the layer-wise momentum query. Additionally, according to <ref> (b), WASNet exhibits good scale adaptability. It can achieve fine results for both large-scale and small-scale text with complex structures. This is because the skeleton-assisted head can obtain the global topological structure of the text through the thinning operation, guiding fine segmentation. 

Once accurate text stroke masks are obtained, downstream text-related generative tasks can demonstrate excellent results. The application effects of text removal, text background replacement, and text style transfer are shown in the supplementary materials.




 §.§ Limitation


Although the proposed synthetic dataset can improve the performance of text segmentation models, the enhancement is limited and does not significantly increase. Even when we further increased the amount of synthetic data, the performance remained unchanged. This could be caused by the bottleneck encountered in the diversity and realism of the synthetic images. In the future, we are considering designing more advanced generative models.



§ CONCLUSION

The paper focuses on a new challenging task of artistic text segmentation. We propose a real dataset for this task to train the models and benchmark the performance. We also construct a synthetic dataset to further enhance the accuracy and generalization ability. In order to meet the challenges of this task, we introduce the layer-wise momentum query to handle the changeable local strokes and the skeleton-assisted head to capture the complex global structure. Experimental results have demonstrated the effectiveness and superiority of our method in the tasks of artistic text segmentation and scene text segmentation. We hope that more researchers can focus on this task in the future and that the dataset we propose can change the experimental paradigm of text segmentation.





§ ACKNOWLEDGEMENTS

This work was supported by the National Science Fund for Distinguished Young Scholars of China (Grant No.62225603).






splncs04




§ MORE DETAILS ON SYNTHETIC DATASET CONSTRUCTION

As stated in the paper, we first construct the training pipeline for a text image generation model, learning to generate text images spatially aligned with text masks. Then we construct an inference pipeline to input new masks and prompts into the trained generation model, generating new text images. Here we add the details of prompt generation in the training and inference pipelines.


 §.§ Training Pipeline

To train the text image generation model such as ControlNet <cit.>, it is necessary to obtain training data of <caption, text mask, text image> triplets. Text masks and text images are from our proposed real dataset. Captions should be detailed descriptions of the text images. To this end, we utilize a large multi-modal model, Monkey <cit.>, to caption the images. Monkey is an open-source model and can handle vision-language tasks with high-resolution input and detailed scene understanding. It performs well on Image Captioning and various Visual Question Answering (VQA) tasks. Therefore, we feed a text image and a prompt “generate the detailed caption in English” to Monkey and let it output a detailed description. The examples of the generated captions are shown in <ref>. We found that, in many cases, Monkey is able to recognize and describe the text in images. To ensure the accuracy of the descriptions and to highlight the importance of the text, we add a sentence after each caption: This image contains the text “text in the image”.







 §.§ Inference Pipeline

During the inference phase, we first need to produce new binary masks of text through the Mask Render introduced in the paper. Moreover, it is crucial to generate new prompts that describe more complex scenes. Combining the masks with rich descriptions of scenes, the trained model can generate new and realistic text images. We use GPT-4 <cit.> to generate the prompts. To ensure that the new prompts and the training prompts are in the same domain, and avoid domain gaps in the images generated by the model, we first provide GPT-4 with 50 caption examples produced by Monkey. Then we ask GPT-4 to mimic the style of these captions and synthesize new prompts. The instruction is Please follow the above caption examples and generate a similar caption, which must contain some double-quoted spaces “ ”. Next, we insert the text corresponding to each new mask into the quotation marks in the new prompt, forming the final prompt. The generated <prompt, mask, image> triplet is shown in <ref>.



§ APPLICATIONS



 §.§ Text Removal

Text removal refers to the process of erasing or deleting text regions from an image. The finer the text mask, the better the erasing performance, as it preserves more background pixels. Therefore, stroke-level text segmentation can greatly benefit this task. Text removal is essentially an image inpainting task, so LaMa <cit.> is employed and the results are shown in <ref>. 






 §.§ Text Background Replacement

Once we have obtained the fine mask of the text, we can freely replace the background of the image, embedding the text into various scenes. We use ControlNet <cit.> to replace the background and <ref> presents the results.





 §.§ Text Style Transfer

Text style transfer is a task that renders text in natural images into artistic text according to a style reference image while keeping the text content unchanged. It usually relies on accurate text masks. We use Intelligent Typography <cit.> as the style transfer model and input the predicted text masks to it. The stylized text is shown in the last column of <ref>.


 Paper Title: DynamoLLM: Designing LLM Inference Clusters for Performance and Energy Efficiency.





















camerareadyfirstpage
  
  
  [C]
    [][12mm][t]13.5cm IEEE International Symposium    
    [][12mm][t]13.5cm IEEE International Symposium
      [][12mm][t]13.5cm IEEE International Symposium
      [][0mm][t]13.5cm IEEE International Symposium
    
    < g r a p h i c s >

    
    < g r a p h i c s >

    
    < g r a p h i c s >

    
  [C]






: Designing LLM Inference Clusters for Performance and Energy Efficiency
    
    

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




camerareadyfirstpage
  empty
plain
  plain










§ ABSTRACT 

The rapid evolution and widespread adoption of
generative large language models (LLMs) 
have made them a pivotal workload in various applications. 
Today, LLM inference clusters receive a large number of queries with strict Service Level Objectives (SLOs).
To achieve the desired performance, these models execute on power-hungry GPUs causing the inference clusters to consume large amount of energy and, consequently, result in excessive carbon emissions.
Fortunately, we find that there is a great opportunity to exploit the heterogeneity in 
inference compute properties and fluctuations in
inference workloads, to significantly improve energy-efficiency.
However, such a diverse and dynamic environment creates a large search-space where different system configurations (e.g., number of instances, model parallelism, and GPU frequency) translate into different energy-performance trade-offs.
To address these challenges, we propose , the first energy-management framework for LLM inference environments.
 automatically and dynamically reconfigures the inference cluster to optimize for energy and cost of LLM serving under the service's performance SLOs. 
We show that at a service-level,  
conserves 53% energy and 38% operational carbon emissions, and reduces 61% cost to the customer, while meeting the latency SLOs.






§ INTRODUCTION




The exponential growth in the adoption of generative large language models (LLMs) has positioned them at the core of numerous technological advancements and applications.
Today, we see use-cases of LLMs in various domains, such as healthcare <cit.>, 
developer productivity <cit.>,
data analytics <cit.>, education <cit.> and other.
As the popularity of LLMs increases among users, the inference clusters receive millions of queries per day <cit.> resulting in large infrastructures with 
sophisticated software and expensive hardware systems.

To meet these ever increasing computing demands,
researchers proposed various software <cit.> and hardware <cit.> techniques. Such techniques improve the performance efficiency of LLM inference clusters.
However, one aspect that has been largely overlooked is the energy consumption of these environments <cit.>.
The substantial energy requirements of serving LLMs running on power-hungry GPUs have emerged as a significant concern.
As these models become integral to various services, minimizing their energy consumption and, consequently, carbon emissions while maintaining high performance is paramount.

To address this gap, this paper 
starts by characterizing the energy-efficiency properties of LLM inference workloads.
Our characterization underscores
that such environments present a distinct set of challenges, divergent from existing energy management schemes tailored for traditional datacenters applications <cit.>.
Specifically, we observe that heterogeneity in LLM inference compute properties and fluctuations in LLM inference workloads
create a dynamic environment with large variations.
Such variations arise from:
(1) requests with varying input/output token lengths,
(2) distinct compute properties of different LLMs,
and (3) different SLOs required by the services using an LLM.

Requests with a large number of input tokens are compute intensive, thus, sensitive to GPU frequency.
Conversely, requests with a few input tokens and many output tokens have low compute, but high memory requirements. Reducing their GPU frequency would save the energy without significantly impacting the performance.
Moreover, 
the number of model parameters also affects the LLM's sensitivity to
the number of GPUs and GPU frequency.
Finally, depending on the service currently using the LLM, the SLO requirements can be strict requiring high-performance configurations,
or loose allowing for lower-performance but more energy-efficient configurations.
Importantly, these characteristics rapidly change due to load fluctuations and dynamic distributions of requests. Such dynamic changes 
cause a system configuration that is energy-efficient at a given point, to quickly become sub-optimal.
This requires a dynamic approach to resource management.


To 
pave the way towards energy-efficient and sustainable LLM inference clusters,
this paper introduces , the first energy-management framework for
LLM inference environments. 
 exploits the unique properties of LLM inference workloads to reduce their energy consumption while meeting the performance SLOs.
The system uses energy-performance profiles of models and their workloads to automatically and dynamically select the energy-efficient configuration.
It leverages multiple 
knobs, including
scaling in/out the number of server instances, model parallelism across GPUs, and GPU frequency scaling.



To handle workload heterogeneity,  maintains differently configured pools of LLM instances that are optimal
for different types of incoming requests.
For instance, compared to a request with many input and output tokens, a request that processes and outputs fewer tokens runs more efficiently on a model parallelized across fewer GPUs running at a lower frequency.
As request distribution varies over time,
 dynamically sizes the pools.
These pools can be merged into fewer pools or divided into multiple pools over time,
providing a balance between right-sizing and fragmentation of resources. 
To efficiently manage the resources,
 
uses a hierarchy of controllers that reduces computation complexity and eliminates centralized bottlenecks.
The controller at each level operates under the conditions imposed by the upper level, computes its dedicated knob, and forwards further constraints to the controllers at
a lower level. Finally, to enable frequent and smooth transition across different configurations,  includes techniques to minimize or hide the reconfiguration overheads.
As a result, the system maintains high levels of efficiency and service quality under changing workload demands.

An evaluation of   with a large GPU cluster running production-level traces from a major cloud provider shows that  is very effective:
it
conserves 53% energy, 38% operational carbon emissions, and reduces 61% cost to the customer, while meeting the latency SLOs.


The contributions of this paper are as follows: 


  * An analysis of the opportunities for energy-efficient LLM serving, rooting in the heterogeneity and fluctuations within the inference workloads.


  * Design and implementation of , a high performance and energy-optimized framework for LLM inference.



  * An evaluation of  on a large-scale platform using production-level traces.
 








§ BACKGROUND






Computational phases of LLMs
Generative LLMs <cit.> are auto-regressive:
they process the whole input in parallel, and serially generate the output tokens.
This property leads to two computationally distinct phases <cit.>.
First is the prefill phase, where the input tokens are computed in parallel.
This is a compute-intensive phase and scales with the number of input tokens. 
Second is the decode phase, where each output token is generated serially, based on all the tokens seen so far.
This is a memory-intensive phase, and scales with the number of output tokens.



Performance metrics for LLMs
To evaluate the performance, we use:
time to first token (TTFT), time between tokens (TBT), and throughput  <cit.>.
TTFT is the latency of generating the first output token;
while TBT is the latency to generate each new output token.
To quantify the energy efficiency, 
we measure the energy consumption 
in Watt-hours (Wh) while meeting certain latency SLOs. 
The SLOs 
vary depending on their use cases for different tasks. For latency-sensitive tasks, both TTFT and TBT are important metrics with strict SLOs.
We define SLOs for TTFT and TBT based on maximum achievable performance, described in <Ref>.




LLM parallelism
A single model can be divided across GPUs
to improve performance and allow larger memory footprints.
LLM inference typically uses pipeline and tensor parallelism.
Pipeline parallelism (PP) partitions the LLM layers among GPUs, while keeping all the operators/tensors of a layer on the GPU.
GPUs then communicate only in between two consecutive stages.
Tensor parallelism (TP) allocates a slice of each layer to each GPU.
This requires aggregation across all the GPU for each layer, in turn needing high bandwidth communication.
TP performs better for GPUs within the same server, connected with high bandwidth interconnects (e.g., NVLink <cit.>), while PP is preferred across servers.
Since most open source models <cit.> fit on 8 GPUs in a single server, we consider only TP in the rest of the paper; the ideas can easily extend to PP.
We denote tensor parallelism across 2, 4 and 8 GPUs as TP2, TP4 and TP8, respectively.




Power and energy in datacenters
A rich body of work explored power/energy efficiency in traditional datacenters <cit.>. 
However, the rapid growth of LLMs 
has posed new challenges 
that have not yet been extensively studied. 
LLM inference workloads comprise a swiftly increasing percentage of datacenter load <cit.>.
This, coupled with the power-dense hardware like DGX A100s and H100s being deployed to serve these workloads makes them power, energy, and carbon-intensive <cit.>.
To effectively address this challenge,
it is important to have a comprehensive framework for managing energy in these systems.












§ OPPORTUNITIES FOR ENERGY EFFICIENCY



To understand the energy-efficiency properties of LLM inference environments, we characterize open-source models <cit.>
on an NVIDIA DGX H100 server <cit.> using vLLM <cit.> inference engine.
We analyze the energy properties of LLMs by varying the request lengths, request load, model, and service SLO.
Additionally, we analyze how the profiled variables change over time in a real-production environment using the invocation traces of two LLM services from :
Coding and Conversation.
The traces include a subset of invocations received by the profiled services
during one week, and  contain the timestamp of the invocation, along with the number of input  and 
output tokens.
These traces are a super-set of our open-source traces for the same services <cit.>.




 §.§ Heterogeneous Energy-Performance Profiles




Request lengths
The prefill and decode phases in an LLM inference exhibit distinct execution behaviors (<Ref>), suggesting that requests of different input and output lengths possess different compute and energy characteristics.
We categorize the requests based on the number of input/output tokens into 9 buckets:
SS (short input, short output), SM (short input, medium output), SL (short input, long output), MS, MM, ML, LS, LM, and LL.
<Ref> shows the thresholds and corresponding TTFT/TBT SLOs.
We set the thresholds for request lengths using the 33^rd, 66^th and 100^th percentiles of the input/output lengths from a trace for a Conversation service from .
We set the SLOs to 5× the latency of a single request running isolated on a system <cit.>.



We use these categories to characterize the energy consumption of different request types running the Llama2-70B <cit.> model with a medium system load of 2000 tokens per second (TPS) under various GPU frequencies and model parallelisms.
<ref> shows our results in the form of a heat map.
Since shorter requests are not computationally intensive, they meet their SLOs with any tensor parallelism, and generally at lower frequencies compared to the rest.
As an example, the least-energy configuration for SS requests is TP2 at 1.2 GHz.
Conversely, LL requests can only run with TP8 without violating the SLO.
With TP8, the least-energy configuration for LL requests is 1.6 GHz.
Note that the lowest power configuration that meets SLOs (TP8 at 1.2 GHz), is not the energy-optimal one due to the increased execution time.
Running all the requests together would require the system to run 
with the most constrained SLO configuration, in this case, as per the LL configuration.
This would make the system energy inefficient.

To exploit this heterogeneity for energy-efficiency, 
the system would need to separate requests based on their input/output lengths, and process different request types with different server configurations.
However, on request arrival, the input length is known, but, due to
the auto-regressive LLM nature, the output length is unknown.
Thus, the system needs to predict the output length.
 will rely on prior work that efficiently performs such operation with relatively high precision <cit.>, and will have a mechanism to mitigate the impact of occasional mis-predictions.



Request loads
In addition to the request length, the incoming load of the LLM inference server drives the compute requirements.
During  periods of low load, the system has a larger SLO slack to exploit and can run the 
requests at low-frequency configurations to save energy.
Conversely, during  periods of high load, the system does not have enough SLO slack, and needs to run at high-frequency configurations.

<ref> shows the energy consumed when running Llama2-70B medium-sized 
input and output  (MM) requests while varying the number of processed prompt tokens per second (TPS).
The system can run low load with any TP at almost any frequency.
Among all feasible configurations, the lowest-energy  configuration is TP4 with 1.2 GHz.
TP8 requires more GPUs to operate in parallel and, thus, consumes more energy.
TP2 uses fewer GPUs but increases the execution time and forces individual GPUs to operate at high 
frequency to meet SLOs, leading to high energy.
Conversely, under high load, the system cannot operate on TP2 and requires TP4 or TP8.
The lowest energy  configuration is TP4 with 2 GHz. Overall,
to minimize the energy consumption while operating under performance constraints, we need to consider the incoming load to set the correct  parallelism and GPU frequency.






Requested model
The diversity of the compute properties of an LLM directly translates into its energy profile.
<ref> shows the energy consumption of different LLMs when running medium-sized requests at medium system load.
Smaller models, such as Llama2-13B and Mixtral-8x7B, can run with any TP (even with a single GPU); their lowest-energy  configuration is TP2 at 1.2 GHz.
Mixtral-8x22B and Falcon-180B are much larger and can only run with TP8.
Their lowest-energy configuration is TP8 at 1.2 GHz.

Compute-bound models with large number of parameters are more sensitive to the GPU frequency and model parallelism. Hence, they often need to operate at high-frequency and high-energy modes.
Sparse models with relatively smaller numbers of parameters tolerate lower frequencies and lower model parallelism. Hence, they  meet the performance requirements even with lower-performance modes.




Service SLO
Different services often use the same model with different SLO requirements <cit.>.
As indicated before, we assume an SLO such that the P99 tail latency is within 5× of the execution time of a request on an unloaded system <cit.>.
However, some services have more relaxed SLOs, at 10× or even 20× of a single request execution <cit.>.
For different SLO requirements, the system may need different energy-optimal configurations.
For example, Table <ref> shows that, with strict SLO (5×), short-input long-output sized LLama2-70B requests at medium load have the optimal configuration at TP4 and 1.2GHz.
However, if with loose SLO (10×), the requests may even operate with TP2 at 1.6GHz.




Insight #1
LLM  workloads are highly heterogeneous in their energy-performance profiles.
To achieve the optimal energy under performance SLOs, different requests (sizes, models and SLOs) need to be processed separately and differently.









 §.§ Dynamic LLM Inference Workloads







Changing request-length distribution
We measure the distribution of request types for Coding  and Conversation services.
<ref> shows the distribution of requests for each workload over a week.
The distribution 
differs across services.
 Conversation has typically longer  outputs and shorter inputs, while  Coding shows the opposite trend.
However,  both services have a significant fraction of each request type, and importantly, the popularity of 
request types changes over time.


As observed earlier, 
different request types require different energy-optimal configurations.
Thus,
the system needs to split its resources into per request-type pools, configure pools individually,
and dynamically adapt the pools' configurations based on the current request distribution.
However, if the system classifies the requests into too few classes, it will not be able to fine-tune the system for best energy.
On the other hand, too many classes may lead to fragmentation and negatively impact  energy efficiency.
Thus, the system has to find the right number of resource pools.
In , we will use historical data to set the number of pools such that requests with distinct SLO requirements (TTFT or TBT bound) and compute properties (compute or memory bound) have separate pools.
Moreover, as the load of a given request type reduces,  will avoid fragmentation by merging the pool with the next available pool that serves longer requests.




Request load fluctuations
 LLM inference workloads, as user-facing applications, exhibit a typical diurnal pattern with peaks during working hours and valleys at night and weekends.
<ref> shows the load in tokens per second of the two workloads over a week.
The load is normalized to the peak of the individual workloads.
The Coding trace shows a clear diurnal pattern, with peaks every day, lower load at night, and much lower load during  weekends.
 Conversation  shows a less extreme, but still significant, diurnal pattern.

The peak load of Conversation is 1.7× and 3.3× higher than its average and valley loads, respectively.
The peak load of Coding is 2.8× and 34.6× higher than its average and valley loads, respectively.
This large slack indicates that the LLM inference servers can frequently operate in a less performant but energy-optimized configuration  without violating the SLO.
Once the load starts building up, the server needs to switch to a more performant mode of operation.
       




LLM service SLO and model diversity
Finally, different services may time-share the same LLM model instance <cit.>.
They may have different SLOs, requiring the 
configuration to be adapted based on the current service-user.
On the other hand, the same service may concurrently use multiple different models <cit.>. This requires different execution plans for the optimal energy consumption of the individual queries.
Thus, it is not trivial for  service providers to operate in an energy-optimal setting while meeting the performance SLOs.




Insight #2
LLM workloads are highly dynamic and, thus, an energy-optimal configuration  can quickly become sub-optimal.
However, the complexity of a large search space requires an automatic and user-transparent configuration selection.



 §.§ Reconfiguration Overheads

To capture the fast changes
in LLM inference workloads, we need to quickly transition between configurations.
However, there are overheads to change (1) number of inference server instances, (2) model parallelism, and (3) GPU frequency.




Changing instance number
To adjust to fluctuating load, it is cost-beneficial to dynamically adjust the number of LLM instances to serve the requests (i.e., scale in and out).
However, the overheads of adding a new inference server are too large to be tolerable on the critical path of inference loads.
<Ref> shows the breakdown of the overheads to:
(1) instantiate a new GPU VM in the cloud (such as H100 VM <cit.>),
(2) initialize the distributed multi-GPU environment (e.g., Ray, MPI),
(3) download the model weights,
(4) setup the inference engine, and
(5) install the weights and key-value cache on the GPUs.
In total, these overheads can take even a few minutes.
Hence, the conventional LLM inference environments typically provision the static number of instances to handle their peak load resulting in heavy underutilization.
In , we will propose techniques to efficiently scale the number of instances (with the current load) while minimizing most of the scale-out overheads.







Changing model parallelism
To modify the model parallelism of an LLM inference server, 
we need to perform two operations.
First, we need to re-shard the model weights and transfer them to the memory of the right GPUs.
Second, the inference engine needs to synchronize the involved GPUs.
Current systems stop the engine, unload the weights from GPUs, load the weights from the host to the new set of GPUs, and re-start the engine 
from the scratch.
This adds intolerable overheads (around 1-2 minutes) if performed on the critical path.
In , we will show how to minimize the re-sharding overheads by smartly mapping the logical to physical GPUs, exploiting inter-GPU direct NVLink connections and moving the weights between GPUs in the background.





Changing GPU frequency
Setting the GPU frequency (e.g., via  <cit.>) incurs non-negligible overheads.
It involves invoking the OS, communicating with the GPU driver via system calls, and performing hardware interactions via firmware.
On average, setting the GPU frequency takes around 50-80ms.
In comparison, one decode iteration of the LLM inference process takes 20-30ms.
Consequently, the time spent adjusting the GPU frequency can significantly impact the overall performance, potentially doubling the latency of individual inference steps.
<Ref> shows the throughput for different request types when constantly running at the highest frequency (1980 MHz) and when re-setting the frequency (to 1980 MHz) in the background on every LLM inference iteration. 
Due to the software overheads, the throughput of LLM inference system significantly drops.
Therefore, optimizing or minimizing frequency changes during LLM inference is crucial for maintaining efficient and responsive performance.





Insight #3 
Transitioning between LLM server configurations incurs significant overheads.
For energy-efficiency, such overheads need to be minimized and considered when computing the energy/performance trade-offs. 











§ : AN ENERGY MANAGEMENT FRAMEWORK FOR LLM INFERENCE CLUSTERS




We use the insights to design , the first energy management framework for LLM inference environments. 
 seamlessly integrates with existing inference platforms, enabling LLM workloads to operate energy-efficiently and cost-effectively while meeting their performance SLOs. 
 has four key principles.
First, it is energy-optimized and SLO-aware, leveraging model profiles to automatically select the most energy-efficient configuration for specific LLMs and inference workloads within their SLO requirements. 
Second,  fine-tunes configurations for heterogeneous LLM workloads by dividing cluster resources into instance pools tailored to specific request types. 
Third,  accommodates fluctuating LLM inference loads by dynamically reconfiguring the chosen organization. 
Finally, to ensure frequent and smooth reconfiguration,  minimizes reconfiguration overheads.



Architecture
<ref> shows the  architecture.
The system is organized hierarchically at the cluster, pool, and instance levels. At each level, the controllers
tune their assigned configuration knob,  and communicate their decisions with the controllers from the upper and lower levels.
The controllers use energy-performance models generated in the profiling phase to determine the number of instances, model parallelization, and GPU frequency for an energy-optimized operation given the current system state.
(1) Cluster Manager receives inference requests, predicts their type, and forwards them to the appropriate instance pool.
Additionally, it periodically re-evaluates how many pools and how many model instances per pool are needed based on the system load.
(2) Pool Manager
schedules the requests to model instances in a manner that minimizes per-pool energy consumption. It also periodically checks if its instances need to be re-sharded into a more energy-efficient parallelization.
(3) Instance Manager schedules the requests to the inference engine and periodically checks if the instance's GPU frequency needs to be adjusted.





 §.§ Configuring Instances for Energy-Efficiency





Generating LLM profiles
When deploying their service to , users specify the LLM used by the service and the expected performance SLOs. 
Then, the system characterizes the model and generates its energy-performance profile.
 profiles the model by running loads of different request lengths at different model parallelisms (TP2, TP4 and TP8) and GPU frequencies (800–1980MHz, with a step of 200MHz).
The system profiles a few load levels, up to the maximum throughput,
and then extrapolates the behavior for the loads in between the measured ones.
The profiling result is a function that takes the load, request length, model parallelism and GPU frequency as inputs and outputs the expected energy consumption and TTFT/TBT latencies.

As many services may use the same model,  can reuse the profiles across services, minimizing the profiling overheads.
Such profiles are stored in a global  repository, and then cached in a cluster-local storage when a given service is deployed in the cluster.


Selecting the energy-optimized configuration
Given the current load and available resources,  uses the generated profiles to minimize energy consumption while staying within performance constraints.
The system formulates this task as an optimization problem for the mixed integer linear programming (MILP) solver.
The solver needs to output how many instances of each tensor parallelism (N^TP_2, N^TP_4 and N^TP_8) are needed, at which frequency they should run (f^TP_2, f^TP_4 and f^TP_8),  and which load should be assigned to each instance (L^TP_2, L^TP_4 and L^TP_8).
We assume that all instances of a given parallelism run at the same frequency and receive fair-share amount of work.

The optimization target of the solver is to minimize the total energy consumption (E),
while the constraints are:
1) the total number of GPUs used by all instance types does not exceed the assigned number of GPUs (N);
2) the load assigned to individual instances sums up to the total expected load (L);
and 3) the expected performance of all instances with the assigned load is within the requirements (SLO).
Functions Energy^TP_i,f_i(L^TP_i) and Performance^TP_i, f_i(L^TP_i) output the expected energy and performance, respectively, when running the load L^TP_i with TP_i parallelism at f_i GPU frequency.
Then, the optimization task can be formulated as:


    min     ( ∑_i (N^TP_i× Energy^TP_i,f_i(L^TP_i)) )   ∀ i ∈{2, 4, 8}
    s.t.     ∑_i i × N ^ TP_i≤ N 
       ∑_i (N^TP_i× L^TP_i) ≥ L 
        Performance^TP_i, f_i(L^TP_i) ≤ SLO


This approach guarantees the energy optimal configuration.
However, it introduces non-negligible overheads (i.e., ∼100s of ms) due to the large search-space for the solver.
Hence, it cannot be used to select the correct system configuration at fine-grain intervals (e.g., every few seconds).
Next we show how to break the task into a hierarchy of subtasks and use an approximation heuristic to reduce the computation complexity.







 §.§ Hierarchical Control for Dynamic Load



 

simplifies computations by assigning specific optimization tasks to individual controllers.
Instead of searching for a globally optimal configuration, controllers set locally optimal values for individual knobs under the constraints imposed by the upper-level controllers and under the assumption that the lower-level controllers operate 
at the highest performance configuration.
This 
allows the controllers to 
operate at varying time scales–from minutes for node adjustments down to seconds for frequency tuning.
The different scales are needed as each operational change involves distinct overheads and energy-efficiency impacts. 


Scale-out/in 
At every epoch (e.g., 30 minutes), the cluster manager
computes the minimal number of nodes per pool that can support the load of a given request type.
The manager
uses a load predictor to forecast the incoming load, PL, for each request type based on its historic data (e.g., via lightweight load templates <cit.>).
Moreover, the manager assumes that all instances
will run at the highest-performance configuration (i.e., TP8 at 1980 MHz).
Then,
if the predicted peak of a given request type is PL, and the maximum load that a node can support for this request type is ML, the manager assigns
[]PL/ML nodes to that pool.
By consolidating the work into a small number of nodes, the system tries to minimize the cost for the user and the idle energy of lightly-loaded GPUs.

Handling fragmentation:
Allocating resources  to handle the peak loads can cause resource underutilization.
If overprovisioning accumulates across pools, the energy efficiency drops.
To prevent such cases,  assigns one instance less to a given instance pool and moves the leftover load to the instance pool of the next larger request type.
The cluster manager uses this information to forward a fraction of the load for a given request type to the appropriate larger instance pool during the next scheduling epoch (e.g., 30 minutes).
In this way, only the instance pool for the largest requests can be overprovisioned minimizing the cluster-level fragmentation.


Shard-up/down
At every epoch (e.g., every 5 minutes), the pool manager decides how 
to split the assigned N GPUs from the cluster manager into correct model parallelism (how many instances to create in the pool) and tensor parallelism (how many GPUs to use for each instance).
The pool manager uses a simplified version of <Ref> assuming that all instances run at the highest GPU frequency (i.e., 1980 MHz).
Thus, the goal is to minimize the energy, while operating with the fixed number of GPUs running at the highest frequency, and controlling only the parallelism knob.

Accounting for the overheads:
 stores the transitioning overheads (scale-out/in, shard-up/down) in an Overhead Table.
This table is integrated with the controllers,
so that when they calculate the energy benefits of new configurations, they can take into account the costs of reconfiguration.
The controllers evaluate whether the energy savings gained from re-configuring justify the associated overheads and downtime.


Reducing downtime:
The reconfiguration cannot occur simultaneously on all instances due to the risk of long downtime.
Instead,  employs a staggered reconfiguration approach, where a subset of instances is reconfigured (e.g., re-sharded) at a time. 
This ensures that while some instances are undergoing reconfiguration, others remain operational to handle ongoing workloads.
The system 
first reconfigures the instances that
have higher
potential energy savings.




Scale-up/down
Finally, at every epoch (e.g., 5s), the instance manager fine tunes the GPU frequency for further energy savings given the assigned model parallelism.
The instance manager uses the performance profile
to first filter-out frequencies that violate the SLO under the current load.
Then, it uses the energy profile
to pick an acceptable frequency that minimizes the energy consumption.




 §.§ Reduced Overheads for Smooth Reconfiguration




To enable frequent reconfiguration,  
proposes a set of techniques to minimize the overheads of (1) scaling-in/out the number of LLM inference servers, (2) sharding-up/down the parallelism of a given instance, and (3) scaling-up/down the GPU frequency of a given instance.



Scaling in/out inference servers
 reduces the overheads of creating a new server instance by implementing several strategies.
First,
it keeps the model weights cached locally within the cluster (shown in <Ref>) avoiding the need to fetch them from a global repository.
Second, 
it
starts VMs from a snapshot with the entire state of the inference engine already initialized,  reducing the boot-up time.
This snapshot includes pre-loaded libraries, GPU drivers and inference engine configurations.
Third, it proactively creates new VMs in the background, outside of the critical path of active workload handling. 
Specifically,  predicts the peak load for the next scheduling epoch and starts the extra VMs before the epoch starts.
By having these VMs ready to go,  can seamlessly offload a fraction of the load to new instances without any noticeable latency impact.





Sharding up/down an instance
To reduce the re-sharding overheads,  optimizes the distribution of weights across GPUs.
We propose two techniques to minimize the data transfers and latency of individual transfers.
First, the system develops a graph matching algorithm that maximizes the amount of weights that remain stationary in their current GPUs.
The algorithm takes current weight distribution and desired tensor parallelism as inputs, and outputs the source and destination GPUs and fraction of weights to be transferred between each source-destination pair.
Specifically, the algorithm constructs a bipartite graph where nodes represent GPUs in the current and next configurations.
Edges between nodes represent potential transfers, weighted by the amount of data to be transferred.
Then, it applies a maximum weight matching algorithm to find the optimal transfer plan that minimizes the total weight of the edges (i.e., minimizes the amount of data transferred).
Second, to reduce the transfer latency, the system uses inter-GPU direct transfers via NVLink, 
allowing them to send fractions of their weights in parallel to other GPUs
without any host intervention.



<Ref> shows an example of re-sharding from TP4 to TP2 and TP8.
Consider first the case when  going from a lower to a higher-level parallelism (TP4→TP8).
In the initial state (TP4), GPUs 0-3 hold a quarter of the model weights each.
In the final state (TP8), all GPUs need to hold an eight of the model weight.
Thus, the first four GPUs already have their required state, and they only need to send half of their weights to the remaining four GPUs.
The four transfers (e.g., GPU0→GPU4, GPU1→GPU5,…) happen in parallel.
Therefore, the re-sharding  requires the time to send 1/8 of the model weights via NVLink (around 50ms in our setup).

Consider now the case when going from a higher to a lower parallelism (TP4→TP2).
In the final state (TP2), two GPUs need to hold a half of the weights each.
As each GPU initially holds a quarter of the weights, we merge the weights from two GPUs to a single one: GPU1 sends weights W2/W3 to GPU0, and GPU3 sends weights W6/W7 to GPU2.
As these two transfers happen in parallel, the total re-sharding time is the time to send 1/4 of the model weights (around 100ms in our setup).
<ref> shows the re-sharding overheads with different source/destination pairs with our optimized approach.
Some configurations quickly switch to other configurations (transition time 0), other changes incur larger overheads (transition time 4T, where T is the time to send 1/8 of the weights).

Moreover, on some transitions, the instance continues serving the requests with the same throughput.
This is the case when scaling from a smaller to a larger tensor parallelism (e.g., TP4→TP8).
The old instance sends a fraction of the weights to the new instance, without increasing its memory footprint.
During other transitions, the current instance needs to operate under lower throughput.
This is the case when the instance scales from a larger to a smaller tensor parallelism (e.g., TP8→TP4).
Some GPUs used by the old instance accept extra weights, reducing their memory capacity to serve new requests.
In general, whenever the GPU memory required to hold model weights increases, the throughput decreases due to the lower memory capacity for the incoming requests.


Finally, after the weights are sent to the correct memories, the inference engine needs to synchronize the GPUs that run the new instance.
State-of-the-art engines, such as vLLM <cit.>, perform this operation in a few 100s of milliseconds to a few seconds.
During this period the instance cannot receive any load, causing noticeable downtime.
To reduce the downtime, 
 allows the old instance to process the requests while the new instance is going through the synchronization process. 
Only when the new instance comes online, the old instance is removed.
This is possible only when the sum of the memory from the old and new instances is below the GPU's memory capacity.
When the sum exceeds the memory capacity (e.g., TP2 and TP4 with 70B parameters), the old instance needs to be shutdown before the new instance is created.
Overall, different transitions incur different overheads and instance downtime requiring a fraction of the load to be shifted to another instance.
 minimizes the overheads, and considers their impact on the overall efficiency when deciding whether to re-configure an instance.








Scaling up/down the frequency
The overheads of changing the GPU frequency are minimized by keeping the NVIDIA System Management Interface () monitor program directly loaded into memory. 
This eliminates the need to reload the program every time a frequency adjustment is required.
Moreover, by running the controller in privileged mode, 
avoids the overhead associated with OS-user interactions, allowing for rapid frequency adjustments.








 §.§ Predictive Scheduling for Request Heterogeneity




To map the heterogeneity of requests to the heterogeneous instance pools,
the cluster controller in  uses an output-length predictor to anticipate the request type and steer requests to the correct instance pool. 
The predictor acts as a proxy model that takes input prompt and classifies the output as short, medium or long.
Based on the predicted output length and known input length, the cluster manager forwards the request to the pool manager being in charge for a given request type.
If the instance pool is currently overloaded, the cluster manager forwards the request to the next available pool for a larger request type.
Once the request arrives to the correct pool, the pool manager needs to pick an instance from the pool.
Specifically, the manager uses the generated models from the profiling step to 
predict energy and response times 
of each instance after potentially adding a new request to that instance.
Then, it chooses the instance that minimizes total energy while staying within per-instance throughput determined by the SLO.


Handling mis-predictions
If the system
over-estimates a request length, the request gets routed to a higher-performance pool.
Hence, it runs with sub-optimal energy, but its latency remains unaffected. Conversely, if a request length is under-estimated, the request is placed to a lower-performance pool, potentially missing its SLOs.
Similarly,
load mis-predictions can result in insufficient resources for a given pool during request bursts.
In both cases, due to some mis-predictions, the system needs to react to the created emergency event.

When an instance manager detects that its queue is building up, indicating that the rate of request processing is lower than the rate of request arrival, it triggers an emergency event.
First, 
the instance manager tries to re-order the requests in its queue and prioritizes those requests that are about to miss their deadline.
Second, if some requests will miss their deadlines even after the reordering, the instance manager ramps up the frequency of its GPUs to the maximum value.
Third, if the backlog persists or worsens, the instance manager re-steers some requests that have not started their execution.
A subset of requests is moved to another instance within the pool via the pool manager. 
Finally, if all the attempts are insufficient, the instance manager resorts to more drastic measures. One such measure involves squashing requests that have been waiting for processing beyond a specified threshold.
This action signals users to retry their requests, allowing the frontend system to redirect these requests to alternative instance pools or retry them later when system load has stabilized.


 §.§  Implementation



We build  on top of vLLM <cit.>, a state-of-the-art LLM inference platform.
However, 's modularity allows it to be integrated with other
platforms, e.g., TensorRT-LLM <cit.>, without modifications.
We implement controllers as lightweight gRPC servers  with low memory and compute requirements.
Cluster and pool managers are hosted  in a dedicated VM for robust management.
Instance managers are co-located with the LLM inference engine instances for low communication overheads.
For output length prediction, we leverage a BERT-based proxy model <cit.>, which provides accurate and efficient classification of incoming requests.
For load prediction, we use a template-based approach that uses historical data to model load patterns over a week <cit.>. 
The pool manager employs Python's PuLP library <cit.> for solving MILP.
 models energy and performance using the interp1d function from the SciPy <cit.> Python library.





§ EVALUATION






 §.§ Evaluation Setup



We run our experiments on servers with 8 H100 GPUs <cit.>.
We show the results for Llama2-70B <cit.>, but other models (i.e., Mixtral <cit.>, Falcon <cit.>, BLOOM <cit.>) follow the same trends.
We set the load using production-level traces:
1 hour open-source traces <cit.> and 1-day and 1-week traces for Coding and Conversation from our fleet.
We compare  to five systems.
SinglePool (a state-of-the-practice baseline) schedules all the requests to the common pool of instances running with TP8 at the highest GPU frequency.
MultiPool separates LLM instances in multiple per-request-type pools.
ScaleInst, ScaleShard, and ScaleFreq additionally scale the number of instances in the pool, model parallelism, or GPU frequency according to the current load, respectively.



 §.§ Cluster-Level Experiments



We first evaluate the system on a cluster of GPU servers using the 1h open source production traces for the Conversation service <cit.>.
We provision the baselines with 12 H100 servers to handle the peak load,
while  scales the number of servers according to the current load.





Energy
<Ref> shows the energy consumption of the cluster for the experiment.
MultiPool increases the energy consumption by 20% over SinglePool, because it allocates a larger number of resources while always operating at the highest-performance configuration. Meanwhile, ScaleInst, ScaleShard, ScaleFreq and  reduce the energy consumption by 4.1%, 7%, 19% and 35%, respectively.
ScaleInst/Shard/Freq reduce the energy by configuring one knob but leave substantial space for further savings.
Finally,  synchronously scales multiple knobs to achieve the lowest energy consumption.
We further breakdown the total energy per request type.
<Ref> shows that longer requests (e.g., LL) and highly-popular requests (e.g., ML) consume dis-proportionally more energy than the other types.







Latency
<Ref> shows the 
TTFT/TBT latencies for each system.
By separating request types into different resource pools, MultiPool 
removes the head-of-line blocking effect and reduces the latencies over SinglePool.
Similarly, ScaleShard and ScaleFreq, and  reduce the tail latency.
However, these systems slightly increase the P50 latency by operating in lower-performance modes when there is available SLO slack.
On the other hand, ScaleInst increases the tail latency due to the large overheads of creating a new inference server on the critical path of users' load.
Overall,  reduces the P99 TTFT and TBT latencies by 5.3% and 11.1% over SinglePool, respectively,
while it increases the P50 TTFT and TBT latencies by
11.4% and 7.6%, respectively.






Power
<Ref> shows the
power consumption across the systems for the cluster (right figure) and on average per-GPU (left figure).
Due to operating in energy-efficient modes,  effectively reduces both cluster and per-GPU power. 
 reduces the P50 and P99 power consumption over the baseline by 43% and 9%, respectively.







Frequency changes
<Ref> shows the average GPU frequency over time for
(1) the whole cluster,
(2) the pool serving short requests, and
(3) the pool serving long requests.
Average frequency is always significantly lower than the maximum allowed frequency (1980 MHz) that is used by the baseline.
 effectively accommodates different request types by operating their pools at different frequencies.





Sharding changes
<Ref> shows the number of GPUs used for different model parallelisms (TP2, TP4 and TP8) for the whole cluster and for the individual pools (SL, ML and LL).
The figure also shows the load over time that a given pool experiences.
Different pools operate with different model parallelisms and  efficiently changes the parallelism as the load changes.







 §.§ Sensitivity Studies



Sensitivity to predictor accuracy
We analyze how the accuracy of the prediction models affects the overall system efficiency.
We introduce bounded errors for the output length misclassification and measure the energy consumption with medium load.
<Ref> shows that the impact of the predictor accuracy is modest for both energy and performance.
Compared to an environment with no error, an environment with an 40% error increases the energy consumption by 13% and TTFT by 7.3%.
The reason for robustness to prediction errors is that  can promptly detect mis-predictions and re-configures the knobs accordingly.





Sensitivity to load
We evaluate  with different system loads.
We generate Low, Medium, and High loads with a Poisson distribution for request inter-arrival times.
<Ref> shows the energy consumption of the five evaluated systems with different load levels.
With Low, Medium, and High load,  reduces the energy of SinglePool baseline by 51%, 40%, and 23.4%, respectively.
As the load increases, the energy savings of  reduce, because the system more frequently needs to operate at higher frequencies with higher levels of model parallelism.





Sensitivity to number of pools
<Ref> shows the energy consumption and performance (TTFT) of  with different number of request pools.
Recall that our chosen design has 9 pools.
By adding too many pools (12 or 16), the system gets fragmented, and the idle energy of GPUs results in the overall energy increase.
Reducing the number of pools (2 or 4) prevents the system from fine tuning the frequency and the model parallelism for specific request types. The performance improves by adding moderately more pools because it helps remove head of the line blocking and introduces more resources for execution. 










 §.§ Long Cluster-Level Experiments



We run longer experiments by running the 1-day traces for the Conversation service.
The trace covers all invocations for a subset of the service's instances during a typical work day.
We run the experiment for 24-hours on 11 H100 servers with SinglePool and scale the number of instances based on the load in .
<Ref> shows the energy consumption over 5-minute intervals for the two systems.
 reduces the energy consumption over the baseline during peak hours (when dynamic power dominates), and during the low utilization times (when idle power dominates).
Over the whole day,  reduces the energy consumption by 42%.






 §.§ Large-Scale Simulations



To generalize our insights to large-scale, we develop a discrete-time simulator that simulates the energy consumption of different systems using production traces.
<Ref> shows the normalized energy consumption for the five evaluated systems using 1-week traces for Conversation and Coding services.
 significantly reduces the energy consumption for both types of services.
 operates in higher energy-efficient modes for the Conversation service due to its typically shorter input lengths (ML dominant request type).
On the other hand, the Coding service has deep valleys during the night and weekends.
Thus,  exploits the periods of low load to save energy.
 reduces the energy consumption over the baseline by 47% and 56% for the Conversation and Coding services,
respectively.





 §.§ Cost and Carbon Emission



User cost
 reduces the operational cost for users
by minimizing the number of GPUs and optimizing their energy efficiency.
The number of GPU servers for the week-long experiments reduces from 40 to 24.6 on average (38.5% cost reduction).
Using the current GPU VM pricing <cit.>,
this saves $1362.7/hour.
By reducing the energy consumption,  reduces the associated energy costs by up to 56%.
As energy costs <cit.> are currently substantially lower than GPU costs, this translates to only $4.4/hour savings.



Carbon emissions
The energy consumption translates into the amount of operational CO_2 emissions.
We use the traces of carbon intensity <cit.> for a week-long period from multiple grids and map the carbon intensity to the energy consumption over time for the SinglePool baseline and .
<Ref> shows the operational carbon emissions over time for the two systems for CAISO <cit.>.
SinglePool and  consume 5t and 3.1t/week of CO_2.
These substantial savings (38%) make a step towards sustainable LLM environments.



§ RELATED WORK




Cluster resource and power management
A rich body of work seeks to improve resource efficiency under the SLO constraints through resource management for a wide range of latency sensitive workloads, such as microservices <cit.> and DL workloads, through effective resource sharing <cit.>, dynamic allocation <cit.>, and hardware reconfiguration <cit.>. Others
focus on approaches that enable safe power management and oversubscription <cit.> leveraging workload characteristics <cit.> and system state <cit.>. 


Energy-efficient workloads
Prior works focused on energy-efficiency for CPU workloads <cit.>, and researchers started exploring unique energy properties of GPU workloads <cit.>. 
Recent schemes build on top and 
manage energy and power consumption for DNN inference and training <cit.> through frequency scaling <cit.>, autoscaling <cit.>, and resource partitioning and mapping <cit.>. We show that improving energy efficiency for LLM inference necessitates a comprehensive view of all available knobs.  is holistic framework that dynamically reconfigures all the knobs considering the diversity and dynamism of requests and loads.


Efficient LLM inference serving
Recent works propose approaches to improve LLM inference efficiency through heterogeneous resources <cit.> and platforms <cit.>, memory and key-value cache management <cit.>, and node- and cluster-level scheduling <cit.>. While these studies focus on improving throughput or latency, we show that optimizing energy efficiency for LLM inference exhibits distinct trade-offs between performance, energy consumption, and overheads and thus requires a comprehensive framework.



§ CONCLUSION



We present , the first energy-management framework for LLM inference clusters. 
 exploits heterogeneity in inference compute properties and fluctuations in inference workloads to save energy. 
The system automatically and dynamically configures the energy-optimal organization of the cluster (number of instances, model parallelism and GPU frequency) while performing under performance guarantees.
 reduces energy, carbon emissions and cost to the customer by 53%, 38% and 61%, respectively.







IEEEtranS





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



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



§ ABSTRACT 

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

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





§ INTRODUCTION



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

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



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


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



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


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


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

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

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

In summary, our primary observations are:


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

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

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




§ SCALING REPEATED SAMPLING



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

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

We consider the following five tasks:


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

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


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


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



 §.§ Repeated Sampling is Effective Across Tasks






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

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



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





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


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


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

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



 §.§ Repeated Sampling Can Help Balance Performance and Cost




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

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


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



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



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






§ CHARACTERIZING THE BENEFITS OF REPEATED SAMPLING




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


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





 §.§ Scaling Laws for Repeated Sampling





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

    log(c) ≈ a k^-b

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

    c ≈exp(a k^-b)

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



 §.§ Similarities in Coverage Curves Across Models







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



§ HARNESSING REPEATED SAMPLING REQUIRES PRECISION



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



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







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


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


We use ArmoRM-Llama3-8B-v0.1 <cit.> as a reward model, which currently has the highest reasoning score on the RewardBench leaderboard <cit.> among open-weight models. We report our results in Figure <ref> as we increase the number of samples. While success rate initially increases with the number of samples for all three methods, it plateaus around 100 samples. Meanwhile, coverage continues to increase with the number of samples and exceeds 95%.

In the case of majority voting, this success rate saturation is easily explainable. As the number of samples increases, the proportion of votes allocated to each answer stabilizes, and therefore the success rate plateaus. For some GSM8K and MATH problems, correct solutions are sampled with a probability of 1% or lower (see Figure <ref>), making them a minority of samples. As the number of samples increases, rare correct solutions will appear for more problems, increasing coverage but not the success rate with majority voting. In order to fully benefit from repeated sampling, sample identification methods must be able to solve these “needle-in-a-haystack” cases and identify rare, correct samples.

Given the poor performance of the reward model, it is reasonable to wonder how “hard’’ it is to verify a candidate solution. With GSM8K and MATH, only a sample’s final answer is used for assessing correctness, with the intermediate chains of thought being excluded. If models generated only non-sensical chains of thought before guessing a correct final answer, verification may not be any easier than solving the problem in the first place. We investigate this question by manually evaluating 105 chains of thought from correct Llama-3-8B-Instruct solutions to GSM8K problems, reporting our results in Table <ref>. 
We find that over 90% of the chains of thought that we graded are faithful, even among problems where correct answers are generated infrequently.
These correct reasoning steps indicate that there is signal for a verifier to exploit when identifying correct samples. Interestingly, during this process we also identified one GSM8K problem that has an incorrect ground truth answer (see Appendix <ref>). 
This incorrect GSM8K problem is also the only one that Llama-3-70B-Instruct did not generate a “correct” sample for across 10,000 attempts.








 §.§ Verifiers and Software Tasks: Two Cautionary Tales



Software development tasks can occupy a middle-ground with respect to available verification tools.
On one hand, the ability to execute and test code allows for a higher degree of automatic verification than is possible with unstructured language tasks.
However, tools like unit tests take a black-box approach to verifying a piece of code and are not as comprehensive as methods like proof checkers.
These imperfections in the verification process can lead to false positives and/or false negatives that are important to consider when applying repeated sampling. Below we provide two examples of software verifier imperfections that we encountered when generating our results from Section <ref>.  



  §.§.§ Flaky Tests in SWE-bench Lite


 
When producing our results on SWE-bench Lite, we identified that 11.3% of problems have flaky test suites that do not produce consistent results when running them on the same candidate solution.
These flaky tests occasionally classify even the dataset's ground-truth issue solutions as incorrect. Additionally, the test suites for some issues can be non-determinstic depending on the candidate solution.
For example, two SWE-bench Lite issues involve manipulating Python sets, which are naturally unordered.
The gold solutions for these issues explicitly order the items in the set and pass the test suites reliably. However, some model-generated candidate solutions do not impose such an ordering, and therefore pass the tests on some “lucky’’ runs and not others.
In Appendix <ref>, we list all of the problem IDs where we identified flaky tests. We also report our SWE-bench Lite results from Figure <ref> with the problematic issues removed, finding similar results to our evaluations on the whole dataset. 



  §.§.§ False Negatives in CodeContests


Each problem from the CodeContests dataset comes with a set of input-output test cases used to asses the correctness of solutions. These test cases are more comprehensive than those from earlier coding benchmarks like APPS <cit.>, cutting down on the frequency of false positive solutions that pass all test cases but do not fully solve the described problem. However, the construction of the CodeContests test suites leads to false negative solutions that are correct but fail the tests.

For some CodeContests problems, the problem description allows for multiple distinct correct outputs for a given test input. However, the corresponding test cases do not handle these scenarios, instead requiring that one particular correct output is emitted. Additionally, many CodeContests test cases have been programmatically generated by mutating original test cases from the problem. Some mutated inputs violate the problem’s input specifications (e.g. a mutated input being zero when the description promises a positive integer). 
These malformed test cases can lead to inconsistent behaviour between different correct solutions.

We assess the prevalence of these issues by running each problem’s test suite on the list of correct solutions that CodeContests provides. Of the 122 problems in the test set that have Python3 solutions, we find that 35 problems have “correct’’ solutions that fail the corresponding tests. Since we do not allow models to view all of a problem’s test cases (and their peculiarities), applying repeated sampling to these problems contains an element of “rolling the dice” to generate a solution that is not only correct, but emits the particular outputs that pass the tests.







§ DISCUSSION AND LIMITATIONS



In this work, we explore repeated sampling as an axis for scaling compute at inference time in order to improve model performance. Across a range of models and tasks, repeated sampling can significantly improve the fraction of problems solved using any generated sample (i.e. coverage). When correct solutions can be identified (either with automatic verification tools or other verification algorithms), repeated sampling can amplify model capabilities during inference. This amplification can make the combination of a weaker model and many samples more performant and cost-effective than using fewer attempts from a stronger, more expensive model.

 

Improving Repeated Sampling: In our experiments, we explore only a simple version of repeated sampling where all attempts to a problem are generated independently of one another using the exact same prompt and hyperparameters. We believe that this setup can be refined to improve performance, particularly along the following directions: 


    
  * Solution Diversity: We currently rely on a positive sampling temperature as the sole mechanism for creating diversity among samples. Combining this token-level sampling with other, higher-level approaches may be able to further increase diversity. For example, AlphaCode conditions different samples with different metadata tags.

    
  * Multi-Turn Interactions: Despite automatic verification tools being available when solving CodeContests and MiniF2F problems, we use only a single-turn setup where models generate a solution without any ability to iterate on it. Providing models with execution feedback from these tools should improve solution quality. We are interested in the tradeoffs associated with multi-turn interactions, since each attempt becomes more expensive, but also may be more likely to succeed.
    
  * Learning From Previous Attempts: Currently, our experiments fully isolate attempts from each other. Access to existing samples, particularly if verification tools can provide feedback on them, may be helpful when generating future attempts.


  

Repeated Sampling and Inference Systems: Repeated sampling is a distinct LLM inference workload from serving chatbot requests. Production chatbot deployments place an emphasis on low response latencies, and adhering to latency targets can force a lower per-device batch size and reduce hardware utilization. In contrast, when sampling many completions to a single prompt, a larger emphasis can be placed on overall throughput and maximizing hardware utilization. Additionally, repeated sampling can benefit from specialized attention optimizations that exploit overlaps in prompts across sequences <cit.>. Repeated sampling inference can therefore be accomplished at a lower cost than naively making many parallel requests to a chatbot-oriented API. These cost savings can further motivate choosing to sample many times from a cheaper model instead of fewer times from a more expensive one. 

 

Verifiers: Our results from Section <ref> highlight the importance of improving sample verification methods when tools for automatically doing so are unavailable. Equipping models with the ability to assess their own outputs will allow repeated sampling to be scaled to far more tasks. Of particular interest is applying repeated sampling to unstructured tasks like creative writing, which can require a more subjective comparison between different samples than the pass-fail tasks we consider. An alternative direction to developing model-based verifiers is to design converters that can make an unstructured task verifiable, for example by formalizing an informal math statement into a language like Lean so that proof checkers can be applied.




§ RELATED WORK


Scaling Inference Compute: Methods that perform additional computation during inference have been successful across many areas of deep learning. Across a variety of game environments, state-of-the-art methods leverage inference-time search that examines many possible future game states before deciding on a move <cit.>. Similar tree-based methods can also be effective in combination with LLMs, allowing models to better plan and explore different approaches <cit.>. Another axis for increasing LLM inference compute allows models to spend tokens deliberating on a problem before coming to a solution <cit.>. Additionally, multiple models can be ensembled together at inference time to combine their strengths <cit.>. Yet another approach involves using LLMs to critique and refine their own responses <cit.>.

 

Repeated Sampling: Previous work has demonstrated that repeated sampling can improve LLM capabilities in multiple domains. One of the most effective use cases is coding <cit.>, where performance continues to scale up to a million samples and verification tools (e.g. unit tests) are often available to automatically score every candidate solution. 
Recently, <cit.> shows that repeated sampling is effective when solving puzzles from the ARC challenge <cit.>, observing log-linear scaling as the number of samples increases.
In chat applications, repeated sampling combined with best-of-N ranking with a reward model can outperform greedily sampling a single response <cit.>.
In domains without automatic verification tools, existing work shows that using majority voting <cit.> or a trained model-based verifier <cit.> to decide on a final answer can improve performance on reasoning tasks relative to taking a single sample. Concurrent with our work, <cit.> finds that using the best available sample improves LLM performance on chat, math, and code tasks, sweeping up to a max of 128 samples.


 

Scaling Laws: Characterizing how scaling affects model performance can lead to more informed decisions on how to allocate resources. Scaling laws for LLM training find a power law relationship between loss and the amount of training compute and provide estimates for the optimal model and dataset size given a fixed compute budget  <cit.>.
<cit.> finds scaling laws in the context of the board game Hex, observing that performance scales predictably with model size and the difficulty of the problem.
Interestingly, they also show that performance scales with the amount of test-time compute spent while performing tree search.
Recently, <cit.> observe scaling laws when augmenting LLMs with external retrieval datasets, finding that performance on retrieval tasks scales smoothly with the size of the retrieval corpus.






§ ACKNOWLEDGEMENTS


We thank Together AI for partially sponsoring the compute for this project, as well as Rahul Chalamala and Ben Athiwaratkun for their help managing this infrastructure. We thank John Yang for his advice and support when running our SWE-bench experiments. Finally, we are grateful to Mayee Chen, Neel Guha, Quinn McIntyre, Jon Saad-Falcon, and Benjamin Spector for  their helpful discussions and feedback throughout this project.


We gratefully acknowledge the support of NIH under No. U54EB020405 (Mobilize), NSF under Nos. CCF2247015 (Hardware-Aware), CCF1763315 (Beyond Sparsity), CCF1563078 (Volume to Velocity), and 1937301 (RTML); US DEVCOM ARL under Nos. W911NF-23-2-0184 (Long-context) and W911NF-21-2-0251 (Interactive Human-AI Teaming); ONR under Nos. N000142312633 (Deep Signal Processing); Stanford HAI under No. 247183; NXP, Xilinx, LETI-CEA, Intel, IBM, Microsoft, NEC, Toshiba, TSMC, ARM, Hitachi, BASF, Accenture, Ericsson, Qualcomm, Analog Devices, Google Cloud, Salesforce, Total, the HAI-GCP Cloud Credits for Research program,  the Stanford Data Science Initiative (SDSI), and members of the Stanford DAWN project: Meta, Google, and VMWare. The U.S. Government is authorized to reproduce and distribute reprints for Governmental purposes notwithstanding any copyright notation thereon. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views, policies, or endorsements, either expressed or implied, of NIH, ONR, or the U.S. Government. 

This work was completed with the support of the Clarendon Fund Scholarships.


plainnat





§ SAMPLING EXPERIMENTAL SETUP




 §.§ Lean Formal Proofs

We report results on the 130 questions in the test set of the https://github.com/rah4927/lean-dojo-mew/blob/main/MiniF2F/Test.leanlean4 MiniF2F dataset that correspond to formalized MATH problems. This dataset is derived from the https://github.com/facebookresearch/miniF2Ffixed version of the original MiniF2F dataset created by <cit.>. We sample with a temperature of 0.5 and do not use nucleus sampling. We generated 10,000 samples per problem. We use proofs of the following 5 theorems from the https://github.com/rah4927/lean-dojo-mew/blob/main/MiniF2F/Validation.leanvalidation set as few-shot examples:

    
  * 
    
  * 
    
  * 
    
  * 
    
  * 



Our prompt consists of:

    
  * Few shot examples.
    
  * Header imports present in each problem in the HuggingFace dataset  dataset, an upload of the lean4 MiniF2F dataset. 
    
  * The theorem definition. In order to avoid leaking information about how to solve the theorem from its name, we replace the name of the theorem with . i ∈{1,2,3,4,5} for the few-shot examples and i=6 for the current problem. 


We set 200 as the max token length for the generated solution. To grade solutions, we use the  library with lean version  . We set a timeout of 10 seconds for every tactic step.

[
    colback=gray!5,
    colframe=gray!75,
    title=Few-Shot Example,
    fonttitle=]
Write a lean4 proof to the provided formal statement. You have access to the standard mathlib4 library.

```import Mathlib.Algebra.BigOperators.Basic

import Mathlib.Data.Real.Basic

import Mathlib.Data.Complex.Basic

import Mathlib.Data.Nat.Log

import Mathlib.Data.Complex.Exponential

import Mathlib.NumberTheory.Divisors

import Mathlib.Data.ZMod.Defs

import Mathlib.Data.ZMod.Basic

import Mathlib.Topology.Basic

import Mathlib.Data.Nat.Digits


open BigOperators

open Real

open Nat

open Topology

theorem theorem1

  Int.floor ((9:ℝ) / 160 * 100) = 5 :=

by (

  rw [Int.floor_eq_iff]

  constructor

  all_goals norm_num

)```

[
    colback=gray!5,
    colframe=gray!75,
    title=Example Prompt,
    fonttitle=]
Write a lean4 proof to the provided formal statement. You have access to the standard mathlib4 library.

```import Mathlib.Algebra.BigOperators.Basic

import Mathlib.Data.Real.Basic

import Mathlib.Data.Complex.Basic

import Mathlib.Data.Nat.Log

import Mathlib.Data.Complex.Exponential

import Mathlib.NumberTheory.Divisors

import Mathlib.Data.ZMod.Defs

import Mathlib.Data.ZMod.Basic

import Mathlib.Topology.Basic

import Mathlib.Data.Nat.Digits


open BigOperators

open Real

open Nat

open Topology

theorem theorem6

  (b h v : )

  (h_0 : 0 < b ∧ 0 < h ∧ 0 < v)

  (h_1 : v = 1 / 3 * (b * h))

  (h_2 : b = 30)

  (h_3 : h = 13 / 2) :

  v = 65 :=

by (






 §.§ CodeContests


We report results on the 140 test set questions that do not include image tags in the problem description. We sample with a temperature of 0.6 and a top-p value of 0.95 following the experiments in CodeLlama <cit.>. We generate 10,000 samples per problem. We use two few-shot examples from the training set that are randomly sampled per-problem. We set 1024 as the max token length for the generated solution. We use the same answer comparison function as <cit.> and use the concatenation of public, private, and generated tests to validate correctness of solutions.

[
    colback=gray!5,
    colframe=gray!75,
    title=Few-Shot Example,
    fonttitle=]
Q: Write python code to solve the following coding problem that obeys the constraints and passes the example test cases. The output code needs to read from and write to standard IO. Please wrap your code answer using ```:

Petya has equal...

...

(Rest of question and input/output examples)

...

A: ```
```

[
    colback=gray!5,
    colframe=gray!75,
    title=Example Prompt,
    fonttitle=]
Q: Write python code to solve the following coding problem that obeys the constraints and passes the example test cases. The output code needs to read from and write to standard IO. Please wrap your code answer using ```:

This is the...

...

(Rest of question and input/output examples)

...

A:






 §.§ MATH
 We report results on 128 randomly selected test-set problems. We sample with a temperature of 0.6 and do not use nucleus sampling. We use the fixed 5 few-shot example from <cit.> for each problem.  We generate 10,000 samples per problem. We set 512 as the max token length for the generated solution. To grade solutions, we use the  functions from LMEval <cit.>.


[
    colback=gray!5,
    colframe=gray!75,
    title=Few-Shot Example,
    fonttitle=]
Problem:

If 𝐀 = 2 and 𝐁 = 12, then find (𝐀𝐁).

Solution:

We have that (𝐀𝐁) = (𝐀)(𝐁) = (2)(12) = 24.
Final Answer: The final answer is 24. I hope it is correct.


[
    colback=gray!5,
    colframe=gray!75,
    title=Example Prompt,
    fonttitle=]
Problem:

What is the domain of the function 
    f(x) = (2x-3)(2x+5)/(3x-9)(3x+6) ?
 Express your answer as an interval or as a union of intervals.

Solution:




 §.§ GSM8K
 We report results on 128 randomly sampled test-set problems. We sample with a temperature of 0.6 and do not use nucleus sampling. We use 5 few-shot examples from the training set that are randomly sampled per-problem. We generate 10,000 samples per problem. We set 512 as the max token length for the generated solution. To grade solutions, we follow LMEval <cit.>.


[
    colback=gray!5,
    colframe=gray!75,
    title=Few-Shot Example,
    fonttitle=]
Question: James decides to replace his car.  He sold his $20,000 car for 80% of its value and then was able to haggle to buy a $30,000 sticker price car for 90% of its value.  How much was he out of pocket?

Answer: He sold his car for 20000*.8=$20000*.8=1600016,000
He bought the new car for 30,000*.9=$30000*.9=2700027,000
That means he was out of pocket 27,000-16,000=$27000-16000=1100011,000

#### 11000

[
    colback=gray!5,
    colframe=gray!75,
    title=Example Prompt,
    fonttitle=]
Question: Mary has 6 jars of sprinkles in her pantry. Each jar of sprinkles can decorate 8 cupcakes. Mary wants to bake enough cupcakes to use up all of her sprinkles. If each pan holds 12 cupcakes, how many pans worth of cupcakes should she bake?

Answer:




§ SWE-BENCH LITE





 §.§ Experimental Setup


For our experiments, we use DeepSeek-V2-Coder-Instruct with the Moatless Tools agent framework (at commit ).
We use Voyage AI <cit.> embeddings for retrieval, the default used by Moatless Tools. We make no modifications to the model or framework, using them entirely as off-the-shelf components.

With this setup, we sample 250 independent completions for each problem using standard temperature-based sampling. To determine the optimal sampling temperature, we conducted a sweep on a random subset of 50 problems from the test set, testing temperatures of 1.0, 1.4, 1.6, and 1.8. Based on these results, we selected a temperature of 1.6 for our main experiments. 



 §.§ Test Suite Flakiness


During our analysis, we identified 34 problems in SWE-bench Lite whose test suites had flaky tests. Using the SWE-bench testing harness provided by the authors of SWE-bench, we tested each solution repeatedly: for some solutions, sometimes the solution was marked as correct, and other times it was marked as incorrect. In 30 of these 34 cases, we observed flakiness even on the correct solutions provided by the dataset authors. Table <ref> lists the problem IDs of the 34 instances with flaky tests.



An additional instance, , was flaky on some machines and not others. The authors of SWE-bench were able to reproduce the flakyness; however, we were unable to. Our prelimiary investigation indicates this specific issue is due to unpinned versions of dependencies in the docker environments that run the unit tests.



Here, we include results on a subset with the problems in Table <ref> removed (266 problems). For the full dataset evaluation, on any problem that has flaky tests, we run the test suite 11 times and use majority voting to determine whether a solution passed or failed. For the evaluation on the subset without flaky tests, all baselines we compare against release which problems they correctly solve, so we simply removed the problems with flaky tests and recomputed their scores.









§ SCALING LAW DETAILS




 §.§ Experimental details


To fit exponentiated power laws to coverage curves, we first sample 40 points spaced evenly along a log scale from 0 to 10,000 and remove duplicates. We then use SciPy's <cit.>  function to find the a and b parameters from Equation <ref> that best fit these points.



 §.§ Additional results



In Figure <ref>, we show additional results fitting power laws to coverage curves for an expanded set of datasets and models. 






§ PRECISION DETAILS



To calculate the Majority Vote, Reward Model + Best-of-N and Reward Model + Majority Vote metrics, we use the same 128 problem subsets for both MATH and GSM8K datasets introduced in Section <ref>.
Each problem corresponds to 10,000 samples for each model we test.
For each verification method, we take 100 random subsets of size k and calculate the success rate using each subset.
We report the mean and standard deviation across subsets in Figure <ref>. 
To calculate the Majority Vote answer, we take the plurality answer in each subset.
For the Reward Model + Best-of-N, we take the answer with the highest score assigned by the reward model.
For the Reward Model + Majority Vote metric, we sum the reward model score across all the samples with the same final answer, and use the final answer with the highest sum.




§ GSM8K INCORRECT ANSWER


As discussed in <ref>, we identify that  https://huggingface.co/datasets/openai/gsm8k/viewer/main/test?row=1042a problem in the GSM8K test set (index 1042 on HuggingFace) has an incorrect ground truth solution.


[
    colback=gray!5,
    colframe=gray!75,
    title=Question,
    fonttitle=]
Johnny's dad brought him to watch some horse racing and his dad bet money. On the first race, he lost $5. On the second race, he won $1 more than twice the amount he previously lost. On the third race, he lost 1.5 times as much as he won in the second race. How much did he lose on average that day?


[
    colback=gray!5,
    colframe=red!75,
    title=Answer,
    fonttitle=]
On the second race he won $11 because 1+ 5 × 2 = <<1+5*2=11>>11 

On the third race he lost $15 because 10 × 1.5 = <<10*1.5=15>>15 

He lost a total of $20 on the first and third races because 15 + 5 = <<15+5=20>>20 

He lost $9 that day because 11 - 20 = <<11-20=-9>>-9 

He lost an average of $3 per race because 9 / 3 = <<9/3=3>>3 

#### 3


The mistake is in the second line of the answer: on the third race, Johnny's dad lost $16.5, not $15, meaning he made $11 and lost $16.5 + $5 = $21.5. So, the answer is an average loss of $3.5 per race, not $3 per race (the answer in the dataset).












 Paper Title: MIST: A Simple and Scalable End-To-End 3D Medical Imaging Segmentation Framework.




MIST



A. Celaya et al.


The University of Texas MD Anderson Cancer Center, Houston TX 77030 Rice University, Houston TX 77005


MIST: A Simple and Scalable End-To-End 3D Medical Imaging Segmentation Framework
    
Adrian Celaya1,2 Evan Lim1 Rachel Glenn1 Brayden Mi1 Alex Balsells1,2 Tucker Netherton1 Caroline Chung1 Beatrice Riviere2 David Fuentes1

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


§ ABSTRACT 
Medical imaging segmentation is a highly active area of research, with deep learning-based methods achieving state-of-the-art results in several benchmarks. However, the lack of standardized tools for training, testing, and evaluating new methods makes the comparison of methods difficult. To address this, we introduce the Medical Imaging Segmentation Toolkit (MIST), a simple, modular, and end-to-end medical imaging segmentation framework designed to facilitate consistent training, testing, and evaluation of deep learning-based medical imaging segmentation methods. MIST standardizes data analysis, preprocessing, and evaluation pipelines, accommodating multiple architectures and loss functions. This standardization ensures reproducible and fair comparisons across different methods. We detail MIST's data format requirements, pipelines, and auxiliary features and demonstrate its efficacy using the BraTS Adult Glioma Post-Treatment Challenge dataset. Our results highlight MIST's ability to produce accurate segmentation masks and its scalability across multiple GPUs, showcasing its potential as a powerful tool for future medical imaging research and development.





§ INTRODUCTION

Medical imaging segmentation is a highly active area of research. Since its introduction in 2015, the U-Net architecture has received nearly 90,000 citations on Google Scholar <cit.>. Other architectures and frameworks like nnUNet have achieved state-of-the-art accuracy in several medical imaging benchmarks like the Brain and Tumor Segmentation (BraTS) and Medical Segmentation Decathlon (MSD) <cit.>. Since nnUNet's introduction in 2018, several innovative approaches for medical imaging segmentation have emerged. These innovations include new architectures like vision transformers <cit.> and loss functions like the boundary and generalized surface losses <cit.>. Despite these recent advances, there remains a lack of standardized tools for testing and evaluating different approaches for medical imaging segmentation. For example, several works like <cit.> report superior performance to nnUNet while <cit.> refutes these results. This inconsistency makes it difficult to evaluate and assess claims of state-of-the-art performance for new research in deep learning-based medical imaging segmentation.


To address this inconsistency, we propose the Medical Imaging Segmentation Toolkit or MIST, a simple, modular, and end-to-end deep learning-based medical imaging segmentation framework that allows researchers to seamlessly train, test, and evaluate existing and new methods on their data in a consistent and reproducible way. MIST has a rule-based data analysis and preprocessing pipeline that it makes available to multiple architectures, loss functions, and training parameters. Additionally, MIST has a standardized and flexible evaluation pipeline that computes multiple metrics like the Dice, Haudorff distance, and surface Dice for any user-defined segmentation class. This standardization of several pipelines allows for fair comparisons between methods and their effect on segmentation accuracy for any given dataset.



§ METHODS
 
In this section, we describe the required data format for MIST, the different pipelines in MIST, and the auxiliary features, such as evaluation, postprocessing, and test-time inference. We also provide the specific settings we use to train on the BraTS challenge data.



 §.§ Data Format
 
MIST uses the same file structure as the BraTS challenge datasets <cit.>. That is, MIST assumes that the train and test (optional) data directories are formatted so that each sub-directory corresponds to one patient and that each imaging modality (i.e., T1, T2, or CT) is its own NIfTI file. The labels should all be in a single NIfTI file in the same sub-directory for each patient. Once the dataset is in the correct format, the next step is to prepare a small JSON file containing the details of the dataset. These include information like the imaging modality, naming conventions for the images and masks, labels, and the final classes that we use for evaluation (i.e., labels 1, 2, and 3 for whole tumor). An example of this JSON file is provided in the documentation page https://mist-medical.readthedocs.io/en/latest/getting_started/here. MIST also provides support for converting datasets that are in CSV or MSD format. 



 §.§ Pipelines
 
MIST contains three main pipelines: Data Analysis, Preprocessing, and Training. These pipelines can all run at once with a single command or individually. We also provide several auxiliary pipelines for evaluation, postprocessing, and test-time inference.



  §.§.§ Data Analysis
 
The data analysis pipeline gathers parameters about the dataset including cropping to foreground, target spacing, patch size selection, normalization parameters, and is described further below. The results of this analysis are saved in a  file.



  
Cropping to Foreground Often, zeros or useless background information surround the region of interest. For example, zeros surround the region outside of the brain in the BraTS datasets, and the table and air surrounding the body in CT imaging do not provide relevant information. From a computational perspective, it is desirable to crop the images to the foreground to make them smaller. Additionally, cropping to the foreground also reduces the sparsity of the dataset. MIST crops each image by first windowing it with its 33rd and 99.5 percentile values, then using an Otsu filter to produce a foreground threshold. We then crop the image according to the binary mask resulting from this threshold (i.e., one if the voxel intensity is greater or equal to the threshold). To determine if we need to use cropping, we check to see if, on average, the resulting volume reduction is at least 20%.



  
Target Spacing To determine the target spacing for a dataset, we first look at the distribution of the voxel spacings along each axis. Our initial target spacing is the median spacing in each direction. However, suppose the resulting target spacing is anisotropic (i.e., the maximum and minimum spacing ratio is greater than three), then we take the 10th percentile spacing along the lowest resolution axis in that case. However, if resampling with this selected target spacing results in any image taking up more than 2 GB of memory, then we increase all components of our target spacing by 25%. We repeat this process until none of the resampled images take up more than 2 GB of memory. This last step is important because it conserves computational resources. Larger images take longer to load from disk and require significantly more resources to manipulate (i.e., patch extraction and augmentation).



  
Patch Size Selection We set the patch size for training from the median resampled image size by taking the nearest power of two less than or equal to each dimension in the median image size up to a maximum patch size. The default value for this maximum patch size is 256×256×256. However, users can adjust this maximum patch size based on their knowledge of their computing resources. 



  
Normalization Parameters If the dataset consists of MR or other non-CT images, then our windowing parameters are the 0.5 and 99.5 percentile values of either the entire image or the non-zero values of the image if, on average, the ratio of the number of non-zero to zero-valued voxels is less than 0.2. We use the mean and standard deviation with the same rules regarding the non-zero values for normalization. On the other hand, if we are dealing with CT images, then we compute the 0.5 and 99.5 percentile values, mean, and standard deviation of all of the voxels over the entire dataset corresponding to regions labeled as non-zero in the ground truth segmentation masks.




  
Other Checks In addition to the previous analysis steps, MIST checks to see that the header information in the images and masks agree and, if there are multiple modalities, that the headers in each image agree. If there is a disagreement between the headers, MIST prints a warning with the patient ID to the console and excludes the example from training. MIST also computes equal-weighted class weights for each class based on the amount of voxels belonging to each class over the entire dataset. These class weights are primarily informational but can be used for weighted loss functions. 



  §.§.§ Preprocessing
 
The preprocessing pipeline takes the information from the analysis pipeline to crop the images based on their foreground mask (if called for), reorient to right-anterior-inferior (RAI), resample to the target spacing, and window and normalize the intensity values. This pipeline can also compute the distance transform maps (DTMs) for each class in the ground truth masks or apply bias correction for MR images. Users also have the option to skip this pipeline if, for example, their NIfTI files contain preprocessed images already. Figure <ref> illustrates the workflow for this pipeline.





  
Reorient & Resample After optionally cropping or bias correcting each image, the preprocessing pipeline will reorient each image and mask to the RAI orientation. For resampling images to the target spacing, we use third-order spline interpolation. We apply one-hot encoding and use linear interpolation to each channel for masks. If the current image spacing is anisotropic, we first resample along the low-resolution axis using nearest-neighbor interpolation and then apply the appropriate interpolation along the other two axes <cit.>. We implement these operations using SimpleITK to optimize the speed and computational performance of these operations <cit.>. 



  
Window & Normalize Before normalizing the voxel intensities in each image, the preprocessing pipeline windows each image by either the image's 0.5 and 99.5 percentile or, if we are using CT images, by the precomputed windowing parameters. After windowing, MIST applies z-score normalization with either the image's mean and standard deviation or the precomputed normalization parameters for CT images. Again, in the case that the ratio of non-zero to zero-valued voxels is less than 0.2, then the mean and standard deviation are computed from the non-zero voxels only, and we multiply the resulting images by the non-zero mask (different from foreground mask) to preserve the zero-valued voxels.



  
Distance Transform Maps A DTM represents the distance between each voxel and an object's closest boundary or edge. The values in a DTM are positive on the exterior, zero on the boundary, or negative in the object's interior. As a result, they can provide helpful information during training. Several loss functions and shape-aware networks use DTMs to improve the accuracy of medical imaging segmentation methods <cit.>. MIST allows users to pre-compute the DTMs for each label in the ground truth segmentation masks. Each channel in the resulting DTM represents the DTM for that label (i.e., channel one is the DTM for label one). If a label does not exist in the image, the DTM will be a 3D array with all values being the diagonal distance of the image. Additionally, MIST allows users to output normalized DTMs whose values range from -1 to 1. 



  §.§.§ Training

Once preprocessing is complete, MIST uses a five-fold cross-validation to train and evaluate a model on the dataset. In other words, for each fold we set aside 20% of the data as an independent test set and use the other 80% of the data for training. By default, MIST sets aside 10% of the training data as a validation set, but users can adjust this percentage. Additionally, users have the option to adjust the number of folds and can even specify custom folds if, for example, they are running leave-one-institution-out cross-validation.


MIST's default behavior is to use all available GPUs. When MIST runs multi-GPU training, it uses data parallelism via PyTorch's DistributedDataParallel package <cit.>. Additionally, MIST uses the Nvidia Data Loading Library (DALI) to handle data loading, patch extraction, and random augmentation. This feature further optimizes its computational performance by running most of the operations on the GPUs. 



  
Network Architectures MIST supports six different network architectures. These architectures include nnUNet (default) <cit.>, U-Net <cit.>, Swin UNETR <cit.>, and PocketNet <cit.>. MIST also provides users with options like adding deep supervision <cit.> or variational autoencoder regularization <cit.> to most of these architectures. Other regularization features like L2 and L1 regularization and gradient clipping are available for all architectures. For more details about the supported architectures (and their differences), please refer to our https://mist-medical.readthedocs.io/en/latest/advanced_topics/#network-architecturesdocumentation page. Finally, MIST's modular design allows users to implement their own architectures and integrate them into MIST with minimal effort. 




  
Loss Functions Like with the choice of network architecture, MIST allows users to choose from a variety of loss functions like Dice with Cross Entropy (default) <cit.>, the Generalized Dice Loss <cit.>, and boundary-based loss functions like the Boundary, Hausdorff, and Generalized Surface losses <cit.>. Again, MIST's modular design allows researchers to implement custom loss functions and seamlessly use them within the MIST framework. 



  
Other Training Options Finally, MIST gives users the flexibility to choose different optimizers like Adam (default) <cit.>, stochastic gradient descent, and AdamW <cit.>. MIST also supports several learning rate schedulers, with a constant learning rate of 0.0003 being the default. Additional features including transfer learning are included in our documentation page https://mist-medical.readthedocs.io/en/latest/here.



  §.§.§ Auxiliary Pipelines




  
Evaluation Once training is complete, the evaluation pipeline launches to assess the accuracy of the predictions with respect to the final classes defined by the user in their dataset JSON file. By default, these metrics are the Dice similarity coefficient and the 95th percentile Hausdorff distance (HD95). Other metrics like the surface Dice and average surface distance are also available. The output of the evaluation pipeline is a CSV file called  where each row shows the accuracy for each class and metric for one patient. The last five rows give the mean, standard deviation, median, and 25th and 75th percentiles for all of the metrics and classes. The evaluation pipeline is also available as a stand-alone pipeline that can run on any set of predictions. 



  
Postprocessing Postprocessing is available as an optional stand-alone pipeline. This pipeline allows users to specify which labels to apply a user-specified set of operations. These operations include the removal of small objects, only keeping the top-k components (adjustable value of k), morphological cleaning, and filling holes with a specified value. Once the postprocessing pipeline finishes, the evaluation pipeline launches and compares the change in accuracy to the baseline results. MIST computes this change in accuracy as a weighted average of the specified metrics. 



  
Test-Time Inference Finally, MIST provides a stand-alone inference pipeline that takes as input one or more MIST models and a CSV or JSON file with the paths to test-time data. This pipeline writes test-time predictions to a user-specified output directory. 



 §.§ BraTS 2024 Training Protocols
 
Our choice of architecture is the Pocket nnUNet with deep supervision (two supervision heads) and with residual convolution blocks <cit.>. MIST automatically selects a patch size of 128×128×128. We use eight NVIDIA H100 GPUs with a batch size of 32 uniformly distributed across the GPUs. Additionally, we use L2 regularization with a penalty parameter equal to 1e-5. Our choice of loss function is the Dice with Cross Entropy loss. While the default number of epochs per fold for MIST is 1,000, we train for 15,000 epochs per fold. We use a cosine learning rate schedule with an initial learning rate of 0.001 <cit.>. Within each fold, we set aside 2.5% of the training data (27 patients) as a validation set to select the best model. Once training is complete, MIST runs inference on the challenge validation data using test time augmentation (flipping along each axis and averaging each prediction) for each model. Inference uses sliding windows with an overlap of 0.5 with Gaussian blending (σ=0.125). The predictions from all five models are averaged to produce a final prediction. All other options and hyperparameters are left at their default values. Please refer to MIST's https://mist-medical.readthedocs.io/en/latest/documentation for these values.



§ RESULTS
 


 §.§ Accuracy

We use MIST and the parameters described in Section <ref> to perform a five-fold cross-validation (CV) with the BraTS Adult Glioma Post-Treatment Challenge data <cit.>. Table <ref> gives the results of this five-fold CV in terms of the Dice score and HD95. Tables <ref> and <ref> give the accuracy of the predictions in the validation set in terms of the same metrics and their lesion-wise versions.









 §.§ Computational Performance

In addition to the accuracy of our predictions, we also report the scalability of MIST for different numbers of A100 and H100 GPUs for multiple batch sizes. Figures <ref> and <ref> show these results for A100 and H100 GPUs, respectively. In these figures, we see that MIST effectively scales for multiple GPUs. This scaling is especially evident for the H100 GPUs. It is also important to note that we observe a roughly two times speed up with the H100 GPUs, which is in line with what Nvidia reports in their comparisons between the two types of graphics cards. 







§ DISCUSSION

Our results indicate that models trained using MIST can produce accurate adult post-operative gliomas, achieving median Dice scores of at least 0.9 for all of the segmentation classes in our five-fold CV. We choose to submit the raw output of our MIST pipeline to the BraTS 2024 competition, omitting the use of post-processing. However, future submissions to the continuous evaluation and competitions will explore different post-processing strategies with MIST by using its built-in post-processing pipeline. Indeed, previous BraTS competitions cite replacing small objects belonging to the ET class with the SNFH label as beneficial <cit.>. Additionally, we may see improvements in accuracy by using information from DTMs using boundary-based losses like the Generalized Surface Loss. Finally, different architectures like the multi-grid inspired FMG-Net or vision transformer-based architectures like Swin UNETR are also available on MIST <cit.>. These architectures may also improve our results.

MIST is able to utilize multiple GPUs to speed up training efficiently. This efficient use of multiple GPUs is especially true for H100s, where we see a nearly six times speed up when going from one to eight GPUs. The ability to take advantage of multiple GPUs and compute nodes effectively will be essential to keep training times manageable as larger datasets like TotalSegmentor become available <cit.>. As the size and quality of publicly available 3D medical imaging datasets increase, the creation of foundation models for medical imaging segmentation will be beneficial to speed up training through transfer learning <cit.>. MIST already supports transfer learning with other MIST models. Given this capability and its scalability, MIST can be instrumental in the creation of future foundation models.

In conclusion, our results show that MIST is capable of achieving accurate segmentation results while also scaling well to multiple GPUs by effectively taking advantage of technologies like PyTorch's DDP and Nvidia's DALI loader. MIST's simplicity, accuracy, and scalability have the potential to make MIST a valuable tool in the development of foundation models for medical imaging tasks like segmentation and radiation treatment planning. MIST's modular design makes it easy for researchers to add new loss functions or architectures to its pipelines. This modularity and easy integration into MIST's existing pipelines allow for new research to be evaluated and compared to existing methods fairly and consistently. MIST is still under active development. We encourage researchers and the broader medical imaging community to contribute to its development by using it, providing feedback, or contributing new features. 



§ AVAILABILITY & ATTRIBUTION

MIST is an open-source package (Apache 2.0 license) and is available on https://github.com/mist-medical/MISTGitHub or https://pypi.org/project/mist-medical/PyPI. Please cite this manuscript, <cit.>, and <cit.> if you use MIST for your own publications.



§ ACKNOWLEDGMENTS

The Department of Defense supports Adrian Celaya through the National Defense Science & Engineering Graduate Fellowship Program. This research was partially supported by the Tumor Measurement Initiative through the MD Anderson Strategic Research Initiative Development (STRIDE), NSF-2111147, NSF-2111459, and NIH R01CA195524


splncs04




 Paper Title: Who should I trust? A Visual Analytics Approach for Comparing Net Load Forecasting Models.
































Who should I trust?A Visual Analytics Approach for Comparing Net Load Forecasting Models



Parts of this work were supported by the U.S. Department of Energy Solar Energy Technologies Office and the Office of Electricity Sensors Program, and performed jointly at the Pacific Northwest National Laboratory under Contract DE-AC05-76RL01830 and at the Lawrence Livermore National Laboratory under Contract DE-AC52-07NA27344).

    Kaustav Bhattacharjee1, Soumya Kundu2, Indrasis Chakraborty3, and Aritra Dasgupta1

1Department of Data Science, New Jersey Institute of Technology, USA. Email: {kb526, aritra.dasgupta}@njit.edu 
2Optimization and Control Group, Pacific Northwest National Laboratory, USA. Email: soumya.kundu@pnnl.gov
3Center for Applied Scientific Computing, Lawrence Livermore National Laboratory, USA. Email: chakraborty3@llnl.gov 

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



§ ABSTRACT 
Net load forecasting is crucial for energy planning and facilitating informed decision-making regarding trade and load distributions. However, evaluating forecasting models' performance against benchmark models remains challenging, thereby impeding experts' trust in the model's performance. 

In this context, there is a demand for technological interventions that allow scientists to compare models across various timeframes and solar penetration levels. 

This paper introduces a visual analytics-based application designed to compare the performance of deep-learning-based net load forecasting models with other models for probabilistic net load forecasting. 

This application employs carefully selected visual analytic interventions, enabling users to discern differences in model performance across different solar penetration levels, dataset resolutions, and hours of the day over multiple months. We also present observations made using our application through a case study, demonstrating the effectiveness of visualizations in aiding scientists in making informed decisions and enhancing trust in net load forecasting models.



visual analytics, AI/ML, net load forecasting




§ INTRODUCTION



Comparing multiple computational models for their performance is crucial for enhancing trust in model outcomes because it provides a basis for evaluating the reliability and consistency of each model <cit.>. By assessing how different models perform under various conditions and scenarios, stakeholders can better understand their strengths, weaknesses, and overall effectiveness. Such comparative analysis helps to identify the most suitable model for specific tasks or applications, thereby instilling confidence in the reliability of the chosen model. This becomes more important while predicting the net load of an electric grid, which is defined as the difference between total electricity demand and generation from behind-the-meter resources like solar and distributed generators, and is influenced by various factors such as weather conditions and time of day <cit.>. 

Accurate net load forecasting enables grid operators, policymakers, and energy providers to make informed decisions regarding energy trade, load distribution, and resource allocation. However, the rise of solar energy generation sources in residential settings has significantly impacted the performance of traditional net load forecasting models, highlighting the need for robust time-series forecasting techniques <cit.>. Collaborating with scientists, we developed a deep-learning model that integrates variables such as temperature, humidity, apparent power, and solar irradiance to achieve strong predictive performance and resilience in the face of missing data <cit.>. However, the model's sensitivity to seasonal variations and noisy inputs underscores the need for further exploration with domain experts. Hence, we developed an interactive tool that empowers users to investigate the model's performance across diverse time periods and input scenarios <cit.>.





However, while we initially recognized the importance of comparing the model's performance with traditional models, feedback from domain experts further emphasized its significance in building trust in model outcomes. Visual analytics can play a pivotal role here, as evidenced by prior research demonstrating its importance in enhancing trust in machine learning models <cit.>.

This is exemplified by the interactive tool we developed, which enabled domain experts to extract valuable insights concerning the model's sensitivity toward temperature and humidity. Moreover, recent discourse, as highlighted in <cit.>, emphasizes the critical role of visual analytics in fostering trust-augmented applications of artificial intelligence and machine learning (AI/ML) within the energy sector. Building upon this, we designed  , incorporating carefully selected visual analytic interventions. These interventions facilitate the comparison of multiple models across various parameters, including solar penetration levels, dataset resolutions, and different hours of the day, enhancing stakeholders' confidence in model performance.



The aim of our application is to build trust through model comparison, primarily comparing our net load forecasting model with a reference model. We enhanced our model to process inputs at varying resolutions and then devised the reference model, which generates predictions by averaging net load ground truths for the last 30 days at the same time point. Despite lacking predictive ML components, this reference model serves as a benchmark in various net load forecasting competitions, including the recent Net Load Forecasting Prize by the National Renewable Energy Laboratory (NREL) and the U.S. Department of Energy Solar Technologies Office (SETO) <cit.>. Through the web interface of , we were able to uncover patterns in the performance that help improve trust in the model outcomes. In this work, we identify the visual analytic tasks that are required to compare these models. These tasks can be extended to compare other similar net load forecasting models, enabling users to make well-informed decisions based on the outcomes of these models.

While previous research has predominantly focused on developing interfaces during the model development phase, tailored to assist model developers, our emphasis lies in the post-hoc evaluation of model performance, specifically addressing the needs of energy scientists and grid operators <cit.>. Other studies have explored the performance of probabilistic net load forecasting models through different visualization charts but have often lacked an integrated interactive interface <cit.>. In contrast, our interactive visual analytic tool provides carefully designed visual cues for comparing model performance across different factors like dataset resolutions and solar penetration levels, thus offering a novel approach in this domain.

In this paper, we first introduce the different models used in this work and the rationale behind choosing them. Subsequently, we outline the identified visual analytic tasks and detail the design decisions guiding the development of our interactive interface. This is followed by some of the observations made by our power scientist collaborator through that demonstrate its efficacy in comparing model performance across multiple facets. Finally, we conclude by sharing insights gained from this development and discussing some of the future research opportunities in this domain. Additionally, a short demonstration video is available https://bit.ly/trustvishere.



§ MODEL DESCRIPTION




We start with a deep learning-based probabilistic model tailored for net load forecasting in high behind-the-meter solar scenarios <cit.>. This model has three key components: a kernelized probabilistic forecasting (kPF) module, an autoencoder (AE), and a long short-term memory (LSTM) network. This model effectively captured complex temporal dependencies and uncertainties inherent in net load data, which is crucial for reliable forecasting in environments with high solar penetration. By incorporating kernel methods into probabilistic forecasting, the model handled non-linear relationships and captured subtle variations in net load influenced by solar energy fluctuations. The autoencoder component enhanced feature extraction and dimensionality reduction, facilitating the LSTM network's ability to capture long-term dependencies and predict future net load values accurately. Experimental results demonstrated the superior performance of this model compared to traditional forecasting models, showcasing its efficacy in addressing the challenges posed by high solar scenarios and advancing the state-of-the-art in net-load forecasting methodologies. 

This model was used in the Net Load Forecasting Prize competition hosted by NREL and SETO <cit.>. However, during this competition, we observed its underperformance on the data provided by the organizers, which had lower resolutions. We discovered that while the autoencoder component excelled with high-resolution datasets (such as 15-minute intervals), it struggled to effectively capture temporal dependencies in lower-resolution datasets (e.g., 1-hour intervals provided by the organizers). Consequently, this limitation led to subpar outcomes generated by the LSTM component. In light of this, we opted to remove the kPF and autoencoder components and instead developed a version of the model solely utilizing the LSTM component. Additionally, fine-tuning the number of layers in the LSTM component yielded significantly improved results in the competition.

The competition employed a reference model to assess model performance across various probability levels. As previously mentioned, this reference model simply utilizes historical input data from the past 30 days to generate probabilistic forecasts for a specific time point. This model can serve as the initial benchmark for assessing the effectiveness of other models. Therefore, in this work, we developed a reference model following the same principles. Since these forecasts are probabilistic in nature, we calculated the Continuous Ranked Probability Score (CRPS) for both the reference model and our model. Subsequently, we computed the Continuous Ranked Probability Skill Score (CRPSS) based on these CRPS scores, evaluating whether our forecast presents an improvement or deterioration compared to the reference forecast <cit.>. A positive CRPSS indicates that the forecast outperforms the reference forecast, whereas a negative value suggests inferior performance. Utilizing these CRPSS values, we compare the performance of our model against the reference model across multiple dates throughout the year, and present these values in . 



§ VISUAL ANALYTICS-BASED DESIGN


Our application employs a visual analytics-based design featuring coordinated views enhanced with visual cues to help users compare model performance. It combines interactive visualization with comparison metrics like CRPSS to improve trust in the model outcomes and also allows the users to probe the model and understand its performance across different solar penetration levels and months. In this section, we highlight the two tasks performed by and how its visual analytics-based design aids in executing these tasks:
T1: Compare model performance across different solar penetration levels and data resolutions: Model performance may vary across different solar penetration levels due to the increased variability in net load data caused by intermittent solar generation. LSTM-based models might struggle to capture and predict these dynamic behaviors accurately. Furthermore, datasets with higher resolutions, such as sub-hourly intervals, enable models to more effectively capture short-term fluctuations and dependencies. Hence, this task essentially involves comparing model performance at different solar penetration levels and data resolutions.
T2: Identify patterns across different timeframes: By analyzing performance across multiple timeframes, power scientists can assess the net load forecasting models' robustness and consistency in capturing both short-term fluctuations and long-term trends. This evaluation helps to identify whether a model's performance is consistent across various temporal scales or exhibits variability or biases at specific time periods. This task relates to identifying patterns in the model's performance across different months, various hours of the day, and different time periods.

Our application implements multiple coordinated views and components in order to fulfill these tasks. Next, we discuss the design of these views and components along with the rationale behind them:
Comparison View: As comparing model performance is the main objective of , we begin with the Comparison View, which utilizes modified box plots to compare models. Figure <ref>a depicts CRPSS values on the y-axis and different data resolutions on the x-axis. The box plots illustrate the distribution of CRPSS values, with the median typically exceeding zero, indicating superior performance of our model over the reference in most cases. Users can utilize the solar penetration level filter to compare performance across various levels (20%, 30%, 50%) (Figure <ref>c) (T1). However, based on initial feedback from domain experts, we recognized the importance of displaying the distribution of CRPSS values. Therefore, we integrated dots representing the CRPSS values, jittered along the x-axis to illustrate their distribution over the entire year. These dots align with insights from the box plot and additionally reveal instances where our model performs worse than the reference on certain dates. This led to the development of the Patterns View, where we can identify the timeframes where the model underperforms compared to the reference.

Patterns View: In this view, utilizes a heatmap to depict performance patterns for each month across different hours of the day (T2). The x-axis represents the 24 hours of the day (0-23), while the y-axis displays the months of the year (Feb - Dec) (Figure <ref>b). Each box in the heatmap denotes the average CRPSS value for each month at each hour, indicated by the color. Darker blues signify more positive CRPSS values, indicating the superior performance of our model compared to the reference at that time. Conversely, darker reds indicate more negative CRPSS values, signifying poorer performance of our model compared to the reference at that time. This diverging color scale aids in easily identifying performance patterns across different time points and facilitates the identification of instances where our model does not outperform the reference. Users can filter specific months to focus on particular time periods and analyze performance accordingly (Figure <ref>e). Additionally, based on feedback from domain experts, we implemented the Sidebar component to allow users to select specific date ranges and analyze performance patterns within those ranges.
Sidebar: The Sidebar facilitates user selection of start and end dates to filter results across all views, enabling focus on specific date ranges (T2) for comparing model performance within those periods (Figure <ref>d). Additionally, based on feedback from domain experts expressing interest in comparing model performance across all solar penetration levels simultaneously, offers a comparison mode toggle in the Sidebar. Enabling this mode updates both views to display box plots and heatmaps for all solar penetration levels side by side, aiding in point-to-point comparison and identification of performance patterns across different solar penetration levels.



§ RESULTS FROM A CASE STUDY

The efficacy of an application can be validated if the application is able to perform the intended tasks effectively. In this section, we show the results through a case study that demonstrates how can be used to compare and select net load forecasting models effectively.

This case study involved a power scientist with over 10 years of experience in power and grid systems. With expertise in nonlinear dynamics, large-scale networks, and distributed control, he played a crucial role in developing the model. His main objective was to assess the model's performance relative to the reference model across various time points and solar penetration levels, which are critical factors to consider before deploying it for a project. We informed him that we had integrated CRPSS values for both the model and the reference model for all dates throughout the year. He then accessed through a browser and examined the distribution of CRPSS values across different data frequencies at a 20% solar penetration level (Figure <ref>a). Notably, he discovered that the model performed better with the lower resolution dataset (1-hour), contrary to expectations for an LSTM-based model (T1). Upon further examination, he noted a marginal difference in the median CRPSS between high and low-resolution datasets (0.85 and 0.89, respectively). Therefore, he enabled the comparison mode through the Sidebar, enabling him to assess the model's performance as solar penetration increased. He noted that the model consistently performed well with the higher resolution dataset (15-min) across all other solar penetration levels, confirming the notion that our model excels with high-resolution datasets in high solar penetration scenarios (Figures <ref>b and  <ref>c). Next, he observed that although the median CRPSS was significantly above zero, the minimum value was negative. Upon inspecting the dots adjacent to the box plots in the Comparison View, he observed that while most dots clustered around the median line, there were a few outliers with negative CRPSS values.

Hence, the scientist sought to understand temporal patterns to identify if any specific hour of the day contributed to the negative values. Consequently, he navigated to the Patterns View within to analyze the CRPSS value distribution in the heatmap at a 20% solar penetration level (Figure <ref>d). Consistent with observations from the Comparison View, most heatmap boxes displayed varying shades of blue, indicating superior model performance compared to the reference across most months and hours. On closer look, the scientist noted a box-like pattern in the heatmap, revealing enhanced performance during morning hours (8 am to 10 am) from April to September, followed by a decline during midday (11 am to 4 pm), and another spike in the evening (4 pm to 6 pm) during these months. This pattern suggested that the model effectively captured diurnal variations in net load data and adjusted predictions accordingly (T2). While similar box-like patterns were evident across other solar penetration levels, the intensity of model performance varied with increased solar penetration. This insight holds practical significance for model selection during deployment, as it suggests the potential use of different models or model ensembles tailored to distinct times of the day, leveraging their respective performances on diurnal patterns. Thus, the power scientist was satisfied that could yield valuable insights regarding the model, aiding informed decision-making. A short demonstration video is available at https://bit.ly/trustvisthis link.


§ CONCLUSION

Efficient model selection for net load forecasting plays a pivotal role in energy planning and grid operations. In this work, we delve into this process, integrating visual analytics with input from domain experts to identify key tasks for comparative model selection. Subsequently, we translate these tasks into an interactive interface, enabling users to assess model behavior across various factors such as solar penetration levels, data resolution, and time of day.


Throughout this endeavor, we gained valuable insights into the model behavior and the challenges posed during the multi-way comparison of models. This collaborative effort underscored the need for interactive tools that facilitate the seamless translation of model insights into actionable decisions. We can argue that our application is a first step towards this direction. As a next step, our plan is to incorporate multiple net load forecasting models into the application and integrate additional metrics for effective comparison of their performance.

Looking ahead, we aim to enhance our application in several ways. Incorporating economic planning and analysis options will allow stakeholders to assess the cost-benefit ratio before model selection, thereby enhancing trust in the outcomes. Additionally, as the energy landscape evolves, our application's flexibility will be pivotal in effectively comparing and selecting forecasting models for real-world applications. In summary, our collaborative endeavor has produced a robust tool for multi-faceted model comparison and has paved the way for informed decision-making through visual analytics in energy planning.








abbrv-doi




 Paper Title: Finch: Prompt-guided Key-Value Cache Compression.




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



§ ABSTRACT 
Recent large language model applications, such as Retrieval-Augmented Generation and chatbots, have led to an increased need to process longer input contexts.
However, this requirement is hampered by inherent limitations. Architecturally, models are constrained by a context window defined during training. Additionally, processing extensive texts requires substantial GPU memory. We propose a novel approach, Finch, to compress the input context by leveraging the pre-trained model weights of the self-attention. Given a prompt and a long text, Finch iteratively identifies the most relevant Key (K) and Value (V) pairs over chunks of the text conditioned on the prompt. Only such pairs are stored in the KV cache, which, within the space constrained by the context window, ultimately contains a compressed version of the long text.
Our proposal enables models to consume large inputs even with high compression (up to 93x) while preserving semantic integrity without the need for fine-tuning. 





§ INTRODUCTION

Large Language Models (LLMs), built upon the Transformer architecture, have delivered breakthroughs in numerous applications. 
With their generalization and reasoning capabilities, models such as ChatGPT have revolutionized fields where extensive input prompts 
are necessary for generating precise responses, such as Retrieval-Augmented Generation, Chain-of-Thought, conversational chatbots, and In-Context Learning <cit.>. 

However, the use of LLMs in production is limited by their increasing requests in terms of GPU memory <cit.>.
First, as the computational complexity grows along with the size of the models, their memory consumption increases. 
Second, this issue becomes more pronounced when LLMs process larger inputs, as demanded by their ever-increasing context size. 
Third, the Key-Value (KV) cache mechanism, typically employed by LLMs to speed up the generation process, prioritizes efficiency by retaining and reusing previously computed KV vectors during attention computation, bypassing 
re-calculations at each token
generation step <cit.>. Nevertheless, this solution comes with the trade-off of increased memory consumption.[It has been reported that OPT-175B (with batch size 128 and sequence length 2048) consumes 325 GB of memory, but its KV cache requires 950 GB <cit.>.] 

To offer more efficient solutions to operate these models, it has been proposed to compress input prompts, exploiting the redundancy in natural language <cit.>. By preserving critical token information while compressing less crucial details, these models reduce the context in a compact description, without noticeably 
degrading the functional accuracy <cit.>. Compression also enables the LLMs to process large inputs that do not fit the model's context size. However, most of these models require a training/fine-tuning process or a large number of calls to an external model for the compression <cit.>.

We revisit the LLMs' generative inference mechanism to deal with the memory constraint problem and the limitations of current solutions in processing large inputs. We propose a novel approach targeting the reduction of the KV cache memory footprint while avoiding resource-intensive retraining or fine-tuning processes. Drawing insights from the patterns inherent in attention modules, and guided by the understanding that not all attention modules engage with every token, our solution compresses the cached vectors, leading to a reduction in memory usage and efficient text generation.





Our approach, termed Finch,[Finch is a small and quick bird, known for its chirp – a complex language for a small animal.] facilitates faster generative inference through adaptive KV cache compression in the Prefill stage. Figure <ref> shows how a long document and the input prompt are processed with a model context size that cannot fit the entire input.  At every step, a document chunk is processed. Finch uses the attention information between the prompt and the document chunk to identify the 
most relevant
KV pairs across different layers.

This information then is stored in the KV cache for the processing of the next input chunk. Our approach dynamically selects what to keep in the KV cache's memory, effectively keeping its footprint constrained, until the Generation stage produces the response. 



Finch, 
incrementally feeds the KV cache with the compressed context without any learning or external summarization module; it can be used in a plug-and-play fashion with any decoder-based model. The 
compression rate is specified by setting the target size of the KV cache as an input parameter constrained by the model context size.







Even with high compression ratios, our method ensures that the 
correctness of the model response 
is preserved. 
We test Finch on two popular benchmarks covering tasks in question answering, summarization, code completion, synthetic tasks and few-shot learning. Compared against the original LLM (without compression) over the SQuAD v2 benchmark <cit.>, Finch achieves comparable generation quality at 2.35x compression and 90% of the reference accuracy score at 3.76x compression, while being faster in terms of end-to-end execution times in most cases.
When compared to the state-of-the-art compression method LongLLMLingua <cit.> 
Finch reports the best quality scores in most of the tasks in LongBench <cit.>, both with Llama 2 and Mistral <cit.>. 

Our method achieves a compression range of 2x to 93x across various tasks, consistently outperforming a truncation baseline in most experiments. Remarkably, Finch even surpasses the performance of the LLMs operating with the full, uncompressed context in certain cases. Finally, in question answering tasks, we also include a RAG baseline, and our method outperforms it in 10 out of 12 experiments.



§ RELATED WORK

We position our work w.r.t. two main topics.
First, we discuss strategies for improving computational efficiency, i.e., making LLMs accessible for real-time applications or use on devices with limited resources. 
Second, we focus on attention patterns in LLMs, as our work shows that those contribute significantly towards optimizing the models to handle larger inputs in a limited context size.

Efficiency Improvements in LLMs. Methods targeting the reduction of inference and fine-tuning costs include models' modification, such as quantization <cit.> and model compression <cit.>.
Other efforts enhance model efficiency for LLMs by eliminating redundant input words based on attention scores <cit.> and compressing the input sequence by augmenting the encoding modules with pooling layers <cit.>. Proposed solutions also involve learning to skip layers in the transformer architecture <cit.> or to select the most critical tokens for performance <cit.>.
Other approaches pursue prompt compression, either by limiting the number of tokens that are processed in inference by learning special “compressed” tokens <cit.> or by pruning and merging tokens <cit.>, e.g., learning thresholds for pruning unimportant ones <cit.>. However, some of these strategies require an additional re-training or fine-tuning phase and others have been designed for encoder models and are not well suited for auto-regressive LLMs such as ChatGPT and Llama <cit.>. 
In contrast with such solutions, our approach condenses auto-regressive LLMs input contexts during the Prefill stage by using the caching mechanism without model re-training and even faster inference. Finally, recent methods focus on optimizing the generation stage to improve efficiency <cit.>. We leave to future work the study of how to use our prompt-guided token selection strategy in such approaches.


The Role of Attention. Our work relies on self-attention to make the most relevant information in a context available in a concise manner.
The development of transformer models provoked studies to unravel the underlying mechanisms of self-attention, e.g., heads prominently pay attention to separator and adjacent tokens <cit.>. 
Our solution capitalizes on the attention mechanism structure to heighten inference efficiency by exploring the KV cache for the most important key, value pairs w.r.t. the given prompt.

Related work evaluates the informativeness of lexical units using a language model and drops less informative content for  compression <cit.>, for example by regarding tokens with lower perplexity as more influential in the inference process.
These techniques 
view LLMs as a compressor for world knowledge and work by further compressing information within prompts <cit.>.
In contrast with these solutions, 
our approach instead optimizes the management of the KV cache during the Prefill stage without requiring a separate LLM.
Other approaches look at how to select the most important tokens in the Prefill stage, but, differently from our method that dynamically identifies the most important tokens, they rely on manually defined policies for token selection <cit.>.


Finally, we focus on a plug-and-play solution for existing models, with an emphasis on limited computing resources. This is in contrast with other solutions that demand more devices to handle a very large input context
 <cit.>.





§ BACKGROUND

Self-attention is foundational in transformer models <cit.>, enabling language understanding and generation capabilities. 

Transformers learn the contextual relationships between words or subwords within a sentence.

Central to this mechanism are three types of vectors — Queries (𝐐), Keys (𝐊), and Values (𝐕) — that are learned from the input embeddings. 
 nolistsep
    
    
  * Queries (𝐐): Represent the current word or token being processed, acting as a point of focus.
    
  * Keys (𝐊): Serve as identifiers, highlighting tokens in the sequence relevant to the query.
    
  * Values (𝐕): Correspond to the actual specific information carried by each token.


    Attention(𝐐, 𝐊, 𝐕) = softmax(𝐐𝐊^𝐓/√(d_k))𝐕

The self-attention mechanism computes the dot product of the Query with all Keys to determine their similarity. 
A softmax function normalizes these scores, creating a distribution that determines how much attention to allocate to each token. The output is a weighted sum of the Values.

In several NLP tasks, 
transformers generate a response sequence from a given context/document and a user prompt. Consider a sequence of tokens representing the context  𝐱^cont∈ℝ^n^cont and a sequence of tokens representing the user prompt 𝐱^que∈ℝ^n^que,

which may also include instructions,

the goal is to enable the model to generate a response sequence 𝐲∈ℝ^a. This process can be divided into two stages. 

 

Prefill Stage.
As a first step, both the context and the prompt sequence are concatenated to form the input sequence 𝐱∈ℝ^n, where:


    𝐱 = [ 𝐱^cont;  𝐱^que ],   and   n = n^cont + n^que



This sequence is then embedded into an embedding matrix 𝐗∈ℝ^n × d, where d denotes the embedding dimension and processed through multiple layers of multi-head self-attention and feed-forward networks, which operate in parallel across the sequence length and attention heads. Each attention layer calculates and stores the corresponding Key and Value
𝐊∈ℝ^n × d,   𝐕∈ℝ^n × d
matrices in a cache for the sake of performance for the subsequent Generation stage.
In the transformer architectures, the 𝐊 and 𝐕 matrices encapsulate historical token information. Unlike other components of the transformer (e.g., feedforward layer or layer norm), which process current inputs independently of their past tokens, the 𝐊 and 𝐕 matrices are the only matrices 
that retain information from previously encountered tokens. Caching the Key and Value matrices for every layer for the context eliminates the necessity to recompute them for each new token generated. 


Generation Stage.
In this step, 
the model iteratively generates new tokens. For each new token, the 𝐪^𝐧𝐞𝐰, 𝐤^𝐧𝐞𝐰, 𝐯^𝐧𝐞𝐰∈ℝ^d are produced at each attention layer with the 𝐤^𝐧𝐞𝐰 and 𝐯^𝐧𝐞𝐰 vectors appended to the existing cache keys and values:


    𝐊←[     𝐊; 𝐤^𝐧𝐞𝐰;       ],   𝐕←[     𝐕; 𝐯^𝐧𝐞𝐰;       ]


Self-attention uses this updated cache to compute attention. Thanks to the stored 𝐊 and 𝐕 matrices, the computational complexity is just O(nd) as opposed to the approach without cache, which has a computational complexity of O(n^2d). Finally, logits are generated and used to predict the next token in the vocabulary, 
e.g., with greedy decoding <cit.>.



§ PROBLEM FORMULATION

As discussed, K and V are the only matrices that retain information about previous tokens. 
We can therefore formulate the problem of compression as reducing the size of these two matrices during the Prefill stage and before the actual answer generation takes place. Specifically, we have to find 𝐊 and 𝐕 where 𝐊,𝐕∈ℝ^k × d such that two properties are satisfied:

    
  * Compression: 
    the target tokens size
    
    k 
    of the compressed 𝐊, 𝐕 matrices should be smaller than the sequence length n^cont of 𝐊^𝐜𝐨𝐧𝐭, 𝐕^𝐜𝐨𝐧𝐭∈ℝ^n^cont× d. 
    
  * Information retention:
    the output 𝐲∈ℝ^a using 𝐊,𝐕 matrices is similar to the output ỹ∈ℝ^a obtained using 𝐊,𝐕, 
    expressed as:
    
    min_𝐊̃, Ṽ f(ỹ, 𝐲)

    where f is a distance function and its choice depends on the task at hand. For example, in question answering, the difference between F1 scores for ỹ and 𝐲 might be used.


We also define 
the compression ratio σ as: 
    
    σ = n^cont/k

In this work, we compress the context  𝐊^𝐜𝐨𝐧𝐭, 𝐕^𝐜𝐨𝐧𝐭 matrices, according to the target tokens size k, while conditioning on the user prompt. This decision is driven by the recognition that the integrity of the user prompt – particularly its instructions for an instruction-tuned model – plays a significant role in the answer generation <cit.>. 
Furthermore, in the 
tasks that we address in this work, the prompt is typically much shorter than the context, making its compression of limited value.





§ METHOD

Our approach aims at compressing contexts into a manageable form for LLMs, particularly when faced with extensive documents and the need to maintain computational efficiency. Our methodology is motivated by the following observation: the softmax of self-attention distributes attention across all elements to varying degrees,

effectively capturing a spectrum of contextual relationships in the data. We hypothesize that the "smooth" distribution of attention may include superfluous information for the given prompt at hand.



 §.§ Adaptive Key-Value Cache Compression

As depicted in Figure <ref>, Finch iteratively processes a document segmented into chunks, each evaluated in conjunction with a user prompt, and uses the self-attention to identify which  K,V pairs to keep in the cache. In analogy to the long-term memory involving the capacity to recall words, concepts, or numbers <cit.>, we say that these pairs can act as the semantic memory for the model.

The document is reduced to its significant elements and processed in the Generation stage.







Document Segmentation.  
The transformer input is constrained by a context window defined during training, denoted as n_max. 
Given the user specified target tokens size k for the KV cache, Finch processes chunks using at most m_max= n_max - k tokens.[We ignore the user prompt size in this discussion as we assume it to be much smaller than the input document size.] 

The input document is partitioned into chunks of size m, which value is constrained by m_max.  
At every Prefill step i, for i>1,





the K,V pairs from the previous step i-1 (the compressed chunk) are added into the tokens reserved for the k target tokens.



This process introduces a trade-off between granularity and throughput. Smaller chunks enable finer granularity in processing, which is beneficial for certain tasks as we highlight in Section <ref>. Conversely, larger chunks (up to m_max) enhance throughput by reducing the number of sequential operations required, thus speeding up the Prefill stage. This trade-off 
is crucial for optimizing performance and is examined in our ablation study. 







Prompt-Guided Layer-wise top r position selection. Our method for selecting the top r (relevant) positions is rooted in the analysis of the attention scores across its layers. We take into account the unique role of each layer for the representation of the input, i.e., early layers might focus on syntactic features, while deeper layers might capture more abstract, semantic relationships <cit.>. As a consequence, for each layer of the transformer, we calculate attention scores (the scaled dot-product attention between Q and K) and determine the context per-token relevance of the chunk with respect to tokens in the user prompt. By acknowledging that relevance varies by layer, we ensure a more holistic compression of the document. 
For example, tokens that are relevant in early layers might be not relevant in deeper layers.
This allows our method to preserve a wide spectrum of information 
without redundancy. 




Our method also takes into consideration the inherent positional bias present in the attention mechanism. In particular, causal language models operate in the principle that each token in a sequence can only be influenced by preceding tokens, not by those that come after it. This is visually represented by a triangular matrix in attention mechanism, where the ability of tokens to "attend" to each other is constrained by their position in the sequence. As a result, early tokens in a sentence have a broader scope of attention compared to later tokens. For example, for the first token, its attention score is maximal since it only considers itself, leading to a score of 1. To address the issue that later tokens in the sequence, which could be equally or more relevant to the question, are not overlooked due to systemic bias, we incorporate a normalization step that adjusts the raw attention scores to mitigate positional bias, ensuring that each token's relevance is equally evaluated.


Consider 𝐀^(l)∈ℝ^H × M × N as the attention scores matrix at layer l, with H attention heads. Here, M and N are defined as: 
    M = m + n^que, N = m + n^que + c


where m is the chunk length and c is the current KV cache length.
The compression process involves several steps as visualized in Figure <ref>.



    
  * Sum over heads: Every Head in a transformer attention layer captures various aspects of the data. 
    
    
    
    We sum the attention scores over the heads to aggregate their contributions, The elements A^(l)sum_ij of 𝐀^(l)sum are defined as: 
    A^(l)sum_ij =   ∑_h=1^H A^(l)sum_hij
    ∀  i ∈{1, …, M},    j ∈{1, …, N}

    
  * Extract prompt-guided submatrix: A submatrix is extracted to focus on the attention scores between prompt tokens and the current document chunk, this includes considering the tokens accumulated in the KV cache, which grows with each iteration:
    
    
    A^(l)cont_i,j =    A^(l)sum_m+i,j
    ∀  i ∈{1, …, n^que},    j ∈{1, …, m + c }

Figure <ref> shows how attention scores for the last layer of Llama 2 evolve in the sequential operations.


    
  * Normalization: Attention scores are normalized to mitigate positional bias, adjusting for non-zero attention scores:
    
    A^(l)norm = A^(l)cont·(count(A^(l)cont≠ 0)/m + c)

    
  * Selection of top r position:
    The final step is to select the top r indices based on the aggregated attention scores over the prompt tokens.
    
    
    
    A^(l)agg_i   = ∑_p=1^n^que A^(l)norm_p,i  ∀  i ∈{1, …, m+c}

    
    𝐭 = top-r(𝐀^(𝐥)𝐚𝐠𝐠, r)

    here, 𝐭 is a vector containing indices of the top r positions with the highest attention scores. 
    The parameter r dynamically updates at each iteration based on the chunk size m, cache length c, and compression rate σ. Specifically, the update rule is given by:
    
    r_it+1 = m_it+1/σ + c_it
 where it denotes the iteration. At the final iteration, r corresponds to the target token size k.







Managing the Cache:  
The key, value pairs for the selected top r positions are preserved within the KV cache due to their significant relevance to the user prompt.
This process involves an adjustment to their positional embeddings. To accurately reflect the tokens' relative positions, we draw inspiration from the mechanisms used in Attention sinks <cit.>. 
For example, given a cache sequence [0, 1, 2, 3, 4, 5] and a relevance ranking [3, 5, 0], we prioritize '3' by moving it three positions to the left, '5' by moving it four positions to the left, and '0' by shifting it two positions to the right, while the others are discarded. For Rotary Position Embeddings
 <cit.>, as in Llama 2, this repositioning involves calculating the cosine and sine required for rotating to earlier or later positions in the sequence.



Compression Output: The final cache, composed of 𝐊 and 𝐕, represents the compressed document, which encapsulates its essence in a condensed form and is used in the Generation stage.








 §.§ Complexity Analysis


    To illustrate the computational benefit of our approach, we report a comparative analysis of complexity metrics between the attention-based Vanilla transformer and Finch. We consider Complexity per Layer according to n (total number of tokens), m (chunk size), d (model's embedding dimension), a (output sequence length), Sequential Operations as the number of times the model is invoked sequentially, Cache Growth per Operation as the increment in cache size c with each sequential operation, and Initial Cache Size at the beginning of the Generation stage (0 at the beginning of the Prefill stage).
Table <ref> shows complexities for the Prefill stage. 
For large n, the Vanilla method has a higher computational complexity due to quadratic relations, while Finch introduces sequential operations that scale according to m, hence demonstrating enhanced efficiency and potential for scalability in processing large sequences (m ≪ n).
Table <ref> shows complexities in the Generation stage, comparing the resource usage when synthesizing the final output. Also in this stage, the benefit for Finch come from the reduced size of the initial cache according to the compression ratio σ.



 §.§ Encoder-decoder

Our presentation of the methods is focused on a decoder-only architecture, as it is increasingly prevalent in NLP applications. While our methodology is experimented with decoder-only models, 
it is equally viable for encoder-decoder models that employ a KV cache mechanism. In such scenarios, during the Prefill stage, 
we can prefill the KV cache enabling the concise representation of context within the decoder. Subsequently, in the Generation stage we can feed the question or instructions to the encoder. The decoder then utilizes cross-attention mechanisms to access this information, along with the compressed context stored in the KV cache to generate the answer. 




§ EXPERIMENTAL SETUP

We evaluate Finch using a variety of datasets and NLP tasks, with a focus on its application to the Llama 2 7B-chat <cit.> and the Mistral 7B-Instruct-v0.2 <cit.> models. 
Experiments are conducted 
with 
4-bit NormalFloat 
Quantization and Double Quantization
 <cit.>. Unless otherwise noted, the experiments are conducted in a zero-shot setting.[
Finch's code and datasets are available at <https://anonymous.4open.science/r/context-compression-EAF6/README.md>].
Experiments are structured around three public datasets and four baseline methods.[Details on the inference hyperparameters and on the chunk size m per every dataset are reported in Tables <ref> and <ref>, respectively, in the Appendix.
]



SQuADv2: For an assessment of Finch's ability to preserve quality when compressing according to Equation <ref>, we use short texts that let us run the entire document as input. We use SQuAD v2 <cit.>, a benchmark which includes both questions that can and cannot be answered with the given documents. We measure how our model maintains or improves its accuracy, despite having reduced context, against two baselines. First, we report for Vanilla, the standard model configuration which has access to the full context. Second, a Truncate strategy that reduces the input to the same size used by Finch. Given a budget, we truncate the input after a number of tokens equal to half the reduced context both from the start and from the end, i.e., we take the beginning and the end of the document. 



LongBench: To assess the robustness of our method with long documents and a variety of tasks, we also evaluate on the LongBench benchmark <cit.>. This is a suite of tasks that involve extended contexts, including single-document question answering (QA), multi-document QA, document summarization, few-shot learning, code completion, and  a synthetic task. The tasks span 16 datasets and presents a challenge due to the length of the input texts; for the size of the output, we use the original values in the dataset (see Table <ref> in the Appendix for details). For this dataset, our model is also compared against a third baseline, LongLLMLingua <cit.>, a state-of-the-art method for compression of long input texts. 
For LongLLMLingua, we use phi-2 <cit.> as the compressor and Llama 2 7B-chat, quantized at 4 bits with double quantization, as the generator. Unlike LongLLMLingua, our method does not use an external model for compression.
For question answering tasks, a natural baseline is a Retrieval Augmented Generation (RAG)  solution <cit.>.
In our implementation of RAG, we segment the long text into chunks of 256 tokens each. To identify the most relevant chunks, we calculate the cosine similarity between the embeddings of these chunks and the embedding of the prompt. We use the all-mpnet-base-v2 model from Sentence Transformers <cit.> for generating these embeddings.



Lost in the Middle: A critical challenge for LLMs is the "lost in the middle" issue <cit.>, where models exhibit degraded performance if relevant information is situated in the middle of long contexts. We evaluate the robustness of our compression technique also in their dataset.





§ RESULTS AND DISCUSSION


We discuss five questions over our results.



1. Does Finch's compression preserve the relevant information?
Our evaluation on SQuAD v2 measures how Finch retains pertinent information in a compressed format. We compare the Vanilla approach (Llama 2 provided with full documents), 
Finch constrained to target tokens size k, 
and the truncation strategy. 

We choose five values of target tokens sizes, corresponding to different average compression ratios; we 
obtain the latter by dividing the average number of tokens in the SQuAD tests (document and prompt) by the average number of tokens that Finch uses according to the given target tokens size. Specifically, 384 target tokens corresponds to an average σ of 1.1x, 256 tokens to 1.53x, 192 tokens to 2.35x, 160 to 3.03x and 144 tokens to 3.76x. 



The results in Figure <ref> show that Finch not only consistently outperforms the truncation strategy across all token lengths but also, in certain cases, exceeds the quality performance of the Vanilla approach. This is evident in the F1 NoAns and Exact Match (EM) NoAns scores, where Finch's ability to prevent responses based on irrelevant or non-existent evidence 

suggests that it eliminates extraneous content that could potentially mislead the model.

The overall EM and F1 scores indicate that Finch 
maintains the integrity of the context as it is compressed. Even as the target tokens size k decreases, Finch holds onto essential information, enabling the model to generate accurate responses with significantly less input data. In this dataset, the loss of quality compared to the full context becomes more significant starting with an average compression of 3.7x.




To further illustrate the impact of our compression, we run the “lost in the middle” experiment, where the position of the information to answer the user question changes within the input document. It has been shown that this position has a significant impact on the model's accuracy <cit.>. We compare again our solution against the original Vanilla model on the dataset from the paper reporting this problem.
Results in Table <ref> show that Finch significantly outperforms the baseline across the different positions, with up to 13.3 absolute points gain when the correct answer is in the first document (Idx 0) and the compression ratio is 4x. The results also show that our method mitigates the original “lost in the middle” issue with 9.8 absolute points difference between the best and worst accuracy for Finch, rather than 15.3 points for Vanilla. 














2. How fast is Finch compared to Vanilla self attention?
Analysis of Finch's efficiency, detailed in Figure <ref>, highlights a reduction on the overall time w.r.t. the Vanilla when the chunk size is greater than 128 on Llama 2. This observation aligns with the complexity study in Section <ref>. Although Finch introduces additional sequential operations in the Prefill stage, these are offset by the reduced complexity per layer, which is contingent on the chunk size m rather than the full context size n. This approach allows Finch to handle each chunk with a complexity of O(mcd + m^2d) as opposed to the Vanilla complexity per layer O(n^2d). With larger chunk sizes, Finch demonstrates improved speed over Vanilla self-attention. In the generation phase, the distinction in performance becomes more pronounced, as in Table <ref>. Finch benefits from a smaller initial cache size, which is a function of the compression ratio σ. Such a configuration is advantageous in real-world applications where the response time is key and the volume of text to be processed is substantial.





3. How does Finch perform on documents larger than the model context?
To study how our method handles long input documents, we focus on the LongBench benchmark. As for the SQuADv2 experiment, we set the target tokens sizes and we feed the input document in chunks, while reserving space for the prompt and the output generation.

We compare Finch also against the state-of-the-art compression model LongLLMLingua.[Results for LongLLMLingua are lower than those reported in their paper, where they use larger models such as ChatGPT <cit.>.] As shown in Table <ref> and Table <ref>, Finch outperforms LongLLMLingua across five of the six tasks on Mistral and four out of six on Llama 2. The benefit of our solution is clear with different datasets and compression ratios, with a boost up to 8.8 absolute points of accuracy for question answering w.r.t. the best baseline (Truncate) on Mistral. Experiments on Llama 2 reports similar patterns, with a an improvement up to 6.3 
points over the best QA baseline. 

Finch outperforms also the Vanilla baseline using the full document as input in the model context in 12 of the 18 experiments (overall results across 6 tasks and 3 target tokens sizes) on Mistral and in 15 over 18 on Llama 2. This is remarkable when considering that the compression ratio varies between 2.23x and 93.17x.

The baselines beat our method in 4 out of 6 experiments in the Synthetic task, where all methods report very low results. We explain this by the limits of the LLM with 7B parameters, since the tasks demands deep contextual understanding. 


Finch shows better performance according to increasing target tokens sizes (512, 1000, 2000). In the

question answering tasks, Finch with a compression at 512 target tokens beats Truncate and LongLLMLingua with 1000 and 2000 target tokens, both with Llama 2 and Mistral.





We use the LongBench datasets also to validate our idea that conditioning the compression guided by the prompt is more effective than analyzing the self attention scores on the entire input (prompt and document) or on the document only. Results in Figure <ref> show that over all the six tasks, the prompt guided solution leads to the best quality.



Finally, Figure <ref> shows how Finch outperforms the RAG baseline both on Mistral and Llama 2 at different compression rates in 10 over 12 question answering experiments. Compressing with Finch, using the LLM KV cache, offers superior reliability w.r.t. a RAG solution, which suffers from increased latency and fragility due to its dependency on external retrieval mechanisms.




4. What is the effect of the chunk size?
Figure <ref> shows the impact of the chunk size m, i.e., the number of tokens into which the input context is divided for sequential processing by the model. Results show nuanced effects on quality performance. Larger chunk sizes (1024) yield better performance in single-document question answering, while smaller sizes (256) are more effective in multi-document settings. This can be attributed to the compression algorithm of retrieving a fixed number of top r tokens per iteration. In noisy multi-document contexts, a smaller chunk size enables better discrimination between relevant and irrelevant content, enhancing overall model performance.
Chunk size has also an impact on the execution times. As expected, larger chunks lead to faster end-to-end execution because of the smaller number of iterations. These positive results are especially important for use cases that require longer outputs generated by the LLMs. As the user requires a bigger output, the space available for input processing gets smaller, thus reducing the size of the chunks in the Prefill stage.




5. What is the benefit in terms of GPU memory?
Table <ref> reports the memory consumed by Finch (different compression rates) and the Vanilla model for the NarrativeQA (LongBench) dataset (truncated at n=4096). Results show that our approach delivers a significant reduction in the initial KV cache size at the beginning of the Generation stage. Unlike the Vanilla model, 
Finch achieves substantial memory savings by reducing the required cache size in proportion to the compression ratio, confirming the 
 results in Table <ref>. This benefit enhances model scalability and makes Finch a practical choice for deployment in resource-constrained environments.




§ CONCLUSION AND FUTURE WORK

We have shown how attention can be used to identify and prioritize important information within the input data, effectively reducing the need for truncation.
Finch tackles the limitations of LLMs in processing large inputs, offering a balance between computational efficiency and maintaining high language model quality. Our solution leverages the pre-trained model weights of the self-attention mechanism to provide an economically feasible method for operating LLMs.

As future work, we envision a dynamic threshold mechanism to avoid that a fixed amount of KV states are selected in every chunk of the Prefill stage, exploiting the fact that some chunks are not relevant and can be compressed more. 
Another interesting research question is about the use of the proposed method to compress the generated output tokens. This extension would be especially valuable in settings where the LLM is requested to generate long outputs, such as chain-of-thought reasoning. Our approach could be used to identify the important tokens to preserve in the generation step - this is aligned with results showing that preserving a fraction of the original context is sufficient to obtain high quality generated outputs <cit.>.

Finally, we are interested in studying how cache compression techniques can be extended to structured data, e.g., for replacing the current data retrieval and filtering solution in table question answering <cit.>.

Acknowledgments. We thank the action editor and the reviewers for their comments which improved the content of this work. We also thank Riccardo Taiello for the insightful discussion on complexity analysis.



acl_natbib



§ APPENDIX





















