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





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



§ ABSTRACT  




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


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

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


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


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


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




§ INTRODUCTION




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




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

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

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

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





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




More specifically, is structured around two stages:

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

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


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




















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

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

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

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

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

Our contributions can be summarized as follows:

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







§ PRELIMINARY























 §.§ Planning Problems


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





 §.§ Planning Problem Implementation





















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

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







 §.§ Large Language Model based Agent

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

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



















































§ METHODOLOGY





  
Problem Definition


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

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



 §.§ Environment Generation







  
Overview


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



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



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



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





 §.§ Task Generation








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

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




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












§ EXPERIMENTS





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



 §.§ Experimental Setup




  
Evaluation Tasks

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

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

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




  
Evaluation Metrics

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



  
Baselines


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









  
Implementation Details

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

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

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

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



 §.§ Evaluation on In-Domain Tasks
























































































































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


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






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

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




 §.§ Robustness





















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



 §.§ Evaluation on Out-of-Domain Tasks










§ RELATED WORK




  
Large Language Model based Agent.

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



  
Planning with Large Language Models.

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




  
Agent Training.

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




  
Environment and Task Generation with Large Language Models.


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



§ CONCLUSION




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

plainnat








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



§ MORE IMPLEMENTATION DETAILS




 §.§ Models



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





 §.§ Natural Language Mapping



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

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

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

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

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

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

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



Your output must strictly follow the provided example.



Example:

PDDL Domain:

“`pddl

(define (domain hanoi)

  (:requirements :strips)

  (:predicates

  (clear ?x)

  (on ?x ?y)

  (smaller ?x ?y)

  )



  (:action move

    :parameters (?disc ?from ?to)

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

               (clear ?disc) (clear ?to))

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

          (not (clear ?to))))

  )

“`

Specification:

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



The actions defined in this domain include:

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



You have the following restrictions on your actions:

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

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

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

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

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

 Natural Language Mapping:

“`python

{

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

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

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

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

}

“`



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

    

PDDL Domain:

{PDDL_Domain}

Specification:

{PDDL_Description}

Natural Language Mapping:






§ MORE STATISTICS ON ENVIRONMENT





 §.§ Environment Specification




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



 §.§ Environment Implementation




























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






 §.§ Diversity Analysis



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



§ EXAMPLES



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



 §.§ Environment Specification


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



The actions defined in this domain include:

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

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

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

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



You have the following restrictions on your actions:

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

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

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

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





 §.§ Environment Implementation


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

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

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

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

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






 §.§ Trajectory Data


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



The actions defined in this domain include:

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

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

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

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



You have the following restrictions on your actions:

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

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

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

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

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

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

(Assistant) Action: jordan tests the recipe almond_butter_bars.

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

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










 Paper Title: Gradient Harmonization in Unsupervised Domain Adaptation.


























































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





































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





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




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

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

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








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

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






























§ INTRODUCTION




























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



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

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

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

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

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


Detailed theoretical support is also provided for the proposed approaches.

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

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


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

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

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





§ RELATED WORK




 §.§ Unsupervised Domain Adaptation (UDA)




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



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

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

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

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

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






 §.§ Multi Task Learning

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



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



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







§ PROPOSED APPROACH





 §.§ Problem Definition

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




 §.§ A General Framework of UDA


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

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


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

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

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

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

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

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






 §.§ Gradient Harmonization (GH)


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

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


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

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

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

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







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

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

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

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

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










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

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

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

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

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

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


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

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

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

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











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

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

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

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


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

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

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

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

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



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

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

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




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

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

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


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

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

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

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






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

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

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


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

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

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



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

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







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

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

The above model can be transformed into

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

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

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

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





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

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












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



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


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

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


















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

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



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



















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

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











 §.§ Essence and Insights of GH


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

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

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

where ⟨·⟩ is the inner product operator.



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

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




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


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







 §.§ Improved Version: GH++

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


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





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



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

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

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

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

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


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

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

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

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

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

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

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

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


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





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

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




It is worth noting that GH++ presents a more favorable optimization scenario for the original task compared to GH. This advantage arises from the fact that the sum of the gradient deviation from the original direction of GH++ is only half of GH.



































 §.§ Equivalent Model of UDA with GH/GH++



By combining the general UDA model (Eq. (<ref>), described in Section <ref>) with gradient harmonization strategies (Eq. (<ref>)/Eq. (<ref>), a well-balanced UDA model can be trained and implemented. However, the gradient aggregation operator in Eq. (<ref>)/Eq. (<ref>) is intricate and has an impact on optimization efficiency. Therefore, we introduce a computation-efficient alternative model that is functionally equivalent to UDA with GH/GH++. For convenience, we can express the proposed gradient harmonization approaches, i.e., Eq. (<ref>)/Eq. (<ref>), as follows:

    g̃=τ_1g_1+τ_2g_2,

where τ_1 and τ_2 are constants that can be calculated by using the original gradients g_1 and g_2. Notably, the gradient g_1 of the original loss ℒ_1(Θ) and the gradient g_2 of the original loss ℒ_2(Θ) can be easily computed, only if the UDA model is fixed.  Then, if GH is chosen to resolve the conflict, τ_1 and τ_2 can be calculated as follows.

    τ_1=1-δ(g_1^Tg_2<0)g_2^Tg_1/‖ g_1 ‖^2,  
       τ_2=1-δ(g_1^Tg_2<0)g_1^Tg_2/‖ g_2 ‖^2.

If GH++ is used to eliminate the conflict, τ_1 and τ_2 can be calculated as follows.

    τ_1=(1+2 δ(g_1^Tg_2<0) sinλ(arccosg_1^Tg_2/|g_1|·|g_2|-π/2)/2),
       τ_1=(1+2 δ(g_1^Tg_2<0) sin(λ-1)(arccosg_1^Tg_2/|g_1|·|g_2|-π/2)/2).

Ultimately, we can derive the equivalent UDA model embedded GH/GH++ by conducting the integral operation on the gradient in Eq. (<ref>). The overall loss of UDA with GH/GH++ model can be represented as



    ℒ̃=∫(τ_1g_1+τ_2g_2)d Θ=τ_1ℒ_1(Θ)+τ_2ℒ_2(Θ),

where ℒ_1(Θ) and ℒ_2(Θ) can be the domain alignment loss and classification loss described in Eq. (<ref>) and Eq. (<ref>), respectively, for unsupervised domain adaptation.
Finally, the objective function of the UDA with GH/GH++ is:

    θ_g,θ_cmin  θ_dmax  ℒ̃=τ_1 ℒ_dom(θ_g,θ_d)+τ_2 ℒ_cls(θ_g,θ_c).



It can be seen that the proposed approaches essentially reweight the original loss function ℒ_1(Θ) and ℒ_2(Θ). The computation of the weights τ_1 and τ_2 needs to use the gradient of the original loss function, as presented in  Eq. (<ref>)/Eq. (<ref>). Finally, a balanced and efficient UDA model is formulated. Without loss of generality, when g_1^Tg_2>0, we can observe τ_1 =1 and τ_2 =1, and the overall loss function is degenerated into ℒ̃=ℒ_1(Θ)+ℒ_2(Θ), i.e., the general UDA model. We describe the training procedure in Algorithm <ref>.

Notably, the proposed approaches are  plug-and-play modules and can be deployed in almost all models with different losses ℒ_1(Θ) and ℒ_2(Θ). In order to clearly observe the computational cost of the equivalent model, we discuss the training speed and recognition accuracy before and after applying GH in Table <ref>.





§ EXPERIMENTS




 §.§ Datasets






Office-31 <cit.> is a mainstream benchmark dataset for visual domain adaptation, which consists of three distinct domains: Amazon (A), DSLR (D), Webcam (W). It totally contains 4,652 images from 31 categories. We evaluate our method in all 6 different transfer tasks across domains.


Office-Home <cit.> is a more challenging and harder benchmark than Office-31. It contains 15.5K images across 65 object categories from 4 different domains: Artistic images (Ar), Clip Art (Cl), Product images (Pr), and Real-World images (Rw). We evaluate our method in all 12 different transfer tasks across domains.

Digits Datasets. We mainly study three datasets: MNIST (M) <cit.>, USPS (U) <cit.> and SVHN (S) <cit.>. MNIST and USPS are two general handwriting recognition datasets involving 10 categories. SVHN is obtained from house numbers in Google Street View images. We conduct experiments in 3 universal tasks, including M→U, U→M and S→M.

VisDA-2017 <cit.> is a simulation-to-real dataset for domain adaptation with two distinct domains: synthetic object images rendered from 3D models and real object images. It contains over 280K images across 12 categories.



DomainNet <cit.> is the largest domain adaptation dataset consisting of about 600K images from six distinct domains, including Clipart (clp), Infograph (inf), Painting (pnt), Quickdraw (qdr), Real (rel), Sketch (skt). There are 48K-172K images categorized into 345 classes per domain. We evaluate our method in all 30 cross-domain tasks.








 §.§ Implementation Details

We compare our method with the following state-of-the-art unsupervised domain adaptation methods: ResNet <cit.>, DAN (Deep Adaptation Networks) <cit.>, DANN (Domain-adversarial Neural Networks) <cit.>, JAN (Joint Adaptation Networks) <cit.>, DRCN (Deep Reconstruction-Classification Networks) <cit.>, CoGAN (Coupled Generative Adversarial Networks) <cit.>, ADDA (Adversarial Discriminative Domain Adaptation) <cit.>, CyCADA ( Cycle-consistent Adversarial Domain Adaptation) <cit.>, CAT (Cluster Alignment with a Teacher) <cit.>, TPN (Transferrable prototypical networks) <cit.>, LWC (Light-weight Calibrator) <cit.>, ETD (Enhanced Transport Distance) <cit.>, TAT (Transferable Adversarial Training) <cit.>, TADA (Transferable Attention for Domain Adaptation) <cit.>, SymNets (Symmetric Networks) <cit.>, BNM (Batch Nuclear-norm Maximization) <cit.>, ALDA (Adversarial-learned Loss for Domain Adaptation) <cit.>, MIMTFL (Mutual Information Maximisation
and Transferable Feature Learning) <cit.>, Meta (Metaalign) <cit.>, MDD (Margin Disparity Discrepancy) <cit.>, CDAN <cit.>, MCD (Maximum Classification Discrepancy) <cit.>, GVB (Gradually Vanishing Bridge) <cit.> and DWL (Dynamic Weighted Learning) <cit.>, SAFN (Stepwise Adaptive Feature Norm) <cit.>, SWD (Sliced Wasserstein Discrepancy) <cit.>, DMRL (Dual Mixup Regularized Learning) <cit.>, CAN (Contrastive Adaptation Network) <cit.>, FixBi <cit.>, FGDA (Feature Gradient Distribution Alignment) <cit.> and ParetoDA (Pareto Domain Adaptation) <cit.>, CGDM (Cross-domain Gradient Discrepancy Minimization) <cit.>, SSRT (Safe Self-Refinement for Transformer-based domain adaptation) <cit.> and TVT (Transferable Vision Transformer) <cit.>.  SSRT and TVT are transformer-based approaches, while other approaches are CNN-based. We use the results in their original papers for fair comparison. When there is no relevant experimental result in the original papers, we implement them according to their released source codes.




 
Furthermore, we investigate the effect of GH/GH++ on five popular UDA approaches, including four CNN-based approaches (i.e., CDAN <cit.>, MCD <cit.>, GVB <cit.> and DWL <cit.>) and one transformer-based approach (i.e., SSRT <cit.>). For CNN-based approaches, we use pre-trained ResNet-50 <cit.> on ImageNet-1K <cit.> as backbone network for Office-31 and Office-Home, and adopt the pre-trained ResNet-101 <cit.> as backbone network for VisDA-2017 and DomainNet. We employ LeNet <cit.> as backbone for Digits datasets. For transformer-based approach, SSRT uses the ViT-base with 16 × 16 patch size pre-trained on ImageNet-21K and ImageNet-1K (denoted by SSRT^∗) as the transformer backbones, respectively. In terms of optimization, we follow the original protocol of the baselines. Note that, in this paper, CDAN means CDAN+E in <cit.>) and GVB refers to GVB-GD in <cit.>. For GH++, we empirically set λ to 0.5.
MCD achieves domain alignment by the game between two classifiers and the generator. 
Therefore, the gradient of the discrepancy loss between the two classifiers is just the gradient of domain alignment loss. We implement our method in PyTorch by following UDA protocol <cit.> that source domain is labeled and target domain is unlabeled. 

For fair comparisons, we independently run all experiments five times and report the average target accuracy. 



 §.§ Results on UDA


























































Tables <ref>, <ref>, <ref>, <ref> and <ref> present evaluation results on Office-31, Office-Home, VisDA-2017, Digits and DomainNet, respectively. Generally, the transformer-based results, such as TVT and SSRT, are much better than CNN-based results, which has been validated in previous work. For TVT and SSRT, the ViT backbone pre-trained on ImageNet-1K is slightly inferior than ImageNet-21K (i.e., TVT^∗&TVT, SSRT^∗&SSRT). Considering that the CNNs are pre-trained on ImageNet-1K, we take SSRT^∗ for fair analysis as default in the following.  

In order to demonstrate the effectiveness and universality of the proposed approaches in UDA, we investigate the effect of GH/GH++ on five popular UDA approaches, including CNN-based approaches (i.e., CDAN, MCD, GVB and DWL) and transformer-based approach (i.e., SSRT). We can observe that the proposed approaches can further improve the baseline, whether based on transformer or CNN. This is attributed to the proposed approaches in alleviating gradient conflict between domain alignment task and classification task in optimization. Meanwhile, it verifies that the proposed approaches can be easily plugged and played in the existing UDA methods. In addition, GH++ usually performs better than GH. The reason is that during removing the gradient conflict between the two tasks, GH++ minimizes the gradient deflection, such that the two tasks can evolve harmoniously during joint training while preserving their individual task-specific optimality.

Specifically, as shown in Table <ref>, when GH/GH++ is applied, CDAN, MCD, GVB, DWL and SSRT outperform the original models by 1.5%/2.2%, 1.0%/1.4%,  0.2%/0.7%, 0.7%/0.8% and 0.7%/0.9% on Avg^ of Office-31, respectively.
As shown in Table <ref>, CDAN, MCD, GVB, DWL and SSRT after applying GH/GH++ outperform the original models by 2.6%/3.2%, 4.5%/5.1%, 0.8%/1.2%, 0.4%/0.6% and 0.6%/0.8% on Avg^ of Office-Home, respectively. As shown in Table <ref>, we can observe that CDAN, MCD, GVB, DWL and SSRT after applying GH/GH++ outperform the original models by 0.9%/0.9%, 2.2%/2.4%,  1.1%/1.3%, 0.6%/0.9% and 0.3%/0.3% on Avg of Digits, respectively.
As shown in Table <ref>, CDAN, MCD, GVB, DWL and SSRT after applying GH/GH++ outperform the original models by 1.4%/3.9%, 4.1%/4.5%, 0.5%/0.8%, 0.7%/0.3% and 0.7%/0.6% on VisDA-2017, respectively.
From the results in Table <ref>, we can observe that CDAN, MCD, GVB, DWL and SSRT after applying GH/GH++ outperform the original models by 3.2%/3.7%, 7.1%/7.3%, 0.7%/1.0%, 1.0%/1.3% and 1.1%/1.4% on Avg of DomainNet, respectively. The proposed approaches generally improve CDAN and MCD more than other baselines because the original models of CDAN and MCD have more obvious optimization conflicts than other baselines. Fig. <ref> also confirms this perspective. Therefore, the proposed approaches have better performance when the inherent gradient conflict is serious. Besides, it is worth noting that the proposed approaches usually achieve greater performance gains on large-scale datasets, i.e., VisDA-2017 and DomainNet, which further indicate the effectiveness and versatility of our method.

























































§ MODEL ANALYSIS AND DISCUSSION



 §.§ Feature Visualization
  Fig. <ref> describes the t-SNE <cit.> visualizations of features learned by MCD (baseline) and MCD+GH on the tasks of U → M and M → U. Fig. <ref> (a) and (c) are visualization features generated by MCD. Fig. <ref> (b) and (d) are visualization features generated by MCD+GH. It can be observed that both features learned by MCD and MCD+GH achieve well-performed global alignment effect with 10 clusters under two tasks. Further, the visualization feature distributions with GH deployed have better clustering effect and have fewer samples distributed across class boundaries, which intuitively boosts the feature discriminability. In addition, visualization results further validate the balance learning ability of our GH approach.



 §.§ Convergence Analysis
  We present the convergence curves of test error with respect to the number of iterations on tasks of U → M and Synthetic → Real as shown in Fig. <ref>. For each subfigure, the blue line represents the test error of different baselines, and the red line represents the test error for baseline+GH (e.g., CDAN+GH). Obviously, compared with baselines, the introduction of GH can further improve the test performance and convergence. This fully indicates that GH plays an active coordination role in practical optimization process, which promotes the cooperation of domain alignment task and classification task towards a benign direction as expected. 

















 §.§ Confusion Matrix Visualization


Fig. <ref> displays the visualizations of confusion matrix for the classifier trained by DWL and DWL+GH. DWL obtains several uncertain predictions with small values while DWL+GH obtains more confident predictions. Comparing with Fig. <ref> (a) and (b), the confusing “Class 1, 2, 7, 8, and 9” are correctly recognized in DWL+GH. From Fig. <ref> (c) and (d), the confusing “Class 8” is corrected in DWL+GH. It implies that optimization conflict during training deteriorates the discriminability in the target domain. Our method preserves the feature discriminability of target samples in the course of harmonic training. DWL+GH improves the between-task balance by coordinating the explicit external task conflict and the implicit internal optimization conflict.





 §.§ Balance Analysis of GH-based UDA
  Fig. <ref> shows MMD distance <cit.> and the max J(W) <cit.> values based on the feature representation learned by MCD+GH. The left vertical axis corresponding to the red curve represents the MMD distance used to measure the alignment degree across domains. The smaller value means the better domain alignment. The right vertical axis corresponding to the blue curve represents the degree of discriminability based on Linear Discriminant Analysis (LDA) <cit.>. The larger max J(W) implies better discriminability, which directly affects the classification result of samples. Fig. <ref>(a) and (b) reflect the degree of domain alignment and discriminability on task of M→U and U→M, respectively. From Fig. <ref>, MCD+GH has a smaller MMD distance and a larger max J(W) value during training. This shows that GH facilitates the coordinated optimality of alignment and classification tasks while maintaining their task-specific optimality. Enhanced and balanced domain-invariant and class-discriminative feature representations can be obtained.



 §.§ Training Speed and Accuracy
 In order to observe the efficiency of the proposed equivalent model more clearly, we present the training speed and classification accuracy before and after applying GH. As shown in Table <ref>, for tasks M→U and U→M on digits datasets, the training speed of MCD+GH is 0.14254 s/epoch and 0.09809 s/epoch longer than MCD, respectively. In other words, MCD+GH takes less training time than MCD in training process, but can get 2.5%/2.7% classification accuracy gain. The computational cost of employing GH is quite low, and thus GH is a powerful and efficient auxiliary tool to facilitate those popular domain adaptation baselines towards more outstanding classification performance.







 §.§ Gradient Inner Product Visualization
  Fig. <ref> presents inner product distributions of two gradients between domain alignment and classification tasks before and after applying Gradient Harmonization for MCD. From Fig. <ref> (a), we observe the acute and obtuse angles between gradients of the two tasks before coordination. The obtuse angles account for about 40% of the total number, which exhibits the identical property as Fig. <ref>. In other words, there exists between-task conflict in the model optimization process but paid less attention. After Gradient Harmonization, as shown in Fig. <ref> (b), the inner products of the two gradients are all positive. That is, the gradient angles between the two tasks are coordinated into acute angles. The proposed GH avoids the optimization conflict by separately adjusting the gradients of the two tasks to achieve the purpose of optimal coordination. Experimental results fully illustrate the effectiveness of the proposed GH.







 §.§ Rationality and Comparison to Other Alternatives



The proposed GH/GH++ aims at altering the gradient angle between two different tasks from an obtuse angle to an acute/vertical angle. 
Whether the same effect can be achieved with other intuitive gradient correction alternatives or not? For instance, one might consider convert an obtuse angle into an acute angle by simply applying a negative sign to one of the two gradients, i.e., transforming either g_1 into -g_1 or g_2 into -g_2. As depicted in Fig. <ref>, it is apparent that the results reveal a significant decline in classification accuracy when using the -g_1&g_2 and g_1&-g_2 methods in contrast to the proposed approaches. While applying a negative sign can indeed transit from an obtuse angle to an acute angle and force the two gradients to be positively correlated, it cannot replicate the excellent classification performance achieved by the proposed GH/GH++.

The rationale behind this observation is that directly applying a negative sign to one of the gradients alters the gradient towards the opposite direction of the original optimal gradient descent direction. This significantly undermines the primary objective of the original task. In other words, the performance deteriorates sharply when the gradient harmonization direction significantly deviates from the original direction. GH/GH++ reasonably adjusts the two gradients to achieve the purpose of optimization coordination on the premise of keeping the harmonic gradient as close as possible to the original gradient descent direction. Therefore, the proposed approaches do not break the task-specific optimality, but pursues a better cooperation, which verifies the rationality of the proposed GH/GH++.








 §.§ Parameter Sensitivity Analysis

To investigate the effect of the parameter λ in GH++, we conduct experiments on three tasks (i.e., M→U, U→M and S→M) based on MCD and GVB by varying λ∈{0, 0.1, 0.3, 0.5,0.7, 0.9, 1}. The results are presented in Fig. <ref>. We can observe GH++ is little sensitive to the scale variety of λ, which indicates that GH++ is robust across different baselines and tasks. Besides, we observe that when λ=0.5, models generally achieve the best performance. In other words, when the gradient deviation of the two tasks is relieved, the performance can be largely improved.
























 §.§ Scalability to Other Multi-task Problems

To further demonstrate the universality and scalability of the proposed approaches, we evaluate GH/GH++ in the object detection and multi-modal interactive retrieval fields, which also involves optimization of multiple objectives.



Dataset. We select widely used benchmarks, i.e., PASCAL VOC 2007 <cit.> and CSS <cit.> for object detection and multi-modal interactive retrieval, respectively.
PASCAL VOC 2007 consists of about 5K trainval images and 5K test images over 20 object categories.
CSS consists of 38K synthesized images with different colors, shapes and sizes. It contains about 19K training image-text pairs and 19K testing image-text pairs, respectively. 



Implementation Details. For object detection, we use the DSSD <cit.> model with 321× 321 inputs as the baseline. We follow the same experimental settings and protocol as <cit.> and adopt Average Precision (averaged AP at IoUs from 0.5 to 0.9) to measure performance. For multi-modal interactive retrieval, we use TIRG <cit.> and the adversarial training version(AT-TIRG) follows the paper <cit.> as the baseline. For evaluation, we follow the same protocols as <cit.> and use retrieval accuracy R@N as our evaluation metric, which computes the percentage of test queries where at least one target or correctly matched image is within the top N retrieved images. Note that we exactly follow the original experimental setups of the baselines. In other words, DSSD+GH/GH++ and AT-TIRG+GH/GH++ indicate that GH/GH++ are inserted directly into the DSSD and AT-TIRG, respectively, without any change in either the backbone or hyper-parameters.





Results on Object Detection. Object detection models <cit.> usually consider two tasks: classification and regression. In this section, we use the proposed GH/GH++ to mitigate the conflict between these two tasks. Table <ref> presents the Average Precision (AP) on the PASCAL VOC 2007 test set. We can observe that GH/GH++ can improve the AP with different bounding boxes, which verifies the effectiveness of the proposed approaches in object detection.








Results on Multi-modal Interactive Retrieval. Recently, <cit.> introduces adversarial training <cit.> into multi-modal interactive retrieval model (abbreviated as AT-TIRG) to improve model robustness, which aims to train a generalized and robust model suitable for both clean samples and adversarial attack samples. During adversarial training, there is conflict in training the clean samples and adversarial attack samples, which results in the limitation of the final performance. In this section, we try to adopt the proposed approaches to solve this problem. The experimental results are provided in Table <ref>. Comparing AT-TIRG+GH with AT-TIRG, we can see the proposed GH can enhance the AT. For the adversarial attack samples, AT-TIRG+GH outperforms the AT-TIRG by 3.4%, 2.9%, and 2.2% on R@1, R@5, and R@10, respectively. AT-TIRG+GH++ outperforms the AT-TIRG by 5.2%, 4.9%, and 3.9% on R@1, R@5, and R@10, respectively. The performance of the clean samples is also improved in most cases. Note that the performance of clean samples is usually close to 100%, so there is little space for improvement.





































§ CONCLUSION

In this paper, we pay attention to the optimization conflict (i.e., imbalance or incoordination) problem between different tasks (i.e., the alignment task and the classification task) in alignment-based unsupervised domain adaptation models. To mitigate this problem, we propose two simple yet efficient Gradient Harmonization approaches, including GH and GH++, which take measures to de-conflict between the gradients of both tasks in optimization. Besides, to facilitate the harmonization during adaptation, we derive the equivalent but more efficient model of UDA with GH/GH++, which becomes a dynamically reweighted loss function of most existing unsupervised domain adaptation models. Further, the essence and insights of the proposed approaches are provided to indicate its rationality. Exhaustive experiments and model analyses demonstrate that the proposed approaches significantly improve the existing UDA models and contribute to achieving state-of-the-art results.
In addition, we have verified that the proposed approaches can be adapted to other problems and areas, such as object detection and multi-modal retrieval, to de-conflict between the gradients of any two tasks in optimization and improve model performance.




























 

  


































IEEEtrans






























































 Paper Title: Universal Approximation Theory: Foundations for Parallelism in Neural Networks.



Turbulence, Waves, and Taylor's Hypothesis for Heliosheath Observations
    M. Nakanotani
    
=======================================================================



§ ABSTRACT 

Neural networks are increasingly evolving towards training large models with big data, a method that has demonstrated superior performance across many tasks. However, this approach introduces an urgent problem: current deep learning models are predominantly serial, meaning that as the number of network layers increases, so do the training and inference times. This is unacceptable if deep learning is to continue advancing. Therefore, this paper proposes a deep learning parallelization strategy based on the Universal Approximation Theorem (UAT). From this foundation, we designed a parallel network called Para-Former to test our theory. Unlike traditional serial models, the inference time of Para-Former does not increase with the number of layers, significantly accelerating the inference speed of multi-layer networks. Experimental results validate the effectiveness of this network.






§ INTRODUCTION

In today's era, deep learning models have undoubtedly taken the lead in the fields of computer vision (CV) and natural language processing (NLP), demonstrating exceptional performance across numerous tasks. In CV, models such as ResNet<cit.> and ViT<cit.> for image recognition, U-Net <cit.> for image segmentation, GAN <cit.> for image generation, and diffusion models like SORA <cit.> for video generation highlight the extensive application of deep learning in CV. These models are primarily based on convolutional neural networks (CNNs), Transformers, or a combination of both. Simultaneously, large language models (LLMs) based on Transformers have become dominant in natural language processing. Overall, deep learning models have shown vast application potential and powerful capabilities so far.

However, this progress has brought forth some challenges: the development trend in deep learning is moving towards larger models and larger datasets. This trend is fundamentally driven by the UAT<cit.>. From the papers UAT2LLM<cit.> and UAT2CV<cit.>, we know that residual-based CNNs and Transformer-based models are specific implementations of UAT, both exhibiting the characteristic of dynamically approximating functions based on inputs. Larger models and datasets are indeed requirements of the UAT theory: deeper network layers can more precisely approximate target functions dynamically based on inputs, while larger and more diverse datasets help ensure the training data closely represents real-world conditions, making the model's approximated functions more reflective of real-world complexities.

However, this trend inevitably leads to increased demands for computational resources and extended training times. As more powerful models are needed, the UAT theory mathematically necessitates larger networks, which in turn require more computational resources. While technological advancements, especially in hardware performance, provide potential ways to alleviate this issue, the increase in time costs remains a fundamental obstacle to scaling serial networks. The improvement in computational resources such as GPUs and TPUs is limited by the speed of electrical signal transmission, whereas the theoretical scale of networks has no such upper limit. When networks reach tens or even hundreds of thousands of layers, inference times can extend to several minutes or even hours for each batch. Current solutions include model optimization techniques like quantization and pruning <cit.>, hardware acceleration with specialized accelerators and edge computing <cit.>, and parallel and distributed computing <cit.> with model parallelization, data parallelization, and distributed computing frameworks. However, these do not fundamentally address the inference delay caused by the increasing number of serial network layers.

We believe the root cause lies in the serial computation mechanism of existing deep learning models, where the output of one layer serves as the input for the next. This design stems from the early development of deep learning networks in the field of CV, where image data exhibit local correlations and spatiotemporal invariance. To accommodate these features, early CV networks primarily employed convolutional layers with small receptive fields. To achieve a larger receptive field, deeper network layers were necessary, leading to early deep-learning architectures like AlexNet and VGG. The introduction of ResNet brought a paradigm shift, as residual networks have the ability to dynamically fit functions based on input, influencing almost all subsequent network designs and solidifying the tradition of multi-layer stacking.



While this design was manageable with fewer layers, the increasing depth and number of parameters in modern networks have created a critical issue: slow computation efficiency due to the need for sequential layer-by-layer calculations during inference. This computation method inevitably leads to delays, which become more severe as the network depth increases, significantly constraining the progress of CV and NLP technologies.

Therefore, it is imperative to explore parallelization techniques for deep learning networks, enabling parallel computation across multiple layers. This approach is key to accelerating the inference process, breaking the time bottleneck, and advancing CV and NLP technologies towards greater efficiency and speed. To develop such parallel computation, we must start from the fundamental theory of deep learning: the UAT. The network must be designed in accordance with UAT principles and must satisfy the requirements outlined in the papers UAT2LLM and UAT2CV, which emphasize the ability to dynamically fit functions based on data. With this in mind and leveraging existing network structures, we designed a parallel network, Para-Former, to validate the feasibility of our theory. Our contributions include:






  * Proposing the theory of parallel deep learning networks and experimentally validating its feasibility.

  * Analyzing the characteristics, advantages, and disadvantages of both serial and parallel networks.

  * Investigating the issues present in deep learning from a data perspective.

 


The structure of this paper is as follows: In Section <ref>, we first present the general form of the Universal Approximation Theorem (UAT). Based on the mathematical form of UAT, we define the parallel network Para-Former in Section <ref>. In Section <ref>, we validate the effectiveness of Para-Former in various ways: we assess its feasibility in Section <ref>, explore the impact of network depth in Section <ref>, and analyze the influence of different datasets in Section <ref>.





§ UAT



Deep learning networks have demonstrated their powerful capabilities in numerous CV and NLP tasks. In the papers UAT2LLMs and UAT2CVs, it has been proven that both Transformer-based networks and residual-based CNNs are specific implementations of the UAT. Therefore, we consider UAT to be the fundamental theory for deep learning. To design parallelized networks, we must start from the principles of UAT. Before formally designing parallel networks, we will first review the basic form of the UAT theory and then propose a network parallelization solution based on this mathematical theory. The UAT was initially proposed by  <cit.>, and its basic form is as follows:





    G(𝐱)=∑_j=1^N α_j σ(𝐖_j^T𝐱+θ_j)

is dense in C(𝐈_n). Here, 𝐖_j ∈ℝ^n and α_j, θ∈ℝ are parameters σ is a sigmoid function. For any f ∈ C(𝐈_n) and ε>0, there exists a function G(𝐱):


    |G(𝐱)-f(𝐱)|<ε   for all   𝐱∈𝐈_n .



This indicates that when N is sufficiently large, a neural network can approximate any continuous function on a closed interval.   <cit.> further proved that this function can approximate any Borel measurable function. Observing Equation (<ref>), where the function G(𝐱) produces scalar outputs in the real number field ℝ, the scenario naturally extends when G(𝐱) maps to ℝ^m, requiring approximation in each dimension. To accommodate such multi-dimensional outputs, we simply need to extend Equation (<ref>): the transformation matrix 𝐖_j is revised to the space ℝ^n × m, the bias term θ_j is recast as a vector in ℝ^m, and α_j is reshaped into a matrix. We call the parameters 𝐖_j and α_j in UAT as weight and θ_j as bias.








§ UAT FOR PARALLEL NETWORK




 §.§ Pre-request for Basic Block of Parallel Network


From Equation (<ref>), we can observe a significant characteristic: if we define different indices j as the indices of layers and the corresponding layer's mathematical form is α_j σ(𝐖_j^T𝐱+θ_j), it is evident that the parameters between different layers do not interact with each other. Therefore, based on (<ref>), we can design networks where parameters across multiple layers are independent, achieving parallel computation, like Figure <ref>. Figure <ref>.a represents a typical sequential network, which clearly requires computation to proceed in order from front to back. In contrast, Figure <ref>.b depicts a parallel network, where there is no parameter interaction between different blocks, allowing them to be considered independent and thus suitable for parallel computation. The blocks in this network should satisfy the formula given in Figure <ref>.b. Clearly, under this condition, the sum of the multiple blocks' outputs adheres to the UAT theory.

If we design the parallel network solely based on Eq. (<ref>), the network will face a limitation once training is complete: the function that can be approximated by the network is fixed, as the parameters in the corresponding UAT are fixed. So to ensure the network can dynamically approximate functions based on input, some or all of the parameters α_j, W_j and θ_j must be influenced by the input (referencing to UAT2LLM and UAT2CV), which can be simply written as α_j(𝐱), W_j(𝐱)  and θ_j(𝐱). 

This allows the network to dynamically adapt and fit the function based on the input while enabling parallel computation. So we could design the basic block in Figure <ref>, like Block_1⋯Block_n are the implementation of α_j(𝐱)) σ(𝐖_j^T(𝐱) 𝐱+θ_j(𝐱)), i=1⋯ n. Some or all the parameters α_j, W_j  and θ_j will change according to input. 






 §.§ Para-former



Based on the previous request on the basic module, we designed the Para-Former which is a parallel network by using the Transformer module as the basic module. We chose Transformer for this design because it is a well-established and widely-used network architecture, excelling in both CV and NLP. Another reason is that its basic mathematical form aligns with the requirements of Section <ref>.

Before formally introducing the structure of Para-Former, let's briefly review the Transformer structure. The basic structure of a Transformer is shown in Figure <ref>, and for convenience, we refer to it as a Transformer module. Typically, it is represented as follows:


    𝐇=   MultiHead(𝐐, 𝐊, 𝐕) 
    
    =   Concat(𝐇̂_1, …, 𝐇̂_h) 𝐖_O



    𝐇̂_i   =Attention(𝐱_i 𝐖_Qi=𝐐_i, 𝐱_i 𝐖_Ki=𝐊_i,
       𝐱_i𝐖_Vi=𝐕_i)
       =softmax(𝐐_i 𝐊_i^T/√(M)) 𝐕_i
       =𝐇_i𝐕_i =𝐇_i[𝐱_i𝐖_Vi]




    FFN(𝐱_i) = 𝐖_i,3σ (𝐖_i,2𝐱_i+𝐛_i,2)+𝐛_i,3


The above representation mainly describes the engineering process rather than its true mathematical form. According to UAT2LLM, the Multi-Head Attention (MHA) and Feed-Forward Network (FFN) components of Transformers can be expressed in matrix-vector form. We denote the matrix-vector form of variables with a prime symbol (e.g., 𝐱_i becomes 𝐱'_i). For more details, refer to UAT2LLM.


    MHA(𝐱'_i) = 𝐖'_i,1𝐱'_i



    FFN(𝐱'_i) = 𝐖'_i,3σ (𝐖'_i,2𝐱'_i+𝐛'_i,2)+𝐛'_i,3


The parameter 𝐖'_i,1 is influenced by the input. As shown in Figure <ref>, the output of a single Transformer layer, 𝐱'_i+1, can be described by Eq. (<ref>) (Details could be found in Section <ref>):



    𝐱'_i+1   = 𝐖'_i,1𝐱'_i + 𝐖'_i,3σ (𝐖̂'̂_̂î,̂1̂𝐱'_i + 𝐛'_i,2) + 𝐛'_i,3



The mathematical form in Eq. (<ref>) meets the requirements of a basic module in UAT. Therefore, we can use the structure in Figure <ref> as the basic module in Para-Former, corresponding to Figure <ref>.a. Besides, we add a symbol of a hat on some parameters to represent the parameters will be affected by the input, like 𝐖̂'̂_̂î,̂1̂. 




    𝐱'_i+2   = 𝐖̂'̂_̂î+̂1̂,̂2̂𝐱'_i + 𝐖̂'̂_̂î+̂1̂,̂3̂σ (𝐖̂'̂_̂î,̂1̂𝐱'_i + 𝐛'_i,2) 
       + 𝐖'_i+1,3σ{𝐖̂'̂_̂î+̂1̂,̂4̂𝐱'_i + 𝐛̂'̂_̂î+̂1̂,̂4̂} + 𝐛̂'̂_̂î+̂1̂,̂5̂


The output of a two-layer Transformer is shown in Eq. (<ref>) (Details could be found in Section <ref>). From Eq. (<ref>), it is evident that a serial configuration of multi-layer Transformers also conforms to the mathematical form of UAT. Therefore, we can extend the Para-Former structure. That is, within one block in Figure <ref>, there can be multiple serially connected Transformer modules, while multiple blocks operate in parallel, as shown in Figure <ref>.b. We refer to the number of Transformer modules within a block as the depth and the total number of Transformer modules across all blocks as the network layer count.





 §.§ Speed-up Ratio


We have proposed a network parallelization scheme, and in this section, we will compare the running speed of this parallelized network with that of traditional serial networks. It is important to note that due to differences in hardware and implementation methods, providing an objective comparison through experiments can be challenging. Therefore, we will focus on a theoretical analysis. Based on the network structure shown in Figure <ref>, the inference time of Para-Former is only affected by its network depth. 

Assume we use a Para-Former with a depth of M and a layer count of L, compared to a traditional Transformer network with a depth of N layers, where L >> N. In this case, the running speed of the parallel network would be N/M times faster than that of the serial network. This is because the speed of the parallel network is solely related to the depth of Para-Former. As the network depth increases, the acceleration effect becomes more pronounced.



§ EXPERIMENTS


Our experimental approach is to first verify the feasibility of the parallel network theory and its adherence to the UAT approximation properties using several commonly used public image classification datasets (Section <ref>). We then explore the impact of network depth on the UAT theory (Section <ref>), followed by an investigation into the influence of the dataset (Section <ref>).





 §.§ Dataset


In pursuit of a comprehensive understanding of the efficacy and practicality of parallel models, we meticulously selected several widely-acknowledged and representative image classification datasets for an integrated comparative analysis. These datasets encompass a diverse array of visual content and span various levels of difficulty, thereby furnishing us with a fertile ground to probe the performance boundaries of parallel models. Specifically, our study centers on the following pivotal datasets:

1. CIFAR-10 & CIFAR-100 <cit.>: These were constructed by Alex Krizhevsky et al. CIFAR-10 encompasses 60,000 32x32 pixel color images across 10 classes, while CIFAR-100 extends this to 100 categories.

2. STL-10 <cit.>: STL-10 comprises 5,000 training and 8,000 testing images distributed among 10 classes, each at a 96x96 pixel resolution. 

3. Flowers-102 <cit.>: Focused on floral categorization, this dataset houses 102 distinct flower species, totalling 8,189 images with a per-species count ranging from 40 to 258. 

4. Oxford-IIIT Pets <cit.>: Tailored for the recognition of pet breeds (cats and dogs), this dataset contains over 7,300 images across 37 categories, characterized by varied animal poses, backgrounds, and lighting conditions. 

In Table <ref>, we provide more details for these dataset. 



Through experimentation on these meticulously chosen datasets, we aspire to thoroughly validate the feasibility and strengths of parallel models across a spectrum of scales, complexities, and domains in image classification tasks. Simultaneously, we seek to elucidate potential limitations and chart avenues for future research endeavors.






 §.§ The Feasibility of Para-Former


To validate the feasibility and characteristics of Para-Former, we conducted training and validation on several commonly used image classification datasets, with all models trained for 500 epochs. The experimental design involved preprocessing inputs using the common ViT techniques of patch extraction and positional encoding. We then used Transformer as the basic model architecture, setting all network depths to 1 and configuring Para-Former from single-layer to multi-layer setups, as shown in Figure <ref>.a. A single-layer setup corresponds to a ViT with a depth of 1. The results of the models are presented in Table <ref>.

It is evident from Table <ref> that as the number of parallel network layers increases, the model accuracy improves progressively (Except for the result of Para-F-1-24 being worse than Para-F-1-12 on the Flower-102 dataset, this variation is within a reasonable range.). According to UAT theory, increasing the network layers is equivalent to increasing N in Equation (<ref>), thereby enhancing the approximation capability of UAT. The results in Table <ref> are in complete agreement with UAT theory. Therefore, we believe that this parallel network approach is entirely feasible.





However, the overall prediction accuracy in Table <ref> is quite low. What could be the reasons for this? We believe there are three main factors affecting the model's predictive performance:

1. Model's Fitting Capability: This is influenced by the model's size and design. Specifically, the value of N in UAT and the degrees of freedom in the parameters. The larger the N value, the stronger the UAT's fitting capability. Parameter freedom refers to the number of parameters affected by the input in UAT, such as a freedom of 1 in Eq. <ref> and a freedom of 6 in Eq. <ref>. The degrees of freedom correspond to the model's ability to dynamically fit based on input. Merely increasing N enhances the model's ability to fit specific functions, but images are often diverse. Therefore, the network must have the ability to dynamically fit the corresponding function based on the input. Higher degrees of freedom mean the model can fit a wide range of features with significant differences. Additionally, the bias term in the network can also dynamically fit based on input using UAT, which we refer to as bias freedom. Models with bias freedom have a stronger dynamic fitting capability. The mathematical model generated by the residual network is an example of UAT with bias freedom, achieving the fitting capability of a multi-layer network with just shallow layers. We will talk about this problem in Section <ref>.

2. Data Influence: Data is another critical factor in deep learning training. Since natural images are usually diverse, limited data can lead to overfitting. For example, in a cat-dog classification task, if dogs in the training set mostly appear against backgrounds of blue skies and grass, these backgrounds might be mistakenly considered part of the dog's features, a phenomenon known as feature coupling. If cats rarely appear in such backgrounds in the training set but do in the test set, they might be misclassified as dogs. Training on a large, diverse dataset means that common backgrounds frequently appear with various objects or animals, such as cats, dogs, and elephants against blue skies, white clouds, houses, and trees. Thus, even when classifying objects, the model inherently learns background information. In this context, fine-tuning on the cat-dog dataset yields better results because the background does not interfere with dog feature recognition. Without sufficient data, even a network with strong fitting capabilities cannot perform well and is prone to overfitting. This problem will be considered in Section <ref>.

Considering these factors, it is clear that simply improving the model architecture is not enough. A robust training dataset is equally important for achieving high predictive accuracy.








 §.§ Model Depth Effect


Based on the conclusions from Section <ref>, we will explore the impact of network depth on predictive performance in this section. According to Eq. <ref> and Eq. <ref>, we know that model depth primarily affects the degrees of freedom of the model parameters in UAT. Therefore, we designed Para-Former models with 24 layers and varying depths of 2, 3, and 6, and compared their predictive performance with that of a ViT with a depth of 6. The results are shown in Table <ref>.

It is evident that all results surpass those of Para-F-1-24 in Table <ref> because the networks in Table <ref> have higher degrees of freedom in their parameters. Notably, Para-F-6-6 consistently outperforms Para-F-2-24 and Para-F-3-24, indicating that the higher the degrees of freedom of the parameters, the stronger the model's performance. However, the model cannot be too deep, so we must balance network depth and layer count. Another issue is that almost all models hit a performance bottleneck. Particularly, the prediction accuracy for OxfordIIITPet and Flower-102 is very poor. This is because these datasets are relatively small and exhibit high data diversity.











 §.§ Dataset Effect







Current deep learning models generally perform well on the aforementioned datasets. However, as shown in Table <ref>, the prediction accuracy of Para-Former is not very high, especially on the OxfordIIITPet and Flower-102 datasets. What is the fundamental reason for this poor performance? In this section, we will explain the relationship between deep learning and data.

First, the results were obtained using the validated ViT model (Para-F-1-1 and Para-Former-6-6), which performed similarly to Para-Former. Given that the ViT model has been extensively validated and is a commonly used model architecture, the effectiveness of the models is unquestionable. So, why is the performance so poor? The root cause lies in the data itself: the model's performance is generally directly proportional to the amount of available data. Take the Flower-102 dataset, for example, which includes 102 classes of flowers but has only 1,020 images in the training set, with some reserved for validation. This is far from sufficient to train deep learning models. Why, then, does the OxfordIIITPet dataset, which has more data per class on average than Flower-102, perform even worse? This is due to the diversity of the data (i.e., the similarity of features within the same class) and the variety of backgrounds.

Figures <ref> and <ref> show some results from the OxfordIIITPet and Flower-102 datasets. Both figures indicate that prediction results are worse for categories with diverse backgrounds or inherent diversity, as seen in Figures <ref>.b and <ref>.b. Conversely, results are relatively better when the background is consistent but the species vary, as shown in Figures <ref>.a and <ref>.a. The best results are achieved when both the background and species variation are minimal, as shown in Figures <ref>.c and <ref>.c. Although the Flower-102 dataset has fewer images, its backgrounds are more consistent and there is less variation among the species. In contrast, animals in the OxfordIIITPet dataset can appear in various forms, while flowers in the Flower-102 dataset are usually photographed from relatively fixed angles, leading to better performance.

Does this mean that deep learning models are not powerful enough or lack generalization and reasoning capabilities? On the contrary, achieving around 25% accuracy on OxfordIIITPet and about 40% on Flower-102 with such limited data actually demonstrates the strength and generalization capability of deep learning models. A common misconception about deep learning is that humans, even if unfamiliar with certain species, can accurately distinguish OxfordIIITPet and Flower-102 images by carefully examining them, thus believing that deep learning lacks generalization or human-like inferencing ability. This viewpoint is flawed. Assuming the human brain functions as a UAT model, it can automatically disregard background information before seeing these images, as this background is frequently encountered in daily life, and the UAT in the human brain has already fitted this information. Therefore, humans only need to focus on specific information and then use it to train their brains to find similarities and differences. For instance, a 10-year-old child’s brain, observing the world at 60 frames per second, would have processed about 1 × 10^10 high-resolution images by the age of 10, many of which include common objects and backgrounds in daily life. Therefore, criticizing deep learning models for being unable to train on small datasets and lacking universality and generalization is unfair.

Thus, we should train on large datasets and then fine-tune on smaller ones. This is also the current deep learning training method, showcasing the powerful generalization capability of deep learning models, especially when there is some correlation between datasets. Following this fine-tuning approach, we fine-tuned the pre-trained ViT model on the aforementioned datasets for just 10 epochs. The results, as shown in Table <ref>, indicate a significant improvement in prediction performance.

Figure <ref> shows the rationality behind fine-tuning. Since we essentially adjust parameters based on the input to produce target results using UAT, if there is a correlation in the data, the corresponding parameters derived from the input (α_i, 𝐖_i, and θ_i; i=1,·,N) will be similar. Therefore, fine-tuning is feasible.








§ CONCLUSION

In this paper, we propose a parallel network based on the Universal Approximation Theorem (UAT) and validate its effectiveness on large datasets. We analyze the characteristics of both serial and parallel networks. Serial networks, especially those based on residual connections, often have higher degrees of parameter freedom, leading to stronger generalization capabilities. However, as the network depth increases, the inference time also becomes longer. On the other hand, parallel networks trade computational resources for time, meaning their runtime is not affected by the number of network layers. This allows for the design of multi-layer parallel networks that can replace serial networks to achieve fast inference. Additionally, we analyze the impact and reasons data have on deep learning. From the UAT perspective, training on large datasets essentially enhances the model's ability to fit data.




















§ THE UAT FORMAT OF TRANSFORMER










 §.§ The UAT Format of Transformer-based ViTs

Eq.(<ref>) and (<ref>) give an example of the matrix-vector format of MHA and FFN. Eq. (<ref>) and  Eq. (<ref>) give the corresponding UAT format of 𝐱_i+1 and 𝐱_i+2. It is obvious that the bias and weight terms in the UAT corresponding to Transformed can change dynamically based on the input.

Due to the complexity of the variables involved, we will combine some terms to align the resulting equations with the standard UAT form. To distinguish these parameter variables, we will keep the original variables unchanged. If a parameter variable is influenced only by other variables, we will denote it with a bar on top, such as 𝐖̅. If it is influenced by the input 𝐱_i, we will mark it with a hat, such as 𝐖̂.

















 Paper Title: Point-supervised Brain Tumor Segmentation with Box-prompted MedSAM.






Point-supervised Brain Tumor Segmentation with Box-prompted MedSAM
    Xiaofeng Liu, Jonghye Woo, Chao Ma, Jinsong Ouyang, and Georges El Fakhri
X. Liu, C. Ma, J. Ouyang and G. El Fakhri are with the Department of Radiology and Biomedical Imaging, Yale University, New Haven, CT 06519.
J. Woo is with the Massachusetts General Hospital and Harvard Medical School, Boston, MA 02114.

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



§ ABSTRACT 
Delineating lesions and anatomical structure is important for image-guided interventions. Point-supervised medical image segmentation (PSS) has great potential to alleviate costly expert delineation labeling. However, due to the lack of precise size and boundary guidance, the effectiveness of PSS often falls short of expectations. Although recent vision foundational models, such as the medical segment anything model (MedSAM), have made significant advancements in bounding-box-prompted segmentation, it is not straightforward to utilize point annotation, and is prone to semantic ambiguity. In this preliminary study, we introduce an iterative framework to facilitate semantic-aware point-supervised MedSAM. Specifically, the semantic box-prompt generator (SBPG) module has the capacity to convert the point input into potential pseudo bounding box suggestions, which are explicitly refined by the prototype-based semantic similarity. This is then succeeded by a prompt-guided spatial refinement (PGSR) module that harnesses the exceptional generalizability of MedSAM to infer the segmentation mask, which also updates the box proposal seed in SBPG. Performance can be progressively improved with adequate iterations. We conducted an evaluation on BraTS2018 for the segmentation of whole brain tumors and demonstrated its superior performance compared to traditional PSS methods and on par with box-supervised methods.


 




§ INTRODUCTION
 



The costly labeling effort in medical image delineation significantly hinder the development of data-driven AI models. There are increasing interests on weakly supervised segmentation to utilize bounding box or even a single point as label for training supervision <cit.>. However, due to the absence of accurate size and boundary guidance, there is a large performance gap between point-supervised medical image segmentation (PSS) and mask/box-supervised counterparts <cit.>. 

The recent progress of vision foundational models has achieved breakthroughs in several weakly supervised tasks, benefited by their strong zero-shot generalizability. For instance, the point-prompt natural image segment anything model (SAM) has been applied to enhance pseudo labels for point-supervised segmentation <cit.>. Recently, the medical image version of SAM (MedSAM) <cit.> has been trained with 1.5 million segmented medical images to achieve generalizable box-prompt segmentation. MedSAM takes both an image slice and a bounding box, i.e., prompt, to predict the possible segmentation mask within the box. Notably, the natural image SAM itself follows point-prompt design. Therefore, it is not straight forward to integrate box-prompt MedSAM to PSS task as <cit.>. In addition, a significant limitation of SAM/MedSAM is the lack of classification ability, resulting in class / structure-agnostic segmentation results <cit.>. They are designed for general use, while fails to accurately delineate the specific lesion or structure as desired <cit.>.  



 

To our knowledge, this is the first attempt to integrate MedSAM to facilitate semantic-aware PSS. We adopt an iterative framework to achieve coarse-to-fine progression of the bounding box. The point prompt is first converted to a box-proposal seed. Then, we configure a semantic box-prompt generator (SBPG) to propose and pick the reasonable box according to semantic similarity as <cit.>. It is followed by a box prompt guided spatial refinement (PGSR) to utilize the generalizable MedSAM to predict the segmentation mask. In addition, the smallest box that covers the mask is further used as the box seed in the next round of SBPG. We do not need to fine-tune the MedSAM to be semantic aware, which is prone to catastrophic forgetting. Notably, only the prompt refiner θ with the encoder part of the segmentor model will be trained, which involves about half the parameters of mask-supervised UNet training. 

We demonstrated its effectiveness in BraTS2018 for PSS of the whole brain tumor segmentation from T2-weighted MRI slices. We show that in testing, with 3 to 5 rounds of iteration, the point-prompt can achieve superior performance to approximate the box-supervised MedSAM <cit.>.








§ METHODOLOGY

In PSS, we are given tuplet {x^i,ρ_c^i,s^i}, in which x^i can be an MR slice, while the point prompt ρ^i_c={ρ^i_x,ρ^i_y,c} indicates the 2D spatial coordinates and the interested class c, e.g., tumor or normal tissue. We would expect the segmentation mask prediction ŝ^i_c to approximate the expert annotation s^i_c. 

To catering the box-prompt MedSAM <cit.>, an initial box proposal seed of class c, i.e., b_c:0^i, is generated from ρ^i_c, which is centered on point location (ρ^i_x,ρ^i_y) with the size of X× Y. We do not expect b_c:0^i fit the region of interests very well as there is no prior information about the size in the initial step. Then, the proposal bag ℬ^i_c={b_c:n^i}_n=1^N is created by scaling b_c:0^i with N different scales. Therefore, the boxes in group ℬ^i_c has strong spatial correlation with ρ_c^i, which avoid the inefficient random proposal in whole image <cit.>. 

To enable the model be aware of the specific semantic class, for example tumor in our task, a parameterized prompt refiner θ with a segmentor encoder and two fully connected layers is applied. Specifically, for the training sample with segmentation mask label s^i_c, we convert s^i_c to the corresponding smallest rectangle bounding box label b^i_c. Then, we store a set of feature f^i_c:n=θ(x^i,b_c:n^i) in the most recent M batches to a memory buffer, and calculate the mean of {f^i_c:nm}_n,m=1^N,M as the prototype 𝒱_c. As conventional multiple instance learning (MIL) based methods <cit.>,
the instance level probability indicating the likelihood of b^i_c:n matches 𝒱_c can be approximated with 
    p(b^i_c:n,𝒱_c)=e^cos(f^i_c:n,𝒱_c)/∑_c=1^C e^cos(f^i_c:n,𝒱_c),
where cos(·) indicates the cosine similarity, and C is the number of segmentation category. With the ground truth b^i_c, θ is trained with the binary cross entropy loss of 

    ℒ=-∑_c=1^C {c^i_1log∑_n^N p(b^i_c:n,𝒱_c) + (1-c^i_1)log(1-∑_n^N p(b^i_c:n,𝒱_c))},
where c^i_1 is the one-hot class label. Therefore, the prompt refiner is optimized to be semantic aware for the specific anatomical structure. The prototype in buffer is also updated with the new model θ. With p(b^i_c:n,𝒱_c) for each b^i_c:n, we pick the highest probability one  in ℬ^i_c as our optimal proposal b^i_c:* in current stage, which is also our bounding box inference.


Then, in the PGSR module, the off-the-shelf MedSAM <cit.> takes both x^i and b^i_c:* to predict the possible segmentation mask ŝ^i_c. Notably, no information about the interested class can be directly informed in MedSAM. Although the MedSAM has strong ability of zero-shot segmentation to delineate the relatively accurate mask within the current box prompt b^i_c:*, the alignment of ŝ_c^i is highly dependent on the precise b^i_c:*. This motivated us to refine b^i_c:* with a more suitable proposal seed b_c:0^i in SBPG. 

In practice, we can simply apply the smallest bounding box b̂^i_c:0 for ŝ^i_c:0, and use b̂^i_c:0 as our box proposal seed in the next round of SBPG. With T rounds of iteration, we expect that the box-prompt can be progressively refined to inform the MedSAM to predict an accurate segmentation mask. 

In particular, in the last round of testing, we do not need to generate b̂^i_c:0. Instead, the T-th round ŝ^i_c:0 is used as our final prediction.









§ EXPERIMENTS AND RESULTS

We evaluated on BraTS dataset for whole tumor segmentation as <cit.>, and use T2-weighted slices as our input. We used the 80%/20% split for training or testing. We generate the smallest bounding box for the mask and use its center as ρ^i_c. Our prompt refiner adopted the encoder part of the segmentor in <cit.>, which is followed by two fully connected layers of 1024 and 256 dimensions, respectively. We empirically set X× Y=21×21 and search proper T∈{1,2,3,5,10}. We followed <cit.> to compare with the conventional point-supervised method of WISE-Net without the use of the vision fundation model. Also, we directly input the ground truth of the bounding box into MedSAM, which can be “upper bound" of the performance of pointwise method. In Tab.<ref>, we can see that our model with T≥5 outperforms WISE-Net by a large margin. The Dice score is close to the box-supervised MedSAM.   





§ CONCLUSION

In this work, we proposed an efficient iterative framework to enable box-prompt MedSAM to point-supervised medical image segmentation. A lightweight prompt refiner with only encoder is specifically trained for the interested structural class. We show that the fixed off-the-shelf large model can flexibly target the specific downstream task without the need of large scale fine-tuning. The medical PSS performance can be largely improved by the advance of vision foundation model, e.g., MedSAM, which show great promise for facilitating the image-guided interventions. 
 


 


§ ACKNOWLEDGMENT

This work is supported by R21EB034911, R01CA165221, R01DC018511, P41EB022544, and NAIRR240016. 

ieee





 Paper Title: Hierarchical Conditioning of Diffusion Models Using Tree-of-Life for Studying Species Evolution.





 



Hierarchically Conditioned Diffusion Model







M. Khurana et al.





 Virginia Tech, Blacksburg VA, USA 

            {mridul, karpatne}@vt.edu Oak Ridge National Laboratory, Oak Ridge TN, USA University of California, Irvine CA, USA Tulane University, New Orleans LA, USA Battelle, Columbus OH, USA Duke University, Durham NC, USA University of North Carolina at Chapel Hill, Chapel Hill NC, USA The Ohio State University, Columbus OH, USA Rensselaer Polytechnic Institute, Troy NY, USA 


Hierarchical Conditioning of Diffusion Models Using Tree-of-Life for Studying Species Evolution
    Mridul Khurana10009-0003-9346-3206 Arka Daw20009-0006-3319-1271 M. Maruf 1 0000-0003-0637-1753 Josef C. Uyeda 1 0000-0003-4624-9680 Wasila Dahdul 30000-0003-3162-7490 Caleb Charpentier 1 0000-0002-9787-7081 Yasin Bakış 40000-0001-6144-9440 Henry L. Bart Jr. 40000-0002-5662-9444 Paula M. Mabee 50000-0002-8455-3213 Hilmar Lapp 60000-0001-9107-0714 James P. Balhoff 70000-0002-8688-6599 Wei-Lun Chao 80000-0003-1269-7231 Charles Stewart 90000-0001-6532-6675 Tanya Berger-Wolf80000-0001-7610-1412 Anuj Karpatne 10000-0003-1647-3534

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



fancy
 
[L]Published as a conference paper at ECCV 2024
[C]




§ ABSTRACT 
A central problem in biology is to understand how organisms evolve and adapt to their environment by acquiring variations in the observable characteristics or traits of species across the tree of life. 
With the growing availability of large-scale image repositories in biology and recent advances in generative modeling, there is an opportunity to accelerate the discovery of evolutionary traits automatically from images. 
Toward this goal, we introduce Phylo-Diffusion, a novel framework for conditioning diffusion models with phylogenetic knowledge represented in the form of HIERarchical Embeddings (HIER-Embeds). 
We also propose two new experiments for perturbing the embedding space of Phylo-Diffusion:  trait masking and trait swapping, inspired by counterpart experiments of gene knockout and gene editing/swapping.
Our work represents a novel methodological advance in generative modeling to structure the embedding space of diffusion models using tree-based knowledge. Our work also opens a new chapter of research in evolutionary biology by using generative models to visualize evolutionary changes directly from images.
We empirically demonstrate the usefulness of Phylo-Diffusion in capturing meaningful trait variations for fishes and birds, revealing novel insights about the biological mechanisms of their evolution. 
[Model and code can be found at https://imageomics.github.io/phylo-diffusion/]
  
  





§ INTRODUCTION









Given the astonishing diversity of life forms on the planet, an important end goal in biology is to understand how organisms evolve and adapt to their environment by acquiring variations in their observable characteristics or traits (, beak color, stripe pattern, and fin curvature) over millions of years in the process of evolution. Our knowledge of species evolution is commonly represented in a graphical form as the “tree of life” (also referred to as the phylogenetic tree <cit.>, see Figure <ref>), illustrating the evolutionary history of species (leaf nodes) and their common ancestors (internal nodes). Discovering traits that are heritable across the tree of life, termed evolutionary traits, is important for a variety of biological tasks such as tracing the evolutionary timing of trait variations common to a group of species and analyzing their genetic underpinnings through gene-knockout or gene-editing/swapping (, CRISPR <cit.>) experiments. However, quantifying trait variations across large groups of species is labor-intensive and time-consuming, as it relies on expert visual attention and subjective definitions <cit.>, hindering rapid scientific advancement <cit.>. 

The growing deluge of large-scale image repositories in biology <cit.> presents a unique opportunity for machine learning (ML) methods to accelerate the discovery of evolutionary traits automatically from images. In particular,  with recent developments in generative modeling such as latent diffusion models (LDMs) <cit.>, we are witnessing rapid improvements in our ability to control the generation of high-quality images based on input conditioning of text or image prompts. This is facilitating breakthroughs in a variety of commercial use-cases of computer vision where we can analyze how changes in the input prompts affect variations in the generated images <cit.>.
We ask the question: can we leverage LDMs to control the generation of biological images of organisms conditioned on the position of a species in the tree of life? In other words, can we encode the structure of evolutionary relationships among species and their ancestors as input conditions in LDMs?
This can help us analyze trait variations in generated images across different branches in the phylogenetic tree, revealing novel insights into the biological mechanisms of species evolution. 


Toward this goal, we introduce Phylo-Diffusion, a novel framework for discovering evolutionary traits of species from images by conditioning diffusion models with phylogenetic knowledge (see Figure <ref>). One of the core innovations of Phylo-Diffusion is a novel HIERarchical Embedding (HIER-Embed) strategy that encodes evolutionary information of every species as a sequence of four vectors, one for each discretized level of ancestry in the tree of life (covering different evolutionary periods). 
We also propose two novel experiments for analyzing evolutionary traits by perturbing the embedding space of Phylo-Diffusion and observing changes in the features of generated images, akin to biological experiments involving genetic perturbations.
First, we introduce Trait Masking, where one or more levels of information in HIER-Embed are masked out with noise to study the disappearance of traits inherited by species at those levels. This is inspired by  gene knockout experiments <cit.>, wherein one or more genes are deactivated or “knocked out” to investigate the gene's function, particularly its impact on the traits of the organism. Second, we introduce Trait Swapping, where a certain level of HIER-Embed in a reference species is swapped with the embedding of a sibling node at the same level, similar in spirit to gene editing/swapping experiments made possible by the CRISPR technology <cit.>. The goal of trait swapping is to visualize trait differences at every branching point in the tree of life that results in the diversification of species during evolution.


Here are the main contributions of this paper. Our work represents a novel methodological advance in the emerging field of knowledge-guided machine learning (KGML) <cit.> to structure the embedding space of generative models using tree-based knowledge. Our work also opens a new chapter of research in evolutionary biology by using generative models to visualize evolutionary changes directly from images, which can serve a variety of biological use-cases. For example, Phylo-Diffusion can help biologists automate the discovery of synapomorphies, which are distinctive traits that emerge on specific evolutionary branches and are crucial for systematics and classification <cit.>. Our proposed experiments of trait masking and swapping can also be viewed as novel image-based counterparts to genetic experiments, which traditionally take years. Our work thus enables biologists to rapidly analyze the impacts of genetic perturbations on particular branches of the phylogenetic tree–a grand challenge in developmental biology <cit.>.
We empirically demonstrate the usefulness of Phylo-Diffusion in capturing meaningful trait changes upon perturbing its embedding for fishes and birds, generating novel hypotheses of their evolution.


§ RELATED WORKS AND BACKGROUND


 

  
Interpretable ML:
Discovering evolutionary traits from images requires the identification and interpretation of fine-grained features in images that define and differentiate species. Several methodologies have recently been developed in the field of interpretable ML for localizing image regions that contain discriminatory information of classes, including ProtoPNet <cit.>, PIP-Net <cit.> and INTR <cit.>. 
Despite their effectiveness and applicability across a wide range of applications, these methods are not directly suited for our target application of discovering evolutionary traits for two primary reasons. First, they are not designed to incorporate structured biological knowledge (, knowledge of the tree of life) in the learning of interpretable features, and thus are unable to provide biologically meaningful explanations of feature differences across groups of species in the phylogenetic tree, which is key to discovering evolutionary traits. Second, since most methods in interpretable ML are designed for the task of classification, it is non-trivial to integrate them in generative modeling frameworks to produce synthetic images with controlled perturbations in the embedding space similar to gene knockout and gene editing/swapping experiments, in contrast to our proposed experiments in Phylo-Diffusion.



  
Phylogeny-guided Neural Networks (Phylo-NN):

A recent work closely aligned with our goal of discovering evolutionary traits directly from images is Phylo-NN <cit.>. Phylo-NN uses an encoder-decoder architecture to represent images of organisms as structured sequences of feature vectors termed “Imageomes”, where different segments of Imageomes capture evolutionary information from varying levels of ancestry in the phylogenetic tree. While Phylo-NN shares several similarities with our proposed framework, Phylo-Diffusion, in terms of motivations and problem formulations, there are also prominent differences. The primary goal of Phylo-NN is specimen-level image reconstruction, whereas Phylo-Diffusion considers a different goal of controlling image generation at the species-level. As a result, Phylo-NN learns a unique Imageome sequence for every organism, enabling us to study the variability in individuals from the same species and the analysis of similarity in Imageome segments learned at shared ancestry levels. On the other hand, Phylo-Diffusion learns a unique embedding for every species and ancestor node in the tree of life, which serves as input conditions to generate distributions of synthetic images. Phylo-Diffusion thus uses hard constraints to ensure that all species with a common ancestor learn the exact same embeddings at their shared ancestry levels, making it easy to analyze trait commonalities and variations across groups of species, in contrast to Phylo-NN. 
Additionally, Phylo-Diffusion allows for perturbations in the embedding space of generative models in biologically meaningful ways inspired by gene knockout and gene editing/swapping experiments, going beyond the capabilities of Phylo-NN. We consider Phylo-NN as a baseline in our experiments to compare its performance with Phylo-Diffusion.



  
Background on Latent Diffusion Models (LDMs):
One of the state-of-the-art approaches in generative modeling is the framework of Diffusion Models <cit.>, which learns a target distribution p(x) by incrementally transforming a noisy sample x generated from a Gaussian distribution 𝒩(0, I) into one that is more likely to be generated from p(x) over a series of timesteps T.
While early frameworks of diffusion models (, DDPM <cit.>, DDIM <cit.> and ADM <cit.>) suffered from high computational costs and long training/inference times, Latent Diffusion Models (LDMs) <cit.> are able to address these concerns to a large extent by operating in a compressed latent space, significantly accelerating their ability to generate high-resolution images. The basic idea of LDMs is to train a separate auto-encoder to map an input image x into its latent representation z_0 = ℰ(x) using encoder ℰ, which when fed to decoder 𝒟 produces a reconstruction of the original image,  x̃ = 𝒟(z̃_0). LDMs employ diffusion models in the compressed latent space z by modeling the conditional probability of the reverse diffusion process as z̃_t-1∼ p_θ(z̃_t-1 | z̃_t, y, t), where y is the input condition. This is implemented using a conditional denoising U-Net backbone ϵ_θ(z_t, y, t) with learnable parameters θ.
LDMs also pre-process y using a domain-specific encoder 𝐄 = τ_ϕ(y) trained alongside the U-Net backbone ϵ_θ to project y into the intermediate layers of ϵ_θ using cross-attention mechanisms. The learnable parameters of LDMs are trained by minimizing the following loss function:

    ℒ(θ,ϕ) = 𝔼_z_t, y, t, ϵ∼𝒩(0, I)[ ϵ̂_θ(z_t, τ_ϕ(y), t) - ϵ^2 ]



§ PROPOSED FRAMEWORK OF PHYLO-DIFFUSION




 §.§ Hierarchical Embedding (HIER-Embed)
 


Phylo-Diffusion uses a novel hierarchical embedding (HIER-Embed) strategy to structure the embedding E of every species node using phylogenetic knowledge. 
As a first step, we consider a discretized version of the phylogenetic tree involving four ancestral levels, level-1 to level-4, where every level corresponds to a different range of time in the process of evolution. (See <Ref> for a detailed characterization of the four ancestry levels for fish species used in this study.) Given a set of n species,
𝒮 = {S_1, S_2, S_3, ..., S_n}, let us represent the position of species S_i ∈𝒮 in the phylogenetic tree at the four ancestry levels as  {S_i^1, S_i^2, S_i^3, S_i^4}, where S_i^l represents the ancestor node of S_i at level-l. Hence, if two species S_i and S_j share common ancestors till level-k, then S_i^l = S_j^l for l=1 to k.
We define the level-l embedding of species S_i as:

    𝐄_i^l = (S_i^l)   ∈ℝ^d',

where (.) is a learnable embedding layer that provides a simple way to store and look-up the trained embeddings of every node. The combined hierarchical embedding (HIER-Embed) of species S_i is obtained by concatenating its embeddings across all four levels as follows:


    𝐄_i = τ (S_i) =  𝐄_i^1,  𝐄_i^2,  𝐄_i^3,  𝐄_i^4  ∈ℝ^d,

where  denotes the concatenation operation and y=S_i is the input condition used in LDMs. Note that different segments of 𝐄_i capture information about the traits of S_i acquired at different time periods of evolution. In particular, we expect the embedding vectors learned at earlier ancestry levels of 𝐄_i to capture evolutionary traits of S_i common to a broader group of species. On the other hand, embeddings learned at later ancestry levels are expected to be more specific to S_i. 
In the following, we present two novel experiments for studying evolutionary traits by perturbing the embedding space learned by HIER-Embed.







 §.§ Proposed Experiment of Trait Masking




The goal of this experiment is to verify if HIER-Embed is indeed able to capture hierarchical information in its level-embeddings such that masking information at lower levels of the embedding only erases traits acquired at later stages of evolution while retaining trait variations learned at earlier levels. In other words, we want to verify that the embeddings learned by HIER-Embed at level-l capture information common to all descendant species that are part of the same sub-tree at level-l. 
<Ref> represents a schematic diagram of the process followed for trait masking. We start with the combined embedding containing information at all four levels, [, , , ]. To examine what is learned at the last level of this embedding, we mask it out by substituting it with Gaussian noise defined as 𝐳_𝐧𝐨𝐢𝐬𝐞∼𝒩(0, I) ∈ℝ^d'. This results in the perturbed embedding [, , , ], effectively eliminating the species-level (or ) information. This masking should prompt the model to generate images that reflect only the information learned up to the third level while obscuring species-level details. We can extend this experiment by incrementally introducing noise at later levels, , at both levels 3 & 4, and so on.



  
Expected Changes in Probability Distributions: Note that when all four level embeddings are used, [, , , ], the generated images are expected to be classified to a unique species S_i. In terms of probability distributions, the probability of predicting species S_i should be distinctly higher than the probability of predicting any other species. However, when we mask out certain level embeddings (, mask out information at level 4), we are intentionally removing information necessary to distinguish species S_i from its siblings species that are part of the same sub-tree (e.g., those that share a common ancestor at level 3). For this reason, we expect the generated images to show higher probabilities of being classified as any of the descendant species of the sub-tree, compared to the other species that are outside of the sub-tree. 
To quantify this behavior, we can measure the
change in probability distributions for species within and outside the sub-tree after masking out an internal node. Since we expect probabilities to increase only for species within the sub-tree, the mean increase in probabilities for within-subtree species should be higher than that for out-of-subtree species, as empirically demonstrated later in the Results Section.



 §.§ Proposed Experiment of Trait Swapping



In trait swapping, we substitute the level-l embedding of a source species with the level-l embedding of a sibling subtree at an equivalent level. 
<Ref> shows a schematic representation of the trait swapping experiment 
where the level-3 embedding (green) of a source species is replaced with its sibling level-3 embedding (yellow). 
Images generated for this perturbed embedding are expected to retain all of the traits of the source species except the swapped embedding, which should borrow traits from the sub-tree rooted at the sibling node. Visualizing trait differences in the generated images before and after trait swapping can help us understand the evolutionary traits that branched at a certain level (, those leading to the diversification of green and yellow sub-trees at level-3 in the example phylogeny of Figure <ref>). 
In terms of the probability distribution, similar to trait masking, we expect to see a drop in the probabilities of the source species (pink), and simultaneously, we expect an increase in probabilities for all the descendent species in subtree at node yellow, red and purple.





§ EVALUATION SETUP




  
Datasets:
We use a collection of fish images as our primary dataset for evaluation. This dataset was procured from the Great Lakes Invasives Network (GLIN) <cit.>  Project, comprising a total of 5434 images spanning 38 fish species. We obtained the phylogenetic tree of fish species from opentree <cit.> python package (see <Ref> for details on the phylogenetic tree). The raw museum images were pre-processed and resized to 256 × 256 pixels and the dataset was partitioned into training and validation sets, following a 75-25 split. We provide additional results on the CUB-200-2011 dataset <cit.> of bird species in <Ref>.



  
Baselines of Conditioning Mechanisms:
(1) Class Conditional: One of the simplest ways of encoding information about a species class is to map class labels y ∈ [1, N_c] to a fixed d-dimensional embedding vector e ∈ℝ^d using a trainable embedding layer. Note that the resulting embeddings are not designed to contain any hierarchical information in contrast to HIER-Embed. (2) Scientific Name Encoding:  The scientific name of a species contains valuable biological information typically comprising of a combination of the genus name and species name. Since species that share their genus name are likely to contain common phylogenetic traits, we use them as a baseline for conditioning LDMs for discovering evolutionary traits. Specifically, we employ a pre-trained frozen CLIP model <cit.> to encode the scientific names of species into fixed d-dimensional embeddings.



  
Training details:
We used d'=128 as the embedding dimension for each level of HIER-Embed, which when concatenated across the four levels produces the combined hierarchical embedding of d=512 dimensions. Phylo-Diffusion uses this d-dimensional embedding to condition LDMs through cross-attention in denoising the U-Net backbone and train LDMs without classifier-free guidance. We used VQGAN <cit.> as the backbone encoder-decoder to achieve the latent representations desired for LDMs with a downsampling factor of 4. All the models with different encoders are trained for 400k iterations, employing the best model checkpoint if convergence occurs early. Additional hyperparameters, such as learning rate, batch size, and U-Net architecture, are detailed in <Ref>.








§ RESULTS




 §.§ Quality of Generated Images

<Ref> compares the quality of generated images of baselines using the metrics of Fréchet Inception Distance (FID) score, Inception Score (IS), and Precision, Recall calculated in the feature space as proposed in <cit.>. Our results show that Phylo-Diffusion is at par with state-of-the-art generative models, achieving an FID of 11.38 compared to LDM's 11.46. We show a sample of generated images in <Ref>, with additional images provided in <Ref>. We also show the robustness of Phylo-Diffusion's results with varying numbers of phylogenetic levels and embedding dimensions in <Ref>.






 §.§ Classification Accuracy


We used a separate model for species classification, specifically a ResNet-18 model <cit.> trained using the same training/validation split as Phylo-Diffusion. The primary objective behind building this classifier is to verify if images generated by Phylo-Diffusion contain sufficient discriminatory information to be classified as their correct species classes. <Ref>  compares the classification F1-scores over 100 samples generated by baseline conditioning schemes. We can see that the synthetic images generated by Phylo-Diffusion achieve the highest F1 score (82.21%), which is quite close to the F1 score of the base classifier on the original test images (85%). We present additional results showing the generalizability of Phylo-Diffusion in classifying generated images to unseen species in <Ref>.



 §.§ Matching Embedding Distances with Phylogenetic Distances


We investigate the quality of embeddings produced by baseline methods by comparing distances in the embedding space with the ground-truth (GT) phylogenetic distances computed from the tree of life, as illustrated in <Ref>.
Ideally, we expect distances in the embedding space of species pairs to be reflective of their phylogenetic distances. 
For Class Conditional, we can see that the distance matrix does not show any alignment with the GT phylogenetic distance matrix.
In the case of Scientific Name Encoding, the distance matrix exhibits notable similarities to the phylogenetic distances, thanks to the hierarchical nature of information contained in scientific names (i.e., genus-name & species-name). However, one limitation of this encoding is its inability to capture inter-genus similarities or differences. In contrast, HIER-Embed shows a distance matrix that closely aligns with the GT phylogenetic distance matrix, validating its ability to preserve evolutionary distances among species in its embedding space.









 §.§ Trait Masking Results



To obtain classification probabilities or logits associated with generated images, we employ the classifier detailed in Section <ref>. For the masked embeddings of subtrees at level 3, defined as [, , ,  ], we analyze the logits of generated images and compare them with those generated without masking. <Ref> demonstrates that for a specific subtree, in this case Lepomis, logits for species within the subtree are higher compared to those for species outside it. This outcome aligns with the expectation that Phylo-Diffusion, when provided with information up to Level 3, can capture overarching characteristics of all species within the given subtree. Additionally, <Ref> presents probability distributions for species within the Lepomis subtree when the full set of hierarchical encodings [, , , ] is provided. It demonstrates that the probabilities are significantly higher for the targeted class, as intended for image generation. After masking, we observe that the generated images are very similar and capture common features of the Lepomis genus. For all our calculations and plots, we generate 100 images for each subtree and node.
<Ref> contains additional histograms that detail logit distributions across all different subtrees at each level, offering comprehensive insights into how the model discriminates and learns the hierarchical structure across different levels and nodes.







  
Quantitative Evaluation of Probability Distrubtions:

To quantitatively evaluate the ability of Phylo-Diffusion to capture hierarchical information and show desired changes in probability distributions after masking, we compute the following metrics. Let us denote the set of all species in the data as 𝒮 and for a given sub-tree at an internal node i of level l, let us denote the subset of descendant species as 𝒮_i^l = {S_1, S_2, …, S_n}. We first compute the reference probabilities P_ref of every species before masking (, by using all four level embeddings). Let us denote the probability of predicting a generated image using all four embeddings of a descendant species S_j ∈𝒮_i^l into species class S_k as P_S_j(S_k). The reference probability of a species S_k can then be given as: 

    P_ref(S_k) = 1/|𝒮_i^l|-1∑_S_j ∈𝒮_i^l ∖ S_k P_S_j(S_k),   if  S_k ∈𝒮_i^l,
    1/|𝒮_i^l|∑_S_j ∈𝒮_i^l P_S_j(S_k),                if  S_k ∉𝒮_i^l.

Note that when S_k is part of the sub-tree, , S_k ∈𝒮_i^l,  P_ref(S_k) is computed by averaging over |𝒮_i^l|-1 probability values since we exclude the case when S_k is used to generate the images. On the other hand, when S_k is outside of the sub-tree, ,  S_k ∉𝒮_i^l, we average over all  |𝒮_i^l| probability values. Given these reference probabilities values before masking, we can compute the change in probability of predicting species S_k after masking as P_diff(S_k) = P_mask(S_k) - P_ref(S_k), where P_mask(S_k) is the probability of predicting a generated image after masking to S_k. We expect P_diff to be larger for descendant species S_k∈𝒮_i^l compared to species that are outside of the sub-tree because of the dispersion of probabilities in a sub-tree as a consequence of masking.  We thus compute the average P_diff for species that belong to subtree 𝒮_i^l as  P_diff^sub(i,l) and species that are outside the subtree 𝒮_i^l as P_diff^out(i,l). 


<Ref> shows the box plot of P_diff^sub and P_diff^out across internal nodes at levels 2 and 3. We observe that species within the subtree exhibit a more pronounced increase in probabilities compared to species outside the subtree, aligning with our expectations. Notably, this trend is consistently observed across both levels 2 and 3. More details of class-wise probability distribution shifts for each level are provided in <Ref>.
The outcomes of these experiments affirm that Phylo-Diffusion effectively identifies unique features at Levels 2 and 3, and captures shared features or traits of any chosen subtree at the internal nodes of the phylogeny.






 §.§ Trait Swapping Results


Figure <ref> shows examples of trait swapping results enabled by Phylo-Diffusion. For the first example (first row of <Ref>), we swap the level-2 embedding of source species Noturus exilis with level-2 embedding of its sibling group Notropis/ Carassius. The goal here is to discover traits of Noturus exilis inherited at level-2 that differentiate it from other groups of species that branched out at this point of time in evolution. We can see that the generated images of the perturbed embedding (center) exhibit the absence of barbels (whiskers) highlighted in purple, while the caudal (or tail) fin is beginning to fork (or split), a trait adopted from Notropis (right). In contrast, other fins such as the dorsal, pelvic, and anal fins highlighted in green remain similar to those of the source species, Noturus exilis (left). This suggests that at level-2, Notropis and Noturus species diverged by developing differences in two distinct traits, barbels and forked caudal fins while keeping other traits intact.




For the second example (<Ref>, row 2), we swap level-2 information of Gambusia affinis (left) with that of Esox americanus (right). The generated images of the perturbed embedding (center) exhibit a more pointed head highlighted in purple, and a slimmer body shape resembling Esox americanus. Notably, the perturbed species retains discoloration at the bottom from the source species highlighted in green. <Ref> presents probability distributions (or logits) of Gambusia affinis before and after trait swapping using the classifier detailed in Section <ref>. We observe a slight decrease in logits for Gambusia affinis and an increase in logits for Esox americanus, consistent with our expectations. In the third row of <Ref>, we swap level-3 information of Lepomis gulosus (left) with that of Morone genus (right). The resulting images from the perturbed embedding (center) capture the horizontal line pattern characteristic of Morone genus, and the dorsal fin highlighted in purple begins to split. All these examples suggest novel scientific hypotheses about differences in evolutionary traits acquired by species at different ancestry levels, which can be validated by biologists in subsequent studies.
Note that our experiments are most effective at levels near the species nodes, specifically at levels 2 & 3, since phylogenetic signal is known to diminish as we move toward the root of the tree <cit.>.
Additional visualizations of trait swapping results are provided in <Ref>.





  
Comparisions with Phylo-NN:
<Ref> compares trait swapping results of Phylo-Diffusion and Phylo-NN for the same set of example species. <Ref> shows trait swapping at level-3 for the source species of Lepomis gulosus (top) and target sub-tree of the Morone genus (bottom). In Phylo-NN, images generated by perturbing the Imageome sequences appear blurry (red circle), while Phylo-Diffusion effectively captures the splitting of dorsal fin (purple circle) and the horizontal stripe pattern of the Morone genus, while maintaining the fin structure of Lepomis gulosus (green circle). Similarly, <Ref> compares trait swapping for Noturus miurus (top) with the target sub-tree of Notropis genus (bottom) at level-2. For Phylo-NN, the perturbed images are almost identical to the source species. However, Phylo-Diffusion shows visible trait differences such as the absence of barbels and the caudal (or tail) fin beginning to fork or split (purple circle), which are traits picked from the target sub-tree of Notropis genus. Note that we had considered the same target sub-tree in <Ref> row 1 and observed similar trait differences in the generated images after perturbation, further validating the ability of Phylo-Diffusion to discover consistent evolutionary traits. We provide additional results comparing Phylo-Diffusion and Phylo-NN trait swapping results in <Ref>.


§ CONCLUSIONS AND FUTURE WORK

In this work, we introduced Phylo-Diffusion, a novel framework for discovering evolutionary traits from images by structuring the embedding space of diffusion models using tree-based knowledge. 
In the future, our approach can be extended to work on other applications involving image data linked with phylogenies or pedigrees.
Our work also has limitations that need to be addressed in future research. For example, while our current work is limited to discretized trees with a fixed number of levels, future works can focus on discovering evolutionary traits at every internal node of the phylogenetic tree with varying levels without performing any discretization. Future works can also attempt to capture convergent changes in evolution, i.e., changes that occur repeatedly in different branches of the tree, and perform ancestral state reconstruction with uncertainty estimates.



§ ACKNOWLEDGMENTS

This research is supported by National Science Foundation (NSF) awards for the HDR
Imageomics Institute (OAC-2118240). We are thankful for the support of computational resources provided by the Advanced Research Computing (ARC) Center at Virginia Tech.

This manuscript has been authored by UT-Battelle, LLC, under contract DE-AC05-00OR22725 with the US Department of Energy (DOE). The US government retains and the publisher, by accepting the article for publication, acknowledges that the US government retains a nonexclusive, paid-up, irrevocable, worldwide license to publish or reproduce the published form of this manuscript, or allow others to do so, for US government purposes. DOE will provide public access to these results of federally sponsored research in accordance with the DOE Public Access Plan (
https://www.energy.gov/doe-public-access-plan).









splncs04




§ SUPPLEMENTARY MATERIALS

Here is a summary of additional details and experiments included in the Appendices.


    
  * <ref>: Hyperparameter Settings and Training Details
    
  * <ref>: Details about Phylogenetic Tree
    
  * <ref>: Additional Details about Trait Masking Experiments
    
  * <ref>: Additional Examples for Trait Swapping Experiments
    
  * <ref>: Additional Comparisons with PhyloNN
    
  * <ref>: Additional Samples of Generated Images
    
  * <ref>: Ablation Results
    
  * <ref>: CUB Dataset Results





§ HYPERPARAMETER SETTINGS AND TRAINING DETAILS
 

<Ref> lists all the hyperparameters for the models trained. We used cross-attention as the conditioning mechanism for all the models and all the models were trained from scratch. At the inference stage, we used DDIM <cit.>sampling with 200 steps. For computing metrics like FID, IS, , we use ADM's <cit.> TensorFlow evaluation script.




All diffusion models require about 7 days to train on a single A100 GPU for both bird and fish datasets. Inference throughput is 0.9 samples/sec using DDIM with 200 steps computed over generating 100 images per class. We do not have any additional overheads in training and inference time compared to LDMs.



§ DETAILS ABOUT PHYLOGENETIC TREE
 
<Ref> shows the phylogeny tree for all the species in the fish dataset along with the information of the four discrete levels used in our study (marked by different colored circles). <Ref> and <ref> list out all the groupings (subtrees) made after discretizing the tree into four levels where the fourth level is the species itself.












§ ADDITIONAL DETAILS ABOUT TRAIT MASKING EXPERIMENTS




  §.§.§ Additional Visualizations of Changes in Probability Distributions after Masking:
 <Ref>, <ref>, <ref>, <ref> and <ref> show additional examples of changes in probability distributions when level-4 information is replaced with noise. In each figure, the first two plots display probability distributions (or logits) of images generated using embeddings from all four levels, [, , , ], of two representative species sharing a common ancestry up to level-3 (highlighted in green). We show that the logits are higher for the targeted species as expected. The third plot logits after masking level-4 embeddings, leading to a dispersion of probabilities across all descendant species within the subtree up to level-3 (highlighted in green). The only exception is <Ref>, where there is some skewness in the logits of descendant species, which is likely due to the data imbalance across classes at higher levels of the tree and also due to biases in the classifier (classifier test accuracy is 85% as reported in <Ref>). In this case, the classifier sometimes misclassifies Notropis boops as Notropis blennius in the first plot and Notropis dorsails as Notropis buccatus in the second plot. Consequently, the third plot for the Notropis subtree shows a higher probability for Notropis blennius. Similarly, <Ref>, <ref> and <ref> provide examples of trait masking where both Level 3 and 4 are replaced with noise, [, , , ]. We observe a similar trend in the dispersion of probabilities across all descendant species within the same subtree at level-2. In all trait masking visualizations, we consistently observe that logits of generated images of species within the subtree (highlighted in green) are higher than for species outside the subtree. This demonstrates Phylo-Diffusion's ability to effectively capture hierarchical information at various levels of the phylogenetic tree.






















  §.§.§ Additional Quantitative Results of Trait Masking:
 Tables <ref> and <ref> show the change in probabilities for different nodes at levels 3 and 2, respectively. We can see that indeed P_diff^sub is larger than P_diff^out for all internal nodes at levels 2 and 3 (except node 5 at level 2), indicating that Phylo-Diffusion is capturing the necessary hierarchical information required for the dispersion of probabilities after masking. <Ref> in the main paper shows the box plots of Tables <ref> and <ref>. <Ref> summarizes this information by showing the average P_diff^sub and P_diff^out for all nodes at a given level. It is important to note that for this experiment, we focus on nodes that have more than one species in the defined subtree.











§ ADDITIONAL EXAMPLES FOR TRAIT SWAPPING EXPERIMENTS




  §.§.§ Additional Visualizations of Trait Swapping Experiments:
 <Ref> illustrates trait swapping for the source species Noturus exilis (left), where the information at Level-2 is swapped with that of a sibling subtree at Node B (right). The image in the center is generated using the trait swapped embedding. This visualization of the perturbed species helps us study the trait changes that would have branched out at level-2 between Node A and Node B.In the generated image (center),  we observe the absence of barbels(whiskers), and the caudal fin (tail) is getting forked (or split) highlighted in pink, which are traits adopted from species in the subtree at B (Notropis). Whereas other fins like the dorsal, pelvic, and anal fin still resemble the source species Noturus exilis highlighted in green. The same is also reflected in the change of probability distribution after perturbations; the probability distribution of source species Noturus exilis decreases and the probability of it being a Notropis increases slightly.

Similarly in <Ref>, for the studied species Lepomis gulosus (left), the information at Level 3 is swapped with subtree at Node B (Morone). The perturbed species generated (center) captures traits from both lineages. The spotted pattern in the body and fins is retained from Lepomis (Node A) but these spots now start to follow the horizontal stripes pattern observed in Morone (species at Node B). Additionally, the dorsal fin highlighted in pink starts to split into two, with the left half retaining the spiny structure from Lepomis highlighted in green. This observation suggests that the species at Node A and B possess distinct traits at this level-3 branching node. The same is reflected in the shift in probability distributions towards species in Node B after trait swapping.








§ ADDITIONAL COMPARISONS WITH PHYLONN


<Ref> compares the trait swapping experiment for Phylo-Diffusion with the PhyloNN baseline, where level-2 information of Gambusia affinis is replaced with that of Esox americanus. In the highlighted pink circle, the face of the image generated after perturbations (center) becomes more pointed, and the body shape flattens to resemble Esox americanus. This perturbed image also retains traits like the caudal (tail) fin and the black-spotted pattern towards the bottom (highlighted in green) from the source species, Gambusia affinis. The differences observed with Phylo-Diffusion are notable, whereas the PhyloNN generates a perturbed image nearly identical to the original, showing no significant changes.

Similarly, <Ref> shows a comparison after replacing level-2 information of Notropis husonius with that of Noturus. For Phylo-Diffusion, the caudal (tail) fin is vibily joining highlighted in pink, resembling the caudal fins of Noturus. This change is analogous to <Ref>, where level-2 information of Noturus was replaced with Notropis (vice-versa), resulting in the caudal (tail) fin getting forked or split. Hence, this helps us understand that at Level-2, the two species diverged to develop different caudal fins. However, for PhyloNN, the generated image after trait-swapping is blurry, and most of the traits still closely resemble close the source species, which is unlikely given that the level-2 embeddings have been replaced.









§ ADDITIONAL SAMPLES OF GENERATED IMAGES


<Ref> shows additional examples of generated images for different species using Phylo-Diffusion. Each row of the figure depicts the generated images for the same species while the different rows represent distinct species. Notably, we observe inter-class variations among species belonging to the same class, such as differences in fish orientation and size.






§ ABLATION RESULTS
 



 §.§ Generalization to Unseen Species: Leave-three-out

As an additional ablation experiment, we conduct a leave-three-out experiment by excluding three species from different subtrees during training to test the model's ability to generalize to new species and situate them in the phylogeny. This experiment involves training the model excluding three species, Notropis blennius, Noturus gyrinus, and Lepomis humilis that belong to different subtrees as seen in <ref>. The generated images from the three subtrees after trait masking closely resemble the actual images of the 3 species, with an F1 score of 95.6 on a classifier trained to discriminate the 3 species. This experiment underscores the robustness and accuracy of Phylo-Diffusion in embedding and generating phylogenetically consistent images.





 §.§ Effect of Varying the Number of Levels in Phylo-Diffusion

To demonstrate the robustness of Phylo-Diffusion, we perform ablation experiments with varying numbers of levels in the discretization of the phylogeny tree. We show that the choice of the number of levels depends on the depth of the phylogenetic tree and the internal nodes to be studied.
We train models with {2, 4, 6,8} levels on the phylogeny shown in <Ref>. <Ref> demonstrates that the model is robust to the choice of the number of levels.





 §.§ Effect of Varying Embedding Dimensions in Phylo-Diffusion

We further evaluate the effect of varying the number of embedding dimensions used in HIER-Embed on the performance of Phylo-Diffusion. In this experiment, we trained Phylo-Diffusion by varying HIER-Embed's dimension in the following range of values: {16,32,64,128,256,512,1024}. <Ref> shows that Phylo-Diffusion is quite robust to the choice of embedding dimension with minimal drop in performance as we reduce the embedding dimension even to small values.






§ CUB DATASET RESULTS
 
To show the applicability of our approach on other datasets with larger and deeper phylogenies, we perform additional experiments on 190 bird species from the CUB-200-2011 dataset (see <Ref>). We selected the set of bird species based on whether we are able to obtain their phylogenetic knowledge from https://birdtree.org/Bird Tree, which are pre-processed similar to the fishes. We removed the background of these images using segmentation masks to focus only on the body of the birds.







  §.§.§ Trait Masking:
 Similar to fishes, <Ref> shows the changes in probability distributions when Level 3 & 4 information is replaced with noise. The first two plots show the logits of images generated for Black-footed albatross and Sooty albatross using embeddings from all the four levels. The third plot shows the dispersion of logits across the three descendant species that are part of the sub-tree defined till level 2, i.e., masking level 3 & 4. We see similar results for CUB as well where the probability of classifying the generated images into any of the descendant species that share a subtree (highlighted in green) is generally greater than the species outside the subtree. 








  §.§.§ Trait Swapping:
 
<Ref> shows an example of the trait swapping experiment on the birds dataset, similar to the experiments for fishes in the main paper. We see that the image generated from the perturbed embedding (center) picks up the trait of black coloration around the eye (purple circle) that is shared by the target sub-tree (right) while traits like pointed beak (green circle) are retained from the source species (left).







 Paper Title: Enhanced Uncertainty Estimation in Ultrasound Image Segmentation with MSU-Net.





Uncertainty Estimation with MSU-Net






R. Banerjee et al.




Carnegie Mellon University, Pittsburgh PA 15213, USA {rohinib,cgmorale,awd}@andrew.cmu.edu


Enhanced Uncertainty Estimation in Ultrasound Image Segmentation with MSU-Net
    Rohini Banerjee* Cecilia G. Morales* Artur Dubrawski
    August 5, 2024
=============================================================================



*These authors contributed equally to this workfootnote

§ ABSTRACT 
Efficient intravascular access in trauma and critical care significantly impacts patient outcomes. However, the availability of skilled medical personnel in austere environments is often limited. Autonomous robotic ultrasound systems can aid in needle insertion for medication delivery and support non-experts in such tasks. Despite advances in autonomous needle insertion, inaccuracies in vessel segmentation predictions pose risks. Understanding the uncertainty of predictive models in ultrasound imaging is crucial for assessing their reliability. We introduce MSU-Net, a novel multistage approach for training an ensemble of U-Nets to yield accurate ultrasound image segmentation maps. We demonstrate substantial improvements, 18.1% over a single Monte Carlo U-Net, enhancing uncertainty evaluations, model transparency, and trustworthiness. By highlighting areas of model certainty, MSU-Net can guide safe needle insertions, empowering non-experts to accomplish such tasks.








§ INTRODUCTION


Trauma, the leading cause of death among young individuals in the U.S. <cit.>, often results in blood loss, which requires rapid fluid resuscitation for vital organ oxygenation. In austere settings, accessing timely medical care can be challenging due to limited access, dangerous conditions, time constraints, or the absence of medical infrastructure, making expertise for optimal needle insertion sites critical. Autonomous robotic systems can assist in intravenous fluid administration when medical experts are unavailable, providing support in emergencies. These systems can also guide non-experts in accurate performance of phlebotomy tasks, empowering them to contribute effectively in dire medical situations.

Ultrasound imaging is widely used for locating vessels to enable fluid resuscitation due to its affordability, speed, safety, and portability, unlike CT or MRI imaging that are not portable and use ionizing radiation <cit.>. Despite advancements in autonomous needle insertion into blood vessels <cit.>, a critical challenge persists: inaccurate predictions can lead to severe consequences. Failing to anticipate vessel structure during intravenous cannulation may result in catastrophic outcomes for critically injured individuals; incorrectly predicting two adjacent vessels as a single one could lead to laceration of the vessel wall and cause hemorrhage upon needle insertion <cit.>. Both medical personnel and automated tools need to assess the certainty of their estimates of vessel location and structure to mitigate this risk <cit.>. 
As healthcare increasingly adopts automation, reliance on Artificial Intelligence (AI) models grows. Building a  model that communicates its uncertainty becomes essential to help users focus their attention and actions on where the model is confident, such as guiding needle insertion accurately within vessel segments. In this paper, we attack the problem of uncertainty estimation of models trained to guide needle insertion tasks. Our contributions include: 

(1) We introduce MSU-Net, a novel Multistaged Monte Carlo U-Net, demonstrating significant advancements in uncertainty quantification; (2) Our model achieves state-of-the-art performance in ultrasound image segmentation, marking the first known improvement in uncertainty estimation for this application.





§ RELATED WORKS



 §.§ Uncertainty Quantification in Deep Learning

Uncertainty quantification is vital in understanding the reliability of AI model predictions. Traditionally, the frequentist approach assumes a single point estimate of network weights, using estimated class likelihoods as confidence measures for predictions. However, studies have shown that these likelihoods often overestimate accuracy <cit.>, and the popular metric used to quantify confidence, expected calibration error, has been criticized for bias and inconsistency <cit.>. This limits utility, motivating the need for alternative approaches to accurately quantify model uncertainty rather than just confidence.

Predictive uncertainty comprises of aleatoric and epistemic components <cit.>. Aleatoric uncertainty accounts for inherent noise in observations, while epistemic uncertainty arises from limited training data and model parameter uncertainty. Recent advancements in Bayesian inference and Bayesian neural networks have provided robust frameworks to quantify both forms of uncertainty by estimating posterior distributions over model weights. Gal and Ghahramani <cit.> introduced Monte Carlo (MC) dropout for Bayesian inference in deep learning, leveraging dropout in convolutional layers for stochastic forward passes to approximate Bayesian variational inference. Bayesian approximation using MC dropout has been extensively applied: Kendall and Gal <cit.> developed Bayesian SegNet for scene understanding, while Dechesne and Lassalle <cit.> used it in U-Net for high-accuracy image segmentation. Seedat <cit.> proposed a human-in-the-loop system using model uncertainty. Yet, single-model architectures are now supplanted by model ensembles due to difficulties in capturing inherent variability.




 §.§ Multistage Neural Network Ensembles

Training multiple individual models through model ensembling offers a range of tools to encourage member model diversity and improve overall accuracy. Diversity can be achieved through several techniques, the most popular being bagging <cit.>, stacking <cit.>, and boosting <cit.>. Non-Bayesian deep ensembles improve predictive uncertainty estimation compared to single model architectures <cit.>, yet remain limited by naïve aggregation strategies such as simple or weighted averaging or majority vote <cit.>. Yang et al. <cit.> addresses drawbacks of traditional ensembles by training a secondary neural network to adaptively assign weights, inspired by stacked generalization. This approach offers flexibility in candidate selection and leverages non-linear modeling capabilities of neural networks. 


Lai et al. <cit.> designed a multistage reliability-based neural network ensemble learning approach for credit card evaluation based on a decorrelation maximization procedure to select diverse models. Their approach combines outputs using different aggregation strategies, achieving superior accuracy over single and hybrid models. Multistage ensembling with U-Nets also proves beneficial; Yin and Hu <cit.> designed Paw-Net, a multistage ensemble for semantic segmentation. Paw-Net demonstrates higher final Intersection-over-Union (IoU) scores by integrating outputs of multiple U-Nets specialized in different classes. 

Our proposed MSU-Net architecture combines the benefits of these appro­aches: candidate U-Nets are first overproduced and then selected based on decorrelation to compose an ensemble. The final stage combines outputs using a Monte Carlo U-Net approach to produce accurate segmentation maps.






§ METHODS



 §.§ Monte Carlo U-Net (MCU-Net)

We introduce stochasticity into our inference process by incorporating dropout layers into our chosen U-Net architecture, thereby enabling MC dropout, as in MCU-Net  <cit.>. MCU-Net is henceforth referred to as our baseline model. Indiscriminately placing dropout after each convolutional layer can lead to reduced fit and poor test performance <cit.>. We opt to situate dropout layers in the decoder section of our U-Net instead. Each decoder block consists of two sets of convolutional layers (3x3 convolution filter, batch normalization, ReLU activation) and an attention block. This enables us to approximate Bayesian inference by conducting T forward passes, or MC samples, of the U-Net during testing. Our empirical results show no significant improvement beyond T=30. Our model outputs both logits and logit variances <cit.> to use for segmentation and epistemic uncertainty maps,



with  [p_t, σ^2_t] =f^ω_t(x) representing the t-th forward pass of MCU-Net with learned weights ω_t. Given sigmoid activation, σ_SIG, the ensemble prediction is  aggregated by averaging individual model outputs, σ_SIG(p_t),  across all T samples. Similarly, the raw epistemic uncertainty map is obtained by averaging logit variances across all MC samples.




 §.§ Multistage U-Net (MSU-Net)

The MSU-Net architecture, illustrated in Fig <ref>, is structured in three main stages: (1) ensembling, (2) candidate selection, and (3) combining. In stage 1, we train multiple models referred to as candidates. Bagging is employed to induce diversity into our training data; given the training dataset of size n_TR, we repeatedly sample with replacement to generate M training subsets, each also of size n_TR. We then proceed to train a total of 15 bootstrapped models <cit.>. Hyperparameter optimization, specifically determining the number of epochs of training for each candidate model, is performed using the VS1 validation set, held out for this purpose. Early stopping is applied to each candidate when validation loss begins to increase, to mitigate model overfitting. 

In stage 2, candidates are selected to form a diverse and efficient ensemble, aiming to reduce computational costs by minimizing correlation. We apply the decorrelation maximization method from <cit.>, using the Brier score loss matrix computed from models f^ω_1_1, f^ω_2_2, ⋯, f^ω_M_M on a separate validation set VS2 within the specified region of interest, replacing binary error with Brier score for finer error measurement. The mean, variance, and covariance of the Brier score loss matrix is then computed in order to construct the correlation matrix R, where the ith row and jth column represents the degree of correlation between f^ω_i_i and f^ω_j_j. We write R in a block form for candidate model f^ω_i_i as


    R [  R_-i   r_i; r_i^T     1 ]


where R_-i is the principle submatrix of R resulting from deleting the ith row and ith column. Subsequently, we compute the plural-correlation coefficient <cit.>

    ρ_i^2 = r_i^T R^-1_-i r_i

by which candidate f^ω_i_i is kept if ρ_i^2 ≤θ for some threshold θ, else discarded, in order to build our final ensemble. Empirical trials testing θ thresholds indicate that K=3 performs almost equally as well as K=15, suggesting that minimal additional training will suffice for our new architecture.






















In stage 3, we train a final MCU-Net combiner taking ensemble member outputs as inputs and outputting segmentation maps. We compute the segmentation and epistemic uncertainty maps of MSU-Net analogous to MCU-Net.





§ EXPERIMENTS



Dataset and Training  
We used a ultrasound scanning system described by Morales et al. <cit.>  
to scan the CAE Blue Phantom anthropomorphic gel model simulating femoral vessels. Equipped with a 5MHz linear transducer, the system can scan up to 5cm in depth, producing 2D transverse ultrasound images. Expert clinicians annotated these images using the Computer Vision Annotating Tool (CVAT) <cit.>, followed by cropping and resizing to 256 × 256 pixels. The dataset was split into training (1392 images), validation (907 images), and testing (856 images). The validation set was further randomly split into two disjoint sets, VS1 and VS2.

For image segmentation, we employed U-Net architectures with a ResNet34 backbone in Pytorch, utilizing the Segmentation Models library <cit.> pretrained on ImageNet. To integrate MC Dropout, dropout layers with optimal rates of 0.4 and 0.5 were added after each ReLU activation  in the decoder <cit.>. Training used a batch size of 8, Adam optimizer with a learning rate of 0.0001, and early stopping based on validation loss stabilization. MSU-Net incorporated bagging for ensembling and the plural-correlation coefficient as a correlation metric.







































Distribution Divergence Estimation
Model quality hinges on uncertainty distribution: low for correct, high for incorrect predictions, enhancing calibration and user trust <cit.>. We utilize the Rényi divergence (RD) statistic as our metric for comparison. It is a generalization of Kullback-Leibler (KL) divergence that measures the dissimilarity between two probability distributions p and q. van Erven and Harremoës <cit.> relay an operational characterization of RD as the number of bits by which a mixture of two codes, p and q, can be compressed.  For some measurable set ℳ_0 that p and q lie in and order α∈ℝ∖{1 }, RD of distribution p from distribution q is defined as:


    R_α(p || q) = 1/α-1ln∫_ℳ_0 p^α(x) q^1-α(x)  dx


A nonparametric estimator of RD that is conditionally L_2-consistent using only k-nearest-neighbor statistics has been proposed to considerably reduce computational effort <cit.>. Let X_1:n_0=(X_1, ⋯, X_n_0) be an i.i.d. sample from a distribution with density p and Y_1:n_1=(Y_1, ⋯, Y_n_1)  an i.i.d. sample from a distribution with density q. We denote ρ_k(i) to be the k-th nearest neighbor of observation X_i in X_1:n_0 and v_k(i) the k-th nearest neighbor of X_i in Y_1:n_1. With B_k, α = Γ(k)^2/Γ(k-α+1) Γ(k+α-1), we can estimate the RD by:


    R_α(p || q) = 1/α-1ln( n_0^-1∑_i=1^n_0((n_0-1) ρ_k(i)/n_1 v_k(i))^1-α B_k,α)


We use RD to measure the ability of our method to distinguish correct predictions (p) from incorrect predictions (q). We conduct permutation testing to assess deterministicity and bootstrapping to obtain confidence intervals for the results. Nearest neighbor estimators are sensitive to perturbations in the underlying distribution, and hence their limited variance cannot be consistently estimated by a naïve Efron-type bootstrap <cit.>. 
Since this behavior may result in a non-negligible positive bias in bootstrap estimates, we instead apply a direct M-out-of-N (MooN) type bootstrap <cit.> for this metric, shown in Algorithm <ref>.









§ RESULTS AND DISCUSSION


We evaluate MSU-Net using quantitative and qualitative metrics. To alleviate heavy class imbalance, we choose to average metrics over a predefined region of interest (ROI), shown in  Fig. <ref>. Fig. <ref> displays model uncertainty distributions on the test set. We employ kernel density estimation with a Gaussian kernel and optimal bandwidth via Silverman’s rule of thumb to visualize difference in means between correct and incorrect prediction uncertainty distributions over 100,000 samples. To quantify separation between these distributions with RD, we select k=4 for k-nearest neighbors and α=0.85 for numerical stability and interpretability <cit.>. We perform B=1000 permutations and bootstrapped samples. For MooN-type bootstrap, γ=0.8 is selected to maintain the largest proportion of original data while achieving the closest coverage probability of 0.95 for 95% intervals. Our final results are displayed in Table <ref>.

 


Fig. <ref>a shows training behavior for both models. MSU-Net shows a more stable convergence and consistently outperforms MCU-Net during validation performed after each epoch. MSU-Net achieves 18.1% better mean IoU and a significant improvement in sensitivity and false negative rate scores at alpha level 0.05, while other metrics remain similar. We additionally utilize precision-recall curves, seen in  Fig. <ref>b, which are resilient to unbalanced classes since they only focus on positive class predictions. Performance results are displayed in Table <ref>.












Permutation tests for MCU-Net and MSU-Net (p ≤ 0.003 for both) indicate that the observed separation between correct and incorrect predictions is not random in both. Yet, at the 95% confidence level, we see no overlap between their 95% confidence intervals, showing that the ability of MSU-Net to distinguish correct from incorrect predictions is significantly better than that of MCU-Net.

Qualitative uncertainty maps visually validate our findings and capture local variations in model performance <cit.>. MCU-Net exhibits indiscriminately low uncertainty in Fig. <ref>, whereas MSU-Net provides more interpretable uncertainty values. MSU-Net shows clearer vessel tops with lower uncertainty and higher uncertanties at the bottom of vessels which is consistent with our assessment. 
MSU-Net yields a higher credibility in assessing its predictive reliability. 



MSU-Net generally outperforms MCU-Net. Although it performs marginally worse at specificity and false positive rate (FPR), the precision-recall curves from Fig. <ref>b show a higher average precision score for MSU-Net at 0.945 than MCU-Net at 0.875 compared to the baseline score of 0.09 for a `no-skill' classifier. As such, at the same level of recall, MSU-Net correctly classifies a greater proportion of pixels that are actually vessels than MCU-Net. Crucially, MSU-Net achieves a considerably lower false negative rate (FNR) than MCU-Net. In this context, failing to anticipate a true vessel can have more catastrophic consequences to a critically injured individual than incorrectly anticipating a vessel. MSU-Net improves credibility not only through higher quality results, but also through more accurate results while avoiding potentially disastrous deficiencies of predictive modeling in the context of medical image segmentation.




§ CONCLUSION

This paper introduces MSU-Net, a multistaged Monte Carlo U-Net for uncertain ultrasound image segmentations. It improves model transparency and trustworthiness compared to standard U-Net by enhancing uncertainty evaluation, despite minimal additional training. Our framework sets a benchmark for future studies, offering qualitative maps for intuitive model assessment. While our preliminary results are limited to phantom data and binary image segmentation, we aim to validate these findings on live animal and human data in future work.





  §.§.§ 
 Authors thank Nico Zevallos, Dr. Michael R. Pinsky, and Dr. Hernando Gomez for gathering experiment data. This work was partially supported by the U.S. Dept. of Defense contracts W81XWH-19-C0083 and W81XWH-19-C0101.



  §.§.§ 

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








 splncs04
 























 Paper Title: Learning Optimal Signal Temporal Logic Decision Trees for Classification: A Max-Flow MILP Formulation.





Recent update of gas-phase chemical reactions and molecular lines of TiO in Cloudy
    Gargi Shaw1,
Gary J. Ferland2,
Phillip Stancil 3,
Ryan Porter 4

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

empty
empty



§ ABSTRACT 
This paper presents a novel framework for inferring timed temporal logic properties from data. 
The dataset comprises pairs of finite-time system traces and corresponding labels, denoting whether the traces demonstrate specific desired behaviors, e.g. whether
the ship follows a safe route or not. 
Our proposed approach leverages decision-tree-based methods to infer Signal Temporal Logic classifiers using primitive formulae.
We formulate the inference process as a mixed integer linear programming optimization problem, recursively generating constraints to determine both data classification and tree structure. 
Applying a max-flow algorithm on the resultant tree transforms the problem into a global optimization challenge, leading to improved classification rates compared to prior methodologies.
Moreover, we introduce a technique to reduce the number of constraints by exploiting the symmetry inherent in STL primitives, which enhances the algorithm's time performance and interpretability.
To assess our algorithm's effectiveness and classification performance, we conduct three case studies involving two-class, multi-class, and complex formula classification scenarios. 







§ INTRODUCTION


The ever-increasing complexity of modern systems has resulted in an urgent need for sophisticated techniques that can help understand and classify temporal behaviors from time-series data. 
Machine learning (ML) <cit.> approaches have emerged as effective tools for this task, particularly in two-class classification, where the objective is to differentiate between desired and undesired system behaviors. 
However, the opaque nature of traditional ML algorithms often hinders interpretability and insight into system dynamics <cit.>.

To address this limitation, formal methods such as Signal Temporal Logic (STL) <cit.> have been gaining traction as a means of specifying temporal properties of real-valued signals. 
STL provides a structured language for expressing complex temporal and logical behaviors, offering both readability and interoperability. 
By formulating temporal properties as logical formulae, STL facilitates the inference of system behaviors from labeled time-series data <cit.>.

While early methods in this domain focused on manual parameter synthesis from predefined template formulae <cit.>, recent advancements have sought to automate this process by inferring both the structure and parameters of STL formulae directly from data <cit.>. 
The authors of <cit.> introduced a fragment of STL called inference parametric signal temporal logic (iPSTL), which enables the classification problem to be formulated as an optimization problem. 
However, this approach faces challenges such as high computational cost due to nonlinear parameter optimization routines and constructing a directed acyclic graph (DAG) based on the ordering of PSTL formulae, which may not necessarily improve classification performance.
Other recent works, such as  <cit.>, have addressed the two-class classification problem by building generative models for each class and deriving a discriminative formula that maximizes the probability of satisfaction for one model while minimizing it for the other. 
Despite the potential of these methods, they require building models of the system under analysis, which entails domain expertise and substantial amounts of data.
On the other hand, decision tree-based frameworks have emerged as promising approaches for efficiently learning STL formulae. 
These frameworks construct trees where each node encapsulates simple formulae optimized from a predefined set of primitives <cit.>. 
Lastly, the authors in  <cit.> use the different layers of neural networks to learn different aspects of STL formulae that allow a compact and efficient binary interpretable classification for level-one STL formulae, and multi-class inference <cit.>.

Despite recent advances, the formulae generated by existing algorithms often suffer from verbosity and complexity, limiting their interpretability and practicality in real-world scenarios. 
In response to these challenges, this paper introduces an innovative approach: optimal max-flow tree data classification using Signal Temporal Logic (STL) inference.
Our method builds upon the concept of tree-based structures, similar to previous works such as  <cit.>, to encode STL primitives recursively and facilitate the classification of labeled datasets. 
However, unlike previous approaches, we leverage this tree structure to formulate a max-flow optimization problem <cit.>, enabling us to simultaneously determine the tree structure and the inferred STL specification.
At each node of the tree, we encapsulate STL primitives and make branching decisions based on the classification of the input data. 
By formulating the classification task as a Mixed Integer Linear Programming (MILP) problem, our approach offers several advantages over previous methods. 
Unlike <cit.>, our approach enables global optimization and can be adapted for multi-class classification scenarios. 
Moreover, we reduce primitive redundancy, resulting in fewer constraints, which improves time performance and reduces complexity.
In contrast to <cit.>, our method considers a broader set of STL primitives and avoids the pitfalls associated with nonlinear optimization, such as getting trapped in local minima. 
By simultaneously maximizing classification accuracy and minimizing formula complexity, our approach ensures optimal performance with enhanced interpretability.
We validate the efficacy of our method through multiple case studies, including two-class and multi-class classification scenarios. 
The results demonstrate significantly improved classification performance and interpretability compared to existing approaches, thereby paving the way for more effective temporal behavior classification and analysis.









The main contributions of this work are, 

    
  * Proposing a novel decision tree-based STL inference algorithm that runs a global optimization and results in an improved classification rate as compared to the existing related approaches.
    
  * Formulating a Mixed Integer Linear Programming encoding that classifies data and generates the rules for growing the tree structure based on STL primitives followed by a global max-flow optimization approach that guarantees high-performance classification.
    
  * Generating a reduction of constraints based on the symmetry of STL primitives, resulting in more efficient and improved time performance.
    
  * Demonstrating three case studies to show the algorithm's capability to handle two-class, multi-class, and complex formulae classification and compare its performance with other approaches.










§ PRELIMINARIES AND NOTATION


Let ℝ denote the set of all real numbers, ℤ the set of integers, 𝔹 the binary set, and ℤ_≥ 0 the set of non-negative integers. 
For a set 𝒮, 2^𝒮 and |𝒮| represent its power set and cardinality.


We have α+S = {α+x| x∈ S}.
The integer interval (range) from a to b is ab. We use I=a and I̅=b.
A discrete-time signal s, with time horizon H ∈ℤ_≥ 0, is defined as a function s: 0H→ℝ^d mapping each time-step to an d-dimensional vector of real values.
The j-th component of x is given by x_j, j∈1d.






 §.§ Signal Temporal Logic


Consider a discrete-time signal s: 0H→𝕄 with values in the compact space 𝕄⊆ℝ^d. 
Introduced in <cit.>, Signal Temporal Logic (STL) is a specification language that expresses real-time properties. Its syntax is defined as follows

    ϕ ::= ⊤| h(s) ≥π|¬ϕ|ϕ_1 ϕ_2 |ϕ_1 ϕ_2 |_Iϕ|_Iϕ,

where ϕ, ϕ_1, and ϕ_2 are STL formulae,
⊤ is the logical True value,
h(s) ≥π is a predicate with h: ℝ^d →ℝ and threshold value π∈ℝ,

¬, , and  are the Boolean negation, conjunction, and disjunction operators.


Discrete-time temporal operators eventually _I, and always _I with I = II̅
a discrete-time interval, I̅≥I≥ 0,
are defined in the usual way <cit.>.
Linear predicates of the form s ∼π, ∼ ∈{>, ≤, <}, follow via
negation and sign change. The logical False value is = ¬⊤.

The (qualitative) semantics of STL formulae over signals s at time k is recursively defined in <cit.> as

    (s, k)  (h(s) ≥π) ≡    h(s(k)) ≥π,
    
        (s, k) ¬ϕ≡    (s, k) ϕ,
    
        (s, k) ϕ_1 ϕ_2 ≡   ( (s, k)ϕ_1 ) ( (s, k)ϕ_2 ),
    
        (s, k) ϕ_1 ϕ_2 ≡   ( (s, k)ϕ_1 ) ( (s, k)ϕ_2 ),
    
        (s, k) _I ϕ≡   ∃ k' ∈ k + I  s.t.  (s, k') ϕ, 
    
        (s, k) _I ϕ≡   ∀ k' ∈ k + I  s.t.  (s, k') ϕ,

where  and  denote satisfaction and violation, respectively.
A signal s satisfying ϕ, denoted as sϕ, is true if (s, 0) ϕ.

In addition to Boolean semantics, STL admits quantitative semantics,
called robustness, that indicates how much a signal satisfies
or violates a specification <cit.>.
The robustness score ρ(s, ϕ, k) is recursively defined as

    ρ(s, ⊤, k)     = ρ_⊤,
    ρ(s, h(s) ≥π, k)    = h(s(k)) - π,
    ρ(s, ¬ϕ, k)     = -ρ(s, ϕ, k),
    ρ(s, ϕ_1 ϕ_2, k)    = min(ρ(s, ϕ_1, k), ρ(s, ϕ_2, k)),
    ρ(s, ϕ_1 ϕ_2, k)    = max(ρ(s, ϕ_1, k), ρ(s, ϕ_2, k)),
    ρ(s,_I ϕ, k)    = min_k' ∈ k+Iρ(s, ϕ, k'),
    ρ(s,_I ϕ, k)    = max_k' ∈ k+Iρ(s, ϕ, k'),

where ρ_⊤ = sup_s, π{|h(s) - π|} is the maximum robustness.



Let s be a signal and ϕ an STL formula.
It holds ρ(s, ϕ, k) > 0 ⇒ (s, k) ϕ for satisfaction
and ρ(s, ϕ, k) < 0 ⇒ (s, k) ϕ for violation.


The time horizon of an STL formula <cit.> is defined as

    ϕ =
    
    0,    ϕ∈{⊤, h(s) ≥π}, 
    ϕ_1,    ϕ = ¬ϕ_1,
    max{ϕ_1, ϕ_2},    ϕ∈{ϕ_1 ϕ_2, ϕ_1 ϕ_2 },
    I̅ + ϕ_1,    ϕ∈{_I ϕ_1, _Iϕ_1 }.




§ PROBLEM FORMULATION



This section introduces the data classification problem by inferring STL specifications.
Let us consider a labeled dataset := {(s^i, ℓ^i)}_i ∈ℐ, where s^i and ℓ^i represent the i-th signal and its corresponding label. 
We denote the set of all possible classification classes as =1||, where ℓ^i=c signifies that the i-th sample is labeled with class c ∈. 
Our aim is to infer STL formulae ϕ_c that encapsulate properties inherent to each data class c ∈. 
The formulae ϕ_c enable the classification of signals not present in the dataset .
Formally, we consider the following multi-class classification problem.
 





Given a labeled data set , find a set of mutually-exclusive STL formulae Φ = {ϕ_c}_c∈ that maximizes the correct classification rate CCR(Φ) where 

    CCR(Φ)=∑_c ∈|{s^i | s^i ϕ_c ∧ℓ^i=c}|/|ℐ|.

such that ϕ_c ϕ_c'≡ for all c, c' ∈ with c≠ c'.








    


§ SOLUTION



In this section, we propose an approach to learning Signal Temporal Logic formulae in the form of decision trees. 
The tree branches from the root to the leaves correspond to the classification process, where leaves are associated with classes and intermediate nodes with decisions in the form of primitive STL formulae.
We infer the structure of the decision tree, the primitives to use for each decision node and their spatial and temporal parameters, and the classes associated with the leaf nodes such that the correct classification rate is maximized and the number of decision nodes is minimized.
We cast the temporal inference problem as a Mixed Integer Linear Programming (MILP) problem using a max flow encoding <cit.>.



Our classification method employs a tree-based structure <cit.>, akin to a binary decision tree but with a unique feature: it incorporates a source node 𝐬 and a sink node 𝐭, as illustrated in Fig. <ref>. 
Furthermore, each node is connected to the sink node. 
We refer to this configuration as a classification tree, defined as follows.


    
    A classification tree 𝒯=(𝒩,ℒ, ℰ, 𝐬, 𝐭) is a Directed Acyclic Graph (DAG) with single source 𝐬 and single sink 𝐭. 
    Nodes in the final layer of the tree are referred to as leaf nodes, denoted by ℒ, while nodes situated between the source and the leaf nodes are internal nodes, denoted by 𝒩 and ℰ⊆(𝒩∪{𝐬})×(𝒩∪ℒ∪{𝐭}) captures relation between nodes.
    We define the sets n = {n' | (n', n) ∈ℰ} and n = {n' | (n, n') ∈ℰ} to represent the parent and children nodes of a given node n, respectively. 


Note that in Fig. <ref>, 𝐬 = ∅, since sink has no parent node and 1 = 𝐬[We abuse notation and write n= n' instead of n= {n'} for readability since a node has at most one parent.].
On the other hand, 𝐭=∅ and for all n∈ℒ we have n= {𝐭}.
An internal node n ∈𝒩 has three children, a left child, a right child, and the sink 𝐭.
In the classification process using 𝒯, data samples that are correctly classified traverse from the source node through the tree to reach the sink. 
Conversely, misclassified samples are blocked from progressing beyond the source node.
The binary decision tree (BDT) with STL primitive nodes is extracted from the classification tree 𝒯.
The root of the STL BDT is the node connected to the source 𝐬, i.e., node 1 in Fig. <ref>.
Not all nodes 𝒩∪ℒ of 𝒯 are part of the inferred STL BDT.
The final structure follows from the optimization result. 
Note that a pre-determined depth value limits the size of the classification tree. 




















To partition the data at each node n∈𝒩, a finite list of possible splitting rules is considered <cit.>. 
We use simple Parametric Signal Temporal Logic (PSTL) formulae <cit.>, called primitives <cit.>, 

defined as follows.



The primitive in the set Ψ consists of a temporal operator, a relational operator, and a predicate function.
Formally, 
Ψ = {Γ (h(s) ∼π) |Γ∈{_I_1, _I_1, _I_1_I_2, ⋯},
∼∈{≥ , <}, h ∈ℋ},
where π∈ℝ is the threshold,
I_j are integer ranges of temporal operators,
and ℋ is a finite set of predicate functions
h:0H→𝕄.
The set of all possible time intervals of temporal operators is
Θ={θ = (I_1,…,I_q) | I_p = I_pI̅_p|I_p ≤I̅_p,  with I, I̅∈0H, ∀ p ∈1q, q∈ℤ_≥ 1},
where q is the number of temporal operators in a primitive.
    
    


In Def. <ref>, π is the spatial parameter,
while the time bounds of temporal operators I_p are the time parameters.

When using a tree 𝒯 for classification, we assign a primitive ψ∈Ψ or a class c ∈ to each node n in 𝒩 and ℒ. Note that a class c ∈ can be assigned at internal or leaf nodes, whereas primitives can only be assigned to internal nodes.
For internal nodes n ∈𝒩 in addition, we also compute the spatial and time parameters of primitive ψ, denoted as π_n and θ_n = (I_1,…,I_q_ψ), where q_ψ is the number of temporal operators in ψ.

For each data sample s^i in the dataset , the following encoding
captures the classification of s^i as a flow from the source 𝐬
to the sink 𝐭.
At the starting node 𝐬, s^i flows down to the top node 1, see Fig. <ref>.
If the node n is a decision node associate with the STL formula ϕ_n = ψ_n(π_n, θ_ψ_n),
then the check s^i ϕ_n is performed.
In case the signal s^i satisfies ϕ_n, the signal flows to the left child
l(n) ∈n. Otherwise, it flows to the right child r(n) ∈n.
The link to the sink node 𝐭 is not used for decision nodes.
The procedure continues from the node reached by s^i.
If the node n is a classification node associated with class c_n,
the check c_n = ℓ^i is performed.
In case the label ℓ^i matches the class c_n, then the signal s^i flows
to the sink 𝐭 and is correctly classified and counted towards the objective.

If a signal cannot be correctly classified, it does not enter the classification tree and, thus, does not count toward the objective as explained later in Sec. <ref>.
In either case, the classification of s^i ends when it reaches the sink 𝐭.
For classification nodes, the links to the right and left nodes are unused.
Moreover, leaves ℒ can only be classification nodes since
they have edges only to the sink 𝐭.






Hence, to achieve accurate data classification, we aim to construct a classification tree 𝒯 that maximizes the Correct Classification Rate (CCR(Φ)). 
Consequently, the flow-based optimization problem solving Pb. <ref> is:








    max      CCR(Φ) 
    s.t.    Flow constraints,
       Node function constraints,
       STL satisfaction constraints.


The objective is to maximize the Correct Classification Rate (CCR) while adhering to flow constraints that ensure the conservation of data flow within the tree.
Constraints related to node functions define primitive predicates, child allocation, and classification. 
The final soundness constraint ensures that the predicted class satisfies the signal. 
The following sections provide a detailed description of these constraints.











 §.§ Flow constraints
























Here, we establish the conservation of data flow within the classification tree 𝒯.
Let z^i_n ∈𝔹 denote a binary variable representing whether data point (s^i,ℓ^i) ∈ traverses node n ∈𝒩∪ℒ. 
It takes a value of one if data point (s^i,ℓ^i) flows through node n and zero otherwise.
We designate the entry of a data point into the source node as z_𝐬^i ∈𝔹. 
Since all internal nodes are connected to a common sink node 𝐭, we introduce z_𝐭,n^i ∈𝔹 which also identifies the node through which data point (s^i,ℓ^i) enters the sink node 𝐭. 
Specifically, z_t,n^i = 1 indicates that data (s^i,ℓ^i) reaches the sink node via node n, and zero otherwise.

Then, the flow constraints are defined as follows



    z_𝐬^i    = z_1^i,   ∀ i∈ℐ,
    
    z_n^i    = z_𝐭, n^i + z_l(n)^i + z_r(n)^i,   ∀ n∈𝒩, ∀ i∈ℐ,
    
    z_n^i    = z_𝐭,n^i,   ∀ n∈ℒ, ∀ i∈ℐ.


The constraint (<ref>) enforces a flow from the source node 𝐬
to enter node n=1 corresponding to the root of the inferred BDT that eventually reaches sink 𝐭 and is correctly classified.
In this case, z_𝐬^i = z_1^i =1.
Otherwise, we have z_𝐬^i = z_1^i = 0, and there is no flow for s^i, and it is thus misclassified.

Next, (<ref>) ensures that data s^i leaving internal node n∈𝒩∪ℒ must enter one of its children nodes or the sink node 𝐭.
The last equation (<ref>) enforces the flow for s^i from leaves ℒ to the sink 𝐭, which is their only child.
These constraints imply that data is either correctly classified, i.e., z_𝐭,n^i = 1 for a node n ∈𝒩∪ℒ via a deterministic tree flow, or it does not enter the tree, i.e., z_𝐬^i = 0.



 §.§ Node function constraints



























In this section, we capture the functionality of the nodes in 𝒯.
Internal nodes 𝒩 are either decision or classification nodes.
Each internal node n ∈𝒩 in 𝒯 either checks a primitive STL formula and creates child nodes or classifies data by directing the flow to the sink.
Leaf nodes n∈ℒ exclusively make classifications by enforcing the flow to the sink 𝐭.
For each node n∈𝒩,
we introduce the binary decision variables b_n^ψ∈𝔹
to capture whether n is a decision node associated with primitive STL function ψ∈Ψ.
Conversely, binary decision variables w_n^c ∈𝔹 indicate that
node n is a classification node for class c ∈.




The constraints governing the node functionality are:



    ∑_ψ∈Ψb^ψ_n+ ∑_c∈𝒞 w_n^c    = 1,   ∀ n∈𝒩, 
    ∑_c∈𝒞 w_n^c    = 1,   ∀ n ∈ℒ.



For internal nodes n∈𝒩, (<ref>) guarantees
that n either performs a decision using a primitive ψ or classifies the data into a class c.
For leaves n∈ℒ, (<ref>) enforces classification, i.e., assignment of a class c.

When imposing a decision using a primitive, time intervals and thresholds are required to complete the STL formula. 

The time parameters variable ξ_θ^n∈𝔹, ∀θ∈Θ,
encompasses all possible valuations of the time parameters within the horizon H. 
The constraint capturing the primitive variable is 

    ∑_ψ∈Ψb^ψ_n = ∑_θ∈Θξ_θ^n,   ∀ n∈𝒩.

If a primitive is selected, i.e., ∑_ψ∈Ψb^ψ_n=1,
then (<ref>) ensures that exactly one set of time parameters θ is selected for the primitive ψ.
Otherwise, no θ is selected.

For classification nodes responsible, data flows to the sink node 𝐭 if the predicted class is correct.
The following constraint enforces this behavior

    z_𝐭,n^i ≤ w_n^c,   ∀ i: ℓ^i =  c ∈𝒞,   n∈𝒩∪ℒ.





 §.§ STL satisfaction constraints


In this section, we connect the decision variables for decision nodes with the satisfaction of primitive formulas.
Additionally, the constraints capture the robustness of signals with respect to the STL formula obtained from a PSTL primitive and computed spatial and time parameters.







  §.§.§ Robustness calculation

To capture the robustness of signals with respect to primitives, we use the following result.



Let s be a signal and ψ∈Ψ be a PSTL primitive.
For any valuations θ and π of the spatial and time parameters of ψ, 
we have

    ρ(s, ψ(π, θ)) = ρ(s, ψ(0, θ)) - π



The proof follows trivially from the structure of primitive formulae in Φ,
and is omitted for brevity.


Using Prop. <ref>, we can precompute the robustness values
for all signals s^i, primitives ψ and time parameters θ,
which we denote as τ_θ^i, ψ.

For example, the robustness of signal s^1 with respect to
ψ_1 = _15 h(s) ≥ 0
is τ_15^1, ψ_1 = max_k∈15 h(s^1(k)).



The robustness ρ_ψ^i, n∈ℝ of data s^i with respect to the formula ψ
with an arbitrary spatial parameter π_n ∈ℝ is captured by

    ∑_θ∈Θξ_θ^nτ_θ^i, ψ + π_n = ρ^i, n_ψ,   ∀ i∈ℐ,  ψ∈Ψ,   n∈𝒩.

We enforce constraint (<ref>) for all internal nodes n∈𝒩
to compute the robustness of every signal with respect to the predicted formula at node n.




  §.§.§ Soundness constraints












If a signal satisfies the primitive formula at a decision node n∈𝒩, it is directed to the left child node l(n).
Otherwise, it proceeds to the right child node r(n).

To represent the sign of robustness, we introduce the binary variables y^i, n_ψ∈𝔹. 
The variable y^i, n_ψ takes the value of one if ρ^i_ϕ≥ 0, and zero if ρ^i_ϕ≤ 0. 
We employ the big-M method to encode this relationship such that the robustness value is not over-constrained when y^i_ϕ=0 as follows


    ρ^i, n_ψ   ≤ M y^i, n_ψ,   ∀ i∈ℐ, ψ∈Ψ,n∈𝒩,
    
         -ρ^i, n_ψ   ≤ M (1 - y^i, n_ψ),   ∀ i∈ℐ, ψ∈Ψ,n∈𝒩.


where M = ρ_⊤∈ℝ_>0 is the constant value of the maximum robustness.





















Next, we need to connect the satisfaction of primitives at a decision node n
with the flow to the left or right child.
Hence, for indicating data flowing to the left child node, we introduce variable z_l(n)^i ∈𝔹 which takes the value of one if two conditions are satisfied simultaneously 
(1) n is a decision node with primitive ψ, i.e., b^ψ_n = 1, and 
(2) the robustness is positive with respect to the formula y^i, n_ψ = 1.
The conditions are enforced by


    z^i_l(n)≤∑_ψ∈Ψκ_ψ^i,n,  ∀ i∈ℐ, ∀ n∈𝒩,


where κ_ψ^i,n= y^i,n_ψ· b^ψ_n is an auxiliary variable that captures the condition
z_l(n)^i=1 if and only if y^i, n_ψ=1 and b^ψ_n=1. 
Note that since both variables are binary, the product is equivalent to κ_ψ^i,n= min{y^i,n_ψ,b^ψ_n}
and is encoded as mixed integer linear constraints.


Data flowing to the right z^i_r(n)∈𝔹 follows a similar process but with opposite conditions as follows:


    z^i_r(n)   ≤ 1-∑_ψ∈Ψκ_ψ^i,n,  ∀ i∈ℐ,  n∈𝒩.
    
            z^i_r(n)   ≤∑_ψ∈Ψb^ψ_n.




 §.§ MILP formulation


The CCR(Φ) of the data set, which is our main objective in the optimization problem, is given by the flow into sink 𝐭.
Thus, we obtain the objective by summing up z_𝐭,n^i as follows:


    CCR(ϕ)=∑_i∈ℐ∑_n∈𝒩∪ℒz_𝐭,n^i/|ℐ|.


Moreover, the optimal flow tree structure enables specifying a penalty
for the number of decision nodes, thereby reducing the complexity of the resulting formula. 
However, this approach may involve a trade-off in performance, as more decisions can lead to better classification, i.e., larger CCR. 
The trade-off is captured by a variable λ∈ [0,1] that serves as a regularization parameter.
This parameter allows users to adjust the emphasis between complexity reduction and maintaining high performance. 
Consequently, the objective function in (<ref>) is modified as follows:

    𝒥=(1-λ)∑_i∈ℐ∑_n𝒩∪ℒz_𝐭,n^i + λ∑_ψ∈Ψb^ψ_n.


The first part of the objective function evaluates the number of correctly classified data, while the second part quantifies the number of decision nodes with λ serving as the adjustment variable.

Finally, problem in (<ref>) becomes the MILP problem:






    max    𝒥
    
      s.t.    (<ref>)  (Flow constraints),
       (<ref>)-(<ref>)  (Node function constraints),
       (<ref>)-(<ref>)  (STL satisfaction constraints).





 §.§ Problem simplification









In this section, we present a simplified version of the problem in (<ref>)
when considering first-level STL primitive formulae, i.e., formulae without nested
temporal operators <cit.>. 







Under these two conditions, the following primitives are equivalent:

    ρ(s, _θ h(s)    ≥π) = - ρ(s, _θ h(s) ≤π) 
    ρ(s, _θ h(s)    ≤π) = -ρ(s, _θ h(s) ≥π)
        ,



From the definition, we have
    
    ρ(_θ h(s)≥π)    = max_t ∈θ{h(s(t)) - π}
       =- (- max_t ∈θ{h(s(t)) - π})
       =- (min_t∈θ{π - h(s(t)) })
       = -ρ(_θh(s) ≤π)
            .

The proof for the second equivalence follows similarly.


Inferring a primitive _θ h(s) ≤π from data is the
same as learning _θ h(s) ≥π.
We remove the primitives containing the always operator and, thus,
reduce Ψ by half.
The simplification can be employed with any set Φ closed under negation.
For example, second level primitives of the form _θ_1_θ_2 and _θ_1_θ_2.






 §.§ STL formula summary

Once the optimization in (<ref>) is completed, we can trace all STL formulae starting from the first layer for each data class. 
Beginning with the first layer, the traversal to a left child from any node signifies the satisfaction of the STL formula associated with that node, whereas progressing to a right child indicates the negation of the STL formula (). 
As we proceed through successive layers, the STL formulae encountered along the path are unified using the logical conjunction (). 
Upon reaching a class prediction node, the aggregated conjunction formula essentially defines the class for that node's prediction. 
Lastly, in cases where a class is represented in multiple nodes across the tree, the distinct STL formulae corresponding to each node are combined using the logical disjunction ().




§ CASE STUDIES


We demonstrate the effectiveness of our method in three case studies. 
The first is the two-class naval surveillance dataset with comparisons to other binary STL classification methods. 
The second case study demonstrates the application of the proposed approach to multi-class STL inference using a four-class trace dataset. 
Finally, the third case study is a designed dataset 

that uses second-level STL primitives for classification.  
The adjustment variable λ is set to 0 to maximize CCR.

All computations of the following case studies were performed on a PC with 20 cores at 3.7 GHz with 32 GB of RAM.
We use Gurobi <cit.> as the MILP solver. 





 §.§ Naval Surveillance

In this case study, we use the naval surveillance dataset from <cit.>. 
The different classes represent two behaviors of the vessel trajectories, as shown in Fig. <ref>. 
All trajectories start from the open sea on the right side of the figure. 
The blue trajectories are normal behaviors that directly head toward the harbor. 
The green and red trajectories are abnormal. 
The red trajectories approach the island and return to the open sea, while the green trajectories veer to the island and then head to the harbor. 
All trajectories contain 61 time steps of x and y positions.





We use the optimal flow tree with a depth of 2 for this dataset and compare other methods from <cit.>. 
All results are computed with an average of 10 runs. 

The classification tree and the resulting decision tree are shown in Fig. <ref> and Fig. <ref>, respectively. 
The colored nodes indicate nodes utilized for classification. 

A typical example of the learned formula from our method for the normal and abnormal trajectories is:

    node1: ϕ_1: _[13,59] y ≤ 23.168 
       node2: abnormal
       node3: ϕ_3:  _[50,59] x ≥ 25.483 
       node6: abnormal    node7: normal
       STL formula for normal behavior c_p: ϕ_1 ϕ_3 
       STL formula for abnormal behavior c_n: ϕ_1ϕ_3

The interpretation of this formula is clear: the y coordinates of normal vessels do not reach the island between 13 and 59 seconds, and the x coordinates of normal vessels eventually enter the harbor between 50 and 59 seconds, while abnormal trajectories are the opposite. 

We compare the results of <cit.> using boosted concise decision trees (BCDTs) with respect to the average runtime and the classification accuracy in Fig. <ref> by varying the number of samples. 
In the BCDTs approach, the variable k is the different number of decision trees used for classification, and d is the depth of the tree. 
The results demonstrate that our approach consistently outperforms the BCDTs approach in obtaining optimal solutions.






 §.§ Trace Dataset

The trace dataset contains four classes of one-dimensional transient behavior from a nuclear industry process <cit.>. 
The length of the data is 275, a total of 200 samples in the dataset, evenly distributed by classes as shown in Fig. <ref>. 

We used a tree with a depth of 3 for this dataset, and the objective achieves 100 percent accuracy and the resulting formulae for the four classes are computed as follows:

    node1:ϕ_1: _[173,181] s≤ 0.74   node2:ϕ_2: _[12,16] s≥ 0
       node3:ϕ_3: _[115,167]s≤ 0.59  node4:ϕ_4: _[21,27]s<0.72
       node5:ϕ_5: _[157,216]s≤ 0.35  node6:ϕ_6: _[230,234]s≥0.88
       node7:ϕ_7:  _[210,210]s<0.69  node8: class 2
       node 9: class 1  node 10 &13 &15: class 3
       node 11 &12 &14: class 4


    class1:    ϕ_1ϕ_2 ϕ_4
    class2:    ϕ_1ϕ_2 ϕ_4
    class3:    
            (ϕ_1ϕ_2 ϕ_5)
             (ϕ_1ϕ_3 ϕ_6)
        (ϕ_1ϕ_3ϕ_7)
    class4:    
            (ϕ_1ϕ_2 ϕ_5)
             (ϕ_1ϕ_3 ϕ_6)
        (ϕ_1ϕ_3ϕ_7)

The results demonstrate the capability for multi-class classification.





















 §.§ Classification using Second-level STL Primitives

Our framework accommodates complex STL formulae. Fig. <ref> is an illustrative example of sample trajectories designed to highlight the expressivity of second-level STL formula.
Current neural network-based approaches are limited to first-level formulae, and adapting them to other primitive sets is non-trivial <cit.>.

Class 1 of the blue trajectory is a triangle wave, and we complement the data with variations of plateaued triangular waves, such as the yellow and green trajectories with plateaus of 2 and 3, respectively. 
Additionally, to ensure that the inferred formula is not trivially defined by one tree level, we have trajectories of constant values at the lower and upper bounds drawn in purple.
Next, we sample different trajectories from each type of wave by shifting the initiation point.
We generate 300 samples with 15 time steps, 100 of which are class 1, and we denote the rest of the data as class 2. 
Lastly, we add a small random noise between -0.01 to 0.01 to all values.


We use a tree with a depth of 2 for classification. The optimization returns 100 percent accuracy, and the results are the following:

    node 1: ϕ_1: _[2, 11]_[0,4] s ≤ 1.0086 
       node 2: class 2
       node 3: ϕ_3: _[0, 10]_[0,2]s ≤ 3.017 
       node 6: class 1    node 7: class 2











ϕ_1 means within the time interval [2,11], there is always a time within any subinterval of length 4 where the signal value is less than or equal to 1.0086. 
Therefore, only the lower constant trajectory satisfies ϕ_1. 
Next, ϕ_3 means that within the time interval [0, 10], there is always an event within any subinterval of length 2 where the signal value is less than or equal to 3.017. 
From the remaining data that does not satisfy ϕ_1, only the class 1 triangular wave satisfies ϕ_3. As a result, the STL formula for class 1 is ϕ_1 ϕ_3. 






§ CONCLUSIONS

In this paper, we propose a decision tree-based STL inference algorithm that works for both two-class and multi-class classification problems. 
Our algorithm formulates the problem as a MILP by leveraging a max-flow approach over a synthesized tree using STL primitives, which allows for the computation of a globally optimal solution, i.e., a high Correct Classification Rate (CCR).
Additionally, we exploit the symmetry of STL primitives to reduce the required constraints for classification, resulting in improved performance and reduced complexity.
We evaluate the performance through three case studies to demonstrate the ability to interpret complex STL formulas and the precision in predictive accuracy.
We show the capacity of our approach to handling two-class and multi-class classification problems with first-level STL primitives and a case for second-level STL primitives.




ieeetr



 Paper Title: Enhancing Multistep Prediction of Multivariate Market Indices Using Weighted Optical Reservoir Computing.
















































label1]Fang Wang
[label1]organization=Department of Physics, Stevens Institute of Technology,
            city=Hoboken,
            postcode=07030, 
            state=New Jersey,
            country=USA

label2]Ting Bu
[label2]organization=Quantum Computing Inc.,
            city=Hoboken,
            postcode=07030, 
            state=New Jersey,
            country=USA

label1,label2]Yuping Huang
yhuang5@stevens.edu




§ ABSTRACT 


We propose and experimentally demonstrate an innovative stock index prediction method using a weighted optical reservoir computing system. We construct fundamental market data combined with macroeconomic data and technical indicators to capture the broader behavior of the stock market. Our approach shows significant higher performance than state-of-the-art methods such as linear regression, decision trees, and neural network architectures including long short-term memory. It captures well the market's high volatility and nonlinear behaviors despite limited data, demonstrating great potential for real-time, parallel, multi-dimensional data processing and predictions.
















Optical reservoir computing Parallel data processing Stock index prediction LSTM Regression Multivariate Forecasting





















§ INTRODUCTION



Over the past decades, machine learning (ML) and deep learning have emerged as a transformative tool in the field of market prediction, particularly for predicting market indices <cit.>. 

The primary challenge to generate accurate prediction is aggregating diverse information for the benchmarks and constructing a reliable model. Techniques including regression and classification, as well as more advanced methods like deep learning and ensemble models, offer sophisticated methods for analyzing data to forecast stock and market trends, <cit.>. Especially, long short-term memory (LSTM) networks, a specialized form of recurrent neural networks (RNNs), have gained significant attention in the realm of stock prediction due to their ability to model temporal dependencies and sequence data effectively <cit.>. 

Together, ML models and LSTM can incorporate diverse data sources, including historical prices, technical indicators <cit.>, and macroeconomic indicators to generate comprehensive and nuanced forecasts.

However, existing ML research on market predictions has suffered several limitations. One significant drawback is that most methods require a large amount of dense historical data for training. This requirement can be resource-intensive and may not always be feasible. Additionally, these models often struggle to accurately capture market peaks and volatility, leading to less reliable forecasts during periods of rapid market changes. As such, these limitations hinder the effectiveness of stock prediction models in providing timely and precise insights, particularly under volatile market conditions.

Recently, reservoir computing (RC) has gained prominence as a streamlined yet effective neural network architecture. It has been demonstrated to accurately predict time-series data with a limited number of training samples <cit.>. 

This architecture holds significant potential as an alternative for traditional, complex neural network structures in processing financial datasets, particularly for real-time data processing when storage capacity for large datasets is limited and/or time-to-prediction is sensitive. An RC system comprises a reservoir that maps inputs into a high-dimensional space and a readout that analyzes patterns from these high-dimensional states <cit.>. In this system, the reservoir architecture remains fixed, and only the readout is trained using simple methods such as linear regression and classification <cit.>. Therefore, reservoir computing offers fast learning and low training costs compared with other RNNs. Additionally, the static nature of the reservoir facilitates hardware implementation using a variety of optical systems, substrates, and devices. As such, optical reservoir computing has garnered increasing attention across diverse research fields<cit.>. These RC systems can be categorized into three main types: free-space optical RC, fiber-based temporal RC, and integrated photonic RC. In comparison, the free-space optical RC offers significant advantages in large parallelism and system optimization flexibility. 



In this paper, we introduce and demonstrate a free-space optical RC system for predicting daily stock index prices. Utilizing a spatial light modulator (SLM), the temporal data from benchmarks is transformed into high-dimensional optical signals in both spatial and temporal domains for high-quality benchmark predictions. Owing to an innovative design of the mapping method and the large number of pixels on the modulator, the system is flexible to implement the training samples with different feature sizes simultaneously. Therefore, the information can be processed in parallel in the optical domain without compromising the time for signals' interconnection or nonlinear effect as in the temporal RC <cit.>. Furthermore, this large-scale parallelism allows multiple-step prediction, by translating temporal correlation in the time-series data to spatial coherence in the free-space optics, and conveniently encoding and reading out in the spatial domain. 

We select seven popular stock indices and gather 500 historical data points for each. Additionally, we incorporate seven macroeconomic and technical analysis features as input variables for the model training. We further optimize the optical RC system by incorporating feature correlation as weights. Our comparative analyses with popular ML pricing models indicate that our system outperforms most ML models for both one-step and multi-step predictions. Both ML and LSTM models serve as significant reference points in experimental financial research and act as essential benchmarks when compared with our optical RC experiment. Unlike previous research requiring extensive historical data, our system efficiently achieves parallel real-time predictions using no more than 500 data points with multi-dimensional features. Notably, the inclusion of feature correlation weights significantly enhances prediction performance by capturing signals and peaks instantly. The feature-weighted optimized RC demonstrates substantial potential for confidential, low-latency edge processing and practical pricing strategies.



§ MODEL




 §.§ Experiment diagram

Figure <ref> illustrates the experimental procedures described in this paper. Initially, we download historical trading datasets from stock benchmarks and perform basic data and feature engineering. Subsequently, we apply optical RC and various popular ML techniques to predict future stock index prices, respectively. We compare different machine learning models, utilizing cross-validation and hyperparameter tuning to select the ML model with the lowest error for each benchmark. Finally, we evaluate our models and compare the predictions from optical RC with the best-performing ML model for each benchmark.




§ DATASETS PREPARATION

This study examines seven stock indices frequently analyzed in investment strategies and portfolio management, as referenced in <cit.>. The primary datasets consist of weekday historical data from May 2023 to September 2023, sourced from Yahoo Finance. These datasets include crucial pricing details, specifically the daily closing prices. Comprehensive ticker symbols and index information are provided in Table I. The closing prices of all the stock indices are inherently noisy, non-linear, and chaotic.














Due to the complexity of the benchmarks, a prediction model based solely on the closing price cannot gather all the necessary information for accurate predictions. Consequently, feature selection is a crucial part of the prediction enhancement. However, it is challenging to extract features from financial data. On the one hand, limiting the features can hinder the predictive model's performance. On the other hand, including all available features from the financial market can lead to a dimensionally complex and difficult-to-interpret model. Additionally, collinearity among multiple variables can negatively impact the model's performance. Some studies select the features rely solely on technical indicators <cit.>, while others incorporate historical data and macroeconomic indicators <cit.>. 







In the above considerations, we select seven main features either from macroeconomic indicators or technical indicators as the prediction inputs, as listed in Table <ref>. These features are chosen based on their high correlation with the closing price <cit.>. As macroeconomic data can significantly impact the stock market, we consider the Cboe Volatility Index (VIX), the Interest Rate (EFFR), the Consumer Sentiment Index (UMCSENT), and the US Dollar Index (DXYNYB). The technical indicators include Moving Average Convergence Divergence (MACD), Average True Range (ATR), and Relative Strength Index (RSI). The details of the features are listed in Table <ref>. The correlation matrix among features is illustrated in Figure <ref>(b). These values are depicted based on the intensity of the color, indicating the strength of the relationship between the given variables. In the following RC prediction, we add the bottom row of the correlation weight (CW) matrix as extra input weights W_cor. For each index, we used 500 days as the training data and 100 days as the testing data. Each day, the training input is the previous day's closing price and the seven selected features. 



 §.§ Optical RC

Our spatial optical RC follows the basic working principle of conventional RC, which can be described by 

    x(t)=f[α x(t-1)W_res+β u(t)W_in],

where x(t) represents the reservoir state at time t. f is the nonlinear function. α and β are the feedback gain and input gain, which will be optimized for the best performance. W_res and W_in are the reservoir nodes' interconnection and input weight, respectively. 
Due to this updating rule of reservoir states, the current state x(t) inherently contains information from the preceding state x(t-1). This endows reservoir computers with the crucial property of short-term memory, a feature for effective time-series data prediction.
In the training stage, the input is injected sequentially into the reservoir, and the resulting reservoir states x for all time instances within the range [1, T] are recorded. Subsequently, the system's output weights (W_out) can be determined by minimizing the error between the predicted results (ŷ=W_outx) and their corresponding true time series (y). The prediction performance for spatial RC and ML models are evaluated using the normalized root mean square error (NRMSE) between the true time series y and the prediction results ŷ on the testing dataset

    NRMSE(y, ŷ) = √(∑_i=0^N - 1 (y_i - ŷ_i)^2)/√(∑_i=0^N - 1 (y_i - y̅)^2),

where y is the mean value of the true time series.





In our optical setup, a continuous-wave 775 nm laser diode serves as the light source, as illustrated in Fig. <ref>. The beam waist is adjusted to 1 mm after collimation. To match the effective area on the spatial light modulator (SLM), a 10× beam expander is used to enlarge the beam size. A power meter is implemented to capture the transmitted light passing through the polarizing beam splitter (PBS) for real-time monitoring of system stability during training and testing.
Before the beam splitter (BS), we employ a half-wave plate to rotate the light's polarization, thereby maximizing the modulation effect on the SLM. Following modulation, the light is focused, and a camera is positioned in the time-fourier domain <cit.> to align the sensor area with the beam size for subsequent information processing. 

The inset in Fig.<ref> illustrates the data flow diagram in our spatial RC. One key here is the combination of input values and the feedback state for spatial light modulation. Here, we use a similar mapping method as the "linear-combination mapping" in our previous paper <cit.>. The difference is with how the overall mapping area is split into several input portals. Instead of equally dividing the SLM into n sections for n input values, we repeat and stack the inputs until all the modulation areas on the SLM are filled. In the case where n=4, as shown in Fig.<ref>, the 4 input values are repeated and stacked sequentially, each with a height and length of 20 SLM pixels. The mapped input is multiplied by an input scale (β) and a random input weight, while the captured camera state is rescaled with a feedback gain (α). The combination is then linearly rescaled again to match with the 0-2π phase modulation on SLM. The nonlinearity of the system is achieved through the camera intensity reading and the intensity saturation rate. Similar to Eq.(<ref>), our spatial RC can be simplified as 

    x(t)=f{[α x(t-1)+β W_in⊙ W_cor⊙ u(t)]W_res}.

Here W_cor is a designed weight for the features of each stock index benchmark. f{} and W_res are the accumulated nonlinear effect and reservoir weight, respectively, through the whole system.



Thanks to the large number of pixels on SLM (400× 400), it can easily function as a multi-input portal, allowing a large amount of information to be encoded into the system at once. In this way, the temporal correlation among the time-series data is translated into spatial coherence of the optical field, whose dynamics implement a complex-number kernel function. Therefore, besides evaluating the 1-step prediction, where the result contains only the next step value, we also assess the system's performance for multi-step ahead prediction using more inputs. In this case, u(t) is the vector containing all the feature information for different time steps. Specifically, we present examples of multi-step prediction with four-step and ten-step ahead predictions, labeled as 4-step and 10-step, respectively. Taking 4-step prediction as an example, the input vector u(t) has 32 values, corresponding to one closing price and 7 features per time step. Accordingly, the prediction horizon is 4 steps containing 4 days of closing prices beyond the most recent input. This encoding method exploits the spatial correlation of the free space optics that occurs naturally in our setup. 






 §.§ Multiple ML and LSTM models

This paper uses LSTM and various regression methods to predict each benchmark <cit.>. The LSTM architecture includes an internal memory that acts like a local variable, storing information as it processes a sequence of inputs. It uniquely handles outputs by recycling the production from one time-step as the input for the subsequent time-step<cit.>. This recycling process enables the LSTM to make informed predictions by considering both the current input and the output from the immediate preceding timestep. Besides LSTM, we also compared multiple regression models for stock index prediction. The regression methods used contain Lasso, ElasticNet, RandomForestRegressor, LinearRegression, Ridge, SGDRegressor, KNeighborsRegressor, DecisionTreeRegressor, and BayesianRidge<cit.>. 

For multi-step prediction, the data preparation for the time-series sequence with features follows the same approach as in optical RC in the previous subsection. The regression methods remain consistent with the one-step prediction, while the model dimensions in the LSTM are adjusted to align with the prediction horizon.






















































  













































§ RESULTS AND ANALYSIS




Figure <ref> (a) and (b) are two examples of 1-step prediction results for stock indices ^NYA and ^N225, respectively, showing the ground truth, the prediction results from RC with and without CW. Also shown are the prediction results of the best ML model, which is the Lasso regression model for both ^NYA and ^N225. 
The NRMSEs of RC prediction for ^NYA with and without CW are 0.104 and 0.108, respectively, whereas the best ML regression result is 0.558, which is over five times worse than the RC prediction. To compare the overall prediction performance with the results for ^N225 in Fig. <ref>(b), the prediction results for ^NYA show a small difference between predictions with and without CW.
The difference for ^N225 between the two is much larger, as illustrated in Fig.<ref>(b), where the NRMSE is 0.688 without CW, and only 0.236 with CW. This indicates that the added feature correlation weights can significantly improve the prediction performance for this benchmark.


In contrast, the predictive performance of ML models is considerably inferior to that of the RC, with the most accurate ML model achieving an NRMSE of 1.309 for ^N225. This is clear in Fig. <ref>, where both the best ML curves for ^NYA and ^N225 in Fig. <ref> (a) and (b) lag behind the price movement trend and deviate significantly from the true values. Conversely, our optical RC predictions capture the stock volatility and nonlinear trends more accurately. These predictions align better with the real-time series stock prices than the best ML regression predictions.



To quantify their performance difference, Figure <ref> lists the error reduction (ER) between the best ML model and optical RC predictions for seven stock indices. The ER is calculated using the formula 

    ER=1- NRMSE_RC/NRMSE_ML,

where NRMSE_RC and NRMSE_ML are the NRMSE of the RC and the best-performing ML methods, respectively. 
The above ER measures the performance improvement by RC compared to the best ML model. A near 100% ER indicates negligible relative error achieved by RC.  
As shown, the RC model significantly outperforms the ML model, with all ER values being positive and more than half of the results exceeding 70% for the different stock indices. Interestingly, for the stock indices ^N225, ^HSI, MME=F, and VTHR, the RC model shows superior performance in 1-step predictions compared to multi-step predictions. For other indices, the RC model achieves comparable or higher ER values in multi-step predictions than in 1-step predictions.


To gain insights into prediction performances between RC and ML models, we plot the prediction results of the two models for ^NYA in Fig. <ref>(a) and (b), respectively. 
Fig.7(a) shows the prediction outcomes from our RC model. The different-step prediction lines fit closely to the ground truth, with only a slight offset around a few peaks for 10-step prediction. 
However, the results of the ML model results in Figure <ref>(b) are more dispersed to compare the results for RC prediction. Additionally, the red rectangular area in Figure <ref>(b) highlights a 'delayed' effect in the multi-step prediction. Larger prediction dimensions result in greater drift from the ground truth. Consequently, as the forecasting horizon extends, prediction accuracy diminishes, leading to greater deviations from actual outcomes.















§ CONCLUSION


In conclusion, we have established an optical reservoir computing system in free space that effectively captures broader market behavior and demonstrates superior performance in stock index prediction. Compared to the existing machine learning methods and long short-term memory neural networks, our system offloads intense computation to an optical system with low energy consumption, showing significant advantages in handling high volatility and nonlinear market behaviors using limited data. This method promises a cost-effective approach to real-time, parallel, multi-dimensional data processing and predictions, with broad applications in financial forecasting. The next steps include moving more data processing into the optical domain, for high-speed, high-quality, and versatile applications in financial analysis. 





















elsarticle-num 
























 Paper Title: A Systematic Review on Long-Tailed Learning.















§ ABSTRACT 
Long-tailed data is a special type of multi-class imbalanced data with a very large amount of minority/tail classes that have a very significant combined influence. Long-tailed learning aims to build high-performance models on datasets with long-tailed distributions, which can identify all the classes with high accuracy, in particular the minority/tail classes. It is a cutting-edge research direction that has attracted a remarkable amount of research effort in the past few years. In this paper, we present a comprehensive survey of latest advances in long-tailed visual learning. We first propose a new taxonomy for long-tailed learning, which consists of eight different dimensions, including data balancing, neural architecture, feature enrichment, logits adjustment, loss function, bells and whistles,  network optimization, and post hoc processing techniques. Based on our proposed taxonomy, we present a systematic review of long-tailed learning methods, discussing their commonalities and alignable differences. We also analyze the differences between imbalance learning and long-tailed learning approaches. Finally, we discuss prospects and future directions in this field.



Long-tailed learning, Long-tailed data, Imbalance learning,  Deep learning, Deep imbalance learning.


A Systematic Review on Long-Tailed  Learning
    Chongsheng Zhang,  Senior Member, IEEE,
        George Almpanidis,
         Gaojuan Fan1,
         Binquan Deng,
         Yanbo Zhang,
         Ji Liu,
         Aouaidjia Kamel,
         Paolo Soda,
        João Gama,  Fellow, IEEE

Chongsheng Zhang, George Almpanidis, Gaojuan Fan, Binquan Deng, Yanbo Zhang and Aouaidjia Kamel are with Henan University, Henan Key Lab of Big Data Analysis and Processing, Kaifeng, China. E-mail: cszhang@ieee.org; almpanidis@ieee.org; {fangaojuan,bqdeng,zhangyanbo,kamel}@henu.edu.cn. 
 Gaojuan Fan is the corresponding author of this work.
Ji Liu is with Baidu Inc., Beijing, China. E-mail: jiliuwork@gmail.com.
Paolo Soda is with University Campus Bio-Medico di Roma, Rome, Italy.  Email: p.soda@unicampus.it.
João Gama is with the Laboratory of Artificial Intelligence and Decision Support, Faculty of Economics, University of Porto, Porto, Portugal. Email: jgama@fep.up.pt. 

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







§ INTRODUCTION


The term “Long Tail Theory” was popularized by Chris Anderson's book  “The Long Tail: Why the Future of Business Is Selling Less of More” <cit.>, in which the shift in business from advertising and marketing a small number of bestseller products to paying attention to the vast amount of niche products was observed and analyzed, since their combined sales/revenues could be as significant as that of the head (most popular) products. It is closely related to “the Pareto Principle”, also known as the “80/20 Rule”, the Zipf's law, or the “Vital Few Rule”. Generally speaking, they are like two sides of the same coin. While the Pareto Principle focuses on the heads/hits and cuts off the rest, the Long Tail Theory aims to capture the tails and emphasizes the combined importance of the tail  classes/niche items.

Long-tailed data is essentially a special type of multi-class imbalanced data with a sufficiently large number of tail (minority) classes. Moreover, the combined importance of these tail classes is very significant, although each tail class by itself only has  a small number of samples (sales). Long-tailed distributed data is a relatively common phenomenon in real-world scenarios, which often needs to be handled by artificial intelligence (AI) systems and applications, e.g. high-speed train fault diagnosis <cit.>, escalator safety monitoring <cit.>, etc. Indeed, AI should not only satisfy the head/hit applications, but also be able to cover a large number of tail cases to ensure the robustness and generalization capability of the techniques and systems.

Long-tailed learning (hereafter referred to as LTL for short) is a sub-domain of artificial intelligence/machine learning that aims to build effective models for applications/tasks having long-tailed distributed data. Its main goal is to significantly improve the recognition accuracy on the tail (rare/minority) classes or cases while maintaining the same or similar accuracy on the head (frequent/majority) classes or cases. Specifically, 1) in object recognition/classification, researchers have proposed many LTL methods <cit.> that can substantially improve the recognition/prediction accuracy on the tail classes/cases, e.g. rare species, defective industrial products. Studies in <cit.> have shown promising results by disentangling deep feature representation learning and classifier training. There are also a large number of loss reweighting approaches that endow different weights to samples of the head and tail classes to adjust the decision boundary under the long-tailed settings <cit.>. 2) in object detection, a few LTL methods <cit.> have also been designed to automatically locate rare objects or cases from images or videos, e.g. locating cracks in industry equipment or detecting the engineering vehicles near natural parks; 3) in image segmentation, researchers have also developed methods <cit.> that can identify and segment the rare objects/cases in images, such as abnormal actions, pathology areas in medical images, etc. In addition, specific data augmentation techniques for LTL have also been devised to mitigate the data scarcity problem in the tail classes <cit.>.



However, due to the rapid development of this field, keeping pace with recent advances in LTL is becoming increasingly difficult. As such, a comprehensive survey of existing methodologies in this field is urgent and beneficial to the community. This motivates us to conduct an in-depth survey of recent advances in long-tailed visual learning to gain insights into their principles and technical aspects in a systematic way. Based on the inherent learning process, we first propose a novel taxonomy that categorizes existing long-tailed visual learning approaches into eight groups, depicted in Figure <ref>, which are data balancing, neural architecture, feature enrichment, logits adjustment, loss function, bells and whistles,  network optimization, and post hoc processing approaches. Based on this new taxonomy, we present a comprehensive survey of state-of-the-art LTL approaches and discuss their ideas and characteristics. Moreover, we make comparisons between imbalance learning and long-tailed learning, in which we elucidate their connections and differences. Finally, we discuss challenges and future research opportunities in this field.

Though a few recent papers <cit.> also provide literature reviews on long-tailed learning, we distinguish our survey from them by the following differences: 1) Existing reviews often employ the conventional taxonomy that categorizes existing LTL approaches into three types, which are the data resampling and augmentation category, the loss reweighting category, and the transfer learning category. However, such a taxonomy does not facilitate a comprehensive understanding of state-of-the-art LTL methods since it can not sufficiently cover LTL's whole learning process. In this work,  we propose a taxonomy that provides a more unified perspective on LTL based on its internal learning process, where we identify four major steps with eight main categories, as shown in Figure <ref>. Using this new taxonomy, we provide a more up-to-date and insightful review of existing LTL methods, discussing their commonalities and alignable differences. 2) We make detailed comparisons between LTL and imbalance learning. 3) We include the latest advances in LTL and point out future directions in this field.

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




  * We propose a  unified taxonomy with eight dimensions to characterize and organize existing long-tailed learning methodologies.


  * We present a comprehensive, taxonomy-guided survey of state-of-the-art long-tail visual learning approaches, focusing on recent advances and trends.



  * We compare long-tailed learning with imbalance learning, elucidate their connections and  differences.


  * We summarize and analyze the results of different LTL methods in different downstream tasks, using the corresponding benchmark datasets.


  * Finally, we discuss a number of future research directions and trends that can be interesting to researchers in this field.



The remainder of this paper is organized as follows. In Section <ref>, we present the background and related survey papers in long-tailed visual learning. In Section <ref>, based on our proposed taxonomy, we present a comprehensive overview of state-of-the-art LTL methods.  In Section <ref>, we provide detailed comparisons between imbalance learning and long-tailed learning, then present a preliminary discussion on long-tailed distribution in Section <ref>.  In Section <ref>, we summarize and analyze the experimental results of different LTL methods in different downstream tasks. We discuss future directions and trends in Section <ref>, and conclude the  paper in Section <ref>. 





§ BACKGROUND AND RELATED WORK
 



 §.§ Background
 
Data imbalance has been a long-standing research issue in both machine learning and computer vision. In the past, researchers commonly used the term “imbalance” to refer to datasets with skewed class proportions <cit.>. However, in the last few years, many deep learning-based methods have emerged for overcoming the data imbalance challenges in computer vision, making this research field even more popular. For instance, <cit.> are among the first methods that adopt deep learning to learn feature representations from imbalanced image data. In the meantime, people have been increasingly using the term “long tail” in these works. Today, these two terms are often used interchangeably.

One reason for the popularity of the “long tail” term in computer vision in recent years is that the training datasets nowadays are often of large scale and present highly imbalanced distributions, and they have a large number of tail (minority) classes. Another reason is the use of deep learning for overcoming the challenges of class-imbalanced distributions and the new problems and research opportunities brought by such deep learning methods. Specifically, the built-in data-driven feature extraction mechanism in deep learning is significantly influenced by the scale of the datasets and their distributions. For image datasets with long-tailed distributions, deep learning-based methods tend to become class-biased in the feature extraction step, hence the need for more advanced  LTL approaches that can effectively extract representative features for the tail classes.

In comparison, traditional computer vision methods use hand-crafted feature extraction methods, which are often model-driven and independent of data distributions. After the feature extraction step, if the resulting features are class-imbalanced, these approaches rely on the subsequent data classification stage to counter such imbalance issues. That is, traditional methods mainly rely on the data classification algorithms/stage for tackling the imbalance problem, where many existing imbalance learning algorithms <cit.> can be straightforwardly applied. Early approaches such as <cit.> fall into this category, which modifies the kernel features of SVM to optimize the class boundary when the extracted features are class-imbalanced.

Besides computer vision, the term “long tail” has also been used in many others fields. 1) In search engines and recommender systems, the long tail issue has been studied in depth, aiming to suggest items or results that are highly relevant to infrequent queries <cit.> or specific users <cit.>. 2) In network traffic modeling,  it is widely accepted that network traffic follows a long-tailed distribution, and different distribution models such as hyper-exponential distribution <cit.> and mixed Erlang model <cit.> have been proposed to model such  long-tailed network traffic data. 





 §.§ Related Work
 

In the literature, there are a few related survey papers on long-tailed learning <cit.>. In <cit.>, the authors propose a taxonomy of three main categories, which are information augmentation, class re-balancing and module improvement.  Their first category consists of data augmentation and transfer learning techniques; the second category contains resampling, class-sensitive learning and logit adjustment methods; the third category consists of  metric learning, classifier design, decoupled training and ensemble learning approaches. In <cit.>, the authors uses a taxonomy for the training stage of LTL that consists of data augmentation, resampling, cost-sensitive loss, multiple experts and transfer learning. In <cit.>, the authors present a taxonomy of data processing, cost-sensitive weighting, decoupled learning and other methods. 

In our taxonomy,  the LTL process is organized into four major modules, containing a total of eight parts/categories. As can be observed from Figure <ref>, the first module is the (input) data pre-processing step, which aims to balance the number of samples of different classes to be used for neural network training via data balancing. The second module is the neural network design  (modeling)  step for effective representation / embedding learning. Based on the internal processing flow of a neural network, it consists of  five parts/categories, which are neural architecture, feature enrichment, logits adjustment, loss function, and bells and whistles (i.e., various training strategies, techniques and practices for model performance improvement). The third module is the (internal) network optimization step, which  aims to effectively update the huge amount of  weights in a network via optimization algorithms and computing the gradients which are the partial derivatives of the loss function with respect to the weights for guiding the update process. The last module is the post-processing (output) step, which adjusts the confidence of the trained long-tailed classifiers or refines their predictions to better fit to the scenarios of real-world applications. 

Comparing with the above surveys <cit.>, our taxonomy is significantly different. First of all, we divide the overall LTL process into four major modules based on the internal learning process of a neural network, including the (input) data pre-processing/balancing step, the neural network design (modeling) step, the  (internal) network optimization step, and the post hoc processing step,  which is more holistic and natural, and easier to understand and accept for those who are new to the LTL field. Second, our taxonomy is more comprehensive, systematic and and in-depth. 1) Previous overview papers in LTL use the term/category “loss reweighting” or “cost-sensitive weighting/learning”, which can not cover other loss functions for the general embedding learning purposes, which we will address in subsection <ref>.  2) Previous overview papers <cit.> lack a systematic discussion on logit adjustment, which is a preliminary step ahead of loss function design. In subsection <ref>, we will specifically illustrate the logits calibration and Softmax margin learning techniques under this category. 3) We categorize weights and gradients rebalancing approaches <cit.> into network optimization <cit.>, which can not be covered by taxonomies in existing survey papers. 












§ LONG-TAILED LEARNING METHODOLOGIES
 

In Figures <ref> and <ref>, we present a new taxonomy for long-tailed learning, in which we divide existing techniques into eight categories, which are 1) data balancing; 2) neural architecture; 3) feature enrichment; 4) logits adjustment; 5) loss functions; 6) bells and whistles; 7) network optimization; and 8) post hoc processing. 

Data balancing approaches utilize data resampling or data augmentation or data synthesis techniques to build class-balanced training sets for training neural network models. Neural architecture approaches design specific network architectures to improve LTL performance. Feature enrichment methods aim to augment the features of the tail samples using a memory bank or the perturbation strategy, or utilize the additional or multi-modal features extracted from pre-trained/foundation  models. Logits adjustment approaches adjust the logit values or enlarge the classification margins to improve LTL performance. Methods in loss function design either upweight the loss values of the tail-class samples or difficult samples, or devise comprehensive loss functions as optimization objectives to enhance the embedding/representation learning effectiveness. “Bells and whistles” contains various strategies and technical details for network training and performance enhancement. Network optimization mainly involves the weights and gradients updating techniques for the internal network optimization.  Finally, post hoc processing methods calibrate the confidence of a  long-tail model to fit to real-world scenarios.  In the following, we will introduce representative techniques in each category.






 §.§ Data Balancing
 

Data balancing aims to increase the volume and diversity of the training samples of the minority (tail) classes to make the training samples class-balanced. Along this line, there are three representative subcategories of approaches, which are a) data resampling, which balances the number of samples in different classes; b) data augmentation, which expands the size of the training data via transformation operations; c) data synthesis, which generates synthetic samples using more advanced techniques such as GANs, distribution estimation/transfer methods, foundation models (such as GPT-4V, DALL-E <cit.>), etc. Table <ref> summarizes representative methods in each subcategory. 






  §.§.§ Data Resampling
 
In machine learning, oversampling/upsampling the minority classes and undersampling/downsampling the majority classes are two commonly used strategies for handling class-imbalanced distributions <cit.>. In deep learning, class-wise sampling (e.g.<cit.>) and instance-wise sampling (e.g. <cit.>) are two frequently used methods for balancing the number of per-class training samples in the mini-batches. Essentially, class-wise resampling may involve both oversampling (the samples of the minority classes) and undersampling (the samples of the majority classes), while instance-wise resampling (i.e. instance-wise resampling) adopts random/uniform sampling (undersampling) to select samples from the original datasets, regardless of their class labels. In <cit.>, the authors show that class-balanced re-sampling can learn discriminative feature representations when the training samples are highly semantically related to their target labels; otherwise, uniform sampling is even better than class-balanced re-sampling.

In <cit.>, the authors first perform clustering within each class, then select samples through a  quintuplet sampling scheme that enforces both inter-cluster and inter-class margins. Such a sampling strategy can alleviate the information loss brought by random under-sampling. 

Existing approaches often treat samples of the same class equally, without taking into account the  difficulty of individual instances. Hard example mining is a commonly adopted technique in deep learning for enhancing the learning capability of the models <cit.>. To improve the recognition performance on datasets with extremely imbalanced image attribute distributions, the authors in <cit.> propose to mine both the hard positives, i.e. samples from the minority classes that are predicted to belong to other classes or with low prediction scores on the ground-truth class, and the hard-negatives, i.e. samples from other classes that are misclassified to one of the minority classes, which can help enlarge the margins between classes.





  §.§.§ Data Augmentation

Data augmentation is a commonly used technique in deep learning for artificially expanding the training set to improve the accuracy of the models while avoiding overfitting. It includes basic image transformations such as cropping and colour space transformations, and more advanced approaches such as image mixing <cit.>.

RandAugment and its variants. AutoAugment <cit.> and RandAugment <cit.> aim to find a group of transformation operations for data augmentation such that the deep model trained by a neural network can obtain better or the best accuracy on the target dataset. AutoAugment automatically searches for the best image transformation choices/policies for a dataset. RandAugment further investigates how to reduce the search space, where straightforward grid search is found to be very effective. Unlike RandAugment, CUrriculum of Data Augmentation (CUDA) <cit.> is designed to find more fine-grained class-wise augmentations, i.e., the number of candidates for sequential augmentation operations for different classes. They reveal that strong/more data augmentations should be applied to samples of the majority (easy) classes instead of the minority (difficult) classes. Therefore, minority classes are given shallow-degree sequential augmentation operations in the training process for curriculum learning. 

Mixup and its variants. Mixup <cit.> is a data augmentation technique that generates synthetic images with soft labels by performing linear interpolations in both the raw input space and label space. The soft label describes the prior inter-class relationship between the two input images. The main difference between Mixup and conventional transformation-based augmentation techniques lies in that Mixup alters both the raw input space and the label space when creating synthetic/augmented images. The idea of Manifold Mixup <cit.> is similar to Mixup <cit.>, with the only difference being that the former interpolates two images in the latent feature space. Remix <cit.> adapts Mixup to the long-tailed setting by assigning a soft label in favour of the minority class, giving a higher label mixing coefficient to the minority classes. Similarly, UniMix <cit.> extends Mixup to the long-tailed setting by adopting a tail-favoured Mixup factor and an inverse sampling strategy to encourage more occurrences of head-tail pairs. Recently, the authors in <cit.> propose BEM which is a novel data augmentation method for long-tailed semi-supervised learning. It uses a sampling strategy based on class-wise entropy to rebalance data uncertainty in data sampling, and the class activation map to achieve more accurate localization to avoid redundant areas in data mixing. <cit.>  proposes a dynamic Mixup approach for data augmentation, which determines the number of per-class images to select in the next training epoch according to the recognition performance in the previous training epochs. <cit.> proposes the Long-Tail Mixed Reconstruction (LMR) algorithm that reconstructs the few-shot videos through weighted combinations of the head-class samples/videos in the batch using Mixup, which reduces overfitting on few-shot classes.




Cutmix and its variants. Cutmix <cit.> is another image data augmentation strategy that replaces a randomly selected region in an image with a patch from another image during image mixing. The soft label of the generated image can be obtained in the same way as Mixup, but the weights to be used in the linear interpolations should be in proportion to the areas of the patches in the generated images. Based on the idea of Cutmix, the Context-Rich Minority Oversampling (CMO) method <cit.> creates synthetic minority images by simply pasting the real minority samples (as the foreground) onto the majority samples (as the background). This way, the rich contexts can be transferred from the majority samples to the synthetic minority samples, making them more diversified for training LTL models. <cit.> proposes a new augmentation method that first extracts the rich background images from the head classes using activation maps then pastes the tail class images upon them to generate tail-class images with diverse contexts. Instead of arbitrarily pairing the foreground and background images like <cit.>, OTmix <cit.> considers the semantic distances between the foreground and background images/distributions and selects the most suitable background image to generate more meaningful samples.

Other augmentation methods. ISDA <cit.> is a data augmentation method that produces diversified augmented samples by translating features along semantically meaningful directions (e.g. colour, backgrounds, visual angles). MetaSAug <cit.> extends ISDA to the long-tailed setting by learning the semantic directions via meta-learning. Major-to-minor Translation (M2m) <cit.> iteratively modifies (translates) a majority sample until it can be confidently classified as a minority sample by a pre-trained classifier built upon the imbalanced dataset.  M2m  leverages the richer information of the majority samples and does not overly use the minority samples, which can thus learn more generalisable features for the minority classes. <cit.> reveals that existing data augmentation methods have side effects since they sacrifice the performance of certain classes. They propose DODA that dynamically maintains class-wise data augmentation preference list by comparing the number of correctly predicted samples of each class in the current and previous epochs. 

Augmentation methods for long-tailed object detection. In <cit.>, the authors present a comprehensive review on the class imbalance problems in object detection. They summarize four types of imbalance problems in object detection: class imbalance, scale imbalance, spatial imbalance, and objective imbalance. <cit.> analyzes the issue of Intersection over Union (IoU) distribution imbalance in two-stage object detection and proposes the positive Region of Interest (pRoI) generator, which increases the number of positive examples, especially for those with greater IoU values, and improves the final object detection performance.

The authors in <cit.> point out that, for long-tailed object detection, image-level resampling alone is not enough to yield a sufficiently balanced distribution at the object level. They thus propose a joint resampling strategy, termed RIO (Resampling at Image-level and Object-level), which augments the minority classes in the current batch with the object-centric RoI features and box coordinates stored in memory over the past batches. MosaicOS <cit.> is another augmentation method that generates pseudo scene-centric images by stitching object-centric images (i.e. ImageNet) for long-tailed object detection. The difference between RIO and MosaicOS is that RIO uses a memory bank to reuse the object-centric RoI features and box coordinates over the past batches; in contrast, MosaicOS straightforwardly generates synthetic scene images by stitching several object-centric images into one image.







  §.§.§ Data Synthesis

While data augmentation techniques generate synthetic images based on image transformations, data synthesis approaches utilize more complex techniques for image generation, including GANs <cit.>, distribution estimation/transfer  <cit.>, pre-trained vision/vision-language models (foundation models <cit.>), etc. 

GAN based data generation. <cit.> finds that information at the lower resolutions tends to be class-independent thus can be shared by all the classes, while class-specific features are usually unveiled at higher solutions. Inspired by such observations, in GAN-based image synthesis for long-tailed data, it proposes to calculate the unconditional (class-independent) objective for low-resolution images, while injecting conditional information at higher-resolution images. 


Distribution based data synthesis. The authors in <cit.> point out that estimating the centroids of the tail classes is much noisier because they have few samples and propose DisRobuLT, which optimizes against the worst-case centroids within a safety hyper ball around the empirical centroid. Rare-class sample generator (RSG) <cit.> is a data synthesis method which obtains the feature displacement of a real frequent-class sample by subtracting its feature vector with the closest cluster centre of the same class to remove the class-relevant information and obtain the divergence value. It next adds the above divergence value of the frequent-class samples directly to the feature vector of each rare-class sample to create synthetic minority samples. The authors in <cit.> propose LEAP, which first uses the angles of the cosine similarity between each feature vector and the centre of the corresponding class centre to model a Gaussian distribution for each class, next forces the variance of the Gaussian distribution for each tail class to be the same as the averaged variance of the head classes, then randomly selects feature vectors using the new distribution of each tail class. LEAP enlarges the variance of the feature distribution of each tail class so that they become more separated from other classes. GistNet <cit.> also proposes to share the geometry of the head classes with the tail classes. However, unlike LEAP, it uses an alternative implementation of the Softmax function by adding the weight vector with a shape parameter that contains the geometry information shared by all the classes. Since the head classes tend to dominate the optimization of the shared parameters, the geometry of the head classes will be naturally transferred to the tail classes.  Overall, RSG forces the intra-class variation information of the tail classes to be the same as the head classes to generate synthetic samples for the tail classes, while DisRobuLT estimates the centroids of the tail classes using the distributional robustness theory. LEAP uses the averaged variance value of all the head classes as the variance of the tail classes to model the distribution of each tail class, whereas GistNet proposes to learn the shared geometry parameters from the head classes and then transfer them to the tail classes. 




Foundation model based data generation. With the rise and prosperity of large pre-trained models/foundation models, e.g.,  CLIP <cit.>, GPT-4V and DALL-E <cit.>, researchers  attempt using such large pre-trained models to generate synthetic data to address the long-tail problem. In <cit.>, the authors leverage Text-to-Image (T2I) model <cit.> to generate diverse images for the tail classes, then refine them by  measuring the cosine similarity between the  image and the corresponding text features extracted via CLIP <cit.>.








 §.§ Neural Architecture

Neural architecture design, i.e., devising an effective neural network structure/model  for a specified task, long with loss function design, are the two preliminary approaches to improving the overall performance of a neural network (model). Along this line, related methods can be divided into subcategories of decoupled learning framework, early-exit framework, knowledge distillation framework, dual streams inference framework, and multi-expert (ensemble) learning framework.






  §.§.§ The Decoupled Learning Framework

Decouple <cit.> is a seminal work that reveals representation learning and classifier training should be carried out separately (i.e. decoupled) for improving long-tailed recognition. Specifically, as shown in Figure <ref>, the proposed decoupled LTL framework first adopts random sampling (instance-based sampling) and the Cross-Entropy loss for training a conventional feature extraction module in the first stage. It then freezes the backbone and uses class-balanced (class-wise) sampling to retrain the classifier in the second stage. Under the same decoupling framework as  <cit.>, <cit.> proposes to train several teacher models upon batches built via instance sampling in the first stage so that features extracted by different models can be more complementary, then distils the knowledge from the teachers to a student model. <cit.> also adopts the decoupling framework to address the class imbalance issue in medical image datasets, the difference is that it utilizes contrastive learning <cit.> to learn the feature space in the first stage. In <cit.>, the authors use class-balanced sampling and conventional random sampling simultaneously to construct mini-batches of data for the feature extraction step.  They reveal that random sampling avoids the risk of overfitting on the tail classes, while class-balanced sampling prevents the training process from being dominated by head classes. 

The above methods have a clear boundary between the two decoupled learning stages. In Bilateral-Branch Network (BBN) <cit.>, the authors propose to use an annealing factor to gradually transits between the two stages, which can adjust the model to learn the global features from the original distributions first, then gradually transit to the tail data modeling and discrimination learning. In Domain Balancing (DB) <cit.>, the authors devise a soft gate to switch between a conventional feature learning module and a tail-class favoured feature learning module. For each tail-class image, the features from the two modules are integrated to improve the LTL performance. Hybrid-SC <cit.> proposes a hybrid network structure composed of a supervised contrastive learning branch to learn better feature representations from the long-tailed data and a  conventional CNN branch with a cross-entropy loss to learn the classifiers. Like BBN, it smoothly transits between the two branches via a  cumulative learning adaptor. The main difference between Hybrid-SC and BBN is that the former uses supervised contrastive learning in the first stage, which can learn better features. In <cit.>, the authors reveal that the inference of visual relationships in a scene faces the problem of the dual long-tailed distributions of the entity (object) and predicate (relation) classes. To address this problem, they propose the DT2 model, which is based on an iterative decoupled training scheme. Like BBN, DT2 first trains (warms up) the model on the original data with uniform sampling, then fixes the backbone part and fine-tunes the classifier on the (remaining) network with the proposed Alternating Class Balanced Sampling (ACBS) method to capture the interplay between the long-tailed distributions of entities and predicates using knowledge distillation.  

<cit.> points out that the performance bottleneck of the two-stage (decoupling) learning framework seems to be the classifier retraining issue in the second stage. It thus proposes DisAlign for confidence calibration to adjust the decision boundary in the second (classifier retraining) stage. In <cit.>, the authors propose the long-tailed class incremental learning problem and present a preliminary method under the two-stage LTL framework that adds a layer of learnable scaling weights to integrate the outputs of the classifier heads for class-incremental learning.




  §.§.§ The Early-Exit Framework
 
In <cit.>, the authors point out that existing approaches to long-tailed recognition reweight samples only by class size, but without considering the fact that there might be easy examples in the minority classes that get incorrectly up-weighted and hard examples in the majority classes that get erroneously down-weighted. They thus employ the early exiting framework (ELF) to early-exit easy examples in the shallow layers of the network to focus on the hard examples in each class in the deeper layers. In implementation, ELF augments a backbone network with several auxiliary classifier branches. During training, it finds and early-exits the correctly and confidently classified examples at the earliest possible exit branch, while harder examples will exit later in the network, sustaining their impact on the overall loss. Similarly, <cit.> proposes the Hardness Aware Reweighting (HAR) framework, which also adopts the early-exiting strategy in the training process. Specifically, HAR augments a backbone network with auxiliary classifier branches, and samples that are confidently and correctly classified at a branch will early exit from that branch.




  §.§.§ The Knowledge Distillation Framework

In <cit.>, the authors propose Self-Supervised Distillation for long-tailed visual recognition (SSD-LT). Adopting the decoupled learning framework, it trains two heads for knowledge distillation, which are the self-distillation head and the classification head, respectively. It thus owns the merits of rebalanced sampling and decoupled training strategy and is reported to obtain state-of-the-art performance in long-tailed recognition. The authors in <cit.> propose the “Distill the Virtual Examples” (DiVE) pipeline for long-tailed recognition. In DiVE, a teacher model is first trained for the long-tailed task with the existing methods, and the predictions from such a teacher model are used as virtual examples. It then uses knowledge distillation to transfer the knowledge from the teacher model to the student model.

<cit.> discusses the use of ViT/Transformers for long-tailed learning. It proposes DeiT-LT which distills knowledge from a CNN teacher to a ViT backbone/student using out-of-distribution images generated via CutMix and Mixup, in which it reweights the distillation loss to enhance the focus on the tail classes.  Recently, for long-tailed object detection, RichSem <cit.> adds a new branch to the detection framework for distilling semantics from CLIP, which can extract rich semantics for the bounding box patches within the image. 




  §.§.§ The Dual Stream Inference Framework
 
<cit.> designs a neural network architecture with two branches: one branch trains a baseline model on the original long-tailed distribution, while the other trains a second model upon the artificially balanced distributions. They add a cross-branch consistency loss based on their logit difference to enforce the two branches to train the multi-label visual recognition model collaboratively. For long-tailed face recognition, Center-dispersed (CD) Loss <cit.> first splits the original data into head-class data and tail-class data, then feeds it into the “unequal-training” framework that consists of a  backbone and two learning branches for the head classes and tail classes, respectively. The first learning branch is based on the head-class data and is used to stabilise face representation, supervised by a noise-resistant loss. The second learning branch is based on the tail-class data and is used to enhance the inter-class discrimination capability. The training of the framework alternates between the two branches. One difference between <cit.>  and <cit.> lies in that the latter trains two models on the split data (head-classes data and tail-classes data) separately, whereas the former trains two models on the original/complete data.

In SimCal <cit.>, the authors investigate instance segmentation over data with long-tail distributions. Upon Mask R-CNN, their proposed method makes three main modifications: i)  it changes the inputs of the original dataset with batches of class-balanced samples, then selects class-balanced object proposals for each class in the region proposal generation (RPN) stage;  ii) in the classification/segmentation stage, it trains a calibrated classifier with Mask R-CNN on the above proposals; iii) in the inference step (after the training is over), it employs the calibrated classifier and the classifier trained on the original dataset using the vanilla Mask R-CNN algorithm to predict the tail and head classes, respectively. Such a dual-head inference framework yields improved performance on the tail classes while still maintaining similar performance on the head classes.




  §.§.§ The Multi-Expert Learning Framework
 
Ensemble learning methods aim to utilize multiple experts/models to make more accurate predictions. Representative ensemble learning methods for LTL include ACE <cit.>, RIDE <cit.>,  CBD <cit.>, LFME <cit.>, BalPoE <cit.>, SHIKE <cit.>, MDCS <cit.> and LGLA <cit.>. 

The authors in <cit.> propose ACE, a multi-expert ensemble network for long-tailed recognition. It aims to achieve biased representation learning and unbiased classifier training in a single stage. The class indices are first ordered descendingly by their number of samples, the first branch (expert) trains a classifier/model with the samples of all the classes; the next branch (expert) first eliminates the current head class, then trains a new classifier with the samples of the rest classes, and so on. Therefore, the number of classes handled by different classifiers decreases monotonically, and the medium-shot or tail classes will have chances to dominate a classifier, thus mitigating the bias towards the majority classes. Overall, ACE can simultaneously improve the recognition accuracy of both the head and tail classes. RIDE <cit.> trains multiple recognition models/experts on randomly sampled subsets and uses their logits mean to make predictions. It also devises a routing algorithm to skip the inferences of certain experts. CBD <cit.> also trains multiple recognition models/experts on randomly sampled subsets, it proposes to distil the knowledge from the multiple teachers/experts into a single student model. Similarly, LFME <cit.> also propose to distil the knowledge from multiple teachers into a unified/single student model. LFME differs from CBD in that the teacher models of the LFME are trained over non-overlapped partitions of classes and samples, while CBD trains multiple teacher models over randomly sampled subsets which may have overlapped classes/samples. Compared to CBD and LFME, ACE <cit.> trains cascaded models/classifiers by iteratively removing the dominating group of classes and the corresponding samples before training the next teacher model, while RIDE  straightforwardly utilize logits mean/average instead of knowledge distillation to ensemble/ally the teacher models. 

The authors in <cit.> propose BalPoE, which is also an ensemble learning framework composed of logit-adjusted experts <cit.>. They prove that the ensemble is Fisher-consistent for minimizing the balanced error. SHIKE <cit.> is a knowledge distillation-based multi-expert ensemble learning framework for long-tailed recognition. The authors find that shallow features of a network (or a shallow network) can better represent the tail classes. They thus propose aggregating the deep features in different experts/branches with different layers of shallow features, respectively. Next, they use knowledge distillation techniques to let the experts learn from each other. Like other ensemble learning algorithms, the averaged predictions from the multi-experts will then be used as the final prediction. Nested Collaborative Learning (NCL) <cit.> also trains multiple experts on all the categories and some of the hard categories, respectively. It then uses knowledge distillation to learn the multiple experts collaboratively. MDCS <cit.> is a multi-expert learning framework that introduces a diversity loss (which is essentially the Balanced Softmax loss <cit.>) to control experts' focus on different categories with an adjustable weight to improve the diversity of experts, and performs  self-distillation between weakly and strongly augmented logits sets for each expert to distill the richer knowledge of weakly augmented instances to reduce model variance. 

Considering the fact that the class distributions of test data are not necessarily uniform but might be long-tailed or even inversely long-tailed, the authors in  <cit.> propose SADE which is a test-agnostic long-tailed recognition approach that trains three experts/models on uniform, long-tailed, and inversely long-tailed class distributions separately, then weighted aggregates them at the inference stage. DirMixE <cit.> is also an multi-expert learning approach, which uses the Monte Carlo method to estimate the mean and semi-variance of a meta-distribution for capturing both the global and local variations to adapt to test-agnostic distributions. LGLA <cit.> trains multiple experts on the entire dataset to extract more discriminative features for different subsets, which uses logit adjustment strategy to enlarge the discrepancy among the subsets. It also learns an expert excelling at inversely long-tailed distribution, next ensembles these experts to achieve superior prediction performance.

<cit.> proposes to leverage the powerful discrimination ability of the large-scale pre-trained vision models <cit.> to tackle long-tailed classification directly. It proposes LPT which first adapts the pre-trained model to the target domain using shared prompt tuning, next adopts the divide-and-conquer strategy for long-tailed recognition by using the pre-trained model and shared prompt to generate a  group-specific prompt set for each group of classes. 

Besides, lately, the Transformer architecture  <cit.> is being increasingly utilized in LTL <cit.>. Transformer based LTL methods, such as  LTGC <cit.> and VL-LTR <cit.> reported significantly better performance than CNN based methods in long-tailed recognition, while Step-wise <cit.>  which adopts the CNN-Transformer combined architecture, also showed superior performance in long-tailed object detection. 




 §.§ Feature Enrichment




Feature enrichment methods enrich/augments features of the tail classes  using fusion-based approaches <cit.>,  perturbation based approaches <cit.>, large pre-trained models/foundation models<cit.>, etc. 




  §.§.§ Fusion Based Feature Enrichment


As depicted in Figure <ref>, the Retrieval Augmented Classification (RAC) method <cit.> uses an external memory module to store features of the tail samples explicitly. Next, for each input image, it finds in the memory bank the Top-K most similar features of the same class, which are then fed into a BERT-like encoder to be transformed into logits, then fused with the logits derived from a conventional backbone. Likewise, the authors in <cit.> propose PFENet, a prior guided feature enrichment network that enriches query features with support features and prior masks to overcome the spatial inconsistency between the query and support samples in few-shot segmentation. LOCE  <cit.> also utilizes a memory module to store features of the tail-class samples for augmenting the features of the tail classes, which works collaboratively with the proposed  Equilibrium Loss (EBL) that adjusts the margins for the tail classes in long-tailed detection. In OLTR <cit.>, the authors combine the features extracted by a conventional CNN pipeline and the features stored in an additional memory module to generate features of the tail classes, which can transfer knowledge between head and tail classes, and increase the sensitivity of the model to novel classes.

<cit.> proposes to disentangle image features into  class-agnostic features (a.k.a. category-related/class-relevant) and class-generic features (a.k.a. category-independent/class-irrelevant)  for  one-shot image generation. It uses an auxiliary memory module to store class-generic features. The class-agnostic features from the input image and the class-generic features from the memory module are then fused and sent to the generator of a vanilla GAN for image generation. Likewise, in <cit.>, the authors point out that some features are shared by the head and tail classes, while other features are class-agnostic (unique to certain classes). The authors thus propose to transfer the class-generic features from the head to the tail classes. <cit.> propose to augment tail classes by grafting the diverse semantic information from head classes, where it replaces a portion of feature maps from tail classes with those belonging to head classes. Such fused features substantially enhance the diversity of tail-class representations. To prevent the neural networks from fitting the head classes first, MFW <cit.> weakens features of the head classes by mixing them with features of other classes in the same mini-batches but without altering the corresponding labels to equalize the training progress across classes.

CoSSL <cit.> devises a tail-class feature enhancement mechanism that gives a higher probability to the tail classes when blending features of unlabeled and labelled data.




  §.§.§ Perturbation Based Feature Augmentation

The authors in <cit.> propose the FASA method for long-tailed instance segmentation that generates class-wise virtual features by random perturbation over a Gaussian distribution on the real samples of each class. 

CC-SAM <cit.> proposes to improve the generalization capability of the LTL models by robust training against class-conditioned perturbations on the features of the last convolutional layer. Its idea is similar to that of LogitAdjust <cit.>. The main difference is that LogitAdjust performs perturbations over logit values in the output layer, whereas CC-SAM adds noises to the features in the last convolutional layer (the last hidden layer).



  §.§.§ Foundation Model Based Feature Enrichment


Recently, foundation models such as CLIP <cit.> and ALIGN <cit.> have demonstrated impressive representation and generalization capabilities. In <cit.>, researchers have  explored using such foundation models for feature-level enhancements. They simultaneously utilize the visual and linguistic representations extracted by the foundation models CLIP <cit.> for long-tailed recognition. For long-tailed object detection, RichSem <cit.> adds a new branch to the detection framework for distilling knowledge from  CLIP, which can extract rich visual and linguistic semantics to enhance the representations of the objects within the image.




 §.§ Logits Adjustment
 
In the output layer, logits contain the raw prediction values of the neural network, which will be used as the input to the Softmax activation function to turn the logit values into a probability distribution over the input classes. Methods that fall into the logits adjustment category either optimize the inherent process in deriving the logit values  <cit.>, or calibrate the logit values in a post hoc manner <cit.>. It also includes softmax margin learning methods <cit.> that enforce margins into the logit values so as to influence the resulting Softmax classifier to have larger classification margins <cit.>. We note that, sometimes it is challenging to clarify the differences between logits calibration and Softmax margin learning. For instance, LogitAdjust <cit.> can be considered to be belonging to both subcategories. In Table <ref>, we summarize the formulas of representative methods in this category.






  §.§.§ Logits Calibration


Inherent Logits Calibration. When deriving logits from the last hidden layer to the output layer, IR-Softmax <cit.> sets the weights for each class as their corresponding class centres in the feature space to avoid the shift between the weights and their centres. The authors in <cit.> propose Balanced MSE that straightforwardly uses the mean square error on each class to replace the corresponding logit element in imbalance regression. 

Post hoc Logits Calibration. In <cit.>, the authors argue that the obscure regions between classes in classification are caused by Softmax saturation. To alleviate this issue, they propose the Gaussian Clouded Logit adjustment (GCL) loss that perturbs the normalized logit with an additive margin that has a random disturbing parameter. In <cit.>, the authors reveal that pseudo-labels in semi-supervised learning are naturally imbalanced. They propose a Debiased Pseudo-Labeling method (DebiasPL) that deducts the probability for each class from the logit. After that, margin learning methods can still be applied. 

Like GCL <cit.>, in BLV <cit.>, category-wise variations are introduced into network predictions (logits) for long-tailed semantic segmentation, where tail classes obtain larger variations by multiplying the random value in a Gaussian distribution with a coefficient that is inversely proportional to the per-class sample size. This simple strategy improves the segmentation performance on the tail categories.

<cit.> finds that a classifier trained on long-tailed/imbalanced data produces skewed class probabilities on a solid color image, rather than uniform probabilities, it thus proposes to adjust the classifier's predicted logits on a test sample by subtracting its predicted logits on a white image, leading to improved recognition performance.

Most LTL methods train a model on long-tailed data but evaluate their performance on uniformly distributed data. To address this issue, the authors in <cit.> propose the LAbel distribution DisEntangling (LADE) method, which regularizes the logits so that they can be explicitly disentangled from the source label distribution and adapt to arbitrary target distributions.




  §.§.§ Softmax Margin Learning

Softmax margin learning algorithms alter the logit values before sending them to the Softmax classifier (Softmax activation function) to improve the discrimination power of the classifiers. Therefore, it can also be considered belonging to the “Post hoc Logits Calibration” sub-category. In label-distribution-aware margin loss (LDAM) <cit.>, when using the logits to calculate the  Softmax output,  the corresponding logit element for the ground-truth class of the input sample will subtract a class-wise margin parameter, which is reversely proportional to the quadratic root of the number of samples in the class. Equalization Loss (EQL) <cit.> removes the logit values of other tail classes in the denominator of the Softmax activation function. AutoBalance <cit.> adopts a generalized cross entropy loss with multiplicative and additive margin parameters, and uses a bi-level optimization procedure to learn these parameters on a validation set directly. The overall ideas of BALMS <cit.> and  LogitAdjust <cit.> are very similar to AutoBalance: the main difference lies in the multiplicative and additive margin parameters used in the Softmax function, as can be seen from Table <ref>. Similarly, in <cit.>, the authors propose the Category-Aware Angular Margin (C2AM) loss for long-tailed object detection, which adds a margin for every non-ground-truth class,  based on the ratio between the weight norms for the ground-truth class and every other class. In <cit.>, the authors deduce that margin-based binary classification provides a closed-form solution for the ideal margin of each object category in long-tailed object detection and propose the Effective Class-Margin (ECM) loss for optimizing margin-based classification.

In <cit.>, the authors propose a progressive margin loss (PML) for facial age classification. They propose a framework that consists of an ordinal margin learning head and a variational margin learning head specially tailored to long-tailed data. The weighted sum of the two margins is adopted to make adjusted predictions.

As a short summary, logits adjustment either intervenes the derivation process of logits, which occurs between  the last hidden layer and the output layer, or modifies the existing logit values, primarily by adding different margin parameters. Ultimately, the calibrated logits will be fed into the Softmax Activation Function, resulting in a better-performing Softmax classifier.




 §.§ Loss Function

An embedding space, also known as the latent space or latent feature space, is the new space produced by a learning algorithm where semantically similar objects are positioned close together, and dissimilar objects are placed far apart <cit.>, as illustrated in Figure <ref>. Roughly speaking, embedding (metric) learning refers to any technique or process that converts an input image into an effective feature representation. Loss function design, along with neural architectures,  are the main approaches to embedding learning and optimization. In fact, the aforementioned feature enrichment techniques, logit adjustment techniques, neural architecture,  as well as training strategies (bells and whistles), are all integral parts of the overall embedding learning process. 






  §.§.§ Loss Function For Embedding Learning
 

In this subsection, we will  introduce various loss functions for deep embedding learning in long-tailed settings.


Long-tailed  Recognition. The authors in <cit.>  propose Affinity Loss which jointly performs max-margin classification and feature space clustering for enhancing inter-class separability and reducing intra-class variations. DASO <cit.> generates unbiased pseudo-labels by blending two complementary pseudo-labels obtained via a linear classifier and the similarity computation between the feature vector of a sample and each class centre, respectively. In DLSA <cit.>, the authors propose to tackle long-tailed recognition by adjusting the label space of the samples via deep clustering. They first cluster the samples according to their likelihood in the feature space. The well-clustered samples are sent to a dedicated Cluster-aided Classifier, which exploits the cluster prior information in classification. The remaining samples are outliers and are sent to the next adjustment module to decrease the training difficulty further. SBCL <cit.> devises a triplet loss like contrastive loss  for long-tailed recognition. It initially divides the head classes into multiple sub-classes  and enforces a sample to be closer to samples of the same sub-class than other samples in the same class, while being closer to samples of a different sub-class under the same class than samples from different classes in the loss function. GLMC <cit.> is based upon the Siamese network architecture: for each pair of samples, it respectively uses the CMO <cit.> and Mixup <cit.> methods to create two synthetic images, then uses both the cosine distances between the features of the two images and the classification loss on the two images to optimize the network.  

In <cit.>, the authors define the distribution overlap coefficient to quantify the dominance severity of the head classes in the representation space and formulate the inter-class discrepancy and class-feature consistency losses to close the gap. The authors in <cit.> find that the distribution of the super-classes in imbalanced data is more balanced than the original data. They thus propose SuperDisco to discover the super-class in a hierarchy of semantic abstraction, then refine the original features of a sample using the corresponding super-class features. 

Contrastive learning  <cit.> has recently gained popularity since it can substantially improve the capability of the neural network in representation learning. The authors in <cit.> propose the K-positive Contrastive Loss (KCL) for LTL, which draws k instances from the same class to form the positive sample set for contrastive learning. Based upon KCL, targeted supervised contrastive learning (TSC) <cit.> adds a new loss term between the samples and the predefined optimal class centres to move the samples closer to the centres of their class and away from centres of other classes.  The PaCo <cit.> loss introduces learnable class-wise centers into the supervised contrastive learning process to make the probability that two samples are a true positive pair become more balanced across different classes, which is beneficial for long-tailed recognition. Model-Aware K-center (MAK) <cit.> sorts the external samples by their empirical contrastive loss expectation value, then selects the Top-N samples to perform K-means clustering and chooses the final samples. The authors in <cit.> propose  Balanced Contrastive Learning loss (BCL) that includes cluster compactness in the contrastive learning process. 

In <cit.>, the authors design a hybrid network comprising a prototypical supervised contrastive (PSC) learning module and a classifier training module. A curriculum is designed to control the weightings of these two modules during network training.  Motivated by the observation in <cit.> that model pruning negatively influences the prediction accuracy on samples of the infrequent classes, the authors in <cit.> design the  Self-Damaging Contrastive Learning (SDCLR) framework, which adds a pruned contrastive learning module to the (original/non-pruned) unsupervised contrastive learning module, to amplify the prediction differences between the pruned and non-pruned models and increase the weights of the infrequent class samples and hard samples in the overall loss. SDCLR increases both the overall accuracy and the balanceness of the learned representations.

<cit.> reveals that existing semi-supervised learning methods obtain low recall performance in the minority classes, but the corresponding precision performance is high. It thus proposes CReST for  semi-supervised LTL that iteratively samples pseudo-labelled data from the unlabeled set to expand the labelled set, then retrains the semi-supervised learning model. SimPro <cit.> is a simple probabilistic framework for long-tailed semi-supervised learning, which explicitly separates the estimation process into conditional and marginal distributions using the Expectation-Maximization algorithm.  ABC <cit.> attaches an auxiliary classifier head to the representation layer of a semi-supervised learning framework to avoid bias towards the majority classes in LTL. ACR <cit.> uses a dual-branch network to adaptively handle various class distributions of unlabeled data, in which it derives the scaling parameter using the bidirectional KL divergence between the estimated class distribution and three anchor distributions. 

Long-tailed  Regression. When quantizing the imbalanced regression target space into interval classes for imbalance regression, HCA <cit.> proposes a coarse-to-fine hierarchical classification strategy which first uses coarse classifiers to provide the range, next selects the corresponding finer classifier in the range, and achieves improved regression performance.

Long-tailed  Object Detection and Segmentation. The authors in <cit.> propose the Balanced Group Softmax method (BAGS), which groups the classes by the number of instances, then trains a separate Softmax classifier for each group to reduce the suppression by the head classes. However, finding a proper group partition number for  BAGS could be laborious and tricky. Thus, the authors in <cit.> devise the Adaptive Class Suppression Loss (ACSL), which utilizes the output confidence to prevent the classifiers of tail categories from being over-suppressed by the samples from head categories, and only the samples from visually confusing categories are kept for discriminative learning. Learning to Segment the Tail (LST) <cit.> also adopts the “divide-and-conquer” strategy for long-tailed segmentation. It divides the original data into several balanced parts and devises a knowledge distillation loss which measures the logit difference of the same class in neighbouring parts to preserve knowledge learned in previous parts. <cit.> adds an Adaptive Hierarchical Representation (AHR) loss to the traditional object detection loss for more discriminative embedding learning. It calculates both the coarse-grained classification loss for distinguishing each cluster clearly, and the fine-grained classification loss within each cluster by adopting adaptive margins. In <cit.>, the authors point out that, in contemporary deep learning-based object detectors, classification and regression are often optimized independently, which leads to inconsistent predictions with high classification scores but low localization accuracy, and vice versa. They thus propose a harmonic mutual-supervision method which devises a regression-aware classification loss to optimize the classification branch and a classification-aware regression loss to supervise the optimization of the regression branch. <cit.> proposes an unsupervised instance segmentation method that uses a pre-trained mask generation network to generate object proposals, next samples mask triplets from the proposals to capture the hierarchical relations and learn the embeddings. 

Considering the advantages of the Gumbel distribution in modeling extreme events with very low probabilities, the authors in <cit.> propose GOL that utilizes the Gumbel activation function for long-tailed object detection and segmentation. Through large-scale experiments, they demonstrate the superiority of Gumbel against Sigmoid and Softmax in detecting/segmenting rare-class objects.




  §.§.§ Loss Reweighting
 

After obtaining the loss on each class or sample, loss reweighting approaches aim at adjusting the importance of their loss values through the weighting factors. Existing approaches along this line can be divided into class-wise reweighting and instance-wise reweighting methods. In Table <ref>, we report the formulas of representative methods in this sub-category.





Class-wise Loss Reweighting. Class-wise loss reweighting methods assign the same weighting factor to instances of the same class or group, as demonstrated in Figure <ref>. CB Loss <cit.> is a class-balanced loss function for long-tailed distributions that reweights the samples by the inverse class frequency, i.e. the weight/factor of a class is inversely proportional to the number of samples in the class. AREA <cit.> proposes to re-weight the loss of each class based on the inverse Pearson correlation of the samples in the same class.  Range Loss <cit.> is also a class-wise reweighting method that consists of an intra-class loss and an inter-class loss: the former is defined as the harmonic mean of the two largest intra-class distances, while the latter is the shortest distance between any two class centres. The Learning to Reweight method (L2RW) <cit.> is a meta-learning algorithm that assigns weights to a batch of samples based on their similarities in gradient directions.  TCR <cit.> is a two-component reweighting loss function that simultaneously computes the class-wise weights using the CB Loss <cit.> and the conditional distribution between the source and target with L2RW <cit.>. 


Instance-wise Loss Reweighting. The above methods assign the same weight to all samples belonging to the same class. In contrast, instance-wise reweighting methods assign different weights to different samples based on their learning difficulty. Focal Loss <cit.> proposes to assign relatively larger weights to samples with high prediction loss, independent of their class labels. In <cit.>, the authors propose the Equalized Focal Loss (EFL) for long-tailed object detection, which modifies the original Focal loss by supplementing the focusing parameter with a new term that is the inverted ratio between the accumulated positive and negative gradients of the input sample's ground-truth class. In <cit.>, the authors propose the Population Stability Index (PSI) to measure the overall trend that an instance experiences loss increase over time. It pays more attention to difficult instances, including those from the majority classes.  In IB Loss <cit.>,  the authors propose a method that reweights each sample differently using the designed influence function to alleviate overfitting in the decision boundary. The authors in <cit.> discover that models trained on long-tailed datasets are more miscalibrated and over-confident than those trained on balanced data. They thus propose the label-aware smoothing soft label method  (MiSLAS) that assigns different weights to the head and tail classes to handle different degrees of over-confidence. 




 §.§ Bells and Whistles

In the context of long-tailed learning, “bells and whistles” refer to various techniques, training strategies and practices that go beyond the neural architectures and loss functions but are also essential for improving model performance. Common “bells and whistles” include network training and  optimization strategies and choices, such as advanced optimizer, learning rate schedulers, regularization, and commonly adopted techniques in deep learning such as data augmentation, pre-training and transfer learning, knowledge distillation, and ensemble learning for improving the performance of deep learning models. 

There are different types of optimization algorithms used in neural networks <cit.>, among them,  Stochastic Gradient Descent (SGD) is the major algorithm for optimizing the parameters (weights and biases) in neural networks for long-tailed learning. Learning rate (LR) schedulers are methods that adjust the learning rate during network training, such as linear decay, step decay (such as CosineAnnealing), time-based decay and exponential decay to optimize model performance over time. Regularization techniques in deep learning are methods used to prevent overfitting and improve the generalization ability of neural networks.  Commonly used regularization techniques include L2 regularization (Weight Decay) <cit.>, dropout, batch normalization, label smoothing (which replaces hard labels with smoothed distributions to reduce the model's confidence and encourage it to be more robust and generalize better) <cit.>, etc. 

Since data augmentation and synthesis, knowledge distillation, ensemble (multi-expert) learning methods have already been addressed in the previous subsections,  we only present a few pre-training and transfer learning techniques, and foundation models based approaches. 

Pre-training and transfer learning for long-tailed learning. SSP <cit.> introduces self-supervised pre-training into class-imbalanced learning to yield good network initialisation, which can alleviate the label bias issue in imbalanced datasets and enhance LTL performance. FBGL <cit.> uses both contrastive pretraining and normalization for improving the feature extractor and classifier for long-tailed recognition, which consists of a new balanced contrastive loss to improve  long-tailed pretraining and a generalized normalization method for the classifier (feature and weight vectors). The authors in <cit.> propose a simple yet effective step-wise learning framework based on Deformable DETR <cit.> that combines fine-tuning and knowledge transfer for long-tailed object detection. In specific, after selecting representative examplars for the head and tails classes,  it first pre-trains a model on all the classes, next learns a head-class model via fine-tuning, then transfers knowledge  from the head-class model to the final model. 

Leveraging large foundation models and multi-modal models for long-tailed learning. Lately, there has been a popular trend in utilizing the powerful feature representation capability of foundation models <cit.> and pre-trained multi-modal models <cit.>  to improve the LTL performance, which can also be considered as a “trick” for long-tailed learning. For instance, integration of textual and visual features  enhances the model's ability to discriminate between classes with fewer examples <cit.>, or enables the model to understand and reason about the semantic relationships between concepts depicted in images and described in texts <cit.>.

Finally, we note that, these strategies and techniques are often collectively used to make LTL models more effective and efficient. 




 §.§ Network Optimization

Network optimization refers to the internal mechanism for updating and optimizing the neural networks to  minimize the loss/objective function. Besides neural architecture and loss function design, optimizing the huge amount of internal learnable parameters of a neural network, mainly the weights and biases, plays a very important role in efficiently and effectively training and optimizing a deep learning model. In particular, this process involves techniques in updating the weights, and the gradients with respect to the loss function and weights.  Existing approaches in this line of research can be divided into weights rebalancing and gradients rebalancing approaches.

We clarify that, “network optimization” refers to techniques that optimize or utilize the weights and gradients of the network for internal network optimization, while “Bells and Whistles” in the above subsection refer to different external training strategies, choices, techniques and tricks for obtaining better learning performance.




  §.§.§ Weights Rebalancing

The authors in <cit.> show an obvious correlation between sampling frequency and the norm of the weight vectors, and propose to raise the weight vector of each class by multiplying the ratio between the number of samples in the most frequent class and the current class to adjust the decision boundary in long-tailed recognition. The authors in <cit.> also reveal that, for long-tailed data, a normally trained deep learning model has imbalanced norms of classifier weights, with the majority classes having dominating weights. They propose to balance network weights using standard regularization techniques such as weight decay (L2 regularization) and notably improve the long-tailed recognition performance. <cit.> discusses why weight balancing <cit.> is effective in two-stage long-tailed learning. It reveals that  weight decay helps achieve better feature extraction, while  class-balanced loss enables implicit logit adjustment in the training. Curvature-Balanced  Regularizer (CR) <cit.> uses the sum of the logarithm of the inverse of the maximum normalized Gaussian curvature mean of each class as a regularizer to enable the model to learn curvature-balanced manifolds. 




  §.§.§ Gradients Rebalancing

<cit.> reveals that long-tailed data suffers from gradient distortion problem in which the overall gradient is shifted toward the head gradient,  it thus proposes a two-phase long-tailed recognition method which first updates model parameters  using the gradient on  the head classes only, next grows the classifier on the added tail classes. The authors in <cit.> propose BSGAL which is a gradient-based generated data contribution estimation method for long-tailed instance segmentation that calculates the loss and gradient of the model on a batch of data with and without augmentation to actively accept or reject the generated data. SeeSaw Loss <cit.> dynamically rebalances the gradients of positive and negative samples for each category with two complementary factors to reduce punishments on tail categories and increase the penalty of misclassified instances. For long-tailed class-incremental learning, <cit.> proposes to separately reweight the gradients for the new tasks and learned tasks, based on the accumulated magnitude of gradients and the volume of the lost training data of each class, respectively. 

Overall, properly tuning the above factors is crucial for the internal optimization of the neural network to achieve high model accuracy and prevent overfitting in long-tailed learning.





 §.§ Post Hoc Processing



As demonstrated in Figure <ref>, post hoc processing methods modify the outputs of the (conventionally) trained models to tackle the imbalance problem at the inference stage, instead of improving the feature representation capability or adjusting the decision boundary at the model training process. The expected calibration error (ECE) <cit.> is often used to measure how well the predicted probabilities match the ground-truth probabilities, which first partitions the confidence scores into several equal-size bins,  then computes the difference between the accuracy and the averaged confidence in each bin.




  §.§.§ Post Hoc Confidence Calibration

Temperature Scaling (TS) <cit.> uses a single scalar parameter T to rescale the logits of a trained neural network, which will then be transformed into calibrated confidence scores using the Softmax activation function. However, this approach has yet to be validated in long-tailed settings. 

For confidence calibration in long-tailed object detection, NORCAL <cit.>  multiplies the prediction (output) scores with a vector of factors that are inversely proportional to the training sample size of each class. 

For confidence calibration in long-tailed object segmentation, the Pairwise Class Balance (PCB) loss <cit.> post-calibrates a sample's prediction probability on each class by multiplying the corresponding pairwise misclassification bias obtained from the confusion matrix. 




  §.§.§ Confidence Aggregation

When a neural network has multiple heads for inference, e.g. the dual streams inference framework and the multi-expert learning framework addressed in subsections <ref> and <ref>, these methods commonly use the logits mean <cit.> or the weighted aggregation of the logits <cit.>, as well as the confidence mean to make the final prediction.





§ COMPARISONS BETWEEN LONG-TAILED LEARNING AND IMBALANCE LEARNING
 

Imbalance learning is a long-standing research direction in machine learning <cit.>.  Long-tailed learning  belongs to imbalance learning, but they also have differences in the learning approaches. While conventional imbalance learning algorithms commonly deal with rational/tabular imbalanced data, long-tailed learning often handles image data, which needs the feature representation learning step as a prerequisite. That is, for image data, LTL needs deep representation learning techniques to extract features from images, in which neural architecture, feature enrichment and logit adjustment techniques are needed. 

However, for relational/tabular data, its features are its attribute  values, therefore,  features are already available before imbalance learning, thus feature representation techniques such as neural architecture, feature enrichment and logits adjustment, are not essential in imbalance learning. But we note that, in recent years,  deep imbalance learning approaches have also emerged, which use neural networks to improve the imbalance learning performance on relational/tabular data. 

In the following, we will mainly compare the data balancing and loss functions/cost-sensitive learning methods, as well as the multi-expert/ensemble learning methods adopted in both LTL and imbalance learning. 




 §.§ Data Balancing





  §.§.§ Data Resampling


In long-tailed learning, only random sampling is utilized, either at class-level or instance-level, referred to as class-wise sampling and instance-wise sampling, respectively. However, in imbalance learning, more advanced undersampling and oversampling techniques have been proposed. 

Undersampling techniques in imbalance learning. In Condensed Nearest Neighbor Rule <cit.>, for each sample from the majority set, if it can not be correctly classified by the Nearest Neighbor (NN) rule, it is kept; otherwise, it is discarded. In <cit.>, the authors propose an intuitive geometric space partition (GSE) based method for imbalance classification which iteratively uses a new hyper-plane classifier to cut the current geometric data space into two partitions, then removes the partition that only contains the majority of samples. 

Oversampling techniques in imbalance learning. ROS  <cit.> is the most straightforward oversampling technique, which  balances the data by replicating the minority class samples. The Majority Weighted Minority Oversampling Technique (MWMOTE) <cit.> first identifies the most important and hard-to-learn minority class samples, then assigns weights to them according to their Euclidean distance to the nearest majority class samples.




  §.§.§ Data Augmentation


For image data, in long-tailed learning, Mixup <cit.>, ReMix <cit.> and UniMix <cit.> create synthetic samples with interpolated feature vectors and labels out of two original samples.

For relational/tabular data, since the features are already available before imbalance learning, it is easy and common to use linear interpolation algorithms to generate synthetic minority samples between two real minority samples of the same class. SMOTE <cit.> creates synthetic minority samples by randomly selecting a new point on the line segment between a pair of neighbouring minority samples. Borderline SMOTE <cit.> extends SMOTE by emphasizing the borderline samples that have both majority and minority class points in their neighbourhood and ignoring the rest of the minority samples. ADASYN <cit.> creates synthetic minority samples according to data density, where a comparatively larger number of synthetic samples are created (using SMOTE) in regions of low density of the minority class than in higher-density regions. SMOTE-RSB <cit.> first uses SMOTE to generate synthetic instances for the minority classes, then employs the rough set theory as a cleaning (undersampling) method to select only the synthetic minority examples that belong to the lower approximation of their class.

In general, the idea of Mixup, ReMix and UniMix is similar to SMOTE  and its variants, but the former interpolates the label space as well, while SMOTE and its variants only handle samples of the same minority classes,  but without modifying the labels of the synthetic samples. 




  §.§.§ Data Synthesis


In LTL, there are distribution based data synthesis methods, such as (RSG) <cit.>, LEAP <cit.> and GistNet <cit.>, GAN based <cit.> and foundation model based data generation <cit.> methods. In comparison, in imbalance learning,  GANs are commonly utilized to generate synthetic tabular data, such as TableGAN <cit.>, CTGAN <cit.>, medGAN <cit.> and QAST <cit.>. 

TableGAN <cit.> was the first attempt to synthesize tables/relational databases using GANs. It adopts the Deep Convolutional Generative Adversarial Network (DCGAN) architecture, but adds a new neural network module called “classifier” and a corresponding loss function to measure the discrepancy between the label of a generated record and the label predicted by the classifier for this record. CTGAN <cit.> models the distribution of the values with a few Gaussian distributions (modes) for each continuous attribute. For discrete attributes, CTGAN introduces a conditional generator and training-by-sampling method to ensure that each category of a discrete attribute gets a fair chance to be included during sampling. This way, it can simultaneously generate a mix of discrete and continuous attributes. Medical Generative Adversarial Network (medGAN) <cit.> uses Auto-Encoders to learn the salient features of discrete attributes and GAN to generate realistic samples in the embedding space. To address the usability of GAN-generated samples in imbalanced learning,  QAST <cit.> designs a semantic pseudo-labelling module to control the quality of the generated features, then calibrates their corresponding semantic labels using a classifier committee consisting of multiple pre-trained shallow classifiers. The above approaches only generate synthetic examples for the minority classes. In <cit.>, the authors propose using GAN to generate a subset of the majority of samples that follow the same distribution as the original dataset. Next, they construct a few balanced subsets, each containing all the minority samples and an equal-sized subset of GAN-generated majority samples. Then, upon each decomposed two-class subset, they use deep metric/embedding learning to build a two-class classification model. These models are then integrated to make the final predictions.  Recently, a diffusion model <cit.> based imbalance learning method, namely TabDDPM  <cit.>,  has been proposed for generating rational samples.




 §.§ Loss Function




  §.§.§ Loss Function for Embedding Learning


The idea behind embedding learning is to map complicated (often high-dimensional) metric spaces into easier (often low-dimensional) metric spaces, where inter-class and intra-class distances can be optimized simultaneously. In LTL, there are many specially designed embedding learning approaches, such as <cit.>, <cit.>, <cit.>, <cit.>, <cit.>, etc.

In imbalance learning over relational data, many traditional (non-deep learning based) embedding learning methods have been proposed to transform the learning space. IML <cit.> is an imbalance learning method that aims to find a more stable neighbourhood space for the testing data by performing metric learning on the data and training sample selection according to the testing data iteratively until the selected training samples in two adjacent iterations are the same. MLFP <cit.>  expands the decision boundaries around the positive samples and learns better metric space using a triplet-based cost function on the false negative and false positive samples. 

In recent years, many deep imbalance learning approaches have also been proposed which use neural networks to transform the metric space and learn better embedding/feature representations. Khan et al. <cit.> propose a cost-sensitive (CoSen) deep neural network method to learn feature representations for both the majority and minority classes. <cit.> addresses the problem of imbalanced data for the multi-label classification problem using deep learning. It introduces incremental minority class discrimination learning by formulating a class rectification loss regularization, which imposes an additional batch-wise class balancing on top of the cross-entropy loss. This can rectify model learning bias due to the over-representation of the majority classes. GAMO <cit.> is an end-to-end deep oversampling model for imbalance classification, which effectively integrates sample generation with classifier training. 




  §.§.§ Loss Reweighting

In LTL, loss-level reweighting methods assign different weights to the loss values of the samples, which include class-wise reweighting algorithms such as CB Loss <cit.> and L2RW <cit.>, instance-wise reweighting algorithms such as IB Loss <cit.>, as illustrated in subsection <ref>. 

In imbalance learning, there are also a few cost-sensitive learning approaches that modify the internal cost function/objective function by assigning different weights to different samples based on their difficulty. Such methods usually assign a higher cost to minority classes to boost the importance of the minority classes during the learning process <cit.>. To adapt SVM to imbalanced data, cost-sensitive SVM (CostSVM) <cit.> adds a class-specific classification penalty in its cost function. 

Overall, the general ideas for loss-level reweighting in LTL and cost-sensitive learning in imbalance learning are very similar, i.e., they all give higher weights to samples of the minority classes or the misclassified samples. 




 §.§ Multi-Expert Learning


Multi-expert learning (ensemble learning) techniques are applicable to both LTL and imbalance learning.  For instance, the idea of ACE <cit.> to train cascaded learning models for LTL is similar to that of BalanceCascade <cit.> and AdaCost <cit.> for imbalance learning. 

There are also a few differences, for instance, i) Knowledge distillation-based ensemble techniques have been adopted in LTL <cit.>, but not in imbalance learning; ii) Dichotomy-based ensemble technique for multi-class imbalance learning <cit.> are not needed in long-tailed learning, since cross-entropy loss based LTL classifiers are naturally multi-class classifiers.




 §.§ Evaluation Metrics


In terms of evaluation metrics,  existing LTL methods commonly assume that the test sets are class-balanced, that is, all the classes have the same number of test instances in the testing/inference stage. Under such an assumption, only the Top-1 accuracy/overall accuracy metric is reported. However, we note that, in real-world applications, the test sets might also be long-tailed, or at least non-uniformly distributed.  In contrast, in imbalance learning the test sets are often considered to be imbalanced, and various evaluation metrics have been proposed <cit.>, including the overall accuracy, G-mean (the geometric mean of the per-class accuracy), Macro-F1, etc. 

In <cit.>, the authors argue that the average accuracy metric used in LTL incurs little penalty on classes with very low per-class recalls. They thus propose three new metrics for evaluating long-tailed recognition algorithms: the geometric mean and harmonic mean of the per-class recall (referred to as G-mean and H-mean for short), and the lowest per-class recall. They also propose the Geometric Mean Loss (GML), which essentially calculates the logarithm of the averaged prediction probability for samples of each class. 

As a short summary, when the test set distributions are non-uniform or skewed, major evaluation  metrics including overall accuracy, G-mean,  Macro-F1 and the lowest per-class recall, should  be adopted in LTL and imbalance learning to more comprehensively assess the performance of the proposed algorithms.





§ LONG-TAILED DISTRIBUTIONS: A DISCUSSION
 




 §.§ Comparing Long-tailed Data with Imbalanced  Data


When referring to long-tailed data, one may neglect the fact that long-tailed data is a special type of imbalanced data, and long-tailed learning (LTL)  is essentially part of imbalance learning (IL). 

The long tail phenomenon occurs in imbalanced data when a small number of classes are highly frequent and are thus easy to model. In the meantime, there are also a substantial number of tail classes that have insufficient samples and are thus difficult to model. Moreover, the combined quantity or influence of all the tail classes is comparable to that of the head classes. Specifically, there are three main differences between long-tailed distributed data and imbalanced data:

First, in long-tailed data, there should be a significantly large number of minority (tail) classes compared to the majority (head) classes. It thus has a very long tail, also named “fat tail” or “heavy tail”. However, in the case of imbalanced data, the number of classes may be as few as two (two-class imbalanced data) or only a few (multi-class imbalanced data). Therefore, long-tailed data must be class-imbalanced data, but class-imbalanced data is not necessarily long-tailed data.
  
Second, since long-tailed data has a very long tail, the overall trend/pattern of event frequency/sample size per class should decrease exponentially from the head classes to the tail classes, but the sub-pattern/trend among the tail classes should be a relatively flat curve with a slowly decreasing pattern. Ordinary class-imbalanced data does not necessarily have this property.
  
Third, the combined quantity or influence of the tail class samples in long-tailed data should be significant, e.g. 60% or 20% of the total quantity/influence of the head classes. This unique property does not hold in ordinary class-imbalanced data.




 §.§ Comparing Long-tailed Distribution with Pareto Distribution


A Pareto distribution is a power-law probability distribution with heavy, slowly decaying tails. In comparison, a “long-tailed distribution”  essentially corresponds to a set of qualitative properties/constraints added to a class-imbalanced distribution rather than being a statistical distribution that can be rigorously modelled. In the following, we empirically study the differences between a Pareto distribution and a long-tailed distribution.




In Figure <ref>, we first depict the empirical rank-size distributions of two real-world datasets, iNaturalist 2018 and LVIS, typically used in long-tailed learning settings. A rank-size distribution gives the empirical relative frequencies  of the data with the classes ordered in descending order in size across the x-axis[When calculating the relative frequencies, for normalisation reasons, we divide the sample size/frequency of each class by the maximum frequency (the frequency of the Top-1 class), instead of the total number of samples of all the classes.], which corresponds to the empirical survival distribution function (ESDF).



    S(x)=[ 1 + ( x-μ/σ) ^ 1/γ] ^ -α , x≥μ , μ∈ℝ , σ , γ > 0, α≥ 1


For comparing a long-tailed distribution with a Pareto distribution, in Figure <ref> we fit each dataset with a Generalised Pareto Type IV (DPIV) survival function (SF), using Equation <ref>, where  μ is the location, σ is the scale (which corresponds to the minimum value of x),  γ is the inequality parameter, and α is the tail index (shape). Robust estimators for the parameters of DPIV have been proposed in the literature. Here, after fixing parameters μ and σ in relation to the empirical distribution of each dataset, we exhaustively vary the values of α and γ with a step of s=0.1 to heuristically find the set of parameter values that minimises the difference between the ESDF and the fitting Pareto survival function, measured in terms of the relative errors/distances between them, which can be observed from the bottom sub-figures of Figure <ref>.

As the Pareto distribution (DP) embodies a heavy (fat) tail, it is natural to present the survival rates of both the theoretical and the empirical distribution on a logarithmic scale. In Figure <ref>, on the two real-world long-tailed datasets iNaturalist and LVIS, we see that the derived Pareto distribution is not guaranteed to consistently fit to all segments of the ESDF of the real long-tailed data.

In summary,  unlike the Pareto distribution, a long-tailed distribution is not a specific statistical distribution that can be precisely modelled, but rather an imbalanced distribution with a set of properties/constraints.




 §.§ When Using the Pareto Distribution to Model the Long Tail Phenomena



In case a Pareto distribution is used to approximately model the long tail phenomena, we need to add the following constraint to the Pareto distribution:

In long-tailed data, the combined quantity/significance of the tail class samples should exceed a user-specified ratio compared to the head classes, e.g. 60%, or 50%, or 20%.

The Cumulative Distribution Function (CDF) of the Pareto Distribution is given by:


    F(x)  = 1 - (x_m/x)^α


where x_m is the scale parameter, e.g. x_m = 10 or 3, which can be considered as the minimum number of samples that a tail class has.  α is the shape parameter, which is essentially the power law slope. 

Based on the Pareto Distribution, the constraint added by the long-tailed data depicted above can be formulated as: 


    F̅(h_l)= P(X ≥ h_l) = (x_m/h_l)^α≤ 1 - r


where h_l is the lower bound frequency of the head classes, i.e. the number of samples of the least-ranked head class, and r is a user-specified threshold which is the ratio between the combined quantity/influence of the tail classes and the total quantity/influence of all the classes.  Different r values can be set in different circumstances; for instance, r can be 60%, 50%, or 20%.  Essentially, for modeling long-tailed data using the Pareto distribution, Equation <ref> adds to the Pareto distribution a joint constraint on the parameters x_m, α, h_l and r.





§ PERFORMANCE ANALYSIS
 




 §.§ Datasets



For long-tailed recognition, benchmark datasets include CIFAR-100-LT <cit.>, ImageNet-LT <cit.>, Places-LT <cit.>, and iNaturalist 2018 <cit.>. For long-tailed object detection and instance segmentation, LVIS <cit.> is the primary real-world dataset. LVIS v0.5 consists of 1,230 categories, whereas LVIS v1.0 contains 1,203 categories, with 100K images and 19.8K images for the training and test sets, respectively. 




 §.§ Evaluation Metrics

For long-tailed recognition tasks, since existing LTL methods all assume that the test set is uniformly distributed, i.e.,  the sample size for each class in the test set is equal, therefore only the overall accuracy metric is needed for performance evaluation. Besides, these methods often separately report the specific performance on three super groups, which are “Many-shot” (if a class has > 100 instances), “Medium-shot” (20-100 instances) and “Few-shot” (< 20 instances).

For long-tailed detection and segmentation tasks,  the main evaluation metric is Average Precision (AP). For long-tailed segmentation, separate performance on three super groups are also reported, which are AP_f for frequent classes (having > 100 instances), AP_c for  common class (having 11-100 instances), and AP_r for rare classes (having < 11 instances).





 §.§ Performance Comparison





  §.§.§ Long-tailed Recognition


Table <ref> summarizes the performance of representative long-tailed classification algorithms on four benchmark datasets CIFAR-100-LT (with two different imbalance ratios (IR) which are 50 and 100), ImageNet-LT,  Places-LT, and iNaturalist 2018. It can be observed that, under the CNN backbone (ResNet), MDCS <cit.>, LGLA <cit.> and SHIKE <cit.> achieve the best overall performance on all the four benchmark datasets. Besides, NCL <cit.>, SADE <cit.>,  RIDE <cit.> and ACE <cit.> also obtain remarkable performance, all of which adopt the multi-expert learning framework, owing to its effectiveness. While under the Transformer backbone <cit.>, due to the utilization of foundation models, LTGC <cit.> and VL-LTR <cit.> achieve significantly better performance than CNN based methods, in particular, their performance improvements over MDCS <cit.> and LGLA <cit.> on ImageNet-LT  are near 20% in terms of overall accuracy, which is very substantial. 

Besides, we see that the performance of loss reweighting methods tend to saturate, since almost all the possible forms have been exhaustively examined, as can be observed from Table <ref>. 

In Table <ref>, we report the experimental settings of representative long-tailed recognition algorithms on the iNaturalist 2018 dataset. 

Overall, on ImageNet-LT and iNaturalist 2018, the state-of-the-art performance is only 82.5% in terms of overall accuracy, while the best overall performance is only 54.1% on Places-LT, which can hardly satisfy the requirements of real-world applications. Therefore, there is still a large room for performance improvement.




  §.§.§ Long-tailed Object Detection and Segmentation

Table <ref> reports the performance of representative long-tailed object detection and segmentation methods on the LVIS v1.0 benchmark dataset. We see that  BSGAL <cit.>, RichSem <cit.>, Step-wise <cit.>, PCB <cit.>, ECM <cit.>, C2AM <cit.> achieve the best long-tailed object detection performance, while BSGAL <cit.>,  GOL <cit.>, ECM <cit.>, PCB <cit.>, C2AM <cit.> and NORCAL <cit.> are among the best performers in long-tailed segmentation. We would like to note that, Step-wise <cit.> is implemented based upon Deformable DETR <cit.> which is essentially a CNN-Transformer combined detection framework, while the other methods are commonly based on Mask R-CNN with ResNet as the backbone. 

Overall, the state-of-the-art performance is below 35.4% in both  long-tailed object detection and segmentation tasks, and the segmentation performance on rare classes is even lower, which is less than 25.4%. Therefore, there is considerable room for future research.



















§ FUTURE DIRECTIONS
  
While numerous approaches have been proposed to address the long tail issue, many research opportunities remain.





 §.§ Federated Long-Tailed Learning

Federated learning enables multiple clients to learn a global model collaboratively without transmitting the local private data on each client to a centralized server <cit.>. The problem becomes more challenging when the training data across the clients is both heterogeneous and long-tailed. 

Fairness refers to whether a model can perform equally for diverse types of categories. Although there are a few preliminary studies <cit.>, ensuring fairness within a federated learning system in long-tailed settings still needs in-depth investigations.




 §.§ Long-Tailed Out-of-Distribution Detection



Out-of-Distribution (OOD) samples are samples that do not match the training data distribution, regardless of whether they belong to known (existing) categories or entirely new (open-set) categories. In <cit.>, the authors point that it is sometimes challenging to distinguish between the in-distribution tail class samples and the OOD samples that belong to entirely new (open-set) categories. To address this issue, <cit.> expands the in-distribution class space by introducing multiple abstention classes and adopts mixture of experts for classification. It also augments the context-limited tail classes using both the head and OOD classes via Cutmix. At test time, the sum of probabilities for the k abstention classes is used for OOD detection. In <cit.>, the authors also investigate the open-set long-tailed recognition problem, in which a long-tailed distribution also contains samples of new/unknown/open-set categories. They propose to simultaneously handle long-tailed classification, few-shot learning and open-set recognition (OOD detection) in a unified framework. 




 §.§ Active Learning for Long-Tailed Distributions


Active learning approaches iteratively select the most informative samples from underrepresented classes for annotation and training to optimize the use of limited labeled data to improve model performance over time. The authors in <cit.>  incorporate class imbalance into active learning by considering both class imbalance and labelling cost. <cit.> proposes a progressive sampling mechanism that explores active learning of OOD samples for enhancing the performance in both long-tailed classification and OOD detection. 




 §.§ Long-Tail Class Incremental Learning


To acquire novel knowledge from long-tailed class-incremental settings, <cit.> proposes to use sub-protypes to discover and store data distributions for sub-tasks with long-tailed distributions, and adopts the knowledge distillation technique to avoid catastrophic forgetting when training a new task. 




 §.§ Long-tailed  Domain Generalization


Domain generalization aims to learn from multiple training domains to generalize to previously unseen target domains. In <cit.>, the authors introduces the maximum square loss which has a linearly increasing gradient to address the long-tailed distribution issue in domain generalization.  While <cit.> focuses on single domain imbalance, <cit.> investigates  the multiple domains LTL problem which aims to learn from multi-domain imbalanced data and generalize to all domains. It proposes a loss for bounding the transferability statistics based on the contrastive loss over the positive and negative cross-domain pairs. 




 §.§ Adversarial Training Under Long-Tailed Distributions


Adversarial training is a defense strategy aimed at enhancing the robustness of machine learning models against adversarial attacks by incorporating adversarial examples into the training process. 
<cit.> investigates the adversarial vulnerability in long-tailed recognition and proposes RoBal, which includes a scale-invariant cosine classifier and a two-stage LTL framework. In <cit.>, the authors reveal that adversarial training under long-tailed distributions suffers from robust overfitting, i.e., the long-tailed model overfits to current adversarial samples. They show that simply combining Balanced Softmax Loss <cit.> and data augmentation can simultaneously alleviate robust overfitting and improve model robustness against adversarial attacks. 




 §.§ Diverse Applications with Long Tail Concerns

Long-tailed data widely exists in many real-world domains, e.g., defect identification <cit.>,  remote sensing <cit.>, nutrition ingredients analysis <cit.>, etc. While we can apply long-tailed learning techniques in these domains, such real-world applications often pose new challenges to be addressed,  which  will advance the long-tailed learning techniques in turn. 





§ CONCLUSIONS
  
In this paper, we propose a new taxonomy for long-tailed learning that divides existing techniques into eight categories, which are data balancing, neural architecture, feature enrichment, logits adjustment, loss function, bells and whistles, network optimization, and post hoc processing techniques. Based on this taxonomy, we systematically review long-tailed learning methods in each category, including the latest advances. We also examine the distinctions between imbalance learning and long-tailed learning approaches. We summarize the experimental results of representative methods, and show that there is still a very large room for further improvements. Finally, we discuss future directions in this field.

IEEEtran




 Paper Title: Non-convolutional Graph Neural Networks.


[
    Bruno S. Dandogbessi^2
    August 5, 2024
==========================





















0.3in












§ ABSTRACT 
Rethink convolution-based graph neural networks (GNN)—they characteristically suffer from limited expressiveness, over-smoothing, and over-squashing, and require specialized sparse kernels for efficient computation.
Here, we design a simple graph learning module entirely free of convolution operators, coined random walk with unifying memory (RUM) neural network, where an RNN merges the topological and semantic graph features along the random walks terminating at each node.
Relating the rich literature on RNN behavior and graph topology, we theoretically show and experimentally verify that RUM attenuates the aforementioned symptoms and is more expressive than the Weisfeiler-Lehman (WL) isomorphism test.
On a variety of node- and graph-level classification and regression tasks, RUM not only achieves competitive performance, but is also robust, memory-efficient, scalable, and faster than the simplest convolutional GNNs.




§ INTRODUCTION: CONVOLUTIONS IN GNNS

Graph neural networks (GNNs) <cit.>—neural models operating on representations of nodes (𝒱) and edges (ℰ) in a graph (denoted by 𝒢 = {𝒱, ℰ}, ℰ⊆𝒱×𝒱, with structure represented by the adjacency matrix A_ij = 1[(v_i, v_j) ∈ℰ])—have shown promises in a wide array of social and physical <cit.> modeling applications.
Most GNNs follow a convolutional scheme, where the D-dimensional node representations 𝐗∈ℝ^|V| × D are aggregated based on the structure of local neighborhoods:

    𝐗' = Â𝐗.

Here, Â displays a unique duality—the input features doubles as a compute graph.
The difference among convolutional GNN architectures, apart from the subsequent treatment of the resulting intermediate representation 𝐗', typically amounts to the choices of transformations (Â) of the original adjacency matrix (A)—the normalized Laplacian for graph convolutional networks (GCN) <cit.>, a learned, sparse stochastic matrix for graph attention networks (GAT) <cit.>, powers of the graph Laplacian for simplifying graph networks (SGN) <cit.>, and the matrix exponential thereof for graph neural diffusion (GRAND) <cit.>, to name a few. 
For all such transformations, it is easy to verify that permutation equivariance (Equation <ref>) is trivially satisfied, laying the foundations of data-efficient graph learning.
At the same time, this class of methods share common pathologies as well:



  
Limited expressiveness.
(Figure <ref>)

<cit.> groundbreakingly elucidates that GNNs cannot exceed the expressiveness of Weisfeiler-Lehman (WL) isomorphism test <cit.>.
Worse still, when the support of neighborhood multiset is uncountable, no GNN with a single aggregation function can be as expressive as the WL test <cit.>.
As such, crucial local properties of graphs meaningful in physical and social modeling, including cycle sizes (Example <ref>) and diameters (Example <ref>) <cit.>, cannot be realized by convolution-based GNNs.



  
Over-smoothing. 
(Figure <ref>)

As one repeats the convolution (or Laplacian smoothing) operation (Equation <ref>), sandwiched by linear and nonlinear transformations, the inter-node dissimilarity, measured by Dirichlet energy, 

    ℰ(𝐗) = 1/N∑_(u, v) ∈ℰ_𝒢 || 𝐗_u - 𝐗_𝐯 ||^2,

will decrease exponentially as the number of message-passing steps l increases <cit.>, 
ℰ(𝐗^(l))
≤
C_1 exp(-C_2 l)
with some constants C_1, 2,
resulting in node representations only dependent upon the topology, but not the initial embedding.



  
Over-squashing. 
(Figure <ref>)

As the number of Laplacian smoothing grows, the reception field of GNNs increases exponentially, while the dimension of node representation, and thereby the possible combinations of neighborhood environment, stays unchanged <cit.>.
Quantitatively, <cit.> quantifies this insight using the inter-node Jacobian and relates it to the powers of the adjacency matrix through sensitivity analysis:

    | ∂𝐗_v ^(l+1)/∂𝐗_u^(0) |
    ≤
    | ∇ϕ |^(l+1) (Â^l+1)_uv,

where ϕ is the node-wise update function, whose Jacobian is typically diminishing.
If this Equation <ref> converges to zero, the node representation is agnostic to the changes happening l edges away, making convolutional GNNs difficult to learn long-range dependencies.



  
Main contributions: Non-convolutional GNNs as a joint remedy.
In this paper, we propose a variant of GNN that does not engage the convolution operator (Equation <ref>) at all.[Code at: <https://anonymous.4open.science/r/rum-834D/>]
Specifically, we stochastically sample a random walk terminating at each node and use the anonymous experiment associated with the random walk to describe its topological environment.
This is combined with the semantic representations along the walk and fed into a recurrent neural network layer <cit.> to form the node embedding, which we call the unifying memory.
We theoretically (Section <ref>) and experimentally (Section <ref>) show that the resulting model, termed random walk with unifying memory (RUM) relieves the aforementioned symptoms and offers a compelling alternative to the popular convolution-based GNNs.





§ RELATED WORKS: WAYS TO WALK ON A GRAPH




  
Walk-based GNNs.
RAW-GNN (<cit.>, compared and outperformed in Table <ref>) also proposes walk-based representations for representing node neighborhoods, which resembles our model without the anonymous experiment (ω_u in Equation <ref>).
CRaWl (<cit.>, outperformed in Table <ref>) also incorporates a similar structural encoding for random walk-generated subgraphs to feed into an iterative 1-dimensional convolutional network.
AWE (<cit.>, Table <ref>) and <cit.>, like ours, use anonymous experiments for graph-level unsupervised learning and temporal graph learning, respectively.
More elaborately, <cit.> and AgentNet (<cit.>, Table <ref>) use agent-based learning on random walks and paths on graphs.



  
Random walk kernel GNNs.
In a closely related line of research, RWK (<cit.>, Table <ref>) employs the reproducing kernel Hilbert space representations in a neural network for graph modeling and GSN (<cit.>, Table <ref>) counts the explicitly enumerated subgraphs to represent graphs.
The subgraph counting techniques intrinsically require prior knowledge about the input graph of a predefined set of node and edge sets.
For these works, superior expressiveness has been routinely argued, though usually limited to a few special cases where the WL test fails whereas they do not, and often within the unlabelled graphs only.

More importantly, focusing mostly on expressiveness, no aforementioned walk-based or random walk kernel-based GNNs address the issue of over-smoothing and over-squashing in GNNs.
Some of these works are also de facto convolutional, as the random walks are only incorporated as features in the message-passing scheme.
Interestingly, most of these works are either not applicable to, or have not been tested on, node-level tasks.
In the experimental Section <ref>, we show that RUM not only outperforms these baselines on most graph-level tasks (Table <ref>) but also competes with a wide array of state-of-the-art convolutional GNNs on node-level tests (Table <ref>).
Moreover, random walk kernel-based architectures, which explicitly enumerate random walk-generated subgraphs are typically much slower than WL-equivalent convolutional GNNs, whereas RUM is faster than even the simplest variation of convolutional GNN (Figure <ref>).



  
State-of-the-art methods to alleviate oversmoothing.

Stochastic regularization. 
DropEdge (<cit.>, Figure <ref>) regularizes the smoothing of the node representations by randomly disconnecting edges.
Its associated Dirichlet energy indeed decreases slower, though eventually still diminishes as the number of layers increases.
Graph rewiring.
<cit.> and GPR-GNN (<cit.>, Appendix Table <ref>) rewire the graph using personalized page rank algorithm <cit.> and generalized page rank on graphs, respectively.
Similar to JKNet <cit.>, they mitigate over-smoothing by allowing direct message passing between faraway nodes.
Constant-energy methods.
<cit.> constrain the pair-wise distance or Dirichlet energy among graphs to be constant.
Nevertheless, the non-decreasing energy does not necessarily translate to better performance, as they sometimes come with the sacrifice of expressiveness, as argued in <cit.>.
Residual GNNs.
Residual connection <cit.> can naturally be added to the GNN architecture, such as GCNII (<cit.>, Table <ref>), to restraint activation to be similar to the input to allow deep networks.
They however can make the model less robust to perturbations in the input.
In sum, these works have similar, if not compromised expressiveness compared to a barebone GCN.









§ ARCHITECTURE: COMBINING TOPOLOGIC AND SEMANTIC TRAJECTORIES OF WALKS




  
Random walks on graphs. 
An unbiased random walk w on a graph 𝒢 is a sequence of nodes w = (v_0, v_1, …) with landing probability:

    P(v_j | (v_0, …, v_i-1)) = 
    1[(v_i, v_j) ∈ℰ_𝒢] / D(v_i),

where D(v_i) = ∑ A_ij is the degree of the node v_i.
Walks originating from or terminating at any given node v can thus be easily generated using this Markov chain.
We record the trajectory of embeddings associated with the walk as
ω_x(w) = (𝐗_i) = (𝐗_0, 𝐗_1, …, 𝐗_l).
In this paper, we only consider finite-long l-step random walk |w| = l ∈ℤ^+.

In our implementation, the random walks are sampled ad hoc during each training and inference step directly on GPU using Deep Graph Library <cit.> (see Appendix Section <ref>).

Moreover, the walk considered here is not necessarily a path, as repeated traversal of the same node v_i = v_j, i ≠ j is not only permitted, but also crucial to effective topological representation, as discussed below.




  
Anonymous experiment.
We use a function describing the topological environment of a walk, termed anonymous experiment <cit.>,  ω_u(w): ℝ^l→ℝ^l  that records the first unique occurrence of a node in a walk (Appendix Algorithm <ref>).
To put it plainly, we label a node as the number of unique nodes insofar traversed in a walk if the node has not been traversed, and reuse the label otherwise.
Practically, this can be implemented using any tensor-accelerating framework in one line ( is the node sequence of a walk) and trivially parallelized [
Note that this particular implementation introduces an intermediate 𝒪(l^2) complexity term, though it is empirically faster than the linear-complexity naive implementation, since only integer indices are involved and thus the footprint is negligible.

]:




  
Unifying memory: combining semantic and topological representations.
Given any walk w, we now have two sequences ω_x(w) and ω_u(w) describing the semantic and topological (as we shall show in the following sections) features of the walk.
We project such sequential representations onto a latent dimension to combine them (illustrated in Table <ref>):

    h(w) = f(ϕ_x(ω_x(w)), ϕ_u(ω_u(w))),

where ϕ_x: ℝ^l × D→ℝ^D_x maps the sequence of semantic embeddings generated by a l-step walk to a fixed D_x-dimensional latent space,
ϕ_u: ℝ^l→ℝ^D_u maps the indicies sequence to another latent space D_u,
and f: ℝ^D_x⊕ℝ^D_u→ℝ^D combines them.
We call Equation <ref> the unifying memory of a random walk.
Subsequently, the node representations can also be formed as the average representations of l-step (l being a hyperparameter) walks terminating (for the sake of gradient aggregation) at that node:

    ψ(v) = ∑_{w}, |w|=l, w_l = v p(w) h(w),

which can be stochastically sampled with unbiased Monte Carlo gradient and used in downstream tasks as such node classification and regression.
We note that this is the only time we perform SUM or MEAN operations.
Unlike other GNNs incorporating random walk-generated features (which are sometimes still convolutional and iterative), we do not iteratively pool representations within local neighborhoods.
The likelihood of the data written as:

    P(y|𝒢, 𝐗) = ∑_{w}, |w|=l, w_l = v p(w) p(y|𝒢, 𝐗, w)

The node representation can be summed 

    Ψ(𝒢) = ∑_v ∈𝒱⊆𝒢ψ(v)

to form global representations for graph classification and regression.
We call ψ in Equation <ref> and Ψ in Equation <ref> the node and graph output representations of RUM.



  
Layer choices. 
Obvious choices to model f include a feed-forward neural network after concatenation, and ϕ_x, ϕ_u recurrent neural networks (RNNs).

This implies that, different from most convolutional GNNs, parameter sharing is natural and the number of parameters is going to stay constant as the model incorporates a larger neighborhood.
Compared to dot product-based, transformer-like modules <cit.>, RNNs not only have linear complexity (see detailed discussions below) w.r.t. the sequence length but also naturally encodes the inductive bias that nodes closer to the origin have stronger impact on the representation. 

The gated recurrent unit (GRU)<cit.> variant is used everywhere in this paper.
Additional regularizations are described in Appendix Section <ref>.



  
Runtime complexity. 
To generate random walks for one node has the runtime complexity of 𝒪(1), and for a graph 𝒪(| 𝒱 |), where |𝒱| is the number of nodes in a graph 𝒢 = {𝒱, ℰ}.
To calculate the anonymous experiment, as shown in Appendix Algorithm <ref>, has 𝒪(1) complexity (also see Footnote <ref>).
If we use linear-complexity models, such as RNNs, to model ϕ_x, ϕ_u, the overall complexity is 𝒪(|𝒱|lkD) where l is the length of the random walk, k the samples used to estimate Equation <ref>, and D the latent size of the model (assumed uniform).
Note that different from convolutional GNNs, RUM does not depend on the number of edges | ℰ | (which is usually much larger than |𝒱|) for runtime complexity, and is therefore agnostic to the sparsity of the graph.
See Figure <ref> for an empirical verification of the time complexity.
In Appendix Table <ref>, we show, on a large graph, the overhead introduced by generating random walks and computing anonymous experiments accounts for roughly 1/1500 of the memory footprint and 1/8 of the wall time.



  
Mini-batches.

RUM is naturally compatible with mini-batching.
For convolutional GNNs, large graphs that do not fit into the GPU memory have traditionally been a challenge, as all neighbors are required to be present and boundary conditions are hard to define <cit.>.
RUM, on the other hand, can be inherently applied on subsets of nodes of a large graph without any alteration in the algorithm—the random walks can be generated on a per-node basis, and the FOR loop in Algorithm <ref> can be executed sequentially, in parallel, or on subsets.

Empirically, in Appendix Table <ref>, RUM can be naturally scaled to huge graphs.



§ THEORY: RUM AS A JOINT REMEDY.


We have insofar designed a new graph learning framework—convolution-free graph neural networks (GNNs) that cleanly represent the semantic (ω_x) and topological (ω_u) features of graph-structured data before unifying them.
First, we state that RUM is permutation equivariant,

For any permutation matrix P, we have

    P𝐗_v(𝒢) = 𝐗_v(P(𝒢)),


which sets the foundation for the data-efficient modeling of graphs.
Next, we theoretically demonstrate that this formulation jointly remedies the common pathologies of the convolution-based GNNs by showing that:
(a) the topological representation ω_u is more expressive than convolutional-GNNs in distinguishing non-isomorphic graphs;
(b) the semantic representation ω_x no longer suffers from over-smoothing and over-squashing.



 §.§ RUM is more expressive than convolutional GNNs.

For the sake of theoretical arguments in this section, we assume that in Equation <ref>:


ϕ_x, ϕ_u, f are universal and injective.



Graph 𝒢 discussed in this section is always connected, unweighted, and undirected.

Assumption <ref> is easy to satisfy for feed-forward assuneural networks <cit.> and RNNs <cit.>. 
Composing injective functions, we remark that h(w) is also injective w.r.t. ω_x(w) and ω_u(w);
despite of Assumption <ref>, our analysis can be extended to disjointed graphs by restricting the analysis to the connected regions in a graph.
Under such assumptions, we show, in Remark <ref> (deferred to the Appendix), that ψ is injective, meaning that nodes with different random walks will have different distributions of representations ψ(v_1) ≠ψ(v_2).
We also refer the readers to the Theorem 1 in <cit.> for a discussion on the representation power of anonymous experiments on unlabelled graphs.
Combining with the semantic representations and promoting the argument from a node level to a graph level, we arrive at:


Up to the Reconstruction Conjecture <cit.>, RUM with sufficiently long l-step random walks can distinguish non-isomorphic graphs satisfying Assumption <ref>.

The main idea of the proof of Theorem <ref> (in Appendix Section <ref>) involves explicitly enumerating all possible non-isomorphic structures for graphs with 3 nodes and showing, by induction, that if the theorem stands for graph of N-1 size it also holds for N-sized graphs.
We also show in Appendix Section <ref> that a number of key graph properties such as cycle size (Example <ref>) and radius (Example <ref>) that convolutional GNNs struggle <cit.> to learn can be analytically expressed using ω_u.
As these are solely functions of ω_x, they can be approximated arbitrarily well by universal approximators.
These examples are special cases of the finding that RUM is stricly more expressive than Weisfeiler-Lehman isomorphism test <cit.>:

Two graphs with N nodes 𝒢_1, 𝒢_2 labeled as non-isomorphic by Weisfeiler-Lehman isomorphism test is the necessary, but not sufficient condition that the representation resulting from RUM with walk length N+1 are also different.

    Ψ(𝒢_1) ≠Ψ(𝒢_2)


Thus, due to <cit.>, RUM is also more expressive than convolutional GNNs in distinguishing non-isomorphic graphs.
This also confirms the intuition that RUM with longer walks are more expressive.
(See Figure <ref> on an empirical evaluation.)
The proof of this corollary is straightforward to sketch—even if we only employ the embedding trajectory ω_x, it would have the effect of performing the function ϕ_x in Equation <ref> on each traversal of the WL expanded trees.



 §.§ RUM alleviates over-smoothing and over-squashing


Over-smoothing refers to the phenomenon where the node dissimilarity (e.g., measured by Dirichlet energy in Equation <ref>) decreases exponentially and approaches zero with the repeated rounds of message passing.

<cit.> relates Dirichlet energy directly with the convolutional operator:
*dirichlet-convLemma 3.1 from <cit.>


    ℰ((1 - Δ̃)𝐗)
    ≤
    (1 - λ)^2 ℰ(𝐗)

where λ is the smallest non-zero eigenvalue of Δ̃, the normalized Laplacian of a graph.

Free of convolution operators, it seems only natural that RUM does not suffer from this symptom (Figure <ref>).
We now formalize this intuition by first restricting ourselves to a class of non-contractive mappings for f in Definition <ref>.


A map f is non-contractive on region Ω if ∃α∈ [1, +∞) such that |f(x) - f(y)| ≥α  || x - y ||, ∀ x, y ∈Ω.

A line of fruitful research has been focusing on designing non-contractive RNNs <cit.>, and to be non-contractive is intimately linked with desirable properties such as preserving the long-range information content and non-vanishing gradients.
From this definition, it is easy to show that, for each sampled random walk in Equation <ref>, the Dirichlet energy is greater than its input.
One only needs to verify that the integration in Equation <ref> does not change this to arrive at:


If ϕ_x, f are non-contractive w.r.t. all elements in the sequence, the expected Dirichlet energy of the corresponding RUM node representation in Equation <ref> is greater than its initial value

    E(ℰ(ψ(𝐗))) ≥ℰ(𝐗).



This implies that the expectation of Dirichlet energy does not diminish even when l → +∞, as it is bounded by the Dirichlet energy of the initial node representation, which is consistent with the trend shown in Figure <ref>, although the GRU is used out-of-box without constraining it to be explicitly non-contractive.




  
RUM alleviates over-squashing is deferred to Appenxix Section <ref>, where we verify that the inter-node Jacobian | ∂𝐗_v ^(l+1) / ∂𝐗_u^(0)| decays slower as the distance between u, v grows vis-à-vis the convolutional counterparts.
Briefly, although RUM does not address the information bottleneck with exponentially growing reception field (the 1 / (Â^l+1)_uv term in Equation <ref>), it nevertheless can have a non-vanishing (nor exploding) gradient from the aggregation function (| ∇ϕ_x |). 



§ EXPERIMENTS


On a wide array of real-world node- and graph-level tasks, we benchmark the performance of RUM to show its utility in social and physical modeling.
Next, to thoroughly examine the performance of RUM, we challenge it with carefully designed illustrative experiments.
Specifically, we ask the following questions in this section, with Q1, Q2, and Q3 already theoretically answered in Section <ref>:
Q1:
Is RUM more expressive than convolutional GNNs?
Q2:
Does RUM alleviate over-smoothing?
Q3:
Does RUM alleviate over-squashing?
Q4:
Is RUM slower with convolutional GNNs?
Q5:
Is RUM robust?
Q6:
How does RUM scale up to huge graphs?
Q7:
What components of RUM are contributing most to the performance of RUM?



  
Real-world benchmark performance.
For node classification, we benchmark our model on the popular Planetoid citation datasets <cit.>, as well as the coauthor <cit.> and co-purchase <cit.> datasets common in social modeling.
Additionally, we hypothesize that RUM, without the smoothing operator, will perform competitively on heterophilic datasets <cit.>—we test this hypothesis.
For graph classification, we benchmark on the popular TU dataset <cit.>.
We also test the graph regression performance on molecular datasets in MoleculeNet <cit.> and Open Graph Benchmark <cit.>.
In sum, RUM almost always outperforms, is within the standard deviation of, the state-of-the-art architectures, as shown in Tables <ref>, <ref>, <ref>, <ref>, as well as in Tables <ref>, <ref>, <ref>, <ref> moved to the Appendix due to space constraint.








  
On sparsity: the subpar performance on the Computer dataset.
The most noticeable exception to the good performance of RUM is that on the Computer co-purchase <cit.> dataset, where RUM is outperformed even by GCN and GAT.
This dataset is very dense with an average node degree (| ℰ | / | 𝒱|) of 18.36, the highest among all datasets used in this paper.
As the variance of the node embedding (Equation <ref>) scales with the average degree, we hypothesize that dense graphs with very high average node degrees would have high-variance representations from RUM.

On the other hand, RUM outperforms all models surveyed in two large-scale benchmark studies on molecular learning, GAUCHE <cit.> and MoleculeNet <cit.>.
The atoms in the molecules always have a degree of 2 ∼ 4 with intricate subgraph structures like small rings.
This suggests the utility of unifying memory in chemical graphs and furthermore chemical and physical modeling.



























  
Graph isomorphism testing (Q1).
Having illustrated in Section <ref> that RUM can distinguish non-isomorphic graphs, we experimentally test this insight on the popular Circular Skip Link dataset <cit.>.
Containing 4-regular graph with edges connected to form a cycle and containing skip-links between nodes, this dataset is artificially synthesized in <cit.> to create an especially challenging task for GNNs.
As shown in Appendix Table <ref>, all convolutional GNNs fail to perform better than a constant baseline (there are 10 classes uniformly distributed).
3WLGNN <cit.>, a higher-order GNN of at least 𝒪(2) complexity that operates on explicitly enumerated triplets of graphs, can distinguish these highly similar 4-regular graphs by comparing subgraphs.
RUM, with 𝒪(| 𝒩 | ) linear complexity, achieves similarly high accuracy.
One can think of RUM as a stochastic approximation of the higher-order GNN, with all of its explicitly enumerated subgraphs being identified by RUM with a probability that decreases with the complexity of the subgraph.




  
Effects of walk lengths and number of samples on performance (Q2, Q3).
Having studied the relationship between inference speed and the walk lengths and number of samples, we furthermore study its impact on performance.
Using Cora <cit.> citation graph and vary the walk lengths and number of samples from 1 to 9, where the performance of RUM improves as more samples are taken and longer walks are employed, though more than 4 samples and walks longer than L > 4 yield qualitatively satisfactory results;
this empirical finding has guided our hyperparameter design.
In Figure <ref>, we also compare the Dirichlet energy of RUM-generated layer representations with not only baselines GCN <cit.> and GAT <cit.>, but also strategies to alleviate over-smoothing discussed in Section <ref>, namely residual connection and stochastic regularization <cit.>, and show that when L gets large, only RUM can maintain Dirichlet energy.
Traditionally, since <cit.> (see its Figure 5 compared to Figure <ref>), the best performance on Cora graph was found with 2 or 3 message-passing rounds, since mostly local interactions are dominating the classification, and more rounds of message-passing almost always lead to worse performance.
As theoretically demonstrated in Section <ref>, RUM is not as affected by these symptoms.
Thus, RUM is especially appropriate for modeling long-range interactions in graphs without sacrificing local representation power.



  
Long-range neighborhood matching (Q3).
To verify that RUM indeed alleviates over-squashing (Section <ref>), in Figure <ref>, we adopt the tree neighborhood matching synthetic task in <cit.> where binary tree graphs are proposed with the label and the attributes of the root matching a faraway leave.
The full discussion is moved to the Appendix Section <ref>.



  
Speed (Q4).
Though both have linear runtime complexity (See Section <ref>), intuitively, it might seem that RUM would be slower than convolution-based GCN due to the requirement of multiple random walk samples.
This is indeed true for CPU implementations shown in Figure <ref> left.
When equipped with GPUs (specs in Appendix Section <ref>), however, RUM is significantly faster than even the simplest convolutional GNN—GCN <cit.>.
It is worth mentioning that the GCN and GAT <cit.> results were harvested using the heavily optimized Deep Graph Library <cit.> sparse implementation whereas RUM is implemented naïvely in PyTorch <cit.>, though the popular GRU component <cit.> have already undergone CUDA-specific optimization.






  
Robustness to attacks (Q5).
With the stochasticity afforded by the random walks, it is natural to suspect RUM to be robust.
We adopt the robustness test from <cit.> and attack by randomly adding fake edges to the Cora <cit.> graph and record the test set performance in Figure <ref>.
Indeed, RUM is much more robust than traditional convolutional GNNs including GCN <cit.> and GAT <cit.> and is even slightly more robust than the convolutional GNN specially designed for robustness <cit.>, with the performance only decreased less than 10% with 10% fake edges added.



  
Scaling to huge graphs (Q6) and Ablation study (Q7)
are deferred to Appendix Section <ref>.



§ CONCLUSIONS


We design an innovative GNN that uses an RNN to unify the semantic and topological representations along stochastically sampled random walks, termed random walk with unifying memory (RUM) neural networks.
Free of the convolutional operators, our methodology does not suffer from symptoms characteristic of Laplacian smoothing, including limited expressiveness, over-smoothing, and over-squashing.
Most notably, our method is more expressive than the Weisfeiler-Lehman isomorphism test and can distinguish all non-isomorphic graphs up to the reconstruction conjecture.
Thus, it is more expressive than all of convolutional GNNs equivalent to the WL test, as we demonstrate theoretically in Section <ref>.
RUM is significantly faster on GPUs than even the simplest convolutional GNNs (Section <ref>) and shows superior performance across a plethora of node- and graph-level benchmarks.



  
Limitations. 
Very dense graphs.
As evidenced by the underwhelming performance of the Computer <cit.> dataset and discussed in Section <ref>, RUM might suffer from high variance with very dense graphs (average degree over 15).
Tottering. 
In our implementation, we have not ruled out the 2-cycles from the random walks, as that would require specialized implementation for walk generation.
This, however, would reduce the average information content in a fixed-length random walk (known as tottering <cit.>).
We plan to characterize the effect of excluding these walks.
Biased walk. 
Here, we have only considered unbiased random walk, whereas biased random walk might display more intriguing properties as they tend to explore faraway neighborhoods more effectively <cit.>.
Directed graphs. 
Since we have only developed RUM for undirected graph (random walk up to a random length is not guaranteed to exist for directed graphs), we would have to symmetrize the graph to perform on directed graphs (such as the heterophilic datasets <cit.>);
this create additional information loss and complexity.



  
Future directions.
Theoretical. 
We plan to expand our theoretical framework to account for the change in layer width and depth to derive analytical estimates for realizing key graph properties.
Applications. 
Random walks are intrinsically applicable to uncertainty-aware learning.
We plan to incorporate the uncertainty naturally afforded by the model to design active learning models.
On the other hand, the physics-based graph modeling field is also heavily dominated by convolutional GNNs.
Inspired by the superior performance of RUM on chemical datasets, we plan to apply our method to the equivariant modeling of n-body physical systems <cit.>.



  
Impact statement. 
We here present a powerful, robust, and efficient learning algorithm on graphs.
Used appropriately, this algorithm might advance the modeling of social <cit.> and physical <cit.> systems, which can oftentimes modeled as graphs.
As with all graph machine learning methods, negative implications may be possible if used in the design of explosives, toxins, chemical weapons, and overly addictive recreational narcotics.





  
Disclosures. 
YW has limited financial interests in Flagship Pioneering, Inc. and its subsidiaries.



  
Funding. 
YW acknowledges support from the Schmidt Science Fellowship, in partnership with the Rhodes Trust, and the Simons Center for Computational Physical Chemistry at New York University.





§ EXPERIMENTAL DETAILS




  
Code availability.
All architectures, as well as scripts to execute the experiment, are distributed open-source under MIT license at <https://anonymous.4open.science/r/rum-834D/>.
Core dependencies of our package include PyTorch <cit.> and Deep Graph Library <cit.>.



  
Hyperparameters.
All models are optimized using Adam <cit.> optimizer and SiLU <cit.> activation functions.
4 random walk samples are drawn everywhere unless specified.
Other hyperparameters—learning rate (10^-5∼ 10^-2), hidden dimension (32 ∼ 64), L2 regularization strength (10^-8∼ 10^-2), walk length (3∼ 16), temperature for ℒ_𝚌𝚘𝚗𝚜𝚒𝚜𝚝𝚎𝚗𝚌𝚢 (0 ∼ 1), coefficient for ℒ_𝚌𝚘𝚗𝚜𝚒𝚜𝚝𝚎𝚗𝚌𝚢 (0 ∼ 1), coefficient for ℒ_𝚜𝚎𝚕𝚏, and dropout probability—are tuned using the Ray platform <cit.> with the default Ax <cit.> search algorithm with 1000 trails or 24 hours tuning budget on a Nvidia A100® GPU.



§ ADDITIONAL TECHNICAL DETAILS




 §.§ Self-supervised regularization. 


The stochasticity encoded in our model naturally affords it with some level of regularization.
Apart from using the consistency loss (ℒ_𝚌𝚘𝚗𝚜𝚒𝚜𝚝𝚎𝚗𝚌𝚢) used in <cit.> for classifications, we further regularize the model by using the RNNs in ϕ_x to predict the semantic representation of the next node on the walk given ω_u and jointly maximize this likelihood:

    ω̂_x_i+1 = g({ω_x_1, ω_x_2, …, ω_x_i}, ω_u | θ);
    ℒ_𝚜𝚎𝚕𝚏(θ) = -log P(ω̂_x_i+1 | θ),

where g(· | θ) is modeled as the sequence output of the RNN ϕ_x in Equation <ref>.
The total loss is modeled as a combination of:

    ℒ(θ) =  
    -log P(y|𝒢, 𝐗, θ)
    + ℒ_𝚜𝚎𝚕𝚏
    + ℒ_𝚌𝚘𝚗𝚜𝚒𝚜𝚝𝚎𝚗𝚌𝚢




 §.§ RUM attenuates over-squashing


Similarly, we can show that if the composing neural networks defy the vanishing gradient problem (w.r.t. the input) <cit.>, the sensitivity analysis in Equation <ref> <cit.> has a lower bound for RUM.


If ϕ_x, f have lower-bounded derivatives, the inter-node Jacobian for nodes u, v separated by a shortest path of length l, RUM with walk length l also has a lower bound:

    | ∂𝐗_v ^(l)/∂𝐗_u^(0) |
    ≥
    | ∇ϕ_x | |∇ f| (Â^l)_uv,

where Â_ij = A_ij / ∑_jA_ij is the degree-normalized adjacency matrix. 

Like the upper bound in Equation <ref>, this lower bound is also controlled by the power of the (normalzied) adjacency matrix, albeit the absence of self-loop will result in a slightly looser bottleneck.
The term (Â^l)_uv corresponds to the probability of the shortest path among all possible walks as a product of inverse node degrees (see Equation <ref>).
There is no denying that the lower bound is still controlled by the power of the adjacency matrix, which corresponds to the exponentially growing receptive fields.
One can also argue that, without prior knowledge, the contribution of the sensitivity analysis by the power of the graph adjacency matrix can never be alleviated, since there are always roughly 1 / (Â^l+1)_uv (assuming uniform node degree) structurally equivalent nodes.
Nevertheless, since ϕ_x is not necessarily an iterative function, we alleviate the over-squashing problem by eliminating the power of the update function gradient term.

Now, we plug in the layer choices of ϕ_x—a GRU <cit.> unit.
Its success, just like that of long short-term memory (LSTM) <cit.>, can be attributed to the near-linear functional form of the long-range gradient.
The term |∇ϕ_x| is controlled by a sequence of sigmoidal update gates, which can be optimized to approach 1 (fully open).
If we ignore the gradient contribution of 𝐗_u^0 to these gates, the non-linear activation function has only been applied exactly once on 𝐗_u^0;
therefore, the gradient | ∂𝐗_v ^(l+1) / ∂𝐗_u^(0)| is neither rapidly vanishing nor exploding.



 §.§ Long-range neighborhood matching (Q3).


Once identifying the target leaf, this task seems trivial;
nonetheless, this piece of information needs to be passed through layers of aggregation and non-linear update and is usually lost in the convolution.
Since, on this binary tree, the receptive field grows exponentially, <cit.> argues that there is a theoretical lower boundary for the layer width D for the convolutional GNN to be able to encode all possible combinations of leaves, which is 2^32D for single-precision floating point ().
This corresponds to the structural Â^l term in Equation <ref> and Equation <ref>.
Evidently, when D=32, as is the adopted experimental setting in  Figure <ref>, this limit is far from being hit.
So we hypothesize that the reason why convolutional GNNs cannot overfit the training set is because of the limitation of the functional forms, which are remedied by RUM, which shows 100% accuracy up to tree depth or problem radius r=5, and a relatively moderate decrease afterward.

Note that when the problem radius exceeds r=7, RUM's performance is not significantly different from the convolutional counterparts.




 §.§ Scalaing to large graphs (Q6).

In Appendix Table <ref>, we apply RUM on an ultra-large graph , that cannot fit easily on a single GPU, and compare RUM with architectures specifically designed for large graphs <cit.>.



 §.§ Ablation study (Q7). 



In Table <ref>, we conduct a brief ablation study where we delete, one by one, the components introduced in Section <ref>.
ω_u = 0 and ω_x = 0 refer to the deletion of the topological and semantic representations of walks, respectively.
Neglecting topological information results in a moderate decrease in performance, whereas neglecting semantic representation is more detrimental.
The ω_u = 0 also resembles <cit.> albeit with different walk-wise aggregation.
This offers a qualitative comparison between our work and <cit.> as no source code was released for this package so no rigorous comparison was possible.
We also see that the regularization methods are helpful to the performance, with self-supervision being more crucial.
We attribute this effect to firstly the relative simplicity of the Cora classification task, and secondly the flexibility of the (overparametrized) RNNs.




§ ADDITIONAL RESULTS






For two nodes v_1, v_2 with distribution of random walks terminating at v_1, v_2 not equal in distribution p(ω_u(w_1)) ≠ p(ω_u(w_2)) or p(ω_x(w_1)) ≠ p(ω_x(w_2)), the node representations in Equation <ref> are also different
ψ(v_1) ≠ψ(v_2).

One way to construct h function in Equation <ref> is to have h(w) positive only where p(ω_u(w_1)) > p(ω_u(w_2));
the same thing can be argued for ω_x.
In other words, one only needs to prove two walks terminating at two nodes p(ω_u(w_1)) ≠ p(ω_u(w_2)) or p(ω_x(w_1)) ≠ p(ω_x(w_2)) are not equal in distribution to verify that RUM can distinguish two nodes.
Conversely, we can also show that p(ω_u(w_1)) = p(ω_u(w_2)) implies that v_1, v_2 are isomorphic without labels—




we refer the readers to the Theorem 1 in <cit.> for a discussion on reconstructing unlabelled graphs using anonymous experiments.












§ MISSING MATHEMATICAL ARGUMENTS.



 §.§ Examples of Theorem <ref>


[Cycle detection.]

A k-cycle 𝒞_k is a subgraph of 𝒢 consisting of k nodes, each with degree two.
The existence of k-cycle can be determined by:

    1(𝒞_k ⊆𝒢)
    = 1[P(
    ω_x_j+1 = ω_x_0,
    ω_x_i≠ω_x_j, ∀ i < j
    )>0]




[Diameter.]

The diameter, δ_𝒢 of graph 𝒢 which equals the length of the longest shortes path in 𝒢, can be expressed as

    δ_𝒢 = 
    l=| ω_x |, ω_x_i≠ω_x_j, ∀ i ≠ jargmax
    | ω_x |







 §.§ Proof of Theorem <ref>



First, we enumerate all possible unlabelled graphs with three nodes satisfying Assumption <ref>—one with two edges, one with three edges.
(Note that there is only one non-isomorphic graph with two nodes.)
Now we consider random walks of length l=3, where

    P(ω_u_3 = ω_u_0) > 0

only stands for the graph with three edges, but not with two edges, just like Example <ref>.

Furthermore, we can also distinguish between the 2-degree node and the 1-degree node in the graph with 2 edges and 3 nodes simply by verifying that

    P(ω_u_1≠ω_u_3) > 0

only stands when v_2 is the 2-degree node.

Moving on to the labeled 3-node graph case, we can reduce the problem to investigate whether RUM can distinguish 3-node graphs that are isomorphic when unlabeled, but non-isomorphic when labeled.
For the three-edged graph, ω_x uniformly samples the labels of three nodes.
For the two-edged graph, suppose the node labels are A, B, C, and we start from the node with 2 degrees B (with nodes bearing A and C labels locally, structurally isomorphic), 

    P(B|ω_x(w_t)) = 
    1, t = 2n, n∈ℕ,
    
    0, t = 2n + 1, n ∈ℕ
    
    P(A|ω_x(w_t)) = P(C|ω_x(w_t)) = 
    0, t = 2n, n∈ℕ,
    
    1/2, t = 2n + 1, n ∈ℕ

If graphs have the same ω_x, they have the same B and the same or swapped A, C.
As such, we have proven Theorem <ref> for graphs with 3 nodes.

Now we prove that Theorem <ref> stand for graphs of N nodes, they also stand for graphs of N+1 nodes, the Reconstruction Conjecture <cit.>

Suppose we have two non-isomorphic graphs with N+1 nodes 𝒢_1 and 𝒢_2 with the same RUM embedding Ψ_1 = Ψ_2.
We enumerate all N + 1 subgraphs with each node deleted for each of these two graphs.
By the Reconstruction Conjecture, at least one pair of subgraphs are non-isomorphic.
For this pair, suppose the deleted vertex is v (ruling out the trivial case where the label or connectivity of v is different for these two graphs), and two remaining subgraphs 𝒢_1^\ v, 𝒢_2^\ v;
since ϕ_x, ϕ_u, f are injective, Ψ_1 = Ψ_2 implies ω_u, ω_x are equal in distribution for 𝒢_1, 𝒢_2.
As such, the walk distribution

    P(ω_x(w), ω_u(w) | w_0 = v, w_i ≠ v, i > 0) 
    
    = P(ω_x(w_0), ω_u(w_0) | w_0 = v) P(ω_x(w_1…), ω_u(w_1…) | ω_x(w_1…), ω_u(w_1…), w_i ≠ v, i > 0)

are also equal in distribution for 𝒢_1, 𝒢_2.

If there is a link between v and the nodes in 𝒢_1, 𝒢_2, we assign a new label to contain both the old label and the connection.
As such, if the second term is not equal in distribution for 𝒢_1, 𝒢_2, we would have 

    Ψ(𝒢_1^\ v) ≠Ψ𝒢_2^\ v),

which is in conflict with the assumption that Theorem <ref> stands for graphs with N nodes.




 §.§ Proof of Lemma <ref>


By the definition of Dirichlet energy (Equation <ref>) and non-contractive mappings (Definition <ref>),

    ℰ(ψ(𝐗))
            = 1/N∑_u, v ∈ℰ_𝒢 || ψ(ω_x(𝐗_u)) - ψ(ω_x(𝐗_v)) || ^2
        = 1/N∑_u, v ∈ℰ_𝒢 || f(ϕ_x(ω_x(𝐗_u))) - f(ϕ_x(ω_x(𝐗_v))) || ^2
       ≥1/N∑_u, v ∈ℰ_𝒢 || ϕ_x(ω_x(𝐗_u)) - ϕ_x(ω_x(𝐗_v)) || ^2
        = 1/N∑_u, v ∈ℰ_𝒢 || ϕ_x(∑_{w}, |w|=l, w_l = u p(w)𝐗_u, u ∈ w) - ϕ_x(∑_{w}, |w|=l, w_l = v p(w)𝐗_v, v ∈ w) || ^2
       ≥1/N∑_u, v ∈ℰ_𝒢 || (∑_{w}, |w|=l, w_l = u p(w)𝐗_u, u ∈ w) - (∑_{w}, |w|=l, w_l = v p(w)𝐗_v, v ∈ w) || ^2
       ≥ℰ(𝐗)

The last inequality is due to the fact that ϕ_x is non-contractive w.r.t. all, and therefore, the last element of the sequence, which are always 𝐗_u and 𝐗_v.

































































































































    

























    



































































 Paper Title: Memorization Capacity for Additive Fine-Tuning with Small ReLU Networks.


An alternative way to decipher the nature of the doubly charmed tetraquark T_cc(3875)^+: its antiparticle
T_c̅c̅(3875)^- photoproduction off nuclei near threshold
    E. Ya. Paryev

Institute for Nuclear Research of the Russian Academy of Sciences

Moscow, Russia
    August 5, 2024
==================================================================================================================================================================



§ ABSTRACT 
  Fine-tuning large pre-trained models is a common practice in machine learning applications, yet its mathematical analysis remains largely unexplored.
In this paper, we study fine-tuning through the lens of memorization capacity. Our new measure, the Fine-Tuning Capacity (FTC), is defined as the maximum number of samples a neural network can fine-tune, or equivalently, as the minimum number of neurons (m) needed to arbitrarily change N labels among K samples considered in the fine-tuning process. 
In essence, FTC extends the memorization capacity concept to the fine-tuning scenario. We analyze FTC for the additive fine-tuning scenario where the fine-tuned network is defined as the summation of the frozen pre-trained network f and a neural network g (with m neurons) designed for fine-tuning. 
When g is a ReLU network with either 2 or 3 layers, we obtain tight upper and lower bounds on FTC; 
we show that N samples can be fine-tuned with m=Θ(N) neurons for 2-layer networks, and with m=Θ(√(N)) neurons for 3-layer networks, no matter how large K is.  
Our results recover the known memorization capacity results when N = K as a special case. 






§ INTRODUCTION



As a branch of machine learning theory, the expressive power of neural networks is investigated for several decades. 
By using the concept of universal approximation, it is shown that neural networks can approximate a large classes of functions, either in the depth-bounded scenarios <cit.> or width-bounded scenarios <cit.>. 
Another line of research focused on the memorization capacity of neural networks <cit.>, exploring the capability of neural networks for memorizing finite samples. 

Meanwhile, with the advent of large language models <cit.> and foundation models <cit.>, the paradigm of pre-training followed by fine-tuning is dominating the machine learning communities. Various empirical results show that a gigantic model pre-trained on large amount of data can be easily fine-tuned to perform well on downstream tasks, given only a small amount of additional data for the target task. 
Compared with the extensive empirical results, mathematical analysis on fine-tuning large pre-trained models remains largely unexplored.

In this paper, we take the first step in understanding the fine-tunability of pre-trained networks through the lens of memorization capacity. 
We focus on the scenario where we fine-tune a pre-trained neural network f on dataset D = {(_i, y_i)}_i=1^K with K samples; here, _i ∈^d and y_i ∈ for all i ∈ [K] where [K] ={1,2,⋯, K}, and we assume _i _j for all i  j.
Let 
    T := { i ∈ [K]: f(_i)  y_i }
be the set of indices of samples that the pre-trained network f does not fit. The cardinality of this set is denoted by N:= |T| ≤ K.
In other words, 

    f(_i)    = y_i

holds for all i ∈ [K] ∖ T, while not guaranteed for i ∈ T.
Our aim is to add a neural network g_θ (parameterized by θ) to the pre-trained network f in a way that the fine-tuned network f + g_θ satisfies

    (f + g_θ) (_i) = y_i,   ∀ i ∈ [K].

See Fig. <ref> for the visualization of the additive fine-tuning scenario we focus on.
This scenario is motivated by recently proposed additive fine-tuning methods  <cit.>, and especially, the side-tuning  <cit.> where a side network g_θ is added to the pre-trained network f. Since our model does not cover other popular fine-tuning methods including LoRA  <cit.>, extending our theoretical results to such popular methods is remained as a future work.
Under such setting, we define the fine-tuning capacity (FTC) of a neural network g_θ as below. 



    The fine-tuning capacity of a given neural network g_θ is the maximum number N satisfying the following property: for all _i ∈^d, y_i ∈, for all T ⊆ [K] satisfying | T | = N, and for any choices of function f satisfying f(_i) = y_i for all i ∈ [K]∖ T, 
    there exists parameter θ such that (f + g_θ)(_i) = y_i all i ∈ [K].











Under such a setting, we establish the upper/lower bounds on FTC, when g is 2-layer ReLU network or 3-layer ReLU network. 
Our main contributions are summarized below:

    
  * We define a new metric called Fine-Tuning Capacity (FTC), which measures the maximum number of samples N^⋆ a neural network with m neurons can fine-tune. Equivalently, we define the minimum number of neurons m^⋆ needed to arbitrarily change N labels among K. FTC can be considered as an extension of memorization capacity, tailored for the fine-tuning scenario. 
    
  * For 2-layer ReLU networks, we establish tight upper and lower bounds on m^⋆, in Theorem <ref>. 
    The upper bound is obtained by a novel neural network construction for fine-tuning.
    Our construction requires less number of neurons than conventional constructions developed in the memorization capacity literature when K ≥ 3N+2. 
    By using our bounds on m^⋆, we also provide an equivalent statement in Corollary <ref>, showing the tight bounds on the fine-tuning capacity N^⋆ .
    
  * For 3-layer ReLU networks, we obtain tight upper and lower bounds on m^⋆ in Theorem <ref>. Our results imply that 
    N samples can be fine-tuned with m = Θ(√(N)) neurons without any dependence on K.
    We also provide an equivalent statement in Corollary <ref>, showing the tight upper and lower bounds on N^⋆ .










§ RELATED WORKS






  
Fine-Tuning

Various methods for efficient fine-tuning are introduced in recent years <cit.>, which fine-tune only a small part of pre-trained models to adapt it for target tasks. 
There are some mathematical analysis on fine-tuning <cit.> or more broadly on transfer learning <cit.>, but none of them analyzed the fine-tunability of large pre-trained models using the lens of memorization capacity.




  
Memorization

One concept relevant to FTC is memorization capacity which measures the ability of memorizing given feature-label pairs {(_i, y_i)}_i=1^K. 
Finding the bounds on the memorization capacity is considered in recent works on various networks <cit.>. 
The effect of memorization in large language models is explored in recent works, both for pre-training <cit.> and for fine-tuning <cit.>.






§ FINE-TUNING CAPACITY








Note that a notion of additive FTC given in Definition <ref> contains the pre-trained network f, but one can confirm that FTC does not depend on f since f(_i)  = y_i holds for all i ∈ [K]∖ T, for every pre-trained network f we are considering. Below we provide an equivalent simpler definition. 



    For a given positive integer K, the fine-tuning capacity (FTC) of a given neural network g_θ is 
    
    N_FTC^⋆(g,K)    := max_N ∈{0, 1,⋯, K} N  such that 
       ∀ T ⊆ [K]  with | T | = N, ∀_i ∈ℝ^d, ∀ z_i ∈ℝ, 
    ∃θ    satisfying 
                g_θ(_i) = z_i    ∀  i ∈ T,
    
                g_θ(_i) = 0    ∀ i ∈ [K]∖ T.


This definition is a generalization of conventional memorization capacity shown below, when the condition is relaxed to a special case, T = [N]. 


    The memorization capacity of a neural network g_θ is 
    
    N_MC^⋆(g)    := max_N ≥ 0 N 
             such that 
    ∀_i ∈ℝ^d,    ∀ z_i ∈ℝ, 
             ∃θ with 
             g_θ(_i) = z_i   ∀ i ∈ [N]




    The memorization capacity and the fine-tuning capacity has a trivial bound: for any neural network g and for arbitrary K > 0, 
    
    N_FTC^⋆(g,K) ≤ N_MC^⋆(g).



Note that FTC is defined as the maximum number of samples N we can fine-tune using a given network g. One can also consider an equivalent definition: the minimum number of neurons m contained in g to successfully fine-tune N samples, which is formally stated below. 


    The minimum number of neurons required for fine-tuning arbitrary N out of K samples, is defined as  
    
    m_FTC^⋆(N,K)    := min_m ≥ 0 m 
          such that 
       ∀ T ⊆ [K]  with | T | = N, ∀_i ∈ℝ^d, ∀ z_i ∈ℝ, 
       ∃ neural network  g_θ with  m  neurons satisfying 
       
                g_θ(_i) = z_i     for all   i ∈ T,
    
                g_θ(_i) = 0     for all  i ∈ [K]∖ T.



Throughout the paper, we use N^⋆ as a short-hand notation for N^⋆_FTC, and use m^⋆ as a short-hand notation for m^⋆_FTC. Our theoretical results provide bounds on m^⋆ and N^⋆, Theorem <ref> is bounding m^⋆, while an equivalent result in Corollary <ref> is bounding N^⋆.













§ FTC OF 2-LAYER FC RELU NETWORKS



A 2-layer fully-connected neural network g_θ:ℝ^d →ℝ with ReLU activation can be represented as

    g_θ(x) = _2σ(_1  + _1) + _2,

which is parameterized by θ = [_1, _2, _1, _2] where _1 ∈^m × d, _2 ∈^1 × m, _1 ∈^m, and _2 ∈. Here, m is the number of hidden neurons, and σ is the ReLU activation. 
The below result states the bounds on m for 2-layer FC ReLU networks.


Let K ≥ 3.

    
  * For all T ⊆ [K], |T|=N, _i ∈ℝ^d, z_i ∈ℝ, i ∈ [K], there exists a 2-layer fully-connected ReLU network g with m neurons satisfying (<ref>) 
    and 
    
    m ≤min{3N+1, K-1}.

    
  * For given T ⊆ [K], |T|=N,  _i ∈ℝ^d, z_i ∈ℝ, i ∈ [K], suppose that (<ref>) holds for some 2-layer fully-connected ReLU network g with m neurons. Then, 
    min{3N, K-2}≤ m.


    Thus, 
    
    min{3N, K-2}≤ m^⋆≤min{3N+1, K-1}.


The proof of this theorem is given in Sec. <ref> and Sec. <ref>.



    If N=K (i.e., fine-tuning changes all labels), then the result of Theorem <ref> reduces to 
    
    m^⋆ + 1 ≤ N = K ≤ m^⋆+2.

    The upper bound of N coincides with the upper bound of memorization capacity studied in the result of <cit.>. On the other hand, the lower bound is consistent with the one given in <cit.>. 




    Below we state the  bounds for the fine-tuning capacity of a 2-layer fully-connected ReLU, directly obtained from the above theorem.


    Suppose K ≥ 3. For given m ∈ℕ, let N^⋆ be the fine-tuning capacity of a 2-layer fully-connected ReLU network g given in (<ref>) with m neurons. 
    
        
  * If K ≥ m+2, then
        
    ⌊m-1/3⌋≤  N^⋆≤m/3.

        
  * If K ≤ m+1, then
        N^⋆ = K.
    


    Suppose that K ≥ m+2. By Theorem <ref>.1, for |T|= ⌊m-1/3⌋, there exists a 2-layer fully-connected ReLU network g where the number of neurons of g is less than or equal to
    
    min{3⌊m-1/3⌋+1, K-1}= 3⌊m-1/3⌋+1 ≤ m.

    Here, K ≥ m+2 yields the first equality while the second inequality follows from the definition of the floor function. This yields the lower bound of (<ref>). On the other hand, due to Theorem <ref>.2 and K ≥ m+2, we conclude the upper bound of (<ref>).

    Suppose that K ≤ m+1. Similarly, by Theorem <ref>.1 for |T| = K, there exists a 2-layer fully-connected ReLU network g where (the number of neurons of g) ≤ K-1 ≤ m. This yields N ≥ K. As N ≤ K always holds due to our construction, we conclude N=K.


In other words, for sufficiently large K, the fine-tuning capacity N does not depend on the size K of the underlying dataset D but the number of neurons m. 



















 §.§ Proof of Lower Bound on m^⋆



We first prove the lower bound in Theorem <ref> when K ≥ 3N+2. 
Define T = {3, 6, 9, ..., 3N}, and define _i = i for all i ∈ [N] for arbitrary vector ∈^d. Let z_i = 2 if i = 6k for some k ∈ and z_i = -1 if i = 6k-3 for some k ∈. See Fig. <ref> when K=14 and N=4.  
Then, g̅_θ(t) := g_θ(t) is a piecewise affine function with at least 3N+1 pieces. Recall that using Theorem 3.3 of <cit.> for 2-layer ReLU network, g̅_θ(t) is having m+1 pieces, since ReLU activation is a piecewise linear function with two pieces. Thus, m ≥ 3N holds. 
Note that given K 
and N,
the above proof scheme specifies T, {_i, z_i}_i=1^K and counts the number of pieces for the piecewise-linear function g̅_θ(t) = g_θ(t), under the setting of  _i = i. 

We use a similar technique for proving the lower bound on m when N ≤ K < 3N+2. Consider revising (T, x_i, z_i) triplet (defined for K ≥ 3N+2 case),
so that K < 3N+2 condition is satisfied. For example, when K = 2N+1, the triplet is revised as T = {2, 4, ⋯, 2N}, z_i  = 2 if i=4k for some k ∈ℕ and z_i  = -1 if i=4k-2 for some k ∈ℕ, as illustrated in Fig. <ref>, which has K-1=8 pieces.
For general K and N satisfying K < 3N+2, one can choose (T, x_i, z_i) triplet such that the corresponding 
g̅_θ(t) = g_θ(t) has K-1 pieces. 
Thus, the number of pieces p(K) of g̅_θ(t) we constructed can be represented as

    p(K) = 
        3N+1,     if  K ≥ 3N+2 
    
        K-1,     if  N ≤ K < 3N+2

This completes the proof. 










 §.§ Proof of Upper Bound on m^⋆









We here establish the upper bound in Theorem <ref>. To be specific,  in Theorem <ref>, we establish the upper bound on m^⋆ in terms of the partition of [K]∖ T. The key idea is to remove all points of  [K]∖ T except for the endpoints of each block as in Figure <ref>. More discussion will be provided below.



Let us introduce some terminology. For a given set I, a partition 𝒫 of I is a set of nonempty subsets P of I such that every element in I is in exactly one of these subsets. We denote P ∈𝒫 by the block of 𝒫.






    For I ⊂ [K], we say that 𝒫(I) is the consecutive partition of I if all consecutive integers in I are included in the same block, i.e., for i,j ∈ I, i and j are in the same block of 𝒫(I) if and only if |i-j|=1.





    If I = {1, 2, 3, 5, 6, 8, 10, 11, 12, 13}, then the consecutive partition is defined as

    𝒫(I) = {{1, 2, 3}, {5, 6}, {8}, {10, 11, 12, 13}},

and the blocks of 𝒫(I) are {1, 2, 3}, {5, 6}, {8}, and {10, 11, 12, 13}.




The following theorem shows that the upper bound of m is given in terms of the size of the blocks in 𝒫([K]∖ T). For proving Theorem <ref>, we find the uniform bound for general datasets using this bound and Lemma <ref>.












Consider the same setting as in Theorem <ref>, and suppose that 

    ^T_1  < ^T_2 < ⋯ < ^T_K

holds for some ∈^d. Then, there exists an m-neuron network 
     g_θ(x) = _2σ(1_m^T  + _1) + _2,
    such that (<ref>) holds, and 
    
    m^⋆≤ K-1 - ∑_P ∈𝒫([K]∖ T)max{ |P|-2,0 }

    Here, _2 ∈^1 × m, _1 ∈^m, and _2 ∈.
     







    To minimize the width m of the network, a smaller number of blocks of bigger sizes would be ideal. If we only have one block in the partition 𝒫([K]∖ T), then Eq. <ref> is given as
    
    m^⋆≤ K-1 - | [K]∖ T | +2 = N + 1.

    This happens when ^T _is for i ∈ [K]∖ T are segregated for some , when [K]∖ T = {i, i+1, ⋯, j} for some i < j in [K].
    Thus, some appropriate projection yields a smaller number of neurons required.


    




    The worst scenario is when every block of 𝒫([K]∖ T) has 2 or less elements, as in Figures <ref> and <ref>. Note that if K - N is much larger than N, then this scenario cannot occur. This is because the number of blocks cannot be larger than min{K-N, N}+1 from Lemma <ref>.



    First, we consider the case when 
    
    |P| ≤ 2  for all  P ∈𝒫( [K]∖ T).

     In this case, it suffices to prove m ≤ K.
    As K neurons can represent K data points, the inequality directly follows from the standard argument as in <cit.> and also shown in Lemma <ref>.

    Let us consider general cases. The main strategy is to remove data points so that (<ref>) holds. Specifically, except for two endpoints of each block P ∈𝒫([K]∖ T), we remove i ∈ [K]∖ T from [K]. Let us denote this new subset of [K] by 
    
    J = {j_1 < j_2 < ⋯ < j_|J|}

    For example, when K=14 and T = {4,7,9,14} as in Fig. <ref>, the consecutive partition 𝒫(I) is given in Eq. <ref>, and thus we remove {2,11,12} from [K] = {1,2, ⋯, 14}, which gives us J = {1,3,4,5,6,7,8,9,10,13,14}.  
    Note that the number of data points in J is 
    
    |J|:= K - ∑_P ∈𝒫( [K]∖ T)max{ |P|-2,0 }.

    Applying Lemma <ref> with A:= {(^T_j, _j)}_j ∈ J and m = |J|-1, there exist _2 ∈^1 × m, _1 ∈^m, and _2 ∈ such that 
    
    h_θ(^T_j) = _2σ(1_m^T _j + _1) + _2 = z_j

    for all j ∈ J. Thanks to h_θ(^T) = g_θ(), we conclude g_θ(_j) = z_j for all j∈ J.
    
    Lastly, for i ∈ [K]∖ J, there exist two endpoints j_i < j_i+1 such that [j_i, j_i+1] ∩ ([K]∖ T) in 𝒫([K]∖ T), and j_i < i < j_i+1. Note that h_θ constructed in Lemma <ref> is linear in [^T _j_i, ^T _j_i+1]. Since g_θ(_j_i) = g_θ(_j_i+1) = 0, we conclude that g_θ(_i) = 0 as desired. 









 

    For m≥1 and A = {(w_i, z_i)}_i=1^m+1  where w_1 < w_2 < ⋯ < w_m+1, w_i ∈ and y_i ∈. There exist _2 ∈^1 × m, _1 ∈^m, and _2 ∈ such that 
    
    h_θ(x) = z_i-z_i+1/w_i-w_i+1(x - w_i) + z_i

    for x ∈ [w_i, w_i+1], i=1,2,⋯, m-1
    and (<ref>) with i=m holds in [w_m, ∞)
    where h_θ(x) = _2σ(1_m x + _1) + _2.



    We prove this by induction. For m=1, choose _1 = -w_1 and b_1  = z_1, then h_θ(w_1) = z_1. Take W_2 = z_1-z_2/w_1-w_2 and we get (<ref>) with i=1 in [w_1, ∞).

    Suppose that the above result holds for m = k. Then, there exist _2 ∈^1 × k, _1 ∈^k, and _2 ∈ satisfying (<ref>) in [w_i, w_i+1] for i = 1,2, ⋯ k-1 and [w_k, ∞). Using this, we choose _2 = (_2, z_k+1-z_k+2/w_k+1-w_k+2) and b_1 = (b_1, -w_k+1). Then, h_θ(x) = _2σ(1_m x + _1) + _2 satisfies  (<ref>) with i=k+1 in [w_k+1, ∞).





    For I ⊂ [K] and the consecutive partition, 𝒫(I), it holds that
    
    |𝒫(I)| ≤min{ |I|, |[K]∖ I| + 1}.



    First, |𝒫(I)| ≤ |I| directly follows from the definition of the partition. 
    
    On the other hand, let a_i and b_i be the end point of each block in 𝒫([K]∖ I):
    
    𝒫([K]∖ I)    = {[a_i, b_i] ∩ [K]: 1≤ i ≤ |[K]∖ I|}.

    Then, 
    
    𝒫(I)    = {[b_i+1, a_i+1-1] ∩ [K] :1≤ i ≤ |[K]∖ I|-1}
       ∪{ [1, b_i-1] ∩ [K], [a_[K]∖ I+1, K] ∩ [K] }

    and thus we conclude that
    |𝒫(I)| ≤ |[K]∖ I| + 1.




    Under the same setting as in Theorem <ref>, we have
    
    ∑_P ∈𝒫([K]∖ T)max{ |P|-2,0 }≥max{K-3N-2, 0 }.

    In particular, J given in (<ref>) satisfies
    
    |J| ≤min{3N+2, K }.




    Recall that ∑_P ∈𝒫([K]∖ T) |P| = |[K]∖ T| =  K-N.
    Using this, we have
    
    ∑_P ∈𝒫([K]∖ T)max{ |P|-2,0 }   ≥∑_P ∈𝒫([K]∖ T) (|P| - 2 )
       ≥ K-N - 2 |𝒫( [K]∖ T)|.

    By Lemma <ref>, the number of blocks in the partition cannot be larger than min{K-N, N+1}.


    
    



    







Proof of upper bound in Theorem <ref>:
    Since _i _j for all i  j, there exists ∈^d satisfying _i^T _j^T all i  j. Without the loss of generality, we assume that (<ref>) holds. Then, using Theorem <ref> and choosing _1 = 1_m^T, there exists
    
    g_θ(x) = _2σ(_1  + _1) + _2,

    satisfying (<ref>) and 
    
    m ≤ K - 1 - ∑_P ∈𝒫([K]∖ T)max{ |P|-2,0 }.

    By Lemma <ref>, we conclude that
        m ≤ K-1 - max{K-3N-2, 0 } = min{3N+1, K-1}. 











§ FTC OF 3-LAYER RELU NETWORK





Now we analyze FTC of 3-layer fully-connected neural network g_θ:ℝ^d →ℝ with ReLU activation. Note that 3-layer network can be represented as 

    g_θ(x) = _3 σ(_2σ(_1 + _1) + _2) + b_3.
 
As before, σ is the ReLU activation and the network is parameterized by θ = [ _1, _2, _3, _1, _2, b_3] where ∈^d_2, _1∈^d_1× d, _2∈^d_2× d_1, _3∈^1× d_2, _1∈^d_1, _2∈^d_2 and b_3 ∈. 
Following the setting considered in the memorization capacity of 3-layer neural network <cit.>, we consider the scenario when z_i ∈ [-1, +1] for all i ∈ [K].
For notational simplicity, we denote z_i = 0 for i ∈ [K]∖ T.
Below theorem summarizes the upper/lower bound on FTC of 3-layer network. 



Let K ≥ 3, T ⊆ [K], |T| = N, and g be a 3-layer FC ReLU network with m neurons.

    
  * For all _i ∈ℝ^d, z_i ∈ℝ, i ∈ [K], there exists g with m neurons satisfying (<ref>) 
    where 
    m ≤min{ 2√(K) + min{ 2√(K), 3N }, 6 √(3N+2)}.

    
  * For given _i ∈ℝ^d, z_i ∈ℝ, i ∈ [K], suppose that (<ref>) holds for some g with m neurons. Then, 
    √(2min{3N, K-2}+1/4) - 1/2≤ m.


Thus, the minimum number of neurons m^⋆ is bounded as

    √(2min{3N, K-2}+1/4) - 1/2
           ≤ m^⋆≤min{ 2√(K) + min{ 2√(K), 3N }, 6 √(3N+2)}




The upper bound in Theorem 5.1 directly shows that m^⋆≤Θ(√(N)).
    The lower bound in Theorem 5.1 indicates two facts: 
    
        
  * If 3N ≤ K-2, then m^⋆≥Θ(√(N)),
        
  * If 3N > K-2, then m^⋆≥Θ(√(K)). Combining this with K ≥ N, we have m^⋆≥Θ(√(N)).
     
    All in all, m^⋆ = Θ(√(N)).




    For given m ∈ℕ, let N^⋆ be the fine-tuning capacity of a 3-layer fully-connected ReLU network g given in (<ref>) with m neurons.
    Then, N^⋆ is bounded as below, for different range of K:
    
        
  * If K ≤⌊m^2/16⌋, then
            
    N^⋆ = K.

        
  * If ⌊m^2/16⌋ + 1 ≤ K < m^2+m+4/2, then
            
    ⌊m^2/108 - 2/3⌋≤ N^⋆≤ K.

        
  * If K ≥m^2+m+4/2, then
            
    ⌊m^2/108 - 2/3⌋≤ N^⋆≤m^2+m/6.

    


    Suppose K ≤⌊m^2/16⌋. By Theorem <ref>.1, for |T|= K, there exists a 3-layer fully-connected ReLU network g where the number of neurons of g is less than or equal to
    
    min{ 4√(K), 2√(K)+3K, 6√(3K+2)} = 4√(K)≤ m.
 This implies that N ≥ K, and since N≤ K always holds, we can conclude N^⋆ = K.

    Suppose K ≥⌊m^2/16⌋ + 1. 
    By Theorem <ref>.1, for |T|=⌊m^2/108-2/3⌋, there exists a 3-layer fully-connected ReLU network g where the number of neurons of g is less than or equal to
    
    min{ 4√(K), 6√(3⌊m^2/108-2/3⌋+2)}
       = 6√(3⌊m^2/108-2/3⌋+2)≤ m,

    where the inequality is derived from the fact that 4√(K)≥ m for 
    K ≥⌊m^2/16⌋ + 1, and 6√(3⌊m^2/108-2/3⌋+2)≤ 6√(3 ( 
    m^2/108-2/3)+2) = m .
    Now we can conclude that N^⋆≥⌊m^2/108-2/3⌋.

    Now we derive the upper bound on N^⋆. 
    If K ≥m^2+m+4/2, then
    
    m    ≥√(2min{3N, K-2}+1/4) - 1/2
       ≥√(min{ 6N, m^2+m } + 1/4) - 1/2
       = min{√(6N+1/4),   m+ 1/2} - 1/2
       = min{√(6N+1/4) - 1/2,   m },

    where the first inequality is from Theorem <ref>.2, and the second inequality is from K ≥m^2+m+4/2.
    We can simplify the above inequalities as m ≥√(6N+1/4) - 1/2. Since this holds for any number of fine-tunable samples N, we have N^⋆≤m^2+m/6.

    If ⌊m^2/16⌋ + 1 ≤ K < m^2+m+4/2, we use a trivial upper bound 
    N^⋆≤ K, which completes our proof. 



 






Fig. <ref> illustrates the results in Corollary <ref>. For different ranges of K, we have either a constant N^⋆ = K, or upper/lower bounds on N^⋆. 



 §.§ Proof of Lower Bound on m




We follow the proof technique used in section  <ref>. Since the number of minimum pieces p(K) provided in Eq. <ref> is still valid for 3-layer neural network, we just need to check how many pieces g̅_θ(t) has. As stated in the proof of Theorem 3.3 of  <cit.>, 
g̅_θ(t) has 2d_1d_2+d_2+1 pieces, where d_1 and d_2 are the number of neurons in layer 1 and 2, respectively. 
Thus, we have

    min{3N, K-2}≤ 2d_1d_2+d_2 ≤m^2/2 + m/2,
 
where the last inequality is from the fact that 2d_1d_2+d_2 is having its maximum value when d_1 = d_2 = m/2. 
Reformulating the above inequality with respect to m completes the proof.






 §.§ Proof of Upper Bound on m


We can prove m ≤ U for an upper bound U by constructing a 3-layer neural network g_θ having U neurons, which satisfies Eq. <ref> for given N and K. Below we construct different types of neural networks satisfying the condition, where each construction gives different upper bounds U_1 = 4√(K), U_2 = 2√(K) + 3N and U_3  = 6√(3N+2). This completes the proof of m ≤min{U_1, U_2, U_3}. 




 
    
       
    

  
Proof of m ≤  4√(K):

    Recall the neural network  g_θ constructed in the proof of Theorem 3.1 of <cit.>, containing √(K) neurons in both 1st layer and 2nd layer. See Fig. <ref> for the illustration of such network, when K=16 and N=1, where the index of fine-tuning sample is T = {7}. In such case z_7  0 from the definition of T.
    Each box (node) in the figure is a neuron, where the curve inside the box represents the activation function of the neuron. In this figure, each neuron uses hard-tanh activation defined as 
    
    σ_H(x) = 
                -1,       t ≤ -1 
    
                t,       -1 < t < +1 
    
                +1,       t ≥ +1,

    where [-1, 1] is the non-clipping region of σ_H, and [-1, 1]^c is the clipping region of σ_H.
    Note that the digit i in the box represents the location where the feature _i for the i-th sample is mapped to. 
    For example, for the first neuron of layer 1, we have α_1^1 (_i)  ∈ [-1,1] for i ∈{1,2,3,4} and α_1^1 (_i) >  1 for i ∈{5, 6, ⋯, 16}, where 
    
    α_j^l () = _l,j + b_l,j
 
    is the input value of node j in layer l, when the input for the network g_θ is given as . Here, the weight matrix and the bias for layer 1 are denoted by _1 = [_1,1^T; ⋯; _1,√(K)^T] and _1 = [b_1,1, ⋯, b_1,√(K)], respectively. 
    Given target {z_i}_i=1^K, the proof of Theorem 3.1 of <cit.> specified parameters θ = [ _1, _2, _3, _1, _2, b_3] satisfying 
        g_θ(_i) = z_i for all i ∈ [K].
    Assigning z_i = 0 for i ∈ [K] ∖ T 
    and reusing these parameters is enough to satisfy the desired condition for fine-tuning in Eq. <ref>. Note that this network uses 2√(K) neurons with hard-tanh activations, which can be converted to a ReLU neural network with 4√(K) neurons, using the fact that one hard-tanh neuron can be expressed with two ReLU neurons. This directly proves that 4√(K) ReLU neurons are sufficient for changing the labels of N ≤ K samples. 


    
    





  
Proof of m ≤  2√(K) + 3N:
    We construct a 3-layer neural network with U_2 = 2√(K) + 3N neurons which successfully fine-tunes N samples. 
    Fig. <ref> illustrates the example of such construction when K=16 and N=1. Here, we assumed T = {7}, the label for 7-th sample is fine-tuned, but similar proof can be applied to arbitrary T with |T| = N =1. 
    Here, our goal is to construct g_θ satisfying   g_θ (_7) = z_7  0 and  g_θ (_i) = 0 for all i  7.
    Our basic idea is, to follow the construction in Fig. <ref>,
    except the activation used in layer 2. The activation in layer 2 is defined as
    
    σ_B (t) = 
                0,     t ∉ [z_7 - δ, z_7 + δ] 
    z_7/δ (t - z_7 + δ),      t ∈ [z_7 - δ, z_7) 
    
                - z_7/δ (t - z_7 - δ),      t ∈ [z_7, z_7 + δ)
 
    for arbitrary δ < min_i  j| z_i - z_j  |/2, which is having a small bump near z_7 as in Fig. <ref>. 
    Among √(K) hard-tanh neurons in layer 2 in Fig. <ref>, we only choose the neuron containing z_7 (the non-zero target label) in the non-clipping region of the activation, the second neuron of layer 2. 
    Given that _2 is the weight matrix for 2nd layer in the construction of Fig. <ref>, the weight matrix for 2nd layer in the construction of Fig. <ref> is defined as _2  := _2,2, the 2nd row of _2. 
    Then, the overall network looks like g_θ() = σ_B ( _2,2σ_H (_1  + _1) + _2), which satisfies g_θ(_7) = z_7 and g_θ(_i)=0 for all i  7.
    Note that the activation σ_B with 4 pieces can be constructed by 3 ReLU neurons. Consider constructing g_θ by adding new neurons (with σ_B activation) in layer 2 for each sample we want to fine-tune, which allocates total 3N ReLU neurons in layer 2. Since layer 1 (as in Fig. <ref>) contains 2√(K) ReLU neurons, one can confirm that our construction contains 2√(K) + 3N ReLU neurons.







  
Proof of m ≤ 6 √(3N+2):



    
It is worth nothing that some techniques developed for showing m ≤ 4N+4 (shown in Appendix) are used to prove m ≤ 6√(3N+2). Thus, we refer to some equations in Appendix during the proof.

We construct a neural network g_θ with 6√(3N+2) ReLU neurons, which fine-tunes N samples. For simplicity, here we provide an example construction when K = 16, N = 3, and T = {2, 4, 9}. By using the definition of J in Eq. <ref>, we have J = { 1, 2, 3, 4, 5, 8, 9, 10, 16 }. 
We will now construct a 3-layer neural network g_θ satisfying

    g_θ(_i) = z_i,    i ∈ J, 
    
    g_θ(_i) = 0,    i ∉ J,

which is illustrated in Fig. <ref>. 
Note that Eq. <ref> can be easily proved by fitting the network using the samples _i with i ∈ J. In the rest of the proof, we will show that for such g_θ satisfying Eq. <ref>, we can prove Eq. <ref>. As an example, we will only show g_θ(_6) = 0, but a similar proof technique can be applied to arbitrary _i with i ∉ J.



 
We construct the first layer by only using samples _i with i ∈ J. We partition 
the |J| samples  into √(|J|) groups, following the trick used  in <cit.> for 3-layer network. 
Let us denote the first index of j-th group as s_j^min and the last index of j-th group as s_j^max. In other words, J is decomposed into √(|J|) groups as J = {s_1^min, ⋯, s_1^max,  }∪⋯∪{ s_√(|J|)^min, ⋯, s_√(|J|)^max}. 


Without loss of generality, we can assume that samples are ordered as ^T_1<^T_2<…<^T_K for some . Let c_i := ^T_i and define c_0 = c_1-ϵ, c_K+1 = c_K+ϵ for arbitrary ϵ > 0. Then, we choose _1, b_1 as in Eq. <ref>, using  and ϵ defined above.

Here, we focus on the relationship between the outputs of layer 1, when the neural network inputs are _5, _6, _8, respectively. 
From the definition of c_i, we have c_6 ∈ (c_5, c_8). Using Eq. <ref> and Eq. <ref>, we have 
α_j^1 (_6) ∈ (α_j^1 (_5), α_j^1 (_8)), meaning that the input α_j^1 of layer 1 (for _6) is bounded by α_j^1 for _5 and α_j^1 for _8. After passing it through the ReLU activation, we also have 

    β_j^1 (_6) ∈ (β_j^1 (_5), β_j^1 (_8)),

for all j ∈{1, ⋯, √(J)}
using the monotonicity of ReLU. Kore precisely, 
we have 
β_1^1(x_6)=+1, -1≤β_2^1(x_6)≤+1, and β_3^1(x_6)=-1. This is illustrated in the first layer in   Fig. <ref>.
Now we move to the construction of the second layer. Recall that the main idea of constructing _2, b_2 is using the linear system in Eq. <ref>. 
Using Eq.<ref> and the fact that we can set all elements of _2 as positive (as shown in <cit.>), Eq. <ref> implies 
 α_j^2 (_6) ∈ (α_j^2 (_5), α_j^2 (_8))
for all j ∈{1, ⋯, √(J)}. Finally, we set the activation function of layer 2 as

    σ_L(t) = 
            t,     if  |t| ≤ 1, 
    
            0,     if  |t| ≥ 1

and arbitrarily increase λ such that β_j^2(_6) = 0. 
Then, the output of the network 
g_θ() = σ_L(_2(σ_H(_1  + _1)) + _2)
satisfies g_θ(_6) = 0.

Now the question is, what is the upper bound on |J|? Recall that Lemma <ref> guarantees that |J|≤3N+2. 
Since σ_H and σ_L can be represented by 2 ReLUs and 4 ReLUs, respectively, 2√(3N+2)+4√(3N+2)=6√(3N+2) ReLU neurons are sufficient for our 3-layer network construction.
    

    





































































































§ EXTENSION TO OTHER NEURAL NETWORKS


For deeper neural networks, a lower bound on N can be obtained in terms of the maximum width d, the number of layers L, and the number of pre-trained samples K. 



For L ≥ 4, K ≥ 3, T ⊆ [K], |T|=N, there exists an L-layer ReLU network with maximum width d satisfying (<ref>) and  
    d ≤   4min{√(3N/√(⌊L-1/2⌋)+5), √(K)} + 2.





The main challenge for proving Proposition <ref> lies in constructing a suitable neural network, which can be addressed by utilizing our 3-layer network from Section <ref> and the construction idea provided in Figure 2 of <cit.>. See Appendix <ref> for the proof of Proposition <ref>.





§ EXPERIMENTS


In this section, we provide experimental results on a synthetic dataset, which supports our theoretical results.
The experimental setup is as follows. 
We first randomly generate K samples D={(_i, y_i)}_i=1^K where the feature and the label of i-th sample have the following distributions: _i ∼ N(0, I_d) and y_i ∼Unif[-1, 1], where the feature dimension is d=10. Then, we train a network ReLU network f that fits the dataset D,  (<ref>) holds, thus having zero mean-squared-error (MSE) loss ℓ = 1/K∑_i=1^K (f(_i)-y_i)^2. 
Considering the fine-tuning scenario, we construct another dataset D'={(_i, y_i' )}_i=1^K as follows. We first initialize D' = D. Then, we randomly choose N out of K samples in D', and re-define the label of the N samples as y_i' ∼Unif[-1, 1].
Finally, we implement the fine-tuning process. Following our additive fine-tuning scenario, we freeze f and train a 3-layer ReLU network g_θ with m neurons, in a way that f+g_θ fits the new dataset D'. We define the fine-tuning loss as ℓ_FT = 1/K∑_i=1^K (f(_i) + g_θ(_i)-y_i')^2.

Figure <ref> shows the fine-tuning loss ℓ_FT for different m and N. As expected, for a given N, the fine-tuning loss decreases as m increases. For each N, the yellow line in the figure shows the minimum m satisfying ℓ_FT(m,N) ≤ 0.04. This yellow line indicates that the required number of neurons to achieve small fine-tuning loss follow the tendency of Θ(√(N)) shown in the red line in the figure, which coincides with our theoretical result in Theorem <ref>.



§ CONCLUSION



We introduced Fine-Tuning Capacity (FTC), a generalization of memorization capacity concept for fine-tuning applications. 
This concept is defined to provide theoretical view on current paradigm of fine-tuning large pre-trained models. 
As an initial step towards analyzing FTC, we focused on the additive fine-tuning scenario where a side network is added to the frozen pre-trained network. 
We obtained upper/lower bounds on FTC for shallow ReLU networks.
For 2-layer network, we showed that fine-tuning N samples is possible by using ReLU networks with m = Θ (N) neurons, irrespective of the size of the pre-trained network and the number of total samples K used during pre-training.
For 3-layer network, the required amount of neurons reduces to m = Θ (√(N)), for practical scenarios where the number of samples N we want to change labels is far less than the number of total samples K used for pre-training.



























 


    JS acknowledges the support of the National Research Foundation of Korea (NRF) grant funded by the Korea government (MSIT) No. RS-2024-00345351. DK was partially supported by the National Research Foundation of Korea (NRF) grant funded by the Korea government (MSIT) (No. RS-2023-00252516) and the POSCO Science Fellowship of POSCO TJ Park Foundation.
    KL was supported by the NSF Award CCF-2339978 and a grant from FuriosaAI.

    


47
    
    
    urlstyle[Barron(1993)]barron1993universal
    Andrew R Barron.
    Universal approximation bounds for superpositions of a sigmoidal
      function.
    IEEE Transactions on Information theory, 390
      (3):0 930–945, 1993.
    
    [Baum(1988)]baum1988capabilities
    Eric B Baum.
    On the capabilities of multilayer perceptrons.
    Journal of complexity, 40 (3):0 193–215,
      1988.
    
    [Bommasani et al.(2021)Bommasani, Hudson, Adeli, Altman, Arora, von
      Arx, Bernstein, Bohg, Bosselut, Brunskill,
      et al.]bommasani2021opportunities
    Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney
      von Arx, Michael S Bernstein, Jeannette Bohg, Antoine Bosselut, Emma
      Brunskill, et al.
    On the opportunities and risks of foundation models.
    arXiv preprint arXiv:2108.07258, 2021.
    
    [Brown et al.(2020)Brown, Mann, Ryder, Subbiah, Kaplan, Dhariwal,
      Neelakantan, Shyam, Sastry, Askell, et al.]brown2020language
    Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla
      Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell,
      et al.
    Language models are few-shot learners.
    Advances in neural information processing systems,
      33:0 1877–1901, 2020.
    
    [Cao et al.(2022)Cao, Prakash, and Hamza]cao2022attention
    Jin Cao, Chandana Satya Prakash, and Wael Hamza.
    Attention fusion: a light yet efficient late fusion mechanism for
      task adaptation in nlu.
    In Findings of the Association for Computational Linguistics:
      NAACL 2022, pages 857–866, 2022.
    
    [Carlini et al.(2022)Carlini, Ippolito, Jagielski, Lee, Tramer, and
      Zhang]carlini2022quantifying
    Nicholas Carlini, Daphne Ippolito, Matthew Jagielski, Katherine Lee, Florian
      Tramer, and Chiyuan Zhang.
    Quantifying memorization across neural language models.
    arXiv preprint arXiv:2202.07646, 2022.
    
    [Chowdhery et al.(2022)Chowdhery, Narang, Devlin, Bosma, Mishra,
      Roberts, Barham, Chung, Sutton, Gehrmann, et al.]chowdhery2022palm
    Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra,
      Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian
      Gehrmann, et al.
    Palm: Scaling language modeling with pathways.
    arXiv preprint arXiv:2204.02311, 2022.
    
    [Cybenko(1989)]cybenko1989approximation
    George Cybenko.
    Approximation by superpositions of a sigmoidal function.
    Mathematics of control, signals and systems, 20
      (4):0 303–314, 1989.
    
    [Du et al.(2020)Du, Hu, Kakade, Lee, and Lei]du2020few
    Simon S Du, Wei Hu, Sham M Kakade, Jason D Lee, and Qi Lei.
    Few-shot learning via learning the representation, provably.
    arXiv preprint arXiv:2002.09434, 2020.
    
    [Englert and Lazic(2022)]englert2022adversarial
    Matthias Englert and Ranko Lazic.
    Adversarial reprogramming revisited.
    Advances in Neural Information Processing Systems,
      35:0 28588–28600, 2022.
    
    [Fu et al.(2021)Fu, Huang, Chen, Tian, and Zhao]fu2021learn
    Cheng Fu, Hanxian Huang, Xinyun Chen, Yuandong Tian, and Jishen Zhao.
    Learn-to-share: A hardware-friendly transfer learning framework
      exploiting computation and parameter sharing.
    In International Conference on Machine Learning, pages
      3469–3479. PMLR, 2021.
    
    [Funahashi(1989)]funahashi1989approximate
    Ken-Ichi Funahashi.
    On the approximate realization of continuous mappings by neural
      networks.
    Neural networks, 20 (3):0 183–192, 1989.
    
    [Giannou et al.(2023)Giannou, Rajput, and
      Papailiopoulos]giannou2023expressive
    Angeliki Giannou, Shashank Rajput, and Dimitris Papailiopoulos.
    The expressive power of tuning only the norm layers.
    arXiv preprint arXiv:2302.07937, 2023.
    
    [Hanin and Sellke(2017)]hanin2017approximating
    Boris Hanin and Mark Sellke.
    Approximating continuous functions by relu nets of minimal width.
    arXiv preprint arXiv:1710.11278, 2017.
    
    [Hardt and Ma(2016)]hardt2016identity
    Moritz Hardt and Tengyu Ma.
    Identity matters in deep learning.
    arXiv preprint arXiv:1611.04231, 2016.
    
    [He et al.(2021)He, Zhou, Ma, Berg-Kirkpatrick, and
      Neubig]he2021towards
    Junxian He, Chunting Zhou, Xuezhe Ma, Taylor Berg-Kirkpatrick, and Graham
      Neubig.
    Towards a unified view of parameter-efficient transfer learning.
    arXiv preprint arXiv:2110.04366, 2021.
    
    [Hornik et al.(1989)Hornik, Stinchcombe, and
      White]hornik1989multilayer
    Kurt Hornik, Maxwell Stinchcombe, and Halbert White.
    Multilayer feedforward networks are universal approximators.
    Neural networks, 20 (5):0 359–366, 1989.
    
    [Houlsby et al.(2019)Houlsby, Giurgiu, Jastrzebski, Morrone,
      De Laroussilhe, Gesmundo, Attariyan, and Gelly]houlsby2019parameter
    Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin
      De Laroussilhe, Andrea Gesmundo, Mona Attariyan, and Sylvain Gelly.
    Parameter-efficient transfer learning for nlp.
    In International Conference on Machine Learning, pages
      2790–2799. PMLR, 2019.
    
    [Hu et al.(2021)Hu, Shen, Wallis, Allen-Zhu, Li, Wang, Wang, and
      Chen]hu2021lora
    Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean
      Wang, Lu Wang, and Weizhu Chen.
    Lora: Low-rank adaptation of large language models.
    arXiv preprint arXiv:2106.09685, 2021.
    
    [Huang(2003)]huang2003learning
    Guang-Bin Huang.
    Learning capability and storage capacity of two-hidden-layer
      feedforward networks.
    IEEE transactions on neural networks, 140
      (2):0 274–281, 2003.
    
    [Huang and Babri(1998)]huang1998upper
    Guang-Bin Huang and Haroon A Babri.
    Upper bounds on the number of hidden neurons in feedforward networks
      with arbitrary bounded nonlinear activation functions.
    IEEE transactions on neural networks, 90 (1):0
      224–229, 1998.
    
    [Ippolito et al.(2022)Ippolito, Tramèr, Nasr, Zhang, Jagielski,
      Lee, Choquette-Choo, and Carlini]ippolito2022preventing
    Daphne Ippolito, Florian Tramèr, Milad Nasr, Chiyuan Zhang, Matthew
      Jagielski, Katherine Lee, Christopher A Choquette-Choo, and Nicholas Carlini.
    Preventing verbatim memorization in language models gives a false
      sense of privacy.
    arXiv preprint arXiv:2210.17546, 2022.
    
    [Kidger and Lyons(2020)]kidger2020universal
    Patrick Kidger and Terry Lyons.
    Universal approximation with deep narrow networks.
    In Conference on learning theory, pages 2306–2327. PMLR,
      2020.
    
    [Kim et al.(2023)Kim, Kim, and Mozafari]kim2023provable
    Junghwan Kim, Michelle Kim, and Barzan Mozafari.
    Provable memorization capacity of transformers.
    In International Conference on Learning Representation (ICLR),
      2023.
    
    [Lu et al.(2017)Lu, Pu, Wang, Hu, and Wang]lu2017expressive
    Zhou Lu, Hongming Pu, Feicheng Wang, Zhiqiang Hu, and Liwei Wang.
    The expressive power of neural networks: A view from the width.
    Advances in neural information processing systems, 30, 2017.
    
    [Malladi et al.(2023)Malladi, Wettig, Yu, Chen, and
      Arora]malladi2023kernel
    Sadhika Malladi, Alexander Wettig, Dingli Yu, Danqi Chen, and Sanjeev Arora.
    A kernel-based view of language model fine-tuning.
    In International Conference on Machine Learning, pages
      23610–23641. PMLR, 2023.
    
    [Maurer et al.(2016)Maurer, Pontil, and
      Romera-Paredes]maurer2016benefit
    Andreas Maurer, Massimiliano Pontil, and Bernardino Romera-Paredes.
    The benefit of multitask representation learning.
    Journal of Machine Learning Research, 170
      (81):0 1–32, 2016.
    
    [Nguyen and Hein(2018)]nguyen2018optimization
    Quynh Nguyen and Matthias Hein.
    Optimization landscape and expressivity of deep cnns.
    In International conference on machine learning, pages
      3730–3739. PMLR, 2018.
    
    [OpenAI(2023)]openai2023gpt
    R OpenAI.
    Gpt-4 technical report.
    arXiv, pages 2303–08774, 2023.
    
    [Ouyang et al.(2022)Ouyang, Wu, Jiang, Almeida, Wainwright, Mishkin,
      Zhang, Agarwal, Slama, Ray, et al.]ouyang2022training
    Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela
      Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al.
    Training language models to follow instructions with human feedback.
    Advances in Neural Information Processing Systems,
      35:0 27730–27744, 2022.
    
    [Oymak et al.(2023)Oymak, Rawat, Soltanolkotabi, and
      Thrampoulidis]oymak2023role
    Samet Oymak, Ankit Singh Rawat, Mahdi Soltanolkotabi, and Christos
      Thrampoulidis.
    On the role of attention in prompt-tuning.
    arXiv preprint arXiv:2306.03435, 2023.
    
    [Park et al.(2020)Park, Yun, Lee, and Shin]park2020minimum
    Sejun Park, Chulhee Yun, Jaeho Lee, and Jinwoo Shin.
    Minimum width for universal approximation.
    arXiv preprint arXiv:2006.08859, 2020.
    
    [Radford et al.(2021)Radford, Kim, Hallacy, Ramesh, Goh, Agarwal,
      Sastry, Askell, Mishkin, Clark, et al.]radford2021learning
    Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh,
      Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark,
      et al.
    Learning transferable visual models from natural language
      supervision.
    In International conference on machine learning, pages
      8748–8763. PMLR, 2021.
    
    [Rajput et al.(2021)Rajput, Sreenivasan, Papailiopoulos, and
      Karbasi]rajput2021exponential
    Shashank Rajput, Kartik Sreenivasan, Dimitris Papailiopoulos, and Amin Karbasi.
    An exponential improvement on the memorization capacity of deep
      threshold networks, 2021.
    
    [Ramesh et al.(2022)Ramesh, Dhariwal, Nichol, Chu, and
      Chen]ramesh2022hierarchical
    Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen.
    Hierarchical text-conditional image generation with clip latents.
    arXiv preprint arXiv:2204.06125, 10 (2):0 3,
      2022.
    
    [Touvron et al.(2023)Touvron, Martin, Stone, Albert, Almahairi, Babaei,
      Bashlykov, Batra, Bhargava, Bhosale, et al.]touvron2023llama
    Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine
      Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale,
      et al.
    Llama 2: Open foundation and fine-tuned chat models.
    arXiv preprint arXiv:2307.09288, 2023.
    
    [Tripuraneni et al.(2020)Tripuraneni, Jordan, and
      Jin]tripuraneni2020theory
    Nilesh Tripuraneni, Michael Jordan, and Chi Jin.
    On the theory of transfer learning: The importance of task diversity.
    Advances in neural information processing systems,
      33:0 7852–7862, 2020.
    
    [Vardi et al.(2021)Vardi, Yehudai, and Shamir]vardi2021optimal
    Gal Vardi, Gilad Yehudai, and Ohad Shamir.
    On the optimal memorization power of relu neural networks.
    arXiv preprint arXiv:2110.03187, 2021.
    
    [Vershynin(2020)]vershynin2020memory
    Roman Vershynin.
    Memory capacity of neural networks with threshold and relu
      activations, 2020.
    
    [Wu et al.(2022)Wu, Zou, Braverman, Gu, and Kakade]wu2022power
    Jingfeng Wu, Difan Zou, Vladimir Braverman, Quanquan Gu, and Sham Kakade.
    The power and limitation of pretraining-finetuning for linear
      regression under covariate shift.
    Advances in Neural Information Processing Systems,
      35:0 33041–33053, 2022.
    
    [Yun et al.(2019)Yun, Sra, and Jadbabaie]yun2019small
    Chulhee Yun, Suvrit Sra, and Ali Jadbabaie.
    Small relu networks are powerful memorizers: a tight analysis of
      memorization capacity.
    Advances in Neural Information Processing Systems, 32, 2019.
    
    [Zaken et al.(2021)Zaken, Ravfogel, and Goldberg]zaken2021bitfit
    Elad Ben Zaken, Shauli Ravfogel, and Yoav Goldberg.
    Bitfit: Simple parameter-efficient fine-tuning for transformer-based
      masked language-models.
    arXiv preprint arXiv:2106.10199, 2021.
    
    [Zeng et al.(2023)Zeng, Li, Ren, Liu, Xu, He, Xing, Wang, Tang, and
      Yin]zeng2023exploring
    Shenglai Zeng, Yaxin Li, Jie Ren, Yiding Liu, Han Xu, Pengfei He, Yue Xing,
      Shuaiqiang Wang, Jiliang Tang, and Dawei Yin.
    Exploring memorization in fine-tuned language models.
    arXiv preprint arXiv:2310.06714, 2023.
    
    [Zeng and Lee(2023)]zeng2023expressive
    Yuchen Zeng and Kangwook Lee.
    The expressive power of low-rank adaptation.
    arXiv preprint arXiv:2310.17513, 2023.
    
    [Zhang et al.(2016)Zhang, Bengio, Hardt, Recht, and
      Vinyals]zhang2016understanding
    Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals.
    Understanding deep learning requires rethinking generalization.
    arXiv preprint arXiv:1611.03530, 2016.
    
    [Zhang et al.(2020)Zhang, Sax, Zamir, Guibas, and Malik]zhang2020side
    Jeffrey O Zhang, Alexander Sax, Amir Zamir, Leonidas Guibas, and Jitendra
      Malik.
    Side-tuning: a baseline for network adaptation via additive side
      networks.
    In Computer Vision–ECCV 2020: 16th European Conference,
      Glasgow, UK, August 23–28, 2020, Proceedings, Part III 16, pages 698–714.
      Springer, 2020.
    
    [Zhang et al.(2022)Zhang, Roller, Goyal, Artetxe, Chen, Chen, Dewan,
      Diab, Li, Lin, et al.]zhang2022opt
    Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui
      Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al.
    Opt: Open pre-trained transformer language models.
    arXiv preprint arXiv:2205.01068, 2022.
    
    
    






 Memorization Capacity for Additive Fine-Tuning
(Supplementary Material)
    E. Ya. Paryev

Institute for Nuclear Research of the Russian Academy of Sciences

Moscow, Russia
    August 5, 2024
====================================================================================================






 

§ PROOF OF M ≤ 4 N+4 (DEFERRED FROM SEC. <REF>) 



    We construct a 3-layer neural network g_θ with 4N+4 ReLU neurons which fine-tune N samples. 
    Note that we can find  such that ^T_i ^T_j for all i  j, since we assume _i _j for all i  j. 
    Without loss of generality, we order samples such that ^T_1<^T_2<…<^T_K. Let c_i := ^T_i and define c_0 = c_1-ϵ, c_K+1 = c_K+ϵ for arbitrary ϵ > 0. 
    Recall that the indices for the samples we want to fine-tune is denoted by T = {T_1, ⋯, T_N}. We define dummy indices T_0 = 1 and T_N+1 = K. Consider 2N+1 groups of disjoint indices,
        
    s_1 = {T_0, ⋯, T_1 - 1},    
            s_2    = {T_1 }, 
    
            s_3 = {T_1 + 1, ⋯, T_2 - 1},    
            s_4    = {T_2 }, 
    ⋯,   
            s_2N   = {T_N}, 
    
            s_2N+1 = {T_N + 1, ⋯, T_N+1},

    where group s_1 is empty when T_1 = 1, and group s_2N+1 is empty when T_N = K.  
    We denote the maximum/minimum element of set s_j as s_j^max and s_j^min, respectively, s_j^max = max s_j and s_j^min = min s_j.
    We place 2N+1 neurons on layer 1, and define parameters for layer 1 as 
    
    _1,j   = (-1)^j-14/c_s_j^max +  c_s_j+1^min - c_s_j-1^max - c_s_j^min^T 
    
            b_1,j   = (-1)^j
            c_s_j^max +  c_s_j+1^min +
            c_s_j-1^max + c_s_j^min/c_s_j^max +  c_s_j+1^min - c_s_j-1^max - c_s_j^min

    for all j =1, ⋯, 2N+1. 
    Under such setting, it can be easily checked that
    
    α_j^1 (_i) ∈ [-1, 1]    for  i ∈ s_j 
    α_j^1 (_i) ∉ [-1, 1]    for  i ∉ s_j

    for all j ∈ [2N+1], where α_j^l () defined in Eq. <ref> is the input value of node j in layer l, when the input for the network is . 
    Fig. <ref> shows the example of such construction, when K=16, N=2 and T = {4,7}. In such case, we have 2N+1 = 5 disjoint groups:
        
    s_1 = {1, 2, 3},   
            s_2    = {4},   
            s_3 = {5,6}, 
    
            s_4    = {7},   
            s_5 = {8, 9, ⋯, 16}.

    As in Fig. <ref>, the input of j-th neuron lie in the non-clipping region α_1^j(_i) ∈ [-1, 1] when i ∈ s_j. For example, the first neuron (j=1) in the first layer has α_1^1 (_i) ∈ [-1, 1] for i ∈ s_1 = {1,2,3}.


    
    
    
    Now we construct layer 2 as below. Let S = ∪_j = 1^2N+1{ s_j^min} be the index set containing the minimum index of each group s_j for j ∈ [2N+1], which is S = {1,4,5,7,8} in the above example. 
    Our goal is to construct _2 ∈ℝ^2N+1 and b_2 ∈ℝ such that a single node in layer 2 locates (1) {_i}_i ∈ S to the desired target z_i, and (2) {_i}_i ∉ S to the clipping region [-1,1]^c. In other words, our desired conditions are
    
    α_1^2(_i)    = z_i,   ∀ i ∈ S, 
    α_1^2(_i)    ∈ [-1, 1]^c,   ∀ i ∉ S.

    Note that the input of the first node in the second layer is represented as
    
    α_1^2(_i) = ∑_j=1^2N+1 W_2,jβ_j^1 (_i) + b_2
 
    where _2 = [W_2,1; ⋯; W_2,2N+1] and
    
    β_j^1 (_i) = σ_H(α_j^1 (_i))
 
    is the output of node j in layer 1, when the input to the network is _i.
    Thus, the first condition in Eq. <ref>
    can be represented as a linear system 
    
    [  _2; b_2 ] 
            = [    z_i_1;        ⋮; z_i_2N+1 ]

    where i_k for k ∈ [2N+1] is defined the as the elements of S = {i_1, ⋯, i_2N+1} and
    
    = [       β_1^1 (_i_1)                ...    β_2N+1^1 (_i_1)                  1;                  ⋮                  ⋱                  ⋮                  ⋮;    β_1^1 (_i_2N+1)                ... β_2N+1^1 (_i_2N+1)                  1 ].

    In our above example, we have 
    
    = [ β_1^1 (_1)         +1         -1         +1         -1          1;         +1 β_2^1 (_4)         -1         +1         -1          1;         +1         -1 β_3^1 (_5)         +1         -1          1;         +1         -1         +1 β_4^1 (_7)         -1          1;         +1         -1         +1         -1 β_5^1 (_8)          1 ]

    Using a similar technique used in <cit.>, this matrix ∈ℝ^(2N+1) × (2N+2) satisfies two conditions:
    
    (1)  rank() = 2N+1,  
        (2)  ∃ν = [ν_1, ⋯, ν_2N+2] ∈null() 
            such that ν_i > 0   ∀ i ∈ [2N+1].

    Thus, the linear system in Eq. <ref> has infinitely many solution in the form of 
    
    [  _2; b_2 ] = μ + λν

    for any scalar λ and a particular solution μ. With the logic used in the proof of Lemma B.1 in <cit.>, we can scale λ sufficiently such that the second condition in Eq. <ref> holds. 
    Thus, by using such weight _2 and bias b_2, the input of layer 2 looks like in Fig. <ref>. Let the neuron in layer 2 has activation function 
    
    σ_T(t) =
            
                t,     if  t < 1 
    
                -1/δ (t-1-δ),     if  1 ≤ t < 1+δ
    
                0,     if  t ≥ 1+ δ

    Then, the output of the network 
    
    g_θ() = σ_T(_2 (σ_H( _1  + _1 ) ) + b_2  ) = σ_T (α_1^2 () )

    satisfies 
    
    g_θ(_i) = 
            
                z_i,       i ∈ T 
    
                0,       i ∈ [K] ∖ T

    by using the definition of σ_T and Eq. <ref>. Note that the first layer of this construction uses 2N+1 neurons with σ_H activation, and the second layer uses 1 neuron with σ_T activation. Since σ_H and σ_T can be converted into 2 ReLU neurons, our construction use total 4N+4 neurons, which completes the proof.



§ PROOF OF PROPOSITION <REF>



Recall the 3-layer network illustrated in Figure <ref>. For given T ⊂ [K], let us denote this 3-layer network satisfying (<ref>) by g_θ, T which has maximum width 4 min{√(3|T|+2), √(K)}. In what follows, we construct an L-layer network based on g_θ, T.


We partition T into ⌊L-1/2⌋ subsets: T_1, T_2, ⋯, T_⌊L-1/2⌋ satisfying |T_i| ≤ |T|/ ⌊L-1/2⌋ + 1 for all i. It can be easily seen that g_θ, T_1(x) + g_θ, T_2(x) + ⋯ + g_θ, T_⌊L-1/2⌋(x) satisfies (<ref>). Using the construction idea provided in Figure 2 of <cit.>, the above function can be represented as an L-layer network, which has a maximum width less than or equal to

    max_i{4 min{√(3|T_i|+2), √(K)} + 2},

and we conclude.







