Paper Title: From Feature Importance to Natural Language Explanations Using LLMs with RAG.






2024
Copyright for this paper by its authors.
  Use permitted under Creative Commons License Attribution 4.0
  International (CC BY 4.0).



MAI-XAI'24: Multimodal, Affective and Interactive eXplainable AI,
  October 19–20, 2024, Santiago de Compostela, Spain










1]Sule Tekkesinoglu[
  orcid=0000-0002-1232-346X,
  email=sule@robots.ox.ac.uk,
  ]
[1]
[1]Oxford Robotics Institute, University of Oxford, Oxford, UK

1,2]Lars Kunze[
  orcid=0000-0001-5302-1938,
  email=lars@robots.ox.ac.uk,
]
[2]Bristol Robotics Laboratory, University of the West of England, Bristol, UK




[1]Corresponding author.





§ ABSTRACT 
  As machine learning becomes increasingly integral to autonomous decision-making processes involving human interaction, the necessity of comprehending the model's outputs through conversational means increases. Most recently, foundation models are being explored for their potential as post hoc explainers, providing a pathway to elucidate the decision-making mechanisms of predictive models. In this work, we introduce traceable question-answering, leveraging an external knowledge repository to inform the responses of Large Language Models (LLMs) to user queries within a scene understanding task. This knowledge repository comprises contextual details regarding the model's output, containing high-level features, feature importance, and alternative probabilities. We employ subtractive counterfactual reasoning to compute feature importance, a method that entails analysing output variations resulting from decomposing semantic features. Furthermore, to maintain a seamless conversational flow, we integrate four key characteristics – social, causal, selective, and contrastive – drawn from social science research on human explanations into a single-shot prompt, guiding the response generation process. Our evaluation demonstrates that explanations generated by the LLMs encompassed these elements, indicating its potential to bridge the gap between complex model outputs and natural language expressions.






  Explainable AI LLMs Subtractive counterfactual reasoning Retrieval-Augmented Generation (RAG) 









From Feature Importance to Natural Language Explanations Using LLMs with RAG
    [
    Received XXX; accepted YYY
============================================================================





§ INTRODUCTION


As we become increasingly reliant on AI applications in our daily lives, it becomes imperative to interact with autonomous decision-making systems in human-understandable terms <cit.>. Regulatory frameworks and standards are also evolving to require transparency and accountability in AI-driven systems <cit.>. The technical necessity for explanations to facilitate model debugging and rectifying potential ethical and legal risks arising from biases and errors has been extensively discussed in the literature <cit.>.

The increasing demand for transparent and interpretable AI models has prompted the exploration of various explanation methods, resulting in significant progress in both inherently interpretable and post-hoc explainability methods <cit.>. While effective, these methods struggle to provide explicit interpretability or intuitive explanations for non-technical users. Consequently, there has been a concerted effort to integrate interpretability approaches with Natural Language Processing (NLP) <cit.>. Recently, Large Language Models (LLMs) have gained acclaim as post-hoc explainers, showcasing their potential to elucidate decisions made by other predictive models. Despite their success, a persistent challenge for these models is the occurrence of object hallucinations—instances where plausible yet incorrect outputs are generated, such as featuring objects that do not exist in the images <cit.>. This raises questions about the fidelity of explanations to the underlying model. To address this issue, we propose a traceable question-answering, which informs LLM responses through an external knowledge source that provides insights into the model's output. By utilizing this data, the LLM acts as a reasoning engine to process the information rather than a source of information. Figure 1 provides a high-level depiction of this process.  

Additionally, to cultivate a seamless conversational experience, we incorporated key elements derived from social science research on human explanations into the system prompt, directing the response generation process. We present our vision for creating this integrated approach, outlining each component, and illustrating its viability through a real-world dataset. This integration has the potential to bridge the gap between intricate model outputs and user comprehension towards advancing the development of human-understandable explanations. 







§ BACKGROUND



 §.§ Recent advances in image explanations

The demand for transparent and interpretable AI models has led to exploring diverse explanation methods for vision tasks <cit.>. Recent years have witnessed significant progress in both intrinsic and post-hoc explainability methods. Gradient-based techniques, which analyse the gradients of the model's output with respect to its input features, have been extensively studied <cit.>. While effective, these methods may fall short in providing intuitive explanations for non-experts. In another line of work, attention mechanisms are proposed by assigning weights to different parts of the input data based on their relevance to the model's decision <cit.>. Similarly, attention mechanisms may lack explicit interpretability, and some argue that attention alone is insufficient. As a result, there have been concerted efforts to integrate these approaches with Natural Language Processing (NLP) <cit.>. Long Short-Term Memory (LSTM) encoders have been employed to extract features from video clips, with LSTM decoders generating descriptive texts based on these features. While effective in image captioning, these methods provide textual descriptions rather than visual explanations. Moreover, researchers have proposed perturbation-based approaches—a model-agnostic method that perturbs input data and observes changes in model predictions. Building on this research, we integrate semantic perturbation analysis with pretrained language models, translating the results into the text modality for enhanced explainability.  



 §.§ LLMs as post hoc explainers

Recent strides in Large Language Models (LLMs), exemplified by GPT-4 (OpenAI), Bard (Google), Claude-2 (Anthropic), and Llama-2 (Meta), mark a transformative era in Natural Language Processing (NLP) research. These models have become pervasive across diverse applications, ranging from machine translation and question-answering to text generation <cit.>. Notably, LLMs have recently gained attention as post hoc explainers, highlighting their potential to elucidate the decisions made by other predictive models <cit.>. Leveraging their in-context learning capabilities by fine-tuning for specific tasks enhances their proficiency in generating contextually relevant explanations <cit.>. Despite their success, a persistent challenge for these models is the occurrence of object hallucinations—instances where plausible yet incorrect outputs are generated, featuring objects that do not exist in the images in vision-based tasks. This issue raises questions about the faithfulness of explanations to the underlying model. One way to address this is by recognizing the role of the semantic map layer in capturing the meaning and context of the physical surroundings <cit.>. By effectively encoding valuable semantic information into LLMs, we can truthfully represent and comprehend the intricate details of the model output within the language space.



 §.§ Characteristics of explanations and XAI


In the context of XAI, the audience for explanations is predominantly human; thus, understanding what makes an explanation human-friendly is paramount. While accuracy is important, the clarity and presentation of explanations are equally crucial in ensuring comprehensibility. According to Miller <cit.>, current work on interpretability in machine learning relies solely on researchers' intuition regarding what constitutes an appropriate explanation for humans. Miller's survey highlights four major characteristics of explanations drawn from findings in social science research on human explanation: explanations are social, causal, selective, and contrastive <cit.>. In the following section, we briefly discuss these characteristics in relation to XAI before demonstrating how we integrate them into our application.



  §.§.§ Explanations are social

An explanation represents an interactive exchange between two roles: the explainer and the explainee, and is governed by certain `rules' <cit.>. Interactive explanations in a dialogue structure consist of three fundamental components: opening statements, clarification questions, and closing statements <cit.>. As such, it should follow the basic rules of conversation, which are captured by Grice's conversational maxims <cit.>. These include principles such as quality, quantity, relation, and manner, which can be understood as `only say what you believe', `only say as much as is necessary', `only say what is relevant', and `say it nicely'. Manner further encompasses various sub-maxims, including avoiding unclear expression, eliminating ambiguity, maintaining brevity to avoid unnecessary wordiness, and presenting information in an orderly fashion. 

Moreover, explanatory agents employing anthropomorphic traits like politeness markers (e.g., “Thank you!”), warm and friendly tones, and empathetic expressions (e.g., “I understand.”) are pivotal for effective human interaction <cit.>. Theory of Mind also plays a crucial role in social explanations.[Theory of mind refers to our ability to understand other people by attributing mental states to them, including beliefs, desires, intentions, emotions, and thoughts that may differ from our own <cit.>.] As an integral part of a dialogue, an explanatory agent should keep track of what has already been explained, maintaining a simple mental model of others. This enables intelligent agents to tailor explanations based on past interactions and evolving contexts.



  §.§.§ Explanations are causal

Causal explanations aim to elucidate the cause-and-effect relationship underlying a decision. When explaining a prediction made by a machine learning model, identifying the most influential features can provide insight into the key causal connections to the prediction. This process, also known as causal inference,  starts with observations (e.g., what if a feature had a different value) and selecting some of those causes as the explanation. Miller divides casual inference into two parts: counterfactuals and abductive reasoning <cit.>.

The reasoner uses a model of hypothetical counterfactual cases to derive an explanation, such as perturbing inputs, to see how they causally impact the prediction (See Section <ref>). Abductive reasoning is inferring the causes from those observations to form the most probable explanation. One can think of abductive reasoning as the following process: observing unexpected or surprising events, generating one or more hypotheses about these events, judging the plausibility of the hypotheses, and selecting the `best' hypothesis as the explanation <cit.>. It is important to note that new information (e.g., different feature perturbations) can alter the most likely explanation.




  §.§.§ Explanations are selected


Explanations are contextual by nature. While an event may have numerous causes, the recipient of the explanation typically cares about only a specific subset relevant to the context <cit.>. The explainer selects a subset based on various criteria called casual selection. Hilton <cit.> argues that explanation selection is used for cognitive reasons since causal chains are often too large to comprehend. The criteria people use for explanation selection include abnormality, relevancy, and simplicity. 

Abnormality is a key criterion for explanation selection; as such, people select unusual causes to explain events <cit.>. In machine learning interpretability, identifying abnormal events corresponds to deviant or higher feature importance values. As pointed out within the conversational maxims, relevancy also plays a role in explanation selection—only say what is necessary and relevant. People select explanations to adhere to these maxims whilst persuading the explainee to new information or viewpoints. Another strong criterion in explanation selection is simplicity, which refers to those that cite fewer causes. People prefer uniquely sufficient causes that bring about the effect without any other causes, but there might be cases with multiple sufficient causes.  




  §.§.§ Explanations are contrastive


Studies have shown that people only request contrastive explanations, essentially `why' questions framed to imply the differences between two possible outcomes <cit.>. As a result, explanations are typically offered in relation to the cause of one event compared to another that did not occur.

Similarly, in the context of ML interpretability, people are not specifically interested in all the factors that led to the prediction but instead in the factors that need to change so that the prediction would also change. Explanations that present some contrast between the instance to explain and a reference point (i.e., a hypothetical instance) are preferable because the cognitive burden of complete explanations is too high. However, most existing work considers contrastive questions but not contrastive explanations, providing two complete explanations individually. This could be because the contrastiveness is application-dependent, given that a reference point may not be apparent in all tasks <cit.>. For instance, in sentiment analysis or anomaly detection, contrastive cases might not be as straightforward as other tasks, such as banking, e.g., loan accepted vs rejected cases.



§ MATERIALS AND METHOD


This section outlines the key elements of the traceable question-answering methodology. We detail the proposed explainability approach for extracting feature importance and the prompting technique to generate accurate and human-friendly explanations. Additionally, we provide an application example illustrating how we establish a connection between the model output, the explainability technique, and language models to produce natural language explanations in an interactive format. 











 §.§ Explanations through subtractive counterfactual reasoning



Subtractive counterfactual reasoning is the process of removing an event to understand its causal influence on an outcome <cit.>. Such reasoning is widely used as an XAI technique, exemplified by occlusion analysis and feature importance attribution, built on the same foundation - quantifying the variation in model output through systematic modifications of input components <cit.>. There are distinctions in terms of how these methods remove features, elucidate model behaviour, and summarize feature importance. Nonetheless, they provide valuable insights into the relative significance of each feature influencing the decision-making process <cit.>. 

This study proposes a decomposition-based approach to examine the output variations by decomposing input values. By measuring such variations, we determine the degree of impact on an outcome and reason about the importance of each feature value. As a function, the model assigns a class label and prediction probability f:x→ f(x) given input y. To measure the effect of each input value, we observe the model's prediction p for y without the knowledge of event T_i=T_k, where T_k is the value of feature T_i. To explicitly represent the absence of information, the feature value T_k is replaced with an undefined value, i.e., NaN (Not a Number). NaN values are treated as invalid or masked elements and are not considered in the computation. We attain an array containing probabilities for all perturbations resulting from decomposition. In order to determine the importance of an individual feature T_i, we begin by identifying the max and min values within this array, which establishes the decision boundaries for the given input. Next, we compute the feature importance as the position of the probability– after decomposing the feature– relative to the max and min probabilities. The computation is expressed as follows:


    importance(T_i) = p(x|y ∖ T_k)_max - p(x|y ∖ T_k)/p(x|y ∖ T_k)_max - p(x|y ∖ T_k)_min


This approximates an importance value between zero and one for each input feature. Considering that evaluating fractions has a higher cognitive load, the values are rounded to the nearest integer to increase comprehensibility. Moreover, this analysis is carried out for the categories following the main prediction simultaneously to generate explanations for the contrastive cases. 





 §.§ Prompting for traceable question-answering with RAG


Traceable question-answering, through a process known as Retrieval-Augmented Generation (RAG), integrates external knowledge sources to furnish responses to user inquiries. This method improves the performance of LLMs by referencing an external knowledge source prior to generating a response (See Figure <ref>). This approach mitigates the problem of `hallucination' and ensures factual consistency, particularly in scenarios where facts may evolve over time <cit.>. Depending on the specific implementation, RAG supports various file formats, including text, PDF files, word documents, structured data (tables), and images. 

In our work, the knowledge repository contains tabular information about model output, features, contrastive cases, and feature importance values obtained through the decomposition-based explanation method described in Section <ref>. This information is saved as the output of the explanation module as a .csv file. Then, the file, the system prompt (See Figure <ref>), and the user input query are fed to the LLM model for response generation. The resulting responses contextualise the feature importance grounded in the models' inner workings, thereby rendering the explanations traceable to the underlying model.   

Furthermore, our aim is to imbue LLM with the characteristics of explanations discussed in Section <ref> to foster a human-friendly interaction. Ideally, LLM-generated responses should adhere to conversational norms, encompassing opening, clarification, and closing statements while appropriately integrating social cues. The model is expected to deduce causal relationships from the observations within the provided knowledge repository, selecting the most pertinent causes to form an explanation and presenting it alongside a contrastive case. To accomplish this goal, we devised a system prompt that provides explicit instructions to the model to generate efficient responses to user queries. The system prompt includes a single-shot example illustrating the desired interaction, potential user inquiries and outlining how the model can effectively address them, enabling the language model to produce consistent and predictable outcomes, as presented in Figure <ref>.


















 §.§ Application example
 

To demonstrate our approach to traceable question-answering, we selected a visual scene understanding task. This involves scene classification, categorizing images based on their content and inferring high-level semantic scenes from low-level visual features. Many critical applications of scene classification include human-robot collaboration, autonomous driving, and other autonomous physical systems that rely on semantic segmentation to inform decision-making processes <cit.>. We experimented with the GoogleNet Places365 model designed specifically for scene recognition and classification tasks <cit.>. The model is trained on the Places365 dataset, containing more than 365 scene categories.  We focus on semantic urban scene understanding with the road segmentation dataset CamVid (Cambridge-driving Labelled Video Database). The dataset helps to understand driving scenes, comprising images capturing street-level views from urban driving scenarios. It includes manual annotations across 32 classes, encompassing elements such as building, driveway, pavement, tree, traffic sign, car, pedestrian, and bicyclist. The Deeplab v3+ deep learning model is trained on the CamVid dataset for semantic segmentation <cit.>. Figure <ref> shows the segmentation result for the example presented in the next section. Since we utilized off-the-shelf models for illustration purposes, the emphasis is not on these models' performances.



Semantic information plays a critical role in autonomous vehicle (AV) decision-making processes, as it captures the meaning and context of the physical surroundings observed in the road scene. This includes understanding the appearance of various elements such as buildings and roads, identifying shapes such as cars and pedestrians, and recognizing spatial relationships to contextualize the scene <cit.>. Particularly in unstructured off-road environments, the accurate understanding of the surroundings concerning semantic classes such as trail, grass, or rock is important for safe and deliberate navigation <cit.>. Moreover, semantic segmentation provides a way to represent and comprehend the details of the environment in the language space. By leveraging the labelled information within the image scene, we can semantically map the crucial components based on their feature importance. This granularity enables a more insightful analysis of feature importance, contributing to a nuanced understanding of the model's decision-making process. Such insights can be helpful for post-incident forensic analysis in case of misreading a scene causing collisions or accidents. In this work, we suppose situations where the user checks in with an AI assistant in an AV system in non-critical driving scenarios while engaging in non-driving related tasks.































Table <ref> reveals the content of the knowledge source - only the first contrastive case is displayed. The effect of removal is the probability when a feature is removed, as listed respectively. Given the system prompt and the data presented in Table <ref>, Figure <ref> showcases the dialogue generated for this case. From a social standpoint, the resulting responses balance between being informative and supportive while maintaining a friendly and engaging tone. It effectively justifies the model's output by emphasising its most significant feature in a concise manner. Additionally, the dialogue delves into contrastive explanations by detailing three consecutive classes of lower likelihood, demonstrating the flexibility to explore different scenarios. When comparing the top prediction to the second most likely class, it highlights the distinguishing features contributing to the output. 













§ EVALUATION RESULTS

In this section, we present our experimental findings. We analyzed nine distinct scene classes — parking lot, street, residential neighbourhood, crosswalk, highway, industrial area, gas station, shopfront, and general store (outdoor)— across different scenarios (totalling 42), where the model displayed varying degrees of confidence in its output. We evaluated the responses in terms of sociability, causality, selectiveness, and contrastiveness generated by GPT-3.5 () and GPT-4 () models. 



 §.§ Sociability


In examining the sociability within the responses, our initial focus centres on the overall tone of the responses. We employed the Valence Aware Dictionary and Sentiment Reasoner (VADER) to gauge the sentiment conveyed within the text, encompassing dimensions of positivity, negativity, and neutrality. Sentiment scores for each response are aggregated to delineate overall sentiment trends. Figure <ref> illustrates a noticeable positive trendline by both models, indicating a prevalent inclination towards positivity in the sentiment expressed. A qualitative inspection of the text reveals a prevalence of polite, helpful, and informative social cues across all responses, supporting this observation. Consequently, the neutral sentiment score remains notably low, underscoring the prevalence of emotional expression throughout the corpus. Exploring negative sentiment scores unveils instances where the models acknowledge their limitations, particularly when unable to retrieve answers from the external knowledge repository. 

 
Label Positive Neutral Negative

GPT3.5   202  5  8
GPT4  195  6  14











Considering the three foundational phases of dialogue—opening, clarification, and closing statements—next, we analyzed the indicators such as politeness markers and emotional expressions within each phase. To facilitate this analysis, we created a distinct dictionary for each phase and computed the term frequency of social cues present in each phase.[Term frequency is defined as the number of times a term occurs in the document over the total number of terms in a document.] Opening statements include greetings such as `hey' and `hello.' Clarification statements are characterized by social cues indicative of helpfulness, continued support, and awareness of its limitations, including terms such as `absolutely', `sorry', `further', `curious', `questions', `help', `feel', and `free'. Closing statements commonly conveyed sentiments of well-wishing or safety, with phrases including terms such as `welcome', `enjoy', `safe', `great', `pleasant', and `glad'.




Figure <ref> shows that both models utilize these terms fairly across the respective phases of the dialogue. Our qualitative assessment of the responses further confirms the contextual relevance of these terms, which align with conversational maxims (See Section <ref>). The higher frequency value observed in closing statements can be attributed to their typically brief nature, which encapsulates most of the social cue terms. Additionally, both models occasionally repeat certain phrases in their responses, albeit with the application of a frequency and presence penalty (0.3).[Frequency penalty discourages repetitive use of words or phrases, while presence penalty promotes diversity in the generated text.] In addition to analysing social cues, we examined part-of-speech (POS) pronouns as social interaction markers, encompassing terms such as `I', `you', `we', `my', `our', and `your.' Quantitative analysis and qualitative evaluation of the results indicate that both models adequately and appropriately applied POS pronouns.


 


 §.§ Expression of Causality


To evaluate the expression of causality, we initially extracted responses indicating cause and effect, primarily found within clarification statements. We then identified key terms related to causality, such as `because', `if', `then', `albeit', `due', `contribute', `influence', `affect', `impact', and `effect', and tallied their occurrences within each answer to measure the presence of causal reasoning. In Figure <ref>, we compare the frequency of these terms used by two models in the response set. Notably, `if' emerges as the most frequent term used by both models, often utilized for expressing counterfactual reasoning (e.g., `if it was not', `if there is/were'), while `influence' typically indicates the cause of a decision. Overall, GPT-4 demonstrated a higher frequency of causality terms and showcased greater diversity in expression compared to GPT-3.5.




Through the qualitative review, we examined existing patterns of dependency parsing signifying causal relationships. Two primary grammatical structures emerged, suggesting cause-effect relationships within the text. The first structure pertains to how the presence or absence of certain features influences the likelihood of specific scenarios. The second structure involves a cause-effect relationship between feature importance and probability, indicating that features with higher importance scores exert a greater impact on probability. It is important to acknowledge that while these grammatical structures and terms are widely used, their accuracy is not always consistent across models. In the subsequent section, we will further evaluate the selection and accuracy of inferred causal relationships.



 §.§ Selectiveness


In this section, we evaluate whether the models effectively detect unusual values within the provided knowledge repository, infer potential causes from those observations, and select the most relevant ones to construct an explanation. Our evaluation encompasses analysing responses across the number of causes given, graded selection, and sort order to assess whether expected causes are addressed correctly. First, we specified the maximum number of uniquely sufficient causes for each case, which are feature importance values >5. The number of causes varies across scenarios; some exhibit only a single highly important feature >5, while others might involve multiple high-importance features with varying degrees of significance. Graded selection is crucial for discerning important values such that the language model must prioritize the most relevant causes when encountered with multiple features of comparable importance levels. Additionally, graded selection entails disregarding features with lower importance scores (<5). Finally, the sort order concerns the sequence of selected causes, which should follow a hierarchy from the most significant and relevant features to those of lesser importance. Thereafter, responses are categorized as fulfilled, partial, subpar, or unfulfilled according to their arrangement with our specifications. 



The findings presented in Table <ref> indicate that GPT-4 outperformed GPT-3.5 in all categories. Notably, GPT-4 had a much higher success rate (76%) than GPT-3.5 (24%) in identifying the most relevant cause (fulfilled). Additionally, GPT-4 had a lower rate of subpar and unfulfilled responses. GPT-3.5 referenced features not present in the knowledge source as potential causes in two instances. These results suggest that GPT-4 is more accurate and selective than GPT-3.5 in identifying the most relevant cause(s) from a set of possible causes. 

Furthermore, an additional criterion for assessing explanation selection is simplicity, which involves utilizing concise language, avoiding technical terminology, and referencing fewer causes. Table <ref> evaluates responses by comparing three metrics: the frequency of technical jargon, response length, and the number of causes cited. Lower values across all three metrics signify simpler responses.



To measure the term frequency of technical jargon, we compiled a list of technical terms in the responses and calculated their frequency within the generated explanations. The list contains terms such as `prediction', `feature importance', `score', `contrastive cases', and `confidence', which may not be easily comprehensible to all users. Table 3 indicates that GPT-3.5 exhibits a slightly higher term frequency (0.0162) compared to GPT-4 (0.0157), suggesting that GPT-3.5 utilizes technical jargon marginally more frequently. The length metric refers to the median lengths of the responses generated by each model, which are quite similar (45 tokens for GPT-3.5 and 44 tokens for GPT-4). In terms of the number of causes cited, GPT-4 cites fewer causes (median of 2) than GPT-3.5 (median of 3), suggesting that GPT-4’s responses tend to be simpler in terms of causal reasoning. Overall, the results suggest that GPT-4 generates simpler responses than GPT-3.5 across all three assessed metrics. 

Regarding the out-of-scope questions posed, such as the number of cars or pedestrians, the intent was to assess whether the model would fabricate an answer. Both models have expressed their inability to provide such information. This demonstrated an awareness of their limitations rather than attempting to generate a response.  




 §.§ Expression of Contrastiveness


Ideally, initial explanations should include a contrast with an alternative case to justify the output. Then the user would investigate it further if desired. In our experiment, we observed that in most cases, models, particularly with GPT3.5, failed to mention alternative scenarios unless explicitly prompted (e.g., inquiries such as “What could it be otherwise if it wasn’t X?” or “How does X differ from other options?”). Nevertheless, to measure the presence of contrastiveness across models, we analysed the occurrence of contrastive terms and comparative phrases within the generated responses. We compiled a dictionary consisting of terms such as `distinguish', `different', `contrast', `compared to', `in contrast', `while', `both', `on the other hand', `whereas', and `conversely', which introduce contrasting explanations. Figure <ref> shows that `distinguish' is the most frequently used term by GPT-3.5, whereas `different' predominates in GPT-4's responses. Other frequently encountered terms include `compared to', `contrast', and `differentiate.' Overall, the results suggest that GPT-4 employs a wider array of contrastive terms than GPT-3.5 in its responses. Our qualitative observations further confirm that while GPT-3.5 tends to adhere to one type of contrastive explanation, GPT-4 explores diverse ways to express such contrasts.




Although both models frequently employ these terms, indicating the generation of alternative viewpoints in their responses, the effectiveness of contrastive explanations varies considerably. Our qualitative review found that well-developed contrasting explanations included alternative scenarios with counterfactuals to strengthen the case. Conversely, weakly constructed contrastive explanations failed to explore `what if' scenarios to elucidate the disparity between the two situations under comparison.





§ CONCLUSION

This study contributes to advancing post-hoc explainability by proposing a traceable question-answering approach. By integrating LLMs with the decomposition-based explainability technique, we generate natural language explanations in an interactive format. Our approach promotes transparency and interpretability in scene-understanding tasks by fusing LLMs with semantic feature importance. The proposed approach has the potential for application across various domains, offering comprehensible explanations for decision-making processes, including critical areas such as medical applications. While this is a promising development, there are still areas for further exploration and improvement. One area that requires attention is the overuse of social cues by LLMs. While polite and helpful phrases can enhance user experience, their misuse can become vexing over time. It is important to evaluate such interactions through human subject studies to understand user perceptions of the explanations provided by the LLM. Future research could also explore dialogue management strategies to address issues related to interaction and personalization. On another note, the dynamic nature of the external knowledge source offers the potential for further expansion with additional observations. This not only ensures that the system remains relevant and up to date but also allows the system to answer different types of questions. Additionally, incorporating multimodal elements, such as combining text with visual explanations, can enhance the effectiveness of the responses (e.g., highlighting the specific elements contributing to the prediction). Another consideration is the estimation of feature importance values, which are subject to perturbations and computational processes. It is important to assess and reconcile differences in feature importance across algorithms to ensure stability and consistency. 
 

 







This work was supported by the Wallenberg AI, Autonomous Systems and Software Program (WASP), funded by the Knut and Alice Wallenberg Foundation. 


















 
  




 Paper Title: Privacy-preserving datasets by capturing feature distributions with Conditional VAEs.



: Augmented Masked Autoencoder Pretraining on Public Brain MRI Data for 3D-Native Segmentation
    Asbjørn Munk1,2* Jakob Ambsdorf1,2* Sebastian Llambias1,2 Mads Nielsen1,2

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



§ ABSTRACT 


Large and well-annotated datasets are essential for advancing deep learning applications, however often costly or impossible to obtain by a single entity. In many areas, including the medical domain, approaches relying on data sharing have become critical to address those challenges. While effective in increasing dataset size and diversity, data sharing raises significant privacy concerns. Commonly employed anonymization methods based on the k-anonymity paradigm often fail to preserve data diversity, affecting model robustness. This work introduces a novel approach using Conditional Variational Autoencoders (s) trained on feature vectors extracted from large pre-trained vision foundation models. Foundation models effectively detect and represent complex patterns across diverse domains, allowing the  to faithfully capture the embedding space of a given data distribution to generate (sample) a diverse, privacy-respecting, and potentially unbounded set of synthetic feature vectors. Our method notably outperforms traditional approaches in both medical and natural image domains, exhibiting greater dataset diversity and higher robustness against perturbations while preserving sample privacy. These results underscore the potential of generative models to significantly impact deep learning applications in data-scarce and privacy-sensitive environments. The source code is available at https://github.com/francescodisalvo05/cvae-anonymizationgithub.com/francescodisalvo05/cvae-anonymization.



 



§ INTRODUCTION


Over the past decade, Deep Neural Networks have made significant progress, impacting a wide range of industries. A key driver of this progress has been the increasing availability of large, well-annotated datasets. However, such datasets are rare in several domains, such as medical image analysis <cit.>. While data-sharing between institutions offers a potential solution, it also raises privacy concerns over personal data. To this extent, data privacy has been a focal point for decades. 

For instance, the k-anonymity paradigm <cit.> was introduced in 2002 to obscure user data in relational databases, ensuring that no individual could be distinguished from at least k-1 others. Building on this idea, the k-Same <cit.> method emerged to address privacy concerns in surveillance technologies. It aggregates disjoint clusters of k data points (i.e., faces) with their centroid, either in the pixel space or in the eigenspace.

With the advent of generative adversarial networks (GANs) <cit.>, new possibilities for generating synthetic images have emerged. However, there is a risk of embedding personal information in the latent space. To address this, methods such as k-Same-Net <cit.> enforced k-anonymity within GANs' latent space during training. Furthermore, GANs have struggled with issues such as mode collapse, where models fail to produce diverse outputs, thereby limiting their generalizability <cit.>. Orthogonally, recent advancements have shifted focus to vision foundation models <cit.>, typically trained in a self-supervised fashion on large unlabeled datasets. These large models using vision transformers <cit.> excel at capturing high-level complex patterns and contextual information across different domains. Moreover, their robustness against image corruptions and distribution shifts <cit.> is an established benefit. Researchers are now proposing foundation models for specific domains including healthcare <cit.>, climate <cit.>, geospatial data <cit.>, and more. The rich and low-dimensional feature representation derived from these models makes them a suitable alternative (to the higher dimensional raw images) for pursuing downstream tasks, such as classification.  

Building on the potential of those highly complex and informative representations, as illustrated in Figure <ref>, we train a Conditional Variational Autoencoder () on the feature space of a pre-trained foundation model to accurately capture the training distribution. The generative process, conditioned on class labels, effectively mimics the original distribution while enhancing privacy and expanding the diversity of the generated feature vectors. After capturing the training data distribution with our , we demonstrate a paradigm that no longer relies on disclosing the original images (or their embeddings) but simply the frozen  decoder allowing the sampling of any required dataset size. Despite tremendous benefits with respect to data and patient privacy, this also substantially reduces the required data exchange from an order of gigabytes (images) to a few megabytes (decoder weights). 

In summary, our contributions are:


	
  * We propose a novel anonymization paradigm using Conditional Variational Autoencoders (s) trained on the feature space of foundation models, which provides a more informative and compact representation than the traditional image space.
	
  * We demonstrate superior feature diversity and classification performance with our  compared to reference approaches across both medical and natural datasets. This also highlights the effectiveness of foundation model's feature representation, even in areas beyond its initial training focus. 
	
  * We show that the diverse feature representation obtained via dynamic sampling and decoding through 's decoder enhances model robustness against perturbations.





§ RELATED WORKS


As digital data usage expands, the need for robust privacy measures becomes critical. The k-anonymity framework, proposed by Sweeney <cit.>, addresses this by ensuring that an individual cannot be distinguished from at least k-1 others within database entries. This method involves generalizing or suppressing specific identifying attributes such as zip code, gender, and other sensitive information. While effective at masking identities, this approach often leads to significant data loss, reducing the utility of the data. With the rise of video surveillance systems in New York, a seminal work <cit.> introduced k-Same for face de-identification, translating these principles to more complex scenarios, including the pixel space and the eigenspace. This is achieved by collapsing groups of k elements into a single representative point, namely the centroid. Despite its application in real-world medical systems <cit.>, the inherited data loss and lower data diversity can substantially affect the generalization <cit.> and robustness <cit.> of machine learning based solutions, particularly in domains where data is already scarce.

Moreover, building on the capabilities of GANs to generate high-quality synthetic samples, k-Same-Net <cit.> applies the k-anonymity paradigm within GANs' latent space for face de-identification. Enhancing this further, k-SALSA <cit.> recently introduced local style alignment to preserve granular details in retinal images, while PLAN <cit.> introduces an auxiliary identity classifier to prevent sample collisions. However, GANs still face significant challenges due to their high computational demands, training instability, and mode collapse <cit.>, which can affect their efficiency and reliability.

Before the advent of GANs, Variational Autoencoders (VAEs) have already demonstrated their generative capabilities. For face de-identification, Yang et al. <cit.> generate realistic yet untraceable images by clustering and re-synthesizing facial features. Taehoon et al. <cit.> enforce privacy by introducing noise into the latent vectors, obscuring identities before reconstruction. Other than the image domain, VAEs have been employed for anonymizing speaker data <cit.> and tabular data <cit.>, as well. Conditional Variational Autoencoders (CVAEs) enhance the versatility of VAEs by allowing the generative process to be conditioned on additional labels or attributes, effectively tailoring the output to specific data characteristics. Hajihassnai et al. <cit.> introduce ObscureNet <cit.>, a CVAE-like architecture to anonymize sensor data, that conditions the latent space on private attributes (, age, gender). The architecture consists of one input encoder, multiple external discriminators (one per attribute), and a conditional decoder. Apart from anonymization, the conditional generative process of CVAEs has been used to augment training data <cit.>, addressing issues such as data scarcity and imbalance. 

In contrast, our approach first enhances the conditional generative process by shifting the generation from the pixel space to the embedding space, utilizing large pre-trained foundation models. This transition enables a more efficient feature generation, thanks to the more compact and informative feature representation. In contrast to ObscureNet <cit.>, we employ one conditional encoder rather than multiple discriminators, which allows our architecture to scale more efficiently with the number of classes. Furthermore, instead of simply augmenting our training dataset, we focus on mimicking its embedding-based training distribution, aiming to replicate a diverse and privacy-aware distribution. We further integrate this generative process during training of downstream tasks, which allows us to dynamically generate new samples until convergence. 



§ METHOD


The proposed approach to capture and anonymize a given training dataset is illustrated in Figure <ref>. Initially, we describe the process of feature extraction through vision foundation models. This is followed by a description of 's training procedure. We then introduce two anonymization strategies: the first generates a persistent synthetic replica of the training dataset, while the second continuously generates data on the fly for downstream tasks. 



 §.§ Feature extraction
 

We employ a large pre-trained model to extract and subsequently store feature embeddings and their associated labels (f_i,y_i) ∈ℱ for each labeled input image (x_i,y_i) ∈𝒳, with an associated categorical class distribution 𝒞. The extracted feature representations f_i represent the positional information of x_i within the feature space of the selected feature extractor. Here, close points are more likely to share similar image patterns while distant points indicate semantic differences, potentially across different classes. This descriptive property of the feature space is beneficial for effectively capturing the training distribution. It is much less pronounced in the original pixel space, which suffers from inherent information redundancy and the challenge of calculating meaningful distances in high-dimensional spaces. Additionally, from a storage perspective, while a conventional 2D RGB image of dimensions 224×224 occupies approximately 150,000 pixels, the corresponding image embeddings typically have around 1,000 components, reducing data storage by several orders of magnitude. For our experiments, we rely on DINOv2 ViT-B/14 <cit.>, one of the most recent open-source foundation models that achieved state-of-the-art performance in a variety of tasks, demonstrating the high quality of their latent feature representations.



 §.§  training
 

Conditional Variational Autoencoders extend traditional Variational Autoencoders by integrating specific attributes, such as class labels, into the data generation process. This extension allows the  to not only produce data that closely resembles the original samples in terms of content but also aligns with the conditioned attributes, thereby improving both the relevance and specificity of the generated data. In our setup, the  is trained to accurately learn and replicate the distribution of a feature dataset ℱ, conditioned on the corresponding labels, i.e., (f_i,y_i=c). This training allows us to subsequently generate new anonymous feature vectors a_j for a chosen class ỹ_j that reflect the original training distribution. We can thus utilize only the pre-trained decoder to sample arbitrarily many samples per class c, adhering to the original categorical class distribution 𝒞. This approach ensures that the generated data maintains consistency with the original dataset while enhancing privacy and utility, likewise.



 §.§ Offline anonymization
 

The first way to anonymize our dataset is via storing a persistent anonymized replica of the data. To this end, once the  is appropriately trained, we generate synthetic feature vectors that reflect the initial data size and class distribution. This anonymization process is described in Algorithm <ref> and requires a pre-trained  model, a number of synthetic samples N and a categorical distribution 𝒞 representing the class probabilities of the original dataset. This ensures the generation of similar class proportions in the synthetic dataset. For each of the desired N data points (step 2), a class label ỹ_j is sampled from 𝒞 (step 3). Next, a latent variable z_j is drawn from a standard normal distribution (step 4) to serve as a stochastic input for data generation. These sampled variables (z_j,ỹ_j) are fed into 's decoder, which generates a new data point a_j conditioned on the class label (step 5). Finally, both the generated data point and its label are then stored (step 6). Upon completion, the algorithm outputs the anonymized datasets, consisting of these synthetically generated feature vectors and their labels, effectively balancing privacy preservation with the utility of the dataset for downstream tasks.





 §.§ Online anonymization



Another way to ensure data anonymization is to eliminate the use of persistent datasets entirely. Our proposed  model exemplifies this approach by replacing the traditional dataset with the 's pre-trained decoder. Specifically, we iteratively generate new data on the fly during the training of our task-specific head, meaning no persistent data is required anymore. This method not only overcomes the need to store or send large volumes of sensitive data but also enables the sharing of a simple model ('s decoder) capable of replicating the training data distribution at an arbitrary incidence rate without actual data exchange. Moreover, in the context of model sharing, such as in federated learning <cit.>, our approach offers additional security benefits. While federated learning allows for task-specific models trained on private data to be shared without exchanging the data itself, it is not without risks; model weights can potentially reveal training data characteristics <cit.>. 

Our approach iteratively generates anonymized and stochastic data. This can introduce an additional layer of security, further mitigating the risks associated with traditional data- and model-sharing approaches.




§ EXPERIMENTAL RESULTS


Initially, we outline the process of feature extraction and the training of our Conditional Variational Autoencoder (). Subsequently, we evaluate the classification performance and feature diversity using our anonymized method compared to k-Same <cit.>, the foundational anonymization method that influenced machine-learning-related applications, such as GANs <cit.>. Additionally, we evaluate the classification performance under inference perturbations, showcasing the robustness of the generated feature representation. Lastly, we perform a qualitative analysis of the initial feature space and the anonymized one, in order to visually comprehend the representativeness of the proposed method. 


 §.§ Feature extraction and  training


For feature extraction, we used the DINOv2 ViT-B/14 foundation model <cit.>, which has been trained on 142 million natural images and outputs image embeddings of size 768. To optimize the training process of our , these embeddings were pre-generated and stored on disk <cit.>. For datasets lacking official validation splits, we extracted a stratified sample (10%) from the training data to ensure representativeness. Note that both validation and test sets were kept non-anonymized, preserving their semantic integrity.  

The  architecture comprises a conditional encoder and a conditional decoder. The encoder first concatenates input features with one-hot class labels and then processes it through two linear layers having 256 and 100 dimensions, respectively. The latent variable z is again concatenated with the one-hot class labels and passed to a symmetric decoder that mirrors the encoder's architecture. We train the  with the Adam optimizer, learning rate of 0.001, and early stopping. In addition to the MSE loss, a KL-divergence loss is employed to enforce a standard normal prior in the latent space weighted with β = 0.1, prioritizing the reconstruction. These hyperparameters are standardized across all experiments to ensure consistent comparisons.



 §.§ Feature diversity and downstream performance



Evaluation metrics   The k-Same method simplifies anonymization by collapsing groups of k data points into their centroid. While effective for anonymization, this approach often results in significant information loss and increased data sparsity. Conversely,  effectively maintains the initial distribution's representation, enabling the generation of more samples than originally used, without the compromise on diversity. Yu et al. <cit.> estimate data diversity through the convex hull and the maximum dispersion, defined as the sum of all pairwise distances within the (anonymized) dataset. Although calculating the convex hull is impractical for datasets with fewer samples than dimensions, it is trivial to demonstrate that  achieves higher maximum dispersion compared to k-Same, as the latter's regression to the centroid inherently reduces pairwise distances. 

To provide a more nuanced comparison, we analyze the average nearest neighbor distance 𝒟 between the original feature set ℱ and the anonymized feature set 𝒜 (both of size N). This is formally defined in Equation <ref>:


    𝒟(𝒜,ℱ) = 1/N∑_j = 1^N d_min(a_j,ℱ)


where d_min(a_i,ℱ) is the minimum Euclidean distance between a_j ∈𝒜 and all f_i ∈ℱ. Intuitively, this is closely related to the mean of the minimum LPIPS <cit.> used in PLAN <cit.>. In fact, the latter measures the average minimum distance between generated samples and their closest real ones, using the activations of a pre-trained network. Moreover, an equally important goal of any anonymization technique is to preserve downstream performance while ensuring data privacy. In our evaluation, considering image classification as the downstream task, we use the area under the receiver operating curve (AUC) as the evaluation metric. This allows us to take into account different class imbalance ratios across the evaluated datasets. Following prior works in feature space evaluation <cit.>, we train a linear layer on top of the anonymized embeddings using the Adam optimizer with a learning rate of 0.001 and early stopping. To fairly compare our method against k-Same, we replicate the same exact class proportions of the original dataset. 

Reference method and datasets   Consistent with prior works <cit.>, we evaluate the performance of our  against k-Same using k ∈{2,5,10,15}. We exclude the previously mentioned GAN-based methods such as k-SALSA <cit.> and PLAN <cit.>, since they applied k-Same within GAN's latent space to generate anonymous synthetic images. Our analysis begins with a wide range of real-world, small-sized, medical datasets to underscore the clinical relevance of our method. From the MedMNIST+ collection <cit.>, we use BreastMNIST with resolution 224×224 (breast ultrasound, binary, 780 samples). Additionally, we include datasets from the MedIMeta collection <cit.>: Skin Lesion (dermatoscopy, multi-class, 1011 samples) and Axial Organ slices (CT, multi-class, 1645 samples). Another challenging dataset is OCTDL <cit.> (OCT imaging, multi-class, 2064 samples), which exhibits a severe class imbalance. To extend our evaluation to the natural image domain, we selected multi-class datasets from the  module [https://pytorch.org/vision/stable/datasets.htmlhttps://pytorch.org/vision/stable/datasets.html], focusing on those with moderate sample sizes. This includes STL-10 <cit.>, DTD <cit.>, Oxford-Pets <cit.>, and FGVC-Aircraft <cit.>, allowing us to cover a wide spectrum of general imaging tasks. 


Results   Figure <ref> presents a scatterplot with AUC on the x-axis and average nearest neighbor distance 𝒟 on the y-axis, positioning the most effective methods towards the top-right corner. To ensure the reliability of our findings, all results are averaged across three different seed runs.   consistently exhibits both high diversity and high downstream performance (𝒟↑ , AUC↑). In contrast, as expected, lower k-Same configurations (k=2,5) achieve comparable downstream performance but lack diversity (𝒟↓ , AUC↑), whereas higher k-Same settings (k=10,15) increase diversity at the expense of performance (𝒟↑ , AUC↓). Interestingly, consistent results are observed with the OCTDL dataset, which is particularly challenging. Specifically, it presents a severe class imbalance, with a majority class of 752 samples and a minority of just 13. Furthermore, the medical results hold true on natural datasets as well, although the differences in classification performance are less pronounced, partly because these datasets were included in the DINOv2 training set <cit.>. 
 




 §.§ Robustness of adaptive data sampling
 

This experiment evaluates the robustness of our proposed  method, which dynamically generates new samples a_j of class ỹ_j ∼𝒞 at every batch. Here, the inherent data diversity provided by our method is expected to contribute positively to model robustness. To evaluate this claim, a common approach involves testing against image corruptions that simulate real-world distortions <cit.>. However, these corruptions in the pixel space might already be mitigated by the inherent robustness of foundation models <cit.>. Therefore, we propose a perturbation test by injecting Gaussian noise into the test feature embeddings. We systematically apply this noise with zero mean and a gradually increasing standard deviation σ = {1, 2, 3}. Here, we evaluate the performance of  against k-Same. Therefore, we first train our  and then use its decoder to continuously generate new samples during training. Instead of only sampling from the latent space with a standard Gaussian distribution (μ=0, σ^2=1) as reported in step 4 of Algorithm <ref>, we sample with variance σ^2 = {0.5,1.0,1.5}. This allows us to vary the prototypicality of the generated vectors. Specifically, since we enforce a standard normal distribution (σ^2=1) during training, a lower sampling variance during data generation yields samples closer to the class prototypes, and vice versa. The classification pipeline follows the previous settings, reporting the average AUC over three seed runs. 


Results   Table <ref> shows the AUC observed on the clean test feature vectors (σ=0) and on those subjected to varying levels of Gaussian noise.  consistently ranks among the top performers. Although  may initially show slightly lower performance in a noise-free environment, it substantially outperforms k-Same as the noise level increases. This is particularly pronounced in the Skin Lesion and OCTDL, where  (σ^2=0.5) outperforms 2-Same at σ=3 by 4.2% and 3.2%, respectively. 
Notably, this lower sampling variance remarkably improves performance and robustness across all datasets, suggesting that a higher prototypicality might be beneficial in settings with limited and imbalanced data. These results underscore the ability of our approach to enhance model robustness across diverse domains, each presenting unique challenges related to data size and data imbalance.




 §.§ Qualitative results



In addition to quantitative metrics, we conduct a qualitative analysis of the feature space, illustrated in Figure <ref>, using t-distributed stochastic neighbor embedding (t-SNE) <cit.>. Notably, our feature extractor was not originally trained on medical data, therefore it may encounter difficulties with tight class separation due to the subtle nature of diagnostic features. For instance, in the BreastMNIST dataset, which consists of breast ultrasound images, the fine-grained details necessary for accurate diagnosis are not always clearly represented. Nevertheless, subtle differences within the class still seem to be detectable, as indicated by the shape of small clusters scattered throughout the 2-dimensional feature space. 
With respect to the evaluated anonymization techniques, the limitations of the k-Same method, such as information loss and data sparsity, become particularly evident. Although DINOv2 ViT-B/14 may not well separate classes in such datasets, our  approach manages to maintain the original distribution. Clearly, unlike k-Same, our method avoids information loss, ensuring our privacy-aware feature distribution closely mirrors the initial one.




§ DISCUSSION AND CONCLUSION
 

Limitations   While our method does not directly expose training data, we did not rigorously investigate formal privacy guarantees. Furthermore, the effectiveness of s heavily depends on the capabilities of the chosen feature extractor. In addition, this approach could further benefit from domain-specific foundation models, as they may be able to capture more nuanced patterns. Lastly, our data-sharing framework assumes that all parties involved use the same feature extractor. However, with the increasing interest over general foundation models <cit.>, this assumption becomes more reasonable and can be justified by the considerable benefits and enhanced downstream performance. 

Conclusion   We demonstrate that conditional generative models, such as s, effectively address challenges related to privacy and data utility across a diverse range of datasets. This approach further shows high robustness even with small sample sizes and severe class imbalances. Future research can build upon our insights, investigating the sensitive role of the sampling variance, exploring further conditional generative models, or extending our approach to be trainable in an end-to-end manner for downstream tasks.




§ ACKNOWLEDGMENTS


This study was funded through the Hightech Agenda Bayern (HTA) of the Free State of Bavaria, Germany.








 Paper Title: ABC Align: Large Language Model Alignment for Safety & Accuracy.





Discretizing Continuous Action Space with Unimodal Probability Distributions for On-Policy Reinforcement Learning
    Yuanyang Zhu, Zhi Wang, Member, IEEE, Yuanheng Zhu, Senior Member, IEEE, Chunlin Chen, Senior Member, IEEE, Dongbin Zhao, Fellow, IEEE


The work was supported by the National Natural Science Foundation
of China under Grant 62006111 and Grant 62073160.
	
Yuanyang Zhu, Zhi Wang and Chunlin Chen are with the Department of Control Science and Intelligent Engineering, School of Management and Engineering, Nanjing University, Nanjing 210093, China, and also with the Research Center for Novel Technology of Intelligent Equipment, Nanjing University, Nanjing 210093, China (e-mail: yuanyang@smail.nju.edu.cn; zhiwang@nju.edu.cn; clchen@nju.edu.cn).
 
 Yuanheng Zhu and Dongbin Zhao are with the State Key Laboratory of Management and Control for Complex Systems, Institute of Automation, Chinese Academy of Sciences, Beijing 100190, China, and also with the School of Artificial Intelligence, University of Chinese Academy of Sciences, Beijing 100049, China (e-mail: yuanheng.zhu@ia.ac.cn; dongbin.zhao@ia.ac.cn)

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



§ ABSTRACT 
Alignment of Large Language Models (LLMs) remains an unsolved problem. Human preferences are highly distributed and can be captured at multiple levels of abstraction, from the individual to diverse populations. Organisational preferences, represented by standards and principles, are defined to mitigate reputational risk or meet legislative obligations. In this paper, we present ABC Align, a novel alignment methodology for LLMs that enables integration of the standards and preferences of a large media organisation into the LLM itself. We combine a set of data and methods that build on recent breakthroughs in synthetic data generation, preference optimisation, and post-training model quantisation. Our unified approach mitigates bias and improves accuracy, while preserving reasoning capability, as measured against standard benchmarks.




§ INTRODUCTION




 §.§ Overview



In this paper, we present a novel approach to the alignment of Large Language Models (LLMs) we call ‘ABC Align’. This alignment is conducted in two settings: the fine-tuning of open-source models, and In-Context Learning (ICL) of closed-source ‘frontier’ models <cit.>. While the techniques and literature of In-Context Alignment (ICA) are less mature than fine-tuning based alignment, we achieve good results across both settings. By design, our methodology is applicable across any range of open- and closed-source models, offering provider independence and flexibility as underlying model capability grows with time. 

We use a variety of data that is both domain-specific and general. These data include news article content, organisational ‘AI Principles’ <cit.>, and human-reviewed question/answer pairs sourced from an internal Retrieval-Augmented Generation (RAG) tool. They are used as input to both generate datasets for fine-tuning, and in the case of the ‘AI Principles’, directly in the frontier model context window itself.

In the fine-tuning alignment setting, we use several techniques from the literature on synthetic data generation, knowledge distillation, and preference optimisation to produce several datasets conditioned on these data. We include results controlling for the impact of each of these methods and data across two open-source models, from Mistral AI and Meta respectively.

We see general performance improvements across industry standard benchmarks. The version of our dataset that combines all data and methods achieves a 23.51% relative performance improvement on Meta’s Llama3-8B model on the TruthfulQA benchmark. Furthermore, these results were achieved using a small dataset of 897 samples, showing substantial improvement over Meta’s own fine-tuned version of the base model, which is trained on more than 10M human-annotated examples <cit.>. 

For the ICA setting, conducted on OpenAI’s GPT4-turbo model, we evaluate the performance of a number of custom system prompts, including the prompt used by our internal RAG tool, and our AI Principles. Most notably in evaluation against the Bias Benchmark for Question Answering (BBQ), a system prompt that uses our AI Principles improves relative performance over baselines by 77.54%.

Aligning LLMs and conducting evaluations are inherently challenging due to difficulty in defining ‘who the alignment is for’, and the need to define benchmarks reflecting a measure of alignment specific to that group, in addition to standard benchmarks. 

Our work aims to show an example of how organisations can answer these questions, integrate their standards into a process of provider-agnostic alignment, hence providing the organisations with choice and agency in how they apply LLMs across a range of use-cases.



 §.§ Background 


Authors’ note: no editorial position is offered or claimed in this publication.
We also do not publish our data, and all LLM prompts are provided to support replication of our work and do not reflect optimisation for production use.

For real-world industrial application, in our case a large public-service media (PSM) organisation, alignment of model outputs to policies and preferences is critical. This is due to charter obligations to provide the public with innovative services <cit.> that mitigate reputational risk while being transparent. 

Many PSMs face additional constraints specific to their operating model, most notably in the link between their sources of funding and restrictions around commercial activity <cit.>. This reality can limit their capacity to support internal R&D efforts around technology, while they simultaneously rely on technology to deliver content to audiences in a highly competitive media marketplace.

LLMs present a unique opportunity for PSMs to improve internal processes. For example, through enhancing content metadata or supporting information retrieval across large corpora. LLMs also carry a unique risk profile <cit.> when working with that content itself, whether it be to transform the content for multi-platform use, or use-cases that sit at the edge of emerging definitions around ‘responsible AI’, such as content generation <cit.>.

Recent progress in model capability that enables processing of content across modalities in a single system – text, audio, video – only serves to enhance this risk profile and magnify impact of the constraints PSMs face, given the computational and operational costs associated with utilising such systems at scale.

Scale is a critical consideration for PSMs both in terms of the audience served and content produced and stored. Leveraging the capabilities of LLMs (and their multimodal derivatives) thus involves a considered trade-off between scale, cost, and safety. 

The Australian Broadcasting Corporation (ABC) has defined a set of AI Principles (see Appendix A) that are designed to guide the adoption and application of AI in a responsible, principled manner. We build on the use of ‘safety’ in the literature, particularly work on ‘Constitutional AI’ <cit.>. We make explicit use of the ABC AI Principles in our methodology, and here include the principles of ‘protecting data’ and ‘mitigating bias’ to broaden the definition of ‘safety’. ‘Accuracy’ maps directly to the AI Principle of the same name. This paper outlines methods that securely leverage the ABC’s data and are evaluated against industry-standard benchmarks that act as proxies for these principles.



 §.§ ABC Align


Our work sits at the intersection of machine learning theory and application. The organisational context in which it is conducted, under the various constraints outlined above, guides our choices of models, data and methods. These methods leverage the unique advantage of media organisations, namely their content, which is created in line with clearly defined policies and principles. We consider these as data valuable for the alignment and evaluation of LLMs, when paired with the approach we outline in this paper.

To best utilise this data, we build on recent work on synthetic data generation aimed at distilling frontier models’ capability into smaller open-source models using ICL <cit.>. Both the open-source and frontier models can be substituted, and the impacts of these changes are measurable. In this paper, we include representative results across several open-source models and a specific frontier model to validate and demonstrate our methods. Our aim is to offer a blueprint that is prescriptive on method only, enabling the replication and use of our work by other PSMs with their own unique data.

Our synthetic datasets, also referred to as augmented datasets, are used by a fine-tuning process that conditions each sample on textual article content, organisational ‘AI principles’, and samples gathered from the live deployment of an internally built assistive tool. This tool enables the capturing of editorial preferences and input at scale, hence also exploiting the ICL abilities of frontier models to guide and align the application to user needs. 

Development of this tool involved the crafting of novel, structured prompts and retrieval architectures, evaluated against user-derived acceptance criteria that were reflective of broader organizational preferences. We defer a lengthier discussion of its implementation to a future publication. 

In addition to our alignment methods in a fine-tuning setting, the results presented in this paper reflect nascent developments in ICA of frontier models, and we expect to build on these in future work, with more detail covered in our 'Future Work' section. We aim to provide a general methodology: general to our organisation’s data, the full complement of available AI models, and application to a range of use-cases specific to media organisations. Such applications can include metadata tagging, sourced information retrieval, and LLM-based approaches to content recommendation <cit.>.

We leave the infrastructure enabling the synthetic dataset creation, knowledge distillation, training and evaluation of LLMs out of scope for this paper. We note that we make extensive use of open-source tools, including Hugging Face's Transformers and TRL <cit.> libraries among others.

Evaluation also forms a key part of our methodology. We build on the work of EleutherAI, specifically their ‘Large Language Model Evaluation Harness’ <cit.> and emerging best practices to measure our alignment efforts against industry-standard benchmarks that reflect the ‘AI Principles’ outlined above. The framework is extensible, enabling custom, organisation-specific evaluations defined in collaboration with non-technical stakeholders. We offer further commentary on this in the ‘Future Work’ section at the end of this paper.



§ RELATED WORK


In this section we outline a brief history of LLM development, and each of the methods we use as part of ABC Align.

Early implementations of the ‘Transformer’ and related attention mechanism were limited in both scale of pre-training data and the number of parameters in the neural network <cit.>. However, the new architecture, with query-, key-, and value-cache to optimise both training and inference, went beyond the capacity of earlier architectures like Long Short-Term Memory (LSTM) and Recurrent Neural Networks (RNNs), both being limited in their ability to process long sequences under fixed computational constraints.

Radfort et al. <cit.> moved beyond masked-language modelling approaches and described the Generative pre-trained Transformer (GPT) architecture, which framed the prediction problem in terms of arbitrary sequence-to-sequence modelling <cit.>. 

Ouyang et al. <cit.> directly improved this ability to follow general instructions via their method of ‘instruction tuning’, framing the model interaction as a ‘request/completion’ dynamic, leading to the conversational-style interfaces seen in recent iterations of GPT-class models.

As this approach was scaled up, the diversity and distribution of training data (consisting of trillions of tokens sourced from web-scale datasets), revealed limitations for the application of these models in a practical setting, where user expectations around accuracy and bias were not met <cit.>.



 §.§ Domain-Specific Large Language Models


LLMs have been developed for application to specific fields or areas of industry, such as biology and finance <cit.>. A canonical example is BloombergGPT <cit.>. Bloomberg has a significant store of domain-specific data (363B tokens) enabling it to perform full pre-training of their own model. Bloomberg’s experience informed the methodology outlined in this paper, as the ABC itself has a large store of data across digital-first content and its archives. 

A complete pre-training process of an ABC LLM is likely to yield a capable, general model, given that digital content alone contains over a billion tokens, and 90+ years of archival content is currently being digitised. However, the expense and risk associated with maintaining a single artifact given the rate of progress in the field <cit.>, have led us to favour a lightweight, modular methodology rather than pre-training a base model from scratch. This would allow us to leverage the value of the ABC’s extensive data but in a way that was agnostic to underlying base models, in a cost-effective and future-proof manner. This has in part been validated by subsequent research, which has demonstrated that OpenAI’s GPT-4 has exceeded the performance of BloombergGPT <cit.>.



 §.§ LLM Alignment


LLM alignment is a rapidly evolving area of research, alongside efforts to further scale up pre-training data and neural network parameter counts. Starting in 2023, early evidence <cit.> emerged that while massive amounts of data were required in the pre-training stage, alignment itself could be facilitated through smaller, high-quality curated datasets with samples numbering in the thousands instead of trillions. In parallel, LLM capabilities were being categorised relative to new evaluation frameworks, specifically around ‘reasoning’ (problem-solving) <cit.>, extending standard definitions of accuracy beyond metrics like ROUGE and BLEU <cit.>.

Our work builds on ‘Constitutional AI’ from Anthropic <cit.> and ‘Less is More for Alignment’ (LIMA) by Meta <cit.>. Constitutional AI focuses on the self-alignment of LLMs based on a ‘constitution’ or set of principles that guide the process of ‘AI feedback’. We adapt this approach to use the ABC AI Principles directly in the creation of our fine-tuning datasets and frontier model alignment prompts. 

LIMA defines the ‘Superficial Alignment Hypothesis’, that states a ‘model’s knowledge and capabilities are learnt almost entirely during pre-training, while alignment teaches it which sub-distribution of formats should be used when interacting with users’. The use of ‘superficial’ here indicates that fine-tuning a subset of parameters, rather than complete pre-training, is the appropriate setting for LLM alignment. Their work examines the impact of a small set of 1000 samples on several alignment metrics, building on <cit.>. Our dataset is thus correspondingly small, high-quality, and informed by our AI Principles.

Our own hypothesis is that given a sub-distribution of fine-tuning and preference data specific to our organisation, we can in turn align LLMs in a way that preserves their general utility, capability, and performance on industry-standard benchmarks that directly map to and reflect our AI Principles. While non-domain-specific data like those from ‘Stack Exchange’, ‘Wikihow’ and ‘Reddit’ can be valuable in an academic fine-tuning setting, their use for alignment in a PSM setting is harder to justify. Consider a user feedback scenario, in the context of PSM obligations outlined in the introduction:


‘Why did your model generate this inaccurate/biased output?’

‘It was aligned using reddit threads!’

A key challenge that the LIMA authors observe is that ‘manually creating diverse prompts and authoring rich responses in a uniform style is laborious’. In our own organisational context, asking editorial staff to label data, even if for a small sample, represents a significant and resource-intensive task.



 §.§ Synthetic Data and Knowledge Distillation: ORCA


The utilisation of ‘synthetic data’ <cit.> represents a novel source of data for training LLMs. It is often cited in wider discussions about the finite quantity of human-created data <cit.>. As with human data, not all synthetic data are created equal. Indeed, low-quality synthetic data, when riddled with hallucinations, inaccuracies, and biases, can significantly impact downstream task performance <cit.>.

To improve the quality of our synthetic data and mitigate these issues, we condition each sample on organisational data: news articles. Beyond the content itself, these articles comply with the ABC’s editorial standards and policies <cit.> at the time of their publication, and provide a real-world example of ‘human annotated data’ that goes well beyond crowd-sourced methods such as or ‘Mechanical Turk’ or similar. We also go beyond algorithms like Self-Instruct <cit.> that use ‘seed instructions’ to generate synthetic data, by providing a unique input for each generated sample.

Knowledge distillation techniques have diversified well beyond the original work of <cit.>. In the context of this paper, the teacher/student dynamic remains, but occurs between frontier model/smaller open-source model. The distillation occurs per-input, given a prompt. The structure of this prompt confines the possible sub-distribution of the synthetic data used by the student model in our fine-tuning process. Here we build on the work of <cit.> on ORCA, where they elicit ‘reasoning traces’ from OpenAI’s GPT-4 based on synthetic samples, with the goal of improving specific capabilities of a smaller 13B parameter model.

We note here that our synthetic data samples are created for the specific purpose of improving underlying model capability rather than meeting editorial standards, following the methods outlined in ORCA. We do not intend on making this data publicly available. This is to avoid any confusion with the published content it is derived from. 




 §.§ Preference Optimisation: ORPO


Soliciting human feedback for alignment efforts proved crucial in the success of tools like ChatGPT <cit.>. In a reinforcement learning setting, numerical instability and sensitivity to training dynamics were seen as a necessary trade-off for improved alignment performance. Model-free techniques such as Direct- and Odds-Ratio Preference Optimisation (DPO/ORPO) provide stable alternatives to the Reinforcement Learning with Human Feedback (RLHF) approach <cit.>.

Following a pre-training stage, models undergo further post-training to make them suitable for general use. Multiple techniques exist to facilitate this, with initial work focused on RLHF <cit.>. This method aims to use human preference data to align pre-trained models and has demonstrated improved results beyond safety and reducing harmfulness <cit.>, which are the key areas of focus for our work.

Recent work on preference optimisation <cit.> aims to elicit similar changes in model behaviour via methods that reduce implementation complexity compared to RLHF. Given the constraints on PSMs (and organisations whose operational focus is not technology itself), these methods provide a way to demonstrate the value of incorporating their own ‘human feedback’ in the alignment of models they use.

In this paper, we make use of an implementation of ORPO that is readily available in Hugging Face’s libraries. ORPO offers good performance when compared with both Supervised Fine-Tuning (SFT) and DPO, including improved computational efficiency for a given batch of training data, making it suitable for the aims of our work.



 §.§ In-Context Alignment


There has been significant recent work on exploiting the ICL capabilities of frontier models <cit.> where the context length has increased by up to an order-of-magnitude, as seen for instance in Google’s Gemini and Anthropic’s Claude series. While performance of fine-tuning in an ICL setting doesn’t yet offer consistent improvements when compared with SFT-based approaches across all tasks, there are specific cases where results are promising <cit.>. 

The work of <cit.> further quantifies the ‘Superficial Alignment Hypothesis’ proposed in LIMA. They measure the token distribution shift between pre-trained and aligned models, specifically the difference in probabilities of predicted tokens given a position in some sample input across the two models. They find that alignment fine-tuning affects a small subset of tokens that they describe as ‘stylistic’, specifically ‘discourse markers, transitional words, and safety disclaimers’. 

Further to the original hypothesis of the LIMA authors, this provides evidence and motivation for the effectiveness of alignment happening in a post-training setting via fine-tuning, and that inference-time alignment may offer comparable results for applications where fine-tuning is unsuitable or not possible. This further validates our decision to not pre-train an LLM, but instead focus on post-training alignment across SFT, PO, and ICA settings in support of our overall aim to provide a lightweight, flexible methodology that considers organisational constraints and the rapid growth of underlying model capabilities.



 §.§ Model Quantisation


As we discussed in the introduction of this paper, PSMs are funding-constrained, thus any training of LLMs needs to be done in an efficient way. Early methods of model quantisation involved performance trade-offs in terms of prediction quality and performance when compared to full-precision training and inference. However, recent advances in model quantisation have enabled rapid and low-cost experimentation. QLoRA <cit.> addresses the shortcomings of earlier methods sufficiently and is our choice for quantisation for the experiments outlined in this paper.



 §.§ Retrieval Augmented Generation & Grounding Synthetic Data


Another key aspect of the ABC Align methodology is the refining of the preference dataset, using user-generated outputs from an internally developed RAG tool. This tool is powered by OpenAI’s GPT-4-32k <cit.> and uses context information retrieved from a corpus of ABC documents to answer user questions. 

Within this context, alignment of the frontier GPT-4 model refers to ensuring that the entire RAG tool produced outputs that were relevant, accurate, and valuable to end users. This ultimately served to reflect whether the application (and thereby the frontier model) met users' expectations, and by proxy their broader organisational values. This effort involved a comprehensive system-wide approach, including not only refining the prompting techniques but also ensuring that the retrieval mechanisms, data sources and generated outputs all complied with the overall expected objectives <cit.>. Evaluation of the system includes using user-acceptance derived metrics, conducting user interviews, and manually reviewing generated samples, ensuring an adequate level of alignment of the frontier GPT-4 model to the ABC's AI Principles.

While we generate ORCA-style synthetic data in this work, RAG grounding of synthetic data to produce reasoning traces that are accurate to, rather than derived from, content data will be explored in future work.



 §.§ Information Theoretic Measures in NLP


A primary focus of this paper is dataset quality, and we demonstrate its impact on LLM performance in multiple fine-tuning settings. While measurement against standard benchmarks is critical, we are also interested in measuring the quality of the data itself, as it is initially transformed during our synthetic data and knowledge distillation process. Information theoretic metrics offer a principled approach to measuring quality and changes in the information that the data represents. For this purpose, we use several well-known metrics, including Shannon entropy, mutual information, and Kullback-Leibler divergence <cit.>.



§ METHODOLOGY



 §.§ Synthetic Data Generation


In this section, we outline two versions of our alignment dataset for full fine-tuning of open-source models, together with the synthetic dataset generation prompts. The inputs to these datasets are news articles, the ABC AI Principles, and samples from our internal RAG tool. 

This tool, while an early prototype, is designed to deliver the capabilities of LLMs directly to users while also providing a means for soliciting editorial feedback and preferences, key to scaling our methods, outlined in the ‘Future Work’ section of this paper.

An automated pipeline was developed to enable the creation of synthetic datasets by augmenting existing data, as outlined in the previous section. The step-by-step process of creating an ABC Align dataset suitable for SFT is as follows:


    
  * Collect news articles as the input.
    
  * Create LLM-generated question/answer pairs relating to each input news article using an ORCA-style prompt.
    
        
  * Prompt: Deduce any reasoning or logical problems from this article:  Generate a complex question and a logical answer that requires step-by-step thinking, and elaborate on this thinking process as part of the answer. Both the question and answer must not refer to the original article. 
    


The resulting dataset is then used to create a dataset suitable for preference optimisation (PO). The answer for each sample is rewritten considering the AI Principles, given two canonical examples of human-reviewed Q&A pairs from our internal RAG tool.


    
  * Rewrite the answers from the SFT dataset
        
  * For ‘chosen’ answers:
        
            
  * Prompt: Here are two high-quality, human-reviewed Question/Answer pairs: Question:  Answer:  Question:  Answer:  Now review a new Question & Answer pair:  – Rewrite the answer to better align with these principles:  Do not reference the principles in your response. Include only the text of the answer.’
        
        
  * For ‘rejected’ answers:
        
            
  * Prompt: Here are two high-quality, human-reviewed Question/Answer pairs: Question:  Answer:  Question:  Answer:  Now review a new Question & Answer pair:  – Rewrite the answer to be unaligned with these principles:  Do not reference the principles in your response. Include only the text of the answer.’
        
    

In all cases, we use OpenAI’s GPT-4-turbo frontier model with a temperature of 0.7 to generate synthetic data for use in a fine-tuning setting. The resulting datasets comprise 897 individual samples.



 §.§ Control Datasets


We include three additional datasets in our experiments. These are OpenORCA <cit.>, LIMA in the SFT setting <cit.>, and UltraFeedback <cit.> for ORPO training.

OpenORCA is a replication of the original ORCA data which was not made public. The dataset is sourced from FLAN-T5 and generated using both OpenAI’s GPT-3.5-turbo and GPT-4 models. The authors aim to reflect the distribution of the dataset outlined in the ORCA paper.

UltraFeedback is an updated version of the dataset used in the ORPO paper that addresses a data contamination issue related to the TruthfulQA benchmark.

The LIMA dataset is the original as published by the authors, and in all cases we randomly take 897 samples for direct comparison with our own datasets.



 §.§ Models and Fine-tuning


We fine-tuned the following pre-trained models: Llama-3-8B <cit.> and Mistral-7B-v0.3 <cit.>, including both the base and instruction-tuned models. We used the default tokenizer of each of the models, and applied Hugging Face’s Zephyr-style chat template (with , , and  tokens).

As described above, we employed a two-step fine-tuning process involving SFT for instruction tuning and ORPO for preference optimisation. To facilitate this, we used TRL’s SFTTrainer and ORPOTrainer <cit.> and scripts from Hugging Face’s Alignment Handbook <cit.>.

ORPO includes an additional training objective, relative ratio loss. It also dynamically penalises ‘disfavoured’ responses, instead of constructing sets of rejected tokens to control for degenerative model behaviour in the SFT setting. 

We trained on single Nvidia A10G GPUs with 24GB memory on cloud infrastructure. We trained for 23 epochs on the SFT dataset, taking around 2.4 hours, and we trained for 10 epochs on the ORPO dataset, taking around 3.8 hours, using a standard checkpoint selection process against selected benchmarks for the models used in our final evaluations. 




 §.§ QLoRA


The SFT/ORPO QLoRA hyperparameters used in this study followed those as set in the Hugging Face Alignment Handbook. An overview of our QLoRA hyperparameters during model fine-tuning are shown in Table <ref>.


Hyperparameter selection for quantisation methods has been shown to influence model performance in numerous ways, where the model may have different optimal ‘rank’ hyperparameters for different metrics, as well as the ratio between alpha/rank <cit.>.

The rank hyperparameter is particularly important since it affects the size of the update matrices, and therefore the number of trainable parameters. It should be noted that these recommendations are evolving, and hyperparameter optimisation may be required for individual use-cases <cit.>.



§ IN-CONTEXT ALIGNMENT


 For our experiments on ICA, we use OpenAI’s GPT-4-turbo model with a temperature of 0.7, and the  evaluation framework. Benchmarking is done in a `' setting rather than traditional multiple choice, due to the limitations around the availability of logits from closed-source APIs.

 Our primary intervention is setting the `' parameter, including three specific prompts, taken directly from our internal RAG tool, the ABC’s AI Principles, and a version of the AI Principles augmented with a Q&A sample sourced from the RAG tool, acting as a canonical example.



 §.§ Dataset Analysis 


This section outlines the information-theoretic and perplexity metrics we used to evaluate the datasets discussed in this study.



 §.§ Shannon Entropy, Mutual Information, KL Divergence


 For dataset analysis, we utilise several information-theoretic metrics on our SFT dataset. Our motivation here is to quantify the changes to average entropy, loss of information, and distribution shift induced by our synthetic data generation process. Our aim is to understand how our methods transform the information present in our input data, as we distil ‘reasoning traces’ suitable for downstream SFT and ORPO fine-tuning. All metrics are normalised.

In Table <ref>, we consider the average Shannon entropy of our input news articles, and the resulting synthetic dataset.




In Table <ref>, we calculate both mutual information and KL divergence across both datasets.



These results demonstrate that, through our process of conditioning synthetic data samples on news articles, we have increased sample complexity relative to the input while keeping shared information low (in-line with our attempt to derive specific patterns with high complexity and abstraction, i.e. ‘reasoning traces’). We have been able to do this while preserving the similarity of the overall distributions, as reflected in the low KL-Divergence score.



 §.§ Perplexity

For perplexity (ppl), we measure the mean ppl per dataset per base model. We calculate the perplexity of all datasets for both pre-trained Llama-3-8B and Mistral-7B-v0.3 models. These metrics are shown in Table <ref> for the ABC Align, LIMA, OpenORCA and UltraFeedback datasets. We note that the ORPO datasets include `chosen' and `rejected' formats, reflecting their preference optimisation nature in line with Hugging Face's Alignment Handbook <cit.>. For the ABC Align ORPO dataset, these were generated by OpenAI’s GPT-4-turbo frontier model to be aligned/unaligned to the ABC AI Principles. 



Interestingly, among the ORPO datasets, the ‘chosen’ format shows lower perplexity than the ‘rejected’ format. Rather than exhibiting equal perplexity for the chosen/rejected formats, this implies that the original model training may have achieved similar results as we aimed to achieve by aligning using the ABC AI Principles. 

We limit our use of ppl to the datasets themselves, rather than the fine-tuning phase. As observed by the LIMA authors, there is an anti-correlation between perplexity measured at training time and generational quality; perplexity can increase without generation quality necessarily decreasing <cit.>. 

Finally, we note that while details of the pre-training datasets of either base model have not been released, data quality is known to impact pre-training performance <cit.>, and indeed Meta does emphasise the use of filtering techniques to select only ‘high-quality’ data <cit.>.



§ MODEL EVALUATION


 In this section we discuss our evaluation strategy, including the mapping between selected benchmarks and the ABC AI Principles. We evaluate both our datasets and model outputs, whereby results are presented in the section below.



 §.§ General Reasoning


In the fine-tuning setting, we want to ensure that we are not training the model to be task-specific. Here we use the ARC dataset, specifically the ‘challenge’ subset, constructed by the Allen AI Institute <cit.>. It comprises 7,787 science questions in a multiple-choice setting aimed towards measuring ‘reasoning capability’. We select this benchmark to ensure that we do not sacrifice general model capabilities while improving bias identification and accuracy, i.e., limiting their application across a range of use-cases – also known as the ‘alignment tax’ <cit.>.



 §.§ ABC AI Principle: ‘mitigating bias’


The Bias Benchmark for Question & Answering (BBQ) <cit.> represents a human-constructed and thorough cross-section of socio-cultural biases in a multiple-choice setting, for the evaluation of a model’s ability to detect bias. In this study, we use , a subset version of the original BBQ, containing samples for evaluation against age, nationality, race, gender, sexual orientation, and physical appearance biases. 

This benchmark serves as a proxy for the ABC AI Principle ‘mitigating bias’, which is defined as ‘We will seek to prevent bias in our AI data or algorithms that could perpetuate and amplify existing inequalities and lead to discrimination’. BBQ is designed specifically to measure whether model outputs contribute to ‘harming marginalised individuals and groups’, thus enabling the identification and mitigation of bias. The list of social biases measured are defined by the US Equal Employment Opportunity Commission <cit.>. We acknowledge the US-centricity of this definition, and outline how we intend to extend our evaluation metrics in a general way, to make them ABC- or other-specific in our ‘Future work’ section at the end of this paper.



 §.§ ABC AI Principle: ‘accuracy’


We use performance on TruthfulQA <cit.> as a proxy for the ABC AI Principle ‘accuracy’. TruthfulQA consists of 817 multiple-choice questions across 38 subjects, including law, health, and politics. This benchmark addresses the limitations of LLMs, where the distribution of their pre-training data may induce the generation of ‘imitative falsehoods’, or answers that sound true without being grounded. TruthfulQA thus represents a robust benchmark for accuracy, where the conditions for a claim to be true are ‘if it describes..[a]..literal truth about the real world’, and the question/answer choices are written by humans. We specifically use the  task.



§ RESULTS

In this section we present the results of our fine-tuning alignment using both SFT and PO datasets.



 §.§ Fine-tuning: SFT (ORCA/LIMA)


In an SFT setting, we evaluate the performance of our ABC Align SFT dataset against data generated using the two methods we build on, ORCA and LIMA respectively. We present our evaluation results in Figure 
<ref> below.





  §.§.§ Evaluation on 


For Mistral-7B-v0.3, we see a decrease in performance compared to the pre-trained model of 0.42% and 0.85% for LIMA and OpenORCA. The Align SFT dataset demonstrates a 3.25% improvement compared to the pre-trained model and those fine-tuned on control datasets. 

For Llama-3-8B, we see improvements over the pre-trained model for LIMA and OpenORCA datasets of 2.78% and 0.59% respectively. The ABC Align dataset demonstrates a 7.17% improvement over the pre-trained baseline.



  §.§.§ Evaluation on 


For Mistral-7B-v0.3, we see a decrease in performance compared to the pre-trained model of 2.17% on the LIMA dataset. For OpenORCA and ABC Align SFT, we see improvements of 12.81% and 14.57% respectively.

For Llama-3-8B, we see an equivalent decrease in performance of LIMA against the pre-trained model of 1.75%. For OpenORCA and ABC Align SFT, we see improvements of 16.34% and 7.20% respectively over the pre-trained baseline.



  §.§.§ Evaluation on 


For Mistral-7B-v0.3, OpenORCA and LIMA improve over pre-trained baselines by 11.01% and 13.69% respectively. ABC Align SFT demonstrates an improvement of 22.64%.

For Llama-3-8B, LIMA improves over the baseline model by 4.56%. OpenORCA and ABC Align improve by 22.57% and 18.26% respectively.



  §.§.§ Analysis


In general, LIMA underperforms in this setting. ABC Align SFT offers the best performance on Mistral-7B-v0.3 across all benchmarks, and on Llama-3-8B for . OpenORCA’s performance exceeds ABC Align SFT on BBQ and TruthfulQA on Llama3-8B.

We see our dataset offers consistent improvements on a specific model, Mistral 7B. Where it underperforms on BBQ & TruthfulQA on Llama-3, it has improved general reasoning capability. The SFT dataset is generated with this aim in mind, as reflected in the ORCA-style prompt. Ideally, we would like to see consistent improvements across a range of open-source models and our three evaluations. These results demonstrate that we need to move beyond SFT and explore PO to achieve this. 




 §.§ Fine-tuning: PO (ORPO)


The PO setting represents the ideal setting for the aims of this paper, where for the ABC Align dataset we have used each of the sources of data unique to our organisation. Preference data in particular is scalable beyond the small number of samples used via our internal RAG tool, to include data sourced from additional workflow integrations. 





  §.§.§ Evaluation on 


For Mistral-7B-v0.3, performance on the pre-trained model decreases by 0.85% for the UltraFeedback (UF) dataset, and ABC Align PO offers a 2.68% improvement. 

For Llama-3-8B, UF and ABC Align PO offer improvements of 2.78% and 6.44% respectively compared to the pre-trained model.



  §.§.§ Evaluation on 


For Mistral-7B-v0.3, UF and ABC Align PO both improve over the pre-trained baseline by 5.26%.

For Llama-3-8B, UF and ABC Align PO both improve over the pre-trained baseline by 7.64% and 7.48% respectively.



  §.§.§ Evaluation on 


For Mistral-7B-v0.3,, we see improvements over the pre-trained baseline for both UF and ABC Align PO of 12.25% and 17.97% respectively.

For Llama-3-8B, we see improvements over the pre-trained baseline for both UF and ABC Align PO of 21.75% and 23.51% respectively.



  §.§.§ Analysis


Here we see ABC Align PO outperforms or offers equivalent performance to the control dataset used by ORPO in all cases. These results demonstrate consistent performance improvements across both base models and all benchmarks. They emphasise the importance of including additional data sources in a PO setting, where ORPO’s penalisation of disfavoured responses better aligns model outputs in a measurable way. 



 §.§ Comparison of Instruction Fine-Tuned (IFT) and SFT/PO Align models


In this section we include results comparing our models and the Instruction Fine-Tuned (IFT) variants of both Mistral-7B-v0.3 and Llama3-8B. Our aim in performing this evaluation is to quantify the impact of large datasets in a fine-tuning setting on the models and benchmarks we have selected.

We note that Mistral do not publish the size of their IFT dataset, while Meta indicates that over 10 million human-annotated samples are included in theirs. We thus restrict our analysis to Llama-3, and include Mistral in our results for completeness. 

We can see in Figure <ref> that our models offer near-equivalent performance gains on . Furthermore, our PO model outperforms the IFT version of Llama3 by 5.2% on , despite our datasets being five orders of magnitude smaller than the IFT dataset. 



The IFT model significantly outperforms both our SFT and PO models on . This result indicates that scaling dataset size and diversity of preference data is likely to offer benefits for specific alignment criteria, represented here by . Identifying socio-cultural biases is a nuanced rather than determinative task, where a wider distribution of samples would be valuable. We elaborate further on this analysis in the ‘Future Work’ section of this paper.  



 §.§ In-Context Alignment


Here we present results of an ICA experiment on a subset of evaluations using data from our ABC Align datasets. This data includes both our internal RAG tool’s system prompt, and prompts designed specifically for the ICA setting, constructed using the ABC AI Principles. These prompts were evaluated using OpenAI’s GPT-4-turbo. It is important to note, as discussed in the ‘In-Context Alignment’ section, that these results are not comparable across the fine-tuning and ICA settings.



We do not include the prompt for our internal RAG tool in this paper. We defer it and a discussion on its construction in a future publication.
 
The prompt referred to as  in Figure <ref> includes our AI Principles themselves, whereas  is an augmentation of the former including a Q&A pair sourced from our internal RAG tool. Please refer to Appendix B for the full v1 prompt.



  §.§.§ Evaluation on 


Compared to the baseline, the improvements across the three versions of our ICA prompts of up to 78.58%.



  §.§.§ Evaluation on 


Results on  demonstrate that our custom prompts reduce performance on this benchmark overall by up to 5.55% across BLEU and ROUGE scores.



  §.§.§ Summary


These results demonstrate that much work remains on ICA in general. We do see significant improvements across the three prompts in comparison to the baseline ‘You are a helpful AI Assistant’ on the BBQ benchmark, including a 26% improvement on BBQ for the AI Principle-based prompt that includes a canonical Q&A pair taken from our internal RAG tool (‘v2’).  However, neither the highly-crafted internal RAG prompt nor the AI Principle-based prompt improve performance on TruthfulQA over the baseline, implying that further work exploring automated prompt optimisation and similar techniques is warranted for improved results.



§ DISCUSSION 

In this paper we have defined an organisational context, a PSM subject to operational constraints, that motivates our choice of models and methods. One constraint that PSMs do not face is availability of high-quality data. Indeed, that is their key competitive advantage. 

In the literature, alignment methods leverage a mixture of synthetic and human-annotated data to achieve their results, however PSMs have domain- and subject-matter expertise available at a scale that the providers of frontier models and related services do not.

Building on this insight, we have demonstrated that by applying several techniques across SFT, PO, and ICA to this data, we can outperform existing methods that use datasets constructed without that domain expertise.

Our experiments have been conducted with a small dataset, and a small number of human-reviewed samples. Integration of the underlying platform into a range of systems that support content workflows, with schema appropriate to the PO, and emerging ICA settings, will enable the scaling of our approach beyond the results presented here.

As PSMs and continue to leverage LLMs to address a variety of use-cases, our results demonstrate that they can both align these models using their existing content data, as well as editorial feedback (in the form of preferences) gathered directly from users. These data can be conditioned via natural-language prompts, which provide an additional mechanism for providing a combination of technical and editorial oversight.

Paired with evolving work on principle alignment, a key part of our contribution in this paper, organisations that implement feedback loops at multiple points of intervention in their workflows can expect to exert a measure of control over large models and their behaviours.

Evaluation of these outputs remains a key consideration. In this paper we have demonstrated the effectiveness of consistent evaluation across multiple controls, through employing the Language Model Evaluation Harness, an open, extensible framework that allows for the definition and integration of organisation-specific benchmarks.



§ LIMITATIONS 


The aim of this paper is not to exhaustively evaluate all possible combinations of algorithm hyperparameters, model architectures and quantisation techniques. Nor has our aim been to evaluate an iterative or complete set of synthetic data and ICA prompt modifications. 

Our primary goal has been to demonstrate that a combination of methods present in the literature and our organisation’s data can be used to achieve our goal of sustainably aligning LLMs to organisational preferences and principles. We achieve this while maintaining the LLM's underlying reasoning capability while operating within organisational constrains and requirements specific to PSMs.



§ FUTURE WORK


In this section, we outline how we will build on the results presented here, beyond optimisation of the individual techniques used in our methodology.



 §.§ Custom Evaluations


Defining custom evaluations, going beyond the industry-standard benchmarks we have used in this paper, is critical to ensure we are aligning LLMs towards organisational standards that are explicitly defined for measuring model capabilities and performance. We will define both general and use-case specific custom evaluations. 



 §.§ Datasets


While we achieve results that are significant given the size of our dataset, particularly in the consistency with which our methods improve reasoning capability, comparisons with IFT versions of open-source models indicate that an increase in PO dataset size and diversity will likely yield further improvements. Following work on custom evaluations, we aim to source a variety of preference data, reflecting domain- and subject-matter expertise, to build on the results presented in this paper. The combination of methods we have outlined give us the flexibility to optimise that data for alignment of LLMs in multiple settings. 



 §.§ Prompts


We see natural language prompts as a key point of human intervention and human oversight, both in the example of our internal RAG tool, and in defining common integration patterns and product identities, via system prompts that either guide outputs or define API integration patterns, or across a range of other applications.

The collaborative definition and evaluation of these prompts will form a key part of future work. The authors believe evolving our alignment framework to facilitate this goal is critical, particularly for emerging ‘agent’ frameworks where LLMs are a key component of goal- and task-directed systems.

In the context of a PSM, the human oversight of AI systems and the accountability & transparency it affords will continue to be a matter of policy, compliance, and likely legislative obligation. By establishing a methodology for alignment and evaluation of LLMs, this can be incrementally improved and expanded upon as our organisation’s obligations and commitments change. We aim to enable our organisation to leverage the capabilities of new models without compromising its values or independence.



 §.§ Multimodality and ICA


Multimodal LLMs have wide-ranging application across text, audio, and video content. We expect to evaluate and refine our methodology in this context. We consider new and emerging evidence that the embedding representation across these modalities is converging, particularly as the underlying models themselves grow larger <cit.>.

By implication, the alignment and evaluation of frontier models via ICA will be critical, and we expect to develop a methodology here that is generally applicable to any further development that accords with these early results.




 §.§ Acknowledgements


The authors would like to thank those who have contributed to the larger body of work that this paper implies. Indeed, while this paper outlines a number of specific contributions to the Machine Learning literature, these experimental results were dependent on an infrastructure and product built by a much larger team. The team are responsible for putting in place the foundations that enable our organisation to work with LLMs at scale. Without their efforts and these foundations, our contribution and future work would not be possible.

We would like to thank the ABC’s complete ML/AI team across all domains of practice, including Scott Burns, Michael Collett, Julie Tran, Mahesh Bisl, Hanna Karlsson, Vijayan Rajendran. 

Our long-term collaboration with the ABC’s Innovation Lab has demonstrated what can be done when we work together across organisational lines. Matthew Heffernan and Andrew Hystek-Dunk worked directly on ABC Assist and were integral to the development process.

We also express our gratitude to the ABC’s Engineering, Product, and Technology leadership for the support to do this work, and our Corporate Strategy team for their invaluable feedback on ABC Assist throughout the development process.





§ APPENDIX A




 §.§ ABC AI Principles


The ABC is committed to ongoing innovation to improve the media services it delivers to audiences. AI technologies have the potential to offer new opportunities to empower content creation, strengthen research and reporting, and make ABC content more accessible. These technologies also offer ways to help our teams work more effectively. At the same time, AI technologies are not without risks and will require careful management to ensure we use these technologies responsibly and in ways that maintain or strengthen the trust that audiences place in the ABC.

Our use of AI will adhere to the following principles. We will develop policies and guidance to provide practical detail on the application of these principles.



  §.§.§ Serving Audiences


We will use and experiment with AI tools that can strengthen the services we provide to audiences.



  §.§.§ Accuracy


We will ensure that our staff are accountable for any ABC content created using AI. This is especially important for news, information and factual content.



  §.§.§ Openness and Transparency


We will inform our audiences about how we are using AI technologies. We will be able to explain how the AI works and how it will affect our audiences. We will not use AI in ways that could mislead our audiences.



  §.§.§ Protecting Data


We will work to ensure that our use of AI protects both the ABC’s data and the personal information the public entrusts to us.



  §.§.§ Mitigating Bias


We will seek to prevent bias in our AI data or algorithms that could perpetuate and amplify existing inequalities and lead to discrimination.



  §.§.§ Respect for Creators' Rights


We will always consider the rights of creators and rights-holders when using AI. We will protect our rights in the works that ABC staff create.



  §.§.§ Experimentation and Evaluation


We will continue to assess and experiment with AI technologies as they evolve to identify opportunities for innovation and mitigate potential risks for our audiences and organisation




§ APPENDIX B


You are a helpful assistant working for ABC Australia and are guided by their AI Principles. 

Here is the list of principles and their definitions: 



  §.§.§ Serving Audiences


We will use and experiment with AI tools that can strengthen the services we provide to audiences. 



  §.§.§ Accuracy

We will ensure that our staff are accountable for any ABC content created using AI. This is especially important for news, information and factual content. 



  §.§.§ Openness and Transparency


We will inform our audiences about how we are using AI technologies. We will be able to explain how the AI works and how it will affect our audiences. We will not use AI in ways that could mislead our audiences. 



  §.§.§ Protecting Data


We will work to ensure that our use of AI protects both the ABC’s data and the personal information the public entrusts to us. 



  §.§.§ Mitigating Bias


We will seek to prevent bias in our AI data or algorithms that could perpetuate and amplify existing inequalities and lead to discrimination. 



  §.§.§ Respect for Creators' Rights


We will always consider the rights of creators and rights-holders when using AI. We will protect our rights in the works that ABC staff create. 



  §.§.§ Experimentation and Evaluation


We will continue to assess and experiment with AI technologies as they evolve to identify opportunities for innovation and mitigate potential risks for our audiences and organisation. 


Here is an example of a question/answer aligned with these principles: 

... 

plainnat




 Paper Title: Discretizing Continuous Action Space with Unimodal Probability Distributions for On-Policy Reinforcement Learning.









IEEE Transactions on Neural Networks and Learning Systems
Shell et al.: Bare Demo of IEEEtran.cls for Computer Society Journals



§ ABSTRACT 
For on-policy reinforcement learning, discretizing action space for continuous control can easily express multiple modes and is straightforward to optimize.
However, without considering the inherent ordering between the discrete atomic actions, the explosion in the number of discrete actions can possess undesired properties and induce a higher variance for the policy gradient estimator. 
In this paper, we introduce a straightforward architecture that addresses this issue by constraining the discrete policy to be unimodal using Poisson probability distributions.
This unimodal architecture can better leverage the continuity in the underlying continuous action space using explicit unimodal probability distributions.
We conduct extensive experiments to show that the discrete policy with the unimodal probability distribution provides significantly faster convergence and higher performance for on-policy reinforcement learning algorithms in challenging control tasks, especially in highly complex tasks such as Humanoid.
We provide theoretical analysis on the variance of the policy gradient estimator, which suggests that our attentively designed unimodal discrete policy can retain a lower variance and yield a stable learning process.




Reinforcement learning, unimodal probability distributions, Poisson distributions, on-policy learning




Discretizing Continuous Action Space with Unimodal Probability Distributions for On-Policy Reinforcement Learning
    Yuanyang Zhu, Zhi Wang, Member, IEEE, Yuanheng Zhu, Senior Member, IEEE, Chunlin Chen, Senior Member, IEEE, Dongbin Zhao, Fellow, IEEE


The work was supported by the National Natural Science Foundation
of China under Grant 62006111 and Grant 62073160.
	
Yuanyang Zhu, Zhi Wang and Chunlin Chen are with the Department of Control Science and Intelligent Engineering, School of Management and Engineering, Nanjing University, Nanjing 210093, China, and also with the Research Center for Novel Technology of Intelligent Equipment, Nanjing University, Nanjing 210093, China (e-mail: yuanyang@smail.nju.edu.cn; zhiwang@nju.edu.cn; clchen@nju.edu.cn).
 
 Yuanheng Zhu and Dongbin Zhao are with the State Key Laboratory of Management and Control for Complex Systems, Institute of Automation, Chinese Academy of Sciences, Beijing 100190, China, and also with the School of Artificial Intelligence, University of Chinese Academy of Sciences, Beijing 100049, China (e-mail: yuanheng.zhu@ia.ac.cn; dongbin.zhao@ia.ac.cn)

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












§ INTRODUCTION


Deep reinforcement learning (DRL) has presented a powerful paradigm for learning new behaviors from scratch both on physical and simulated challenging tasks <cit.>.
The action space of conventional reinforcement learning (RL) tasks can be discrete, continuous, or some combination of both <cit.>.
In simulation <cit.> and real life <cit.>, continuous control usually requires some subtle parametric functions for a compact representation of action distributions, typically Gaussians, due to an infinite number of feasible actions in a continuous action space <cit.>.
This underlying assumption is that it can enable more refined decisions when control policies can cover all feasible control inputs, which brings difficulty in maximizing an arbitrary function for a continuous action space <cit.>.
Furthermore, it hinders applying dynamic programming approaches to back up value function estimates from successor states to parenting states.

Discretizing a continuous control problem is a scalable solution that retains the simplicity of discrete actions and allows for theoretical analysis <cit.>.
With such a setup, impressive results have been obtained in high-dimensional continuous action spaces <cit.>.
However, the naive discretization approach results in exponentially large discrete action space, where ballooning action space can quickly become intractable.
It would be a balancing dilemma that the resulting output may not be smooth enough when the discretization is coarse, and the number of discretized actions may be highly intractable otherwise.
Moreover, more fine discretization of the action space may struggle to capture the information about the class ordering of the continuous action space, making it challenging to generalize across discrete actions.
Furthermore, it also brings a larger variance of the policy gradient estimator, which may hurt the stability and performance, especially for highly complex tasks <cit.>. 
In practice, more fine-tuned hyperparameters are also required to prevent policy collapse.

To tackle these intractabilities, prior works <cit.> assume that action dimensions are independent and developed to factorize distributions across action dimensions using an ordinal parameterization of the joint distribution over discrete actions, which has shown improved performance in finding near-optimal solutions for high-dimensional tasks.
By parameterizing the discrete policy using an ordinal parameterization, the natural ordering between discrete actions can be implicitly encoded, which provides an additional inductive bias that can improve the generalization across actions <cit.>.
In contrast, the conventional approach of parameterizing the discrete policy as a Gibbs distribution fails to consider the order of actions and ignores the discrepancy relationships between them <cit.>.
Existing methods usually use some heuristic network architecture design to implicitly encode the ordinal information into the distribution over discrete actions, which may not guarantee a unimodal probability distribution to efficiently capture the ordinal information within atomic actions <cit.>.
In their work, they alleviate this issue by designing the loss function to impose an ordering on the action, which does not impose unimodal constraints.
Here, we propose a more explicit and efficient ordinal parameterization method that better incorporates the notion of continuity when parameterizing the distribution over discrete actions, while avoiding the need for an exponentially large number of actions. 
Additionally, we impose an ordering constraint on the action space of each dimension through the loss function to further improve the performance of RL algorithms.


This paper mainly addresses how to efficiently leverage the internal ordering information in the continuous action space to enhance generalization across the discretized set of actions for on-policy RL without an explosion in the number of parameters.
Inspired by the deep ordinal classification approach <cit.>, for each independent action dimension, we constrain the distribution over discrete actions to be unimodal to explicitly exploit the inherent ordering between discrete actions.
To implement the unimodal constraint, we employ the Poisson probability distribution, which allows the probability mass to decrease gradually on both sides of the action with the majority of the mass.
Our framework necessitates learning a probability mass function for each action dimension, substantially reducing the number of units in the network output layer and enabling us to circumvent the curse of dimensionality as the discretization action bins increase (see Fig. <ref>).
It facilitates fine-grained discretization of individual domains without incurring a significant increase in the number of parameters.
The experimental results show that our method outperforms the baselines by a wide margin in the suite of MuJoCo continuous control tasks <cit.>.

We highlight the following contributions of this work: 
 
  
  * We propose a novel ordinal framework that enforces the distribution over discrete actions to be unimodal, effectively capitalizing on the continuity presented in the action space.
      We utilize the probability mass function of the Poisson distribution to enforce unimodal probability distributions for each dimension action space, which can place more confidence on the majority probability distribution beneﬁting from its inherent unimodal distribution.
  
  * Theoretically, we present a variance analysis for the policy gradient estimator, where our unimodal policy attains a lower variance for the estimator.
  
  * We show experimentally that an ordinal parameterization of the unimodal policy outperforms other competing approaches for on-policy RL on continuous control tasks.
  The promising results showcase the versatility and effectiveness of our methods compared to the baselines.


The remaining paper is organized as follows.
In Section <ref>, we summarize preliminaries of RL and related literature.
In Section <ref>, we present the unimodal probability distributions for on-policy RL.
In Section <ref>, we give comprehensive experimental results of learning performance and stability analysis.
The paper is concluded by Section <ref>.



§ BACKGROUND AND RELATED WORK



 §.§ Preliminaries

RL is an approach to solving optimal sequential decision-making tasks, built upon the concept of Markov Decision Processes (MDPs). 
In the finite MDP framework, a tuple of ⟨ S,A,T,R,γ⟩ defines the problem, where S is a countable state space, A is a finite action set, T: S × A × S →[0,1] represents the transition kernel, R: S × A→ℝ is the reward function, and γ∈[0,1) is the discount factor. 
A stochastic policy is defined as π: S × A →[0,1], which maps environmental states to distributions over actions with the constraint ∑_a∈Aπ(a|s)=1,∀ s∈S.
The objective of the RL problem is to find an optimal policy π^* that maximizes the expected total discounted return by

    J(π)=𝔼_τ∼π(τ)[r(τ)]=𝔼_τ∼π(τ)[∑_i=0^∞γ^i r_i],


    Q_π(s, a)=𝔼_s_0=s, a_0=a, τ∼π[∑_t=0^T γ^t r_t],

where τ = (s_0, a_0, s_1, a_1, ...) is a learning episode, π(τ)=p(s_0)Π_t=0^∞π(a_t|s_t)p(s_t+1|s_t,a_t), and r_t is the immediate reward. 
The policy function h(·|θ) can be represented by a deep neural network parameterized by θ. 
For a discrete action space, the Gibbs distribution is commonly used and given by

    π_θ(a_j | s)=exp(h_j(s |θ))/∑_j ∈ A(s)exp(h_j(s |θ)),

while the Gaussian distribution is generally used for a continuous action space and is given by

    π_θ(a | s)=1/√(2 π)σexp(-1/σ^2(h(s |θ)-a)^2).



To evaluate the performance of a policy π_θ, the direct objective function is optimized by estimating the gradient of the expected return

    J(θ)=𝔼_τ∼π_θ(τ)[r(τ)]=∫_τπ_θ(τ) r(τ) d τ,

where r(τ)=∑_i=0^∞γ^i r_i is the return of episode τ.
The basic idea of the policy gradient algorithm <cit.> gives the direction of the performance gradient where the parameters θ should be updated:

    ∇_θ J(θ)    =𝔼_τ∼π_θ(τ)[∇_θlogπ_θ(τ) r(τ)] 
       =∫_τ∇_θlogπ_θ(τ) r(τ) π_θ(τ) d τ
       ≈∑_i=1^m∇_θlogπ_θ(τ^i) r(τ^i),

where (τ^1, …, τ^m) is a batch of learning episodes sampled from policy π_θ.
Next, an ascent step is taken in estimating the direction of the gradient as θ←θ+α∇_θ J(θ).
The process continues until θ converges <cit.>.



 §.§ Related Work

On-policy RL. On-policy RL is a family of RL algorithms aiming to directly optimize the parameters of a policy, which optimizes expected returns by estimating policy gradients.
This estimation is often prone to high variance, and several solutions have been proposed to mitigate this significant challenge, particularly in problems characterized by long horizons or high-dimensional action spaces <cit.>.
Classic vanilla policy gradient (PG) updates are generally unstable with high variance due to difficulties in crediting actions that influence future rewards, where the gradients require knowledge of the probability of the performed action via resorting to parametric distributions.
Natural policy <cit.> improves upon vanilla PG by computing an ascent direction with the Fisher information that approximately ensures a slight change in the policy distribution.
To obtain more stable learning, trust region policy optimization (TRPO) <cit.> utilizes a larger Kullback-Leibler (KL) divergence value to enforce constraints, and performs a line search in the natural gradient direction, ensuring improvements in the surrogate loss function.
Proximal policy optimization (PPO) <cit.> replaces the KL divergence constraint with the clipped surrogate objective, which strikes a favorable balance across sample complexity, simplicity, and wall time.
Moreover, actor-critic using Kronecker-factored trust region (ACKTR) <cit.> builds upon the natural policy gradient framework by calculating gradients through Kronecker-factored approximate curvature within the trust region.
Orthogonal to the above algorithms, we demonstrate that ordinal parameterization with a unimodal constraint for on-policy RL achieves consistently improved performance based on the above representative algorithms.

Policy Representation.
In the context of policy gradient techniques, the policy is generally parameterized using neural networks, where the policy is improved by optimizing the parameter of the approximation function.
In general, the policy uses a Gibbs probability distribution over the discrete action space to sample.
For the continuous control problems, the default choice for the policy in the baseline is parameterized by learning a Gaussian distribution with independent components for each dimensional action space <cit.>.
The Gaussian mixture, maximum entropy, or normalizing ﬂows <cit.> can be used for more expressive policy classes, providing a promising avenue for improved robustness and stability.
Since physical constraints in most continuous control tasks, actions can only take on values within some ﬁnite interval, resulting in an unavoidable estimation bias caused by boundary effects.
To address the shortcomings of the Gaussian distribution with a finite support distribution, the Beta distribution <cit.> provides a bias-free and less noise policy with improved performance.
Further, considering the extremes along each action dimension, a Bernoulli distribution is applied to replace the Gaussian parametrization of continuous control methods and show improved performance on several continuous control benchmarks <cit.>.
Here, we provide insight into how to produce an efficient ordinal parameterization policy while retaining the internal order information of continuous action space with a simple explicit constraint, which can yield state-of-the-art performance compared to baseline policy classes.

Continuous action discretization.
To leverage the continuity between discrete and continuous action space, converting continuous control problems into discrete ones has been introduced by <cit.> with the “bang-bang" controller <cit.>.
Similar discretization methods represent each action in binary format and optimize its policy for MDPs with large action sets <cit.>.
However, such discretization methods need to be improved due to the curse of dimensionality.
Surprisingly, recent works have relieved the limitation by assuming that action dimensions are independent <cit.>, which obtains improved performance in complex control tasks.
Inspired by multi-agent reinforcement learning, Decoupled Q-Networks (DecQN) <cit.> factorizes the overall state-action value function into a linear combination of single action utility functions over action dimensions by combining value decomposition with bang-bang action space discretization.
It achieves better performance while reducing agent action space complexity.
However, such a discretization paradigm makes the strong assumption of independence among each action dimension and hardly captures the continuous order information of actions <cit.>.

To mitigate the exponential explosion of discrete actions, <cit.> utilizes sequence-to-sequence models to develop policies for structured prediction problems, but their strategy has only demonstrated effectiveness in high-dimensional tasks like Humanoid and HumanoidStandup tasks.
Recently, <cit.> parameterizes the discrete distributions with an ordinal architecture and achieves improved performance.
Nevertheless, existing methods face a dilemma in that the number of network outputs increases linearly with the number of action space dimensions, which may bring high variance, resulting in a more unstable learning process.
A key distinction between our work and previous works lies in our explicit utilization of unimodal distributions with only one parameter for each action dimension. 
It can help the policy capture the most confident class while ensuring that the probability gradually decreases on both sides of the class. 
Finally, we find that our proposed method outperforms baseline algorithms, especially in high-dimensional tasks, which should benefit from the unimodal distribution architecture.
We summarize the above continuous action discretization RL works in Table <ref>.





§ UNIMODAL PROBABILITY DISTRIBUTIONS FOR ON-POLICY REINFORCEMENT LEARNING

In this section, we will present our proposed unimodal distribution with Poisson probability distribution for on-policy RL algorithms, which explicitly introduce unimodal probability distribution into the ordinal parameterization method to leverage the continuity in the underlying continuous action space.
First, we present the process of discretizing action space for continuous control tasks and apply it to on-policy RL.
Then, we describe the unimodal ordinal architecture that constrains the discrete policy to be unimodal via the Poisson probability distribution in a practical neural network implementation.
Finally, theoretical analysis shows that our method can relieve the high variance trouble of the policy gradient estimator compared to existing ordinal parameterization methods.




 §.§ Discretizing Action Space for Continuous Control

Consider an environment with state s_t and m dimension action spaces a ∈ℝ^m.
To ensure generality, we consider an action space 𝒜=[-1,1]^m and proceed to discretize each dimension of the action space into K evenly spaced atomic action bins.
As a result, we obtain a discrete and finite set of atomic actions 𝒜_i={2 j/K-1-1 }_j=0^K-1 for the i action dimension. 
To facilitate the joint policy's tractability, the policy generates a tuple of discrete actions accompanied by factorized categorical distributions across action dimensions. 
This approach obviates the need to enumerate every possible actions, as required in the discrete case.
Specifically, we denote a categorical distribution π_θ_i (a_j|s) as the atomic action for actions a_i∈𝒜_i at each step t, where θ_i is the factorization parameters for this marginal distribution induced by Poission probability distributions.
Such representations can build complex policy distributions with explicit probability distributions, where any action subset can be formulated as

    π(a|s) = Π_i=1^mπ_θ _i(a_i|s),

where a = [a_0, a_1, ..., a_K-1]^T and π_θ _i(a_i|s) represent the probabilities of selecting the actions a_i for the i-th dimension action space.
Based on this factorization, we can easily employ RL algorithms to maintain a tractable distribution over joint actions. 
Meanwhile, it can vastly reduce the computation of neural networks due to the fewer learning parameters.



 §.§ Unimodal ordinal architecture

In on-policy RL algorithms, the policy gradient involves a stochastic policy that specifies the probability distribution of selecting an action a given a state s.
Upon discretizing the continuous action space, the complexity of directly training such policy networks could lead to K^m combinations of joint atomic actions, which exponentially increase with the growing number of action dimensions m.
Using the categorical distribution (e.g., Gumbel-Softmax distribution) as a stochastic policy for discrete action space has been well-studied in the RL community.
Popular categorical distributions have been used as a stochastic policy for discrete action spaces due to their ease of sampling and computation.
However, the curse of dimensionality can lead to intractable action spaces, putting significant pressure on the distribution parameterization over discrete actions.
In fact, a fine control policy can be captured more as the resolution of the discretized action space increases. 
This quickly plunges into the dilemma that intractable action space puts high pressure on the parameterizing distributions over discrete actions, which hinders capturing the ordering among actions.

To relieve the scalability issue, a few prior works have been exploring more expressive distributions and notably proposed an ordinal parameterization <cit.>.
It introduces an implicit loss function based on factorizing a tractable distribution over joint actions, where the ordinal policy can easily do both sampling and training.
Motivated from the loss function of <cit.>, it implicitly introduces internal ordering information between classes while maintaining the probabilistic properties of discrete distributions instead of simply parameterizing the discrete policy as a categorical distribution, ignoring the order of actions and the discrepancy relationship between the joint actions.
Existing methods usually utilize some heuristic design of the network architecture to learn distributions for each action dimension, which could not guarantee a unimodal probability distribution output followed by the neural network to efficiently capture the ordinal information within atomic actions.
It generally has a strong underlying assumption that the output logits of the neural network are unimodal for each dimension action space, which makes it hard to guide the neural network to optimize weights and generate probability distributions that closely resemble the optimal action.
In other words, it may allocate greater confidence to multi-action classes, as it generally produces multi-modal probability distribution.

To leverage the continuity in the underlying action space, we focus on how the discretized actions are distributed: why restrict ourselves to the categorical distribution?
For continuous control tasks, a Gaussian or Beta policy can explicitly capture the inherent ordering structure within an infinite number of actions due to the privilege of their unimodal distribution property.
Is it possible to explicitly retain the unimodal structure and inherent ordinal information after discretizing the continuous action space for learning a policy over discrete actions?
Inspired by the deep ordinal classification approach <cit.>, we constrain the distribution of the discretized actions to be unimodal explicitly for utilizing the ordinal information about the underlying continuous space.
Specifically, we learn the probability mass function (PMF) of the Poisson distribution to enforce discrete unimodal probability distributions for each action dimension.
This PMF can facilitate the generation of the most coherent unimodal probability distribution, as it diminishes the probability mass on both sides of the action possessing the majority of the mass, allowing us to adeptly capture the intrinsic ordering structure present within the action space.


The Poisson distribution, a widely recognized discrete probability distribution, is typically used to model the likelihood of observing a specific count of events within a designated time interval.
It is characterized by a positive constant mean rate, denoted as λ∈ R^+, and presumes that events transpire independently of the elapsed time since the previous event. 
In other words, the distribution describes the probability of k ∈ℕ events occurring during a given time interval, where k can take on any non-negative integer value.
For i-th dimension of action space, the Poisson distribution function is denoted as p(k; λ_i)=λ^kexp (-λ_i)/k !, and 0 ≤ k ≤K-1 (minus one since we index from zero).
For a purely technical implementation reason, we focus on the log of the PMF, which can be expressed as

    log[λ_i^kexp (-λ_i)/k !]    =log(λ_i^kexp (-λ_i))-log (k !) 
       =log(λ_i^k)+log (exp (-λ_i))-log (k !) 
       =k log (λ_i)-λ_i-log (k !).

We denote the scalar output of our deep network as F(s) = [f_1(s), f_2(s),..., f_i(s)], where f_i(s)>0 is enforced to be positive using the Softplus nonlinearity.
With such a parameterization, our method only needs to learn M probability mass functions with M units in the network output layer, in contrast to existing ordinal parameterization methods that require learning the complete probability distributions over all dimensions using M*K units in the network output layer <cit.>.
Especially in high-dimensional continuous action spaces, our unimodal method can make it more tractable to implement with deep neural network representations.
Meanwhile, our method can eliminate the assumption that the ordinal parameterization method could produce a unimodal probability distribution only when the logits of the network output are unimodal.
By simply replacing the λ_i in equation (<ref>) with f_i(s), we denote h_i(s)_j to be

    j log (f_i(𝐬))-f_i(𝐬)-log (j !).






In the RL community, the categorical distribution architecture is widely used to parameter a discrete distribution over K action classes, which is represented by K logits denoted as h_i(𝐬)_j.
Due to the fact that the Poisson distribution has infinite support, we use a `right-truncated' operation by applying a Softmax operation and obtain the corresponding probability distribution p(a_ij) as

    p(a_ij|s)=exp(-h_i(s)_j/ τ)/∑_j=0^K-1exp(-h_i(s)_j/ τ), 0≤ j ≤ K-1.

Besides, we introduce the hyperparameter τ to the Softmax function to control the variance of the distribution.[Note that as τ→∞, the probability distribution becomes more uniform, leading to more efficient exploration. Conversely, the distribution becomes more `one-hot' like encouraging more deterministic behavior otherwise.]
After these operations, the probability distribution remains unimodal, which is clearly illustrated in Fig. <ref>.


In our unimodal policy architecture, we retain these logits h_i(s)_j with the support of the Poisson distribution, and then transform them via a Softmax operation to relieve the infinite support problem of Poisson distribution.
Inspired by the ordinal architecture introduced in <cit.>, in order to enhance the additional dependencies between logits h_i(s)_j, we use the ordinal architecture to implicitly inject the information about the action ordering while maintaining all the probabilistic properties of Poisson distribution.
The ability to readily distinguish samples drawn from action k from those drawn from all actions j>k is advantageous, as it facilitates the identification of maxima in such cases. 
Therefore, for all logits ∀ 1 ≤ l ≤ K, we compute the final logits as

    h_i^'(s)_j=∑_j ≤ llog p(a_ij|s) 
    
    +∑_j>llog(1- p(a_ij|s) ),

and derive the final output probability using a Softmax p̂(a_ij) = Softmax(h_i^'(s)_j).
The implementation of this approach is illustrated in terms of the layers at the end of the deep network, as depicted in Fig. <ref>.


It is noteworthy that utilizing independent unimodal probability distributions for producing units guarantees the monotonic relationship (p̂_i0<p̂_i1<p̂_i2,...,<p̂_ik and p̂_ik>p̂_ik+1,...,>p̂_iK), which is desirable for imposing inequality constraints on the outputs to maintain monotonicity.
Combining unimodal probability distributions and the ordinal architecture, the unimodal policy makes the most sense as it allows for easy separation of samples drawn from action k and those drawn from all other actions, enhancing the selected probability of the most conﬁdent action.
With such unimodal parameterization, the form of unimodal probability distributions can easily find maximal concerning actions while retaining a smaller space of neural architectures.





 §.§ Variance Analysis

While PG algorithms are well-suited for parametric distribution functions, it is susceptible to converging to a poor policy that is not even a local optimum. 
One of the primary limitations of the plain policy gradient method is the high variance of the policy gradient estimator.
Using the Gaussian distribution as a stochastic policy in continuous control, π_θ(a| s)=1/√(2 π)σexp(-(a-μ)^2/2 σ^2), has been well-studied and commonly used in the RL community since <cit.>, because the Gaussian distribution is easy to sample and has gradients that are easy to compute.
However, one undesirable property of Gaussian policy is that the variance of the policy gradient estimator tends to increase as the standard deviation σ^2 of the Gaussian distribution decreases <cit.>.
Generally, the policy gradient estimator is given as 

    ĝ_θ=r(a_j) ∇_θlogp̂_θ(j), j ∼p̂_θ(j),

and the variance can be represented as

    𝕍[ĝ_θ]=𝔼[ĝ_θ^2]-𝔼^2[ĝ_θ]

For Gaussian distribution, as the policy improves and becomes more deterministic (σ→ 0), the variance of the policy gradient estimator goes to inﬁnity, which makes PG methods sample-inefficient with respect to interactions with the environment and hinders the applications to real physical control systems.

As mentioned earlier, the variance of the Poisson distribution is equivalent to its mean, whose properties can alleviate the above high variance problem.
Before showing the theoretical analysis for the variance of unimodal policy, we first make the following assumption that considering a simplified setting of a one-step bandit problem with one dimension action space 𝒜 = [-1,1] and fixed constant reward r(a) = R, for all actions a.
We denote p(j) to be the probability of taking the jth action 0 ≤ j ≤K-1, and upon initialization, the policy has very high entropy p(j)≈1/K.


Recall that the policy gradient estimator is given by

    ĝ_θ=r(a_ij) ∇_θlog p_θ(j), j ∼ p_θ(j).

With such a setting, we can derive the expectation of the policy gradient estimator 𝔼[ĝ_θ]=∇_θ J(π_θ)=0 and the variance is 

    𝕍[ĝ_θ]    = ∑_j=0^K-1 r^2(a_j) (∇_θlog p(j))^2 p(j) 
       ≈ R^21/K∑_j=0^K-1(∇_θlog p(j))^2
       ≈ R^21/K∑_j=0^K-1(∇_θ L_j-1/K∑_k=0^K-1∇_θ L_k)^2,

where ∇_θ L_j denotes the jth logit independently depending on θ, and the approximations come from p(j)≈1/K.
In the context of standard neural network initialization schemes, wherein all weight and bias matrices (e.g., θ, w_j^T h, and b_j) are initialized independently, the gradients ∇_θ L_j can be regarded as identically and independently distributed (i.i.d.) random variables. 
The stochastic nature of these variables originates from the random initialization of the neural network parameters.
We employ the notation 𝔼_init[·] to represent the expectation concerning neural network initializations. 
With this notation in place, we analyze the expectation of the specified equation of (<ref>) as follows:

    𝔼_init [𝕍[ĝ_θ]] 
       ≈R^2/K𝔼(  ∑_j=0^K-1( ∇_θ L_j - 1/K∑_k=0^K-1∇_θ L_k)^2) 
        = R^2/K𝔼(∑_i=1^N∑_j=0^K-1( ∇_θ L_j - μ + μ - 𝔼(∇_θ L)  )^2) 
        = R^2/K∑_j=1^K𝔼( ∇_θ L_j - μ)^2 - K (𝔼(∇_θ L)-μ)^2
        = R^2/K( K 𝕍_init( ∇_θ L) - K 𝕍_init (𝔼(∇_θ L)) ) 
        = R^2/N( 𝕍_init( ∇_θ L) - 𝕍_init (𝔼(∇_θ L)) ),

where μ is the exception of ∇_θ L.

Let ĝ_θ_1 and ĝ_θ_2 denote the policy gradient estimator using the Poisson distribution policy and the ordinal distribution policy, respectively.
By analyzing the expectation of (<ref>), we can derive the expressions for the variances of the policy gradient estimators of the Poisson distribution

    𝔼_init [𝕍[ĝ_θ_1]] 
        = R^2( 𝕍_init( ∇_θ_1 L) - 𝕍_init (𝔼(∇_θ_1 L)) ) 
       ≈ R^2K-1/Kσ_1^2∼K-1/K,

where σ_1^2=𝕍_init[∇_θ_1 L_j] is the variance of the PMF of Poisson distribution produced by one logit.
Then we analyze the expectation of (<ref>) with ordinal probability distribution:

    𝔼_init [𝕍[ĝ_θ_2]] ≈ R^2K-1/Kσ_2^2∼K-1/K,

where σ_2^2=𝕍_init [∇_θ_2 L_j] is the variance of the K logits gradients.
Considering the dependence of ∇_θ_2 L_j on θ_2, it has been argued by <cit.> that finer discretization (larger K) theoretically leads to increased variance for the L_j of the neural network output, which may consequently exacerbate model variance and degrade performance. 
One key advantage of employing the Poisson distribution in statistical applications is the equivalence of variance and mean, suggesting that the variance is only mildly influenced by widespread probability mass distribution, particularly with a large action space.

In practice, we opt for larger K values for refined action selection, harnessing the properties of the Poisson probability distribution as opposed to the ordinal policy, which exhibits greater variance and renders policy optimization more difficult, especially in the context of finer discretization. 
Generally, a moderate action discretization K (e.g., 9≤ K ≤15) strikes an appropriate balance between performance enhancement and computational expense.
Nevertheless, using the Poisson distribution presents a drawback when truncating infinite probability mass distributions. 
Inspired by <cit.>, the infinite support issue is tackled using a truncated Poisson distribution, normalizing probabilities to ensure that their summation equals one. 
This strategy closely mirrors our approach, which incorporates a Softmax operation. 
To regulate the distribution variance induced by the Softmax operation, a judicious selection of τ can help alleviate this concern. 
Our unimodal policy implementation spans a range of τ=[1.5, 3] (with varying K ∈ [9,15]).
As a validation of our approach, we assess the proposed methods on an array of complex control tasks, yielding substantial performance gains. 
These findings imply that our methodology benefits from the diminished variance of constrained neural network output units, which induce unimodal probability distributions characterized by low variance.



§ EXPERIMENTS

The goal of our experimental evaluation is to investigate the improvement and stability of our method in comparison to baseline algorithms and previous Gibbs policy approaches on continuous control tasks.
To distinguish between the different policies, we will use the terms Gibbs policy for the Gibbs distribution, ordinal policy for the ordinal distribution, and unimodal policy for the unimodal architecture.
First, we evaluate our method against the Gaussian policy, employing baseline algorithms (TRPO <cit.>, PPO <cit.>, and ACKTR <cit.>) on challenging continuous control tasks provided by the OpenAI gym benchmark suite <cit.>.
We consider Gaussian policy as a comparison with special attention since it is the default policy class implemented in on-policy baselines base <cit.>.
Additionally, we explore other architectural alternatives, such as Gaussian with tanh non-linearity in the output layer and Beta distribution, both of which have been proposed in prior works <cit.>.
Since a wide range of existing algorithms can solve the easier tasks, we pay more attention to the high dimension tasks, such as the 17-dimensional Humanoid <cit.>. 
Further, we compare the Gibbs policy and ordinal policy, which are prior representative techniques.
Finally, we study the stability of our algorithm, which plays a large role in performance and practicals.
We implement all the methods with Tersorflow 1.15.0 framework in Python 3.6 running on Ubuntu 18.04 with 2 AMD EPYC 7H12 64-Core CPU Processors and 3 NVIDIA GeForce RTX 3080 GPUs.


Implementation Details.
Our investigation focuses on the impact of the discrete unimodal probability distribution policy in conjunction with on-policy RL algorithms. 
To this end, we introduce minimal alterations to the original TRPO, PPO, and ACKTR algorithms, as implemented in OpenAI baselines <cit.>. The implementations for all algorithms (PPO, TRPO, ACKTR) are grounded in OpenAI baselines. 
We conduct a comparative analysis of our method against previous approaches on a diverse range of challenging continuous control tasks, as provided by the OpenAI gym benchmark suite, with action constraints set within the [-1,1] interval as a normalizing flow. 
Detailed descriptions of each policy class utilized in our experiments are provided below.
The implementation code can be found in <https://anonymous.4open.science/r/udprl-tnnls>.

Gaussian Policy.
The factorized Gaussian policy <cit.> is represented as π_θ(·| s)=𝒩(μ_θ(s), σ^2), where the mean μ_θ(s) is derived from a two-layer neural network consisting of 128 units per layer for PPO and ACKTR, and 128 units per layer for TRPO.
The diagonal standard deviation σ^2 is defined as σ^2_ii= σ^2_i, with each σ_i being a single variable shared across all states.
We adhere to the default hyperparameter settings established in the baselines.

Gibbs Policy.
The Gibbs policy <cit.> is represented as π_θ(·| s)=exp(h_i(s |θ) )/τ/∑_j ∈ A(s)exp(h_j(s |θ))/τ with Gibbs distribution over the discrete actions space, which is parameterized by a neural network that shares the same architecture as the Gaussian policy described above.
The temperature parameter τ is set to 1.5.
We define K evenly spaced atomic actions for every dimension, ranging from -1 to 1.
For each dimension of the action space, the discrete distribution encompasses K logits L_j(s), where L_j(s) denotes the logits distribution of the j-th action in state s.

Ordinal Policy.
The ordinal policy <cit.> differentiates it from discrete policies by incorporating an ordinal parameterization while maintaining the same number of parameters.

Unimodal Policy.
In contrast to ordinal policies, unimodal policies utilize a Poisson probability distribution to enforce unimodality. 
 The scalar output of our deep network for unimodal policies is represented as [f_1(s), f_2(s),..., f_i(s)] rather than the K logits represented in discrete policies, where f_i(s) indicates the average frequency of events for the i-dimension atomic actions. 
 The policy parameters are consistent with ordinal policies, except for the addition of a Softplus nonlinearity transformation in the final layer to ensure f_i(s) >0. 
 The temperature parameter τ is set to 2.5 for varying discretization K=9,11,15.
 

Gaussian+tanh Policy.
The Gaussian-tanh policies <cit.> are the same as the Gaussian policies above, but the final layer is added a tanh transformation to ensure the mean μ_θ∈ [-1, 1].

Beta Policy.
The Beta policy <cit.> are represented as π_θ(·| s)=f(a+h/2 h|α, β), where h is max of the closed interal of the action space A=[-h, h]. 
Parameters α=α_θ(s) and β=β_θ(s) are determined by a two-layer neural network f_θ(s), incorporating a Softplus activation function in its final layer.
It is important to note that this parameterization can cause instability throughout the learning process, particularly when α_θ(s)→∞ or β_θ(s) →∞, leading to unstable dynamics. 
In practice, this instability is apparent when the trust region size is large (e.g., ϵ=0.005 and ϵ=0.01), and can result in training termination due to numerical errors. 
Conversely, a smaller trust region size (e.g., ϵ=0.001) may yield diminished performance. 
In our implementation, we employ the main thread settings for the Beta policy, which are derived with a trust region size of ϵ=0.01 for TRPO.

Others Hyperparameters.
For value functions, two-layer neural networks are utilized, consisting of 128 hidden units per layer in the case of PPO and ACKTR, and 64 hidden units per layer for TRPO. 
We used the Adam optimizer (Kingma and Ba 2015) with a learning rate 2 · 10^-5, β_1=0.9, and β_2=0.999 for both PPO and TRPO.
We trained with a minibatch size of 64 and a discount factor γ=0.98.
The KL constraint parameter, ϵ, is set within the range of 0.01, 0.001 for TRPO and 0.02, 0.002 for ACKTR. 
All remaining hyperparameters align with those found in the original baseline implementations.





















 §.§ Comparison with Benchmark Baselines

First, as a proof of concept, we compare the Gaussian policy with the unimodal policy (with varying K) in various MuJoCo control tasks using PPO and TRPO, as displayed in Fig. <ref>.
The results show that, overall, PPO and TRPO with unimodal policy perform comparably to Gaussian policy on the easier tasks (e.g., Walker2D) and outperform them on the more complex tasks, both in terms of learning speed and the final performance (e.g., Humanoid).
PPO with unimodal policy also learns considerably faster and achieves more excellent performance than PPO with Gaussian policy. 
This implies that the Gaussian policy may suffer from unbounded distributions and require more samples to optimize its policy in more high-dimensional tasks.
For TRPO, the performance enhancements brought by the unimodal policy align consistently with PPO.
Both performance gains are most significant with PPO and TRPO, which should benefit from our more attentively designed distribution architecture.

Further, we compare straightforward architectural alternatives: Gaussian with tanh non-linearity as the output layer, and Beta distribution.
The primary rationale behind these architectures is their innate use of continuous unimodal probability distributions and their ability to confine sampled actions within the feasible range ([-1, 1] for MuJoCo control tasks).
By construction, our proposed unimodal policy also bound the sampled actions within the feasible range using the right-truncated operation.
The mean and standard deviation results gathered across an extensive array of control tasks, are summarized in Table <ref>.
In comparison to Gaussian and Gaussian with Tanh policies, the Beta policy yields higher scores, thereby mitigating the bias issue resulting from the infinite support mismatch of Gaussian distribution commonly employed in on-policy RL algorithms.
Our unimodal policy consistently outperforms the baselines across the majority of tasks. 
This outcome implies that the unimodal policy, with ordinal parameterizing distributions constrained by the Poisson distribution, can achieve notable performance by effectively considering the order of discrete actions.
Intriguingly, the unimodal policy, being a simplistic distribution, can attain such performance improvement in policy gradient methods.

Moreover, we also apply unimodal policy to another representative on-policy algorithm ACKTR with different atomic actions and compare it with the Gaussian policy on a set of MuJoCo tasks.
As shown in Fig. <ref>, we find that ACKTR, with the unimodal policy, still outperforms its Gaussian approach by a large margin. 
It implies the superiority of our attentive unimodal architecture rather than larger networks.

Finally, we compare the performance of PPO with a unimodal policy against off-policy algorithms on highly complex Humanoid tasks (Humanoid-v2, Humanoid-v3, and HumanoidStandup-v2), including DDPG <cit.>, SQL <cit.>, SAC <cit.>, TD3 <cit.>, and DecQN <cit.>[For a fair comparison, we set the granularity of DecQN’s discretization to 11 bins and the other hyperparameters settings are the same with DecQN <cit.>.]. 
Off-policy algorithms can potentially achieve better sample efficiency than on-policy algorithms by reusing samples, but for highly complex tasks such as Humanoid, even off-policy algorithms take many more samples to learn due to instability and less informative off-policy samples.
DecQN obtains higher competitive performance than the other baselines, which should benefit from the versatility of simple decoupled discrete control via leveraging a linear factorization over action dimensions.
In our experiments, PPO with unimodal policy achieves comparable or even better results than off-policy baselines, demonstrating its competitiveness in general complex applications. 
Our results show that PPO with unimodal policy is a promising method for complex control tasks, with performance comparable to state-of-the-art off-policy methods.




 §.§ Comparison with Discrete Policy

In a similar vein to our work, the extension of discretizing continuous action space is applied to solve high-dimensional control tasks.
We evaluate the unimodal policy alongside the discrete and ordinal policy architectures on tasks with K=11.
As shown in Fig. <ref>, the unimodal policy still significantly outperforms the Gibbs policy.
When compared to the ordinal policy, the unimodal policy learns a superior policy and achieves better performance due to the unimodal probability distribution.
Under a unimodal constraint, it is guaranteed that the two atomic actions on either side of the majority-selected action receive the next greatest amount of probability mass, which can result in a near-optimal policy during the early stages.
By this construction, the unimodal policy has fewer parameters than the discrete or ordinal policy when they share the same encoding architecture.
The results are predicted by our analysis showing that unimodal probability distributions of policy can still obtain high performance and better policies.
Although the Gibbs or ordinal policy can achieve enhanced performance, it is generally noisier without effectively utilizing the ordering information of discrete actions, relying more on the probability distribution of network output. 
By incorporating Poisson probability distribution constraints, the unimodal policy can more efficiently capture the continuity in the underlying continuous action space for generalization across discrete actions, especially in the high-dimension complex action space tasks.



 §.§ Stability and HyperParameter Analysis

PG methods, being generally defined over parametric distribution functions, are susceptible to converging to sub-optimal extrema, where policy entropy is usually utilized to enhance exploration by preventing premature convergence to suboptimal policies. 
We assess the sensitivity of policy entropy and performance with varying numbers of atomic actions K per dimension, presenting the results for PPO in Fig. <ref>.
As shown in Fig. <ref>, our proposed unimodal policy performs better with a larger number of atomic actions K during the learning process.
With a larger K, we can execute more precise actions through fine discretization, ultimately reaching high-quality final policies. 
From the perspective of policy entropy, finer discretization enables more efficient exploration with larger entropy.
By examining both the performance and policy entropy in Fig. <ref>, we observe that a higher absolute value of policy entropy yields better performance with more refined discretization, consistent with findings from <cit.>. 
Additionally, we also note that a moderate action bin K (e.g., 9 ≤ K ≤ 15) is sufficient for striking an appropriate balance between performance improvement and computational cost. 
The more minor variance of policy gradients tends to be more stable with finer discretization (larger K), potentially benefiting from the limited network output units inducing lower variance in the unimodal distributions.

Further, we study the performance of unimodal policy with varying temperature τ with 1.8, 2.4, and learnable bias, where the learnable τ is clipped to [1.5, 3] as a bias.
As shown in Fig. <ref>, our method can achieve robust performance under different temperatures τ.
Note that a lower temperature τ could reduce exploration and encourage more deterministic behavior, which may lead to suboptimal policies without exposure to diverse environments.
The learned τ is more stable under different fixed temperatures τ, though it does not show any significant gain over simply learning it as a bias.
One advantage of this technique is that the network can estimate its uncertainty on a per-example basis. 
Compared to their fixed temperature, it can better trade-off exploration versus exploitation under parameterized conditions.





§ CONCLUSION AND LIMITATIONS

In this paper, we introduce a straightforward technique to enforce a unimodal ordinal probabilistic policy distribution using Poisson distributions. 
Given the inherent ordering between discretized actions in continuous control tasks, this property is crucial to consider in discrete stochastic policy gradients. 
Despite its simplicity, our approach significantly enhances the performance and stability of baseline algorithms, particularly in high-dimensional tasks with complex dynamics. 
Furthermore, the unimodal ordinal policy allows probability distributions to behave more sensibly in benchmark continuous control tasks.
For limitations, since we mainly aim to refine off-the-shelf on-policy RL algorithms, it is not sufficient to conclude that our method can transcend off-policy algorithms.
In a practical view, our unimodal policy is more suitable for improving policy for high-dimension control tasks in bounded time, while the current form does not provide an optimality guarantee.
For future work, we will focus on designing more attentive unimodal distributions and extend them to multi-agent reinforcement learning.

IEEEtranN
























































 Paper Title: DKL-KAN: Scalable Deep Kernel Learning using Kolmogorov-Arnold Networks.


Entropy production due to spacetime fluctuations
    Lucas C. Célerihttps://orcid.org/0000-0001-5120-8176
    < g r a p h i c s >

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




§ ABSTRACT 

The need for scalable and expressive models in machine learning is paramount, particularly in applications requiring both structural depth and flexibility. Traditional deep learning methods, such as multilayer perceptrons (MLP), offer depth but lack ability to integrate structural characteristics of deep learning architectures with non-parametric flexibility of kernel methods. To address this, deep kernel learning (DKL) was introduced, where inputs to a base kernel are transformed using a deep learning architecture. These kernels can replace standard kernels, allowing both expressive power and scalability. The advent of Kolmogorov-Arnold Networks (KAN) has generated considerable attention and discussion among researchers in scientific domain. In this paper, we introduce a scalable deep kernel using KAN (DKL-KAN) as an effective alternative to DKL using MLP (DKL-MLP). Our approach involves simultaneously optimizing these kernel attributes using marginal likelihood within a Gaussian process framework. We analyze two variants of DKL-KAN for a fair comparison with DKL-MLP: one with same number of neurons and layers as DKL-MLP, and another with approximately same number of trainable parameters. To handle large datasets, we use kernel interpolation for scalable structured Gaussian processes (KISS-GP) for low-dimensional inputs and KISS-GP with product kernels for high-dimensional inputs. The efficacy of DKL-KAN is evaluated in terms of computational training time and test prediction accuracy across a wide range of applications. Additionally, the effectiveness of DKL-KAN is also examined in modeling discontinuities and accurately estimating prediction uncertainty. The results indicate that DKL-KAN outperforms DKL-MLP on datasets with a low number of observations. Conversely, DKL-MLP exhibits better scalability and higher test prediction accuracy on datasets with large number of observations. 







§ INTRODUCTION


 
During the 1990s, before the emergence of deep learning, Gaussian Processes (GP) (<cit.>)  became popular as alternatives to neural networks (NN) because of their ability to quantify uncertainty and their often superior performance. 
During this period, benchmarking datasets primarily consisted of tabular data, leading to the perception that representation learning could be entirely bypassed, focusing solely on the relationship between vectors in the input space. However, in the early 2010s, the community's focus shifted towards more intricate domains such as sequences (<cit.>), images (<cit.>), and speech (<cit.>), driven by a series of works leveraging the inductive biases of domain-specific NNs. This shift posed a challenge for GPs, primarily due to the difficulty in tailoring kernels that express useful inductive biases for these diverse domains. In response to these challenges, researchers proposed methods for learning GP kernels directly from data (<cit.>), aiming to improve the adaptability of GPs to diverse datasets by automatically capturing relevant features.

GP (<cit.>) are favored in Bayesian modeling due to their transparent  interpretability and reliable uncertainty quantification. Typically, these models involve a few kernel hyperparameters optimized with respect to the marginal likelihood.
Nonetheless, for many common choices, the kernel remains fixed, restricting the capability of GP to learn representations from the data that could improve predictions, thereby functioning primarily as smoothing mechanisms.. 
This restricts the applicability of GP to data that is high-dimensional and highly structured.
Conversely, deep neural networks (DNN) (<cit.>) are renowned for learning strong representations which are utilized for predicting unseen test inputs. Despite their good performance, deterministic NN tend to produce overconfident predictions (<cit.>) and lack trustworthy uncertainty estimates. The Bayesian approach to NN aims to mitigate these challenges; however, despite recent progress in variational inference and sampling techniques, inference in Bayesian NN remains difficult due to complicated posteriors and the huge number of parameters. Additionally, Bayesian NNs generally require multiple forward passes to obtain several samples of the predictive posterior to average over. Consequently, it is logical to combine the uncertainty representation advantages of GPs with the representation-learning advantages of NNs to leverage the strengths of both. This gave rise to deep kernel learning (DKL) (<cit.>) which uses an NN to extract meaningful representations that are subsequently fed into a standard GP, allowing for end-to-end learning of all model parameters. Both kernel hyperparameters and network parameters can be jointly trained by using varational inference or maximizing the marginal likelihood. In <cit.>, DKL has shown to outperform conventional shallow kernel learning methods, such as the radial basis function (RBF) kernel and standard NNs on a wide range of applications.

Multilayer perceptrons (MLP) (<cit.>), a cornerstone of modern deep learning, consist of multiple layers of nodes, including an input layer, one or more hidden layers, and an output layer. These networks are capable of learning complex non-linear relationships within data, as evidenced by the universal approximation theorem (<cit.>). Despite their success, MLP are often criticized for their lack of interpretability <cit.>, susceptibility to overfitting, scalablity, and issues related to vanishing or exploding gradients during training.

Kolmogorov-Arnold Networks (KAN) (<cit.>) have garnered significant attention lately as a novel neural network architecture aimed at addressing some of the inherent limitations of MLP. The foundation of KAN lies in the Kolmogorov-Arnold representation theorem, which asserts that any multivariate continuous function can be expressed as a finite sum of univariate functions and addition operations. This theorem, initially proven by Kolmogorov in 1957 to solve Hilbert's 13th problem (<cit.>), laid the groundwork for various neural network designs that leverage its decomposition principles.
Unlike MLP, which rely on linear weights and fixed activation functions, KAN use learnable activation functions, typically implemented as splines or B-splines. This design choice allows KANs to adaptively shape their activation functions to better fit the underlying data, thereby improving both accuracy and interpretability. 

The main contributions of this paper are as follows. First, we introduce a scalable deep kernel using KAN (DKL-KAN) as an alternative to the traditional DKL using MLP (DKL-MLP). 
Additionally, we present a comparative analysis of two DKL-KAN configurations: one matching the neuron and layer count of DKL-MLP, and another with a similar number of trainable parameters, ensuring a rigorous evaluation against the baseline. Our method is tested across various benchmark UCI regression datasets (<cit.>), focusing on both computational training time and prediction accuracy. Additionally, we also investigate the efficacy of DKL-KAN in modeling discontinuities and accurately calculating uncertainty bounds. To address the challenges of large datasets, we use Kernel Interpolation for Scalable Structured Gaussian Processes (KISS-GP) (<cit.>) for low input dimensions and KISS-GP with scalable kernel interpolation for product kernels (SKIP) (<cit.>) for high-dimensional inputs. 
The remainder of the paper is organized as follows. Section <ref> discusses the related works. Section <ref> delves into the brief explanation of KAN. This is followed by explanation of GP in Section <ref>. The findings of this study are detailed in Section <ref>, and finally, the conclusions are succinctly presented in Section <ref>




§ RELATED WORKS





 §.§ Gaussian Process



GPs, traditional nonparametric modeling tools and NNs, known for their flexibility in fitting complex data patterns, have an intriguing equivalence (<cit.>). As derived by <cit.>, in the limit of infinite width, a one-layer NN corresponds to a GP. This equivalence implies that for infinitely wide networks, an i.i.d. prior over weights and biases can be replaced with a corresponding GP prior over functions, allowing exact Bayesian inference for regression using NNs. 

The concept of DKL was first presented by <cit.>, building upon the framework of KISS-GP (<cit.>).
This approach combines deep NN with traditional covariance functions using local kernel interpolation to create complicated kernels. Although DKL excels in approximating intricate kernels for GP, it exhibits notable drawbacks. It does not support stochastic gradient training and is inadequate for classification tasks. In this regard, <cit.> introduced the stochastic variational deep kernel learning  approach, which allows the model in performing classification tasks, multi-task learning, and stochastic gradient optimization. <cit.> extended the general framework by introducing manifold GP, which involve applying a kernel function to a parameterized transformation of the input data. This transformation can be parameterized using NN. The study of intricate kernel learning has also been a focus within certain specialized fields. For instance, <cit.> used a recurrent neural network to acquire kernels characterized by recurrent patterns. In contrast, <cit.> explored using NN to learn a shared kernel function within the context of personalized federated learning. There has also been initiatives aimed at the regularization of models based on DKL. <cit.> introduced a latent-variable framework that integrates a stochastic encoding of inputs to allow regularized representation learning. Recently, <cit.> introduced an innovative methodology for training DKL. This method uses the uncertainty quantification from a GP with an infinite-width deep NN kernel to steer the optimization of DKL.

The flexibility of deep NN can however lead to overfitting in DKL models, as identified by <cit.>. They discovered that DKL often correlates all input data points in an effort to minimize the effect of the complexity penalty term in the marginal likelihood. This challenge was tackled by using full Bayesian inference over the network parameters using Markov Chain Monte Carlo (MCMC) techniques. However, MCMC is inefficient when dealing with high-dimensional posteriors. To address this, <cit.> has introduced amortized variational DKL, which leverages deep NN to generate embeddings for the kernel function and to learn variational distributions over inducing points, thereby reducing overfitting by locally smoothing predictions.



 §.§ Kolmogorov-Arnold Networks



Recent research has proposed several extensions and variants of KANs to improve their performance and reduce computational complexity. For instance, Fourier KAN (<cit.>), Wavelet KAN (<cit.>), RBF-KAN (<cit.>),  fractional KAN (<cit.>), BSRBF-KAN (<cit.>), rKAN (<cit.>) and sineKAN (<cit.>) replace B-splines with other function bases to achieve more accurate solutions. Another notable variant is the Chebyshev KAN (<cit.>), which use orthogonal polynomials as univariate functions. Each of these variants aims to address specific limitations or improve certain capabilities of the original KAN architecture. KANs have also shown promise in the context of solving partial differential equations and operator learning (<cit.>). Reference <cit.> combined KANs with physics-informed neural networks (PINNs) to address the 2D Poisson equation, achieving superior performance compared to traditional methods. Similarly, <cit.> introduced DeepOKAN, an RBF-based KAN operator network, for solving 2D orthotropic elasticity problems. These studies highlighted the potential of KANs to tackle complex scientific computing tasks more effectively than MLPs. The efficacy of KANs has also been explored in various other applications, including computer vision (<cit.>), medical image segmentation (<cit.>), and learning graphs (<cit.>. However, despite these advancements, the scalability and computational efficiency of KANs remain significant challenges (<cit.>).



§ KOLMOGOROV-ARNOLD NETWORKS



KANs represent an innovative class of neural networks that leverage the theoretical insights provided by the Kolmogorov-Arnold representation theorem. This theorem provides a theoretical basis for expressing any multivariate continuous function as a finite composition of univariate continuous functions and additions. By exploiting this theorem, KANs aim to overcome some of the limitations associated with traditional neural network architectures, particularly in terms of function approximation and interpretability. This theorem, also known as the superposition theorem, is a fundamental result in approximation theory. It asserts that any continuous multivariate function on a bounded domain can be represented as a composition of a finite number of univariate functions and linear operations. Assuming a smooth known function f: [0,1]^k→ℝ, continous univariate functions ϕ_i,j: [0,1] →ℝ, and continous functions ψ_j: ℝ→ℝ, the theorem states:

    f(x) = ∑_j=1^2k+1ψ_j(∑_i=1^kϕ_i,j(x_i)),

for some k. Irrespective of the chosen functions, an L layer deep KAN can be expressed in matrix form as follows:

    f̂(x) = (Ψ_L-1∘⋯∘Ψ_1 ∘Ψ_0)(x),

where each Ψ_l represents the function matrix of the l^th KAN layer. If n_l denotes the number of neurons in layer l, then the functional matrix Ψ_l that links neurons in layer l to those in layer l+1 is defined as follows:

    Ψ_l=([        ϕ_1,1(·)        ϕ_1,2(·)               ⋯     ϕ_1, n_l(·);        ϕ_2,1(·)        ϕ_2,2(·)               ⋯     ϕ_2, n_l(·);               ⋮               ⋮               ⋮               ⋮;   ϕ_n_l+1, 1(·)   ϕ_n_l+1, 2(·)               ⋯ ϕ_n_l+1, n_l(·) ]) ,

where ϕ_i,j signifies the activation function connecting neuron i in layer l to neuron j in layer l+1. Let 𝒟 denote the derivative operator and ϵ be a constant based on the behaviour of f(.). The convergence of this approximation is characterized by:

    max _|γ| ≤ msup _x ∈[0,1]^k𝒟^γ{f̂(x) - f(x)}<ϵ.

For more details, please refer to <cit.>.



The first KAN paper (<cit.>) introduced the residual activation function ϕ(x), which is defined as the sum of the base function b(x) = silu(x) and the spline function s(x), each multiplied by their respective weight matrices w_b and w_s. This spline function s(x) is represented as a linear combination of B-splines B_i(x):

    s(x) = ∑_i c_i B_i(x),

where c_i are trainable constants. The activation function is configured with w_s=1 and s(x) ≈ 0, and w_b is initialized using Xavier initialization. The residual activation function is then given by:

    ϕ(x) = w_b b(x) + w_s s(x).


Due to the high computational cost of KANs, a more efficient variant was introduced in <cit.>. This version follows the methodology of <cit.> but optimizes the computations using B-spline basis functions combined linearly, thereby reducing memory usage and simplifying the computational process. The authors also replaced the incompatible L1 regularization on input samples with L1 regularization on weights, incorporated learnable scales for activation functions, and changed the initialization of the base weight and spline scaler matrices to Kaiming uniform initialization. These modifications significantly improved performance on the MNIST dataset. In this paper, we utilize this efficient KAN for comparison with MLP.


§ GAUSSIAN PROCESSES



GP is a probabilistic non-parametric approach which defines a prior over functions and can be utilized for regression tasks. A GP is exhaustively described by its mean function, m(x), and a covariance function known as the kernel function, k(x, x'). Importantly, for any finite set of inputs, a GP will output a multivariate Gaussian distribution. In the context of GPR, the underlying assumption is that the function to be approximated is a sample from a GP. The mean function, m(x), and covariance function, k(x, x'), are represented by the equations:

    m(x)    = 𝔼[f(x)], 
    
    k(x, x')    = 𝔼[(f(x) - m(x))(f(x') - m(x'))].

In practice, it is common to use a constant mean function, i.e., m(x) = c, and focus on defining the covariance function. The covariance function, measures the similarity between data points. One common choice is the standard radial basis function (RBF) kernel. Assuming σ^2 as the signal variance, l_d representing the length scales for each dimension, and D indicating the number of dimensions, we have:

    k(x, x') = σ^2 exp(-∑_d=1^D (x_d - x'_d)^2/2l_d^2).

The length scales l_d control the relevance of different dimensions. When l_d is small, changes in the d-th dimension have a large effect on the covariance, making that dimension more relevant. 

Given a dataset 𝒟 of n input vectors, X={𝐱_1, …, 𝐱_n}, each having dimension D, these index an n × 1 vector of targets 𝐲=(y(𝐱_1), …, y(𝐱_n))^⊤. Assuming the presence of additive Gaussian noise, the relationship between the target y(𝐱) and a function f(𝐱) can be modeled as:

    y(𝐱) | f(𝐱) ∼𝒩(y(𝐱) ; f(𝐱), σ_y^2),

where σ_y^2 is the variance of the noise. Assuming all covariance matrices implicitly rely on the kernel hyperparameters θ, we have the predictive distribution of the GP at the n_* test points indexed by X_* to be expressed as:

    𝐟_* | X_*, X, 𝐲, θ, σ_y^2 ∼𝒩(𝔼[𝐟_*], cov(𝐟_*)), 
    𝔼[𝐟_*]=m(X_*)+k_θ(X_*,X)[k_θ(X,X)+σ_y^2 I]^-1𝐲
    cov(𝐟_*)=k_θ(X_*,X_*)-k_θ(X_*,X)[k_θ(X,X)+σ_y^2 I]^-1 k_θ(X,X_*),

where,

    𝐟_*=f_*(X)=[f_*(𝐱_1), …, f_*(𝐱_n)]^⊤.

Here, k_θ(X_*,X) is an n_* × n matrix of covariances between the GP evaluated at X_* and X. Additionally, m(X_*) designates the n_* × 1 mean vector, and k_θ(X,X) represents the n × n covariance matrix evaluated at training inputs X. 
The conditional probability density of the target variable 𝐲, given the parameters θ and input data X, can be expressed as:

    log p(𝐲|θ, X) ∝ -[𝐲^⊤(k_θ(X,X) + σ^2 I)^-1𝐲 + log|k_θ(X,X) + σ^2 I|].




 §.§ Deep kernel


A deep kernel integrates a neural network with a traditional kernel function, improving the model's ability to capture complex patterns in data. 
The neural network effectively compresses and distills the input space, mapping it to a lower-dimensional representation that is more informative for the GP model. 
The transformation induced by the neural network can allow the model to fit more complicated, non-smooth, and anisotropic functions, unlike the standard kernels like RBF, which model smooth and isotropic functions.
Also, by mapping a high-dimensional input space to a lower-dimensional one using a neural network, we can mitigate the curse of dimensionality. This makes the GPR more tractable and can improve its generalization by focusing on the most informative dimensions of the data.




Assuming ϕ(x;β) is a nonlinear mapping given by a neural network parameterized by β, and k(x_i, x_j' |α) representing the standard RBF kernel parameterized by α, we transform the inputs as:

    k(x_i,x_j'|α) → k(ϕ(x_i; β), ϕ(x_j'; β)|α, β).

We investigate the effectiveness of using KANs for DKL in this study. Fig. <ref> shows the schematic of DKL-KAN with three inputs and one output.
We simultaneously optimize all deep kernel hyperparameters, denoted as γ = (α, β), by maximizing the log marginal likelihood ℒ of the exact GP. To compute the derivatives of ℒ with respect to γ, we use the following chain rule:

    ∂ℒ/∂α=∂ℒ/∂  k_γ(X,X)∂ k_γ(X,X)/∂α,   ∂ℒ/∂β=∂ℒ/∂  k_γ(X,X)∂  k_γ(X,X)/∂ϕ(x_i; β)∂ϕ(x_i; β))/∂β .


The implicit derivative of the log marginal likelihood with respect to covariance matrix  k_γ(X,X) is expressed as

    ∂ℒ/∂ k_γ(X,X)=1/2(k_γ(X,X)^-1𝐲𝐲^⊤ k_γ(X,X)^-1 - k_γ(X,X)^-1),

where the noise covariance σ^2 I is included within the covariance matrix and is considered part of the base kernel hyperparameters α.

A significant computational challenge in DKL is resolving the linear system (k_γ(X, X) + σ^2 I)^-1𝐲. Additionally, within the realm of kernel learning, evaluating the logarithm of the determinant log|k_γ(X, X) + σ^2 I| as part of the marginal likelihood presents another difficulty. The conventional method uses the Cholesky decomposition of the n × n matrix k_γ(X, X), which requires 𝒪(n^3) operations and 𝒪(n^2) storage. After completing inference, the computation of the predictive mean entails a cost of 𝒪(n), while calculating the predictive variance incurs a cost of 𝒪(n^2) for each test point 𝐱_*.

 


 §.§ Scaling to large datasets for low input dimensions



Due to the cubic scaling of computations with the number of data points in the previously mentioned approach, we use structural kernel interpolation (SKI) using local cubic kernel interpolation to construct GP models for large datasets (<cit.>). Many strategies to scale GP models for large datasets belong to the class of inducing point methods (<cit.>). These methods replace the exact kernel with an approximation to accelerate computations. For m inducing points, the approximate kernels lead to computational complexity of 𝒪(m^2 n + m^3) and storage complexity of 𝒪(m n + m^2) for GP inference and learning (<cit.>), resulting in a predictive mean and variance cost of 𝒪(m) and 𝒪(m^2) per test case. 
Nevertheless, when using a limited number of inducing points, these techniques experience a considerable decline in accuracy (<cit.>). In contrast, methods that exploit structural properties, like Kronecker (<cit.>) and Toeplitz (<cit.>), leverage the covariance matrix to achieve scalable inference and learning without relying on approximations, thereby preserving both scalability and predictive precision.
However, the necessity of an input grid limits the applicability of these methods to most problems. SKI (<cit.>) integrates the benefits of both inducing point and structure-exploiting methods, explaining how the accuracy and efficiency of an inducing point method depend on the number of inducing points, kernel choice, and interpolation method. By selecting different interpolation strategies for SKI, novel inducing point methods can be devised. KISS-GP (<cit.>) uses SKI with local cubic and inverse distance weight interpolation to create a sparse approximation to the cross-covariance matrix between inducing points and original training points. Hence to ensure scalability, we replace all occurrences of k_γ(X, X) with the KISS-GP covariance matrix

    k_γ(X, X) ≈ W k_γ(U, U) W^⊤,

where W represents a sparse matrix of interpolation weights, containing only four non-zero elements per row for local cubic interpolation. The term k_γ(U, U) is a covariance matrix derived from the deep kernel, evaluated over m latent inducing points U=[𝐮_i]_i=1 … m. The computational complexity for KISS-GP training is 𝒪(n + h(m)), with  h(m) being nearly linear in m. The ability to have m ≈ n enables KISS-GP to achieve near-exact accuracy in its approximation, maintaining a non-parametric representation while ensuring linear scaling in n and 𝒪(1) time per test point prediction. This approach is especially efficient when the input dimensionality is less (generally less than four), as the cost of grid construction escalates exponentially with the amount of data.



 §.§ Scaling to large datasets for high input dimensions


To tackle the aforementioned problem, we use SKIP as introduced by <cit.>. SKIP is designed to handle large datasets with high-dimensional inputs efficiently. It comprises of two primary elements: SKI and the product kernel structure. The product kernel structure in SKIP mitigates the curse of dimensionality inherent in SKI. For data characterized by d input dimensions and denoting ⊗ as element-wise multiplication, the kernel matrix within the product kernel structure is represented as:

    k_γ(X,X) = k_γ(X,X)^(1)⊗…⊗ k_γ(X,X)^(d).

This formulation ensures that SKIP maintains linear time complexity, even when dealing with high dimensional inputs.





§ EXPERIMENTS





 §.§ Data normalization



All the training datasets were normalized using empirical cumulative distribution function technique prior to the initiation of the training process. This normalization technique transforms data into a uniform distribution, which can be particularly useful if the original distribution is unknown or if the data does not follow a Gaussian distribution which is a common assumption in many machine learning algorithms. It is also more robust to outliers as it ranks the datapoints rather than directly scaling their values. 



 §.§ Results



The GP models and deep kernel is programmed using GPyTorch library respectively with a PyTorch backend. We set the loss function as negative of exact marginal log likelihood and the optimizer used is Adam with tuned hyperparameters. We set the learning rate as 0.075 with an exponential decay rate of 0.997 for smooth convergence of the loss. The GP models are trained for a total of 2500 epochs with a patience of 1000 epochs. All training is performed on an NVIDIA A100-SXM4-40GB GPU (1 GPU with 512 GB RAM). The codes for this paper will be made available at https://github.com/shrenikvz/dkl-kanhttps://github.com/shrenikvz/dkl-kan upon publication.



  §.§.§ UCI regression datasets



We evaluate our models on a diverse set of benchmark UCI regression problems with varying sizes and characteristics. The base kernel is fixed as the RBF kernel. For DKL approaches, we reduce the dimensionality to 2 before feeding it to the base kernel. The neural network architecture for DKL-MLP consists of three hidden layers with 1000, 500, and 50 neurons, respectively. To fairly compare DKL-KAN with DKL-MLP, we train two models using KANs: DKL-KAN1 and DKL-KAN2. DKL-KAN1 has the same number of neurons and layers as DKL-MLP, while DKL-KAN2 has approximately the same number of trainable parameters as DKL-MLP. DKL-KAN2's architecture includes three hidden layers with 256, 128, and 64 neurons, respectively. We utilize RMSE as the metric to evaluate the test prediction accuracy.











Given that GP models could be trained in a finite time for datasets with fewer than 20,000 observations, we use SKI exclusively for datasets exceeding 20,000 observations. For these larger datasets, we use KISS-GP for DKL-MLP, DKL-KAN1, and DKL-KAN2, as the neural network maps to two output features, allowing for the use of SKI. Conversely, we use KISS-GP with product structure kernel (SKIP) for the standard GP, since the input features are not reduced before being fed to the base kernel.

As illustrated in Table <ref>, DKL-KAN1 shows superior performance on datasets with a low number of observations. Additionally, we observe that increasing the complexity of the KAN architecture leads to notable improvements in test prediction accuracy. However, for datasets with a large number of observations, DKL-MLP outperforms DKL-KAN, suggesting that the current iteration of KANs is less effective for large datasets.  However, we do see DKL-KAN1 outperform all other GP models for  dataset.

Table <ref> illustrates the RMSE performance of these GP models with an increased number of partitions (i.e., 20) for datasets with a large number of observations. Due to increased partitions (which means using lesser data to train on), we did not consider using SKI/SKIP for generating these results. It is evident that for these large datasets, increasing the number of partitions leads to better test prediction accuracy of DKL-KAN1 over DKL-MLP, further validating our earlier assertion that KANs in their current version are not scalable to large datasets.

There are also several instances where a standard GP shows relatively better performance than DKL. Table <ref> presents the total number of trainable parameters for these models, while Table <ref> provides the average computational training time (in seconds) for each model on specific datasets. Given that we used efficient KANs in this study, the training time of DKL-KANs is comparable to that of other GP models (Fig. <ref>).



  §.§.§ Step function



To evaluate the effectiveness of DKL-KAN in handling discontinuities and accurately estimating prediction uncertainty, we consider the step function depicted below

    y     = f(x) + w,    w ∼𝒩(0,0.01^2), 
    
    f(x)     = 0     if  x ≤ 0 
     1     if  x > 0

For the training process, 100 input points are sampled from a standard normal distribution, whereas the test set consists of 500 datapoints uniformly distributed between -5 and 5. Given that we are modeling a step function, a neural network with a single hidden layer of 6 neurons is used for both DKL-MLP and DKL-KAN, with the reduced latent dimension set to 2. Fig. <ref> illustrates the predictive mean and uncertainty bands when different GP variants are tested on the test set. We observe a relative decrease in epistemic uncertainty in regions where training datapoints are present for GP, DKL-MLP, and DKL-KAN. The standard GP, which assumes a smooth function structure, struggles to model the discontinuity effectively. Due to DKL's joint supervised learning of input transformation and regression, the neural network captures the discontinuity in the reduced 2D feature space for both MLP and KAN (Fig.<ref>), facilitating a smooth mapping from reduced feature space to output that the GP can handle (<cit.>). Additionally, we observe that DKL-MLP maintains low uncertainty across the entire function domain as explained in <cit.>. In contrast, DKL-KAN exhibits increasing epistemic uncertainty in regions where training data is not present. The behavior of DKL-MLP with respect to uncertainty bounds is not ideal, as, in an optimal scenario, epistemic uncertainty should rise in areas lacking training data. DKL-KAN, however, offers a balanced solution by accurately capturing discontinuities and appropriately estimating prediction uncertainties.
  





















§ CONCLUSIONS



In this study, we investigate the effectiveness of incorporating KANs into DKL. To ensure a fair comparison with DKL-MLP, we train two GP models using KANs: one with an identical number of neurons and layers as the MLP (DKL-KAN1) and another with approximately same number of trainable parameters (DKL-KAN2). Our findings indicate a relatively better performance of KANs when using a larger number of neurons for a fixed number of hidden layers. Moreover, DKL-KANs demonstrate superior performance compared to other GP models on datasets with a low number of observations. Conversely, DKL-MLP exhibit better scalability and higher test prediction accuracy on datasets with a large number of observations. This suggests that, in their current form, KANs are less effective for training on larger datasets. Additionally, DKL-KAN accurately captures discontinuities and prediction uncertainties, particularly in regions lacking training data. The ability of DKL-KAN to increase epistemic uncertainty where data is sparse ensures a more reliable model performance in practical applications.





§ ACKNOWLEDGEMENTS


This work is supported by RTX Technology Research Center.

unsrtnat





 Paper Title: MART: MultiscAle Relational Transformer Networks for Multi-agent Trajectory Prediction.





 



MART: Multiscale Relational Transformer Networks






Lee et al.




School of Integrated Technology (SIT), Gwangju Institute of Science and Technology
(GIST), Cheomdan-gwagiro 123, Buk-gu, Gwangju 61005, Republic of Korea {lsj2121, junseoklee, yeon_guk, smterry0928, kyoobinlee}@gist.ac.kr

MART: MultiscAle Relational Transformer Networks for Multi-agent Trajectory Prediction
    Seongju Lee10000-0002-1712-0282 Junseok Lee10000-0001-5212-2657 Yeonguk Yu10000-0003-2147-4718 Taeri Kim10000-0002-5139-7963 
Kyoobin Lee1Corresponding author.0000-0003-4299-4923
    August 5, 2024
======================================================================================================================================================================================




§ ABSTRACT 
    Multi-agent trajectory prediction is crucial to autonomous driving and understanding the surrounding environment. Learning-based approaches for multi-agent trajectory prediction, such as primarily relying on graph neural networks, graph transformers, and hypergraph neural networks, have demonstrated outstanding performance on real-world datasets in recent years. However, the hypergraph transformer-based method for trajectory prediction is yet to be explored. Therefore, we present a MultiscAle Relational Transformer (MART) network for multi-agent trajectory prediction. MART is a hypergraph transformer architecture to consider individual and group behaviors in transformer machinery. The core module of MART is the encoder, which comprises a Pair-wise Relational Transformer (PRT) and a Hyper Relational Transformer (HRT). The encoder extends the capabilities of a relational transformer by introducing HRT, which integrates hyperedge features into the transformer mechanism, promoting attention weights to focus on group-wise relations. In addition, we propose an Adaptive Group Estimator (AGE) designed to infer complex group relations in real-world environments. Extensive experiments on three real-world datasets (NBA, SDD, and ETH-UCY) demonstrate that our method achieves state-of-the-art performance, enhancing ADE/FDE by 3.9%/11.8% on the NBA dataset. Code is available at <https://github.com/gist-ailab/MART>.
    





§ INTRODUCTION


Multi-agent trajectory prediction, which forecasts the future trajectories of agents based on their past movements, holds significance in understanding human behavior and has applications in several fields, including robotics <cit.> and autonomous driving <cit.>. This task is complex and challenging for several reasons. For instance, some pedestrians move in formation, often as part of a group, introducing interdependencies and shared decision-making. These group dynamics influence individual paths, making predicting the trajectory of a single agent without considering the broader context challenging. Consequently, researchers have explored sophisticated methods to capture individual interactions and group behavior <cit.>.




Previous studies have focused on modeling pair-wise interactions in trajectory prediction using graph neural networks (GNNs) <cit.> and transformer networks <cit.>. Some studies have employed graph convolutional networks <cit.> to account for pair-wise relations in crowded scenes <cit.>. Other methods, such as those based on graph attention networks <cit.>, have been proposed in <cit.>. In addition, message-passing neural networks <cit.> have been employed in trajectory prediction tasks <cit.>. Recently, Xu <cit.> introduced a motion prediction model under Euclidean geometric transformations based on equivariant GNNs <cit.>. In addition, Yu <cit.> proposed a graph transformer network that considers pair-wise and temporal relations by incorporating an adjacency matrix into the attention mechanism.

To cooperate with individual and group behaviors, GroupNet <cit.>, a multiscale hypergraph message-passing neural network, was proposed for multi-agent trajectory prediction. As an extension of GroupNet <cit.>, DynGroupNet <cit.> captures time-varying interactions at the pair and group scales. However, transformer machinery, an effective relationship-learning mechanism, is not fully incorporated into their methods. Bae <cit.> proposed the GP-Graph, which considers intra- and inter-group relationships by partitioning the most likely behavioral agents in the same group. However, round-trip node-edge communication, crucial for relational reasoning <cit.>, is ignored in GP-Graph.


In this study, we propose a MultiscAle Relational Transformer (MART) network, a hypergraph transformer architecture for multi-agent trajectory prediction. To address social relationships in hypergraph transformer architecture, we introduce the MART encoder (MARTE), comprising a pair-wise relational transformer and a Hyper Relational Transformer (HRT) to enhance the capabilities of the relational transformer <cit.>. We introduce Hyper Relational Attention (HRA) into HRT, incorporating hyperedge features into the attention mechanism. As hyperedge features involve group information, HRT captures group behavior while promoting attention weights to focus on group-wise relations.


Since group relations are not explicitly defined in real-world scenarios, pioneering studies have introduced methods for reasoning group relations <cit.>. However, they have limitations. In <cit.>, the focus is on non-overlapping group relations that cannot handle complex relationships. For <cit.>, the number of agents in a group requires a manual definition, and as the number of group scales increases, more encoders are required, resulting in higher computational costs. To address the aforementioned limitations, we propose an Adaptive Group Estimator (AGE). Inspired by the Straight-Through Estimation (STE) technique <cit.> used in binarized neural networks <cit.>, we utilize adaptive thresholding to include highly correlated agents within the same group. Our approach estimates overlapping group relations, recognizing that agents can be correlated with more than one group. Moreover, unlike <cit.>, the manual definition of the number of agents in a group is unnecessary in our method.


To validate the effectiveness of MART, we conduct extensive experiments on three real-world datasets (NBA, SDD <cit.>, and ETH-UCY <cit.>) and demonstrate that our method achieves state-of-the-art (SOTA) performance. Notably, on the NBA dataset, which is the most challenging and complex, MART shows a subjective performance improvement by reducing the ADE by 3.9% and FDE by 11.8% compared to the previous SOTA method, EqMotion <cit.>. The ablation study on the interaction encoder demonstrates that hypergraph transformer architecture is beneficial for multi-agent trajectory prediction. The proposed MARTE, based on a hypergraph transformer, outperforms other SOTA GNN-based and transformer-based encoders, as illustrated in Figure <ref>. Additionally, MART exhibits computational efficiency compared to recent SOTA models. The main contributions are as follows:



    
  * We introduce MART, a Multiscale RelAtional Transformer network for multi-agent trajectory prediction. MART captures relations at individual and group scales through a hypergraph transformer mechanism via MARTE. The proposed MARTE promotes the attention weights to emphasize interactions between agents, particularly focusing on group-wise relations.
    
  * We propose the AGE module to infer group relations. The AGE module outperforms existing group reasoning methods in both performance and efficiency by considering overlapping group relations and dealing with complex interactions through adaptive thresholding.
    
  * We experimentally demonstrate that MART achieves SOTA performance on three real-world datasets (NBA, SDD, and ETH-UCY). Moreover, MART is computationally efficient compared to the recent SOTA models.




§ RELATED WORK


Trajectory Prediction.
Traditional trajectory prediction methods are based on heuristics, energy potentials, and simple machine learning-based methods <cit.>. Berndt <cit.> proposed using the global vehicle position with a hidden Markov model to predict object trajectory in a road traffic environment. Subsequently, recurrent neural network (RNN)-based methods <cit.> have been proposed to capture the long-term behavior of agents. For instance, Vemula . <cit.> used an RNN-based prediction model with attention that captured the relative importance of each person to predict agents' trajectories in the crowd. Generative adversarial network (GAN)-based methods <cit.> and variational autoencoder-based methods <cit.> have also been studied to generate future trajectory distribution. Similarly, <cit.> proposed trajectory-prediction methods based on a diffusion model. In addition, transformer structures <cit.> have been used to capture the importance of the agents' spatial and temporal correlations. Our method is a hypergraph transformer network that captures individual and group behaviors for trajectory prediction.


Transformers with Graph.
Recently, transformer-based structures have been proposed for representing graphs <cit.>. Specifically, Chen <cit.> addressed concerns regarding transformer architecture, in which the node representation generated by the transformer fails to capture the structural similarity between them, and solved it by incorporating structural information into the self-attention module. In addition, relational transformer <cit.>, which uses an attention module to capture edge information effectively, has been demonstrated to perform diverse graph-structure tasks. For trajectory prediction, Li <cit.> proposed a graph-based spatial transformer to capture the spatial interactions between pedestrians. Herein, we introduce MART to effectively capture individual and group behaviors via transformer mechanisms.



§ BACKGROUND AND PROBLEM FORMULATION



 §.§ Trajectory Prediction

Multi-agent trajectory prediction predicts the future trajectories of multiple agents based on their past trajectories. Assume that N agents exist in a multi-agent system. Specifically, let 𝐗∈ℝ^N × T_p × 2, 𝐘∈ℝ^N × T_f × 2, and 𝐘∈ℝ^N × T_f × 2 represent past trajectories, ground-truth future trajectories, and predicted future trajectories in 2D spatial coordinates, respectively. The objective is to minimize the error between the predicted and ground-truth future trajectories. We design a trajectory prediction model denoted by 𝐘 = ℳ(𝐗). Here, ℳ produces predicted future trajectories that closely align with the ground-truth future trajectories.




 §.§ Relational Transformer

The relational transformer (RT) <cit.>, an extension of the transformer <cit.>, incorporates edge vectors as fundamental components to achieve SOTA performance in various tasks involving graph-structured data. Mathematically, for a given input-directed graph, we denote an attributed graph as 𝒢=(𝒩, ℰ). Here, 𝒩 represents an unordered set of node vectors 𝐧_i∈ℝ^d_n, where 𝐧_i is the i-th node vector and d_n is the node dimension. ℰ represents the set of edge vectors 𝐞_ij∈ℝ^d_e, where 𝐞_ij represents the directed edge pointing from j-th to i-th node and d_e is the edge dimension.

Similar to message-passing neural networks <cit.>, RT comprises two processes: node and edge updates. In the node update, RT introduces relational attention (RA), which incorporates edge vectors into the scaled dot-product attention. Formally, the query, key, and value vectors in RA are defined as follows:

    [ 𝐪_ij = (𝐧_i𝐖_n^Q + 𝐞_ij𝐖_e^Q) ∈ℝ^d_n,; 𝐤_ij = (𝐧_j𝐖_n^K + 𝐞_ij𝐖_e^K) ∈ℝ^d_n,; 𝐯_ij = (𝐧_j𝐖_n^V + 𝐞_ij𝐖_e^V) ∈ℝ^d_n, ]

where 𝐖_n∈ℝ^d_n × d_n and 𝐖_e∈ℝ^d_e × d_n are the weight matrices used to project the node and edge vectors onto the query, key, and value vectors, respectively. The remainder of the node update is the same as the transformer mechanism. Formally, the updated node vector after (l+1)-th layer 𝐧_i^(l+1)∈ℝ^d_n is defined as follows:

    𝐧_i^(l+1) = ℱ_n(softmax_j( 𝐪_ij^(l)𝐤^(l)_ij^⊤/√(d_n)) 𝐯_ij^(l), 𝐧_i^(l))

where ℱ_n(·) represents a transformer-like node update function of the form [Add & Norm]–[FeedForward]–[Add & Norm]. In this study, we omitted the equation for multi-head attention for simplicity.

For the edge update, the edge vector after (l+1)-th layer 𝐞_ij^(l+1)∈ℝ^d_e is produced by the edge update function ℱ_e(·) as follows:

    𝐞_ij^(l+1)=ℱ_e(𝐦_ij^(l), 𝐞_ij^(l)).

Here, 𝐦^(l)_ij∈ℝ^d_e represents the message vector that aggregates the information from adjacent updated nodes and their edges. Function ℱ_e(·) denotes the transformer-like edge-update function, which has an identical operation as ℱ_n(·). The message vector 𝐦^(l)_ij is derived as follows:

    𝐦^(l)_ij = ReLU([𝐞_ij^(l);𝐞_ji^(l);𝐧_i^(l+1);𝐧_j^(l+1)]𝐖_m),

where [· ; ·] denotes the concatenation operation and 𝐖_m ∈ℝ^(2d_e + 2d_n)× d_h is a weight matrix for the message function.






§ METHOD

As illustrated in Figure <ref>, MART comprises four components: feature initialization, AGE, MARTE, and decoder. The subsequent sections offer a detailed description of the core components of MART—feature initialization, AGE module, and MARTE—as well as the overall process.



 §.§ Feature Initialization
 
We first transform the past input trajectories into agent embeddings to represent the initial node and edge features. Because edge relationships are often not explicitly defined in real-world environments, we represent edge features by aggregating the features of the adjacent nodes. In this paper, we consider pair-wise edge and hyperedge features.



Node Feature Initialization.
The past trajectory of i-th agent 𝐗_i∈ℝ^T_p × d_i is embedded via 𝐧^(0)_i = ℱ_NI(𝐗_i), where 𝐧_i^(0)∈ℝ^d_n is the initial node feature of the i-th agent and ℱ_NI(·) represents two layers of multilayer perceptron (MLP). d_n denotes the node dimensions.



Edge Feature Initialization.
To consider pair- and group-wise relations, we initialize pair-wise edge and hyperedge features, respectively. To form pair-wise edge features, two adjacent nodes are concatenated and passed through the pair-wise edge feature initializer ℱ_PI(·): 𝐞^(0)_ij = ℱ_PI([𝐧^(0)_i; 𝐧^(0)_j]), where 𝐞^(0)_ij is the initial pair-wise edge feature from j-th to i-th agent and ℱ_PI(·) is MLP.

The initial hyperedge feature is formed by averaging the initial node features of the corresponding hyperedge. The aggregated node features are fed to the hyperedge initializer ℱ_HI(·): 𝐡^(0)_i = ℱ_HI(1/|𝒩_i|∑_n_j∈𝒩_i𝐧^(0)_j), where ℱ_HI is implemented by MLP, 𝒩_i is a set of node features of the i-th hyperedge, and |𝒩_i| is its cardinality. In the following section, we propose a group reasoning method to address the lack of predefined hypergraph topology in real-world environments.




 §.§ Adaptive Group Estimator

Real-world environments often conceal group relationships, which are crucial for understanding group behavior. Thus, several works have proposed methods for reasoning about group relations <cit.>; however, these methods have limitations, as mentioned in the Introduction. To overcome the aforementioned limitations, we propose the Adaptive Group Estimator (AGE) to effectively infer complex group relationships through adaptive thresholding, thereby grouping closely related agents.

Mathematically, given a hypergraph 𝒢 = (𝒩, ℋ), a group relation can be expressed as a group incident matrix 𝐆∈ℝ^|𝒩| × |ℋ|, where 𝐆_ij∈{0, 1}. Here, ℋ is an unordered set of hyperedge features, and |ℋ| is its cardinality. When 𝐆_ij=1, the i-th agent belongs to the j-th hyperedge; however, when 𝐆_ij=0, it does not. In this study, we assume that the number of hyperedges is set to |𝒩| to include agents highly correlated with the j-th agent (ego agent) in the same group. Thus, the (i,j)-th element of the group incidence matrix 𝐆 is defined as follows:

    𝐆_ij = 𝒰(𝐀_ij - Θ),

where 𝐀_ij = 𝐧_i^⊤𝐧_j/(||𝐧_i||_2||𝐧_j||_2)∈ℝ^N × N is the (i,j)-th element of the affinity matrix, Θ∈ (-1.0, 1.0) is a learnable threshold, and 𝒰(·) is a unit step function defined as follows:

    𝒰(x) = 
        1,    if x ≥ 0.
    
        0,    otherwise.


As the unit step function is not differentiable, we use the STE trick <cit.> inspired by binarized neural networks <cit.>. We compute the group incidence matrix 𝐆 using Eq. <ref> for the forward path. In the backward path, we estimate the gradient of 𝒰(·) to enable backpropagation. The estimated gradient of 𝒰(·) is defined as Eq. <ref>. Note that this estimation is derived based on Eq. <ref>, which approximates 𝒰(·). By introducing Θ into the AGE module, it selectively groups more correlated agents into the same group. The supplementary material describes the details of the function selection.














 §.§ Multiscale Relational Transformer Encoder

We introduce the MultiscAle Relational Transformer Encoder (MARTE) into our network to capture social interactions at multiple scales via hypergraph transformer architecture. The detailed architecture of MARTE is shown in Figure <ref>. PRT and HRT address individual and group behaviors, respectively.



Pair-wise Relational Transformer.
PRT is based on RT <cit.>, an effective relational learner, as described in Section <ref>. However, RT deals only with pair-wise relationships. Therefore, in the following section, we introduce HRT as an extension of RT to accommodate hypergraph structures, ensuring synergy between architectures. 




Hyper Relational Transformer.
We propose HRT to capture group behavior via a hypergraph transformer mechanism. To achieve this, we introduce Hyper Relational Attention (HRA) into HRT, extending the relational attention mechanism to accommodate hypergraph structures. Additionally, we introduce a message function to update the hyperedges in accordance with these structures.


(Node update for HRT) To extend the mechanism of relational attention to a hypergraph structure, we integrate hyperedge information into the query, key, and value by aggregating the hyperedge features to which the corresponding agent belongs. The average operation is used for aggregation for stable training. The aggregated hyperedge features in query, key, and value work as the relative positional information that includes group information, formulated as follows:


    [ 𝐪_i = ( 𝐧_i𝐖_n^Q + 1/|ℋ_i|∑_h_k∈ℋ_i𝐡_k𝐖_h^Q) ∈ℝ^d_n,; 𝐤_j = ( 𝐧_j𝐖_n^K + 1/|ℋ_j|∑_h_k∈ℋ_j𝐡_k𝐖_h^K) ∈ℝ^d_n,; 𝐯_j = ( 𝐧_j𝐖_n^V + 1/|ℋ_j|∑_h_k∈ℋ_j𝐡_k𝐖_h^V) ∈ℝ^d_n, ]


where 𝐖_h∈ℝ^d_e × d_n is the weight matrix used to project the hyperedge features onto the query, key, and value vectors, ℋ_i is a set of hyperedges to which the i-th agent belongs, and |ℋ_i| is its cardinality. These relations of belonging are derived from 𝐆. Consequently, the updated node feature 𝐧_i^(l+1) is formulated as follows:

    𝐧_i^(l+1) = ℱ_n(softmax_j( 𝐪_i^(l)𝐤^(l)_j^⊤/√(d_n)) 𝐯_j^(l), 𝐧_i^(l)).


(Hyperedge update for HRT) The i-th hyperedge feature after the (l+1)-th HRT layer, denoted as 𝐡_i^(l+1)∈ℝ^d_e, is defined as follows:

    𝐡_i^(l+1) = ℱ_h(𝐦_i^(l), 𝐡_i^(l))

where ℱ_h(·) denotes the hyperedge update function based on feedforward network, and 𝐦_i^(l) is a message vector. For the message vector 𝐦_i^(l), we formulate a message function to accommodate the hyperedge structure. More specifically, this function is formulated using the hyperedge feature and the mean of the updated node features belonging to the corresponding hyperedge. Mathematically, the message vector 𝐦_i^(l)∈ℝ^d_e is defined as follows:

    𝐦^(l)_i = ReLU([𝐡_i^(l);1/|𝒩_i|∑_n_j∈𝒩_i𝐧_j^(l+1)]𝐖_m),

where 𝐖_m∈ℝ^(d_e + d_n) × d_h denotes the weight matrix to project post node-update features onto the message vector.




 §.§ Overall Process of MART

We detail the trajectory prediction process of MART in the following sequence: feature initialization, encoding, decoding, and loss functions. The first three processes represent the forward path of MART in matrix formula, and the loss function serves as the objective function for model training.




Feature Initialization.
As described in Section <ref>, our model first initializes the node, pair-wise edge, and hyperedge features for N agents as:

    [   𝐍^(0) = {𝐧^(0)_1, 𝐧^(0)_2, ..., 𝐧^(0)_N }∈ℝ^N × d_n,; 𝐄^(0)={𝐞^(0)_11, 𝐞^(0)_12, ..., 𝐞^(0)_NN}∈ℝ^N^2 × d_e,;     𝐇^(0)={𝐡^(0)_1, 𝐡^(0)_2, ..., 𝐡^(0)_N }∈ℝ^N × d_e. ]

The pair-wise edge features 𝐄 have N^2 × d_e dimensions because we consider a fully connected directed graph with self-loops.




Encoding Process.
MARTE encodes pair- and group-wise node features from the initial features. Formally, the features produced by MARTE are described as follows:

    [     𝐍^(L)_pair = 𝒫ℛ𝒯(𝐍^(0), 𝐄^(0)),; 𝐍^(L)_group = ℋℛ𝒯(𝐍^(0), 𝐇^(0), 𝐆), ]

where 𝐍^(L)_pair∈ℝ^N × d_n represents the node features at the pair-wise scale from the L-th PRT layer and 𝐍^(L)_group∈ℝ^N × d_n denotes the node features at the group scale from the L-th HRT layer. We use 𝒫ℛ𝒯(·) to refer to the PRT operation and ℋℛ𝒯(·) for the HRT operation.




Decoding Process.
To predict future trajectories, we utilize multiple prediction heads as the decoder to obtain k predictions, as in <cit.>. Three types of features are concatenated: 𝐍^(0), 𝐍^(L)_pair, and 𝐍^(L)_group, then fed into the multi-head decoder ℱ_D. Formally, the trajectories obtained from the k-th prediction head can be described as 𝐘̂_k = ℱ_D, k([𝐍^(0);𝐍^(L)_pair;𝐍^(L)_group]), where ℱ_D, k(·) denotes a decoder with three layers of MLP.




Loss Function.
We trained our model using the variety loss proposed in social-GAN <cit.> to optimize the best prediction. With the given k-th predicted trajectories resulting from k-th prediction head 𝐘̂_k, the minimum of the L2-norm among k predictions is utilized for backpropagation. Mathematically, the loss function is formulated as ℒ=1/NT_f∑_n=1^N∑_t=1^T_fmin _kY^(n,t) - Y_k^(n,t)_2, where (n,t) denotes the indices for the time step t for the n-th agent.




§ EXPERIMENTS


Datasets.
For the evaluation of MART, we utilize three datasets for multi-agent trajectory prediction: NBA SportVU (NBA), Stanford Drone (SDD) <cit.>, and ETH-UCY <cit.>. Further details about the datasets and evaluation scheme are described in the supplementary material.


Implementation Details.
The node and edge dimensions, denoted by d_n and d_e, respectively, are set to 64. Eight heads are used for multi-head attention. The number of layers in MARTE, denoted by L, is 4. For all datasets, our network is trained using the Adam optimizer <cit.>. A learning rate decay of 0.5 is applied every 100 epochs. The network is trained using an NVIDIA GeForce RTX 3090. All our implementations are on PyTorch 2.1.0 <cit.>. Further details about the implementation can be found in the supplementary material.


Evaluation Metrics.
To assess the trajectory prediction performance, we employed two metrics, minADE_k and minFDE_k, as in <cit.>. ADE and FDE represent the Euclidean distances between the predicted and actual trajectories, considering the entire path and endpoints, respectively. Consequently, minADE_k and minFDE_k are the smallest values among k computed ADEs and FDEs from the k predicted trajectories.





§ RESULTS AND DISCUSSION




 §.§ Comparison with State-Of-The-Art (SOTA) Methods






The performance of the proposed model is compared with the SOTA methods on the NBA, SDD, and ETH-UCY datasets using two metrics: minADE_20 and minFDE_20, which are described in Section <ref>.

NBA Dataset. Table <ref>-(a) presents the performance comparison to eight SOTA methods on the NBA dataset. Compared to the previous best method EqMotion <cit.>, our model reduces minADE_20 and minFDE_20 at 4.0 s by 3.9% and 11.8%, respectively.

SDD Dataset. We compared our method with eight SOTA methods on the SDD dataset, as listed in Table <ref>-(b). Our method enhances the minADE_20 by 7.8% compared to EigenTrajectory (ET) <cit.> and shows comparable performance in minFDE_20 compared to LED <cit.>.

ETH-UCY Dataset. Performance comparison with eight SOTA methods on the ETH-UCY dataset is listed in Table <ref>-(c). Our model yields results comparable to those of a previous SOTA method, LED <cit.>.

Computational Efficiency of MART.  The computational efficiency of MART is evaluated in comparison with recent SOTA methods in terms of model parameters and multiple-add cumulation (MAC) for inference. These metrics are computed with 10 agents in the ETH-UCY setting. As shown in Table <ref>, our method demonstrates computational efficiency compared to previous SOTA methods such as LED <cit.> and EqMotion <cit.>. Despite STAR <cit.> having only 1.0M model parameters, it has a larger MAC count than ours due to its iterative architecture.

Discussion.
MART demonstrates SOTA or comparable performance in trajectory prediction compared to previous methods. Our approach is more computationally efficient, with 7.3x fewer parameters and 347.3x fewer MACs than the previous SOTA method, LED <cit.>. These results indicate that MART is highly competitive for accurate trajectory prediction in real-time applications.




 §.§ Ablation Study





Effect of MARTE.
Ablation studies are conducted on different interaction encoders to demonstrate the effectiveness of MARTE, which is a core component of MART. Five interaction encoders are evaluated in our prediction system for a fair comparison. The model specifications, such as the number of layers and heads for transformer <cit.> and relational transformer <cit.>, are the same as those for MARTE. As illustrated at the top of Table <ref>, MARTE outperforms both transformer-based encoders (transformer and relational transformer) and GNN-based encoders (NMMP, GroupNet, and EqMotion). On the other hand, the results demonstrate that transformer-based encoders perform better than GNN-based encoders. Although the relational transformer exhibits commendable performance, it overlooks the importance of group-wise interaction, which is essential for accurate trajectory prediction. Thus, the findings underscore that MARTE enhances performance via a hypergraph transformer, marking a novel approach in trajectory prediction.


Effect of AGE.
To validate our AGE module, it is compared to previously proposed group estimation methods <cit.> by integrating them into MART. As shown at the bottom of Table <ref>, the AGE module performs better than the other methods. The group estimation of GP-Graph <cit.> leads to lower performance in complex scenes like the NBA datasets due to non-overlapping group estimation. GroupNet <cit.> performs similarly to our method in ADE but requires manual group size setting. Moreover, as the number of group scales increases, more encoders are required, resulting in higher computational costs.





 §.§ Qualitative Results
 





Trajectory Prediction.
In Figure <ref>, the predicted trajectories from MART (ours) are compared with two recent SOTA methods, LED <cit.> and EqMotion <cit.>, along the ground truth (GT) trajectories on the NBA dataset. MART outperforms the previous methods, resulting in accurate trajectory predictions. Further visualization can be found in the supplementary material.

Group Estimation.
Figure <ref> displays examples of group estimation results on the NBA dataset, and Figure <ref>-(c) shows the matrix for unique groups. Each group reflects a specific strategy. For example, in Group #1, defenders employ man-to-man marking on offensive players during zone defense (See supplementary material for details).

Attention Weights Visualization.
As illustrated in Figure <ref>-(a) and -(b), the attention weights of PRT and HRT focus on the pair-wise and group-wise relations, respectively. Notably, for HRT, the attention weights in HRT show the same accentuation as the group incidence matrix displayed on the right side of Figure <ref>-(b). Consequently, the visualization demonstrates that the proposed MARTE promotes the attention weights to emphasize interactions between agents, particularly focusing on group-wise relations.











§ CONCLUSION


Conclusion.
This study proposes MART, a multiscale relational transformer network for multi-agent trajectory prediction. MART considers interactions at both individual and group scales via the hypergraph transformer mechanism. Within the MART framework, we introduce the MARTE and AGE modules. MARTE extends the capabilities of the relational transformer and promotes attention weights to emphasize interactions between agents, particularly focusing on group-wise relations. The AGE module outperforms pioneering group estimation modules in terms of performance and efficiency. Our comprehensive evaluation across three real-world datasets demonstrates that MART achieves state-of-the-art performance in multi-agent trajectory prediction, showcasing its potential to advance the field.

Limitation and Future Work.
Our model does not fully account for the temporal context of trajectories because it incorporates the past trajectory into the agent embedding. In future work, we plan to extend the model to incorporate spatio-temporal relationships. In addition, we aim to explore a trajectory prediction model with multi-modal inputs (e.g., trajectory inputs with scene context) because scene context, such as maps or images, is important in path prediction.




§ ACKNOWLEDGEMENTS

This research was supported by the Future Challenge Program through the Agency for Defense Development funded by the Defense Acquisition Program Administration (No. UC200015RD).






splncs04




section


 Supplementary Material 



§ DETAILED MODEL ARCHITECTURE


We provide detailed descriptions of the components in the model architecture, which were briefly explained earlier. For clarity, the bias vectors are omitted in this study. The components are detailed as follows: the node/pair-wise edge/hyperedge feature initializer ℱ_NI(·) / ℱ_PI(·) / ℱ_HI(·), the transformer-like node / pair-wise / hyperedge update function ℱ_n(·) / ℱ_e(·) / ℱ_h(·), and the decoder ℱ_D,k(·).



    
  * ℱ_NI(·): It comprises two fully connected layers. The sinusoidal positional encoding <cit.> blends the temporal information before the last layer. The equations are as follows:
    
    𝐧^(0)_i = Flatten(PosEnc(𝐗_i𝐖_NI,1))𝐖_NI,2,

    where Flatten(·) denotes the flattening operation that transforms a multi-dimensional tensor into a one-dimensional tensor. PosEnc(·) denotes the sinusoidal positional encoding <cit.> along the temporal dimension. 𝐖_NI,1∈ℝ^d_in× d_n and 𝐖_NI,2∈ℝ^T_p d_n × d_n are weight matrices for node feature initialization.
    
    
  * ℱ_PI(·) & ℱ_HI(·): Two layers of MLP with ReLU activation and the hidden dimension d_h.

    
  * ℱ_n(·): Assume that the output of multi-head relational or hyper-relational attention for the i-th agent is 𝐚^(l)_i, and given the node feature for the i-th agent 𝐧^(l)_i, the updated node feature 𝐧^(l+1)_i via ℱ_n(·) is described as follows:
    
    [ 𝐮^(l)_i = LayerNorm(𝐚_i^(l)𝐖_1 + 𝐧_i^(l)) ∈ℝ^d_n,;            𝐳^(l)_i = FeedForward(𝐮^(l)_i) ∈ℝ^d_n,;  𝐧^(l+1)_i = LayerNorm(𝐳_i^(l) + 𝐮_i^(l)) ∈ℝ^d_n, ]

    where 𝐖_1 ∈ℝ^d_n × d_n is a weight matrix, LayerNorm(·) denotes the layer normalization <cit.>, and FeedForward(·) represents two layers of MLP with ReLU activation and a hidden dimension of d_h.

    
  * ℱ_e(·) & ℱ_h(·): Given the message vector 𝐦^(l)_i for the i-th edge feature 𝐞^(l)_i, transformer-like edge update function is described as follows:
    
    [ 𝐮'^(l)_i = LayerNorm(𝐦_i^(l)𝐖_2 + 𝐞_i^(l)) ∈ℝ^d_e,;           𝐳'^(l)_i = FeedForward(𝐮'^(l)_i) ∈ℝ^d_e,; 𝐞^(l+1)_i = LayerNorm(𝐳'^(l)_i + 𝐮'^(l)_i) ∈ℝ^d_e, ]

    where 𝐖_2 ∈ℝ^d_h × d_e is a weight matrix.
    
    
  * ℱ_D,k(·): Three layers of MLP with ReLU activation and hidden dimensions of [d_D, d_D / 2], where d_D denotes the decoder dimension.





§ FURTHER EXPERIMENTAL DETAILS



 §.§ Datasets

NBA SportVU Dataset (NBA). The NBA collected the NBA SportVU dataset using the SportVU tracking system during the 2015–2016 NBA season. The raw data is in JSON format and contains frame information, such as the IDs and positions of the players on the court, team IDs, and the game period. Consequently, the trajectories of the 10 players and the ball during actual basketball games were recorded and included in the dataset. Following <cit.>, we used 32.5k samples for training, 5.0k samples for validation, and 12.5k samples for testing. Additionally, we utilized the past 2.0 seconds (10 frames) of trajectories to forecast the next 4.0 seconds (20 frames).


ETH-UCY dataset. The ETH-UCY dataset <cit.> includes five subsets: ETH, HOTEL, UNIV, ZARA1, and ZARA2, each comprising various pedestrian motion scenes captured at a frame rate of 2.5Hz. Based on previous studies <cit.>, our model predicts the future trajectories of multiple agents for 4.8 s (12 frames) using the given 3.2 s (8 frames) of past trajectories. For evaluation, we train on four sets and assess the performance on the unused set, utilizing the leave-one-out method.

Stanford Drone Dataset (SDD). The SDD dataset <cit.> is a large-scale pedestrian motion dataset, recorded from a bird's eye view on a university campus and consisting of 20 scenes. We utilize the standard train-test split based on a previous study <cit.>. Then, given the 3.2 seconds (8 frames) of past trajectories, we predict the future trajectories for 4.8 seconds (12 frames).



 §.§ Baseline Methods

We describe the baselines compared in Section <ref>.

    
  * NMMP <cit.>: This method introduces neural motion message-passing to model interactions among traffic actors for motion prediction in vehicles and social robots.
    
  * STAR <cit.>: This model is a spatio-temporal graph transformer focusing on trajectory prediction for crowd motion dynamics by leveraging attention mechanisms.
    
  * PECNet <cit.>: The model introduces the predicted endpoint-conditioned network for human trajectory prediction. This method models the stochastic endpoints and proposes a novel non-local social pooling layer.
    
    
    
    
  * GroupNet <cit.>: This method proposes a multiscale hypergraph message-passing neural network for multi-agent trajectory prediction to capture individual and group behaviors.
    
  * MemoNet <cit.>: This model is a novel instance-based trajectory prediction framework designed to predict future intentions by establishing direct connections between present scenarios and previously observed instances.
    
  * MID <cit.>: This approach, termed motion indeterminacy diffusion (MID), employs a parameterized Markov chain and a transformer-based diffusion model to progressively refine trajectory predictions.
    
  * NPSN <cit.>: This method addresses stochastic pedestrian trajectory prediction variability by introducing the Quasi-Monte Carlo method for improved sampling and proposing the Non-Probability Sampling Network (NPSN).
    
  * GP-Graph <cit.>: GP-Graph is a pedestrian trajectory prediction model that captures intra-group and inter-group interactions by utilizing a group assignment module to cluster pedestrians into proximity-based groups.
    
  * DynGroupNet <cit.>: This method, an extension of GroupNet <cit.>, is a dynamic-group-aware network designed to enhance the modeling of interactions among agents for trajectory prediction, capable of capturing time-varying interactions at both the pair-wise and group scales.
    
  * ET <cit.>: The paper presents EigenTrajectory (ET), a method for enhancing pedestrian trajectory prediction by creating a compact ET space for pedestrian movement representations.
    
  * EqMotion <cit.>: EqMotion is a motion prediction model that is equivariant under Euclidean geometric transformations while considering invariant interaction reasoning.
    
  * LED <cit.>: LED is a diffusion-based trajectory prediction model that offers real-time, accurate, and diverse predictions. This method accelerates inference speed by bypassing numerous denoising steps with a leapfrog initializer.





 §.§ Model Input

We utilize the absolute and relative positions as model inputs for the NBA dataset, resulting in an input dimension d_in of 4. For the ETH-UCY and SDD datasets, the relative positions are used as input, establishing the input dimension d_in at 2.



 §.§ Model Hyperparameter Search




Table <ref> lists the hyperparameters considered within the model architecture. Utilizing these candidates, we conducted a hyperparameter search on the NBA dataset, with the number of layers (L) fixed at 1. Subsequently, the optimal value for L was established, as detailed in Section <ref>.






 §.§ Model Training
 




We used an Adam optimizer <cit.> to train our network. The specific training hyperparameters, including the learning rate, batch size, and the number of training epochs, are detailed in Table <ref>. We applied a learning rate decay of 0.5 per 100 epochs for all datasets. Our training and evaluation were conducted on an NVIDIA GeForce RTX 3090. Our implementations were on PyTorch 2.1.0 <cit.>.




§ SUPPLEMENTARY EXPERIMENTS




 §.§ Effect of the Number of Layers

Table <ref> displays the effect of the number of layers (L) in MART. This experiment was conducted on the NBA dataset using the same condition described in Section <ref>. The performance of MART gradually improves until L=4, after which it begins to decline. This experiment illustrates why we selected four layers.





























 §.§ Ablation for Estimated Gradient within AGE
 
Figure <ref> illustrates the candidates for the estimated gradient within AGE. Table <ref> lists the ablations for the estimated gradient within AGE. The modified triangular-shaped estimator performs better than others. The original STE <cit.> and the triangle-shaped estimator <cit.> approximate the gradient of Sign(·). Because we estimate the unit step function 𝒰(·), we have made slight modifications to them, as illustrated in Figures <ref>-(a) and -(b).






§ FURTHER QUALITATIVE RESULTS




 §.§ Trajectory Prediction

Figure <ref> presents additional qualitative result comparisons on the NBA dataset. As depicted in the figure, our method yields more accurate predictions than recent state-of-the-art (SOTA) methods, including LED <cit.> and EqMotion <cit.>, across a larger number of samples. Figure <ref> shows qualitative results of MART on the ETH-UCY dataset. In some cases, we observed that the predicted trajectory does not satisfy the scene context, such as obstacles and maps.









 §.§ Group Estimation


This section provides a detailed description of the group estimation results presented in Section <ref>. Figure <ref> displays the visualization of group estimation results on the NBA dataset, while Figure <ref> shows the visualization of the group incidence matrix 𝐆 and the unique groups corresponding to Figure <ref>. Each group reflects a specific strategy as follows:

    
  * Group #1 (Zone Defense Group): Blue team defenders with player IDs 5, 7, and 9 employ man-to-man marking on red team offensive players with player IDs 0, 4, and 1, respectively.
    
  * Group #2 (Global Att&Def Group): The red team player with ID 2 dribbles the ball while being marked by the blue team player with ID 6. The player with ID 3 from the red team is connected to the player with ID 2 through an attacking link. Blue team players with IDs 6 and 8 are marking the red team players with IDs 2 and 3, respectively.
    
  * Group #3 (Local Att&Def Group): The red team player with ID 2 is dribbling with the ball. The blue team player with ID 8 is marking the attacking red team player with ID 3.
    
  * Group #4 (Local Att&Def Group): The blue team player with ID 6 defends against the attacking red team player with ID 2.








 Paper Title: Identifying the Hierarchical Emotional Areas in the Human Brain Through Information Fusion.



1
.001


Identifying the Hierarchical Emotional Areas in the Human Brain Through Information Fusion


Z. Huang et al.


mode=title]Identifying the Hierarchical Emotional Areas in the Human Brain Through Information Fusion

1,2]Zhongyu Huang


huangzhongyu2020@ia.ac.cn


Conceptualization, Data curation, Formal analysis, Investigation, Methodology, Software, Visualization, Writing – original draft


[1]organization=Laboratory of Brain Atlas and Brain-Inspired Intelligence, State Key Laboratory of Multimodal Artificial Intelligence Systems, Institute of Automation, Chinese Academy of Sciences,
            city=Beijing,

            postcode=100190,
            country=China

[2]organization=School of Artificial Intelligence, University of Chinese Academy of Sciences,
            city=Beijing,

            postcode=100049,
            country=China

1]Changde Du

Funding acquisition, Supervision, Validation, Writing – review & editing

3]Chaozhuo Li

Investigation, Methodology, Writing – review & editing

[3]organization=Key Laboratory of Trustworthy Distributed Computing and Service (MoE), Beijing University of Posts and Telecommunications,
            city=Beijing,

            postcode=100876,
            country=China

1,2]Kaicheng Fu

Conceptualization, Data curation

1,2]Huiguang He

huiguang.he@ia.ac.cn

Funding acquisition, Project administration, Supervision, Writing – review & editing


[1]


[1]Corresponding author


§ ABSTRACT 
The brain basis of emotion has consistently received widespread attention, attracting a large number of studies to explore this cutting-edge topic.
However, the methods employed in these studies typically only model the pairwise relationship between two brain regions, while neglecting the interactions and information fusion among multiple brain regions—one of the key ideas of the psychological constructionist hypothesis.
To overcome the limitations of traditional methods, this study provides an in-depth theoretical analysis of how to maximize interactions and information fusion among brain regions.
Building on the results of this analysis, we propose to identify the hierarchical emotional areas in the human brain through multi-source information fusion and graph machine learning methods.
Comprehensive experiments reveal that the identified hierarchical emotional areas, from lower to higher levels, primarily facilitate the fundamental process of emotion perception, the construction of basic psychological operations, and the coordination and integration of these operations.
Overall, our findings provide unique insights into the brain mechanisms underlying specific emotions based on the psychological constructionist hypothesis.





  * Identify hierarchical emotional areas to study brain mechanisms underlying emotion.

  * Conduct an in-depth theoretical analysis based on information fusion and graph theory.

  * Develop a novel framework to improve emotion decoding using the identified areas.

  * Demonstrate the potential of identified areas by cross-dataset emotion decoding.





Emotion Human brain Hierarchical emotional areas Multi-source information fusion Graph machine learning Psychological constructionism


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






§ INTRODUCTION



Emotions are complex psychological and physiological states that arise in response to stimuli from internal feelings or external environments.
Generally, these external stimuli are perceived in diverse manners, such as visual experiences, audio signals, and textual information.
Our brains respond to these stimuli by generating neural activity in various brain regions, thereby producing corresponding psychological and physiological states.
Therefore, the fusion of multi-source information from either diverse manners or various brain regions contributes to understanding the brain mechanisms underlying emotion <cit.>.

The research on the brain basis and neural representation of emotion mainly focuses on two hypotheses: locationism and psychological constructionism.
The locationist hypothesis suggests that discrete emotion categories consistently and specifically correspond to distinct brain regions, implying a one-to-one mapping between brain regions and emotion categories <cit.>.
On the contrary, the psychological constructionist hypothesis suggests that discrete emotion categories are constructed from more general brain networks, i.e., the interactions among various brain regions, rather than specific brain regions <cit.>.
While both hypotheses are of equal age, there is relatively little evidence supporting the locationist hypothesis.
In contrast, a large number of studies <cit.> demonstrate that the neural representation of emotion is distributed, leading to the conclusion that the psychological constructionist hypothesis has greater potential for revealing the brain basis of emotion.

To investigate the brain basis of emotion, many previous studies have employed various methods to identify patterns of brain activity associated with specific emotions, such as multivariate pattern analyses <cit.>, Bayesian approaches <cit.>, and brain networks <cit.>.
However, these methods either fail to leverage the relationships between brain regions or only model the pairwise relationship between two brain regions, while neglecting the interactions and information fusion among multiple brain regions—one of the key ideas of the psychological constructionist hypothesis.
As a result, their findings on the brain regions involved in emotion encoding vary, and even a consistent conclusion on the brain mechanisms for specific emotions has yet to be established.

To reveal the brain mechanisms underlying emotion based on the psychological constructionist hypothesis, it is essential to consider the interactions and information fusion among multiple brain regions at a higher level.
From a network neuroscience perspective, the brain regions involved in emotion encoding present a modular structure <cit.>.
Modules with similar functions are coupled together, forming a unique area with a specific function.
We refer to this unique area as the emotional area, similar to the functional areas in the human brain.
Moreover, a series of influential studies <cit.> demonstrates that brain networks further exhibit hierarchically modular organizations.
The main module in a brain network plays a dominant role, with a set of sub-modules cooperating to perform subordinate functions.
In addition to these sub-modules, there exist sub-sub-modules and so forth, each playing an increasingly subtle role.
Accordingly, in the context of emotion encoding, the entire brain network can be decomposed into hierarchical subnetworks that correspond to hierarchical emotional areas.

Since emotions are complex subjective experiences, they include multiple psychological operations and cognitive processes, thereby involving various brain regions <cit.>.
When an individual experiences emotional stimuli, the entire brain network is decomposed into hierarchical subnetworks.
Each of these subnetworks corresponds to an emotional area and performs distinct functions, contributing to different psychological operations and cognitive processes involved in emotion encoding.
Meanwhile, subnetworks in various emotional areas interact and collaborate to create a comprehensive brain network for emotion encoding.
Based on the hierarchically modular organization of the brain network, this analysis aligns perfectly with both the key idea of the psychological constructionist hypothesis and the professional knowledge in affective science.
Therefore, this work aims to investigate the hierarchical emotional areas in the human brain through multi-source information fusion and graph machine learning methods.
Our ultimate goal is to reveal the brain mechanisms underlying specific emotions based on the psychological constructionist hypothesis.

Overall, the key innovations and main contributions of our research can be summarized as follows:

    
  * We introduce the concept of hierarchical emotional areas to understand the brain mechanisms underlying emotion, aligning perfectly with both the key idea of the psychological constructionist hypothesis and the professional knowledge in affective science.
    
  * We provide an in-depth theoretical analysis of how to maximize interactions and information fusion among brain regions. Building on the results of this analysis, we identify hierarchical emotional areas in the human brain.
    
  * We develop a novel framework, Hierarchical Emotion Network (HEmoN), which improves emotion decoding tasks by exploiting the identified hierarchical emotional areas.
    
  * To verify the rationality and effectiveness of our approach, we conduct experiments on two large-scale datasets, both containing audio-visual multi-source naturalistic stimuli. Furthermore, we accomplish a challenging task: identifying hierarchical emotional areas on one dataset and applying them to cross-dataset emotion decoding on another dataset.




§ RELATED WORK


Emotions arise from the activations of specialized neuronal populations in the human brain.
Therefore, researchers have devoted themselves to exploring the neural representation of emotion to gain a deeper understanding of emotion.
One of the most prominent research areas is emotion encoding and decoding, which aims to identify the mapping patterns between emotional stimuli and their corresponding brain responses.
In essence, encoding and decoding are complementary operations: encoding investigates how to transform stimuli into brain responses, while decoding investigates how to predict information about the stimuli from brain responses <cit.>.

In laboratory settings, researchers usually measure and collect the brain signals of subjects when they experience emotional stimuli.
By analyzing the collected brain signals and the experienced emotional stimuli, researchers can delve into emotion encoding and decoding.
Commonly used brain signals include functional magnetic resonance imaging (fMRI), electroencephalography (EEG), and magnetoencephalography (MEG).
Due to its high spatial resolution, fMRI allows precise localization of the active brain regions involved in emotional processing <cit.>, providing an excellent opportunity to study the underlying brain mechanisms for specific emotions <cit.>.
On the other hand, traditional emotional stimuli include images, especially facial expressions <cit.>, and mental imagery <cit.>.
However, these stimuli are abstract, simplified, and, in many ways, fail to reflect the complexity and dynamics inherent in real-life behaviors and stimuli <cit.>.
To investigate real-life emotional experiences, researchers are increasingly turning to naturalistic paradigms, such as watching movies <cit.>, listening to stories <cit.>, or experiencing audio-visual multi-source stimuli <cit.>.
A growing body of evidence suggests that these naturalistic stimuli benefit from the fusion of multi-source information, incorporate richer and more dynamic experiences, evoke stronger emotions, and provide a more reasonable approximation to freeform cognition in our daily lives <cit.>.

Given their simplicity, efficiency, and strong interpretability, a large number of studies <cit.> employ linear regression or multivariate statistical methods to develop models for emotion encoding or decoding.
Unfortunately, these methods are unable to leverage the rich relationships between brain regions.
To overcome this limitation, some pioneering works <cit.> use fMRI data to construct functional brain networks <cit.>, which can explicitly model the pairwise relationship between brain regions.
A series of works <cit.> uses diverse network measures to assess the significance of each region of interest in the brain network, thereby studying the mechanisms of emotion encoding.
Another series of works <cit.> maps various emotion categories into distinct brain networks with more distinguishable graph structural features, thereby improving the performance of emotion decoding.
In addition, several recent advances <cit.> have made remarkable achievements in interpreting emotions through various graph machine learning methods.
<cit.> investigate the connectivity within the brain network and reveal a general emotion pathway connecting neural nodes involved in basic psychological operations.
<cit.> design a graph neural network framework for analyzing brain networks and utilize it to infer salient brain regions for emotional tasks.
<cit.> construct an emotion-brain bipartite graph to model the relationships between emotions and brain regions, thereby improving the representation learning process.
<cit.> study the encoding rules of emotional instances and discover that they are represented in an extensive network encompassing sensory, prefrontal, and temporal regions.

Although these works have employed graph machine learning methods to investigate emotional information in the brain more effectively, their primary focus remains on identifying key brain regions or capturing pairwise relationships between two objects.
These approaches have essentially deviated from the key idea of the psychological constructionist hypothesis, which emphasizes the interactions and information fusion among multiple brain regions throughout the entire methodology.
In addition, these works have rarely considered the hierarchically modular organization of the brain network, neglecting important facts in network neuroscience.
Consequently, we aim to investigate the hierarchical emotional areas in the human brain by combining graph machine learning methods with information fusion, thereby overcoming the above challenges and revealing the brain mechanisms underlying specific emotions.



§ PRELIMINARIES




 §.§ Graph theory


Let G = (𝒱, ℰ) ∈𝒢 be a graph with vertex set 𝒱 = {v_1, v_2, …, v_N_G} and edge set ℰ = {e_1, e_2, …, e_M_G}, where N_G = |𝒱| and M_G = |ℰ| represent the number of vertices and edges in G, respectively.
Let A ∈ℝ^N_G × N_G be the adjacency matrix of G, and X ∈ℝ^N_G × c_0 be its node feature matrix.
For each node v ∈𝒱, x_v ∈ℝ^c_0 denotes its node feature vector, and d_v denotes its degree (i.e., the number of 1-hop neighbors).
The set of v's 1-hop neighbors is represented as 𝒩(v).

A walk in graph G is a finite sequence of alternating vertices and edges, such as v_0, e_1, v_1, e_2, …, e_m, v_m, in which each edge e_i = (v_i-1, v_i).
A walk may contain repeated edges.
A path is a walk in which all vertices, and hence all edges, are distinct (except, possibly, v_0 = v_m).
We denote a path P with length m as its vertex sequence (v_0, v_1, …, v_m).
If v_0 = v_m and m ≥ 2, this path is called a cycle.

In an unweighted graph G, the shortest path between two vertices is the path with the minimum length between them, and its length is called the (shortest path) distance between these two vertices.
We use dist(u, v) to denote the distance between a pair of vertices u and v.
The longest shortest path in graph G is the shortest path with the maximum length between any pair of vertices.
The diameter of graph G is the length of the longest shortest path, defined as

    𝔡 = max{dist(u, v) |∀ u, v ∈𝒱}.

In this paper, we use the normal d to denote the node degree and the Gothic 𝔡 to denote the graph diameter.

A tree is a special type of graph that is connected and acyclic.
Trees have several essential properties; for instance, there exists exactly one path between any pair of vertices.
Additionally, a tree with N_T vertices has exactly N_T-1 edges.
A forest is a collection of disjoint trees, and a tree is a special type of forest with only one connected component.



 §.§ Psychological constructionist hypothesis


The psychological constructionist hypothesis assumes that emotion categories are common-sense categories, with instances emerging from the combination of more basic psychological operations, such as core affect, conceptualization, executive attention, language, etc.
These psychological operations are the common ingredients of all mental states.

The core affect is a mental representation of bodily sensations, enabling an organism to determine whether something in the environment has motivational salience, e.g., whether it is beneficial, detrimental, approachable, or avoidable.
The conceptualization is a process that imbues sensations from the body or the external world with meaning in a given context by drawing on representations of previous experiences.
These representations are activated by the current physical and psychological situation.
The executive attention is a process that selectively enhances some representations while suppressing others.
It can regulate the activity in other processes, such as core affect, conceptualization, or language use.
The language word serves as an “essence placeholder,” which helps integrate feelings, behaviors, and expressions into instances of a meaningful emotion category.



§ HIERARCHICAL EMOTIONAL AREAS: THEORETICAL ANALYSIS AND METHODOLOGY


In this section, we present the identification of hierarchical emotional areas and their application in emotion decoding.
We start by extracting a brain tree from the brain network to overcome its challenges and facilitate the capture of hierarchical information.
Furthermore, we provide an in-depth theoretical analysis based on this brain tree to maximize interactions and information fusion among brain regions, aligning with the key idea of the psychological constructionist hypothesis.
Building on the results of this analysis, we identify the hierarchical emotional areas and present corresponding technical implementations.
Finally, we develop a novel framework to improve emotion decoding tasks by exploiting the identified hierarchical emotional areas.
The pipeline of this study is illustrated in Figure <ref>.





 §.§ Extracting the brain tree



We first construct a brain network, G = (𝒱, ℰ), using the collected fMRI signals.
In this network, nodes are defined as brain regions of interest (ROIs), generally derived from an existing brain parcellation, such as the Power Atlas <cit.> and the Destrieux Atlas <cit.>.
The edges are defined as the functional connectivity between these ROIs, with functional connectivity commonly computed as the pairwise correlation of the fMRI time series between two ROIs.

However, conventional brain networks may suffer from various limitations, such as threshold problems <cit.>, spurious and noisy connections <cit.>, scaling effects <cit.>, and methodological biases <cit.>.
Fortunately, many studies <cit.> have demonstrated that the spanning tree can effectively overcome these challenges and capture the essential properties of complex brain networks, leading to great success in a wide range of applications.
Moreover, previous research <cit.> has indicated that the tree structure is an efficient tool for capturing hierarchical information from data, aligning with our goal of identifying hierarchical emotional areas.
Therefore, we propose to extract a brain tree from the original brain network using the spanning tree algorithm.

Specifically, given a brain network, we first sort all edges in descending order according to their weights.
The construction process of the brain tree begins by using the edge with the highest weight.
Subsequently, we select the edge with the highest weight from the remaining edges.
If the newly selected edge forms a cycle with the previously selected edges, it is discarded; otherwise, it is added to the brain tree.
This process is executed iteratively until there are N_G-1 edges in the brain tree.
Finally, all edge weights in the brain tree are set to a value of one.
In this way, we transform the brain network G = (𝒱, ℰ) into a brain tree T = (𝒱, ℛ), where 𝒱 denotes the vertex set of T, identical to that of G, and ℛ denotes the edge set of T.



 §.§ Theoretical analysis: Interactions and information fusion in the brain tree


As discussed in Section <ref>, the psychological constructionist hypothesis emphasizes the interactions and information fusion among multiple brain regions.
Since these interactions can be essentially considered as the information influence that one brain region (i.e., node) exerts on another, we start our analysis by quantitatively describing the information influence between an arbitrary pair of nodes, u and v.
Based on previous literature <cit.>, we take the random walk as a general rule for information propagation in the brain tree and define node influence as follows.


    Let h_v^(0)∈ℝ^c_0 be the initial feature vector of node v, and h_u^(k) be the representation vector of node u after k steps of the random walk.
    The node influence I(u, v) of node v on node u after k steps of the random walk is defined as the norm of the Jacobian matrix ∂ h_u^(k) / ∂ h_v^(0):
    
    I(u, v) = ‖∂ h_u^(k)/∂ h_v^(0)‖,

    where the norm ‖·‖ is an arbitrary induced norm[An induced norm is a matrix norm induced by a vector norm, defined as ‖ A ‖ = sup{‖Ax‖|‖ x ‖ = 1 }.].


The node influence quantifies how changes in the information of node v affect that of node u.
The following theorem presents an analytical expression of the node influence for further analysis.



    Given a tree T, the node influence I_T(u, v) of node v on node u after k = dist(u, v) steps of the random walk is
    
    I_T(u, v) = | a_uv_1· a_v_1v_2⋯ a_v_k-1v/∑_i ∈𝒩(u)a_ui·∑_j ∈𝒩(v_1)a_v_1j⋯∑_l ∈𝒩(v_k-1)a_v_k-1l|,

    where the only path between nodes u and v is denoted as (u, v_1, v_2, ⋯, v_k-1, v);
    a_v_iv_j is the edge weight between nodes v_i and v_j, and 𝒩(u) is the set of u's 1-hop neighbors.


Detailed proof of Theorem <ref> is provided in Appendix <ref>.
Since all edges in the brain tree have a uniform weight of 1, the result of Theorem <ref> can be further simplified to

    I_T(u, v) = 1/d_u d_v_1⋯ d_v_k-1

where d_u is the degree of node u.
Eq. (<ref>) indicates that node influence decreases as the distance k between nodes and node degrees increase, consistent with intuition.
More precisely, node influence I_T(u, v) decreases exponentially with the increase in the distance between nodes u and v, and decreases polynomially with the increase in node degrees on the path connecting u and v.

To facilitate effective interactions among brain regions, we aim to maximize the node influence I_T(u, v) in the brain tree.
This objective can be achieved by solving the following optimization problem:

    max_d_u, d_v_i     
        I_T(u, v) = 1/d_u d_v_1⋯ d_v_k-1
    s.t.     
        d_u  ≥ 1 
       
        d_v_i≥ 2, ∀ i = 1, 2, …, k-1

Clearly, the optimal solution is d̂_u = 1, d̂_v_i = 2, which corresponds to the path between nodes u and v.
Since trees have an essential property that there exists exactly one path between any pair of nodes, the path between u and v is also the shortest path.
Accordingly, the optimal solution motivates us to propagate information along the shortest path, thereby facilitating interactions among brain regions.

Furthermore, it is essential to consider the fusion of information among multiple brain regions, enabling them to form a unique emotional area with a specific function.
The above analysis indicates that information should be propagated along the (shortest) path.
Consequently, we continue our analysis by quantifying the total amount of information carried by a path and defining path information as follows.


    Let P denote a path of length m with vertex sequence (v_0, v_1, …, v_m).
    The path information I(P) of path P is defined as the sum of the node influences between every pair of nodes on the path:
    
    I(P) = ∑_i,j=0 
     i ≠ j^m I_P(v_i, v_j).



The path information quantifies the extent of information fusion among nodes on path P.
We aim to maximize the path information, thereby enhancing the fusion of information among brain regions as much as possible.
The following theorem gives the path with the maximum path information.



    Given a tree T, the longest shortest path P̂ has the maximum path information
    
    I(P̂) = ∑_k=1^𝔡(𝔡-k+1) ·2/2^k-1,

    where 𝔡 is the diameter of tree T.


Detailed proof of Theorem <ref> is provided in Appendix <ref>.
Theorem <ref> indicates that the longest shortest path has the maximum amount of information among all possible paths in a brain tree.
Consequently, the theoretical analysis in this subsection demonstrates that propagating information along the longest shortest path on a brain tree maximizes interactions and information fusion among brain regions.
In the next subsection, we will exploit the longest shortest path to further identify the emotional area and present the corresponding implementations.



 §.§ Identifying the hierarchical emotional areas



In the natural world, the tree trunk is an essential component of a tree.
The remaining parts of a tree after removing the tree trunk are also regarded as (sub)trees, each with its own “(sub)tree trunk.”
Thus, natural tree trunks exhibit hierarchical organization throughout the entire tree.
Inspired by this interesting natural phenomenon, we decompose the entire brain tree into hierarchical trunks and leverage them to identify the hierarchical emotional areas.

Specifically, the 1st-level trunk, 𝔱_1, serves as the main trunk of the entire brain tree, and there is usually only one main trunk at the first level.
The trunk at the remaining ℓth (ℓ≥ 2) level, 𝔱_ℓ, is considered a branch of the entire brain tree, and there are usually multiple branches at the ℓth level.
We denote the set of trunks at the ℓth level as 𝒯_ℓ.
Our previous theoretical analysis suggests that propagating information along the longest shortest path on a brain tree maximizes interactions and information fusion among brain regions.
Therefore, we designate the longest shortest path as the main trunk 𝔱_1.
In addition, the brain regions corresponding to all nodes on this main trunk constitute the main emotional area, denoted as 𝒜_1.
Once the main trunk is obtained and removed from the brain tree, the remaining part becomes a forest.
This forest contains at least one (often multiple) connected component(s), and every component is a (sub)tree <cit.>.
To identify the trunks at the ℓth (ℓ≥ 2) level, we search for the longest shortest path in each connected component of the forest after removing the trunks of all previous (ℓ-1) levels.
Similarly, the brain regions corresponding to all nodes on these ℓth-level trunks constitute the ℓth-level emotional area, denoted as 𝒜_ℓ.
This process runs iteratively and ends once all nodes in the tree have been removed.
Figure <ref> in Appendix <ref> provides a concrete example to illustrate the process above.



Suppose there are L levels in total.
Algorithm <ref> describes the detailed process of identifying hierarchical emotional areas.
Here, F_ℓ represents a forest obtained after removing the trunks of all previous ℓ levels and the isolated nodes.
Line 2 initializes the forest F_0 to tree T = (𝒱, ℛ).
Line 5 considers all ℭ_ℓ-1 connected component(s) in F_ℓ-1.
Line 6 searches for the longest shortest path in the 𝔠th connected component of F_ℓ-1 to identify an ℓth-level trunk 𝔱_ℓ,𝔠.
Lines 7 and 14 identify the emotional area, 𝒜_ℓ, at the ℓth level.
Lines 8-9 remove all edges in 𝔱_ℓ,𝔠 from F_ℓ-1, and Lines 10-11 remove all isolated nodes.
Line 15 updates F_ℓ based on the results of Lines 9 and 11.
If there are no remaining nodes, Line 16 terminates the loop, and finally, Line 17 returns the emotional areas at every level.



 §.§ Practical application: Emotion decoding



We further develop a novel framework, Hierarchical Emotion Network (HEmoN), to improve emotion decoding tasks by exploiting the identified hierarchical emotional areas.
As previously analyzed, we identify hierarchical trunks by iteratively searching for the longest shortest path.
Given that nodes on the longest shortest path (i.e., a trunk) form a long sequence, we use the Long Short-Term Memory (LSTM) <cit.> to effectively capture long-term dependencies and recognize complex patterns, thereby facilitating the representation learning process of emotion decoding.
Specifically, after obtaining the trunks at every level according to Line 13 of Algorithm <ref>, we apply an LSTM to learn the representation for each trunk along its corresponding trunk path, formulated as

    h_T^(ℓ) = ∑_𝔱_ℓ∈𝒯_ℓLSTM(  x_v_0^(𝔱_ℓ),  x_v_1^(𝔱_ℓ), ⋯,  x_v_k_𝔱_ℓ^(𝔱_ℓ))

where 𝒯_ℓ is the set of trunks at the ℓth level,
and 𝔱_ℓ is a single trunk in 𝒯_ℓ;
let the length of 𝔱_ℓ be k_𝔱_ℓ;
x_v_i^(𝔱_ℓ)∈ℝ^c_0 denotes the feature vector of node v_i on trunk 𝔱_ℓ,
and h_T^(ℓ)∈ℝ^c_ℓ denotes the trunk representation at the ℓth level.
Since nodes on all trunks in 𝒯_ℓ constitute the ℓth-level emotional area, h_T^(ℓ) also serves as the representation of the emotional area at the ℓth level.
Then, we combine all the trunk representations to create a representation of the brain tree:

    h_T = ∑_ℓ=1^L  W^(ℓ) h_T^(ℓ)

where W^(ℓ) represents a learnable weighting matrix for the ℓth level, optimized along with other model weights during the model training process.
For emotion classification tasks, the brain tree is associated with an emotion category label y_T from C emotion categories.
Accordingly, we set W^(ℓ)∈ℝ^C × c_ℓ and further input h_T into the softmax function to obtain the predicted class.
For emotion regression tasks, the brain tree is associated with an emotion rating vector y_T ∈ [0,a]^C containing C emotion categories, where a represents the maximum rating and is usually set to 100.
Accordingly, we set W^(ℓ)∈ℝ^C × c_ℓ and further input h_T into the logistic function to obtain the predicted ratings.
The above representation learning process is illustrated in Appendix <ref> through a concrete example.



§ EXPERIMENTS


In this section, we conduct comprehensive experiments on two large-scale datasets, both containing audio-visual multi-source naturalistic stimuli.
The entire experimental evaluation consists of two stages.
In the first stage (Section <ref>), we identify hierarchical emotional areas based on the overall experience of all emotions in Dataset 1 <cit.>.
In the second stage (Section <ref>), we use the proposed model, HEmoN, which builds on the identified hierarchical emotional areas, to perform cross-dataset emotion decoding on Dataset 2 <cit.>.
Furthermore, we provide a thorough analysis and discussion of the experimental results in Section <ref>.
Our code is publicly available at <https://github.com/zhongyu1998/HEmoN>.





 §.§ Dataset description


We conduct comprehensive experiments on two large-scale datasets <cit.> to demonstrate the effectiveness and superiority of our proposed approach.
Both datasets contain audio-visual multi-source naturalistic stimuli and use a similar experimental paradigm for data acquisition.
Specifically, they collect fMRI responses when subjects watch emotionally charged movies (i.e., emotional stimuli).
To obtain the corresponding emotion ratings for these stimuli, they recruit multiple human raters who are not involved in the fMRI experiments.
These raters label each emotional stimulus according to specific emotion categories, assigning a score to each emotion category for every stimulus.
The final emotion rating for a specific emotion category of a stimulus is determined by averaging the scores across these raters.
As for fMRI responses, we use the fMRI data preprocessed by <cit.> and <cit.>.
Here, we provide an overview of the basic information and fMRI data preprocessing steps for each dataset.
Further details on emotion ratings, fMRI data acquisition, and fMRI data preprocessing are available in their original literature.



  
Dataset 1
This dataset contains fMRI responses collected from 15 subjects when they experience emotional stimuli, featuring 3599 pairs of {fMRI responses, emotion ratings} for each subject.
The emotional stimuli are composed of movie segments from Forrest Gump, sourced from the StudyForrest project (<http://studyforrest.org>), and contain six basic emotions <cit.>: happiness, surprise, fear, sadness, anger, and disgust.
The fMRI data are preprocessed using the Advanced Normalization Tools (ANTs) <cit.> and the Analysis of Functional NeuroImages (AFNI) <cit.>.
This process mainly involves the following steps:
spike correction, slice timing correction, motion correction, rigid body transformation, registration to T1-weighted images, phase distortion correction, spatial smoothing with a Gaussian kernel, and physiological noise filtering.



  
Dataset 2
This dataset contains fMRI responses collected from 8 subjects when they experience emotional stimuli, featuring 5400 pairs of {fMRI responses, emotion ratings} for each subject.
The emotional stimuli are composed of movie clips sourced from the video-sharing site Vimeo (<https://vimeo.com/jp>) and contain 80 emotion categories.
The complete list of these 80 emotion categories and their sources are available in the original literature <cit.>.
The fMRI data are preprocessed using the Statistical Parameter Mapping toolbox (SPM8, <https://www.fil.ion.ucl.ac.uk/spm/software/spm8/>).
This process mainly involves the following steps:
motion correction, z-score normalization, median filter convolution, long-term trend removal, and registration to T1-weighted images.



 §.§ Implementation details


We initialize the learnable weighting matrix in Eq. (<ref>) and all model weights using the Xavier initializer <cit.>.
Then, we pre-map the emotional stimuli into 64 dimensions with a linear transformation and use a batch of 32 emotional stimuli as the model input.
Our HEmoN model consists of 3 LSTM layers (Eq. (<ref>)) and a final combination layer (Eq. (<ref>)).
Each LSTM layer contains 64 hidden units, with a dropout <cit.> ratio of 0.2 applied to the outputs of each LSTM layer except the last one.
The model is trained using the Adam optimizer <cit.> with an initial learning rate of 0.0005.
If the performance on the validation set does not improve after 10 epochs, the learning rate is reduced by a factor of 0.5.
The training is stopped when the learning rate reaches the minimum value of 2E-5 or when the number of epochs reaches the maximum value of 300.

Our model is implemented with Python version 3.6, PyTorch version 1.4.0 <cit.>, and PyTorch Geometric (PyG) version 1.7.0 <cit.>.
All experiments are conducted on a Linux server equipped with an Intel® Xeon® CPU E5-2650 v4 and 8 NVIDIA TITAN RTX GPUs.



 §.§ Identification of hierarchical emotional areas



To identify the hierarchical emotional areas, we conduct experiments on Dataset 1 and use the Power Atlas <cit.> to define ROIs.
The Power Atlas comprises 264 putative functional areas (i.e., ROIs) distributed across 13 functional systems and one uncertain system.
All the systems are listed at the top of Figure <ref>.
The fMRI time series for each ROI is obtained by averaging the activity of voxels within a fixed-radius sphere of 5mm, centered on the ROI coordinates defined by the Power Atlas.
Subsequently, we construct the brain network, extract the brain tree, and identify hierarchical emotional areas according to the methodologies described in Sections <ref> and <ref>.
We present the identification results for the overall experience of all basic emotions, as well as those for the single experiences of the most representative positive and negative emotions: happiness and sadness.
The identified hierarchical emotional areas for these three types are illustrated in Figure <ref>.



 §.§ Performance of emotion decoding



To perform cross-dataset emotion decoding, we build the HEmoN model by exploiting the identified hierarchical emotional areas and evaluate its performance on Dataset 2.
Let y_i and ŷ_i denote the true and predicted emotion rating vectors for the ith stimulus, respectively, where ŷ_i is calculated according to the steps described in Section <ref>.
Additionally, let the scalars y_ij and ŷ_ij denote the true and predicted ratings for the jth emotion category in the ith stimulus, respectively.
We take the Mean Absolute Error (MAE), which is widely used in regression tasks, as our evaluation metric.
The MAE value α is calculated as:

    α = 1/N_s∑_i=1^N_s∑_j=1^C | y_ij - ŷ_ij|

where N_s denotes the number of stimuli.
Since most emotion categories have ratings close to or containing zeros, Eq. (<ref>) performs the sum operation ∑_j=1^C instead of the average operation 1/C∑_j=1^C across all C emotion categories.

Since Dataset 2 has already been split into training and test sets, we adopt the default split and perform emotion decoding between the fMRI responses to each emotional stimulus and its corresponding emotion ratings.
Specifically, in each experimental run, we use the training set consisting of 3600 pairs to train the model and then evaluate its performance on the test set consisting of 1800 pairs.
We repeat this process 10 times and report the evaluation result based on the test MAEs of these 10 runs.
Furthermore, we compare HEmoN with multiple competitive baselines, including
Fully-connected Neural Network (FNN),
Graph Convolutional Network (GCN) <cit.>,
Graph Isomorphism Network (GIN) <cit.>,
BrainNetCNN <cit.>,
BrainGNN <cit.>,
Brain Network Transformer (BrainNetTF) <cit.>,
and Graph-enhanced Emotion Decoding (GED) <cit.>.
The evaluation results on Dataset 2 are summarized in Figure <ref>.





 §.§ Ablation study


The hierarchical emotional areas and the information fusion rule of brain regions play pivotal roles in our HEmoN model.
Therefore, we conduct ablation studies to investigate their impacts on model performance.
Specifically, we replace the hierarchical emotional areas with a single 1st-level emotional area to assess the impact of hierarchical emotional areas.
In addition, we perform a depth-first traversal on the brain tree and fuse information among brain regions based on the traversal result instead of hierarchical emotional areas to assess the impact of the information fusion rule.
The corresponding results are shown in Figure <ref>, where the suffixes “-EA1” and “-DFT” represent the 1st-level emotional area and the depth-first traversal, respectively.



§ DISCUSSION



Figure <ref>a reveals that for the overall experience of all basic emotions, the 1st-level emotional area mainly includes sensory/somatomotor, auditory, default mode, visual, and subcortical systems.
These functional systems involve various senses, including sight (visual), sound (auditory), and touch (sensory/somatomotor), enabling individuals to perceive external emotional stimuli and thus creating necessary conditions for the construction of core affect.
In addition, a small portion of this area includes salience and ventral attention systems, which also contribute to the perception of external stimuli and the construction of core affect.
Moreover, the default mode and subcortical systems play key roles in distinguishing experiences of discrete emotions <cit.>, thereby promoting the recognition, identification, and perception of emotions.
In summary, the 1st-level emotional area primarily facilitates the fundamental process of emotion perception.
The 2nd-level emotional area mainly includes sensory/somatomotor, cingulo-opercular task control, auditory, default mode, visual, fronto-parietal task control, salience, subcortical, and cerebellar systems.
This emotional area primarily facilitates the construction of basic psychological operations.
For instance, the anterior cingulate cortex in the cingulo-opercular task control system and the amygdala in the subcortical system both play pivotal roles in core affect <cit.>.
In addition, a small portion of this area includes memory retrieval, ventral attention, and dorsal attention systems, which contribute to the construction of conceptualization and executive attention.
Moreover, Broca’s and Wernicke’s areas contained in this area are essential for the production and comprehension of human language <cit.>.
It is worth noting that our finding of the 2nd-level emotional area is consistent with a previous study <cit.>, which identified a general emotion pathway connecting neural nodes involved in these basic psychological operations.
The 3rd-level emotional area mainly includes sensory/somatomotor, cingulo-opercular task control, default mode, visual, fronto-parietal task control, salience, and dorsal attention systems.
These functional systems typically participate in advanced cognitive processes, thereby facilitating the coordination and integration of the above basic psychological operations.

Figure <ref>b reveals that for a single experience of happiness, the 1st-level emotional area mainly includes sensory/somatomotor, default mode, visual, and subcortical systems.
In addition, a small portion of this area includes salience and ventral attention systems.
These functional systems are similar to those found in all basic emotions, with the exception of the auditory system.
A possible explanation for the absence of the auditory system is that visual information is more dominant than auditory information in the perception of happiness.
This conjecture is supported by an interesting study <cit.>, which demonstrates that happiness is more easily perceived visually than auditorily.
The 2nd-level emotional area mainly includes sensory/somatomotor, cingulo-opercular task control, default mode, visual, fronto-parietal task control, salience, subcortical, ventral attention, and cerebellar systems.
In addition, a small portion of this area includes auditory, memory retrieval, and dorsal attention systems.
The activation of happiness in the brain is associated with multiple brain regions within these functional systems, such as the anterior cingulate cortex, prefrontal cortex, and insula in cingulo-opercular task control and salience systems <cit.>.
The 3rd-level emotional area mainly includes sensory/somatomotor, auditory, default mode, visual, fronto-parietal task control, salience, ventral attention, and dorsal attention systems.

Figure <ref>c reveals that for a single experience of sadness, the 1st-level emotional area mainly includes auditory, default mode, memory retrieval, visual, fronto-parietal task control, ventral attention, and dorsal attention systems.
These functional systems differ significantly from those found in all basic emotions.
Compared with the basic emotions, these functional systems incorporate memory retrieval, fronto-parietal task control, ventral attention, and dorsal attention systems, while excluding sensory/somatomotor and subcortical systems.
This distinction may stem from the fact that sadness can impair our attention, memory, and decision-making abilities <cit.>.
As a result, the above functional systems are the first to be affected when perceiving sadness.
The 2nd-level emotional area mainly includes sensory/somatomotor, cingulo-opercular task control, auditory, default mode, visual, fronto-parietal task control, salience, and subcortical systems.
The activation of sadness in the brain is associated with multiple brain regions within these functional systems, such as the anterior cingulate cortex, amygdala, and hippocampus in cingulo-opercular task control and subcortical systems <cit.>.
The 3rd-level emotional area mainly includes sensory/somatomotor, cingulo-opercular task control, auditory, default mode, visual, fronto-parietal task control, salience, subcortical, ventral attention, dorsal attention, and cerebellar systems.

Furthermore, the emotion decoding results in Figure <ref> show that our proposed model, HEmoN, consistently outperforms advanced models such as BrainNetTF and GED, achieving outstanding performance in all cases.
Competitive baselines, including BrainNetCNN, BrainGNN, and BrainNetTF, utilize convolutional neural networks, graph neural networks, and transformers to analyze brain networks, respectively.
However, they all neglect the information fusion among multiple brain regions.
Therefore, the superior performance of HEmoN in emotion decoding tasks highlights the importance of integrating information from multiple brain regions.
In addition, as expected, we observe an obvious performance degradation in each ablation of HEmoN-EA1 and HEmoN-DFT, demonstrating the effectiveness of both the hierarchical emotional areas and the proposed information fusion rule.
Interestingly, HEmoN-EA1 is inferior to HEmoN-DFT in all cases, further showing the potential of the identified hierarchical emotional areas.

Overall, the experimental evaluation of cross-dataset emotion decoding demonstrates the significant potential and strong generalization capability of the identified hierarchical emotional areas.
However, it is worth noting the individual differences in decoding performance among various subjects.
Investigating and addressing these individual differences constitutes an important topic for future research.



§ CONCLUSION


In this paper, we investigate the hierarchical emotional areas in the human brain through multi-source information fusion and graph machine learning methods.
Furthermore, we provide preliminary insights into the brain mechanisms underlying specific emotions based on the psychological constructionist hypothesis.
To gain a deeper understanding of the roles of emotional areas at various levels in emotion encoding, it is promising to conduct meta-analyses using a large number of existing literature in future research.






§ DECLARATION OF COMPETING INTEREST


The authors declare that they have no known competing financial interests or personal relationships that could have appeared to influence the work reported in this paper.



§ DATA AVAILABILITY


The data are publicly available on the Open Science Framework (OSF) at <https://osf.io/tzpdf>.



§ ACKNOWLEDGMENTS


The authors would like to express their sincere gratitude to Prof. Dan Zhang for his helpful discussions and valuable suggestions, which greatly contributed to improving the overall quality of this paper.
They also extend their thanks to Giada Lettieri, Naoko Koide-Majima, and Shinji Nishimoto for providing the essential data for this research.



equationsection



§ PROOF OF THEOREM <REF>




According to the propagation rule, we have

    h_u^(k) = ∑_w ∈𝒩(u)( a_uw/∑_i ∈𝒩(u)a_ui h_w^(k-1))

where a_uw is the edge weight between nodes u and w, and 𝒩(u) is the set of u's 1-hop neighbors.

After k iterations of expansion as in Eq. (<ref>), h_u^(k) can be represented as

    h_u^(k)   = 1/∑_i ∈𝒩(u)a_ui∑_w ∈𝒩(u)a_uw·1/∑_j ∈𝒩(w)a_wj∑_x ∈𝒩(w)a_wx⋯
       =1/∑_l ∈𝒩(y)a_yl∑_z ∈𝒩(y)a_yz h_z^(0)


Abbreviate the node influence I_T(u, v) as I.
By substituting Eq. (<ref>) into the definition of node influence, we obtain

    I
           = ‖∂ h_u^(k)/∂ h_v^(0)‖
       = ‖∂/∂ h_v^(0)( 1/∑_i ∈𝒩(u)a_ui∑_w ∈𝒩(u)a_uw·1/∑_j ∈𝒩(w)a_wj∑_x ∈𝒩(w)a_wx⋯. 
       = ‖∂/∂ h_v^(0)( .. 1/∑_l ∈𝒩(y)a_yl∑_z ∈𝒩(y)a_yz h_z^(0)) ‖
       = ‖∂/∂ h_v^(0)( 1/∑_i ∈𝒩(u)a_ui a_uv_1·1/∑_j ∈𝒩(v_1)a_v_1j a_v_1v_2⋯.. 
       = ‖∂/∂ h_v^(0)( ... 1/∑_l ∈𝒩(v_k-1)a_v_k-1l a_v_k-1v h_v^(0)) ‖
       = ‖∂ h_v^(0)/∂ h_v^(0)‖·| a_uv_1/∑_i ∈𝒩(u)a_ui·a_v_1v_2/∑_j ∈𝒩(v_1)a_v_1j⋯a_v_k-1v/∑_l ∈𝒩(v_k-1)a_v_k-1l| 
       = | a_uv_1· a_v_1v_2⋯ a_v_k-1v/∑_i ∈𝒩(u)a_ui·∑_j ∈𝒩(v_1)a_v_1j⋯∑_l ∈𝒩(v_k-1)a_v_k-1l|


According to the properties of trees, there exists exactly one path between nodes u and v; we denote this path as P = (u, v_1, v_2, ⋯, v_k-1, v).
Eq. (<ref>) holds because the partial derivative of nodes that are not in P becomes 0, leading to their elimination.
Eq. (<ref>) holds because the Jacobian matrix ∂ h_v^(0) / ∂ h_v^(0) is the identity matrix I_c_0, and its induced norm ‖ I_c_0‖ equals 1.




§ PROOF OF THEOREM <REF>




Without loss of generality, we consider an arbitrary path P of length m with vertex sequence (v_0, v_1, …, v_m).
According to Eq. (<ref>), the node influence I_P(v_i, v_j) of any pair of nodes v_i and v_j with a distance of k in path P is

    I_P(v_i, v_j) = 1/2^k-1, s.t. |i - j| = k

In addition, there are a total of (m-k+1) pairs of nodes with a distance of k (1 ≤ k ≤ m) in path P.
By using the definition of path information and Eq. (<ref>), we obtain

    I(P)
           = ∑_i,j=0 
     i ≠ j^m I_P(v_i, v_j)
       = 2 ·∑_k=1^m ∑_i=0^m-kI_P(v_i, v_i+k)
       = ∑_k=1^m (m-k+1) ·2/2^k-1

Thus, finding a path with the maximum path information in tree T is equivalent to solving the following optimization problem:

    max_m      I(P) = ∑_k=1^m (m-k+1) ·2/2^k-1
    s.t.      1 ≤ m ≤𝔡
        m ∈ℤ^+

where 𝔡 is the diameter of tree T.

Consider a finite sequence (x_m)_m=1^𝔡 that corresponds to the objective function Eq. (<ref>), where

    x_m = ∑_k=1^m (m-k+1) ·2/2^k-1

Since

    x_m+1 - x_m = ∑_k=1^m+12/2^k-1 > 0, ∀ m ∈ℤ^+

We can conclude that sequence (x_m)_m=1^𝔡 is strictly increasing.
Accordingly, the optimization problem described in Eq. (<ref>) is equivalent to the following one:

    max_m      m 
    s.t.      1 ≤ m ≤𝔡
        m ∈ℤ^+

Clearly, the optimal solution is m̂ = 𝔡, which corresponds to the longest shortest path P̂.
In conclusion, the longest shortest path P̂ has the maximum path information I(P̂) = ∑_k=1^𝔡(𝔡-k+1) ·2/2^k-1.




§ ILLUSTRATION OF THE HEMON MODEL





In this section, we illustrate the process of identifying hierarchical trunks and the representation learning process of our proposed HEmoN model through a concrete example, as shown in Figure <ref>.
Let T = (𝒱, ℛ) denote the brain tree in Figure <ref>, where

    𝒱 = {.
            v_0, v_1, …, v_6, u_0, u_1, u_2 .}, 
    ℛ = {.
            (v_0, v_1), (v_1, v_2), …, (v_5, v_6), 
        (u_1, u_2), (v_2, u_0), (v_3, u_1) .}.


According to the methodology described in Section <ref>, we designate the longest shortest path as the main trunk 𝔱_1.
As a result, we obtain the set of 1st-level trunk 𝒯_1 = {𝔱_1}.
Then, we remove 𝒯_1 from T, and the remaining part becomes a forest F_1.
This forest contains two connected components: connected component 1 and connected component 2, with each component being a (sub)tree.
To identify the trunks at the second level, we search for the longest shortest path in each connected component of the forest.
Accordingly, we obtain the 2nd-level trunk 𝔱_2,1 in connected component 1, the 2nd-level trunk 𝔱_2,2 in connected component 2, and thus the set of 2nd-level trunks 𝒯_2 = {𝔱_2,1, 𝔱_2,2}.

Since there are no remaining nodes after removing 𝒯_2 from F_1, we stop the process of identifying hierarchical trunks and decompose T into two levels of trunks, i.e., 𝒯_1 = {𝔱_1} and 𝒯_2 = {𝔱_2,1, 𝔱_2,2}.

After obtaining all the trunks at each level, we learn the trunk representation h_T^(ℓ) for each level ℓ, which also serves as the representation of the emotional area at the ℓth level.
According to Eq. (<ref>), we have

    h_T^(1)   = LSTM(  x_v_0^(𝔱_1),  x_v_1^(𝔱_1), ⋯,  x_v_6^(𝔱_1)) 
     h_T^(2)   = LSTM(  x_u_0^(𝔱_2,1),  x_v_2^(𝔱_2,1)) + 
       =LSTM(  x_v_3^(𝔱_2,2),  x_u_1^(𝔱_2,2),  x_u_2^(𝔱_2,2))


Finally, according to Eq. (<ref>), we combine both the trunk representations to create a representation of the brain tree T:

    h_T = ∑_ℓ=1^2  W^(ℓ) h_T^(ℓ) =  W^(1) h_T^(1) +  W^(2) h_T^(2).



model1-num-names






 Paper Title: Certifying Robustness of Learning-Based Keypoint Detection and Pose Estimation Methods.















xushengl@andrew.cmu.edu
0000-0003-4342-7234


twei2@andrew.cmu.edu
0000-0003-2505-4585


siminliu@andrew.cmu.edu
0009-0009-1495-6920


ziweiwa2@andrew.cmu.edu
0000-0001-9225-8495


  Carnegie Mellon University
  5000 Forbes Avenue
  Pittsburgh
  Pennsylvania
  USA
  15213



Luis.E.Mattei-Mendez@boeing.com

taylor.s.loper@boeing.com

joshua.neighbor@boeing.com

  The Boeing Company
  PO Box 3707
  Seattle
  Washington
  USA
  98214-2207



fhutchin@nrec.ri.cmu.edu

cliu6@andrew.cmu.edu
0000-0002-3767-5517


  Carnegie Mellon University
  5000 Forbes Avenue
  Pittsburgh
  Pennsylvania
  USA
  15213













§ ABSTRACT 
  This work addresses the certification of the local robustness of vision-based two-stage 6D object pose estimation. The two-stage method for object pose estimation achieves superior accuracy over the single-stage approach by first employing deep neural network-driven keypoint regression and then applying a Perspective-n-Point (PnP) technique.
    Despite advancements, the certification of these methods' robustness, especially in safety-critical scenarios, remains scarce. This research aims to fill this gap with a focus on their local robustness on the system level—the capacity to maintain robust estimations amidst semantic input perturbations. The core idea is to transform the certification of local robustness into a process of neural network verification for classification tasks. The challenge is to develop model, input, and output specifications that align with off-the-shelf  verification tools. To facilitate verification, we modify the keypoint detection model by substituting nonlinear operations with those more amenable to the verification processes. Instead of merely injecting random noise into images, as is common, we employ a convex hull representation of images as input specifications to more accurately depict semantic perturbations. Furthermore, by conducting a sensitivity analysis, we propagate the robustness criteria from pose estimation to keypoint accuracy, and then formulating an optimal error threshold allocation problem that allows for the setting of a maximally permissible keypoint deviation thresholds. Viewing each pixel as an individual class, these thresholds result in linear, classification-akin output specifications. Under certain conditions, we demonstrate that the main components of our certification framework are both sound and complete, and validate its effects through extensive evaluations on realistic perturbations. To our knowledge, this is the first study to certify the robustness of large-scale, keypoint-based pose estimation given images in real-world scenarios. 







<ccs2012>
   <concept>
       <concept_id>10010520.10010553</concept_id>
       <concept_desc>Computer systems organization Embedded and cyber-physical systems</concept_desc>
       <concept_significance>500</concept_significance>
       </concept>
   <concept>
       <concept_id>10010147.10010178.10010224.10010245</concept_id>
       <concept_desc>Computing methodologies Computer vision problems</concept_desc>
       <concept_significance>300</concept_significance>
       </concept>
 </ccs2012>


[500]Computer systems organization Embedded and cyber-physical systems
[300]Computing methodologies Computer vision problems





Certifying Robustness of Learning-Based Keypoint Detection and Pose Estimation Methods
    Changliu Liu
    
======================================================================================





§ INTRODUCTION

In the realm of computer vision, vision-based 6D object pose estimation, i.e., 3D rotation and 3D translation of an object with respect to the camera,  serves as a pivotal method for identifying, monitoring, and interpreting the posture and movements of objects through images <cit.>

This technology is fundamental in granting machines the ability to comprehend the physical environment, finding its utility in diverse domains such as robotics <cit.>, augmented reality <cit.>, and human-computer interaction <cit.>. The evolution of deep learning and the adoption of neural networks, particularly convolutional neural networks (CNNs), have markedly surpassed traditional techniques that depend on manually engineered features. Within the spectrum of learning-based approaches, a distinct classification exists: single-stage methods directly estimate the 6D pose from an image <cit.>. Conversely, a more widespread and accurate category of methods employs a two-stage strategy, initially regressing sparse keypoints <cit.> or dense pixels <cit.> from the image, followed by the utilization of a Perspective-n-Point (PnP)-based strategy for pose estimation through established 3D-2D point correspondences.

Despite the increasing efforts to boost the empirical robustness of these methods against challenges like occlusions, fluctuating lighting, and varied backgrounds, the focus on validating or certifying the reliability of vision-based pose estimation systems remains minimal. The absence of performance assurances for these frameworks raises concerns about their integration into safety-critical applications. In this work, our objective is to certify the robustness of learning-based keypoint detection and pose estimation approaches given input images. We focus on the aspect of local robustness, which refers to the ability to maintain consistent performance or predictions when the input data is perturbed around a given input point. The core question is determining whether pose estimation stays within an acceptable range when the input image undergoes perturbations. To the best of our knowledge, this study is the first one to certify the robustness of large-scale, keypoint-based pose estimation problem encountered in the real world.

Given the crucial role of neural networks in learning-based visual pose estimation, certifying their robustness is inherently linked to neural network verification <cit.>. This area of verification has attracted significant attention in recent years, driven by the paradox of widespread adoption of neural network solutions without adequate assurance of their reliability, primarily due to their opaque nature <cit.>. What distinguishes our problem from existing neural network (NN) verification efforts is our focus on system-level properties rather than on verifying attributes of isolated neural networks, which typically relate directly to the NN outputs <cit.>. For example, in robustness verification of classification models, the objective is to ensure the predicted class remains unchanged despite variations in input. Our research, however, targets the pose estimation framework, wherein the NN constitutes only one component of the entire system. The verification encompasses system-wide requirements, necessitating not just the evaluation of the keypoint detection model's robustness but also that of the PnP-based method.





 §.§ Overview of the Approach

To certify the robustness of the two-stage keypoint-based pose estimation framework, our main idea is to convert the local robustness certification of pose estimation into a standard neural network verification problem for classification networks. The primary challenge involves crafting three verification components—model, input, and output specifications—in a manner that is compatible with existing verification tools. A graphical overview is provided in Fig. <ref>. 


  
Model modification for verification. Considering the keypoint detection model involves complex nonlinear operations such as the softmax function, a variant of the model that is more amenable to verification is created, and an analysis is conducted to understand the properties that are maintained between the original and this modified model. 


  
Input specification through convex hulls. To account for realistic semantic perturbations in input images, we define the input space as a convex hull of possible perturbed images, which captures variations in a mathematically rigorous manner, allowing for a linear representation of input perturbations. Such a specification outperforms existing methods that simply introduce random independent noise into images.


  
Output specification via sensitivity analysis. The core of connecting system requirements with the neural network's output lies in conducting a sensitivity analysis of the downstream PnP method. By understanding how variations in detected keypoints affect the estimated pose, it's possible to translate system-level pose accuracy requirements into error thresholds for keypoint detection. By treating each pixel as a separate class, these thresholds are then used to define classification-like linear output specifications.



 §.§ Contributions

Our contributions can be summarized as follows:

    
  * We propose a local robustness certification framework for the learning-based keypoint detection and pose estimation pipeline;
    
  * We analyse the soundness and completeness properties of this certification framework;
    
  * We demonstrate the method's efficacy through validation on a real-world scale keypoint-based pose estimation problem.























§ RELATED WORK




 §.§ Formal Verification of Neural Networks

The objective of verifying neural networks involves ensuring they meet certain standards of safety, security, accuracy, or robustness. This essentially means determining the truth of a specific claim about the outputs of a network based on its inputs. In recent years, there has been a significant influx of research in this area. For comprehensive insights into neural network verification, one can refer to <cit.>. Verification techniques are generally divided into three main groups: reachability-based approaches, which perform a layer-by-layer analysis to assess network output range <cit.>; optimization methods, which seek to disprove the assertion <cit.>; and search-based strategies which combine with reachability analysis or optimization to identify instances that contradict the assertion <cit.>. In 2020, VNN-COMP <cit.> launched as a competition to evaluate the capabilities of advanced verification tools spanning a variety of tasks, including collision detection, image classification, dataset indexing, and image generation. However, these methods treat deep neural networks in isolation, concentrating on analyzing the input-output relationship.

Concurrently, there is research focused on the system-level safety of cyber-physical systems (CPS) incorporating neural network components, particularly within the system and controls domain. They broadly fall into two categories. The first category <cit.> focuses on ensuring the correctness of neural network-based controllers, taking their input from the structured outcomes of the state estimation module, regardless of the state estimation module is based on perception or not. Neural network controllers of this type generally consist of several fully connected layers, making them relatively straightforward to verify. The second category focuses on validating the closed-loop performance of vision-based dynamic systems that incorporate learning-based components. Among these, studies such as <cit.> examine LiDARs as the perception module, processed by multi-linear perceptrons (MLPs) with a few hidden layers. Other approaches, primarily applied to runway landing and lane tracking, deal with high-dimensional inputs from camera images, employing methods like approximate abstraction of the perception model <cit.>, contract synthesis <cit.>, simplified networks within the perception model <cit.>, or a domain-specific model of the image formation process <cit.>. 
Nevertheless, studies directly dealing with high-dimensional inputs from camera images are still limited due to the images' high dimensionality and unstructured data nature, in contrast to structured robot states such as position and velocity.





 §.§ Certification of Keypoint Detection and Pose Estimation Methods

The investigation of certification methods for pose estimation is relatively limited. <cit.> introduced a certifiable approach to keypoint-based pose estimation from point clouds by correcting keypoints identified by the model, ensuring the correctness guarantee of the pose estimation. <cit.> expanded on this by integrating the correction concept with ensemble self-training. In a similar vein, By propagating the uncertainty in the keypoints to the object pose, <cit.> created a keypoint-based pose estimator for point clouds that is provably correct and is characterized by definitive worst-case error bounds. The above work focuses on point clouds as opposed to images. <cit.> applied an advanced deep learning technique for uncertainty quantification to assess the uncertainty (i.e., the predicted distribution of a pose) in multi-stage 6D object pose estimation methods. In contrast, <cit.> aimed to design a neural network that processes camera images to directly predict the aircraft's position relative to the runway with certifiable error bounds. Of all these studies, <cit.> is the most similar to our work, which focuses on the verification of keypoint detection, excluding the examination of the PnP method for system-wide assurances. Their approach verifies the neural network in isolation and is limited to very slight perturbations, failing to encompass realistic semantic variations.




§ BACKGROUND

In this work, we represent scalars and scalar functions by italicized lowercase letters (x), vectors and vector functions by upright bold lowercase letters (), matrices and matrix functions by upright bold uppercase letters (), and sets and set functions with calligraphic uppercase letters ().



 §.§ Keypoint-based Pose Estimation

The keypoint-based approach consists of two steps to estimate the 6D pose from a 2D model image. First, a neural network is tasked with predicting the 2D locations of keypoints, whose 3D locations are predefined within the object model. A common strategy involves the use of heatmap regression, wherein ground-truth heatmaps are created by placing 2D Gaussian kernels atop each keypoint. The heatmap pixel values are interpreted as the likelihood of each pixel being a keypoint. These heatmaps are then used to guide the training through an ℓ_2 loss. The detection network can be divided into two parts. A backbone network, denoted by b, inputs a 2D image to produce unnormalized heatmaps, which is first followed by a softmax layer that transforms unnormalized heatmaps into normalized ones, and then by argmax operations, or another layer of differentiable spatial to numerical transformation (DSNT) <cit.>  for keypoint extraction. We refer to the part after the softmax (including) as the head network. The entire network is represented by

=  () = h∘b(), 

where ∈^H× W× C represents a 2D RGB image with dimensions being H× W× C, and ∈^K× 2 denotes the 2D coordinates of K keypoints. Here, ∘ denotes function composition. To enhance accuracy and robustness, it's often essential to preprocess the input image  before it is passed to the network, such as resizing and color normalization. Denote this preprocessing step by _0, leading to the equation =  () = h∘b∘_0().  In what follows, we omit the preprocessing step unless it is critical to consider it.



The second step employs the Perspective-n-Point (PnP) algorithm, which executes a nonlinear least squares (NLS) optimization to estimate the 6D pose from established 3D-to-2D correspondences. An illustration of the pose estimation  for an airplane is presented in Fig. <ref>. Let ∈^3×3 denote the camera intrinsic parameter matrix, _k ∈^3 denote the 3D coordinate of keypoint k, where k = 1, …, K, and _k ∈^2 denote the corresponding 2D coordinate. These 3D-2D correspondences are formed through the following perspective projection model:

    λ_k [ _k;  1 ]
       =  (_k + ),

where ∈^3×3 and ∈^3 represent the rotation matrix and translation vector, respectively, that establish the transformation between the object and camera coordinate systems, with λ_k representing the scaling factor. The objective of the PnP method is to approximate this transformation, represented by  and , by minimizing the ℓ_2 norm of the reconstruction error:

    ⟨, ⟩  =   ,       ∑_k=1^K (_k + ) -   λ_k [ _k;  1 ]^2_2.


Let the notation ⟨, ⟩ = (, ) represent the PnP procedure, where ∈^K× 3 denotes the 3D coordinates of keypoints. Note that the keypoint-based pose estimation framework described above is the basic version. Numerous adaptations have been developed to increase accuracy and robustness <cit.>, particularly in challenging conditions such as occlusions, varying viewpoints, and different lighting scenarios. This paper concentrates on validating the robustness of this foundational pipeline, marking a crucial step towards certifying the effectiveness of more intricate keypoint-based pose estimation techniques. Hereafter, we use Ψ to represent this pipeline, which includes keypoint detection followed by the application of a PnP method, that is, ⟨, ⟩ = Ψ() = (, ()).








 §.§ Verification of Neural Networks

Consider a multi-layer neural network representing a function , which takes an input ∈_⊆^d_0 and produces an output ∈_⊆^d_n, where d_0 is the input dimension, and d_n is the output dimension. Any non-vector inputs or outputs are restructured into vector form. The verification process entails assessing the validity of the following input-output relationships defined by the function :
∈⇒ = () ∈,
where sets, ⊆_ and ⊆_, are referred to as input and output constraints, respectively. 

In the context of confirming the robustness of a classification network, the goal is to ascertain that all samples within a proximal vicinity of a specified input x_0 receive an identical classification label. Assuming the target label is i^∗∈{1, …, d_n}, the specification for verification is that y_i^∗ > y_j for every j not equal to i^∗. The constraints on inputs and outputs are established accordingly:
= {| - _0_p ≤ϵ}, 
 = {| y_i^∗ > y_j,  ∀ j ≠ i^∗},
where ϵ represents the maximum permissible deviation in the input space. The metric used to quantify disturbance can be any ℓ_p norm.

Neural network verification algorithms can generally be categorized into three main types: reachability analysis, optimization, and search. NN verification essentially seeks to transform the nonlinear model checking problem into piece-wise linear satisfiability problems, and it can be applied to various nonlinearities, including ReLU and, more recently, softmax <cit.>. Two pivotal attributes—soundness and completeness—are of critical importance. A verification algorithm is sound if it only confirms the validity of a property when the property is indeed valid. It is complete if it consistently recognizes and asserts the existence of a property whenever it is actually present. There is a trade-off between computational complexity and conservativeness (or in-completeness).






§ PROBLEM FORMULATION


The most common type of input specification involves limiting the ℓ_p-norm of the variation to a threshold, that is, - _0 _p≤ϵ. However, ℓ_p perturbation,  often characterized by a small value ϵ,  is not a correct mathematical description of more realistic perturbations, as the independent nature of pixel-wise perturbations falls short in creating perturbations that reflect semantic correlations between pixels, such as large variations in lighting, weather conditions, and the effects of camera motion blur. Another approach for input perturbations is to directly add a generative model that perturbs the input image before the original neural network, such as <cit.>, and then verify them together. However, the verification result can be biased by the generative model used and it is not user friendly due to the difficulty in controlling the changes made to the image. To address these shortcomings, we adopt a strategy based on the convex hull, which involves combining a seed image with a collection of perturbed images. These perturbed images can be created through different methods including simulators and learning-based generative models. Also, the convex hull specification can directly enable users to specify the perturbed images, which makes the perturbation specifications more user friendly.


Given a seed image _0 and a set of n perturbed images {_1, …, _n}, the convex hull constituted by these images is defined by the set of all their possible convex combinations. Mathematically,

    = { |  = ∑_i = 0^n α_i _i,   s.t. α_i ≥ 0, ∑_i=0^n α_i  = 1.}.


Images within the convex hull ∈ result from varying degrees of continuous blending among the provided images.  Through the convex combination of perturbation instances, we can model environmental and sensor-related perturbations, including changes in brightness, contrast, weather conditions, motion blur, and dust on lens. Convex hull perfectly captures the entire range of brightness or contrast shifts, and does a reasonable approximation of color shifts, provided they are small. However, this approach does not capture translational perturbations of the object or perturbations related to camera movements, such as changes in the viewpoint. Preliminary methods exist that characterize pixel-wise perturbations caused by camera motion <cit.>, but these methods cover only a very limited range. Investigating robustness against such perturbations will be addressed in future research.  A collection of perturbed images and sampled images from the convex hull can be found in Fig. <ref> in Section <ref>. 

Given two rotation matrices  and , we define the function _r (, ) = |() - () |, where (·) represents the function that converts rotations to Euler angles, and |·| denotes the element-wise absolute value. Thus, _r measures the rotational difference along each axis. Similarly, define _t (, ) =  | - |, which calculates the translational difference per axis.





 Given a convex hull representation  consisting of a seed image _0 and n perturbed images, and a keypoint-based pose estimation framework Ψ.
 
 
 
 
 The problem is to certify whether the pose estimation framework Ψ is robust to any image within the set . Mathematically,

    [ _r (,  );  _t (, ) ]≤[ ϵ_r; ϵ_t ]  s.t. ⟨, ⟩ = Ψ(),  ∀ ∈,

where ⟨, ⟩ =  (, ) represents the pose estimated given the ground-truth 3D keypoint coordinates  and ground-truth 2D keypoint coordinates   for _0, and ϵ_r ∈^3 and ϵ_t ∈^3 denote the error thresholds specified by the user for rotation and translation. We refer to  and  as the nominal transform, and  and  as the perturbed transform. In simpler terms to describe Eq. (<ref>), the variation between the nominal transform and the perturbed transform is kept within these predefined thresholds.



If a pre-processing step, such as resizing or color normalization, is applied to an image before it is input into the keypoint detection model, then the convex hull should be constructed using images that have also undergone these pre-processing steps. Consequently, any image within the convex hull is represented as = ∑_i=0^n α_i _0(_i), where _0 denotes the pre-processing function. 




 The vector of error thresholds ϵ_r and ϵ_t can specify requirements from system designers in an interpretable manner. For example, such requirement can be setting the error thresholds for translation to be within 1 meter. The difficulty presented by Problem <ref> is that, the error thresholds ϵ are defined with respect to pose estimation rather than the direct outputs of neural networks—keypoints—and there are nonlinear mappings between keypoints and pose estimation. 
 
To address Problem <ref>, the core of our approach is to  transform the local robustness certification problem into a standard neural network verification problem for classification networks. This involves adapting three verification components—model, input, and output specifications—to be compatible with existing verification tools. As detailed in Section <ref>, we have modified the keypoint detection model to handle complex operations like the softmax function, simplifying the verification process. The input space is the a convex hull of possible perturbed images, which captures semantic variations  compared to traditional methods that introduce random noise. Section <ref> outlines how the output specification is established through a sensitivity analysis of the PnP method, detailed further in Section <ref>. This analysis helps translate system-level pose accuracy requirements into error thresholds for keypoint detection. Finally, Section <ref> is dedicated to analyzing the soundness and completeness of our proposed verification framework.





§ FORMULATION OF VERIFICATION OF THE NEURAL NETWORK


To convert Problem <ref> into an equivalent problem that can be verified using existing NN verification tools, we introduce approaches to modify the NN model and specify the output constraints.



 §.§ Verification-Friendly Keypoint Detection Model


As introduced in Section <ref>, a softmax layer follows the backbone network to generate probabilistic heatmaps. This step introduces verification challenges due to the highly nonlinear exponential function in softmax. To address this, we propose to maintain the backbone network but replace the head network with an alternative one. A critical insight is that during the inference phase, the unnormalized heatmaps output by the backbone network b already provide significant clues about the locations of keypoints. The subsequent operations within the head network h, such as softmax and DSNT, essentially serve to aggregate this information across the entire heatmap. This aggregation allows for the generation of predictions in an average manner. In essence, the peaks within the heatmaps play a pivotal role in these predictions, and their locations remain unchanged between unnormalized and normalized heatmaps. Consequently, by focusing on the characteristics of these peaks, we can bypass some of the complexities introduced by the head network's nonlinear operations while still capturing the essential information required for accurate keypoint detection. 


By appending an average pooling layer followed by an argmax layer as the new head network, referred to as h', to the backbone network, we create a proxy model, denoted by proxy = h' ∘b. We refer to the original model as the target model. The introduction of this proxy model simplifies the transformation of the verification problem into one akin to classification, where each pixel is treated as a unique category, as illustrated in Fig. <ref>. The average pooling layer offers two advantages. Firstly, its downsampling effect significantly reduces the number of categories (pixels), thereby simplifying the complexity of verification. Secondly, by aggregating local features within each pooling region, it ensures a more accurate representation than applying argmax directly. 







The selection of pooling parameters should be guided by the error threshold on pose estimation ϵ_r and ϵ_t, as elaborated in Section <ref>. This leads to a scenario where the pooling parameters assigned to each heatmap, corresponding to individual keypoints, vary. This variation arises because keypoints have differing levels of influence on the accuracy of pose estimation. Consequently, keypoints with lesser influence on pose estimation are allocated a larger permissible error range, whereas those that play a critical role in pose estimation are subjected to stricter error ranges. It is worth noting that, with certain assumptions, the target model and proxy model are equivalent in terms of verification. The analysis on how assurances regarding the proxy model's performance can be extended to the target model is conducted in Section <ref>.


































 
 


















 








































 
 















 



 §.§ Output Specification: Polytope Representation

For the coordinate _k of the k-th keypoint, where _k represents its 2D coordinate, let _k denote its averaged ground-truth coordinate following the application of the average pooling and argmax layers. To ensure the classification result is consistent across perturbations, the output specification necessitates that the value at v̅_k exceeds the values at all other entries. This requirement can be encapsulated by a half-space-based polytope, represented as _k _k ≤_k with



    _k = 
    [  1  0  … -1  …  0;  0  1  … -1  …  0;  ⋮  ⋮  ⋱  ⋮  ⋮  ⋮;  0  0  … -1  …  0;  ⋮  ⋮  ⋮  ⋮  ⋱  ⋮;  0  0  … -1  …  1;    ]
    ,   _k = 
    [    y_1;    y_2;      ⋮; y_v̅_k;      ⋮;    y_n ]
    ,   _k = 
    [ 0; 0; ⋮; M; ⋮; 0 ],

where the v̅_k-th column of _k is populated with (-1)'s, and all diagonal elements set to 1 with the exception of the element at (v̅_k, v̅_k), the components of _k are set to zero except for the v̅_k-th element, which is assigned a significantly large integer M. Consequently, the specification for all keypoints is denoted by ≤ with

    𝐀 = 
    [ 𝐀_1   0   0   …   0;   0 𝐀_2   0   …   0;   0   0 𝐀_3   …   0;   ⋮   ⋮   ⋮   ⋱   ⋮;   0   0   0   … 𝐀_K ]
    ,    = 
    [ _1; _2; _3;  ⋮; _K ]
    ,    = 
    [ _1; _2; _3;  ⋮; _K ].

The construction of Eq. (<ref>) presupposes verifying all keypoints. However, should there be a keypoint that does not require verification, potentially due to being an outlier, the matrices associated with it are then adjusted accordingly: _k = 𝐈 and _k = [M, …, M]^T.





























§ ERROR PROPAGATION AND DETERMINATION OF POOLING PARAMETERS

To calculate the parameters for average pooling, the initial step involves mapping the error thresholds ϵ_r and ϵ_t from the pose estimation onto the keypoints' error thresholds that may be correlated to each other, which is accomplished by employing sensitivity analysis methods.

Subsequently, we formulate an optimal error threshold allocation problem, aiming to assign independent error thresholds to each keypoint. This step is crucial for determining the allowable errors for each keypoint individually, ensuring that the overall pose estimation adheres to predefined error limits. The idea is graphically depicted in Fig. <ref>.










 §.§ Error Propagation via Sensitivity Analysis

Sensitivity analysis for nonlinear optimization explores the impact of first-order changes in the optimization parameters on the locally optimal solution <cit.>. It involves analyzing an unconstrained parameterized nonlinear optimization problem expressed as min_ F(; ), with  representing the parameter vector. Here, let z represent the optimal objective value given the parameter vector , i.e.,
z () = min_ F(; ).
Sensitivity analysis <cit.> connects the first-order derivatives ∂ and δ z with ∂:

    [ -F_; -F_ ]∂  
        = [ F_ -1; F_  0 ][   ∂; ∂ z ].

In the context of pose estimation, the keypoint coordinates  are considered as parameters and the pose to be estimated represents the decision variables, the notation for partial derivatives ∂ is substituted with the notation for deivation δ. Consequently, Eq. (<ref>) is transformed into

    [ -G_; -G_ ]δ
        = 
        [ G_ -1; G_  0 ][   δ; δ z ],

where ∈^6 = [, ] represents the 6D pose vector, with the 3D rotation matrix  expressed in the axis-angle or Euler angles representations form , and G corresponds to the NLS optimization objective outlined in Eq. (<ref>). Note that δ∈^2K. The following notation is introduced

    _ :=  [ -G_; -G_ ]∈^7× 2K,  _ :=  [ G_ -1; G_  0 ]∈^7×7.


If _ is invertible, then it follows that
_^-1_δ =  [   δ, δ z]^T.
    
    
    
By extracting the rows in the matrix _^-1_ that correspond to δ and naming the resulting matrix _, we obtain the linear mapping 

_δ = δ.

This leads to the establishment of a linear relationship between  2D keypoints errors and 6D pose errors. Upon specification of the keypoints  and the transform ξ, the matrix _ becomes determined. 

Let Ξ represent the set of tolerable errors, indicating the range within which deviations of the perturbed transform from the nominal transform are acceptable, meaning δ∈Ξ. Let δ denote the derived tolerable keypoint  errors that has the following property:

 δ =  {δ | _δ   =  δ∈Ξ}.

Note that δ is an approximation of actual tolerable keypoint errors that lead to tolerable pose errors δξ. As the pose error bounds decrease, the accuracy of the PnP linearization improves, narrowing this disparity. In what follows, we assume that Ξ is defined as a polytope represented by linear inequalities, specifically Ξ= {δ | _ξδ≤_}. This inequality can be established based on a user-defined thresholds ϵ_r and ϵ_t as outlined in Problem <ref>. By the linear mapping, the set δ can also be characterized as a polytope:
    δ = {δ | _ξ_δ≤_} = {δ | _δ≤_},
where _ =  __ and _ = _.  While larger set  δ is desirable, it is an approximation of actual tolerable errors due to the linearization of sensitivity analysis. To mitigate the risk of over-approximation, we  introduce a scaling factor to reduce the size of  set δ in the next section.



 §.§ Optimal Error Threshold Allocation

The tolerable errors of keypoints are interrelated within the derived polytope set δ. In this section, we allocate the tolerable errors across each keypoint to ensure their tolerances are independent, as current verification tools lack the capability to verify dependencies between  keypoints. The allocated error threshold  aims to be maximized, provided that the pose deviation remains tolerable.

For a given error vector δ, we can uniquely define an axis-aligned, axis-symmetric hyper-rectangle, denoted as (δ) ⊂^2K, centered at the origin with each element δ_i representing the half-length of its sides. Axis symmetry indicates that the error thresholds are identical in opposite directions. We establish the following problem for optimal error threshold allocation:

    max_δ∈_+^2K      w_1 ∏_k=1^2Kδ_k + w_2 Δ
    s.t.     (δ)  ⊆δ, 
         δ_k≥Δ,     k = 1, …, 2K.


Constraint (<ref>) ensures that the hyper-rectangle is encompassed within the set of tolerable errors δ, while constraint (<ref>) sets a minimum threshold for the side lengths. The goal defined in objective (<ref>) is to optimize a weighted sum of two objectives: the first seeks to identify the largest possible axis-aligned, axis-symmetric hyper-rectangle within a convex polytope, and the second strives to extend the minimal side length as much as feasible.

To address the optimal error threshold allocation problem, we translate constraint (<ref>) into a series of linear inequalities, drawing from the method in <cit.> for identifying the largest axis-aligned hyper-rectangle inscribed within a convex polytope. Considering a hyper-rectangle = {∈^d |≤≤} and a convex polytope = {|≤}, <cit.> demonstrates that ⊆ if and only if

    ∑_j=1^d (_ij^+ _j - _ij^- _j) ≤_i,    i = 1, …, n.

where _ij^+ = max{_ij,0} and _ij^- = max{-_ij,0}. Given the axis symmetry in our scenario, it follows that _j = -_j. Substituting _j with δ_j,  with _, and  with _, we can reformulate Ineq. (<ref>) as

    ∑_j=1^2Kδ_j (_ ij^+  + _ ij^-) ≤_ i,    i = 1, …, 6.

Given that |_ ij| = _ ij^+ + _ ij^- where |·| returns element-wise absolute values, it follows that

    ∑_j=1^2Kδ_j |_ ij|≤_ i,    i = 1, …, 6.

In matrix notation, Ineq. (<ref>) is expressed as
|_| δ≤_.
Ultimately, the problem of optimal error threshold allocation, encompassing objectives (<ref>) through constraints (<ref>), is equivalent to

    max_δ∈_+^2K      w_1 ∏_k=1^2Kδ_k + w_2 Δ
    s.t.     κ  |_| δ≤_, 
         δ_k≥Δ,    k = 1, …, 2K,


where _+ represents the set of non-negative integers, and κ∈_+ is the scaling factor that controls the actual size of polytope. The effect of the scaling factor κ on the allocation is discussed in Section <ref>. The solution, denoted by δ^*, returns the maximum tolerable errors for all keypoints while ensuring that the pose errors remains within the error thresholds. This solution will determine the parameters for average pooling in Section <ref>.






































































































































 §.§ Determination of Average Pooling Parameters



A widely accepted practice to train the keypoint detection model is to regress the heatmaps centered on the ground truth keypoint. Building on this approach, we propose the following assumption and support it with evaluation results detailed in Section <ref>.


   The softmax layer in the target model _target produces a distribution with a peak at its center and exhibits symmetry along axes that intersect this peak.


Upon determining the independent error threshold for each keypoint, we proceed to define the parameters for average pooling in the proxy model _proxy. Focusing on a particular keypoint with coordinates = (_h, _v), and considering the error thresholds δ^*_h and δ^*_v allocated for horizontal and vertical directions respectively, we define the set 

    _δ = { (h, v) | | h - _h| ≤δ^*_h  | v - _v| ≤δ^*_v }.

Only pixels within _δ are considered acceptable potential predictions for the keypoint. The average pooling parameters for this keypoint are selected to ensure every pixel in _δ is contained by the same pooling patch. This guarantees that all these pixels are pooled into the same after-pooling pixel as the keypoint, thereby belonging to the same class.

We begin by examining how predictions are generated in the proxy model proxy. Based on Assumption <ref>, the unnormalized heatmap returned by the backbone network _b exhibits symmetry along axes passing through this peak, as softmax operation will not change the overall shape of the unnormalized heatmap. To infer the keypoint prediction, we concentrate on the average pooling patch that predominantly overlaps with the vicinity of the peak. Since the average pooling patches also exhibit axial symmetry, patches closer to the peak yield larger averaged pixel values. Consequently, the pooling patch with its center closest to the peak results in the predicted keypoint. This insight enables us to transform the determination of the location of the maximal pixel value after pooling into identifying the closest pooling patch to the peak of the unnormalized heatmap under Assumption <ref>. Next, we compute the average pooling parameters for each keypoint to ensure that: (i) the ground-truth keypoint aligns with the center of an average pooling patch, and (ii) all pixels in proximity to this pooling patch fall within the allocated error thresholds from the ground truth keypoint, making them suitable to act as the peak of the unnormalized heatmap. Let = (s_h, s_v) denote the stride, = (k_h, k_v) the kernel size, and = (p_h, p_v) the patch size.





  
Determination of stride and kernel parameters. 
The initial step involves determining the stride and kernel parameters, denoted as s_h and k_h, respectively, as in Fig. <ref>. We focus on the horizontal axis (similar logic applies to the independent vertical axis). Let  represent the pooling patch centered at a keypoint, with _-1 and _+1 denoting the adjacent patches to the left and right, respectively. Additionally, let b_-1 and b_+1 denote the perpendicular bisectors between these patches. The distance between these two perpendicular bisectors corresponds to the horizontal stride. To position the ground-truth keypoint at the center, an odd stride is required. As long as the peak of the unnormalized heatmap lies between b_-1 and b_+1, the pooling patch  is the closest to the peak, resulting in the predicted keypoint. Let δ^*_h denote the allocated error threshold. To ensure that peaks falling between b_-1 and b_+1 remain permissible with respect to the error threshold, the threshold should be no less than half the stride distance from the ground-truth keypoint on either side, i.e.,
δ^*_h ≥1/2 (s_h - 1),
which leads to the inequality s_h ≤ 2δ^*_h + 1. Choosing s_h =  2 δ^*_h + 1 minimizes the post-pooling size. It's worth noting that there are no constraints on the kernel parameter. For simplicity, we set the kernel k_h to be equal to the stride s_h.



















  
Determination of the padding parameter.
The next step involves determining the padding parameter to center a pooling patch at the ground-truth keypoint. Padding can be added from the left (horizontal) or from the top (vertical) to achieve this. Considering the horizontal axis where the kernel is equal to the stride, the idea is to identify the closest pooling patch to the keypoint from its right side and then shift it left by the distance between the keypoint and its center, as illustrated in Fig. <ref> and <ref>. Assuming k_h = s_h, let r_h be the number of pixels the keypoint is from the left side of the closest pooling patch, calculated as r_h ≡_h k_h. (i) If r_h ≤k_h + 1/2, the closest patch center to the right of the keypoint, belongs to the pooling patch containing the keypoint. (ii) Otherwise, the closest patch center to the right comes from the pooling patch immediately to the right of the one containing the keypoint. The horizontal padding parameter p_h is determined as k_h+1/2 - r_h if r_h ≤k_h+1/2, otherwise, as 3k_h+1/2 - r_h.









Note that the shift is not unique because multiples of the stride can be added to p_h.

 
 
 
 
 
 
 
 
 
 
 
 




§ THEORETICAL ANALYSIS

With the method we introduced (to reformulate the problem into a verifiable form using off-the-shelf verification tools), it is important to understand whether false analysis results will be introduced. To answer that question, we perform the following analysis.



 §.§ Properties of the Proxy Model

We examine two target models. The models ending with a DSNT layer and an argmax layer are represented by dsnt and argm, respectively. The term target is used to represent either target model. The notation ^k is employed to identify the k-th predicted keypoint. Recall that _k denote the averaged ground-truth coordinate of the k-th keypoint after the application of average pooing and argmax layers. Before providing theoretical results, we establish formal definitions for both the soundness and completeness related to the proxy model, which are visually represented in Fig. <ref>.



    A proxy model is deemed sound if, for cases where the proxy model's predicted keypoint aligns with the averaged ground-truth keypoint _k, the deviation of the target model's prediction from the ground truth _k does not exceed the error threshold δ_k^*. Specifically, 
    
    proxy^k() = _k ⇒|target^k() - _k |≤δ_k^*,   for  k = 1, …, K.

     Here, |·| indicates the computation of element-wise absolute differences.


    A proxy model is deemed complete if the difference between the target model's prediction and the ground truth  _k is confined within the error threshold δ^*_k, then the predicted keypoint by the proxy model matches the averaged ground truth _k. That is, 
    
    |target^k() - _k |≤δ_k^* ⇒proxy^k() = _k,   for  k = 1,…, K.




If Assumption <ref> holds, the proxy model is sound, that is,

    proxy^k() = _k    ⇒|dsnt^k() - |≤δ_k^*,     for   k = 1, …, K. 
    proxy^k() = _k    ⇒|argm^k() - |≤δ_k^*,     for   k = 1, …, K.



To prove Theorem <ref>, we first prove the equivalence of two target models.
    
      If Assumption <ref> holds, the two target models produce the same predictions, i.e.,

    dsnt^k() =  argm^k(),      for    k = 1, …, K.

    
    




 Before delving into the proof, we first introduce how DSNT extracts a keypoint from a normalized heatmap. Let  represent a normalized heatmap of dimensions m × n, which is generated by a softmax layer. Define  as an m × n matrix where each entry _ij is calculated by _ij = 2j-(n+1)/n. Consequently, each column of  contains identical values. These values fall within the range of (-1, 1), with the center value at 0. To determine the column coordinate _h of the keypoint, DSNT computes the sum of the element-wise product between  and , expressed as _h = ∑_i, j_ij_ij. The row coordinate is obtained in a similar way.

Without loss of generality, due to independence of horizontal and vertical axes, we will show for the horizontal direction. Considering a symmetric heatmap with dimensions m × n, we initially position the center of the distribution at the m+1/2-th row and the n+1/2-th column, effectively placing the distribution's center at the heatmap's midpoint. Upon applying the DSNT process, the resulting keypoint coordinate is (0, 0). This outcome is attributed to the distribution's symmetry along both the horizontal and vertical axes that intersect at its center, with the DSNT operation cancelling out symmetric values relative to the origin. Subsequently, we consider a horizontal displacement of the distribution by c columns, where a positive c indicates a shift to the right. The column coordinate _h, as determined by DSNT following the horizontal shift, is then













    ∑_i, j_ij_ij 1= ∑_i, j'_ij'(_ij' +  2c/n) 2= ∑_i, j'_ij'_ij' + ∑_i, j'_ij'2c/n 3=   0 + 2c/n∑_i, j'_ij' 4= 2c/n


In 1, the pixel located at (i, j') is the corresponding pixel at (i, j) prior to the horizontal shift, with both pixels having identical values, denoted as _ij = _ij'. The value 2c/n represents the  difference across every c columns within the  matrix .
For 3, the expression ∑_i, j'_ij'_ij' equals zero, reflecting the distribution's initial centering at the heatmap's midpoint. In 4, the probability over distribution sums up to 1. The calculation of the DSNT column coordinate as 2c/n∈ (-1, 1) aligns it with the heatmap's (n+1/2+c)-th column. On the other hand, the argmax operation in argm^k returns the  distribution's center, located at the (n+1/2+c)-th column following the shift. Consequently, we establish that dsnt^k() = argm^k(), thereby completing the proof.


According to Proposition <ref>, it is sufficient to validate Eq. (<ref>) that connects proxy with argm to prove Theorem <ref>. The distinction between the proxy and argm models lies solely in their method of handling the unnormalized heatmap, specifically whether they use a softmax or an average pooling layer as depicted in Fig. <ref>. With this, we are poised to validate Theorem <ref>, a conclusion that naturally follows from the way we identify average pooling parameters.

Considering the k-th keypoint _k, if proxy^k() = _k, it implies that average pooling parameters ensure that the closest pooling patch __k to the ground-truth keypoint _k is exactly centered at _k.

 The selection of stride parameters ensures that the heatmap's peak is within a δ^*_k distance from the center of __k, which is also the ground-truth coordinate _k. Given that the softmax layer preserves the peak's position, the prediction argm^k precisely returns the peak's location, leading to the conclusion that argm^k() is within a δ^*_k distance from _k, i.e.,
 |argm^k() - |≤δ_k^*.






If Assumption <ref> holds, the proxy model is complete, that is,

    |dsnt^k() - |≤δ_k^*    ⇒proxy^k() = _k,    for   k = 1, …, K. 
    |argm^k() - |≤δ_k^*     ⇒proxy^k() = _k,    for   k = 1, …, K.
 


Based on Proposition <ref>, we focus exclusively on the target model argm. The selection of pooling parameters guarantees that every pixel within the δ^*_k proximity of the ground-truth keypoint—eligible to be recognized as the predicted keypoint—shares the same closest pooling patch. This closest pooling patch is identical to that of the ground-truth keypoint. As a result, these pixels yield an averaged pixel that is the same as that of the ground-truth keypoints.



 


 §.§ Properties of the Optimal Error Threshold Allocation

In the previous section, we examined the properties of the proxy model after determining the error threshold for keypoints. This section focuses on the properties of the optimal error threshold allocation. We introduce the concepts of probabilistic soundness and probabilistic completeness to measure the relationship between tolerable errors in pose and keypoints, as depicted in Fig. <ref>. 

Considering the set of acceptable pose errors Ξ, let δ_Ξ represent the corresponding set of acceptable keypoint errors. Note that δ_Ξ might not be polytopic, due to the non-linear nature of the PnP method, even though Ξ is polytopic. Let μ(·) denote the measure of a set.
    
    Probabilistic soundness is defined by the ratio μ(δ_Ξ∩)/μ(), which is the fraction of keypoint errors in  that result in tolerable pose errors within Ξ.
    
    Likewise, considering the set of tolerable errors on keypoints , let Ξ_ represent the corresponding set of tolerable errors on poses.
     
     Probabilistic completeness is defined as the ratio μ(Ξ_∩Ξ)/μ(Ξ), which represents the fraction of permissible pose errors in Ξ that result in errors on keypoints within .
    
Deriving an exact or lower bound for these two metrics is difficult, due to the challenges in making definitive conclusions using linearization-based sensitivity analysis, as the accuracy of local linearizations varies with instances, such as varying poses. In Section <ref>, we perform statistical tests to approximate these two values, offering insights into the practicality of our framework.
 
    













































































































































 







§ EVALUATION RESULTS





 §.§ Verified Keypoint Detection Model and Perturbations



  §.§.§ Verified keypoint detection model

  The verified keypoint detection model, illustrated in Fig. <ref>, is modified from the proxy model proxy. This verified model includes three components:


  
1. Backbone model  The dataset contains 7,320 images, each with dimensions of 1920 × 1200. 


    
  * CNN. The architecture includes 5 convolutional layers and an equal number of deconvolutional layers, designed to handle inputs of size 64×64×3. The model comprises 39 layers and contains around 6.57× 10^5 trainable parameters.
    
  * ResNet-18. This model includes 8 residual blocks  and  processes inputs of size 256×256×3. It is composed of 84 layers and possesses approximately 1.2× 10^7 trainable parameters, which is approximately an order of magnitude larger than existing works <cit.>.








 
 





  
2. Average pooling Consider that the parameters for average pooling are tailored based on optimally allocated error thresholds that may vary across different keypoints, thereby necessitating distinct pooling parameters for each. To address this issue, the second component includes the division of a multi-channeled unnormalized heatmap layer into individual channels. Each of these channels is then processed through its own average pooling layer. This design allows for simultaneous verification of all keypoints, eliminating the need to verify each keypoint individually.



  
3. Flattening and concatenation. The third component involves the flattening of the previously splitted channels, which are then concatenated into a single-dimensional format for verification.



  §.§.§ Perturbations

 Out of 7000 images, we randomly sample 200 images as seed images, adding local or global perturbations to them.


  
Local object occlusions
We created a set of 40 realistic semantic disturbances featuring personnel and vehicles typically encountered at an airport. These disturbances are depicted in Fig. <ref>. To create perturbed images, we randomly selected 20 objects as patches, each up to 150 pixels in size, positioned randomly on the seed images, with each perturbed image having one patch. The perturbed images are classified into two groups: overlapping and non-overlapping, based on whether the patches overlap with the airplane in the image. There are 4000 perturbed images in total, including 893 overlapping and 3107 non-overlapping images. The convex hull's complexity is adjusted by changing the number of perturbed images m, where m ranges from 2 to 4. These perturbed images are randomly selected, allowing us to systematically assess how the addition of semantic disturbances affects the robustness and performance of the system. Note that the convex hull is comprised solely of either overlapping or non-overlapping images. A collection of perturbed images and sampled images from the convex hull are presented in Fig. <ref>.



  
Local block occlusions Given a seed image, we generate 4 perturbed images. In each perturbed image, 
a 3 × 3 square is placed either away from the airplane or centered over a randomly selected keypoint to create non-overlapping or overlapping images, respectively. All pixel values within this square are randomly selected from the range [0, 255]. A convex hull is formed using these 5 images per seed image, of which all perturbed images are non-overlapping or overlapping ones.



  
Global perturbations
We change each pixel's value through two types of global perturbations: brightness and contrast. For brightness, a variation value b ∈ℤ is applied such that each pixel's value increases by b, that is, I' = (I + b), where I represents the original pixel values, I' the new pixel values, and  ensures the values remain within the range [0, 255]. For contrast perturbation, a variation value c ∈ℝ adjusts each pixel's value by a percentage c, formulated as I' = (I × (1 + c)).  The convex hull is constructed as follows. For a positive value b ∈ℤ_+, two perturbed images are created for b and -b, respectively. These images act as vertices of the convex hull. Along with the seed image, this approach facilitates verification of the model's robustness to any brightness variation within the range [-b, b]. The same methodology applies to contrast variations c ∈ℝ_+. We examine the effects for b values of {1, 2} and c values of 5× 10^-4, 5× 10^-3, 1× 10^-2}.










 

 
 
 
 
 
 




 §.§ Validity of Assumption <ref>


To assess the degree of symmetry and unimodality exhibited by the normalized heatmaps, we use four metrics: skewness, the disparity between mean and median, the disparity between mean and mode (peak), and the Interquartile Range (IQR), which is the difference between the third quartile (Q3) and the first quartile (Q1). A skewness near zero and a minimal difference between the mean and median indicate a symmetrical distribution. Combined with small skewness and close mean and median, a small IQR and close mean and mode suggest unimodality. These metrics are calculated individually for each image dimension. The evaluation uses the ResNet-18-based backbone model  across 3000 seed images and 4000 perturbed images with object occlusions and is presented in Fig. <ref>. The findings reveal skewness values approximately zero and mean and median values that are closely aligned, especially when considering the heatmap dimensions of 256 × 256. The difference between mean and mode and IQR values are also minor relative to the total range of 256. Thus, these statistical outcomes demonstrate that the normalized heatmaps predominantly feature unique peaks and exhibit axis symmetry, supporting Assumption <ref>.







 


 §.§ Probabilistic Soundness and Completeness of Optimal Error Threshold Allocation



  §.§.§ Optimal tolerance allocation

Note that the outcomes of allocating optimal tolerance are influenced by images and defined thresholds for acceptable pose errors. In our experiments, we manipulate these thresholds by incorporating a scaling factor α∈_+, resulting in adjusted error thresholds αϵ_r and αϵ_t.  The pose error thresholds are set as ϵ_r = α· [10^∘, 10^∘, 10^∘] and ϵ_t = α· [4, 4, 20]. The threshold factor α is varied across the values {0.5, 1.0, 1.5, 2.0}. We maintain the scaling factor κ at 1.0 as specified in Eq. (<ref>), and assign the weights w_1 and w_2 values of 1 and 5, respectively. Our analysis includes 3000 images, each contributing 5000 sets of perturbed keypoints, accumulating a total of 1.5 × 10^7 samples. The mean and standard deviation of the tolerance allocated in the horizontal and vertical directions are presented in Fig. <ref>. As the allowable pose errors grows, the tolerance allocated per keypoint increases, predominantly uniform in both the horizontal and vertical axes, with the exception of the 10-th and 11-th keypoints, which are symmetrical with respect to the body axis of the plane (see Fig. <ref>), exhibit larger tolerances horizontally, while the 5-th and 23-th keypoints, aligned along the body axis, have larger tolerances vertically. 

 


  §.§.§ Probabilistic soundness
 In this part, we intend to evaluate the probabilistic soundness of optimal error threshold allocation by verifying if the pose estimation errors resulting from perturbed keypoints fall within the predefined error thresholds. One approach is to perform random sampling within the hyper-rectangle (δ^*). However, given the high dimensionality of , which is 2^2K with K exceeding 20 in our scenario, this method demands an extraordinarily large sample size to achieve sufficient coverage. To address this issue, we only evaluate on vertices. For a seed image we randomly choose vertices from the hyper-rectangle (δ^*), which we then add to the coordinates of the ground-truth keypoints to create perturbed keypoints. Mathematically, in matrix form,
= ⊙δ^* + ,
where the elements of matrix  are randomly set to either -1 or 1, the symbol ⊙ represents element-wise multiplication, and  denotes the perturbed keypoints. This allows the perturbation of keypoints to reach the maximum tolerable errors. 

The simulation setup is largely the same as the previous section. We vary both the scaling factor κ in Eq. (<ref>) and the threshold factor α, and compute the ratio of samples where the pose estimation errors remain within thresholds over the total number of samples.  We report the ratio of samples resulting in tolerable pose errors. The results are summarized in Table <ref>. As we can see, only a small number of samples result in poses that surpass the error threshold. There is a noticeable trend where, vertically, as κ increases, the ratio of acceptable samples rises due to the shrinking of the keypoint error threshold polytope. Horizontally, increasing α results in a lower ratio, as linearization becomes less accurate when the pose moves away from the reference point.
 





















 























  §.§.§ Probabilistic completeness
 
To compute completeness, we randomly sample from the set of tolerable pose errors Ξ, use the perspective projection model (<ref>) to determine keypoints, and then compute the reprojection error relative to the ground truth. We verify whether these reprojection errors remain within the permissible error bounds on keypoints, represented by the hyper-rectangle . This simulation follows the same setup as used for probabilistic soundness. The results are presented in Table <ref>. Vertically, as κ increases, completeness diminishes due to the reduction in the size of , leading to more samples exceeding these bounds. On the other hand, there is no discernible trend when viewed horizontally.

When examining Table <ref>, a trade-off between soundness and completeness is evident. Our framework exhibits better soundness, implying that the set of allocated keypoint errors, , is relatively small compared to the ground-truth set δ_Ξ. This smaller size of  accounts for the reduced completeness observed in the results. To support this, Fig. <ref> illustrates the reprojection errors relative to the optimally allocated bounds (rectangle) for keypoints 1, 3 and 5 from a specific image under the conditions α = 0.5 and κ = 1.0. These reprojection errors are calculated by randomly generating pose errors within tolerable ranges and calculating the difference between ground-truth keypoints and those obtained from perspective projection model (<ref>). The illustration shows that a significant number of samples exceed the computed bounds, contributing to the overall low completeness observed across all keypoints. Therefore, if the verification tool returns a hold, it implies that the pose estimation is robust with high probability. On the other hand, if the tool returns a violation, we cannot draw any definitive conclusions. The conservativeness in error allocation mainly stems from the need to decouple the dependencies among keypoints.





















 §.§ Verification of Local and Global Perturbations

In this section, we assess the robustness of the pose estimation method when subjected to various levels of perturbations. We aim to answer three key questions:

    
    
    
  * How computationally efficient is the resulting neural network verification problem?
    
  * How accurate is the proposed certification method for robust pose estimation?




  §.§.§ Metric


We measure the performance using verification times and verified rates. Verified rate is defined as the proportion of cases where the verification algorithm confirms robustness against those where seed images produce acceptable pose estimation errors.   Another critical measure is verification accuracy which is defined as the proportion of cases where the verification algorithm confirms robustness against those that are indeed robust. However, determining the exact number of truly robust instances is impractical. Consequently, the verified rate serves as a lower bound of verification accuracy as the instances with acceptable pose estimation errors from seed images exceed those that are truly robust.













  §.§.§ Verification results


 We employ the verification toolbox  () <cit.>, which accepts convex hulls as input specifications. is the state-of-the-art verifier that supports a wide range of verification algorithms and is the most user-friendly to extend. It follows a branch-and-bound strategy to divide and conquer the problem efficiently. Two parameters guide this process:  determines the division of an unknown branch into smaller branches for further refinement, and  dictates the approach to navigating through the branch. We set  to use breadth-first search and   to bisect the branch. The computing platform is a Linux server equipped with an Intel CPU with 48 cores running at 2.20GHz and 376GB of total memory, approximately 150GB of which is available owing to multiple users. Additionally, the server includes 4 NVIDIA RTX A4000 GPUs, each with 16GB of memory.

 
 


 
 



  
Local object occlusions for the CNN-based model The results presented in Tab. <ref> indicate that for non-overlapping images, verification time increases as the number m of perturbed images forming the convex hull rises. As the error bounds for pose estimation expand, with threshold factor α increasing from 1.0 to 1.5, the time decreases. A similar effect is observed when the scaling factor κ decreases from 1.5 to 1.0. This reduction occurs because the optimally allocated error thresholds on keypoints expand with increasing α and decreasing κ, as illustrated in Fig. <ref>, which results in fewer nodes. A similar trend is observed for overlapping images. We emphasize that the verification process for overlapping images requires more time than for non-overlapping ones, given the same number of perturbed images and identical pose estimation error bounds, indicating that the keypoint detection model is effective in suppressing disturbances external to the airplane. 

In reference to the verified rates displayed in Tab. <ref>, for non-overlapping images, the rate remains stable regardless of the number of perturbed images, given the same κ and α. Conversely, there is an increase in the verified rate with a decrease in κ and an increase in α, as larger allocated keypoint error thresholds or larger allowable pose error thresholds  result in more images being verified as robust. In the case of overlapping images, a notable trend is the decline in the verified rate when the number of perturbed images increases. This is because a rise in the number of perturbed images means more objects are overlaid on the airplane, which compromises the accuracy of predictions and, in turn, decreases the number of images verified as robust.






  
Global perturbations for the CNN-based model A similar trend to that seen with local perturbations emerges, as indicated in Tables <ref> and <ref>. The verification can handle contrast variations c of 1% and brightness variations b of 2/255. Greater variations in contrast and brightness lead to longer verification times, whereas a larger α reduces verification time and increases the verified rate. 



  
Local block occlusions for the ResNet-18-based model We set κ = 1.0 and α = 1.5. For the convex hull comprised of non-overlapping images, the verification time is 314.4±227.0 in seconds, with a verification rate of 93.5%. Conversely, for the convex hull consisting of overlapping images, the verification time significantly escalates to 1571.7±1213.5 seconds, and the verification rate is 94.0%. 


 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 


 






























 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

 

 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

 
 
 
 
 
 
 
 
 
 
 
 
 
 
 









 







§ CONCLUSIONS

In this study, we introduce a framework designed to certify the robustness of learning-based keypoint detection and pose estimation methods. Given system-level requirements, our approach transforms the certification of PnP-based pose estimation into the standard verification for classification, allowing us to leverage off-the-shelf tools. The evaluation results demonstrated that our framework can handle realistic semantic perturbations compared to existing methods. We emphasize that our certification framework is general for safety-critical applications that depend on accurate keypoint detection. These include airport runway detection for automatic landing, pedestrain pose estimation for autonomous driving, and anatomical landmark identification for robot-assisted surgery. Future directions of this framework could include: 1)	Expanding the input specifications to represent more perturbations, such as the translational movement of objects. 2) Reducing the conservativeness caused by independent error allocation among keypoints.



§ ACKNOWLEDGEMENT

This material is based upon work supported by The Boeing Company. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views of The Boeing Company.


ACM-Reference-Format




 Paper Title: Beat this! Accurate beat tracking without DBN postprocessing.




Exploring atmospheric neutrino oscillations at ESSnuSB
    O. Zormpa
    August 5, 2024
======================================================



§ ABSTRACT 
We propose a system for tracking beats and downbeats with two objectives: generality across a diverse music range, and high accuracy.
We achieve generality by training on multiple datasets – including solo instrument recordings, pieces with time signature changes, and classical music with high tempo variations – and by removing the commonly used Dynamic Bayesian Network (DBN) postprocessing, which introduces constraints on the meter and tempo.
For high accuracy, among other improvements, we develop a loss function tolerant to small time shifts of annotations, and an architecture alternating convolutions with transformers either over frequency or time.
Our system surpasses the current state of the art in F1 score despite using
no DBN.
However, it can still fail,
especially for difficult and underrepresented genres, and performs 
worse on continuity metrics, so we publish our model, code, and preprocessed datasets, and invite others to beat this.






§ INTRODUCTION






Beat tracking is the task of estimating the temporal locations of musical beats in an audio signal. It is often combined with the downbeat tracking task, which targets a higher metrical level: tracking the beginning of each measure.
Despite being one of the long-standing problems in the Music Information Retrieval (MIR) field, it still attracts attention and several approaches were proposed in recent years <cit.>. Most recent work follows a common pipeline: the audio files are transformed into some spectrogram-like representation, then a deep neural network predicts frame-wise beat and downbeat probabilities, which are postprocessed to obtain the final predictions. The most widely used postprocessing technique is the Dynamic Bayesian Network (DBN) in the form proposed by Böck et al. <cit.>. It addresses four tasks: variable threshold peak-picking, forcing the tempo to stay in a certain range (i.e., limiting the allowed distance between beats), limiting sudden tempo changes, and (for downbeat tracking) ensuring that the downbeat falls every n beats, where n is constant for a piece of music and is selected from a limited list of values.

We argue that the DBN is a problematic component because it is inherently bound to fail for several music pieces: pieces with time signature changes, pieces whose tempo falls outside of the tempo range (or that slow down/speed up outside the tempo range), and pieces whose number of beats per measure are not included in the list of supported values. 
Moreover, it has a fixed parameter controlling allowed tempo variations, although we can expect, for example, classical music to have bigger tempo variability than rock music. 
Finally, even the hypotheses of having periodic beats and downbeats may be invalid, for example, for songs where the players make mistakes or audio tracks containing multiple concatenated songs.

Still, the DBN performs well on most pieces commonly used to train and evaluate beat tracking systems: music with a constant time signature of 3/4 or 4/4 and a stable, medium tempo. This can be seen from the default DBN parameters which most systems use,[We could verify that <cit.> use these parameters, since their code is publicly available, and we assume <cit.> do as well, since they do not mention any details in their paper.] i.e., tempo range [55, 215] BPM, beats per measure [3,4], and a tempo variability optimised on pop, rock and dance datasets. Pieces outside these specifications are likely to be mispredicted by the system, but form a minority in typical datasets. Therefore, in terms of evaluation metrics, it usually does not pay to remove the DBN.
However, working in these “simplified” conditions blocks research from solving corner cases in existing datasets and targeting more challenging or diverse data.


Our first goal is thus to replace the DBN with minimal postprocessing, free of the aforementioned musical assumptions.
A recent attempt to remove the DBN was made by Chen and Su <cit.>. However, their system may not look appealing to practitioners requiring beat tracking for downstream tasks, or researchers seeking a system to improve, as its accuracy falls clearly behind DBN-based ones.

Our second goal is to provide a powerful basis for practitioners and researchers.
The current best-performing system (which uses a DBN) from Hung et al. <cit.> falls short in this regard, as its code or a pretrained model is not public, its architecture is very complex, and (to the best of our knowledge) the results could not be reproduced by others.

In this paper, we present an open-source system that obtains new state-of-the-art F1 scores without a DBN.
It is based on a rotary transformer <cit.> applied on spectrograms, with the following novelties:
(1) We design a frontend alternating convolutions with a transformer variant by Lu et al. <cit.> that attends alternatively over frequency bins or time frames.
(2) We train with a shift-tolerant binary cross-entropy (BCE) that can cope with small deviations in the beat/downbeat annotations, and with weights on beat/downbeat frames to balance their relative scarcity.
(3) We propose an approach that encourages downbeat predictions to be a subset of beat predictions, and
(4) a data augmentation masking input segments to encourage the network to consider a longer context.
 Our code, pretrained models, and preprocessed datasets are openly available.[<https://github.com/CPJKU/beat_this>]



§ RELATED WORK

The currently best-performing model (on the GTZAN dataset <cit.> commonly used for evaluation) is by Hung et al.<cit.> and serves as a point of comparison. It uses a complex neural network architecture named SpecTNT which alternates computing frequency-related features with a frequency-oriented transformer, and processing a virtual extra frequency band with a time-oriented transformer. This runs in parallel with a more widely used Temporal Convolutional Network (TCN, a fully convolutional network with dilated convolutions), and the outputs of the two networks are merged for the final predictions. Unfortunately, the approach is not open source, and to the best of our knowledge, no other research group has managed to reproduce its results. Moreover, it still uses the DBN, which, as argued in the introduction, limits the system's generality.

Although no other work could reach the accuracy reported by Hung et at., two other recent beat tracking papers brought new interesting ideas <cit.>. Both perform instrument separation (with a pretrained network) and feed the separate stems 
(bass, drums, vocals, other, for <cit.> also piano)
into the model, mixing their information in cross-instrument attention blocks. While this approach is very reasonable from a music perception standpoint, it reduces the generality of the system, since it assumes that the input pieces will contain such instruments, at least to some extent. Another proposal of <cit.> is the use of dilated attention, following the successful use of dilated convolutions in beat-tracking architectures to 
increase the receptive field without adding computations. We find that flash attention <cit.> enables us to train with dense attention over a satisfactorily large input size, and leave experiments with dilated attention for future work.

Chen and Su <cit.> try to remove the DBN and propose a set of improvements. The most impactful is to replace the BCE with the Dice <cit.> and Focal <cit.> loss, inspired by their common usage for medical image segmentation. While these losses improve results,
possibly due to their inherent ability to handle unbalanced classes,
we found that a BCE with weights on the positive (beat and downbeat) classes outperforms them. We suspect this is because, in contrast to medical image segmentation, our positive examples are single frames, and the Dice and Focal loss perform better when the area of positive predictions is larger <cit.>. 

Another proposal by Chen and Su is to predict the phase of the beat/downbeat instead of a single binary value, following <cit.>. Although this seems promising, the results on both papers (and ours) do not show any consistent improvement.

Many recent approaches <cit.> use the additional task of tempo prediction (a single tempo target for each input excerpt) in a multi-task setting. While this improves their results, it goes against our goals of generality, since it assumes an (almost) constant tempo in the file excerpt, which is not the case for many kinds of music.

Other recent papers do not align with the goal of this paper: <cit.> explores the usage of different time resolutions between the input audio and output predictions (only addressing beats); <cit.> performs unsupervised beat tracking; <cit.> focuses on online beat tracking;  <cit.> notices the problems of the DBN for music with tempo changes, and proposes a different postprocessing method targeted specifically to classical music; <cit.> focuses on fine-tuning existing systems, and changing the DBN parameters for targeting specific underrepresented genres.





§ METHOD

Our beat tracker
is based on a neural network with ∼20 M parameters.
It
starts from 30 seconds of mono audio sampled at 22.05 kHz and converts it to a 128-bin mel spectrogram from 30 Hz to 10 kHz, with a window size of 1024 and hop size of 441 samples (yielding 50 frames per second), and magnitudes scaled via ln(1 + 1000 x) (similar to ln(max(10^-3, x)), but maps silence to zero).
These hyperparameters were optimised in preliminary experiments.
Our model processes this into frame-wise beat and downbeat probabilities, followed by minimal post-processing to derive beat and downbeat locations.



 §.§ Model



Our model (Figure <ref>) processes a T×128 spectrogram into T×2 probabilities; T being the number of input frames (1500 in case of 30 s).
It consists of three components: a frontend converting the spectrogram into a sequence of feature vectors, a transformer processing these vectors, and two task heads computing the output probabilities.




  §.§.§ Frontend

The frontend's role is to integrate information across the 128 frequency bands into feature vectors.
Typically, this is done via 2d convolutions gradually reducing the number of bands to 1 while increasing the number of channels <cit.>.

We adopt this, but found it helps to interleave convolutions with Partial Transformers, which treat the time and frequency axis independently.
Overall, our frontend consists of a stem, three identical blocks, and a linear projection.

The stem (Figure <ref>, top right) starts with a batch normalisation that processes each frequency band separately to homogenise them, followed by a 2d convolution of 3×4 kernels, regular batch normalisation, and GeLU nonlinearity. The convolution is strided to reduce the number of frequency bands to a fourth and creates 32 channels.

Each block (Figure <ref>, middle right) consists of two partial transformers, a strided 2d convolution halving the number of bands while doubling the number of channels, batch normalisation and GeLU.
The first partial transformer is frequency-directed, i.e., it processes the T×F×C tensor by treating each time frame as a sequence of length F (the number of bands), the second one is time directed and treats each frequency band as a sequence of length T (the number of frames), an idea adopted from the Band Split RoFormer <cit.>.
Each transformer has a head size of 32 (one head in the first frontend block, two in the second, four in the third), rotary positional embedding <cit.>, a sigmoid gate per head <cit.>,
and includes a usual pointwise feedforward network with a hidden size of four times the channel count.




After three frontend blocks, the resulting T×4×256 tensor (4 bands, 256 channels) is reshaped to a T×1024 tensor and linearly projected to 512 features.



  §.§.§ Transformer

The transformer makes up the bulk of our model's parameters and compute.
It consists of 6 stacked transformer blocks processing the 512-dimensional vectors with 16 heads of size 32, rotary positional embedding, sigmoid gating, and a pointwise feedforward network of 2048 hidden units.
This matches the configuration in the frontend transformer blocks, but as it processes a single sequence of 512-dimensional feature vectors, it is a regular transformer over time without separately considering a frequency dimension.
Its goal is to map the 512 input features to a space that relates to beats and downbeats.
Due to the attention mechanism, the transformer's receptive field covers the full sequence, and it could therefore produce an output that has characteristics that we want in the beat predictions, for example, regularity.



  §.§.§ Task Heads

The output of the final transformer block is processed by two linear layers, one for beats and one for downbeats. Initially, we used the common approach of passing their output into 2 sigmoid functions to produce a probability for each input frame, and then threshold this probability at 0.5 to obtain "hard" beat predictions. However, we observe that this sometimes produces downbeat predictions not coinciding with a beat prediction, which is allowed under the evaluation metrics but is a musically invalid and unusable output.
This problem is solved when using a DBN to jointly process beats and downbeats. However, we noticed that several works, e.g., <cit.>, use two independent DBNs to predict beats and downbeats (and others <cit.> do not specify). To our surprise, this leads to better metrics, but it severely limits practical use.

To mitigate this problem, we propose a Sum Head that sums the output of the beat and downbeat layers, and treats this as the beat logits (for prediction and training). This is a very simple way of helping the network produce a beat when there is a downbeat (though it does not enforce that; a highly negative output of the beat layer can still counter the downbeat layer). We explored other ways of aggregating the beat and downbeat logits, like taking their maximum, but this hampered training due to the sparser gradients. On the GTZAN dataset, the sum head almost halves the percentage of downbeats that are more than 70 ms away from the closest beat, from 1.1% to 0.62%, compared to directly using the output of the linear layers. We observe that the remaining unmatched downbeats are in pieces with very erratic predictions that would be incorrect anyway.

Some systems circumvent the problem by using a 3-way classifier (beat vs. downbeat vs. none) instead of the two binary classifiers (beat vs. none, and downbeat vs. none). However, to be able to train on datasets that do not include downbeat annotations, we stick to binary classifiers.




















 §.§ Postprocessing

To obtain beat/downbeat locations, we pick all frames assigned the highest beat/downbeat probability within a neighborhood of ±3 frames (± 70 ms), and probability > 0.5. In case adjacent frames are assigned the same probability, we report their center. Finally, we move all downbeat predictions to the closest beat prediction to correct the remaining mismatches described in the previous section.
For music pieces longer than 30 s, we concatenate predictions over non-overlapping 30-second excerpts.




 §.§ Loss




The model is trained by gradient descent on a loss function that compares the frame-wise beat and downbeat predictions with frame-wise binary annotations.
The usual loss for binary classification is Binary Cross-Entropy, L_bce(𝐲, ŷ) = -∑_t y_t log(ŷ_t) + (1-y_t) log(1-ŷ_t).
Training with BCE leads to unconfident predictions since the problem is heavily imbalanced.
To counter this, we can weight positive examples by a factor w as L_wbce(𝐲, ŷ, w) = -∑_t w y_t log(ŷ_t) + (1-y_t) log(1-ŷ_t).
We found that setting w to the number of negative examples divided by the number of positive examples (over the training set) yields the best result and is crucial when not using a DBN.

Another problem persists: annotations are not precise down to our spectrogram resolution, due to annotators' impreciseness, players' asyncronicity, or simply the limits of human perception.
This is taken into account during the evaluation, e.g., the typically used F1 score accepts predictions in a ± 70 ms window around labels.
During training, the BCE loss punishes close positive predictions (Figure <ref>, left) even though they may be correct, thus creating two problems: training is slower and the network learns to predict wide “blurred” peaks.
This is commonly addressed by adding two extra positive labels around each annotation weighted by 0.5, but this only mitigates the former problem without helping with the latter.
Instead, we max-pool predictions over time (7 frames, stride 1) before comparing them to the labels.
In this way, only the largest positive prediction ± 3 frames from each label is considered (Figure <ref>, right).
The loss for negative examples is ignored ± 6 frames from each label, as this is how far a max-pooled prediction 3 frames away from a label spreads.
Denoting max-pooling of k frames with m_k(·), we can formalise our Shift-tolerant weighted BCE as:
L_st(𝐲, ŷ, w) = -∑_t w y_t log(m_7(ŷ)_t) + (1 - m_13(𝐲)_t) log(1 - m_7(ŷ)_t).



 §.§ Data Augmentation


Masking. To encourage the model to not only rely on local information for its predictions, we mask 0 to 6 areas of 0.5 to 2 s. Each masked area is randomly divided into 5 to 10 parts which are randomly reordered. This destroys local correspondence between audio and beats without changing local input statistics, and works better than zero masking as employed in SpecAugment <cit.>. We assume our approach makes it harder for the network to learn a dedicated behaviour for masked areas.

Pitch and time. We speed up and slow down every song by 20, 16, 12, 8, and 4%, and transpose by at most +6 and -5 semitones.
We precompute these augmentations (using Pedalboard <cit.>) so experiments become reproducible without access to the original audio.
To limit storage use, tempo and pitch augmentations are not combined, giving 22 variations for each song.
We verified that our limited tempo augmentation gives comparable results to the commonly used approach by Böck and Davies <cit.> of performing on-the-fly augmentations by randomly changing the hop size of the STFT, at the advantage of not requiring audio access.






§ EXPERIMENTS

We perform 8-fold cross-validation experiments on multiple datasets, compute results on the test-only GTZAN dataset, and do an ablation study.

We use the standard beat-tracking metrics: F1, CMLt, and AMLt 
and compute them using the  package <cit.> with default parameters.[This includes “trim_beats” of 5s that discards the first 5 seconds during the evaluation, which is a choice we do not necessarily approve of, but we use it to be consistent with what seems the standard way of evaluating.] CMLt and AMLt are called continuity metrics, and only consider a beat/downbeat as correct if both it and the previous beat/downbeat are correct; AMLt also accepts different metrical levels such as half or double time, and offbeats <cit.>.
The metrics and their settings match those used by Hung et al. <cit.>; this enables a comparison with their reported results, though it is unclear which 8-fold datasplit they use,
and we cannot run any statistical significance tests since their code is not reproducible. Therefore, any comparison needs to be taken only as an indication.




 §.§ Datasets

We train and validate with several datasets: Simac<cit.>, SMC <cit.>,
Hainsworth <cit.>, Ballroom <cit.>, HJDB <cit.>, Beatles <cit.>, Harmonix <cit.>, 
RWC <cit.> (classical, pop, royalty-free, and jazz), 
TapCorrect <cit.>, JAAH <cit.>, Filosax <cit.>, ASAP <cit.>, Groove MIDI <cit.>, GuitarSet <cit.>, Candombe <cit.>. The first two datasets contain only beat annotations, all others both beat and downbeats. We discard one Beatles piece which does not contain downbeat annotations and one with empty beat annotations, resulting in a total of 4556 tracks. For comparison, Hung et al. <cit.> train with only the first 7 datasets reported above (Simac to Harmonix), plus RWC pop, for a total of 3144 pieces (when assuming the same handling of missing annotations).
We use the GTZAN <cit.> dataset (993 pieces discarding one unannotated track and 6 tracks that miss downbeat annotations) for testing only.

We use only the backing tracks of Filosax without the saxophone solos. For ASAP, we discard the tracks that contain the “rubato” beat annotations. In Groove MIDI, we keep all pieces that are longer than 20 seconds and use the provided audio renderings. We only use the comping tracks of GuitarSet, discarding the solos. 

We employ the 8-fold cross-validation splits published by Böck and Davies <cit.> for the datasets they used and produce new ones for
our added datasets, ensuring different versions of the same piece are not spread across folds, and stratifying by metadata when possible.
We also produce a new single split with ∼ 15% of the pieces on each dataset as validation (again taking care of different versions of the same piece).



 §.§ Training


We train for 100 epochs with gradient accumulation over 8 batches of size 8,[This enables training with under 8 GiB of GPU memory.] AdamW optimizer <cit.>, weight decay of 0.01 (excluding biases and learned norms), learning rate warm-up <cit.> of 1000 steps to a maximum of 0.0008, and cosine annealing.
During each epoch, we randomly sample 30 seconds of each piece, and pad if the total piece duration l is less than 30 seconds. 
We draw k samples from pieces that are longer than 30 seconds, following the equation k = round(α· l/30 )



with α=0.65, since we observe it leads to faster convergence than using one random sample per piece, or l/30 non-overlapping samples. On average, this yields ∼ 3 samples per piece. During training, every time a sample is drawn, we randomly select a precomputed augmentation described in Section <ref>, and apply masking.
The full training takes around 8 hours on a single NVIDIA RTX 2080 Ti, 6 hours on A40, and 4 hours on A100.

During our experimentation, we found that to achieve good results without a DBN, we need our network to be overconfident in its predictions.
This may seem to violate usual deep learning practice, but can be explained by a closer look at the beat tracker's desirable output. We do not want our network to produce probabilities close to 0.5 when unsure, since this will lead to random oscillations between positive and negative predictions, and thus erratic beats. Instead, we want it to give steady, high-probability predictions even when unsure, exactly like the DBN would.
To achieve this, we keep training even after the validation loss starts increasing,


which would typically indicate overfitting.

Indeed, we see that the validation F1 score continues to improve even with increasing validation loss.

This means that even with our modifications, the BCE loss is not a good indicator of the F1 score, and further research into alternative losses may be valuable.

The reader may wonder why, once we obtain our well-performing network, we do not use the DBN to increase the metrics even more. 
By having overconfident predictions, we reduce the benefits of such a postprocessing method. With a degree of simplification, we can imagine the DBN as using a model's most high-confident predictions to infer beats in low-confidence areas. By avoiding the low-confidence predictions, we are disrupting this mechanism. 



 §.§ Main Results


We report the results on our 18 datasets in the commonly used 8-fold cross-validation setting: each dataset is split into 8 parts, we jointly train on 18·7 parts (all but one per dataset) and predict on the remaining 18, after 8 such runs we covered all pieces and average metrics over pieces by dataset. We observe that our model outperforms Hung et al. <cit.>, except for Harmonix and RWC Pop (downbeat). 
In our results, the lowest downbeat performance is obtained in the ASAP and RWC Classical datasets, confirming the well-known difficulty of beat-tracking classical music <cit.>. Performance on SMC (where only beat annotations are accessible) is also very low, consistent with the outcomes of other systems, highlighting the substantial room for improvement that beat tracking systems continue to hold.





We also report the results on the GTZAN dataset in Table <ref>. We compute these results with a single model trained on the entirety of our training-val dataset (note that we do not perform any early stopping or other techniques for which the validation dataset may still be necessary). All our runs are computed 3 times with different random seeds, and we report the means and standard deviations of the computed metrics over the 3 seeds.
We notice that even when training on the reduced collection of datasets by Hung et al. (third row in the table), we still outperform their F1 score without a DBN, proving the effectiveness of our design choices. 
Our main model has ∼20 M parameters, 5 times more than Hung et al. with 4 M, so we also show the results for a smaller model with the hidden dimension of the main transformer blocks reduced from 512 to 128, and the number of heads from 16 to 4. This small model has ∼2 M parameters and still gives SOTA F1 scores.

Disappointingly, we notice that the continuity metrics (CMLt and AMLt) are lower than those of Hung et al. From qualitative inspections of the results, we notice that for complex or underrepresented pieces, our network introduces non-periodic beats, which drastically lower the continuity metrics. We are then brought to wonder why our network cannot learn a supposedly obvious behaviour, such as only producing periodic-like output, and we can propose two explanations. Firstly, our loss does not specifically penalise non-periodic predictions, but treats each beat individually. This results in a discrepancy between what is preferred by continuity metrics and what the network learns to predict in difficult parts to minimise the loss. Secondly, our datasets contain several non-periodic annotations, some due to quality issues (see Section <ref>), some in correctly annotated pieces such as tapcorrect_10 or beatles_Wild-Honey-Pie, where a 2/4 measure in the middle of a 4/4 piece disrupts the assumption of periodicity for downbeats. Finally, one could also question the generality of the AMLt metric as a tool to quantify double/half-time errors, since the computations of different metrical levels assume that the time signature and the tempo do not change and that a measure can always be divided into 2 or 3 parts.

Using a DBN increases our CMLt downbeat performance by correcting some of the (wrongly) non-periodic outputs, but it reduces our F1 performance, by changing other otherwise correct predictions that fall outside the DBN assumptions. The AMLt score does not increase since our network is overconfident in its predictions, and the DBN cannot easily switch to another metrical level. 




 §.§ Ablation Studies




We ablate multiple components of our model on the single split described in Section <ref>.
We perform every experiment 3 times with different seeds and report the mean and standard deviation on the validation set in Table <ref>.
The usage of our Sum Head shows little impact, but we use it to have a musically valid output, rather than to increase the F1 score. Pitch, mask, and tempo augmentations help, in this order of importance.
The usage of partial transformers in our frontend proves more effective than only having convolutions.
Our most impactful design choice is the weighted shift-tolerant loss. Using a normal BCE with positive example weights results in decreased performance, which decreases even further when the weights are removed.



 §.§ Notes on Dataset Quality

While exploring the datasets, we found multiple problems in the annotations, and we think this hinders the development of better models, especially for downbeat predictions. Even the GTZAN dataset, which is commonly used for evaluation, is not immune to quality problems. Some of them are evident and not debatable, like jazz_00000, jazz_00002, jazz_00083, blues_00015, reggae_00095, classical_00077, rock_00067.
Furthermore, there are pieces where even for experts it would be hard to agree on a unique beat and downbeat annotation, like metal_00081 or classical_00056, and multiple annotations would be necessary. 
Finally, some pieces question the primary assumption of beat tracking, i.e., that there is a beat/downbeat to track, like pop_00064, and jazz_00003.



§ CONCLUSIONS AND OUTLOOK

In this paper, we presented a new beat tracking system which obtained a state-of-the-art F1 score on a very diverse set of music, with minimal assumptions about the tempo, time signature, and their changes over time. Remarkably, we do not use the DBN postprocessing, which was employed by all recent high-accuracy models. However, removing the DBN hurts the CMLt and AMLt metrics. A study on how this trade-off affects human perception, alternative metrics, and a direct comparison with DBN-based models on complex pieces is left for future work.

We emphasise that beat tracking is not a solved problem, even for commonly targeted genres such as rock or electronic music, especially for the downbeat tracking task. We provide an open-sourced model that can be used as a starting point, and we invite future researchers to improve it. Potential directions are: reducing the model parameters, developing new losses that enforce periodicity during training, using other data augmentation techniques to make the system more robust to multiple sound conditions, fine-tuning it on specific genres, and training on larger datasets. The contribution of people with musical expertise will also be essential, as we think that correcting the existing commonly used datasets, and producing new annotated data for underrepresented genres is a crucial step for further development.



§ ACKNOWLEDGEMENTS

This work was supported by the European Research Council (ERC) under the EU's Horizon 2020 research & innovation programme, grant agreement No. 101019375 (Whither Music?), and the Federal State of Upper Austria (LIT AI Lab). The computational results presented were achieved in part using the Vienna Scientific Cluster (VSC).






 Paper Title: An Effective Dynamic Gradient Calibration Method for Continual Learning.



[
An Effective Dynamic Gradient Calibration Method for Continual Learning














equal*


Weichen Linequal,yyy
Jiaxiang Chenequal,yyy
Ruomin Huangcomp
Hu Dingsch





yyySchool of Data Science, University of Science
 and Technology of China, Anhui, China.
compDuke University 
schSchool of Computer Science and Technology, University of Science and Technology of China, Anhui, China

Hu Dinghuding@ustc.edu.cn





Machine Learning, ICML

0.3in
]










 

§ ABSTRACT 
Continual learning (CL) is a fundamental topic in machine learning, where the goal is to train a model with continuously incoming data and tasks. Due to the memory limit, we cannot store all the historical data, and therefore confront the “catastrophic forgetting” problem, i.e., the performance on the previous tasks can substantially decrease because of the missing information in the latter period. Though a number of elegant methods have been proposed, the catastrophic forgetting phenomenon still cannot be well avoided in practice. In this paper, we study the problem from the gradient perspective, where our aim is to develop an effective algorithm to calibrate the gradient in each updating step of the model; namely, our goal is to guide the model to be updated in the right direction under the situation that a large amount of historical data are unavailable. Our idea is partly inspired by the seminal stochastic variance reduction methods (e.g., SVRG and SAGA)  for reducing the variance of gradient estimation in stochastic gradient descent algorithms. Another benefit is that our approach can be used as a general tool, which is able to be incorporated with several existing popular CL methods to achieve better performance. We also conduct a set of experiments on several benchmark datasets to evaluate the performance in practice. 




§ INTRODUCTION


In the past years, Deep Neural Networks (DNNs) demonstrate remarkable performance for many different tasks in artificial intelligence, such as image generation <cit.>, classification  <cit.>, and pattern recognition <cit.>. Usually we assume that the whole training data is stored in our facility and the DNN models can be trained offline by using Stochastic Gradient Descent (SGD) algorithms <cit.>.  

However, real-world applications often require us to consider training lifelong models, where the tasks and data are accumulated in a streaming fashion <cit.>. For example, with the popularity of smart devices, a large amount of new data is generated every day. A  model needs to make full use of these new data to improve its performance while keeping old knowledge from being forgotten.  Those applications motivate us to  study the problem of continual learning (CL) <cit.>, where its goal is to develop effective method for gleaning insights from current data while retaining information from prior training data.





  A significant challenge that CL encounters is  “catastrophic forgetting” <cit.>, wherein the exclusive focus on the current set of examples could result in a dramatic deterioration in the performance on previously learned data. This phenomenon is primarily attributed to limited storage and computational resources during the training process; otherwise, one could directly train the model from scratch using all the saved data. To address this issue, we need to develop efficient algorithm for training neural networks from a continuous stream of non-i.i.d. samples, with the goal of mitigating catastrophic forgetting while effectively managing computational costs and memory footprint.

  A number of elegant CL methods have been proposed to alleviate the catastrophic forgetting issue <cit.>. One representative CL approach is referred to as “Experience Replay (ER)” <cit.>, which has shown promising performance in several 

continual learning scenarios <cit.>. Roughly speaking, the ER method utilizes reservoir sampling <cit.> to maintain historical data in the buffer, then extract new incoming training data with random samplings for learning the current task. 

Though the intuition is simple, the ER method currently is one of the  most popular CL approaches that incurs 



moderate computational and storage demands.

Moreover, several recently proposed approaches suggest that the ER method can be combined with knowledge distillation to further improve the performance; for example, 
the methods of DER/DER++ <cit.> and X-DER <cit.> preserve previous training samples alongside their logits in the model as the additional prior knowledge.

Besides the ER methods, there are also several other types of CL techniques proposed in recent years, and please 

refer to <ref> for a detailed introduction.



 §.§ Our Main Ideas and Contributions


  Though existing CL methods can alleviate the catastrophic forgetting issue from various aspects, the practical performances in some scenarios are still not quite satisfying <cit.>.
  
  In this paper, 

we study the continual learning problem from the gradient perspective, and the rationality behind is as follows. In essence, an approach for avoiding the catastrophic forgetting issue in CL,  e.g., the replay mechanisms or the regularization strategies, ultimately manifests its influence on the gradient directions during model updating <cit.>. If all the historical data are available, one could compute the gradient by using the stochastic gradient descent method and obviously the catastrophic forgetting phenomenon cannot happen. The previous  methodologies  aim to approximate the gradient by preserving additional information and incorporating it as a constraint to model updates, thereby retaining historical knowledge. However,   the replay-based methods in practice are often limited by storage capacity, which leads to a substantial loss of historical data information and inaccurate estimation of historical gradients <cit.>.

Therefore, our goal is to develop a more accurate gradient calibration algorithm in each step of the continual learning procedure, which can directly enhance the training quality. 

 We are aware of several existing CL methods that also take into account of the gradients <cit.>, but our idea proposed here is fundamentally different. 
 We draw inspiration from the seminal stochastic variance reduction methods (e.g., SVRG from <cit.> and SAGA from <cit.>), which are originally designed to reduce the gradient variance so that the estimated gradient can closely align with the true full gradient over the entire dataset (including the current and historical data). These variance reduction methods have been extensively studied in the line of the research on stochastic gradient descent method <cit.>; their key idea is to leverage the additional saved full gradient information to calibrate the gradient in the current training step, which leads to significantly reduced gradient variance comparing with the standard SGD method.  This intuition also inspires us to handle the CL problem. In a standard SGD method, the variance between the obtained gradient and the full gradient is due to the “batch size limit" (if the batch size has no bound, we can simply compute the full gradient).  Recall that the challenge of CL is due to the “buffer size limit”, which impedes the use of full historical data (this is similar with the dilemma encountered by SGD with the “batch size limit”). So an interesting question is 
 
 

 Can the calibration idea for “batch size limit” be modified to handle “buffer size limit”? Specifically,  is it possible to develop an effecitve method to compute a SVRG (or SAGA)-like calibration for the gradient in  CL scenarios?


  Obviously, it is challenging to directly implement the SVRG or SAGA algorithms in continual learning because of the missing historical data. 

Note that <cit.> proposed a streaming SVRG (SSVRG) method that realizes the SVRG method within a given fixed buffer, but unfortunately it does not perform quite well in the CL scenarios (as shown in our experimental section). One possible reason is that SSVRG can only leverage information within the buffer and fails to utilize all historical information. 



In this paper, we aim to apply the intuition of SVRG to handle the “buffer size limit” in CL, and our contributions can be summarized as follows:




    
  * First, we propose a novel two-level dynamic algorithm, named Dynamic Gradient Calibration (DGC),   to maintain a gradient calibration in continual learning. DGC can effectively tackle the storage limit and leverage historical data to calibrate the gradient of the model at each current stage. 
   
   Moreover, our theoretical analysis shows that our DGC based CL algorithm can achieve a linear convergence rate. 

    
   
   
    
  * Second, our method can be conveniently integrated with most existing reservoir sampling-based continual learning approaches (e.g., ER <cit.>, DER/DER++ <cit.>, XDER <cit.>, and Dynamic ER <cit.>), where this hybrid framework can 
    
    induce a significant enhancement to the overall model performance. 
    
    Note that storing the gradient calibrator can cause an extra memory footprint; so a key question is whether such an extra memory footprint can yield a larger marginal benefit than simply taking more samples to fill the extra memory? 
    In Fig <ref>, we illustrate a brief example to answer this question in the affirmative; more detailed evaluations are shown in the experimental part. 

    
    
  * Finally, we conduct a set of comprehensive experiments on the popular datasets S-CIFAR10, S-CIFAR100 and S-TinyImageNet; the experimental results suggest that our method can improve the final Average Incremental Accuracy (FAIA) in several CL scenarios by more than 6%. Moreover, our improvement for a larger number of tasks is more significant than that for a smaller number. Furthermore, our adoption of the SVRG-inspired DGC calibration method leads to enhanced stability in minimizing the loss function throughout the parameter optimization process.
    
    
    



 §.§ Related Work





We briefly overview existing important continual learning approaches (except for the ones mentioned before).  We also refer the reader to the recent surveys <cit.> for more details.



  
  A large number of CL methods are replay based, where they often keep a part of previous data through approaches like reservoir sampling <cit.>. Several more advanced data selection strategies focus on optimizing the factors like the sample diversity of parameter gradients or the similarity to previous gradients on passed data, e.g.,  GSS <cit.> and   GCR <cit.>.  
  
  Experience replay can be effectively combined with knowledge distillation. For example, <cit.> proposed to distill colliding effects from the features for new coming tasks, and ICARL <cit.> proposed to take account of the data representation trained on old data. MOCA <cit.> improves replay-based methods by diversifying representations  in the  space. Another replay-based approach is based on generative replay, which obtains replay data by generative models <cit.>.
  



  
  Another way for solving continual learning is through some deliberately designed optimization procedures. For example, the methods GEM <cit.>, AGEM <cit.>, and MER <cit.> restrict parameter updates to align with the experience replay direction, and thereby preserve the previous input and gradient space with old training samples. Different from saving old training samples, <cit.> proposed to adapt parameter updates in the orthogonal direction of the previously saved gradient.   The method AOP <cit.> projects the gradient in the direction orthogonal to the subspace spanned by all previous task inputs, therefore it only keeps an orthogonal projector rather than storing previous data.






  
  To mitigate the problem of forgetting, we can also augment the model capacity for learning new tasks. <cit.> tried to enhance model performance by employing meta-learning techniques  when dynamically extending the model.
  

The method ANCL <cit.> proposes to utilize an auxiliary network to achieve a trade-off between 
plasticity 
and 
stability. Dynamic ER <cit.> introduces a novel two-stage learning method that employs a dynamically expandable representation for learning knowledge incrementally. 



§ PRELIMINARIES



We consider the task of training a soft-classification function f(· ;θ):𝒳→𝒴, where 𝒳 and 𝒴 respectively represent the space of data and the set of labels, and θ is the parameter to optimize. Without loss of generality, we assume 𝒴={1, 2, ⋯, K}. So 
 f(·; θ) maps each x∈𝒳 to some f(x; θ)∈ℝ^K. 
 
 To find an appropriate θ, the classification function  f is usually equipped with a loss function ℓ (f(x;θ), y), which is differentiable for the variables x and θ (e.g., cross-entropy loss). To simplify the notation, we use ℓ(x,y,θ) to denote ℓ(f(x;θ), y).  Given a set of data P={(x^i, y^i)| 1≤ i≤ n} = 𝒳_P×𝒴_P ⊂𝒳×𝒴, the training process is to find a θ such that the empirical risk of ℓ(x,y, θ), i.e., ∑^n_i=1ℓ(x^i,y^i, θ), is minimized. We define the full gradient of ℓ(x,y, θ) on P as 

    𝒢(P, θ)    ≜∇_θℓ(𝒳_P, 𝒴_P, θ) = 1/n∑^n_i=1∇_θℓ(x^i, y^i, θ).




 §.§ Continual Learning Models


In this paper, we focus on two popular CL models: Class-Incremental Learning (CIL) <cit.> and Task-Free Continual Learning (TFCL) <cit.>.


  In the setting of CIL, the training tasks come in a sequence {𝒯_1, 𝒯_2, ⋯, 𝒯_T} with disjoint label space; each time spot t ∈{1,2,⋯,T} corresponds to the task 𝒯_t with a training dataset {(x_t^i,y_t^i)| 1≤ i≤ n_t}. With a slight abuse of notations, we also use 𝒯_t to denote its training dataset. Also, we use 𝒴_t to denote the corresponding set of labels {y_t^i| 1≤ i≤ n_t}. 


Although the training data for individual tasks 𝒯_t is independently and identically distributed (i.i.d.), it is worth noting that the overall task stream {𝒯_1, 𝒯_2, ⋯, 𝒯_T} does not adhere to the i.i.d. assumption due to the evolving label space over time.  We define the overall risk under this CL setting 
  at the current time spot t 
 as follows: 

    ℓ_CL^t (θ) ≜1/t∑_c = 1^t(x, y) ∼𝒯_c𝔼[ℓ(x,y,θ)].








  The setting of TFCL is similar to CIL, where the major difference is that the task identities are not provided in neither the training nor testing procedures. So the TFCL setting is more challenging than CIL because the algorithm is unaware of task changes and the current task identity. In the main part of our paper, we present our results in CIL; in our supplement, we  explain how to extend our results to TFCL.







 §.§ Variance Reduction Methods




A comprehensive introduction to variance reduction methods is provided in <cit.>. Here 
we particularly introduce 

one representative variance reduction method SVRG <cit.> 

, which is closely related to our proposed CL approach. 

The high-level idea of SVRG is to construct a calibration term to reduce the variance in the gradient estimate. The complete optimization process can be segmented into a sequence of stages. We denote the training data as P = {(x^i,y^i)}_i=1^n⊂𝒳×𝒴,  and denote the parameter at the beginning of each stage as θ̃.  The key part of SVRG is to  minimize the variance in SGD optimization by computing 

an additional term μ̃:


    μ̃   ≜𝒢(P, θ̃) =  1/n∑_i = 1^n∇_θℓ(x^i,y^i, θ̃).

In each stage, 
SVRG applies the standard SGD with the term μ̃ to update the parameter θ: 



let

    v^k     =  ∇_θℓ(x^i, y^i,θ^k) - (∇_θℓ(x^i, y^i, θ̃) -μ̃), 
    θ^k+1    = θ^k - η v^k,

where (x^i,y^i) ∈ P is the sampled training data, θ^k represents the  parameter at the k-th step of SGD, and η is the learning rate. Since 𝔼[∇_θℓ(x^i, y^i, θ̃^̃k̃)] = μ̃, v^k is an unbiased estimate of the gradient 𝒢(P,θ^k). Subsequently, the term “∇_θℓ(x^i, y^i, θ̃) - μ̃” in Eq (<ref>) can be regarded as a calibrator to reduce the variance of gradient estimation and achieve a linear convergence rate <cit.>, which is faster than directly using ∇_θℓ(x^i, y^i,θ^k).















§ OUR PROPOSED METHOD



In this section, we propose 
the Dynamic Gradient Calibration (DGC) approach which maintains a gradient calibration during the learning process. A highlight of DGC is that it utilizes the whole historical information to obtain a more precise gradient estimation, and consequently relieves the negative impact of catastrophic forgetting. In <ref>, we introduce ER and analyze the obstacle if we directly combine ER with SVRG. 

In <ref>, we present our DGC  method for addressing the issues discussed in <ref>. In <ref>, we explain how to integrate 
  DGC with other  CL  techniques.



 §.§ Experience Replay Revisited and SVRG



 First, we overview the classical ER <cit.> method as a baseline for the CIL setting. ER employs the reservoir sampling algorithm to dynamically manage a buffer (denoted as ℳ_t) at time t, which serves to store historical data. At each time spot t (and assume the current updating step number of the optimization is k), ER updates the model parameter  θ^k_t following the standard gradient descent method: 

    θ_t^k+1 = θ_t^k - η· v_t^k,

where η is the learning rate and v_t^k is the calculated gradient. If not using any replay strategy, v_t^k is usually calculated on a randomly sampled training data (x^k_t, y^k_t) ∈𝒯_t. It is easy to see that this simple strategy can cause the forgetting issue for shifting data stream since it does not contain any information from the previous data. 

Hence the classical ER algorithm takes a random sample 
(x̅^k_t, y̅^k_t) from the aforementioned buffer ℳ_t (who contains a subset of historical data via reservoir sampling), and computes the gradient:


    v_t^k =1/t∇_θℓ(x^k_t, y^k_t, θ_t^k) + t-1/t∇_θℓ(x̅^k_t, y̅^k_t, θ_t^k).



(1) For simplicity, we assume that the data sets of all the tasks have the same size. So the obtained v_t^k in (<ref>) is an unbiased estimation of the full gradient 𝒢(⋃_c=1^t𝒯_c,θ) at the current time spot t. If they have different sizes, we can simply replace the coefficients “1/t” and “t-1/t” by “|𝒯_t|/∑^t_c=1|𝒯_c|” and “∑^t-1_c=1|𝒯_c|/∑^t_c=1|𝒯_c|”, respectively. (2) Also, we assume that the batch sizes of the random samples from 𝒯_t and ℳ_t are both “1”, i.e., we only take single item (x^k_t, y^k_t) and (x̅^k_t, y̅^k_t) from each of them. Actually, we can also take larger batch sizes and then the Eq (<ref>) can be modified correspondingly by taking their average gradients. 













  Now we attempt to apply the SVRG method to Eq (<ref>).  

Our objective is to identify a more accurate unbiased estimate of the gradient at the current time spot t so as to determine the updating direction.


At first glance, one possible solution is to adapt the streaming SVRG method <cit.> to the CL scenario. 
We treat ℳ_t as the static data set in our memory, and apply the 
 SVRG technique to calibrate the gradient “∇_θℓ(x^k_t, y^k_t, θ_t^k)” and “∇_θℓ(x̅^k_t, y̅^k_t, θ_t^k)” in Eq (<ref>). 

 Similar to the procedure introduced in <ref>, we denote the parameter at the beginning of the current stage s as θ̃_t,s. For simplicity, when we consider the update within stage s, we just use θ̃_t to denote θ̃_t,s. Similar with Eq (<ref>), we define the  terms
 





    μ̃≜𝒢(ℳ_t, θ̃_t),   ṽ≜𝒢(𝒯_t, θ̃_t).


Then we can calibrate the gradient by ∇_θℓ(x_t,y_t,θ̃_t) and ∇_θℓ(x̅_t, y̅_t, θ̃_̃t̃) (which serves as the similar role of ∇_θℓ(x^i, y^i, θ̃ in (<ref>)); the new form of v_t^k becomes

    v_t^k     =1/t(∇_θℓ(x_t, y_t, θ_t^k) - ∇_θℓ(x_t,y_t,θ̃_t) + ṽ) + 
       t-1/t( ∇_θℓ(x̅_t, y̅_t, θ_t^k) -  ∇_θℓ(x̅_t, y̅_t, θ̃_̃t̃) + μ̃).


  
Obviously, if the buffer ℳ_t contains the whole historical data (denote by 𝒯_[1:t) = ⋃_c=1^t-1𝒯_c), the above approach is exactly the standard SVRG. However, because ℳ_t only takes a small subset of 𝒯_[1:t), this approach still cannot avoid information loss for the previous tasks. In next section, we propose a novel two-level dynamic algorithm to record more useful information from 𝒯_[1:t), and thereby reduce the information loss induced by ℳ_t. We also take the approach of Eq (<ref>) as a baseline in <ref> to illustrate the advantage of our proposed approach. 



 §.§ Dynamic Gradient Calibration





To tackle the issue discussed in <ref>,  we propose a novel two-level update approach “Dynamic Gradient Calibration (DGC)” to maintain our calibration term. Our focus is designing a method to incrementally update an unbiased estimation for 𝒢(𝒯_[1:t), θ_t^k). To illustrate our idea clearly, we decompose our analysis to two levels: (1) update the parameter during the training within each time spot t; (2) update the parameter at the transition from time spot t to  t+1 (i.e., the moment that the task 𝒯_t has just been completed and the task 𝒯_t+1 is just coming). 











  (1) How to update the parameter during the training within each time spot t. We follow the setting of the streaming SVRG as discussed in <ref>: 
  the training process at the current time spot t is divided into a sequence of stages;  the model parameter at the beginning of each stage is recorded as θ̃_t. 
  To illustrate our idea for calibrating the gradient v^k_t in Eq (<ref>), we begin by considering an “imaginary” approach: 
  
  we let
 

    v_t^k     =1/t( ∇_θℓ(x_t, y_t, θ_t^k) -∇_θℓ( x_t, y_t,θ̃_t) + 𝒢(𝒯_t,θ̃_t)) 
        + t-1/tΓ,

 where 

    Γ= ∇_θℓ(x̅_t, y̅_t, θ_t^k)- (∇_θℓ(x̅_t, y̅_t, θ̃_t)  - 𝒢(𝒯_[1:t), θ̃_t )). _Calibration from the previous parameter θ̃_t

Different from Eq (<ref>), we compute v^k_t based on the full historical data 𝒯_[1:t), which follows the same manner of SVRG. However, a major obstacle here is that we cannot obtain the exact Γ since 𝒯_[1:t) is not available. This motivates us to design a relaxed form of (<ref>).

We define a surrogate function to approximate Γ, which can be computed through recursion. Suppose each training stage has m≥ 1 steps, then we define

    Γ_DGC(t, k)     = ∇_θℓ(x̅_t, y̅_t, θ_t^k) - 
       (∇_θℓ(x̅_t, y̅_t, θ̃_t)  - Γ'_DGC(t) ), 
    Γ'_DGC(t)     = Γ_DGC(t, m+1),

 in the stage. Note that the term “∇_θℓ(x̅_t, y̅_t, θ̃_t)” in (<ref>) can be computed by the previous parameter θ̃_t during training, so we do not need to store it in buffer. For the initial  t=1 case (i.e., when we just encounter the first task),  we can directly set Γ'_DGC(1)=0. We update the function Γ'_DGC(t) at the end of each training stage in (<ref>), and use the function Γ_DGC(t, k) to approximate Γ in (<ref>). Comparing with the original formulation of Γ in (<ref>), we only replace the term “𝒢 (𝒯_[1:t), θ̃_t )” by “Γ'_DGC(t)”. Also, we have the following lemma to support this replacement.  The detailed proof of lemma <ref> is provided in our supplement.




    𝔼[Γ'_DGC(t) ] = 𝒢 (𝒯_[1:t), θ̃_t )



   We utilize the term “∇_θℓ(x̅_t, y̅_t, θ̃_t)  - Γ'_DGC(t)” of (<ref>) as the calibrator for each updating step, thereby preserving the unbiased nature of the gradient estimator and reducing the variance of gradient estimation.

  (2) How to update the parameter at the transition from  time spot t to t+1. 

 At the end of  time spot t, we update the recorded θ̃_t to θ_t^m+1, and the data 𝒯_t from time t should be integrated into the historical data. In this context, it is essential to update the calibrated gradient accordingly:







    Γ'_DGC(t+1) = 1/t((t-1) ·Γ'_DGC(t) +  𝒢(𝒯_t, θ̃_t) ) .





 The complete algorithm is presented in Algorithm <ref>. Compared with the conventional reservoir sampling based approaches, we only require the additional   
 
  storage for keeping Γ'_DGC(t), and so that the gradient Γ_DGC(t, k) in (<ref>) can be effectively updated by the recursion. Moreover, our method can also conveniently adapt to TFCL where the task boundaries are not predetermined. In such a setting, we can simply treat each batch data (during the SGD) as a “micro” task at the time point and then update the gradient estimation via Eq (<ref>).
  
  
  The detailed algorithm for TFCL is placed in our appendix. 

Similar to the theoretical analysis of SVRG <cit.>, under mild assumptions, the optimization procedure of our DGC method at each time spot t can also achieve linear convergence. We denote the optimal parameter at time spot t as θ_*≜min_θℓ_CL^t(θ). Then we have the following theorem. 













  
 Assume that f(x;θ) is  L-smooth and γ-strongly convex; the parameters 

m ≥10L^2/γ^2 and η = γ/10L. Then we have a linear convergence in expectation for the DGC procedure at time t:

    𝔼[θ̃_t,s+1-θ_*_2^2] ≤1/2^s𝔼[θ̃_t,1-θ_*_2^2]

where θ̃_t,s represents the initialization parameter at the beginning of the s-th stage at time spot t.


  The  proof of Theorem <ref> is provided in appendix. This theorem indicates that the gradient calibrated by our DGC method shares the similar advantages with SVRG. For instance, when updating each task 𝒯_t, the loss function has a smoother decrease (we validate this property in <ref>).





 §.§ Combine DGC with Other CL Methods




Our proposed DGC approach can be also efficiently combined with other CL methods. 
As discussed in Section <ref>, a number of popular CL methods rely on the reservoir sampling technique to preserve historical data in   buffer <cit.>. 



For these methods, such as DER and XDER, we can conveniently combine the conventional batch gradient descent with 

the DGC calibrated gradient estimator Γ_DGC(t,k) defined in <ref>, so as to obtain a more precise gradient estimator with reduced variance based on Eq (<ref>):

    v_t^k    =1/t(∇_θℓ(x_t, y_t, θ_t^k) - ∇_θ(x̅_t,y̅_t,θ̃_t) + 𝒢(𝒯_t,θ̃_t)  )+ 
       t-1/t[ αΓ_DGC(t,k) + (1-α)∇_θℓ(x̅_t, y̅_t, θ^k_t)],

where α is a given parameter to control the proportion of the two unbiased estimations Γ_DGC(t,k) and ∇_θℓ(x̅_t, y̅_t, θ^k_t) of the gradient 𝒢(𝒯_[1:t),θ_t^k). 
According to the theoretical analysis in SSVRG <cit.>, 

the selection of α should be related to 1/L, where L is the smoothness coefficient of the model f(x;θ), i.e., 

    L = max_θ_1,θ_2 ∈Θ|∇_θ f(x;θ_1) - ∇_θ f(x;θ_2)|/|θ_1-θ_2|,


where Θ represents the parameter space. The experimental study on the impact of α is placed in our supplement. In <ref>, we show that the amalgamation of the CL method and our DGC calibration procedure can yield a more precise update direction, and consequently enhance the ultimate model performance. 

























§ EXPERIMENTS



We conduct the experiments to compare with

various baseline methods across different datasets. We consider both the CIL and TFCL models. 










  Datasets We carry out the experiments on three widely employed datasets S(Split)-CIFAR10, S-CIFAR100 <cit.>, and S-TinyImageNet <cit.>. S-CIFAR10 is the split dataset by 
partitioning CIFAR10 into 5 tasks, each containing two categories; similarly, S-CIFAR100 and S-TinyImageNet are the datasets by respectively partitioning 
 CIFAR100 and TinyImageNet  into 10 tasks, each containing 10 (S-CIFAR100) and   20 (S-TinyImageNet) categories. 
 

  Baseline methods We consider the following baselines. (1) Replay-based methods: ER <cit.>, DER++
  <cit.>, XDER <cit.>, MOCA <cit.>, GSS <cit.>, GCR <cit.>, HAL <cit.>, and ICARL <cit.>. (2) Optimization-based methods: AGEM <cit.>, AOP <cit.>, and SSVRG <cit.>.
 
  (3) Dynamic architecture method: Dynamic ER <cit.>.  We integrate DGC with ER, DER++, XDER, and Dynamic ER, and assess their performances in CIL. For  TFCL, we consider its combination with ER and DER++.
For convenience, we use “DGC-Y” to denote the combination of  

 DGC  with a CL method “Y”. For example, DGC-ER denotes the method combining ER and DGC methods.
 

  
  
  

  Evaluation metrics

We employ the Average Accuracy (AA) <cit.> and the final Average Incremental Accuracy (FAIA) <cit.> to assess the performance. These two metrics are both widely used for continual learning. Let a_k,j∈ [0,1] (k≥ j) denote the classification accuracy evaluated on the testing set of the task 𝒯_j after learning  𝒯_k. The value AA at time spot i is defined as

    AA_i ≜1/i∑_j=1^i a_i,j. 

In particular, we  name the value AA_T as the final Average Accuracy (FAA).
The final AIA is defined as 

    FAIA≜1/T∑_i=1^TAA_i.




We also use Final Forgetting (FF) from <cit.> to measure the forgetting of the model throughout the learning process
(the formal definition of FF and its numerical results are placed in the supplement).
Each instance of our experiments is repeated by 10 times.









 §.§ Results in CIL



  Hyper-parameters selection In our implementation, we fixed the values of epoch and batch size, which implies that the total number of optimization steps (i.e., the value s× m) is also fixed. In our experiments, we set the value m=200, so the value of s in Algorithm <ref> is also determined.



 Except for the Dynamic ER (which is dynamically expanded), all other testing methods have constant storage limits. The results shown in Table <ref> reveal that our DGC method can bring improvements to the combined methods on the testing benchmarks.
 For the sake of clarity, we also underline the best-performing method except Dynamic ER and DGC-Dynamic ER in Table <ref>.
Among the methods with constant storage limits, DGC-ER achieves the best results on S-CIFAR10 when the buffer size is 2000, and DGC-XDER achieves the best results on S-CIFAR100 and S-TinyImageNet.
It is worth noting that our DGC method can also be conveniently integrated with existing advanced dynamic expansion representation techniques, such as Dynamic ER,  
which demonstrates the improvements  to certain extent, e.g.,  it achieves an improvement more than 6% on S-CIFAR10/100 with buffer size 2000.   We also record the training time of these baseline methods in our supplement.









  In the subsequent experiment, we investigate the performance of DGC compared to other baseline methods with varying the number of tasks. Throughout the experiment, we maintain a constant buffer size of 2000. As outlined in Table <ref>, our results demonstrate that DGC can bring certain improvements to ER, DER++, and XDER with setting the number of tasks to be 5 and 20.
  
Moreover, similar to the previous research in <cit.>,   the results shown in Table <ref> also imply that increasing the number of tasks could exacerbate the catastrophic forgetting issue. This phenomenon occurs because the model faces a reduced volume of data on each specific task, and thereby necessitates the capability of retaining the information of historical data to guide the model updates. 
As can be seen, the improvement obtained by DGC for the case of 20 tasks usually is more significant compared with the case of 5 tasks. For example, DGC-DER++ achieves a 3.94% improvement with 20 tasks versus a 1.18% improvement with 5 tasks, while DGC-XDER exhibits a 3.08% improvement with 20 tasks and only 0.42% with 5 tasks. These results highlight the advantage of DGC for mitigating catastrophic forgetting by effectively utilizing historical data through gradient-based calibration. 





 §.§ Results in TFCL



 
We then conduct the experiments in TFCL. The curves shown in Figure <ref> depict the average AA_t evolutions with varying the “implicit” task number t on S-CIFAR100. We call it “implicit” since the value t is not given during the training, which means that the design of the algorithm cannot rely on task boundaries or task identities. Therefore, we only compare those baseline methods that are applicable to TFCL model. The results suggest that DGC can bring improvements to ER and DER++ on AA_t for almost all the ts; in particular, DGC-DER++ achieves the best performance and also with small variances. 
Through approximating the full gradient, the GCR and SSVRG methods can relieve the catastrophic forgetting issue to a certain extent, but they still suffer from the issue of storage limit, which affects their effectiveness for estimating the full gradient. Consequently, these methods exhibit performance downgrade and larger variance in Figure <ref>. Comparing with them, the approaches integrated with DGC illustrate more consistent and stable improvements.
More detailed results for TFCL are available in our supplement.



 §.§ Smoothness of Training with DGC





  Since our proposed DGC approach stems from the variance reduction method of SGD, a natural question is whether it also improves the smoothness during the training process.
  We  compare the loss trajectories on the entire training dataset before and after implementing DGC.  From Figure <ref> we can see that the classical ER method has erratic fluctuations in loss during the training process; this could cause some practical problems in a real-world CL scenario, e.g., 
  we may need to pay more effort to carefully adjust the learning rate and determine the stopping condition.
 
  In contrast, the  DGC method has a smoother reduction in loss, and ultimately yields lower loss values.
  



§ CONCLUSION


In this paper, we revisit the experience replay method and aim to utilize historical information to derive a more accurate gradient for alleviating catastrophic forgetting. Inspired by the variance reduction methods for SGD, we introduce a new approach “DGC” to dynamically manage a gradient calibration term in CL training. Our approach can be conveniently integrated with several existing continual learning methods, contributing to a substantial improvement in both CIL and TFCL. Moreover, the improved stability of training loss reduction can also ease our practical implementation. 


 

§ ACKNOWLEDGEMENTS

We want to thank the anonymous reviewers and Xianglu Wang for their helpful comments. The research of this work was supported in part by the National Key Research and Development Program of China 2021YFA1000900, the National Natural Science Foundation of China 62272432, and the Natural Science Foundation of Anhui Province 2208085MF163. 
 





§ IMPACT STATEMENT

This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.


icml2024











§ OMITTED PROOFS






    𝔼[Γ'_DGC(t) ] = 𝒢 (𝒯_[1:t), θ̃_t ).




We prove it by induction. When t = 1,  both the two-side terms are 0, so the equation holds. Assume the equation holds when t = t_1, and we consider the change of  Γ'_DGC(t) at the time spot t_1.

Firstly, Γ'_DGC(t_1) is updated at the end of every stage.  After the update, Γ'_DGC(t_1)=Γ_DGC(t_1,m+1). Meanwhile, we have

    𝔼[Γ_DGC(t_1,m+1)] 
        = 𝔼[∇_θℓ(x̅_t_1, y̅_t_1, θ_t_1^m+1) - (∇_θℓ(x̅_t_1, y̅_t_1, θ̃_t_1)  - Γ'_DGC(t_1))] 
       = 𝒢(𝒯_[1:t_1),θ_t_1^m+1) - 𝒢(𝒯_[1:t_1),θ̃_t_1) + 𝒢(𝒯_[1:t_1),θ̃_t_1) 
       = 𝒢(𝒯_[1:t_1),θ_t_1^m+1) 
       = 𝒢(𝒯_[1:t_1),θ̃_t_1)  ,

where the first equation comes from Eq (<ref>) and the second equation is based on the inductive hypothesis.

  Secondly, in the update from  time spot t to t+1,  we have

    𝔼[Γ'_DGC   (t_1+1)] 
       = 𝔼[1/t_1((t_1-1) ·Γ_DGC(t_1,m+1) +  𝒢(𝒯_t_1, θ̃_t_1) )] 
       = 1/t_1((t_1-1) ·𝒢(𝒯_[1:t_1),θ̃_t_1) +  𝒢(𝒯_t_1, θ̃_t_1) ) 
       = 𝒢(𝒯_[1:t_1),θ̃_t_1+1)

 based on Eq (<ref>).
This implies that our conclusion also holds when t = t_1+1.



























































































































































Assume that f(x;θ) is  L-smooth and γ-strongly convex; the parameters 

m ≥10L^2/γ^2 and η = γ/10L. Then we have a linear convergence in expectation for the DGC procedure at time t:

    𝔼[θ̃_t,s+1-θ_*_2^2] ≤1/2^s𝔼[θ̃_t,1-θ_*_2^2]

where θ̃_t,s represents the initialization parameter at the beginning of the s-th stage at time spot t.





Our proof  is inspired by the idea of <cit.>. To simplify the notations, we define 

    v_t^k = ∇ψ_x_t(θ_t^k)-∇ψ_x_t(θ̃_t)+μ̃
 through combining Eq (<ref>) and (<ref>) (we replace “Γ” in Eq (<ref>) by (<ref>)), where

    ∇ψ_x_t(θ_t^k)    ≜1/t∇_θℓ(x_t,y_t,θ_t^k) + t-1/t∇_θℓ(x̅,y̅,θ_t^k) , 
    ∇ψ_x_t(θ̃_t)    ≜1/t∇_θℓ( x_t, y_t, θ̃_t) + t-1/t∇_θℓ(x̅,y̅,θ̃_t) , 
    μ̃   ≜1/t𝒢(𝒯_t, θ̃_t) + t-1/tΓ^'_DGC(t) .


Before proving Theorem 3.3, we provide the following key lemmas first.

 

    𝔼[v_t^k_2^2] ≤ 2 L[𝔼[||θ_t^k-θ_*||_2^2]+𝔼[||θ̃_t-θ_*||_2^2]].




    𝔼[v_t^k_2^2] 
    ≤    2 𝔼[∇ψ_x_t(θ_t^k)-∇ψ_x_t(θ_*) _2^2]
        +2 𝔼[∇ψ_x_t(θ̃_̃t̃)-∇ψ_x_t(θ_*)  -∇ℓ_CL^t(θ̃_t)_2^2] 
    
    =     2 𝔼[ ∇ψ_x_t(θ_t^k)-∇ψ_x_t(θ_*)_2^2]
       +2 𝔼 [∇ψ_x_t(θ̃_t)-∇ψ_x_t(θ_*) 
       -𝔼[∇ψ_x_t(θ̃_t)-∇ψ_x_t(θ_*)] _2^2] 
    ≤    2 𝔼[∇ψ_x_t(θ_t^k)-∇ψ_x_t(θ_*)_2^2]
       +2 𝔼[∇ψ_x_t(θ̃_t)-∇ψ_x_t(θ_*)_2^2] 
    ≤    2 L[𝔼[||θ_t^k-θ_*||_2^2]+𝔼[||θ̃_t-θ_*||_2^2]],

where the first inequality comes from the fact  a+b_2^2≤ 2a_2^2+2b_2^2  for any vectors a,b  and  Lemma <ref>,  the second inequality is based on the fact  𝔼[ξ-𝔼ξ_2^2]=𝔼[ξ_2^2]-𝔼 [ξ]_2^2≤𝔼[ξ_2^2]  for any random vector  ξ, and the third inequality is based on L-smooth.



 

    𝔼[θ_t^k-θ_*_2^2] ≤ (1-2ηγ+2Lη^2)𝔼[||θ_t^k-1-θ_*||_2^2] 
       + 2Lη^2𝔼[||θ̃_t-θ_*||_2^2].



Note that 𝔼 [v_t^k]= 𝔼 [∇ψ_x_t(θ_t^k)] =∇ℓ_CL^t(θ_t^k-1), and so we have

    𝔼[θ_t^k-θ_*_2^2] 
    
    =    θ_t^k-1-θ_*_2^2-2 η(θ_t^k-1-θ_*)^⊤𝔼 [v_t^k-1]+η^2𝔼[v_t^k-1_2^2] 
    ≤   θ_t^k-1-θ_*_2^2-2 η(θ_t^k-1-θ_*)^⊤∇ℓ_CL^t(θ_t^k-1)
        +2 L η^2[𝔼[||θ_t^k-1-θ_*||_2^2]+𝔼[||θ̃_t-θ_*||_2^2]] 
    ≤   θ_t^k-1-θ_*_2^2-2 ηγ ||θ_t^k-1-θ_*||_2^2
        +2 L η^2[𝔼[||θ_t^k-1-θ_*||_2^2]+𝔼[||θ̃_t-θ_*||_2^2]],

where the first inequality comes from  Lemma <ref>, and the second inequality is based on the γ-strong convexity of  ℓ_CL^t(θ). We take the expectation on θ_t^k-1 from the above inequality and then have

    𝔼[θ_t^k-θ_*_2^2] 
       ≤𝔼 [θ_t^k-1-θ_*_2^2]-2 ηγ𝔼[||θ_t^k-1-θ_*||_2^2] 
        +2 L η^2[𝔼[||θ_t^k-1-θ_*||_2^2]+𝔼[||θ̃_t-θ_*||_2^2]] 
        = (1-2ηγ+2Lη^2)𝔼[||θ_t^k-1-θ_*||_2^2] + 2Lη^2𝔼[||θ̃_t-θ_*||_2^2].






We now proceed to prove the theorem. To simplify notation, we denote 𝔼θ_t^k-θ_*_2^2 as p_k.

Based on lemma <ref>, we have

    p_k≤(1-2 ηγ+2 η^2 L) p_k-1+4 η^2 L  p_k-1.

By setting  η = γ/10L and m ≥10L^2/γ^2, we have

    p_m+1   ≤(1-γ^2/10 L^2) p_m≤(1-γ^2/10 L^2)^m p_1
       ≤exp(-γ^2/10 L^2 m) p_1≤1/2 p_1,

i.e

    𝔼[θ_t^m+1-θ_*_2^2] ≤1/2𝔼[θ_t^1-θ_*_2^2].

Subsequently, we consider the update during different stage s, so that the initial parameter θ̃_t,s is θ_t^1  and updated parameter θ̃_t,s+1 = θ_t^m+1  is selected after all of the updates have been completed. Then Eq (<ref>) becomes

    𝔼[θ̃_t,s+1-θ_*_2^2] ≤1/2𝔼[θ̃_t,s-θ_*_2^2].

Through recursively applying  the above equation from stage 1 to stage s+1, we have 

    𝔼[θ̃_t,s+1-θ_*_2^2] ≤1/2^s𝔼[θ̃_t,1-θ_*_2^2].





§ DETAILED DGC PROCEDURE IN TFCL

We consider the data streams {𝒯_1, 𝒯_2, ⋯ 𝒯_T} with each 𝒯_i being the batch training data in TFCL. Algorithm <ref> shows how to implement the DGC procedure in this model. 












§ DETAILS FOR EXPERIMENTAL SETUP


  Architecture and hyperparameter Most of the comparison methods in our experiments use the implementation of Mammoth [https://github.com/aimagelab/mammoth]. For the Dynamic ER <cit.> method we use the implementation of PyCIL [https://github.com/G-U-N/PyCIL], and for the GCR method we use the implementation of Google-research [https://github.com/google-research/google-research/tree/master/gradient_coresets_replay]. To be consistent with the selection of these open source frameworks, all the methods use ResNet18 <cit.> as the base network. All the networks are trained from scratch. 
For the hyperparameter selection of different methods, we directly use the original hyperparameters used in these open-source frameworks, which are obtained by using grid search on 10% of the training set.


  Training details 
  
  To maintain consistency, we utilize the default parameter settings and network architectures provided by the framework. As in previous studies <cit.>, random crops and horizontal flips are used as data augmentation in all experiments. All methods that do not incorporate DGC are optimized using standard SGD. In our DGC calibrated method, we empirically set m = 200. 
 
  To fairly compare the methods with constant storage limits, we uniformly train for 50 epochs in each task on S-CIFAR10 and S-CIFAR100, and 100 epochs in each task on S-TinyImageNet. The batch size is all set to  be 32. According to the experimental description in the Dynamic ER method, 
  
  we train the first task for 200 epochs on all datasets and train all subsequent tasks for 170 epochs;    the batch size is set to be 128.



§ OTHER EXPERIMENTAL RESULTS



  Selection of α We explore the impact of different α values to the performance of our DGC method. In the experiment, we fix the buffer size to 2000. Table <ref> shows that DGC can improve ER under different α values. 
  
  In all the experiments of <ref>, we fix the value of α to be 1e^-3.

  Selection of m As stated in Theorem <ref>, when m is sufficiently large, our DGC method has a linear convergence in expectation. The results in Table <ref>  show that when m is greater than 100, DGC-ER can significantly improve the performance of ER. Table <ref> also suggests that our method is relatively robust to the selection of m. 





  FF and FAA results in CIL 
  we use Final Forgetting (FF) <cit.> to measure the forgetting of the model throughout the learning process.
Following the previous work, it is defined as FF≜1/T-1∑_j=1^T-1(max _k ∈{1, …, T-1} a_k,j-a_T,j).
  Table <ref> and Table <ref> show that DGC based methods can reduce forgetting while improving the performance in all the cases.  For most instances, our DGC based methods achieve the best performance.






Detailed results in TFCL Table <ref> shows that our DGC method can improve the performance in TFCL. The DGC-DER++ method achieves the best performance in all the cases.

  Training Time We set b=500 and conduct experiments on S-CIFAR100 to measure the training time of several algorithms on the first three tasks in the CIL scenario. Table <ref> shows that our method and GCR have roughly the same training time, which is higher than AGEM but lower than other baselines. The results on the remaining tasks and other datasets are similar. Therefore, the increase caused by our method on training time is not significant, compared with most of the baselines.  






 Paper Title: Conformal Trajectory Prediction with Multi-View Data Integration in Cooperative Driving.


Beyond full statevector simulation with 
    Andrea Pasquale^1,2, Andrea Papaluca^3, Renato M. S. Farias^1,4, Matteo Robbiati^2,5, Edoardo Pedicillo^1,2,
and Stefano Carrazza^1,2,5
    August 5, 2024
===========================================================================================================================================



§ ABSTRACT 
Current research on trajectory prediction primarily relies on data collected by onboard sensors of an ego vehicle. With the rapid advancement in connected technologies, such as vehicle-to-vehicle (V2V) and vehicle-to-infrastructure (V2I) communication, valuable information from alternate views becomes accessible via wireless networks. The integration of information from alternative views has the potential to overcome the inherent limitations associated with a single viewpoint, such as occlusions and limited field of view. In this work, we introduce V2INet, a novel trajectory prediction framework designed to model multi-view data by extending existing single-view models. Unlike previous approaches where the multi-view data is manually fused or formulated as a separate training stage, our model supports end-to-end training, enhancing both flexibility and performance. Moreover, the predicted multimodal trajectories are calibrated by a post-hoc conformal prediction module to get valid and efficient confidence regions. We evaluated the entire framework using the real-world V2I dataset V2X-Seq. Our results demonstrate superior performance in terms of Final Displacement Error (FDE) and Miss Rate (MR) using a single GPU. The code is publicly available at: <https://github.com/xichennn/V2I_trajectory_prediction>.









§ INTRODUCTION



Trajectory prediction plays a critical role in autonomous driving. Typically, we rely on the on-board sensors of an ego vehicle to gather environmental information necessary for performing various autonomous driving tasks. However, with the rapid advancement in connected technologies, such as vehicle-to-vehicle (V2V) and vehicle-to-infrastructure (V2I) communication, valuable information from alternate views becomes accessible via wireless networks. The integration of information from alternative views has the potential to overcome the inherent limitations associated with a single viewpoint, such as occlusions and limited field of view. We categorize the data obtained solely from ego vehicle on-board sensors or infrastructure as single-view information, whereas the accessibility to multiple viewpoints is referred to as multi-view data. Figure <ref> depicts an intersection scenario where the AV faces a potential left-turn collision due to its field of view being obstructed by a large truck. Roadside cameras, strategically positioned to have an unobstructed view of the entire intersection, provide a comprehensive overview of the traffic situation. This enhanced perspective allows the AV to receive critical information and effectively avoid the impending collision.

From a trajectory predictive modeling perspective, significant research efforts have been dedicated to leveraging single-view datasets collected by the ego vehicle <cit.>. These efforts typically involve modeling the temporal dependencies, agent-agent interactions and agent-lane relations. However, the challenge arises when dealing with multi-view data, particularly in how to effectively fuse the information due to overlapping field of views.

Existing work on multi-view data fusion in the context of cooperative driving predominantly focuses on collaborative perception tasks, with limited research addressing trajectory prediction. <cit.> pioneered the creation of the first V2I real-world dataset for trajectory prediction studies. They manually fused multi-view data using trajectory association and stitching techniques at each frame, followed by the application of single-view trajectory prediction models. Although this approach is intuitive, it fails to fully exploit the motion behavior captured by data from each view, resulting in suboptimal outcomes. To address this limitation, <cit.> proposed a novel approach that encodes trajectory information from each view as independent graph nodes, thereby minimizing information loss. They formulated the node association problem as a graph link prediction task and introduced a cross-attention module to fuse node embeddings from associated nodes across different views. While their model can be trained end-to-end, the node association process necessitates pretraining, leading to the formulation of two optimization objectives.

We introduce a trajectory prediction framework that utilizes multi-view data without the need for explicit association between different perspectives. Rather than developing a specialized multi-view model, our approach seamlessly integrates with state-of-the-art single-view trajectory models, maximizing the utility of existing research efforts. No special training strategies are required. We can easily leverage the pretrained single-view trajectory models to expediate the training. Specifically, for trajectory data from each single-view, we employ established graph neural network (GNN) based models, such as LaneGCN <cit.>, HiVT <cit.>, to capture temporal dependencies, agent-agent interactions, and agent-lane relations. Then we utilize a cross-graph attention module to fuse the node embeddings from different views. The fused final embeddings will then go through a multimodal decoder to get future trajectory predictions. 

Existing works have modeled the multi-modality explicitly by introducing anchors <cit.>, mixture models <cit.>, or implicitly through latent variables such as Conditional Variational Auto-encoder (CVAE) or generative models <cit.>. The implicit models often face the issue of mode collapse, therefore we will model the multi-modality by a mixture model built upon MLP. It is a common strategy to assign a higher score to the modality closer to the ground truth during training. This strategy, however, may encounter robustness issues during inference. Figure <ref> presents two examples where the top-scored prediction is not the closest to the ground truth trajectory.



 This arises from the inherent randomness in the ground truth data, posing challenges for the model to learn the priorities of multiple predictions. Relying on the scores to rank the predictions can lead to hazardous results in safety-critical scenarios. Additionally, while most mixture models like GMM provide both point estimates and variance estimation, the reliability of these uncertainty quantification (UQ) remains uncertain. Unfortunately, this aspect is often overlooked in the model evaluation process. However, UQ can offer valuable insights for downstream decision-making processes. To address these challenges, we introduce conformal prediction (CP) <cit.>, as a post-hoc module to our framework. CP creates statistically rigorous uncertainty intervals for the predictions that are guaranteed to contain the ground truth with a user-specified probability. 

To summarize, our contributions in this work lie in twofolds:
1) We introduce V2INet, a novel trajectory prediction framework designed to model multi-view data by extending existing single-view models. Unlike previous work, our model supports end-to-end training, enhancing both flexibility and performance.
2) We introduce conformal prediction as a post-hoc module to calibrate the prediction results. This results in statistically rigorous uncertainty intervals, significantly enhancing the reliability of the predictions. These calibrated predictions are particularly beneficial for downstream decision-making tasks, such as motion planning.



§ RELATED WORK



 §.§ Cooperative Driving Dataset

Collaborative perception has been the most extensively studied area in collaborative driving <cit.>. It leverages the wireless communication technologies, such as V2V and V2I, to share information, which enables perception beyond line-of-sight and field-of-view hence overcoming common perceptual shortcomings with individual perception, such as blind spot, occlusions and long-range issues <cit.>. Many datasets have been collected in both simulation environment <cit.> and real-world <cit.> to facilitate the study. A comprehensive comparison can be found in Table <ref>. While majority of them focused on the upstream tasks such as detection, tracking and segmentation, to enable well-informed decision-making for autonomous vehicles, however, it's critical to also incorporate V2X data for predicting the behavior of surrounding traffic participants. V2X-Seq/Forecasting <cit.> and V2X-Traj <cit.> present two real-world datasets specifically designed for cooperative trajectory prediction tasks. The experiments in V2X-Seq/Forecasting demonstrate that leveraging the infrastructure-side trajectories can enhance the trajectory prediction performance. V2X-Traj aims for more general V2X scenarios, extending to V2V cooperation.




 §.§ Cooperative Information Association and Fusion







The main challenge in cooperative trajectory prediction lies in the multi-view data source fusion. In cooperative scenarios, vehicles gather safety-related data using sensors like radar, lidar, cameras, and GPS. This data is standardized into Basic Safety Message (BSM) format, ensuring compatibility across vehicles and infrastructure <cit.>. BSM messages, containing crucial information such as position and speed, are broadcasted periodically. Upon receiving BSM messages, vehicles combine this data with their own sensor data to enhance accuracy. Advanced algorithms have been developed to associate and fuse the multi-source data. Given the temporal and spatial dimensions of the collected trajectories, there has been research focusing on the communication delays alignment <cit.> and pose errors alignment <cit.>. Based on the aligned data, <cit.> utilized CBMOT <cit.>, a multi-object tracking method, to fuse the infrastructure and ego-vehicle trajectories at each single frame and then trained the network taking in the fused dataset. While straightforward, this method fails to capture the motion behavior provided by the infrastructure across all time steps, leading to suboptimal results. In contrast, <cit.> encoded trajectory information from each view as independent graph nodes, formulating the association process as a graph linking problem. While effective, this approach necessitates separate training procedures. Building upon their work, we propose the utilization of a cross-graph attention mechanism to fuse multi-view information, eliminating the need for additional training processes.


 §.§ Uncertainty Quantification

In the trajectory prediction task, numerous sources of uncertainty exist, including inherent multi-modality, partial observability, short time scales, data limitations, intention type imbalances, and domain gaps. Most existing trajectory prediction models address uncertainty by maximizing the likelihood of an assumed distribution, such as Gaussian or Laplace <cit.>. However, trajectories with the largest likelihood are often nonsensical <cit.>. Alternatively, some research focuses on approximating Bayesian inference for deep learning models using techniques like Monte Carlo dropout <cit.>, which involves performing stochastic forward passes through the network and averaging the results. Among these approaches, <cit.> stands out as the only work that incorporates collaborative uncertainty among agents into the modeling process to guide the ranking of multimodal trajectories by uncertainty, albeit requiring special training strategies. However, none of the predicted uncertainties from these methods offer finite sample coverage guarantees, which is suboptimal for safety-critical applications such as vehicle trajectory prediction.

Conformal prediction (CP) <cit.> has emerged as a widely adopted uncertainty quantification method, owing to its simplicity, generality, theoretical rigor, and low computational overhead. Notably, CP is agnostic to the underlying model and data distribution, making it highly versatile. It seamlessly integrates with any pre-trained model to deliver statistically valid prediction regions. Of particular relevance to our multimodal trajectory prediction task is recent progress in generalizing CP to time-series forecasting. For instance, <cit.> introduced the Copula conformal prediction algorithm for multivariate, multi-step time series forecasting, applicable to any multivariate multi-step forecaster. Additionally, <cit.> focused on generating non-conformity score functions that yield multimodal prediction regions with minimal volume. Moreover, <cit.> employed CP to generate statistical uncertainty intervals from Gaussian mixture model outputs, obtaining separate prediction intervals corresponding to each GMM component prediction. Furthermore, <cit.> proved the validity of CP on graph data. Inspired by their work, we explore the potential of applying CP to the multimodal trajectory prediction comparing different CP methods.



§ PROBLEM FORMULATION


At the scenario level, We have trajectory data 𝒯 from both the vehicle and infrastructure viewpoints, denoted as 𝒯^𝒱 and 𝒯^ℐ, respectively. While 𝒯^𝒱 and 𝒯^ℐ share overlapping information where their fields of view intersect, 𝒯^ℐ also provides complementary information, being free from occlusions. Our modeling objective is to utilize the information from 𝒯^ℐ to improve the accuracy of trajectory prediction based solely on 𝒯^𝒱. 

The trajectory prediction task involves leveraging historical trajectories 𝒯^𝒱_h ∈ℝ^N^𝒱× T_h × a_h and 𝒯^ℐ_h ∈ℝ^N^ℐ× T_h × a_h, alongside contextual information, typically HD maps denoted as ℳ, to forecast future trajectories 𝒯^𝒱_f ∈ℝ^N^𝒱× T_f × a_f. Here, N^𝒱 and N^ℐ represent the number of observed actors from the vehicle and infrastructure perspectives, respectively. T_h denotes the historical time horizon, and T_f is the prediction horizon. a_h and a_f represent the number of node features which we consider the vehicle center location defined by its x- and y-coordinates. Notably, the trajectory data from both views are defined within the same coordinate system. For the HD map, we opt for a vectorized representation due to its lightweight nature and efficiency <cit.>. This vector map is depicted by lane centerlines, which are composed of lane segments. We denote it as ℳ∈ℝ^N^l × a_l, where N^l is the number of lane segments and a_l is the number of lane attribute.

We formulate the overall probabilistic distribution as P(𝒯^𝒱_f | 𝒯^𝒱_h, 𝒯^ℐ_h, ℳ^𝒱, ℳ^ℐ). Driven by the critical safety demands inherent in trajectory prediction, we aim to incorporate uncertainty quantification to preempt any potentially consequential model failures. Let 𝒟 be the set of scenarios of the form (𝒯^𝒱_h, 𝒯^ℐ_h, ℳ^𝒱, ℳ^ℐ, 𝒯^𝒱_f). We split the dataset into training 𝒟_train, validation 𝒟_val, calibration 𝒟_cal and test 𝒟_test. One black-box deep learning model is trained on 𝒟_train and evaluated on 𝒟_val. We achieve the uncertainty quantification by conformal prediction in a post-hoc way on 𝒟_cal.  

At the agent level, we denote the features and labels of agent i from the vehicle view as X^𝒱 = 𝒯^𝒱_h,i and Y^𝒱 = 𝒯^𝒱_f,i for brevity. In real-world scenarios, future trajectories may exhibit multimodal behavior, often approximated by mixture models, resulting in Ŷ^𝒱∈ℝ^K × T_f × a_f, where K represents the number of mixtures or modes. Given a new agent sample X^𝒱_test from 𝒟_test, we seek to construct the prediction intervals 𝒞(X^𝒱_test) ∈ℝ^K × T_f × a_f × 2 such that it covers the ground truth label Y^𝒱_test under a predefined coverage rate leveraging conformal prediction.

CP proceeds in three steps. First, we define a nonconformity score A: 𝒳×𝒴∈ℝ^K × T_f × a_f to quantify how well Y conforms to the prediction at X. Typically, we choose a metric of disagreement between the prediction and the ground truth as the non-conformity score, such as the Euclidean distance. Second, given the predefined miscoverage rate α, we compute the 1-α quantile of the non-conformity scores on the calibration set {A(X_1, Y_1),⋯,A(X_n, Y_n)}, where n is the number of agents. The resulting quantile is denoted as Ĥ∈ℝ^T_f × a_f. Last, when presented with a new test agent X^𝒱_test, CP constructs the prediction interval 𝒞(X^𝒱_test)={Y^𝒱∈𝒴: A(X^𝒱_test, Y^𝒱) ≤Ĥ}.




§ METHODOLOGY



 §.§ Overview

Our method V2INet consists of two key components, predictive modeling and post-hoc conformal prediction. An overview of our proposed model is illustrated in Figure <ref>. 
We first represent the scenario data collected from both views as graphs. A single view encoder is then applied separately to each graph, encoding various information such as agent-agent interactions, temporal dependencies, and agent-lane information. Subsequently, The vehicle-view embedding is fused with the infrastructure-view embedding through a cross-graph attention module. Finally, the updated vehicle-view embedding pass through a multi-modal decoder, providing multimodal predictions for all the agents of interest. The post-hoc conformal prediction is then applied at the agent level to construct valid prediction intervals given a predefined coverage rate.




 §.§ Scene Representation

We adopt an ego-centric coordinate system utilizing vectorized representation as first introduced in <cit.>. First, data from both views are transformed such that the ego vehicle is centered at the origin, with its heading aligned along the positive x-axis. Each trajectory is then characterized as a sequence of displacements  {Δ𝐱^t}^0_t=-(T_h-1), where Δ𝐱^t∈ℝ^2 is the 2D displacement from time step t-1 to t. Similarly, each lane segment is represented as Δ𝐱^l ∈ℝ^2, which captures the 2D displacement from the starting coordinate to the end coordinate of lane segment l. We then construct scenario graphs consisting of actor nodes and lane nodes for both views separately. The edge attribute is the absolute relative position between two nodes.



 §.§ Single View Encoder

To encode the spatiotemporal information, including agent-agent interactions, temporal dependencies, and context information captured by agent-lane relations for data from each view, existing graph-based models offer effective solutions. Models like LaneGCN <cit.>, HiVT <cit.>, and VectorNet <cit.> incorporate these components and can be readily employed. This encoding step can be performed efficiently at the edge device, such as the vehicle's on-board computers and the roadside unit (RSU) for infrastructure data, before broadcasting, thereby minimizing computational overhead.

We exemplify here with HiVT <cit.> where only attention mechanism is employed, with the adoption of a rotation-invariant representation. To attend to the local information, at each timestamp, the surrounding actors' information is aggregated and the embeddings at each time stamp subsequently go through a transformer to capture the temporal dependencies. After obtaining the spatio-temporal embeddings, the surrounding lane information is aggregated for each actor. The update actor node embeddings is then passed to a global interaction module which aggregates the local context of different actors and updates each actor’s representation to capture long-range dependencies and scene-level dynamics. Here's the proofread version:

We denote the final embedding for actor node i from the vehicle view as 𝐡_i, and actor node j from the infrastructure view as 𝐡_j. Both 𝐡_i and 𝐡_j are in ℝ^d_h, where d_h is the embedding dimension.



 §.§ Cross-graph Attention Module

To effectively utilize information from the infrastructure side, we employ a cross-graph attention module that aggregates information captured by the infrastructure view encoder. Specifically, the embedding from the vehicle view 𝐡_i is transformed into the query vector, while the embedding from the infrastructure view 𝐡_j alongside the relative position at the last observed time step 𝐱_i^t=0 - 𝐱_j^t=0, are utilized to compute the key and value vectors. We denote 𝐡_ij=(𝐡_j, 𝐱_i^t=0 - 𝐱_j^t=0) representing the concatenation of node and edge attribute from infrastructure node j to vehicle node i:

    𝐪_i=𝐖^Q_f𝐡_i,    𝐤_ij=𝐖^K_f𝐡_ij,   𝐯_ij=𝐖^V_f𝐡_ij

where 𝐖^Q_f, 𝐖^K_f, 𝐖^V_f∈ R^d_k × d_h are learnable matrices for linear projection and d_k is the transformed dimension. The resulting query, key and value vectors are then taken as input to the scaled dot-product attention block:

     
           α_ij = softmax(𝐪_i^T/√(d_k) ·[{𝐤_ij}_j ∈𝒩_i]),

           𝐦_i = ∑_j ∈𝒩_iα_ij𝐯_ij,

           𝐠_i = sigmoid(𝐖^gate[𝐡_i,𝐦_i]),

           𝐡̃_i = 𝐠_i ⊙𝐖^self𝐡_i+(1-𝐠_i)⊙𝐦_i
    

where 𝒩_i is the set of agent i's neighbors, 𝐖^gate and 𝐖^self are learnable matrices and ⊙ denotes element-wise product. We followed the structure in HiVT to fuse the infrastructure information with a gating function. The attention block supports multiple heads. Finally, we apply another MLP block to obtain the final fused embedding 𝐡̃_i ∈ℝ^d_h for agent i from the vehicle view. 


 §.§ Mixture Model Based Decoder and Learning

There are two widely used mixture models for describing the multimodal trajectories, Gaussian and Laplacian. Previous methods <cit.> have found that the ℓ_1-based loss function derived from the Laplace distribution usually leads to superior prediction performances, as it is more robust to outliers. Hence, we will parameterize the future trajectories following Laplace distribution. For each agent, the decoder receives the final embedding as inputs and outputs K possible future trajectories and the mixing coefficient of the mixture model for each agent. The decoder are consisted of three MLPs, one for predicting the future locations μ^t_i,k∈ℝ^2 for agent i and its mode k at each time step t, one for predicting the associated uncertainty 𝐛^t_i,k∈ℝ^2 assuming independence of the x- and y-coordinates, the last one followed by a softmax is for producing the scores for each mode. 

To ensure the prediction diversity <cit.>, instead of optimizing all the predicted trajectories, only the mode k̃ closest to the ground truth is optimized. The closeness here is defined as the average Euclidean distance between ground-truth locations and predicted locations across all future time steps. The Loss includes both regression loss and classification loss

    
    J = J_reg + ε J_cls


Here, ε is the weight of the classification loss. We employ the negative log-likelihood as the regression loss:

    
    J_reg = -1/nT_f∑^n_i=1∑^T_f_t=1log𝐏(𝐲_i^t - 𝐲̂_i^t | μ̂^t,k̃_i,𝐛̂^t,k̃_i)



where 𝐏(|) is the probability density function of Laplace distribution and μ̂_i^t,k̃, 𝐛̂_i^t,k̃ are the mean and uncertainty estimates of the best mode k̃.

For J_cls, cross-entropy loss is applied:

    
    J_cls = -1/n∑^n_i=1∑^K_k=11_k logπ_k


where 1_k is a binary indicator if k is the best mode, π_k is the mixing coefficient of mode k. The model is trained on 𝒟_train and evaluated on 𝒟_val.



 §.§ Post-hoc uncertainty quantification module

Uncertainty quantification methods are evaluated on two key properties: validity and efficiency. Validity is established when the predicted confidence level exceeds or equals the probability of events falling within the predicted range, while efficiency refers to minimizing the size of the confidence region. We utilize conformal prediction to obtain both valid and efficient prediction intervals. The standard conformal prediction methods typically operate with scalar point estimates for regression problems. However, since our output consists of a multimodal multivariate time-series, we need to make certain adaptations to accommodate this complexity. Following the steps outlined in Section <ref>, we proceed with non-conformity score function definition, quantile computation and prediction interval construction.


  §.§.§ Non-conformity Score Functions

As emphasized in <cit.>, the usefulness of the prediction sets is primarily determined by the score function, we adopt three score functions A and compare their performance.

Z-score. As the decoder returns both mean and variance predictions, we define the Z-score function as:

    Z = |Y -Ŷ|/B̂


where Z ∈ℝ^K × T_f × 2 and B̂ = {{b̂^̂t̂,̂k̂}_t=0^T_f}_k=0^K. Here, we compute scores separately for x- and y-coordinates, reflecting the observation that motion uncertainty varies significantly across different dimensions.

L2-norm. Next, we consider the Euclidean distance, the most commonly used metric in regression problems:

    L_2 = ||Y -Ŷ||_2

where L_2 ∈ℝ^K × T_f.

L1-norm. Recognizing that the L2-norm disregards dimension differences, we consider L1-norm:

    L_1 = ||Y -Ŷ||_1

where L_1 ∈ℝ^K × T_f × 2.



  §.§.§ Quantile Computation

Given a predefined miscoverage rate α∈ [0,1], the 1-α quantile of the non-conformity scores is calculated on the calibration set 𝒟_cal. Conventionally, the quantile is determined as follows: Ĥ=quantile({A(X_1, Y_1),⋯,A(X_n, Y_n)}, (1-α)(1+1/n)) where A(X_i, Y_i) is a scalar, and n denotes the total number of agents in 𝒟_cal. However, in our case, we deal with multimodal time-series scores, which are multivariate when using Z-score and L1-norm, hence necessary adaptations are needed. For brevity, let Γ_i = A(X_i, Y_i), with Γ_i ∈ℝ^K × T_f × 2 for Z-score and L1-norm, and ℝ^K × T_f for L2-norm.

To enhance the efficiency of our prediction intervals, we focus on computing the quantile using only the mode k̃ that exhibits the smallest average Euclidean distance to the ground truth trajectory. This reduces the computation to multivariate time-series. Within this framework, we investigate two established methods: CF-RNN <cit.> and CopulaCPTS <cit.>.

CF-RNN. Since the time-series predictions are obtained from the same embedding, this work proposes the application of Bonferroni correction to the calibration scores to maintain the desired miscoverage rate α. Specifically, the original α is divided by T_f, yielding Ĥ=quantile({Γ_i}_i=0^n, (1-α/T_f)(1+1/n)) for single-variate case.

CopulaCPTS. As implied by its name, this method models the joint probability of uncertainty for multiple predicted time steps using a copula. The calibration set 𝒟_cal is split into two subsets: 𝒟_cal-1, which estimates a Cumulative Distribution Function (CDF) for the nonconformity score of each time step, and 𝒟_cal-2, utilized to calibrate the copula. The copula function captures the dependency between time steps and can enhance the efficiency of prediction intervals.

For multivariate scores generated by the L1-norm and Z-score functions, we adopt the Bonferroni correction for each dimension, as inspired by CF-RNN. Specifically, we use α/2 as the miscoverage rate for both x- and y-coordinates.



  §.§.§ Prediction Interval Construction

We utilize the obtained quantile Ĥ to form the prediction intervals for new examples in 𝒟_test:

    𝒞(X^𝒱_test)={Y^𝒱∈𝒴: A(X^𝒱_test, Y^𝒱) ≤Ĥ}

Specifically, for Ĥ obtained from the Z-score, we have:

    𝒞(X^𝒱_test)=[Ŷ^𝒱_test - B̂Ĥ, Ŷ^𝒱_test + B̂Ĥ]

and for Ĥ obtained from the L2-norm and L1-norm, we have 

    𝒞(X^𝒱_test)=[Ŷ^𝒱_test - Ĥ, Ŷ^𝒱_test + Ĥ]



§ EXPERIMENTS



 §.§ Experimental Setup

In this section, We introduce the specifics of the dataset, the evaluation metrics and the implementation details including hardware, hyperparameters, etc.


  §.§.§ Dataset

We evaluate the proposed model on the publicly available large-scale and real-world V2I dataset V2X-Seq <cit.>, which provides the trajectories of agents from both vehicle and infrastructure sides, along with vector map data. V2X-Seq consists of 51,146 V2I scenarios, where each trajectory is 10 seconds long with a sampling rate of 10 Hz. The task involves predicting the motion of agents in the next 5 seconds, given initial 5-second observations from both infrastructure and vehicle sides. The dataset has been split into train and validation. Our trained model is evaluated on the validation set, allowing for comparison with existing models. For post-hoc conformal prediction, we divide the validation set into calibration set and test set at a 4:1 ratio. For a discussion on the calibration data size, please refer to <cit.>.



  §.§.§ Evaluation Metrics
 
Model metrics. For model evaluation, the standard metrics in motion predictions are adopted, including minimum Average Displacement Error (minADE), minimum Final Displacement Error (minFDE), and Miss Rate (MR), where errors between the best predicted trajectory among the K=6 modes and the ground truth trajectory are calculated. The best here refers to the trajectory that has the minimum endpoint error. The ADE metric calculates the L2 distance across all future time steps and averages over all scored vehicles within a scenario, while FDE measures the L2 distance only at the final future time step and summarizes across all scored vehicles. MR refers to the ratio of actors in a scenario where FDE are above 2 meters. 

Conformal prediction metrics. We assess validity and efficiency for each method. Validity is evaluated by reporting independent and joint coverage on the test set, aiming for coverage levels close to the desired confidence level 1-α. The independent coverage for each agent is calculated as:

    ind. coverage_1-α = 1/T_f∑_X,Y ∈𝒟_test1(Y ∈𝒞(X))

We identify the maximum independent coverage among the K modes for each agent. For all metrics, the final reported value is the average among all agents in 𝒟_test.

For efficiency, we calculate the average size of the predicted 2D area across all time steps for the mode k̃ with the maximum independent coverage:

    size = 1/T_f||𝒞(X)_k̃||

The 2D area is defined as ellipsis for Z-score and L1-norm, and circle for L2-norm.

The joint coverage is defined as if there exists a mode k such that the truth values across all time steps fall into the confidence region:

    joint coverage_1-α = 1(∃k∈ K: ∀t∈ T_f: Y_k,t∈𝒞(X))




  §.§.§ Implementation Details

Model. The model was trained for 64 epochs on an Nvidia V100S GPU with 32GB memory using AdamW optimizer <cit.>. Hyperparameters including batch size, initial learning rate, weight decay and dropout rate are 32, 1e-3, 1e-4 and 0.1, respectively. The learning rate is decayed using the cosine annealing scheduler <cit.>. Our model employs the original setting of HiVT with 64 hidden units for the single view encoder and 1 layer of cross-graph attention module with 8 heads. The radius of all local regions is 50 meters. The number of prediction modes is set to 6.

Conformal prediction. With the pretrained model and datasets 𝒟_cal and 𝒟_test, we execute all conformal prediction methods on the CPU. We evaluate the methods on the three defined score functions at three different α levels: 0.2, 0.1, 0.05. The optimization step in CopulaCPTS remains consistent with the original work.



 §.§ Results Analysis

In this section, we first examine the model's prediction performance, assessing it from both quantitative and qualitative perspectives. Next, we showcase the effectiveness of the post-hoc uncertainty quantification method.


  §.§.§ Model Performance

Comparison with state-of-the-art. We benchmark our proposed model against state-of-the-art models using the V2X-Seq dataset, as detailed in <cit.>. Results are summarized in Table <ref>. Both TNT <cit.> and HiVT <cit.> are single-view models. <cit.> evaluated them under two settings: Ego, where only the vehicle-view data is utilized, and PP-VIC, which employs a two-stage method with both vehicle view and infrastructure data. Specifically, in PP-VIC, data from both views are fused offline with some tracking method and then the stitched trajectories were fed into a single-view model. We retrained the HiVT model under both Ego and the PP-VIC setting with 64 epochs and present our results in Table <ref>. Comparing the Ego and PP-VIC results, it's uncovered by <cit.> that integrating information from the infrastructure side can enhance prediction accuracy. V2X-Graph represents the current state-of-the-art model, employing a graph link prediction module to associate two-view data, followed by fusing the embeddings of the associated nodes using attention mechanism. However, this node association module requires pre-training, leading to a two-stage training process.  Our method demonstrates the best performance in terms of minFDE and MR. Without the explicit node association from both views, our attention based fusion module can attend to the most relevant nodes from both views through learning, which significantly simplifies the modeling framework and facilitates ease of training, all while achieving better results.  


Qualitative results. We visualize our prediction results and six representative scenarios are shown in Figure <ref>. To maintain clarity and simplicity in the visualization, each scenario displays the ground truth and multimodal prediction results only for the target agent, although predictions for all agents are accessible. In the first column, scenarios S1 and S4 depict the target agent turning right. In the second column, the target agent is shown going straight, leaving (S2) and approaching (S5) an intersection. The last column presents scenarios where the target agent is merging (S3) and turning left (S6). It's shown that our model predictions capture the multimodal behavior. In scenarios S2 and S5, it's exhibited in the form of different velocity profiles. Since lane information is integrated via an attention mechanism rather than as rigid constraints, off-road and road rule-violating predictions and may occur, as shown in S4 and S6. Notably, there are more uncertainties when agent making turns, as evidenced by the dispersion among the predictions. Furthermore, it's important to highlight that the prediction with the highest probability does not always align with the ground truth, as demonstrated in S1 and S3, resulting from the inherent model and data uncertainties. 



We further compare the prediction results from model HiVT Ego, HiVT PP-VIC and ours on selected scenarios, as shown in Figure <ref>. It's observed that HiVT Ego, relying solely on vehicle-view data, consistently predicted incorrect modes. In contrast, both HiVT PP-VIC and our model, incorporating additional data from the infrastructure view, demonstrate improved accuracy in capturing the ground truth.





  §.§.§ Post-hoc Uncertainty Quantification Performance

Quantitative comparison. We show in this section that conformal prediction method produces more valid and efficient confidence regions than the model predictions. The evaluation results are shown in Table <ref>.


The Mixture column presents UQ results derived directly from the multimodal predictions. We first examine the metrics of independent and joint coverage. It's evident that the model predicted intervals fail to meet the validity requirements across all specified miscoverage rates. Interestingly, CF-RNN yields overly conservative results, maintaining coverage levels around 0.99 regardless of the miscoverage rate. This could indicate that Bonferroni correction, often employed with shorter horizons, might not be suitable for our predicted horizon of 50. CopulaCPTS achieves coverage levels very close to the desired values across all specified miscoverage rates α.

Next, we inspect the metric "size", which serves as an indicator of UQ efficiency. Benchmark the size from model predictions in the Mixture column, CF-RNN produces extremely large confidence region which is expected considering the overly conservative coverage observed earlier. CopulaCPTS yields comparative and even smaller size compared to the benchmark. 

Finally, we examine the impact of different score functions. In both CF-RNN and CopulaCPTS, Z-score appears as the least efficient among all three functions. Since Z-score is computed using the model-predicted uncertainties B̂, it will always be more inefficient than the benchmark size in the "Mixture" column. One plausible explanation provided in <cit.>, is that although widely used, B̂ in the Z-score function is not directly related to the quantiles of the label distribution. Therefore, it may not be the most suitable score function option in our case. On the other hand, both L2-norm and L1-norm demonstrate significantly higher efficiency, particularly within the CopulaCPTS method. L1-norm, in particular, considers variations among different dimensions, leading to enhanced efficiency. Taking into account both validity and efficiency requirements, we have highlighted the best results among all methods in red.

Qualitative analysis. We visualize the uncertainty prediction from the mixture model and the CopulaCPTS, as shown in Figure <ref>. Three scenarios are selected in each column. We can see that the mixture model predictions fail to cover the ground truth while CopulaCPTS achieves the coverage at the specified α. 
 





§ CONCLUSION

We have presented a novel model framework with multi-view data integration in the cooperative driving setting. Our proposed model is straight forward and can be built upon any existing graph-based single-view models. It has demonstrated its effectiveness and advantages over existing benchmarks. Moreover, we have incorporated a post-hoc uncertainty quantification module, providing valid and efficient confidence regions, which is crucial in safety-critical tasks such as trajectory prediction. 

The proposed framework has certain limitations. From the model's perspective, we currently treat all road agents as the same type. However, in the public V2X-Seq dataset, there are different vehicle types, such as trucks, vans, buses, motorcycles, etc. Future work should address the different characteristics of these vehicle types to enhance model performance. Moreover, better methods for encoding the lane information to eliminate off-road and road-rule-violating predictions should be investigated. From the uncertainty quantification perspective, we simplify the quantile computation on the multimodal prediction into computation on the single best mode. This approach loses valuable distributional information. Therefore, exploring score functions that consider the distribution could lead to more accurate uncertainty quantification for multimodal results. Furthermore, we evaluate the uncertainty assuming independence among the agents. Future work can incorporate agent correlations based on the graph structure to better reflect the underlying uncertainty relations and provide more insightful results.

unsrt  
  




