Paper Title: Automated Review Generation Method Based on Large Language Models.



1,2,6]Shican Wu
1,2,6]Xiao Ma
1,2]Dehui Luo
1,2]Lulu Li
5]Xiangcheng Shi
1,2,3]Xin Chang
1,2]Xiaoyun Lin
1,5]Ran Luo
1,2]Chunlei Pei
1,2]Zhi-Jian Zhao*
1,2,3,4,5]Jinlong Gong*

[1]School of Chemical Engineering and Technology; Key Laboratory for Green Chemical Technology of Ministry of Education, Tianjin University, Tianjin 300072, China
[2]Collaborative Innovation Center of Chemical Science and Engineering (Tianjin), Tianjin 300072, China
[3]Haihe Laboratory of Sustainable Chemical Transformations, Tianjin 300192, China
[4]National Industry-Education Platform of Energy Storage, Tianjin University, 135 Yaguan Road, Tianjin 300350, China
[5]Joint School of National University of Singapore and Tianjin University, International Campus of Tianjin University, Binhai New City, Fuzhou 350207, Fujian, China
[6]These authors contributed equally: Shican Wu, Xiao Ma
[*]Corresponding author. Email: zjzhao@tju.edu.cn; jlgong@tju.edu.cn





Automated Review Generation Method Based on Large Language Models
    [
    August 5, 2024
=================================================================



§ ABSTRACT 
    Literature research, vital for scientific advancement, is overwhelmed by the vast ocean of available information. Addressing this, we propose an automated review generation method based on Large Language Models (LLMs) to streamline literature processing and reduce cognitive load. In case study on propane dehydrogenation (PDH) catalysts, our method swiftly generated comprehensive reviews from 343 articles, averaging seconds per article per LLM account. Extended analysis of 1041 articles provided deep insights into catalysts' composition, structure, and performance. Recognizing LLMs' hallucinations, we employed a multi-layered quality control strategy, ensuring our method's reliability and effective hallucination mitigation. Expert verification confirms the accuracy and citation integrity of generated reviews, demonstrating LLM hallucination risks reduced to below 0.5% with over 95% confidence. Released Windows application enables one-click review generation, aiding researchers in tracking advancements and recommending literature. This approach showcases LLMs' role in enhancing scientific research productivity and sets the stage for further exploration.




§ INTRODUCTION


In scientific research, peer-reviewed academic literature serves as a dense and reliable medium for information dissemination, enabling researchers to push the boundaries of human knowledge by building on previous work<cit.>. The clarity and rigor of scientific language constrain information dissemination, making it a key carrier in the research process for entity description, concept extraction, information transfer, and consensus building. This ensures that in the transmission and evolution of knowledge, both the sender and receiver of information construct highly consistent models of the referenced objects and concepts on the cognitive level. For instance, the development of industrial catalysts requires a thorough understanding of the materials' structure, chemical properties, and reactivity, considering their activity, selectivity, and stability<cit.>. This necessitates leveraging foundational catalytic theories and reaction mechanisms detailed in literature. However, the rapid pace of literature publication has outstripped researchers' capacity to assimilate information<cit.>, highlighting the need for tools to efficiently analyze and integrate data, thus avoiding redundant discoveries and broadening research perspectives, ultimately facilitating scientific research processes such as catalyst development.

Natural language processing (NLP), a branch of machine learning (ML), with core tasks like co-reference resolution and semantic analysis<cit.>, is well-suited as such a tool for literature understanding and integration. NLP technology has evolved through feature engineering, neural network architecture engineering, pre-trained language model development, and prompt-based learning<cit.>. Recent years have seen NLP applied in catalysis literature for extracting synthetic methods<cit.>, materials and properties<cit.>, key reaction parameters<cit.>, and reactions<cit.>. However, these studies often focus on isolated aspects, which limits their transferability and requires prior domain knowledge and programming skills, making it challenging for newcomers. Naturally, we considered integrating literature information in the form of reviews, which can also be generalized to more disciplines. However, early attempts at automated review generation often treated it as a multi-document summarization (MDS) task<cit.>, relying on existing reviews and citation networks<cit.>, thus may not keep pace with rapid advancements in research. Additionally, bibliometric analyses can overlook recent studies due to insufficient citations; focusing solely on abstracts or citations rather than full texts<cit.> may miss critical details. Focusing on extractive summarization rather than integrated generation<cit.>, or filling in sentence templates<cit.> can omit significant information or be redundant.

Since November 2022, Large Language Models (LLMs) like ChatGPT, representing breakthroughs in NLP, have demonstrated unprecedented language comprehension abilities<cit.>. LLMs excel in zero-shot and few-shot learning, common sense, logical reasoning, and versatility across NLP tasks<cit.>, serving as a second brain for researchers that processes and comprehends extensive scientific literature without additional foundational knowledge<cit.>. However, LLMs' creative reprocessing of understood information differs from search engines and poses the challenge of "hallucinations", a prevalent but unresolved issue in the industry<cit.>. Research reveals that calibrated LLMs inevitably generate hallucinations<cit.>, a phenomenon unavoidable in any computable LLM, regardless of model architecture, learning algorithm, prompt techniques, or training data<cit.>. Hallucinations in LLMs, which often emerge from statistical biases or noise in the training data and strategies for handling ambiguous information, manifest as baseless false information, contextually misaligned or unrelated responses<cit.>. Hallucinations are exacerbated in specialized domains by limited data exposure, leading LLMs to produce seemingly credible yet misleading outputs with specialized terminology, posing significant risks in scientific research where accuracy is critical. Even advanced LLM like GPT-4 achieve only 73.3% accuracy in professional contexts<cit.>, risking inaccurate academic conclusions and misdirected research, causing considerable time and resource losses<cit.>. Ensuring hallucination mitigation is crucial for the scientific integrity and reliability of automated review generation. Galactica<cit.>, which claimed review generation capabilities, was withdrawn due to hallucination concerns, highlights this necessity. Therefore, this study emphasizes mitigating hallucinations and enhancing reliability to safeguard the quality of LLM-generated reviews.

Addressing existing methodological limitations and leveraging LLMs' potential, this study develops an LLM-based, efficient, and comprehensive automated review generation approach. It features an end-to-end pipeline for literature retrieval, reading, summary distillation, and coherent text organization, underpinned by a multi-tier quality control strategy to counter LLM hallucination risks. LLMs' adeptness at information refinement and knowledge building enables accurate entity and concept extraction from texts, bolstering literature handling, refinement, and knowledge construction capabilities of researchers. By exploiting computational strengths in storage and parallel processing, LLMs overcome human cognitive constraints, easing researchers' cognitive load and aiding in the rapid identification of research trends. This automated review generation method thus presents a panoramic field view, offering a one-click solution for researchers lacking relevant background. Below is a comparative analysis of automated and traditional review generation methods (see Table 1):



The automated review generation method based on LLMs holds substantial scientific significance in research. It enhances literature processing efficiency and quality, fosters new knowledge discovery, and stimulates innovation. This method is invaluable in advancing the scope and depth of contemporary scientific research. As an innovative literature processing tool, it could become integral to scientific research infrastructure, significantly promoting scientific research progress.



§ RESULTS




 §.§ Automated retrieval


In our study, automated review generation essentially reprocesses the retrieved information. The process hinges on efficiently retrieving and extracting pertinent information from extensive scientific literature, with the review's quality and scope directly tied to the retrieval process's comprehensiveness and accuracy. We utilized SerpAPI for automated retrieval on Google Scholar, focusing on propane dehydrogenation (PDH) catalysts, covering literature from 1980 to 2024 in top-tier(Q1) chemistry and chemical engineering journals according to the 2022 Chinese Academy of Sciences division table.

The automated retrieval yielded 1420 initial results from Google Scholar. To address the challenge of irrelevant or duplicate findings, we implemented a dual-level filtering process. The first level employed quick filtering of abstracts and titles to remove obviously irrelevant documents, serving as a rapid but less precise narrowing method. The second level involved deeper LLM-based analysis of full texts, offering higher accuracy albeit at a slower pace. This coarse-to-fine screening method, reminiscent of high-throughput screening, enabled us to efficiently and accurately identify literature pertinent to our research. The initial screening shortlisted 343 articles as related to our topic. Subsequent LLM evaluation further confirmed 238 of these articles as relevant.



 §.§ Implementation and analysis of one-click automated review generation


Using PDH catalysts as an example and building on the aforementioned automated retrieval, we have effectively produced high-quality, specialized review articles. By focusing on top-tier journals, we ensured the retrieval of articles with significant academic impact, offering an accessible starting point for users new to the domain. For those with domain familiarity, the program allows the specification of a custom journal list to refine article selection.

We evaluated the efficacy of this method by contrasting two strategies for constructing review topics: one based on existing reviews and another using LLM-generated topics (see Table 2). The examples showcased in subsequent sections and the Supplementary Information are based on outlines derived from existing reviews.



Key benefits of the generated reviews, irrespective of the topic construction method, include:


    
  * Content Accuracy: The content has been manually checked by experts in the relevant field, with no errors in knowledge, correct referencing of cited literature, and a length and citation count that align with conventional review standards. Specific examples of generated content can be found in the Supplementary Information.
    
  * Customizable Research: Enables the addition of specific questions to tailor research focus and refine review specificity.
    
  * Forward-Looking Insights: Each topic includes a section on "Comprehensive understanding and prospective outlook", providing profound insights and innovative suggestions by LLM. This allows the LLM to engage in divergent thinking and integration beyond the information domain provided by the literature.


To improve accessibility, we developed a Python3 graphical user interface (GUI), enabling straightforward, one-click review generation on Windows, requiring no programming skills or domain knowledge.



 §.§ Data mining and visual analysis


In this study, the data mining module was deployed for comprehensive analysis in the PDH catalysts domain, examining literature from 1980 to 2024 within the chemistry and chemical engineering journals ranked Q1, Q2, and Q3 by the 2022 Chinese Academy of Sciences. Out of 1041 articles filtered by abstracts and titles, 839 were pinpointed as pertinent to PDH research via LLM selection. Leveraging LLMs for data extraction and subsequent analysis, we provided insightful conclusions on catalysts' composition, structure, and performance. This approach not only highlighted PDH research trends but also explored the maximum performance of individual factors and the synergistic effects between variables.

For instance, a statistical analysis of the annual publication numbers by catalyst types (see Figure 1 (a)) and sources of performance enhancement (see Figure 1 (b)) showed a surge in alloy research since 1995 and a spike in single-atom catalyst studies post-2015, and primarily driven by advancements in structural composition. This trend underscores the PDH field's evolving focus and hints at fresh avenues for catalyst development, including synthesis methods. In our analysis of the impact of promoter elements (see Figure 1 (c)) and support materials (see Figure 1 (d)) on catalyst performance, including selectivity and stability, we identified that promoter elements such as Zn, Sn, and La, as well as support materials like alumina and zeolites, can achieve notable peak performance, which signaled pathways for catalytic innovation. The combination analysis, for instance, of active site elements with composition elements (see Figure 1 (e)) and alloy structure types with preparation methods (see Figure 1 (f)), revealed that multi-metal systems generally outperform single-metal systems, especially when promoter elements like Sn, Zn, In are used to enhance the performance of Pt-based catalysts. Moreover, impregnation-prepared nanometallic catalysts exhibited superior conversion rates and selectivity, while single-atom alloys showed high selectivity but lower conversion rates.

This comprehensive analysis reveals the nuanced interplay between variables, guiding future research towards optimizing catalyst performance, aiding researchers in achieving the optimal performance balance in catalyst design and optimization. It suggests selecting Pt-based catalysts for maximum selectivity or metal oxides for enhanced conversion rates, and conducting deeper exploration into single-atom and nanostructured catalysts, which show promise in exceeding the efficacy of conventional catalysts. These insights not only showcase the diverse characteristics and performance benchmarks within the PDH domain but also highlight LLMs' utility in scientific exploration, providing researchers with real-time domain understanding and progress perception, thereby fostering catalyst development. This holistic approach empowers researchers to refine catalyst design and optimization effectively, aligning with industrial needs.





 §.§ Hallucination mitigation


To address the challenge of hallucinations in LLMs, a high priority has been placed on the detection and prevention of such phenomena. In the entire automated review generation process, we adopted a multi-level filtering and verification quality control strategy, similar to the concept of retrieval-augmented generation (RAG)<cit.>, to mitigate and correct hallucinations:



  §.§.§ Prompt design and task decomposition

Firstly, we utilized strict and clear text summary guiding prompts, aimed at enhancing the scientific rationality of LLM's outputs and ensuring accuracy and reliability in its analysis and generation processes. Notably, the task of automated review generation aligns well with the strengths of LLMs —information extraction and text generation capabilities. LLMs can rapidly and accurately extract core information from a vast array of literature and integrate it into a coherent and rigorous review text. To enhance efficiency and quality, we deconstructed the core of the review writing process, namely literature reading and summarization, into a series of text summarization tasks. This approach is adopted because summaries generated by LLM significantly surpass manually crafted and fine-tuned model-generated summaries in terms of fluency, factual consistency, and flexibility<cit.>. By establishing a list of questions, we directed the model to extract relevant content from the literature and respond based on this content, subsequently conducting a comprehensive analysis of all literature citations and responses. Ultimately, the LLM generates high-quality paragraphs closely related to the topic. Additionally, we employed a single-round, segmented generation strategy to avoid truncation limitations of approximately 8K output length. By reasonably segmenting long texts for generation, we not only ensured that the output was completed in a single conversational round but also provided finer parallel granularity to improve generation efficiency. In practice, we divided the 35 questions into 5 groups, ensuring that the generation results for each group could be successfully completed within the 8K limit of the LLM. This granularity avoids efficiency drops due to a high proportion of shared content and identical prompt frameworks, thereby enhancing processing speed while ensuring the quality of text generation.



  §.§.§ Hallucination filtering and verification

To mitigate and rectify hallucinations, we employed a layered filtering and verification approach:


    
  * Text format filtering: Noting that hallucinations often disrupt text formatting, we applied a predefined XML format template to filter out disarrayed texts.
    
  * DOI verification: DOIs, a combination of symbols and numbers lacking direct semantic linkage to context, present a challenge in generation and are prone to hallucinations. Yet, the precise reference nature of DOIs allows for verification. Through strict DOI verifications on generated content, we suppressed hallucinatory content from advancing further, ensuring each generated conclusion is traceable to its original source.
    
  * Relevance verification: Within the RAG system, documents related in semantics but lacking correct answers are particularly detrimental<cit.>. We scrutinized each response in the knowledge extraction phase to ensure its relevance, eliminating off-topic answers with relevant keywords.
    
  * Self-consistency<cit.> verification: For text summarization, where a definitive correct answer exists, recognizing that the stochasticity of hallucinations means correct answers should recur more frequently across iterations, we employ aggregation from repeated queries to effectively suppress hallucinations.
    
  * Full data stream traceability mechanism: By using DOIs as key reference identifiers for each piece of generated content and mandating citations for every conclusion, we enable review readers to easily trace back to the original literature, supporting verification and deeper exploration in topics of interest.




  §.§.§ Effectiveness of hallucination mitigation

In evaluating the effectiveness of hallucination mitigation, we employed a confusion matrix to classify outcomes according to whether the LLM provided content and its pertinence to the original text, differentiating between two types of inaccuracies: false positives, which include fabricated or inconsistent information, and false negatives, referring to overlooked or partially extracted content. Our focus was primarily on reducing false positives, while adopting a relatively tolerant stance on false negatives.

Substantial progress was made in mitigating hallucinations. During paragraph generation, only 36% of outputs met criteria following format and DOI validations. This was achieved through 9 repetitions of generating 35 paragraphs, cumulatively resulting in 875 generations. Analyzing 343 relevant articles, we executed 1715 information extractions across 35 questions, yielding 8575 responses and ultimately aggregating to 2783 valid information combinations. Impressively, 84.80% of these outcomes were confirmed by the LLM as 100% consistent with the aggregated results (see Table 3 and Figure 2 (a)), affirming the model's reliability. This method also establishes an rough benchmark for hallucination ratio, facilitating the selection and evaluation of LLMs.

Upon conducting manual verification on 25 articles each from the knowledge extraction and data mining stages, we calculated the accuracy, false positive rate, 95% confidence interval of the false positive rate, precision, recall, F1 score, and consistency (see Table 3). The 95% confidence interval for the false positive rate was provided by the statsmodels library in Python3. The results are detailed in the following table.





The data comparison underscores the efficacy of self-consistency verifications, revealing a substantial decrease in hallucinations, i.e., false positive content, while also compensating for some false negatives, where information was not fully extracted (see Figure 2 (b)). In the knowledge extraction phase, critical for review content, our manual sampling found no fabricated conclusions by LLMs (see Figure 2 (a)), attesting to our method's scientific integrity and reliability. From the sampling results, we are over 95% confident that the likelihood of hallucinations in this part is less than 0.5% (see Table 3). Analysis of false positives in the post-aggregation data mining phase revealed hallucinations typically involved correct numerical extraction but with errors in units or definitions. False negatives mainly stemmed from LLMs' inability to comprehend highly abstract expressions, reflecting a general LLM's limited understanding of highly specialized scientific concepts. The incidence of hallucinations in knowledge extraction was significantly lower than in data mining, as answering questions did not involve converting units and concepts, thus avoiding the most challenging part of testing an LLM's grasp of scientific knowledge. Domain-specific models enhanced by domain-adaptive pretraining (DAPT) <cit.>are poised to mitigate this issue. Opting not to fine-tune LLMs for specific domains in this study prioritizes out-of-the-box functionality and multi-domain generalization, utilizing a general LLM as the base. Comparisons between RAG and fine-tuning effects in specific domains indicate that RAG sustains efficacy with contextually new knowledge and offers a significantly lower initial cost<cit.>, aligning with our objective to support researchers' entry into diverse fields efficiently.

Considering the stringent accuracy requirements in research, increasing the number of repetitions can significantly reduce the probability of hallucinations appearing in aggregated results. Binomial probability calculations indicate that theoretically, a model with 79.09% accuracy yields aggregated prediction accuracies of 93.49%, 96.12%, and 97.64% after five, seven, and nine independent predictions, respectively, aligning with our sampling results (see Table 3). Detailed sampling outcomes and calculations are available in the SI.



On this foundation, every conclusive description in the generated reviews is supported by literature references and has been verified by relevant field researchers through tracing the cited literature, confirming that all literature references are correctly linked to the original publications and that the descriptions in the generated reviews correspond to those in the original publications.

This multi-layered strategy for hallucination control has built an effective verification system, ensuring the scientific integrity and reliability of the automated review generation. Furthermore, through a full data stream traceability mechanism, the authenticity and practicality of the content are further strengthened. This not only provides a secondary means of hallucination mitigation but also allows researchers to delve into original research papers for more precise and detailed academic information while accessing fast, automated research reviews. The strategy also implements a kind of literature recommendation mechanism. Since each content segment includes related DOIs, researchers can quickly locate specific original literature based on their interests and research needs, enabling deeper academic exploration.



§ DISCUSSION


In this study, we introduce an innovative LLM-based automated review generation method, adeptly addressing two key scientific challenges: streamlining literature review efficiency and significantly reducing LLM hallucination risks. This modular, comprehensive, end-to-end solution integrates modules for literature search, topic formulation, knowledge extraction, and review composition, transforming an extensive corpus of scientific literature into coherent, detailed, and error-free reviews tailored to specific research themes. Notably, our advanced data mining module offers experienced users an in-depth field overview, exploiting the LLM's analytical prowess. Additionally, a user-friendly one-click program on Windows platforms significantly simplifies the review generation process.

A pivotal achievement of our method is its capacity to surpass traditional human resource limitations. Our rigorous quality assurance solution, encompassing format filtering, DOI verification, relevance verification and self-consistency verifications, ensures high reliability and traceability throughout the data processing pipeline. Expert evaluations with a case study of PDH catalysts confirm the method's efficacy, with reviews paralleling manual ones in length and citations, but without hallucinations and with impeccable reference accuracy. Through rigorous testing, including the analysis of 875 LLM outputs from a sample of 25 articles, we demonstrate over 95% confidence in reducing the hallucination probability to below 0.5% (see Table 3).

Our method's modular design offers excellent reusability and scalability. Individual modules like literature search, topic formulation, and knowledge extraction can serve various research purposes, like literature tracking, research topic discovery, and data mining datasets construction. Future enhancements will focus on augmenting LLM's comprehension of scientific concepts through pan-scientific field fine-tuning, elevating the method's overall utility. Planned upgrades include improving multimodal processing, automating scientific inquiries, personalizing text generation, and delving deeper into specific research areas.

In summary, our method signifies a major advancement in scientific research tools, offering rapid access to field breakthroughs and developments. It's set to transform the landscape of scientific research, with far-reaching implications for knowledge base construction, literature recommendation, and structured academic writing, heralding a new era in scientific research productivity and interdisciplinary collaboration.



§ METHODS


The method for constructing review articles consists of four parts: literature search, topic formulation, knowledge extraction and review composition, along with an additional data mining module for experienced users (see Figure 3).





 §.§ Literature search


Initially, a list of journals designated for the set review topic's subject area is obtained from journal classification tables. Then, literature containing specified keywords within these selected journals is retrieved via search engine's API. This is followed by a preliminary filter, checking each title and abstract for intersections with a selected list of keywords. Literature with intersections is saved, and those of a review nature are marked (see Figure 3 (i)). Our method supports various types of textual literature, including journals, patents, conference papers, books, etc. This means that any content in textual form can be included in the search scope, further expanding the application scenarios and coverage of our method. In our example, using "propane dehydrogenation" as a keyword, we retrieved 343 publications in top-tier Chemistry and Chemical Engineering journals (2022 Chinese Academy of Sciences classification), including 14 reviews, after filtering titles and abstracts with keywords like "propane dehydrogenation", "PDH", "ODH", "Oxidative Dehydrogenation", etc., through SerpAPI on Google Scholar.



 §.§ Topic formulation


There are two approaches to constructing review topics: one involves LLM directly drafting the outline, and the other is based on LLM refining and drafting outlines from existing literature reviews. After obtaining a list of topics, additional topics can be manually added and sorted as needed (see Figure 3 (ii)). In our example, the Claude2 model generated an outline including 12 topics directly, and another with 9 topics and 35 guiding questions based on existing review articles (see Table 2).



 §.§ Knowledge extraction


Based on the obtained list of topics, the LLM generates a list of questions for extracting information from literature, corresponding to each review topic. After repeating this process for multiple times for each article, all answers are concatenated. The LLM then determines whether the answers are relevant to the questions and aggregates them (see Figure 3 (iii)). In our example, in the case of PDH, after transforming the 35 guiding questions into questions for extracting information from literature, the Claude2 model was used to extract information from 343 top-tier articles five times, leading to the aggregation of 8575 responses into 2783 valid information combinations.



 §.§ Review composition


After associating each article's answers with their source DOI, paragraphs are generated and integrated for each topic. The LLM generates review paragraphs from all the answers combined, followed by summarization and outlook. After repeating multiple times, the LLM scores the generated paragraphs, selecting the best ones for each topic to form a preliminary draft of the full review. The full text is then polished with the help of the LLM, adjusting and checking citation formats to produce the final draft (see Figure 3 (iv)). In our example, each question's answers from various articles were combined into JSON format information groups, inputted into the Claude2 model for paragraph generation, integrated to form smooth paragraphs, repeated 9 times, scored by the Claude2 model based on criteria (as shown in SI), and polished to produce the final draft.



 §.§ Data mining


Based on the automated review generation method described above, we proposed a data mining method based on LLMs, catering to users with some domain knowledge. This method is almost identical to the knowledge extraction steps (see Figure 3 (ii)), effectively extracting and aggregating specific data from a large volume of literature. Users first define specific data extraction targets, which may include but are not limited to catalyst types, chemical compositions, and performance characteristics. On the established literature dataset, the LLM parses each article, extracting the user-defined target data multiple times and outputting in XML format. Similar to the knowledge extraction process (see Figure 3 (ii)), the LLM aggregates results from multiple extractions to finalize each article's information for each extraction target. The extracted data often require manual cleaning and processing, including correcting extraction errors, standardizing data formats, and removing redundant information to facilitate subsequent statistical analyses. The cleaned data is then further integrated and analyzed to form visual charts. The code for the cleaning process and chart statistics can be generated by GPT4, requiring no programming background from the user. In our example, in the case of PDH catalysts, relevant literature from tiers one, two, and three was downloaded using the literature search module, filtered through abstracts and titles totaling 1041 articles, of which 839 were deemed PDH-related by the Claude2 model. After data cleaning and processing through Python3 programming, the extracted data included catalyst types, composition elements, active species elements, promoter elements, support materials, alloy structural types, alloy preparation methods, propane partial pressure, reaction temperature, inlet flow rate, selectivity, conversion, selectivity and other key indicators, covering seven categorical variables such as structure and element composition and three continuous variables related to reaction conditions. Subsequently, corresponding charts were generated through code execution. Initially, we tallied the annual publication numbers of various catalyst influencing factors, represented in line or Gantt charts. Furthermore, we calculated the average of the top five selectivity and stability for all influencing factors across all catalyst data, visualized in radar charts to showcase the peak performance achievable by a specific factor. Lastly, through pairwise combination analysis of influencing factors, we produced 45 bivariate correlation bubble charts, intuitively demonstrating how different variable combinations affect overall catalyst performance. These bubble charts use bubble color intensity, size, and border thickness to represent the levels of selectivity, conversion, and stability, respectively.



§ DATA AVAILABILITY

tocsectionData availability

Our study leverages a dataset compiled from scientific literature acquired through our institution's subscription. Due to copyright considerations, the dataset itself cannot be made publicly available. However, we ensure that our research's integrity and reproducibility do not rely on direct access to these proprietary documents. Instead, we provide extensive documentation on the dataset's structure, the criteria used for literature selection, and the analysis methods applied, enabling interested researchers to reconstruct a similar dataset from publicly available resources or their institutional subscriptions.

Furthermore, to facilitate a deeper understanding of our research process and promote further exploration and innovation, we have made all intermediate data, excluding the copyrighted full-text articles, publicly available on GitHub [<https://github.com/TJU-ECAT-AI/AutomaticReviewGenerationData>]. This repository includes the prompts used in our study and the corresponding responses generated by the large language model. By sharing these resources, we aim to provide valuable insights into our methodology and encourage other researchers to build upon our work, advancing the field of natural language processing and its applications in scientific literature analysis.



§ CODE AVAILABILITY

tocsectionCode Availability

The custom code developed for this research is central to our conclusions and is made available to ensure transparency and reproducibility of our results. The codebase, including all relevant custom scripts and mathematical algorithms, has been open-sourced under the Apache 2.0 license and is accessible via our GitHub repository at [<https://github.com/TJU-ECAT-AI/AutomaticReviewGeneration>]. We encourage users to review the license for any usage restrictions that may apply. As stated in the text, all LLMs invoked in this article are Claude2.

It is important to note that our published graphical user interface (GUI) leverages certain APIs for functionality, which, due to legal and regulatory requirements, necessitate that users provide their own API keys. This requirement is detailed in the documentation accompanying the code repository to assist users in setting up and utilizing the GUI effectively.



§ ACKNOWLEDGMENTS

tocsectionAcknowledgments

The authors acknowledge the financial support from the National Key R&D Program of China (2022YFE0102000), the National Natural Science Foundation of China (No. 22121004, U22A20409), the Haihe Laboratory of Sustainable Chemical Transformations, the Program of Introducing Talents of Discipline to Universities (BP0618007) and the XPLORER PRIZE. The authors also appreciate the generous computing resources provided by the High Performance Computing Center of Tianjin University.



§ AUTHOR CONTRIBUTIONS

tocsectionAuthor contributions

J.G and Z.Z. conceived and supervised the project. S.W. and X.M. designed the research. S.W. led the manuscript drafting. D.L. and L.L. validated the scientific reliability of the method. X.S. and X.C. innovated at the intersection of computational models and catalytic science. X.L., R.L., and C.P. conducted data analysis and validation. All authors wrote and revised this manuscript.



§ COMPETING INTERESTS

tocsectionCompeting interests

The authors declare no competing interests.



§ ADDITIONAL INFORMATION

tocsectionAdditional information

Supplementary information is available for this paper.

Correspondence and requests for materials should be addressed to J.G.

50=0pt

    1
    Ermel APC, Lacerda DP, Morandi MIW, Gauss L. Literature reviews: modern methods for investigating scientific and technological knowledge. Springer Nature (2021).

    2
    Nørskov JK, Bligaard T, Rossmeisl J, Christensen CH. Towards the computational design of solid catalysts. Nature chemistry 1, 37-46 (2009).

    3
    Zhao Z-J et al. Theory-guided design of catalytic materials using scaling relationships and reactivity descriptors. Nature Reviews Materials 4, 792-804 (2019).

    4
    Resasco J, Abild-Pedersen F, Hahn C, Bao ZN, Koper MTM, Jaramillo TF. Enhancing the connection between computation and experiments in electrocatalysis. Nature Catalysis 5, 374-381 (2022).

    5
    Taylor HS. A theory of the catalytic surface. Proceedings of the Royal Society of London Series A, Containing Papers of a Mathematical and Physical Character 108, 105-111 (1925).

    6
    Mou TY et al. Bridging the complexity gap in computational heterogeneous catalysis with machine learning. Nature Catalysis 6, 122-136 (2023).

    7
    Lawrence S. Free online availability substantially increases a paper's impact. Nature 411, 521 (2001).

    8
    Lok C. Speed reading: scientists are struggling to make sense of the expanding scientific literature. Corie Lok asks whether computational tools can do the hard work for them. Nature 463, 416-419 (2010).

    9
    Khurana D, Koli A, Khatter K, Singh S. Natural language processing: state of the art, current trends and challenges. Multimed Tools Appl 82, 3713-3744 (2023).

    10
    Liu PF, Yuan WZ, Fu JL, Jiang ZB, Hayashi H, Neubig G. Pre-train, Prompt, and Predict: A Systematic Survey of Prompting Methods in Natural Language Processing. Acm Comput Surv 55, 1-35 (2023).

    11
    Wang L et al. A corpus of CO(2) electrocatalytic reduction process extracted from the scientific literature. Sci Data 10, 175 (2023).

    12
    Plata DL, Jankovic NZ. Achieving sustainable nanomaterial design though strategic cultivation of big data. Nat Nanotechnol 16, 612-614 (2021).

    13
    Cruse K et al. Text-mined dataset of gold nanoparticle synthesis procedures, morphologies, and size entities. Sci Data 9, 234 (2022).

    14
    Batra SR. Emerging materials intelligence ecosystems propelled by machine learning. Nature Reviews Materials, (2020).

    15
    Suvarna M, Vaucher AC, Mitchell S, Laino T, Perez-Ramirez J. Language models and protocol standardization guidelines for accelerating synthesis planning in heterogeneous catalysis. Nat Commun 14, 7964 (2023).

    16
    Vaucher AC, Zipoli F, Geluykens J, Nair VH, Schwaller P, Laino T. Automated extraction of chemical synthesis actions from experimental procedures. Nat Commun 11, 3601 (2020).

    17
    Boschen I. Evaluation of JATSdecoder as an automated text extraction tool for statistical results in scientific reports. Sci Rep 11, 19525 (2021).

    18
    Nandy A, Terrones G, Arunachalam N, Duan C, Kastner DW, Kulik HJ. MOFSimplify, machine learning models with extracted stability data of three thousand metal-organic frameworks. Sci Data 9, 74 (2022).

    19
    Olivetti EA et al. Data-driven materials research enabled by natural language processing and information extraction. Appl Phys Rev 7, (2020).

    20
    Huang S, Cole JM. A database of battery materials auto-generated using ChemDataExtractor. Sci Data 7, 260 (2020).

    21
    Court CJ, Cole JM. Magnetic and superconducting phase diagrams and transition temperatures predicted using text mining and machine learning. Npj Computational Materials 6, (2020).

    22
    He TJ et al. Similarity of Precursors in Solid-State Synthesis as Text-Mined from Scientific Literature. Chemistry of Materials 32, 7861-7873 (2020).

    23
    Kim E et al. Machine-learned and codified synthesis parameters of oxide materials. Sci Data 4, 170127 (2017).

    24
    Vaucher AC, Schwaller P, Geluykens J, Nair VH, Iuliano A, Laino T. Inferring experimental procedures from text-based representations of chemical reactions. Nat Commun 12, 2573 (2021).

    25
    Kim E, Huang K, Saunders A, McCallum A, Ceder G, Olivetti E. Materials Synthesis Insights from Scientific Literature via Text Extraction and Machine Learning. Chemistry of Materials 29, 9436-9444 (2017).

    26
    Mahbub R, Huang K, Jensen Z, Hood ZD, Rupp JLM, Olivetti EA. Text mining for processing conditions of solid-state battery electrolyte. Electrochemistry Communications 121, 106860 (2020).

    27
    Scheffler M et al. FAIR data enabling new horizons for materials research. Nature 604, 635-642 (2022).

    28
    Pesciullesi G, Schwaller P, Laino T, Reymond JL. Transfer learning enables the molecular transformer to predict regio- and stereoselective reactions on carbohydrates. Nat Commun 11, 4874 (2020).

    29
    Ma C, Zhang WE, Guo M, Wang H, Sheng QZ. Multi-document summarization via deep learning techniques: A survey. Acm Comput Surv 55, 1-37 (2022).

    30
    Nikiforovskaya A, Kapralov N, Vlasova A, Shpynov O, Shpilman A. Automatic generation of reviews of scientific papers. In: 2020 19th IEEE International Conference on Machine Learning and Applications (ICMLA). IEEE (2020).

    31
    Mohammad S et al. Using citations to generate surveys of scientific paradigms. In: Proceedings of human language technologies: The 2009 annual conference of the North American chapter of the association for computational linguistics (2009).

    32
    Agarwal N, Reddy RS, Kiran G, Rose C. Towards multi-document summarization of scientific articles: making interesting comparisons with SciSumm. In: Proceedings of the workshop on automatic summarization for different genres, media, and languages (2011).

    33
    Jaidka K, Khoo C, Na J-C. Deconstructing human literature reviews–a framework for multi-document summarization. In: proceedings of the 14th European workshop on natural language generation (2013).

    34
    Kasanishi T, Isonuma M, Mori J, Sakata I. SciReviewGen: A Large-scale Dataset for Automatic Literature Review Generation. arXiv preprint arXiv:230515186, (2023).
    
     35
    Liu Y et al. Summary of chatgpt-related research and perspective towards the future of large language models. Meta-Radiology, 100017 (2023).

    36
    Naveed H et al. A comprehensive overview of large language models. arXiv preprint arXiv:230706435, (2023).

    37
    White AD. The future of chemistry is language. Nature Reviews Chemistry 7, 457-458 (2023).

    38
    Zhang Y et al. Siren's Song in the AI Ocean: A Survey on Hallucination in Large Language Models. arXiv preprint arXiv:230901219, (2023).

    39
    Sanderson K. GPT-4 is here: what scientists think. Nature 615, 773-773 (2023).

    40
    Kalai AT, Vempala SS. Calibrated Language Models Must Hallucinate. arXiv preprint arXiv:231114648, (2023).

    41
    Xu Z, Jain S, Kankanhalli M. Hallucination is inevitable: An innate limitation of large language models. arXiv preprint arXiv:240111817, (2024).

    42
    Wu S et al. A comparative study of open-source large language models, gpt-4 and claude 2: Multiple-choice test taking in nephrology. arXiv preprint arXiv:230804709, (2023).

    43
    Taylor R et al. Galactica: A large language model for science. arXiv preprint arXiv:221109085, (2022).

    44
    Lewis P et al. Retrieval-augmented generation for knowledge-intensive nlp tasks. Advances in Neural Information Processing Systems 33, 9459-9474 (2020).

    45
    Truhn D, Reis-Filho JS, Kather JN. Large language models should be used as scientific reasoning engines, not knowledge databases. Nature Medicine, 1-2 (2023).

    46
    Pu X, Gao M, Wan X. Summarization is (almost) dead. arXiv preprint arXiv:230909558, (2023).

    47
    Cuconasu F et al. The power of noise: Redefining retrieval for rag systems. arXiv preprint arXiv:240114887, (2024).

    48
    Wang X et al. Self-consistency improves chain of thought reasoning in language models. arXiv preprint arXiv:220311171, (2022).

    49
    Gururangan S et al. Don't stop pretraining: Adapt language models to domains and tasks. arXiv preprint arXiv:200410964, (2020).

    50
    Gupta A et al. RAG vs Fine-tuning: Pipelines, Tradeoffs, and a Case Study on Agriculture. arXiv preprint arXiv:240108406, (2024).





 Paper Title: Synthetic Simplicity: Unveiling Bias in Medical Data Augmentation.












Indian Institute of Technology Madras, India 
 GE HealthCare, Bangalore, India
	
	

Krishan Agyakari Raja Babu et al.



Synthetic Simplicity: Unveiling Bias in Medical Data Augmentation
    Krishan Agyakari Raja Babu1
 Rachana Sathish2Mrunal Pattanaik2
Rahul Venkataramani2{rahul.venkataramani}@gehealthcare.com
    
=============================================================================================================================



§ ABSTRACT 
Synthetic data is becoming increasingly integral in data-scarce fields such as medical imaging, serving as a substitute for real data. However, its inherent statistical characteristics can significantly impact downstream tasks, potentially compromising deployment performance. In this study, we empirically investigate this issue and uncover a critical phenomenon: downstream neural networks often exploit spurious distinctions between real and synthetic data when there is a strong correlation between the data source and the task label. This exploitation manifests as simplicity bias, where models overly rely on superficial features rather than genuine task-related complexities. Through principled experiments, we demonstrate that the source of data (real vs. synthetic) can introduce spurious correlating factors leading to poor performance during deployment when the correlation is absent. We first demonstrate this vulnerability on a  digit classification task, where the model spuriously utilizes the source of data instead of the digit to provide an inference. We provide further evidence of this phenomenon in a medical imaging problem related to cardiac view classification in echocardiograms, particularly distinguishing between 2-chamber and 4-chamber views. Given the increasing role of utilizing synthetic datasets, we hope that our experiments serve as effective guidelines for the utilization of synthetic datasets in model training.









§ INTRODUCTION


The acute challenge of data scarcity in medical imaging has spurred the development and application of various methods, including transfer learning, self-supervised learning, and data augmentation. More recently, synthetic data, either alone or combined with real data, has emerged as a popular technique to address the inherent data scarcity in medical imaging. The growing prominence of synthetic data is largely attributed to the efficiency of advanced generative techniques, such as generative adversarial networks (GANs), diffusion models  <cit.>, which produce high-quality visual images with relatively straightforward training. As visual generative models continue to improve, researchers are increasingly exploring the utility of generated synthetic samples to augment real data  <cit.>, aiming to achieve fairness and generalization.

Since the images produced by generative models appear extremely realistic, deep learning practitioners have increasingly adopted the use of synthetic data as a substitute of real data. In this paper, we investigate whether synthetic images possess statistical signatures, which may perhaps be inconspicuous visually, specific to their generation algorithms, distinguishing them from those found in real images. These distinctive features could potentially be exploited by downstream algorithms. It is well-known that discriminative deep learning methods tend to exploit the simplest features while disregarding other relevant aspects—a phenomenon referred to as shortcut learning due to “simplicity bias”  <cit.>. Medical imaging is particularly prone to these scenarios as data for a particular class could be hard to obtain (e.g. pathologies). Since the presence of “simplicity bias" can lead to sub-optimal performance during deployment, it is necessary for model developers to be cognizant of the perils of indiscriminate use of synthetic data during model development.

In summary, our contributions are as follows:

	
  * We conceive the possibility of high correlations between source of data and target resulting in deep learning model exploiting non-task specific features leading to poor performance during deployment.
	
  * We expose, with experiments on controlled datasets covering both natural images (digit classification) and medical imaging tasks (echocardiograph view classification), the problem of models exploiting the source of dataset (real vs. synthetic) as opposed to the task-specific causal feature.






§ RELATED WORK


Diffusion Probabilistic Models (DDPM),  <cit.>, have demonstrated remarkable performance in generating high-quality image samples. DDPMs leverage iterative denoising processes to model complex data distributions effectively. The core idea of DDPM revolves around iteratively applying a denoising process to generate samples from a target distribution. To enhance control over semantic details, recent advancements in generative modeling have focused on conditioning models on object-level segmentation masks. A semantic diffusion model (SDM) was proposed in a recent study  <cit.> for generating synthetic 2D echocardiograms. In this study, we adopt the SDM framework as our model for generating synthetic images.


The promise shown by DDPMs, SDMs and similar models have spurred the utilization of synthetic data  to enrich the training dataset <cit.>, anomaly detection <cit.>, super resolution <cit.> etc.
Due to the increasing utilization of synthetic datasets in medical imaging, it is imperative to critically analyze the potential pitfalls carefully. Particularly, the phenomenon of simplicity bias has been observed in discriminative deep learning methods, where models tend to exploit spurious, low-level statistical patterns that correlate significantly with the target rather than high-level semantic features, leading to poor out-of-distribution (OOD) performance . This issue is exacerbated in scenarios where synthetic data is used to augment real datasets. Particularly, three implications of the simplicity bias have been discussed by <cit.>—lack of robustness, lack of reliable confidence estimates, and suboptimal generalization.  While the dangers of shortcut learning in the context of healthcare <cit.> have been discussed, there is no prior work in the context of spurious features manifested due to utilization of synthetic data.




§ METHODOLOGY


In order to explore the manifestation of synthetic simplicity in medical data augmentation, we consider the binary classification problem on a dataset D = {(X_1, y_1),(X_2, y_2),…,(X_N, y_N)} comprising of N samples, where input images X∈ℝ^H × W and labels y ∈{-1, +1}.Each image can be considered to possess M latent features represented as f ∈ℝ^M which can be leveraged for classification. One of the features can be assumed to be a binary feature f_k ∈{0, 1}, where k denotes the feature index, representing whether a given sample is real or generated synthetically. We intentionally set the correlation between this binary feature and the output label to be high during training and evaluate during the testing phase when the correlation is lower. The test set consists of samples of the form (X, y = 1, f_k = 0), indicating the low correlation between y=1 and f_k=0. We vary the degree of correlation using the parameter α to explore different scenarios of synthetic simplicity.

So, with this configuration, the model during training has the potential to learn the simpler but spurious feature of real vs. synthetic. Since the correlation between the task label and data source is quite high, this will result in excellent performance during training. However, during evaluation when the correlation is lower, the performance may degrade if the simplicity bias is manifested. This test allows us to systematically investigate if the neural network utilizes the rest of the features f_i ∖ k or is predisposed to utilize feature f_k due to it being simpler than other features.




§ EXPERIMENTS




 §.§ Synthetic Data Engineering
In this section, we will describe the dataset generation for both the digit classification and the echocardiographic view classification problem. The classification network was trained using real data augmented by the synthetic data in varying proportions to study the impact of the synthetic simplicity. Let D^Tr_R,C_1 denote the set of all real images of class 1, D^Tr_S,C_1 denote the set of all synthetic images of class 1, D^Tr_R,C_2 denote the set of all real images of class 2, and D^Tr_S,C_2 denote the set of all synthetic images of class 2. The composition of the training data is thus defined as follows:


    |D^Tr| = α |D^Tr_R,C_1| + (1 - α) |D^Tr_S,C_1| + (1 - α) |D^Tr_R,C_2| + α |D^Tr_S,C_2|


where α∈ [0, 1] is a hyperparameter and α|D^Tr_R,C_1| + (1-α) |D^Tr_S,C_1| = (1-α)|D^Tr_R,C_2| + α|D^Tr_S,C_2|=k_Tr. The test set has an equal number of real and synthetic images for both classes, i.e., |D^Ts_R,C_1| = |D^Ts_S,C_1| = |D^Ts_R,C_2| = |D^Ts_S,C_2|=k_Ts.

We design the experiments to include a mix of real and synthetic images of class 1 and 2 in the training dataset D^Tr, controlled by the hyperparameter α.
When α=1, the training dataset D^Tr consists entirely of real images of class 1 (D^Tr_R,C_1) and synthetic images of class 2 (D^Tr_S,C_2). When α = 0, the training dataset D^Tr consists entirely of synthetic images of class 1 (D^Tr_S,C_1) and real images of class 2 (D^Tr_R,C_2).

For values of α between 0 and 1, the training dataset D^Tr is a mix of real and synthetic images of both classes 0 and 1, with α determining the specific proportion. If the model is able to maintain consistent performance regardless of the value of α, it can be considered robust. Otherwise, it may be affected by synthetic simplicity.



  §.§.§ Downstream Network
- In the case of digit classification task, we employed a simple Convolutional Neural Network (CNN) <cit.> as the downstream task. The model was trained on 6000 images each of digits 2 and 4 (k_Tr=6000) , using the Adam optimizer with a learning rate of 0.001, a batch size of 128, and the cross-entropy loss criterion, over the course of 50 epochs. Evaluation was conducted on a separate dataset comprising 1000 real and synthetic images for each digit 2 and 4 (k_Ts=1000).

Similarly, for the echocardiogram view classification task, we utilized ResNet-18 <cit.> as the downstream network. The model was trained on 400 views each of the two-chamber and four-chamber echocardiograms (k_Tr=400) using the Adam optimizer with a learning rate of 0.001, a batch size of 64, and the cross-entropy loss criterion. The training was conducted for 15 epochs. Evaluation was performed on a separate test set containing 100 real and synthetic views each for the two-chamber and four-chamber echocardiograms (k_Ts=100).




  §.§.§ Diffusion Model
- To generate synthetic digits, we employed the open-source generative diffusion  <cit.> model. And, in order to generate synthetic samples for echocardiograms, we utilized a pre-trained modified Semantic Diffusion Model (SDM) provided by  <cit.>. The synthetic samples were created based on the conditioning of semantic segmentation maps within the SDM framework.

We describe the two datasets used for experiments in the sections below: 



 §.§ MNIST Digit Classification

We utilized the publicly available "Modified National Institute of Standards and Technology" (MNIST) <cit.> handwritten digits dataset for our experiment. It consists of 70,000 28×28 black-and-white images of handwritten digits. For our binary classification experiment, we focused on the “2" vs./ “4" classes in the MNIST dataset. Figure <ref> depicts the real and synthetic samples of the digits 2 and 4. 





 §.§ CAMUS Echocardiogram View Classification




After investigating the issue of synthetic simplicity in the digit classification task, we demonstrate the phenomenon on a representative medical imaging task of cardiac view classification from echocardiographs. We utilized the publicly available echocardiography dataset "Cardiac Acquisitions for Multi-structure Ultrasound Segmentation" (CAMUS)  <cit.> for this experiment. The CAMUS dataset comprises four types of cardiac views from over 500 subjects: two-chamber end-systolic (2CH ES), two-chamber end-diastolic (2CH ED), four-chamber end-systolic (4CH ES), and four-chamber end-diastolic (4CH ED). We focused on the "4CH ED" and "2CH ED" classes in the CAMUS dataset for classification task. Figure <ref> shows real and synthetic samples of 4CH ED and 2CH ED echocardiogram views. 



§ RESULTS AND DISCUSSION






Experimental results for the CAMUS and MNIST classification tasks are presented in Figures <ref> and <ref>, respectively. The following observations can be noted:


	
  * Poor Performance in Extreme Cases of α: In the extreme case of α = 1 (high positive correlation), the training dataset consists solely of real 2CH images and synthetic 4CH images. Under this condition, the model performs well on real 2CH images (Fig. <ref>(a)) and synthetic 4CH images (Fig. <ref>(d)). However, due to the simplicity bias, the model's performance on synthetic 2CH images (Fig. <ref>(b)) and real 4CH images (Fig. <ref>(c)) is nearly zero. 
	Similarly, in the extreme case of α = 0 (high negative correlation), the model performs poorly on real 2CH images (Fig. <ref>(a)) and synthetic 4CH images (Fig. <ref>(d)). The MNIST results in Fig. <ref> show a similar trend.
	
	
  * Implication of Synthetic Simplicity in Medical vs./ General Domains: We can observe that the drop in test accuracy for CAMUS in all four cases is significantly steeper and more sudden compared to MNIST, where the drop is more gradual.  We hypothesize that this variation between echocardiographic classification vs./ digit classification is due to the fact that the complexity of the former task is arguably more complex, which biases the network to rely more heavily on the simpler feature - source of data. This indicates that the impact of synthetic simplicity has more severe implications in the medical context compared to general contexts. The complexity of the feature set in medical image data likely contributes to this pronounced effect. This underscores the urgent need to address simplicity bias in the medical domain, where the stakes are considerably higher due to the critical nature of the applications. Ensuring that models generalize well and do not overly rely on simplistic and spurious features is crucial in medical imaging to avoid potentially harmful misclassification. 
	
	
  * Case of Balanced Augmentation: When the correlation between the task labels and the data source is set to around 0.5, representing a balanced case, the test accuracy is quite consistent and is not significantly affected by the value of α. This can be observed in both the CAMUS (Fig. <ref>) and MNIST (Fig. <ref>) results. This observation suggests that in order to mitigate the simplicity bias  while using the synthetic data, we need to adopt a balanced approach when augmenting synthetic data with real data to achieve robust evaluation. A balanced approach ensures that the model does not overly rely on spurious features that are easy to learn but irrelevant to the actual task. By maintaining an intermediate level of correlation, the model is encouraged to utilize more complex and relevant features, leading to better generalization.
	
	
  * Persistence of Synthetic Simplicity Across Different Downstream Networks:
	In our study, we employed a basic CNN architecture for digit classification and a more intricate ResNet model for echocardiogram view classification as the downstream network. Despite ResNet's increased complexity, we consistently observed synthetic simplicity across both tasks. This phenomenon persisted irrespective of the network's architecture and depth. Even with its deeper layers, ResNet tended to exploit simplistic and spurious features inherent in synthetic data, leading to diminished performance during testing. These findings highlight the pervasive nature of synthetic simplicity in downstream tasks and emphasize the need for cautious integration of synthetic data for augmentation purposes.






§ CONCLUSION


We clearly demonstrate that downstream tasks can leverage the binary feature of the dataset source (real v/s synthetic) for tasks such as classification. Therefore, it is crucial for practitioners to use synthetic data carefully during model development. This caution helps prevent the inclusion of spurious features and the generation of artificially high accuracy values on a test set derived from the same distribution. We hope that this work will stimulate more interest in understanding the risks associated with using synthetic data without domain expertise and encourage further research.







 splncs04
 



 Paper Title: Accelerating Full Waveform Inversion By Transfer Learning.





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

1113



 



§ ABSTRACT



Full waveform inversion (FWI) is a powerful tool for reconstructing material fields based on sparsely measured data obtained by wave propagation. For specific problems, discretizing the material field with a neural network (NN) improves the robustness and reconstruction quality of the corresponding optimization problem. We call this method NN-based FWI. Starting from an initial guess, the weights of the NN are iteratively updated to fit the simulated wave signals to the sparsely measured data set. For gradient-based optimization, a suitable choice of the initial guess, i.e., a suitable NN weight initialization, is crucial for fast and robust convergence. 

In this paper, we introduce a novel transfer learning approach to further improve NN-based FWI. This approach leverages supervised pretraining to provide a better NN weight initialization, leading to faster convergence of the subsequent optimization problem. Moreover, the inversions yield physically more meaningful local minima. The network is pretrained to predict the unknown material field using the gradient information from the first iteration of conventional FWI. In our computational experiments on two-dimensional domains, the training data set consists of reference simulations with arbitrarily positioned elliptical voids of different shapes and orientations. We compare the performance of the proposed transfer learning NN-based FWI with three other methods: conventional FWI, NN-based FWI without pretraining and conventional FWI with an initial guess predicted from the pretrained NN. Our results show that transfer learning NN-based FWI outperforms the other methods in terms of convergence speed and reconstruction quality. 




Keywords: 
transfer learning, neural network, deep learning, adjoint optimization, Full waveform inversion






§ INTRODUCTION



Civil infrastructure is often deteriorated due to exposure to loads, temperature fluctuations, natural disasters, and potential faults during construction. Assessing the damage present in the infrastructure requires non-destructive testing (NDT) <cit.>. In recent years, attempts have been made to exploit the potential of full waveform inversion (FWI) in NDT. Originating in seismology <cit.>, FWI infers material models of the earth's interior by fitting simulated seismic waves to measured data sets from active or passive sources <cit.>. The approach is adopted in engineering applications <cit.> utilizing ultrasonic waves to reveal information about unknown defects such as voids, cracks, and inclusions. Extracting detailed information about damage location and severity using FWI is difficult, as the problem is ill-posed and computationally expensive <cit.>.


Deep learning-based approaches are investigated to enrich FWI, potentially offering a chance to overcome the challenges mentioned above. There are several methods to incorporate neural networks (NNs) into FWI, the most common being using NNs in a supervised manner. In this context, the NN learns to map time series signals, such as displacements or pressures at sensor locations within the domain, to the corresponding field of material coefficients, often referred to as distribution, which describes wave velocities or density <cit.>. Another approach uses NNs to enrich the signals with low-frequency data <cit.>. Kleman et al. <cit.> use an NN to predict the actual velocity distribution by inputting the velocity distribution after the first iteration of conventional FWI. A common challenge with supervised approaches is the need for large training data sets. To address this limitation, generative adversarial networks can be utilized <cit.>. However, due to the absence of an accuracy measure in the supervised learning approach, the resulting predictions are often unreliable, thereby limiting their practical applicability.


Another way to take advantage of deep learning methods in FWI is to use them with conventional FWI <cit.>. Herrmann et al. <cit.> show that inflating the optimization space by using a massively overparameterized NN to predict the coefficients of a piecewise constant material field can serve as an implicit regularization of the optimization problem. Solving the corresponding problem with the Adam optimizer results in fast convergence to physically reasonable local minima. One reason for this may be the inherent bias of NN to learn the low-frequency content of the data <cit.>. Therefore, the algorithm converges to an image that has fewer high-frequency artifacts. The gradients used for this case are calculated first by the continuous adjoint method, where the derivative of the cost function with respect to the coefficients of the piecewise constant material field is calculated. Then, the gradient of the cost function with respect to the NN parameters is determined by automatic differentiation (see the gradient computation in <Ref>). In the sequel, we will refer to this approach of using the NN as a parameterization of the material field as NN-based FWI.


It is possible to further improve NN-based FWI using transfer learning. Transfer learning is a technique in machine learning where an NN trained on one data set is used for a different but related learning task. By doing so, the NN starts the downstream task with weights already adjusted to recognize relevant features, significantly reducing the training time. Simply put, transfer learning uses the knowledge from previous data to accelerate learning for related tasks. Transfer learning has already been utilized in various fields, such as for the classification of text <cit.> and images <cit.>, machine translation <cit.> and robotics <cit.>. This approach can also be used in conjunction with NN-based Full Waveform Inversion (FWI), where pretraining the NN provides a suitable initial guess for the optimization task. The training data set for this includes wave simulations on samples with arbitrary defects. In a recent study by Muller et al. <cit.>, their NN was pretrained directly on sensor data to predict an appropriate initial velocity distribution. A similar approach is described in <cit.>.


In this paper, although a similar idea is used, a significant improvement in the pretraining of the NN is proposed. Instead of simply using raw sensor data to predict an initial material distribution as in <cit.>, the gradient computed by the adjoint method with respect to the material coefficients of the piecewise constant representation (obtained during the first iteration with a homogeneous material) is utilized. This simplifies the architecture of the NN because the gradient of the conventional FWI from the first iteration has the same shape as that of the domain. The gradient contains an approximation of the shape of the damage (<Ref>), which renders the learning task easier. Furthermore, using pretrained NN removes the dependency of the convergence of the NN on the initial weight initialization. The paper at hand thereby proposes a method for accelerating the NN-based FWI <cit.> by pretraining the NN. Using the pretrained NN for FWI is called transfer learning with NN-based FWI. The proposed method is compared with three other approaches: conventional FWI, NN-based FWI without pretraining, and conventional FWI with an initial guess from the pretrained NN.


The remainder of the paper is organized as follows: In <ref>, the background of conventional and NN-based FWI is discussed, including the proposed transfer learning approach. In <ref>, we demonstrate the superior performance of the transfer learning NN-based FWI with the help of several numerical examples. Finally, <ref> concludes the paper's main findings. 



§ METHODOLOGY
 




 §.§ Physical model


In the sequel, the scalar wave equation is used. It describes the propagation of the scalar wave field u(x,t) in space x and time t in a medium whose spatial distribution of the density is given by ρ(x). The spatial distribution of the wave speed is denoted as c(x), the external force as f(x,t). With this notation, the 2D scalar wave equation reads

    ρ(x) ü(x, t)-∇·(ρ(x) c(x)^2 ∇ u(x, t)) = f(x,t),   x=[x,y] ∈Ω, t ∈𝒯,

where Ω is the 2D spatial domain and 𝒯 = [0, T_max] the time interval of interest. Appendix A of <cit.> gives a detailed derivation of this equation. Bürchner et al. <cit.> show that using mono-parameter FWI to only invert for density leads to better identification of defects with strong contrasts such as voids. Therefore, a dimensionless scaling function γ(x) is introduced to parameterize the scalar wave equation and the following optimization problem. Assuming ρ_0 and c_0 correspond to undamaged background material, the unknown material field γ(x) scales the density, i.e., ρ(x) = γ(x)ρ_0, while the wave speed remains constant, i.e.,  c(x)=c_0. This leads to the density-scaled version of the scalar wave equation, which is used throughout the paper,


    γ(x)ρ_0ü(x,t)-∇·(γ(x) ρ_0 c_0^2 ∇ u(x, t)) = f(x,t),   x=[x,y] ∈Ω, t ∈𝒯.


The initial and boundary conditions are


    u(x, 0) =u̇(x, 0)=0     on Ω,



    ∂ u(x, t)/∂ x =0      on Γ_y×𝒯,


    ∂ u(x, t)/∂ y =0      on Γ_x×𝒯.





 §.§ Numerical configuration



In the paper, all investigations are performed with the same numerical configuration. The setup is shown in <Ref>.



Considering a 2D domain of 0.1×0.05, a finite differences (FD) solver is used to solve the corresponding wave equation (as discussed in detail in Appendix A of <cit.>). The domain is excited with a two-cycle sine burst at the positions x^s. Introducing the central frequency k_c with the corresponding angular frequency ω = 2 π k_c and amplitude A_0 (<Ref>), the external force reads

    f(x,t) =
        
        A_0sin(ω t)sin^2(ω t/2 n_c) δ(x - x^s)    for  0 ≤ t ≤2 π n_c/ω
    
            0    for 2 π n_c/ω < t

where n_c = 2 is the number of cycles and x^s is the source position of the underlying experiment.


The undamaged density and wave speed of the material is  2700kg/m^3 and  6000 m/s. The central excitation frequency is 500. The synthetic reference data is generated with simulations on a grid with 512 × 256 grid points and 2000 time steps with dt=3e-8. For the FWI, a different grid configuration has been used to avoid the “inverse crime” <cit.>. The domain is discretized with 256 × 128 grid points, and the time integration is performed with 1000 time steps and a time step length of dt=6e-8. The sources and sensors are positioned along the top edge of the domain. There are four sources in total and the simulation with the excitation from one sensor is equivalent to carrying one experiment. The sensors are distributed symmetrically about the center of the plate, 36 grid points apart. 24 sensors spaced 6 grid points apart are placed between the rightmost and leftmost sources. The excitation of one source is equivalent to conducting one experiment. 



 §.§ Conventional FWI


Using conventional FWI, the problem at hand is parameterized by a set of coefficients γ̂ representing the scaling function γ(x) at all finite difference grid points. Introducing a misfit functional ℒ(γ̂) the corresponding optimization problem is described by

    γ̂^* = γ̂argmin ℒ (γ̂),

where γ̂^* is the optimal solution. The misfit functional sums up the squared error at all N^r receiver positions x^r for all experiments N^s

    ℒ (γ̂; u^o) = 1/2∫_T∫_Ω∑_s=1^N_s∑_r=1^N_r[ (û_s(γ̂;x,t)-u_s^o(x^r,t))^2 δ(x-x^r)  ]dΩ dt,

where û_s(γ̂;x,t) denotes a single simulated wave field given a set of coefficients γ̂ and u_s^o(x^r,t) a observed wave field. Measurements across all experiments are summarized by u^o={u_s^o}_s=1^N_s. 




For the computation of the derivative of the cost function with respect to the material coefficients, the adjoint method is used. For a detailed derivation, see e.g. <cit.>. Given a wave field u, the derivative with respect to one coefficient γ̂_i is

    d ℒ/d γ̂_i = ∑_s=1^N_s∫_𝒯[-ρ_0 u̇_s,i^†u̇_s,i + ρ_0 c_0^2 ∇ u_s,i^†·∇ u_s,i] dt,

where the subindex i denotes the solution of the corresponding primal and adjoint solution at the grid point i. The derivatives at all grid points form the gradient is ∇_γℒ(γ̂; u^o) and can readily be used in gradient-based optimization.



 §.§ NN-based FWI




  §.§.§ Hybrid gradient computation


In the NN-based FWI, an NN parameterized by θ is used to predict the scaling function coefficients, i.e., γ̂(θ). To this end, consider an NN N_γ that predicts the material coefficients at all grid points: 

    γ̂(θ) = N_γ(θ;k),

where θ contains all parameters of the NN and k is the input as shown in <Ref>. Therefore, the optimization problem becomes

    θ^* = θargmin ℒ (γ̂(θ); u^o).

The gradient of the cost function with respect to the NN parameters ∇_θ can be computed with the chain rule

    ∇_θℒ(γ̂; u^o) = ∇_γ̂ℒ(γ̂; u^o)_ adjoint method·∇_θγ̂_automatic
    differentiation.

As described in detail in <cit.>, the first term ∇_γℒ(γ̂(θ); u^o) is obtained through the adjoint method, while the latter ∇_θγ̂(θ) with automatic differentiation. The NN used to predict the density scaling function γ (<Ref>) has a generator type architecture <cit.>. The details of the generator architecture can be found in the Appendix <Ref>. The input k to the generator network is a tensor of size 128×8×4 filled with random values sampled from a normal distribution with zero mean and variance one. The output is the spatial distribution of the density scaling function, a tensor of size 1 × 256 × 128 equivalent to the size of the FD grid. The Adam optimizer <cit.> available in the PyTorch library <cit.> is used to solve the optimization problem. The NN-based FWI method is summarized in <Ref>.






  §.§.§ Transfer learning NN-based FWI



Transfer learning NN-based FWI uses a pretrained U-Net (<Ref>) to parameterize the density scaling function γ. The knowledge from a training data set is transferred to provide good weight initialization of the NN for the downstream task of FWI.



Pretraining: In <cit.>, the pretraining was carried out using observed wavefields as inputs and the true density distribution as the output. This already improved NN-based FWI. Yet, further improvements are possible by using the adjoint gradient ∇_γℒ(γ̂_0; u^o) from the first iteration as the input where the material field γ̂_0 is equal to 1 throughout the entire domain. This simplifies the design of the NN architecture since the shape of the gradient has the identical dimension as the discretized domain of computation itself, thus allowing the choice of the same shape for input and output (<Ref>). The gradient also already contains an approximation of the damage. This causes the initial guess to be closer to the sought density scaling function γ^true. Therefore, the training data set is D = {∇_γℒ(γ̂_0; u^o_i),γ^true_i}_i=1^N where N is the number of training samples. The supervised pretraining is carried out using the mean squared error (MSE) over all N samples as loss function given as
 
    ℒ_D = 1/N∑_i=1^N‖γ^true_i-γ̂_i‖_2^2,

 where γ̂ is the output of the NN N_γ. For this purpose, the NN, as shown in <Ref>, is used. Note that a different architecture is used in transfer learning NN-based FWI as compared to NN-based FWI, since the input tensor for pretraining is the same size as the output tensor, which was not the case earlier. Supervised training aims to minimize the average loss over the entire training data set. Pretraining the U-Net in such a manner requires one full gradient computation of the first conventional FWI iteration. Therefore, preparing the dataset for pretraining is more expensive than the one carried out in <cit.> since one adjoint simulation has to be carried out in addition to every forward simulation. However, it leads to better results as demonstrated in <Ref> because it provides a better initialization of the NN. 


A U-Net-based convolutional NN architecture <cit.> is employed for pretraining. More details about this architecture are in the Appendix (<ref>). The U-Net employs skip connections (<Ref>) to address the vanishing gradient problem. As previously mentioned, the input tensor is the adjoint gradient from the first conventional FWI iteration ℒ(γ̂_0; u^o), and the output is the predicted density scaling function γ̂. The U-Net is trained on a data set containing one ellipsoid-shaped damage parameterized using the semi-major and minor axes a and b, the center x_c, y_c, and the rotation θ, see <Ref>. The mean squared loss function is used along with RMSprop <cit.> optimization to train the network. The training of the U-Net is carried out with 800 samples for 100 epochs using a batch size of 80 and requires around 240 seconds on a Nvidia Quadro RTX 8000. More details about the optimal number of epochs and samples used for pretraining can be found in Appendix <ref> and <ref>, respectively.




 
 
Once the U-Net is trained, it is used for the downstream task of FWI. The optimization then updates the weights of N_γ based on the FWI loss function <Ref>. The updated <Ref> is:

 
    γ̂(θ^p) = N_γ(θ^p;∇_γℒ(γ̂_0; u^o))
        .


In the following, the pretrained U-Net is used to perform NN-based FWI (<Ref>). The input to the pretrained network is the adjoint gradient from the first conventional FWI iteration. However, the loss function is changed to <Ref>, and the gradient used to update θ^p is used for the NN-based FWI outlined in <Ref>.



  §.§.§ Conventional FWI with initial guess



For comparison, the initial guess provided by the pretrained network is used with conventional FWI. Therefore, the prediction γ̂ from the pretrained network N_γ(θ^p;∇_γℒ(γ̂_0; u^o)) is used as starting model.



  §.§.§ Importance of NN initialization

Weight initialization for NNs is important to achieve fast convergence and good accuracy. Using the weight initialization suggested by Glorot in <cit.> is common practice. This avoids the vanishing gradient problem but does not form an optimal starting point for FWI, as the random weights lead to the noisy initialization patterns depicted in the upper left part of <Ref>. On the contrary, an expected outcome is that the material remains undamaged almost everywhere except for locally confined flaws. An approximation to undamaged material can be provided by additionally initializing the weights of the last layer using a standard normal distribution with a narrow standard deviation of 0.01 and setting a rather high bias of, e.g. three. This removes most of the noise from the initialization of the density scaling function and provides a better starting point for NN-based FWI. As a last step, a sigmoid function maps the values to the range between zero and one. Furthermore, setting the weight of the last layer causes the method to be less dependent on the initialization of the weights. <Ref> shows the effect of the weight initialization on the reconstruction of a damage case within 35 iterations. This weight initialization is used for the generator network for NN-based FWI (<Ref>) as well as for the U-Net employed for transfer learning NN-based FWI (<Ref>). In the case of transfer learning NN-based FWI, the weights of the last layer are initialized in this way during pretraining, which improves the convergence of the training process.





§ RESULTS AND DISCUSSION
 




 §.§ Damage cases


To compare all four methods, i.e., conventional FWI (<Ref>), NN-based FWI (<Ref>), conventional FWI with an initial guess from pretraining (<Ref>), and transfer learning NN-based FWI (<Ref>), four damage cases are used, see <Ref>. The aim is to quantify the methods with respect to their reconstruction quality and convergence speed. The red color represents the region where the density scaling function has a value of 1, and the blue color represents a void/damage, i.e., a value of 10^-5 the scaling function. 



The four case studies are presented to assess the performance of transfer learning NN-based FWI. Then, the four methods are compared quantitatively using an average of over 100 test cases. Apart from the four case studies, we also compare the methods quantitatively with an average performance over 100 test cases. These cases were not part of the training data set used to pre-train the U-Net. All methods were executed on an Nvidia Quadro RTX 8000, and each case required ∼100seconds to complete.





 §.§ Case 1


<Ref> compares the four methods for the first damage case. The test case is similar to but not a sample from the data used for the pretraining of the network described in <Ref>. Therefore, it is expected that the pretrained network will perform better than other methods for this case. The first column in <Ref> shows that conventional FWI recovers the damage within 10-15 iterations. However, the reconstruction suffers from artifacts, which only reduce slightly with increasing iterations. The second column shows the individual snapshots of NN-based FWI without pretraining, which recovers the damage within 10 iterations with good accuracy and without any artifacts in the reconstruction.
The third column shows the results of the corresponding iterations for conventional FWI with an initial guess from the pretrained network. As expected, providing a good initial guess improves the adjoint method's reconstruction and helps recover the damage faster than without pretraining. Finally, the fourth column shows the results from transfer learning NN-based FWI. The reconstruction of the damage is much faster; within 5 iterations, the reconstructed shape of the damage is very accurate, leading to the best results for this case.  



 §.§ Case 2


<Ref> shows the comparison of the four methods for a damage case containing a rectangular-shaped hole and a circular hole. Note that the training set does not include a square hole; only elliptical holes were included (<Ref>). The first column shows conventional FWI, which recovers the shape approximately within 20-25 iterations. However, there are strong artifacts in the area just above the rectangular damage. The NN-based FWI recovers the rectangular-shaped hole and the circular hole within 20 iterations with good accuracy. The conventional FWI with an initial guess from the pretrained network performs much better than without pretraining and achieves a good reconstruction within 20 iterations with less artifacts in the reconstructed image. Nevertheless, the circular hole is not reconstructed very well. Finally, the transfer learning NN-based FWI performs the best. It achieves an almost perfect reconstruction within 10 iterations without any artifacts. As compared to the NN-based FWI without pretraining, the reconstructed shape is close to the true shape, thus demonstrating that the pretrained network performs very well for cases slightly outside the distribution of data on which it is trained.





 §.§ Case 3


For the third case study (<Ref>), three holes of different sizes need to be reconstructed. This renders the task to be quite different from the training data. The conventional FWI can roughly recover the rightmost damage while failing to recover the damages on the left side within 35 iterations. There are a lot of artifacts in the reconstruction and it is difficult to infer the correct shape and location of the damages from the reconstruction. The NN-based FWI can recover one damage shape within 15 iterations but struggles to reconstruct the other two. However, in this case, the reconstruction is better than conventional FWI, although the reconstructed shape of the damages is not accurate. The conventional FWI with an initial guess from the pretrained NN, delivers a similar reconstruction of the damage without pretraining. By contrast, the transfer learning NN-based FWI can quickly recover all the locations of the holes with good accuracy within 20 iterations. It improves the reconstructed shape of the damages in the subsequent iterations. The damaged shape is recovered more accurately than other methods, and the artifacts are negligible. Therefore, the transfer learning NN-based FWI performs best among all the methods in this case.  




 §.§ Case 4: Is transfer learning NN-based FWI always the best?


In very few cases, the transfer learning NN-based FWI does not perform as well as the other methods. The initial guess provided by the network is too inaccurate in some cases where the sought damage is very different from those of the training data. One such example is shown in <Ref>. The figure depicts two intersecting rectangular damages with three circular holes with a decreasing radius located on a straight line. Furthermore, there are two additional circular holes on the left and the right sides, resulting in six different damages to be recovered. This is a difficult case for FWI in general, specifically for the pre-trained network, because the training data included holes with only a single damage. Additionally, the smaller circular holes are located in the bigger circular hole's shadow region, making it more challenging to reconstruct the second and third holes.
 

The conventional FWI can partially recover the rectangular-shaped damage and one hole on the side. The circular holes behind the rectangular damage are not recovered. The NN-based FWI delivers a fast reconstruction of the whole damaged region using only 25-30 iterations to provide high accuracy. In the case of the conventional FWI with an initial guess from the pretrained network, the initial guess is inaccurate. The recovered damages are partially accurate for the rectangular-shaped damages. The circular holes on each side are also partially recovered. For the transfer learning  NN-based FWI, the method could only recover the shape of the rectangular-shaped damage and missed all the other damages. Upon increasing the total number of iterations, it is observed that it can recover the circular holes on both sides after 55 iterations. Still, it cannot recover the two intersecting rectangular damages with three circular holes behind each other with a decreasing radius ever after 100 iterations. This is one of the few cases where the damage reconstruction is unsatisfactory. However, the NN-based FWI without pretraining performs the best in this case, clearly indicating the advantage of using NN to discretize the material field. 









 §.§ Comparison over 100 test cases



FWI is carried out for 100 test cases to compare the four methods quantitatively, each undergoing 35 iterations. The cost function and the mean squared error between the reconstructed and true density scaling function over the iterations are used as a measure of comparison. Results are depicted in <Ref>.





The cost function for the transfer learning NN-based FWI reduces the fastest till 20 iterations. After 35 iterations, the cost is comparable for NN-based FWI and the transfer learning NN-based FWI. However, the mean squared error is the lowest for transfer learning NN-based FWI, followed by NN-based FWI. The mean squared error at the start of the iteration is the same for conventional FWI with an initial guess from the pretrained NN and transfer learning NN-based FWI. However, the MSE for conventional FWI without and with initial guess from pretrained NN increases slightly over the iterations due to the artifacts in the reconstructed density scaling function. Conversely, the MSE consistently decreases for NN-based FWI and transfer learning NN-based FWI, indicating good reconstruction of the density scaling function with fewer artifacts.
 

In the examples shown here (as well as in <cit.>), the reconstruction of a void using NN-based FWI contains considerably fewer artifacts. This is because NNs contain an inherent spectral bias <cit.> and prefer to learn the low-frequency content of data. On the contrary, the frequencies of the artifacts in the conventional FWI reconstructions are rather high. Therefore, the reconstruction of the NN-based FWI contains a particular kind of high-frequency filter that surprisingly does not inhibit the formation of sharp changes from material to void and vice versa, and, yet, does not exhibit artifacts. 



 §.§ Can transfer learning FWI recover from bad initial guess?


When using a transfer learning framework, a potential issue arises if the pretrained network provides an incorrect initial guess. This can occur when the data given to the pretrained network is different from the training data. To this end, case 2 from <Ref> is revisited. To emulate a bad pretraining, an incorrect initial guess is set manually for Case 2 (<Ref>) with a subsequent inversion. In Figure <ref>, it is evident that the initial guess in the first iteration differs strongly from the target damage shape. Although the transfer learning FWI required roughly 70 iterations (compared to the 10 iterations with Transfer learning FWI and 25 iterations with NN-based FWI experienced in <Ref>), it accurately recovers the damage. Similar results were observed for many cases, including Case 1, whereas for some more complicated cases, such as Case 3 and Case 4, the recovered shape was not as good as before. In general, a greater number of iterations is needed to satisfactorily recover the damages when the initial guess is inaccurate. 






 §.§ Drawbacks and limitations

Transfer learning NN-based FWI also suffers from some drawbacks that are common to any machine learning-based method. It requires extra effort to tune all the hyperparameters involved in the process, such as the learning rate, the learning rate scheduler, gradient clipping, cost scaling, and the number of samples for pretraining. This can be a time-intensive process. 


As in any other machine learning-based approach, carefully selecting training data is crucial for pretraining the network. Data that is too similar leads to overfitting, with the consequence that the network will not generalize well for out-of-sample data. However, selecting training data can be challenging, particularly if the type of expected damage is unknown. As the domain gets bigger, the task of generating sensible training data for pretraining, as well as hyperparameter tuning, will become more challenging. This indicates potential for further improvement in the architecture of NN models, which are less sensitive to the choice of hyperparameters.




§ CONCLUSION


NN-based FWI <cit.> uses an NN to parameterize the density scaling function. It relies on wave signals emitted at sources and received by the sensors to reconstruct flaws. This paper demonstrates that it is possible to improve the NN-based FWI using transfer learning, as introduced in <cit.>. The transfer learning is achieved with a U-Net, which is pretrained in a supervised manner — mapping the gradient from the first iteration of conventional FWI to the true density scaling function. The pretraining alleviates the strong dependency of the NN on the initialization of the NN weights. The combination of the NN ansatz in FWI and transfer learning through supervised pretraining greatly improves the convergence and reconstruction quality over conventional FWI. This includes conventional FWI relying on the same initial guess obtained through the pretrained NN. 


In addition to evaluating the improvements compared to the original transfer-learning-based scheme presented <cit.>, the transfer learning NN-based FWI is thoroughly compared to conventional FWI both with and without using the same pretrained network to provide a good starting point. In all cases, the presented transfer learning NN-based FWI outperforms conventional FWI. It might be argued that appropriate regularization techniques <cit.> lead to similar results as the NN-based FWI for conventional FWI. These can, however, equally be applied to both the NN-based FWI and transfer learning FWI. Whether additional regularization boosts the performance of transfer learning FWI and if similar results are obtainable with conventional FWI is the subject of future study.



§ ACKNOWLEDGEMENTS

The authors gratefully acknowledge the funding provided by the Georg Nemetschek Institut (GNI) through the joint research project DeepMonitor, which finances Divya Shyam Singh and Qing Sun, as well as the Geothermal-Alliance Bavaria (GAB) by the Bavarian State Ministry of Science and the Arts (StMWK), which finances Leon Herrmann. Furthermore, we gratefully acknowledge funds received by the Deutsche Forschungsgemeinschaft under Grant no. KO 4570/1-1 and RA 624/29-1 which support Tim Bürchner. Felix Dietrich would like to acknowledge the funds received by the Deutsche Forschungsgemeinschaft - project no. 468830823.



§ DECLARATIONS

Conflict of interest No potential conflict of interest was reported by
the authors.



§ DATA AVAILABILITY

We provide a PyTorch implementation of all methods in <https://doi.org/10.5281/zenodo.13150916>.










§ NEURAL NETWORK ARCHITECTURES






 §.§ Generator network architecture


<Ref> shows the detailed architecture of the generator network used for the NN-based FWI method. The generator network upsamples the input tensor of size 128 × 8 × 4 through its layers to match the finite differences grid of the domain (1 × 256  × 128). 




 §.§ U-Net architecture





Since both the input and the output data are a 2D matrix, a U-Net-based convolutional NN architecture is employed; see <ref>. The U-Net <cit.> architecture consists of a blocked structure where each block consists of a Convolution layer, a Batch normalization layer, and an activation layer repeated twice, followed by a Max pooling layer. The input and output tensor is 256 × 128 in shape. The skip connections (<Ref>) are used to overcome the vanishing gradient problem.


A small parametric study was carried out to find the best parameters for pretraining the U-Net and the optimum number of samples and epochs required for pretraining. It's worth noting that the optimum performance for the complete method is the performance of the pretrained network for the downstream task, i.e., carrying out the NN-based FWI. After each pretraining task, the network was run for 10 test cases, and the average was compared to the performance of each pretrained network. It was found that the optimum performance was achieved by continuing the pretraining further after the validation error had stopped reducing for the pretraining task. A similar observation was reported in <cit.>.




 §.§ Hyperparameter tuning

The hyperparameters for each method are tuned separately using 35 iterations to maximize the individual performance (Table <ref>). A learning rate scheduler with a polynomial decay (β·epoch+1)^α was also used in this process with α=-0.5 and β=0.2.





  §.§.§ Optimal number of epochs


To find the optimum number of epochs for pretraining, the network was trained for 50, 100, 200, and 500 epochs. This was done with three different numbers of samples for pretraining, 50 samples, 200 samples, and 500 samples, to find out if the change in the number of samples is a factor in the optimum number of epochs for pretraining. The hyperparameters used for this parametric study are listed in <ref>.


 
 Hyperparameter     Value 

  Learning rate     8 × 10^-4 
 
 Gradient clipping     5 × 10^-5 

α     -0.5  

β     0.2 

Batch size     No. of samples / 10 

 
tableHyperparameters used for pretraining the U-Net







<Ref> shows the result of a parametric study for a U-Net pretrained using 800 samples and 50,100,200,300,500 epochs for pretraining averaged over 100 test cases. The graph also looks similar for 50 and 200 samples, with 100 epochs performing the best for carrying out transfer learning NN-based FWI. Increasing the number of epochs further reduces the accuracy. The initial guess of the network trained for 100 epochs was the best compared to other cases, and the error reduced much faster than in other cases. Therefore, 100 epochs have been used to pretrain the U-Net.



  §.§.§ Optimal number of samples
 
A parametric study was conducted to determine the optimal number of data sets used for pretraining the U-Net. This study aims to determine the effect of increasing the number of samples on the performance of the U-Net for NN-based FWI. The number of samples used for pretraining the U-Net is 200, 300, 400, 500 and 800. The hyperparameters used for pretraining the network were kept constant for each case. These hyperparameters are shown in <Ref>. <Ref> depicts the mean squared error at each iteration averaged over 100 test cases. Each curve represents the corresponding amount of samples used for pretraining. <Ref> demonstrates the clear pattern that more samples improve the performance of transfer learning NN-based FWI. As can be seen, the MSE error is the least for pretraining using 800 samples for the downstream task. Therefore, the network pretrained from 800 samples is used for final comparison with other methods. 














10
url<#>1urlprefixURL href#1#2#2 #1#1

park2000impedance
G. Park, H. H. Cudney, D. J. Inman,
  http://dx.doi.org/10.1061/(asce)1076-0342(2000)6:4(153)Impedance-based
  health monitoring of civil structural components, Journal of Infrastructure
  Systems 6 (4) (2000) 153–160.
https://doi.org/10.1061/(asce)1076-0342(2000)6:4(153)
  doi:10.1061/(asce)1076-0342(2000)6:4(153).
<http://dx.doi.org/10.1061/(asce)1076-0342(2000)6:4(153)>

lynch2007overview
J. P. Lynch, http://dx.doi.org/10.1098/rsta.2006.1932An overview of
  wireless structural health monitoring for civil structures, Philosophical
  Transactions of the Royal Society A: Mathematical, Physical and Engineering
  Sciences 365 (1851) (2006) 345–372.
https://doi.org/10.1098/rsta.2006.1932
  doi:10.1098/rsta.2006.1932.
<http://dx.doi.org/10.1098/rsta.2006.1932>

kim2007health
S. Kim, S. Pakzad, D. Culler, J. Demmel, G. Fenves, S. Glaser, M. Turon,
  http://dx.doi.org/10.1145/1236360.1236395Health monitoring of civil
  infrastructures using wireless sensor networks, in: Proceedings of the 6th
  international conference on Information processing in sensor networks - IPSN
  ’07, IPSN ’07, ACM Press, 2007.
https://doi.org/10.1145/1236360.1236395
  doi:10.1145/1236360.1236395.
<http://dx.doi.org/10.1145/1236360.1236395>

tarantola1984inversion
A. Tarantola, http://dx.doi.org/10.1190/1.1441754Inversion of seismic
  reflection data in the acoustic approximation, GEOPHYSICS 49 (8) (1984)
  1259–1266.
https://doi.org/10.1190/1.1441754
  doi:10.1190/1.1441754.
<http://dx.doi.org/10.1190/1.1441754>

fichtner2010full
A. Fichtner, http://dx.doi.org/10.1007/978-3-642-15807-0Full Seismic
  Waveform Modelling and Inversion, Springer Berlin Heidelberg, 2011.
https://doi.org/10.1007/978-3-642-15807-0
  doi:10.1007/978-3-642-15807-0.
<http://dx.doi.org/10.1007/978-3-642-15807-0>

menke2018geophysical
https://www.sciencedirect.com/science/article/pii/B9780128135556099949Copyright,
  in: W. Menke (Ed.), Geophysical Data Analysis (Fourth Edition), fourth
  edition Edition, Academic Press, 2018.
https://doi.org/https://doi.org/10.1016/B978-0-12-813555-6.09994-9
  doi:https://doi.org/10.1016/B978-0-12-813555-6.09994-9.
<https://www.sciencedirect.com/science/article/pii/B9780128135556099949>

seidl2017full
R. Seidl, E. Rank, http://dx.doi.org/10.1063/1.4974657Full waveform
  inversion for ultrasonic flaw identification, in: AIP Conference
  Proceedings, Author(s), 2017.
https://doi.org/10.1063/1.4974657
  doi:10.1063/1.4974657.
<http://dx.doi.org/10.1063/1.4974657>

He2021NumericalUF
J. He, J. Rao, J. D. Fleming, H. Nath Gharti, L. T. Nguyen, G. Morrison,
  http://dx.doi.org/10.1088/1361-665X/ac0f44Numerical ultrasonic full
  waveform inversion (fwi) for complex structures in coupled 2d solid/fluid
  media, Smart Materials and Structures 30 (8) (2021) 085044.
https://doi.org/10.1088/1361-665x/ac0f44
  doi:10.1088/1361-665x/ac0f44.
<http://dx.doi.org/10.1088/1361-665X/ac0f44>

burchner2023immersed
T. Bürchner, P. Kopp, S. Kollmannsberger, E. Rank,
  https://www.sciencedirect.com/science/article/pii/S0045782523000166Immersed
  boundary parametrizations for full waveform inversion, Computer Methods in
  Applied Mechanics and Engineering 406 (2023) 115893.
https://doi.org/https://doi.org/10.1016/j.cma.2023.115893
  doi:https://doi.org/10.1016/j.cma.2023.115893.
<https://www.sciencedirect.com/science/article/pii/S0045782523000166>

BKK23
T. Bürchner, P. Kopp, S. Kollmannsberger, E. Rank,
  https://www.sciencedirect.com/science/article/pii/S0045782523004103Isogeometric
  multi-resolution full waveform inversion based on the finite cell method,
  Computer Methods in Applied Mechanics and Engineering (2023) 116286https://doi.org/https://doi.org/10.1016/j.cma.2023.116286
  doi:https://doi.org/10.1016/j.cma.2023.116286.
<https://www.sciencedirect.com/science/article/pii/S0045782523004103>

virieux2009overview
J. Virieux, S. Operto, http://dx.doi.org/10.1190/1.3238367An overview
  of full-waveform inversion in exploration geophysics, GEOPHYSICS 74 (6)
  (2009) WCC1–WCC26.
https://doi.org/10.1190/1.3238367
  doi:10.1190/1.3238367.
<http://dx.doi.org/10.1190/1.3238367>

virieux2017introduction
J. Virieux, A. Asnaashari, R. Brossier, L. Métivier, A. Ribodetti, W. Zhou,
  http://dx.doi.org/10.1190/1.9781560803027.entry66. An introduction to
  full waveform inversion, Society of Exploration Geophysicists, 2014, pp.
  R1–1–R1–40.
https://doi.org/10.1190/1.9781560803027.entry6
  doi:10.1190/1.9781560803027.entry6.
<http://dx.doi.org/10.1190/1.9781560803027.entry6>

yang2019deep
F. Yang, J. Ma, http://dx.doi.org/10.1190/geo2018-0249.1Deep-learning
  inversion: A next-generation seismic velocity model building method,
  GEOPHYSICS 84 (4) (2019) R583–R599.
https://doi.org/10.1190/geo2018-0249.1
  doi:10.1190/geo2018-0249.1.
<http://dx.doi.org/10.1190/geo2018-0249.1>

wu2019inversionnet
Y. Wu, Y. Lin, Inversionnet: An efficient and accurate data-driven full
  waveform inversion, IEEE Transactions on Computational Imaging 6 (2020)
  419–433.
https://doi.org/10.1109/TCI.2019.2956866
  doi:10.1109/TCI.2019.2956866.

Wang2020VelocityMB
W. Wang, J. Ma, http://dx.doi.org/10.1190/geo2018-0591.1Velocity model
  building in a crosswell acquisition geometry with image-trained artificial
  neural networks, GEOPHYSICS 85 (2) (2020) U31–U46.
https://doi.org/10.1190/geo2018-0591.1
  doi:10.1190/geo2018-0591.1.
<http://dx.doi.org/10.1190/geo2018-0591.1>

rao2023quantitative
J. Rao, F. Yang, H. Mo, S. Kollmannsberger, E. Rank,
  https://www.sciencedirect.com/science/article/pii/S0022460X22006010Quantitative
  reconstruction of defects in multi-layered bonded composites using fully
  convolutional network-based ultrasonic inversion, Journal of Sound and
  Vibration 542 (2023) 117418.
https://doi.org/https://doi.org/10.1016/j.jsv.2022.117418
  doi:https://doi.org/10.1016/j.jsv.2022.117418.
<https://www.sciencedirect.com/science/article/pii/S0022460X22006010>

jin2018learn
Y. Jin, W. Hu, X. Wu, J. Chen,
  http://dx.doi.org/10.1190/segam2018-2997901.1Learn low-wavenumber
  information in fwi via deep inception-based convolutional networks, in: SEG
  Technical Program Expanded Abstracts 2018, Society of Exploration
  Geophysicists, 2018.
https://doi.org/10.1190/segam2018-2997901.1
  doi:10.1190/segam2018-2997901.1.
<http://dx.doi.org/10.1190/segam2018-2997901.1>

sun2020extrapolated
H. Sun, L. Demanet, https://doi.org/10.1190/geo2019-0195.1Extrapolated
  full-waveform inversion with deep learning, Geophysics 85 (3) (2020)
  R275–R288.
http://arxiv.org/abs/https://pubs.geoscienceworld.org/geophysics/article-pdf/85/3/R275/5018257/geo-2019-0195.1.pdf
  arXiv:https://pubs.geoscienceworld.org/geophysics/article-pdf/85/3/R275/5018257/geo-2019-0195.1.pdf,
  https://doi.org/10.1190/geo2019-0195.1
  doi:10.1190/geo2019-0195.1.
<https://doi.org/10.1190/geo2019-0195.1>

yang2022deep
T. Yang, X. Chai, W. Nie, J. Yu,
  https://library.seg.org/doi/abs/10.1190/iwmg2021-35.1Deep-learning-based
  low-frequency reconstruction for full-waveform inversion, 2022, pp.
  134–137.
http://arxiv.org/abs/https://library.seg.org/doi/pdf/10.1190/iwmg2021-35.1
  arXiv:https://library.seg.org/doi/pdf/10.1190/iwmg2021-35.1, https://doi.org/10.1190/iwmg2021-35.1 doi:10.1190/iwmg2021-35.1.
<https://library.seg.org/doi/abs/10.1190/iwmg2021-35.1>

li2022high
Y. Li, A. Alali, T. Alkhalifah,
  https://www.earthdoc.org/content/papers/10.3997/2214-4609.202210778High-wavenumber
  extrapolation for band-limited fwi results using deep learning 2022 (1)
  (2022) 1–5.
https://doi.org/https://doi.org/10.3997/2214-4609.202210778
  doi:https://doi.org/10.3997/2214-4609.202210778.
<https://www.earthdoc.org/content/papers/10.3997/2214-4609.202210778>

kleman2023full
C. Kleman, S. Anwar, Z. Liu, J. Gong, X. Zhu, A. Yunker, R. Kettimuthu, J. He,
  http://dx.doi.org/10.1115/1.4062092Full waveform inversion-based
  ultrasound computed tomography acceleration using two-dimensional
  convolutional neural networks, Journal of Nondestructive Evaluation,
  Diagnostics and Prognostics of Engineering Systems 6 (4) (Mar. 2023).
https://doi.org/10.1115/1.4062092
  doi:10.1115/1.4062092.
<http://dx.doi.org/10.1115/1.4062092>

araya2019deep
M. Araya-Polo, S. Farris, M. Florez,
  http://dx.doi.org/10.1190/tle38110872a1.1Deep learning-driven
  velocity model building workflow, The Leading Edge 38 (11) (2019)
  872a1–872a9.
https://doi.org/10.1190/tle38110872a1.1
  doi:10.1190/tle38110872a1.1.
<http://dx.doi.org/10.1190/tle38110872a1.1>

he2021reparameterized
Q. He, Y. Wang, http://dx.doi.org/10.1190/geo2019-0382.1Reparameterized
  full-waveform inversion using deep neural networks, GEOPHYSICS 86 (1) (2021)
  V1–V13.
https://doi.org/10.1190/geo2019-0382.1
  doi:10.1190/geo2019-0382.1.
<http://dx.doi.org/10.1190/geo2019-0382.1>

zhu2022integrating
W. Zhu, K. Xu, E. Darve, B. Biondi, G. C. Beroza,
  http://dx.doi.org/10.1190/geo2020-0933.1Integrating deep neural
  networks with full-waveform inversion: Reparameterization, regularization,
  and uncertainty quantification, GEOPHYSICS 87 (1) (2021) R93–R109.
https://doi.org/10.1190/geo2020-0933.1
  doi:10.1190/geo2020-0933.1.
<http://dx.doi.org/10.1190/geo2020-0933.1>

herrmann2023use
L. Herrmann, T. Bürchner, F. Dietrich, S. Kollmannsberger,
  http://dx.doi.org/10.1016/j.cma.2023.116278On the use of neural
  networks for full waveform inversion, Computer Methods in Applied Mechanics
  and Engineering 415 (2023) 116278.
https://doi.org/10.1016/j.cma.2023.116278
  doi:10.1016/j.cma.2023.116278.
<http://dx.doi.org/10.1016/j.cma.2023.116278>

Muller2023
A. P. O. Muller, J. C. Costa, C. R. Bom, M. Klatt, E. L. Faria, M. P.
  de Albuquerque, M. P. de Albuquerque,
  http://dx.doi.org/10.1093/gji/ggad215Deep pre-trained fwi: where
  supervised learning meets the physics-informed neural networks, Geophysical
  Journal International 235 (1) (2023) 119–134.
https://doi.org/10.1093/gji/ggad215
  doi:10.1093/gji/ggad215.
<http://dx.doi.org/10.1093/gji/ggad215>

rahaman2019spectral
N. Rahaman, A. Baratin, D. Arpit, F. Draxler, M. Lin, F. Hamprecht, Y. Bengio,
  A. Courville, https://proceedings.mlr.press/v97/rahaman19a.htmlOn the
  spectral bias of neural networks, in: K. Chaudhuri, R. Salakhutdinov (Eds.),
  Proceedings of the 36th International Conference on Machine Learning, Vol. 97
  of Proceedings of Machine Learning Research, PMLR, 2019, pp. 5301–5310.
<https://proceedings.mlr.press/v97/rahaman19a.html>

nigam2000text
K. Nigam, A. K. Mccallum, S. Thrun, T. Mitchell, Machine Learning 39 (2/3)
  (2000) 103–134.
https://doi.org/10.1023/a:1007692713085
  doi:10.1023/a:1007692713085,
  http://dx.doi.org/10.1023/a:1007692713085[link].
<http://dx.doi.org/10.1023/a:1007692713085>

Qasim2022
R. Qasim, W. H. Bangyal, M. A. Alqarni, A. Ali Almazroi,
  http://dx.doi.org/10.1155/2022/3498123A fine-tuned bert-based
  transfer learning approach for text classification, Journal of Healthcare
  Engineering 2022 (2022) 1–17.
https://doi.org/10.1155/2022/3498123
  doi:10.1155/2022/3498123.
<http://dx.doi.org/10.1155/2022/3498123>

shaha2018transfer
M. Shaha, M. Pawar, Transfer learning for image classification, in: 2018 Second
  International Conference on Electronics, Communication and Aerospace
  Technology (ICECA), 2018, pp. 656–660.
https://doi.org/10.1109/ICECA.2018.8474802
  doi:10.1109/ICECA.2018.8474802.

Hussain2018
M. Hussain, J. J. Bird, D. R. Faria,
  http://dx.doi.org/10.1007/978-3-319-97982-3_16A Study on CNN Transfer
  Learning for Image Classification, Springer International Publishing, 2018,
  p. 191–202.
https://doi.org/10.1007/978-3-319-97982-3_16
  doi:10.1007/978-3-319-97982-3_16.
<http://dx.doi.org/10.1007/978-3-319-97982-3_16>

Barret2018
B. Zoph, D. Yuret, J. May, K. Knight,
  http://arxiv.org/abs/1604.02201Transfer learning for low-resource
  neural machine translation, CoRR abs/1604.02201 (2016).
http://arxiv.org/abs/1604.02201 arXiv:1604.02201.
<http://arxiv.org/abs/1604.02201>

Surafel2018
S. M. Lakew, A. Erofeeva, M. Negri, M. Federico, M. Turchi,
  http://arxiv.org/abs/1811.01137Transfer learning in multilingual
  neural machine translation with dynamic vocabulary, CoRR abs/1811.01137
  (2018).
http://arxiv.org/abs/1811.01137 arXiv:1811.01137.
<http://arxiv.org/abs/1811.01137>

Tang2021
H. Tang, L. Notash, http://dx.doi.org/10.1115/DETC2021-69006Neural
  network based transfer learning for robot path generation, in: Volume 8B:
  45th Mechanisms and Robotics Conference (MR), IDETC-CIE2021, American Society
  of Mechanical Engineers, 2021.
https://doi.org/10.1115/detc2021-69006
  doi:10.1115/detc2021-69006.
<http://dx.doi.org/10.1115/DETC2021-69006>

Pereida2018
K. Pereida, M. K. Helwa, A. P. Schoellig, Data-efficient multirobot, multitask
  transfer learning for trajectory tracking, IEEE Robotics and Automation
  Letters 3 (2) (2018) 1260–1267.
https://doi.org/10.1109/LRA.2018.2795653
  doi:10.1109/LRA.2018.2795653.

Makondo2018
N. Makondo, B. Rosman, O. Hasegawa, Accelerating model learning with
  inter-robot knowledge transfer, in: 2018 IEEE International Conference on
  Robotics and Automation (ICRA), 2018, pp. 2417–2424.
https://doi.org/10.1109/ICRA.2018.8461218
  doi:10.1109/ICRA.2018.8461218.

kollmannsberger2023transfer
S. Kollmannsberger, D. Singh, L. Herrmann, Transfer learning enhanced full
  waveform inversion*, in: 2023 IEEE/ASME International Conference on Advanced
  Intelligent Mechatronics (AIM), 2023, pp. 866–871.
https://doi.org/10.1109/AIM46323.2023.10196158
  doi:10.1109/AIM46323.2023.10196158.

wirgin2004inverse
A. Wirgin, The inverse crime (2004).
http://arxiv.org/abs/math-ph/0401050
  arXiv:math-ph/0401050.

badrinarayanan2017segnet
V. Badrinarayanan, A. Kendall, R. Cipolla, Segnet: A deep convolutional
  encoder-decoder architecture for image segmentation, IEEE Transactions on
  Pattern Analysis and Machine Intelligence 39 (12) (2017) 2481–2495.
https://doi.org/10.1109/TPAMI.2016.2644615
  doi:10.1109/TPAMI.2016.2644615.

Kingma2014AdamAM
D. P. Kingma, J. Ba, https://arxiv.org/abs/1412.6980Adam: A method for
  stochastic optimization (2014).
https://doi.org/10.48550/ARXIV.1412.6980
  doi:10.48550/ARXIV.1412.6980.
<https://arxiv.org/abs/1412.6980>

NEURIPS2019_9015
A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen,
  Z. Lin, N. Gimelshein, L. Antiga, A. Desmaison, A. Kopf, E. Yang, Z. DeVito,
  M. Raison, A. Tejani, S. Chilamkurthy, B. Steiner, L. Fang, J. Bai,
  S. Chintala, Pytorch: An imperative style, high-performance deep learning
  library, in: Advances in Neural Information Processing Systems 32, Curran
  Associates, Inc., 2019, pp. 8024–8035,
  
<http://papers.neurips.cc/paper/9015-pytorch-an>
<-imperative-style-high-performance-deep-learning-library.pdf>.

Ronneberger2015
O. Ronneberger, P. Fischer, T. Brox,
  http://dx.doi.org/10.1007/978-3-319-24574-4_28U-Net: Convolutional
  Networks for Biomedical Image Segmentation, Springer International
  Publishing, 2015, p. 234–241.
https://doi.org/10.1007/978-3-319-24574-4_28
  doi:10.1007/978-3-319-24574-4_28.
<http://dx.doi.org/10.1007/978-3-319-24574-4_28>

Graves2013GeneratingSW
A. Graves, https://arxiv.org/abs/1308.0850Generating sequences with
  recurrent neural networks (2013).
https://doi.org/10.48550/ARXIV.1308.0850
  doi:10.48550/ARXIV.1308.0850.
<https://arxiv.org/abs/1308.0850>

glorot2010understanding
X. Glorot, Y. Bengio,
  https://proceedings.mlr.press/v9/glorot10a.htmlUnderstanding the
  difficulty of training deep feedforward neural networks, in: Y. W. Teh,
  M. Titterington (Eds.), Proceedings of the Thirteenth International
  Conference on Artificial Intelligence and Statistics, Vol. 9 of Proceedings
  of Machine Learning Research, PMLR, Chia Laguna Resort, Sardinia, Italy,
  2010, pp. 249–256.
<https://proceedings.mlr.press/v9/glorot10a.html>

Arridge2019
S. Arridge, P. Maass, O. Öktem, C.-B. Schönlieb,
  http://dx.doi.org/10.1017/S0962492919000059Solving inverse problems
  using data-driven models, Acta Numerica 28 (2019) 1–174.
https://doi.org/10.1017/s0962492919000059
  doi:10.1017/s0962492919000059.
<http://dx.doi.org/10.1017/S0962492919000059>

neyshabur2020being
B. Neyshabur, H. Sedghi, C. Zhang,
  https://proceedings.neurips.cc/paper_files/paper/2020/file/0607f4c705595b911a4f3e7a127b44e0-Paper.pdfWhat
  is being transferred in transfer learning?, in: H. Larochelle, M. Ranzato,
  R. Hadsell, M. Balcan, H. Lin (Eds.), Advances in Neural Information
  Processing Systems, Vol. 33, Curran Associates, Inc., 2020, pp. 512–523.
<https://proceedings.neurips.cc/paper_files/paper/2020/file/0607f4c705595b911a4f3e7a127b44e0-Paper.pdf>






 Paper Title: FSSC: Federated Learning of Transformer Neural Networks for Semantic Image Communication.


	
	
	
	 
	
FSSC: Federated Learning of Transformer Neural Networks for Semantic Image Communication
    
		
			Yuna Yan1, 
			Xin Zhang1, 
			Lixin Li1, 
			Wensheng Lin1, 
            Rui Li2,
            Wenchi Cheng3,
			and Zhu Han4 
		1School of Electronics and Information, Northwestern Polytechnical University, Xi’an, China, 710129
        2Samsung AI Center, Cambridge, UK
        3State Key Laboratory of Integrated Services Networks, Xidian University, Xi'an, China, 710071
		4Department of Electrical and Computer Engineering, University of Houston, Houston, TX, 77004
	
    August 5, 2024
==================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================



§ ABSTRACT 
In this paper, we address the problem of image semantic communication in a multi-user deployment scenario and propose a federated learning (FL) strategy for a Swin Transformer-based semantic communication system (FSSC). Firstly, we demonstrate that the adoption of a Swin Transformer for joint source-channel coding (JSCC) effectively extracts semantic information in the communication system. Next, the FL framework is introduced to collaboratively learn a global model by aggregating local model parameters, rather than directly sharing clients' data. This approach enhances user privacy protection and reduces the workload on the server or mobile edge.
Simulation evaluations indicate that our method outperforms the typical JSCC algorithm and traditional separate-based communication algorithms. Particularly after integrating local semantics, the global aggregation model has further increased the Peak Signal-to-Noise Ratio (PSNR) by more than 2dB, thoroughly proving the effectiveness of our algorithm. 




Semantic communication, federated learning, swin Transformer, privacy protection.




§ INTRODUCTION

Wireless communication is embracing an explosive growth of data traffic semantic nature, generated by various contemporary multi-media applications such as AR/VR, video streaming, gaming and telehealth. Traditional communication systems take the source messages (e.g. images and texts) and first encode these messages with source coding algorithms on the application layer, in order to remove redundant information. Subsequently, when the compressed information reaches the physical layer, it is then further channel-encoded by injection of redundant bits, to be resilient to the noise in the communication channel. There are three major factors that make this current approach inefficient in the 5th Generation Mobile Communication Technology(5G) and future generations of mobile networks. Firstly, although modern channel coding algorithm on its own has theoretically reached the Shannon limit on the physical layer <cit.>, when jointly considering the source coding and channel coding, the Shannon's optimality theory only holds when the message blocks are asymptotically infinitely long, which is impractical in the real-world <cit.>. Secondly, the separation-based approach<cit.> compresses and then upscales messages, thus inherently introducing extra latency, which lead to inferior performances for major time-sensitive 5G/6G and upcoming applications, e.g. tactile internet and autonomous driving. Third, channel coding strategies treat all bits of information equally, missing the opportunity to utilize the abundant semantic information available in the messages, especially in image and video contents.

To address these issues in the traditional communication regime, semantic communication<cit.> a.k.a. joint source and channel coding, has been proposed and studied, which focuses on the transmission of meaning rather than bits, helping to escape the “Shannon trap”<cit.>. 
Compared to conventional communication methods, semantic communication is centered around the relevant semantic features essential for transmitting information, thereby eliminating unnecessary redundancy. This approach leads to a significant reduction in data transmission latency and bandwidth consumption. To enhance the versatility of semantic communication systems, we propose the employment of a Swin Transformer-based semantic communication (STSC) system that primarily focuses on the reconstruction of global image signals. Its unique architecture enables the image-oriented semantic communication system to effectively capture deeper semantic features from potential semantic information and represent complex semantic relationships within the input data. In addition, the Swin Transformer is capable of reducing computational complexity and speeding up both model training and inference compared to traditional Transformer models. As a result, it is well-suited for handling large-scale semantic communication, especially in multi-user deployments.



However, image-oriented semantic communication targets applications that often involves users' private data, such as their faces, home interior, license plates, etc., which cannot always be uploaded to the cloud server. Moreover, most of the prior art on image-based semantic communication are point-to-point systems, which do not fit in many applications where multiple users are involved. In this paper, we propose a multi-user federated learning (FL) framework for semantic communication, showcasing the potential of efficient learning of a global semantic communication model that combines the learned knowledge from multiple users, without direct access to users' private data. We demonstrate that the proposed FL semantic system, when powered with a hierachical transformer architecture, i.e. the Swin Transformer, achieves a better image transmission quality at a faster convergence rate compared to the single-user.

In summary, the main contributions of this paper are summarized as follows:


	
  * We propose a Swin Transformer based semantic communication architecture, which aims at minimizing the difference between the transmitted and received images, so as to effectively communicate image semantic information.
	
  * In contrast to the centralised approach of sharing the user data directly to a cloud server in order to train a transformer neural network for semantic communication, we consider a cross-silo FL setting that employs multiple users to train locally their own model using local data, and share the local updates with a global server, thereby preserving users' data-privacy.
	
  * Through the use of well-established dataset for computer vision, simulation results show that our proposed framework surpasses the representative JSCC methods, and further boosts peak signal-to-noise ratio (PSNR) by over 2dB on top of aggregating local semantics.


The rest of this paper is organized as follows. In Section <ref>, we introduce the framework of the image semantic communication system based on FL and formulate the corresponding research problem. We then present the federated learning strategy for a Swin Transformer-based semantic communication (FSSC) in Section <ref>. In Section <ref>, we present numerical simulation results and showcase the semantic communication performance of FSSC as well as the convergence of the system. Finally, the paper is concluded in Section <ref>.




§ SYSTEM SETUP


In this paper, we consider a multi-user FL semantic communication system, as shown in Fig. <ref>, which consists of a cloud server and multiple clients. An FSSC client can be an edge server or a terminal device. When the client is an edge server, it is able to perform computational tasks, hence the training and inference of STSC happens on the client locally. When the client is a light-weight terminal device equipped with only limited or no computational power, the client should be served by a corresponding edge server which is trusted by the client to share data with in order to perform training. The federated training process is detailed as follows:

	
  * The cloud server initialize the STSC model for semantic communication, and distribute it to all participating clients. Each client will then perform n epochs of local training using its own data. Once the local training is completed, all clients transmit the updated model parameters to the cloud server, and the server will then aggregate the updated model parameters following Federated Averaging <cit.>, which completes one round of the federated learning. This process is repeated until the model converges;
	
  * During deployment of a trained STSC model, a user will initiate a communication request and sends a semantic message to the client. The client will then process this message by performing an inference using the trained STSC model with the message sample as the joint source-channel encoding process, and then transmit the output of the STSC model through the uplink channel to the receiver;
	
  * The transmitted message will then go through a noisy channel which typically add noisy corruption to the message. Once the receiver receives the message, it will then use the decoding part of the STSC model to recover the message. Note that the receiver in this setting could be either a cloud server, or a client.




Obviously, semantic information extraction and reconstruction of images are considered to be the key to the success of this system. Therefore, in this paper, the average mean square error (MSE) between the original input image and the reconstructed image is used as the loss function, which is defined as,


    MSE=1/L∑_i=1^n(x_i-x̂_i)^2,

where x_i is the original image vector, x̂_i is the reconstructed image vector, and L  represents the length of the image vector. A smaller MSE indicates that the reconstructed image is closer to the original image and the quality of image reconstruction is better.




§ IMAGE SEMANTIC ENCODER AND DECODER NETWORK


In this section, we first present an edge architecture for Swin Transformer-based semantic communication (STSC), which aims to capture long-term hierarchical semantic information from images. Additionally, to enhance the accuracy of semantic information extraction and ensure user privacy, we introduce the federated learning strategy for a Swin Transformer-based semantic communication (FSSC), which enables the training of edge models using FL across multiple devices hosted on cloud servers. As a result, our proposed scheme can effectively learn semantic representations from images transmitted by different users, and ensure reliable communication on various channels.



 §.§ STSC Model

Fig. <ref> shows the semantic communication network structure based on Swin Transformer. In this paper, joint source channel coding is adopted for semantic communication, in which the Swin Transformer module  <cit.> is adopted for encoding and decoding. Meanwhile, a non-trainable fully connected layer is used to simulate the physical channel.

At the transmitter, a set of N images 𝐗={x_i}_i=1^N is given, where x_i ∈ℝ^3 × H × W denotes the i-th image, and H and W denote height and width of the image, respectively. Firstly, the images are divided into non-overlapping patches through patch partion module. The size of patch is 4 × 4  and the vector dimension is converted to (H / 4,  W / 4,48). 
Then, the feature block sequence is linearly embedded. And through the learnable embedding matrix E, the feature blocks can be projected into an  embedding representation of arbitrary dimension C. After passing through the Swin Transformer block, the vector dimension is transformed to  (H / 4,  W / 4,C), 
and the whole process is called Stage 1. After the embedded representation, the feature blocks are fed into the Swin Transformer module, as shown in Fig. <ref>. As can be seen from the figure, the first patch merging layer and Swin Transformer block are combined into Stage 2. The patch merging layer concatenates each group of adjacent patches of size 2 × 2, so that the number of patch tokens becomes 1 / 4 of the original, i.e., H / 8×  W / 8. Meanwhile, the dimension of patch token is expanded by four times, i.e., 4C. In order to reduce the output dimension and realize the downsampling of the feature map, the patch merging layer performs a fully connected operation (implemented by a 1 × 1  convolutional layer). After this process, the dimension of the concatenated feature patch is reduced from 4C to 2C. Then, the patch goes through the swin Transformer block for feature transformation. Finally, the output dimension becomes  (H / 8,  W / 8,2C). The output x̃ passes through a fully connected layer to facilitate transmission in the channel. Specifically, it can be expressed as,

    s=W_1 x̃+b_1,

where, W_1 is the weight parameter matrix of the full connection, and  b_1 is the bias of the full connection.

In order to realize the joint training of encoder and decoder, a layer of fully connected neural network is used to simulate the channel. The neural network is actually a pair of input-output mappings, where the mapping relationship is determined by the neuron weight W_n and the bias b_1. Here, the weight W_n represents the channel gain, and the bias b_n is a random variable added to simulate the noise in the channel. The variance of the variable corresponds to the power of the channel noise, and its value mainly depends on the SNR and transmission power. After transmission in the physical channel, the signal s receives interference from channel noise and becomes the signal y when it reaches the receiver. 

At the receiver, the decoder consists of the Swin Transformer blocks. The function is to decode the signal y to restore the signal x̂.




 §.§ FL Training

In this paper, we develop the FSSC algorithm, by applying STSC to a FL framework, which can extract semantic information accurately and carry out semantic communication while protecting user privacy.While Fig. <ref> depics the overal architecture of FSSC, we detail below the specific training process of FSSC algorithm.

We use Fedavg<cit.> as the aggregation algorithm for our federated training. Assume that N clients participate in FL training. D_k  denotes the local Dataset of client E_k, and |D_k| is the size of the corresponding dataset. Thus, the training loss of the client based on the model in the local sample is

    Loss^k(ω)=1/|D_k|∑_i=1^|D_k|Loss_i^k(ω),

where  Loss^k(·) is the training Loss of the local client E_k, and Loss_i^k(·) is the loss of the sample x_i corresponding to the client. In the FSSC algorithm, Loss(·) is represented in (<ref>).

Therefore, the global federated training loss can be obtained by weighted average of the loss of each client according to the dataset size. The specific formula is as follows.


    Loss^∘(ω)=∑_k=1^N |D_k|/|D_o|Loss^k(ω),

where Loss^∘(ω) is the global loss after cloud aggregation, and |D_o| is the global dataset size. The purpose of the FL training is to find the global optimal model that minimizes the sum of training losses over all data. The specific federated training process is as follows.

Initialization. The dataset is non-uniformly split and distributed on each client. The semantic communication network based on Swin Transformer is initialized on the cloud server, and the initialized global model parameters W^o are broadcast to each client to build a local STSC network.

Model training. In the communication round t, the client participating in the training will conduct local training on STSC by using the local data set. And the stochastic gradient descent method is used to update the local model parameters. For each batch i, the process can be expressed as,


    ω_i+1^k ←ω_i^k-η∂Loss(ω_i^k ; b)/∂ω_i^k,

where  η is the learning rate. Each client saves its own trained local parameters and uploads them to the cloud server.

Parameter aggregation. The cloud server aggregates model parameters from different clients, and updates W_t^o to W_t+1^o . The aggregation mechanism shown in the following equation. Through this process, the update of the global shared model can be completed.


    w_t+1←∑_k=1^K |D_k|/|D_o| w_t+1^k,

where K is the total number of clients participating in training. This aggregation mechanism provides a certain guarantee for the stability of the system from the perspective of model convergence. Specifically, a certain degree of client-side fluctuations can be accepted without affecting the final convergence of the model.

Model convergence. The parameters of the global model W_t+1^o  are downloaded to each client. Then, the next round of training begins. Repeat this step until a specified global training cycle is reached or the model converges.



§ EXPERIMENT AND NUMERICAL RESULTS





 §.§ Simulation Settings

In this paper, we use the CIFAR-10 dataset<cit.> to evaluate the performance of the proposed FSSC algorithm. This dataset contains 50,000 training images and 10,000 test images, each with a size of 32 × 32. In addition, the number of clients is set to 3, and each client is assigned a different amount of non-overlapping training data. All clients use the same set of validation and test images. The training and testing environment are Windows 11+ CUDA 12.1, and the deep learning framework is Pytorch2.2.1.



In the experiments, the process of downloading and uploading model parameters between the client and the cloud server is regarded as an ideal situation. The focus is on the effects of global models obtained through local training and federated aggregation. The following introduces the initialization parameter settings for the semantic communication network based on Swin Transformer. Firstly, the network hyper-parameter C is set to 32. The network training all uses MSE as the loss function, the number of communication round is set to 60, the batch size of the data is 64, the learning rate is 0.001, and the Adam optimizer<cit.> is used. In all the tests, the compression ratio (CR) of the images is set to 0.33. 

To demonstrate the superiority of our method, we choose the typical JSCC <cit.> and the traditional communication algorithm as the comparison algorithm. The traditional communication algorithm adopts JPEG <cit.> as the source coding method, the channel coding adopts LDPC <cit.>, and the modulation mode and the demodulation mode are Quadrature Amplitude Modulation(QAM). And The modulation order is set to 4. 



 §.§ Simulation Results

In order to evaluate the semantic representation of the FSSC model, the experiments in this section compared the performance results of the local model algorithm STSC, JSCC and the traditional communication algorithms on AWGN channels, and verified the changes of the PSNR of the three methods under different SNR conditions, as shown in Fig. <ref>. It can be seen from Fig. <ref> that the PSNR of STSC and JSCC gradually increase with the increase of SNR. The STSC algorithm outperformed JSCC in terms of PSNR, with an average improvement of over 5 decibels. This significant advantage demonstrates the superior efficiency of the STSC algorithm in semantic feature extraction and reconstructing image data, resulting in a more accurate representation of the original image. At the same time, it can be noted that the PSNR of JPEG+LDPC+QAM remains unchanged when the channel conditions are very bad or greatly improved. This is because when the SNR is low, the traditional communication algorithm basically cannot transmit any semantic information. However, when the SNR is high, the PSNR reaches the performance saturation of the traditional communication algorithm. Thus, the accuracy of the proposed algorithm in this paper is significantly better than that of the traditional communication algorithms, especially when the SNR is improved, the performance improvement is very obvious. This is because while LDPC codes and QAM modulation can improve the robustness of data transmission, the JPEG and BPG compression algorithm are lossy and can lead to irreversible information loss. Combining the JPEG, LDPC and QAM techniques may reduce the error resilience of the overall system, leading to degraded image quality or data integrity in the presence of transmission errors. In comparison, the STSC algorithm retains the semantic information to the greatest extent, which ensures the effective transmission of images.



In order to evaluate the convergence of the proposed STSC model under the FL framework, the convergence analysis experiment of the FSSC model is also carried out in this section. Taking SNR=12 as an example, Fig. <ref> shows how the target loss varies as the training rounds/time increases in the case of FL training and local training, where the target loss is related to the degree of image reconstruction. As can be seen from the figure, on AWGN channels, the target loss of the communication model decreases rapidly in a short period of time, then fluctuates for a period of time, and finally remains unchanged. In other words, the model eventually reaches convergence. Moreover, the convergence efficiency of the FL training model is higher than that of the local training of the client. At the same time, the final training loss of global training model is relatively lower. This is because the local model uses a local single data set for local training, which makes it impossible to fully learn the global semantic distribution characteristics of the data. However, FL training has the advantages of data integration and improving model performance, and thus solves this problem well. 


Fig. <ref> shows how the PSNR change with the channel SNR for the global model and the local model without federated aggregation after model convergence. From the figure, it can be observed that with the increase of SNR, the PSNR of the global model and the local model without joint aggregation also significantly improved. In contrast to local models devoid of collaborative aggregation, global models can further enhance the PSNR by 2-3dB. This is attributed to the global model's ability to learn richer noise models and finer image semantic features through effective information aggregation and synchronization. Consequently, under high signal-to-noise ratio conditions, its PSNR can continue to improve, demonstrating superior generalization and noise resistance capabilities. This is because the image reconstruction is better as the channel condition improves. Moreover,  FSSC avoids the sharing of sensitive data, thus protecting the privacy of users and motivating users to participate in semantic communication under the FL framework to a certain extent.



§ CONCLUSION


In this paper, we proposed a semantic communication architecture for image transmission over wireless channels in multi-user cases, called FSSC. In this architecture, the local STSC semantic communication system model accurately extracts semantic information on the respective dataset and carries out the semantic communication model training. Then, joint training is performed through FL parameter aggregation to minimize the target loss function of the reconstructed image. Simulation results show that the algorithm can converge effectively. This approach is able to combine the semantic information exists in the diverse data from clients, and distributes the computational workload from edge to users, without compromising the user's privacy. Under the premise of privacy protection, the proposed algorithm can greatly reduce the amount of data required by each client while ensuring the fidelity in recovered images. Therefore, the FSSC proposed in this paper is a promising candidate scheme for multi-user image semantic communication system. 

00

ref1
K. Lu, Q. Zhou, R. Li, Z. Zhao, X. Chen, J. Wu, and H. Zhang, “Rethinking modern communication from semantic coding to semantic communication,” IEEE Wireless Communications, vol. 30, no. 1, pp. 158-164, February 2023.

ref2
E. Bourtsoulatze, D. Burth Kurka and D. Gündüz, “Deep Joint Source-Channel Coding for Wireless Image Transmission," IEEE Transactions on Cognitive Communications and Networking, vol. 5, no. 3, pp. 567-579, September 2019.

ref3
N. Farsad, M. Rao, and A. Goldsmith, “Deep learning for joint source-channel coding of text,” in 2018 IEEE international conference on acoustics, speech and signal processing (ICASSP), Calgary, AB, Canada, April 2018. 

ref4
H. Xie, Z. Qin, G. Y. Li, and B. H. Juang, “Deep learning enabled semantic communication systems” IEEE Transactions on Signal Processing, vol. 69, pp. 2663–2675, April 2021. 

ref5
K. Choi, K. Tatwawadi, A. Grover, T. Weissman, and S. Ermon, “Neural joint source-channel coding,” in Proceedings of the 36th International Conference on Machine Learning (ICML), Long Beach, CA, United states, May 2019. 

ref6
J. Xu, B. Ai, W. Chen, A. Yang, P. Sun, and M. Rodrigues, “Wireless image transmission using deep source channel coding with attention modules,” IEEE Transactions on Circuits and Systems for Video Technology, vol. 32, no. 4, pp. 2315-2328, April 2022. 

ref7
Claude E Shannon, “A mathematical theory of communication,” The Bell system technical journal, vol. 27, no. 3, pp. 379–423, July 1948.

ref8
D. Gündüz, Z. Qin, I. E. Aguerri, H. S. Dhillon, Z. Yang, A. Yener, K. K. Wong, and C.-B. Chae, “Beyond transmitting bits: Context, semantics, and task-oriented communications,” IEEE Journal on Selected Areas in Communications, vol. 41, no. 1, pp. 5–41, November 2022.

ref9
M. Kountouris and N. Pappas, “Semantics-empowered communication for networked intelligent systems,” IEEE Communications Magazine, vol. 59, no. 6, pp. 96–102, June 2021.

ref10
W. Yang et al., “Semantic Communications for Future Internet: Fundamentals, Applications, and Challenges,” IEEE Communications Surveys & Tutorials, vol. 25, no. 1, pp. 213-250, November 2023.

ref11
Z. Liu et al., “Swin Transformer: Hierarchical Vision Transformer using Shifted Windows,” in IEEE/CVF International Conference on Computer Vision (ICCV), Montreal, QC, Canada, October 2021.

ref12
K. Yang, S. Wang, J. Dai, K. Tan, K. Niu and P. Zhang, “WITT: A Wireless Image Transmission Transformer for Semantic Communications,” in IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), Rhodes Island, Greece, June 2023.

ref13
H.B. Mcmahan, E. Moore, D. Ramage, S. Hampson, and B. A. y. Arcas, “Communication-Efficient Learning of Deep Networks from Decentralized Data,” in International Conference on Artificial Intelligence and Statistics (AISTATS), Fort Lauderdale, Florida, USA, February 2017.

ref14
A. Krizhevsky, “Learning multiple layers of features from tiny images,” Master's thesis, University of Tront, April 2009.

ref15
D. P. Kingma and J. Ba, “Adam: A method for stochastic optimization,” arXiv preprint arXiv:1412.6980, December 2014.


ref16
G. K. Wallace, “The JPEG still picture compression standard,” IEEE Transactions on Consumer Electronics, vol. 38, no. 1, pp.xviii-xxxiv, February 1992.

ref17
R. Gallager, “Low-density parity-check codes,” IRE Transactions on Information Theory, vol. 8, no. 1, pp. 21–28, January 1962.








 Paper Title: Tamper-Resistant Safeguards for Open-Weight LLMs.



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




[1]Equal contribution. ^Equal advising. Correspondence to mailto:rishubt2@illinois.edurishubt2@illinois.edu.
[2]Part of work done during an internship at the Center for AI Safety.





§ ABSTRACT 
Rapid advances in the capabilities of large language models (LLMs) have raised widespread concerns regarding their potential for malicious use. Open-weight LLMs present unique challenges, as existing safeguards lack robustness to tampering attacks that modify model weights. For example, recent works have demonstrated that refusal and unlearning safeguards can be trivially removed with a few steps of fine-tuning. These vulnerabilities necessitate new approaches for enabling the safe release of open-weight LLMs. We develop a method, called , for building tamper-resistant safeguards into open-weight LLMs such that adversaries cannot remove the safeguards even after thousands of steps of fine-tuning. In extensive evaluations and red teaming analyses, we find that our method greatly improves tamper-resistance while preserving benign capabilities. Our results demonstrate that tamper-resistance is a tractable problem, opening up a promising new avenue to improve the safety and security of open-weight LLMs.






§ INTRODUCTION







The most capable open-weight large language models (LLMs) released over the past year now rival closed-source frontier models <cit.>. The availability of open-weight LLMs for anyone to download and use has yielded numerous benefits, including lowering costs for end users and enabling academic research on safety and security <cit.>. However, as these models become increasingly powerful, many have raised concerns that they could be repurposed by malicious actors to cause harm, motivating research on how to safeguard these models against malicious use.


Existing open-weight models often adapt safeguards designed for closed-weight models served through APIs <cit.>. These safeguards include refusal mechanisms and preference-based training, and they have provided substantial robustness against input-based jailbreaking attacks. However, recent work has demonstrated these safeguards are trivially defeated by attacks that edit model weights, breaking down after only a handful of fine-tuning steps <cit.>. This poses a serious problem for open-weight models, because adversaries have full access to model weights and can tamper with
built-in safeguards.

The vulnerability of open-weight models to tampering attacks poses risks for model developers as well. Under background tort law, AI developers must exercise reasonable care, meaning they have an obligation to take reasonable precautions to prevent foreseeable harm. If malicious actors can easily customize models to cause critical harm, model developers may inadvertently violate reasonable care standards and become open to liability under existing law. Thus, there is an urgent need for more robust safeguarding techniques that can withstand tampering attacks.

In this work, we study the problem of tamper-resistant safeguards for LLMs. This problem is depicted in Figure <ref>. Unlike existing research on LLM safeguards, we focus on attacks that modify model weights, which we refer to as tampering attacks. This problem has been considered very challenging and by some intractable, as no method has yet provided substantial robustness to these attacks. However, making progress on this problem would provide a valuable tool to regulators and model developers by ameliorating the dual-use dilemma of open-weight models <cit.>. 

To demonstrate that progress on this problem is possible, we develop the first LLM safeguards that obtain strong robustness against a wide variety of tampering attacks. Our approach allows developers to add a safeguard such that tampering attacks cannot easily remove the safeguard, while preserving the general capabilities of the LLM. We achieve this by performing adversarial training against tampering attacks, leveraging approaches from meta-learning. We identify various crucial factors that enable our method to work, including the choice of tamper-resistance loss, the selection of train-time adversaries, and the two-stage approach that we use for building in safeguards.



We apply our method to develop tamper-resistant unlearning and refusal safeguards. In experiments, we demonstrate that our safeguards are far more robust to tampering attacks than prior methods. We stress-test our safeguards with extensive red teaming evaluations against  test-time adversaries, demonstrating resistance to fine-tuning attacks up to 5,000 steps. We hope our results foster future work on this important problem. Our experiment code and models are available at <https://github.com/rishub-tamirisa/tamper-resistance>.





§ RELATED WORK



  
Adversarial attacks on LLMs. Due to the extensive pre-training distribution of modern LLMs, they are prone to generating harmful content <cit.>. To mitigate this, many LLMs undergo fine-tuning to implement safeguards <cit.>, using methods such as reinforcement learning from human feedback (RLHF) <cit.> and direct preference optimization (DPO) <cit.>. While effective for normal use, these safeguards have been shown to be brittle, breaking down under jailbreak attacks <cit.> or a handful of fine-tuning steps on “uncensored” data <cit.>. This suggests current techniques for LLM alignment are inadequate, raising security concerns after deployment.

r0.5
    

    
    < g r a p h i c s >

    Comparison of our method to 12 baseline safeguards. Unlike prior methods, provides far greater tamper-resistance at similar levels of general capability, measured via MMLU. Tamper-resistance is computed as the normalized error on WMDP Biosecurity, Chemical Security, and Cybersecurity questions <cit.>, averaged across up to  fine-tuning attacks.
    
    







  
LLM safeguards. Since the discovery of these attacks, many safeguards have been proposed to defend against them. Against jailbreak attacks, defenses include system-level defenses <cit.> that modify or filter model inputs or outputs and model-level defenses such as adversarial training <cit.>. Alternatively, some works explore machine unlearning as a way to remove harmful knowledge entirely with techniques such as influence functions <cit.>, maximizing loss on forget sets <cit.>, or modifying representations <cit.>. However, jailbreaking defenses are not fully robust to adaptive adversaries <cit.>, and existing unlearning methods are not robust to adversaries with access to model weights <cit.>.





  
Robust safeguards. Several works have explored the tamper-resistance of unlearning methods for image classification <cit.>. For bidirectional BERT-style models, <cit.> proposed a meta-learning approach for robustly preventing models from learning harmful tasks. In concurrent work, <cit.> proposed a method extending this approach to small-scale vision classifiers and diffusion models. Recently, <cit.> discussed the potential for robust unlearning in LLMs to improve the safety of open-source models, and <cit.> proposed evaluation metrics for robust unlearning in LLMs. To the best of our knowledge, no methods have been proposed for autoregressive LLMs that are robust to tampering attacks. 

Several concurrent works have explored ways of defending LLM refusal mechanisms against fine-tuning. <cit.> add a perturbation loss to keep new embeddings close to the original, while <cit.> regularize harmful representations to look random. Unfortunately, the authors acknowledge that these approaches have highly limited robustness. We corroborate this in our comparisons, finding that they lack robustness to the tampering attacks in our evaluations.


























§ TAMPER-RESISTANT SAFEGUARDS






 §.§ Threat Model


We assume the defender releases an LLM with weights θ_G and a safeguard G applied. The defender's goal is to design G such that θ_G obtains high values on (θ_G) and (θ_G). Moreover, the defender seeks to preserve a high value of (θ_G) after the adversary's move. We consider a compute-bounded adversary with unrestricted access to θ_G, enabling attacks that directly modify θ_G. We refer to these as “tampering attacks.” The adversary's goal is to obtain a model θ'_G that minimizes the safety metric given reasonable compute limits, such as fine-tuning for 1,000 to 5,000 steps. We assume the adversary will not spend a significant fraction of the compute required to pre-train the LLM, since at that point they could train their own model without safeguards.









 §.§ Problem Definition and Metrics



We describe a general notation for quantifying the tamper-resistance of safeguards. Define G, θ_G, , and  as in the threat model. Let  denote a compute-bounded adversarial attack that maps θ_G to θ'_G, with stronger attacks obtaining lower values of (θ'_G). We say that a safeguard G is tamper-resistant if its post-attack (θ_G') is high across a broad range of strong test-time adversarial attacks 𝒜_test.

Note that θ_G often modifies an underlying θ that lacks safeguards, often through a fine-tuning procedure. Additionally, strong tamper-resistance can be obtained if the safeguard simply overwrites θ with noise, but this model would no longer be useful. Thus, maintaining a high (θ_G) is crucial, and evaluation of a safeguard must consider both its tamper-resistance and how well it preserves general capabilities. 


We focus on two common safeguard domains: weaponization knowledge restriction and harmful request refusal. In each domain, we define safety and capabilities test metrics, which we use alongside test-time adversaries to evaluate tamper-resistant safeguards.



  
Weaponization knowledge restriction. In weaponization knowledge restriction, safeguards prevent the model from producing text about weaponization knowledge, while preserving capabilities for benign knowledge domains. Existing safeguards of this nature include representation engineering methods like circuit breaking <cit.>. The  is defined as error on a forget set, and the  is defined as accuracy on a retain set. Specifically, we consider the problem of restricting biosecurity, chemical security, and cybersecurity knowledge, and evaluate the resulting model on the Weapons of Mass Destruction Proxy (WMDP) benchmark <cit.>. WMDP contains 3,668 multiple-choice questions, spanning biosecurity, chemical security, and cybersecurity knowledge. Importantly, WMDP questions do not evaluate hazardous knowledge directly, but instead measure proxy expert-level knowledge for each hazardous domain, such that suppressing the expert-level knowledge would also suppress the hazardous knowledge. We define the forget set as the respective hazardous knowledge subject in WMDP, and retain set as the complement of the given subject in MMLU <cit.>, a multi-task question-answering benchmark spanning 57 tasks across a variety of knowledge domains.



  
Harmful request refusal. In the harmful request refusal setting, safeguards prevent the model from producing “harmful” outputs. We define the  as the complement of average Attack Success Rate (ASR) of various jailbreaking attacks, while the  captures the conversational abilities of θ_G. Specifically, we use a static set of test cases from HarmBench, an automated red-teaming framework for measuring prompt jailbreak robustness in LLMs, to evaluate jailbreak ASR <cit.> after tampering attacks. We use MT-Bench, a multi-turn question-answering benchmark graded by an LLM judge, to evaluate conversational abilities <cit.>.


    


 §.§ Red Teaming




To properly measure the robustness of tamper-resistant safeguards, we conduct red-teaming with up to  adversaries, including many that are unseen at training time. In our evaluations, we subject our method to adversaries with varying compute budgets, access to held-out datasets, and diverse hyperparameters. For fine-tuning adversaries, we vary the learning rate, learning rate scheduler, optimization algorithm, and batch size. Many of these adversaries were fixed during early experiments, with some added over time as we found attacks that broke intermediate versions of our method. Extensive stress testing of this nature is critical for obtaining confidence in a tamper-resistant safeguard. For research on developing these safeguards, extensive red teaming also allows measuring incremental progress, using the number and strength of existing attacks one can defend against as a robustness metric.












§ SAFEGUARD TAMPER-RESISTANCE TRAINING









To obtain tamper-resistant safeguards, we propose a new method outlined in Algorithm <ref> inspired by adversarial training and meta-learning to directly strengthen LLM safeguards against tampering attacks, called Tampering Attack Resistance (). We identify unique properties of this adversarial training regime and leverage them to improve robustness.

Our method for training tamper-resistant safeguards consists of two phases: (1) model safeguarding and (2) tamper-resistance training.



 §.§ Model Safeguarding
 The method begins by including an initial safeguard G into a base model θ. For example, initial safeguards for knowledge restriction can be drawn from a wide variety of existing methods, including circuit breaking <cit.> or constrained gradient ascent for a particular knowledge domain. Similarly, we can include a refusal safeguard by performing RLHF <cit.> or DPO <cit.> on refusal completions. Importantly, these initial safeguards do not need to be tamper-resistant. Empirically, we find that this safeguarding step is important for obtaining strong tamper-resistance.







 §.§ Tamper-Resistance Training






Starting from θ_G_0, we train the tamper-resistant θ_G using a novel adversarial training procedure. Namely, we train against a set of tampering attacks 𝒜_train, where the defender's objective is to maximize a proxy  after applying an adversarial attack ∼𝒜_train to θ. Since it may not be feasible to differentiate through , we draw on insights from prior work in meta-learning, defining (θ_G) = θ'_G = θ_G + '(θ_G) as a perturbation on top of initial parameters, where backpropagation through ' is approximated with a straight-through estimator <cit.>. 








We focus on supervised fine-tuning (SFT) adversaries where  applies several steps of optimization to θ_G, which allows straight-through estimation through ' to benefit from the setting and approximations of first-order MAML <cit.>. However, we note key differences in our approach from standard meta-learning and prior methods <cit.>. In particular, traditional meta-learning techniques seek to obtain a model initialization that is close to optimality on multiple test distributions. In our setting, we seek to obtain an initialization that is far from optimality on multiple adversaries' test distributions. Novel to our approach in this new setting is the use of a tamper-resistance loss in the “outer loop” that differs from the fine-tuning adversary's loss function and serves to maximize the proxy safety metric. We depict this structure in Algorithm <ref>, and explain the objective below.





  
Impeding the adversary's loss. The aim of tamper-resistance training is to prevent adversaries with large compute budgets from reducing the  at test-time. In adversarial training for tamper-resistance, we define a tamper-resistance loss ℒ_TR that counters . We operationalize our goal of avoiding adversary optimality as searching for θ such that ℒ_TR is minimized for (θ).

Empirically, we find that the choice of tamper-resistance loss ℒ_TR significantly affects this goal. Prior work <cit.> negates the loss of a fine-tuning adversary, in which the aim is to arbitrarily maximize the adversary's loss throughout fine-tuning. This formulation has two issues: (1) maximizing a cross-entropy loss can cause divergence; (2) empirically we observe that when using this objective against fine-tuning adversaries, the model learns to explode the adversary's loss for the first few inner loop steps, while loss at later steps remains low. In Figure <ref>, we show the difference in choosing ℒ_TR to be a clamped negative cross-entropy loss vs. negative entropy loss for weaponization knowledge restriction. For the latter, ℒ_TR is eventually satisfied for all inner loop steps. For harmful request refusal, we choose ℒ_TR to be the DPO loss <cit.>. We provide further detail on the choice of ℒ_TR in both settings in Appendix <ref>.



  
Tamper-resistance objective. We now describe the general proxy objective used for preventing a tampering attack from recovering weaponization knowledge or harmful behavior. For a given , let 𝒟_TR and ℒ_TR respectively be a dataset and loss function such that minimizing ℒ_TR(θ_G, 𝒟_TR) serves as a proxy objective for maximizing the (θ_G). We define 𝒟_retain and ℒ_retain correspondingly for (θ_G). The defender's objective is to solve the following optimization problem:

    min_θλ_TR·𝔼_∼𝒜_train[  ℒ_TR((θ); 𝒟_TR) ] + λ_retain·ℒ_retain(θ; 𝒟_retain),
 where ℒ_TR is a tamper-resistance loss that counters (θ). The ℒ_retain term is a representation engineering <cit.> inspired retain loss for preserving performance on the capabilities proxy dataset 𝒟_retain, given by

    ℒ_retain(θ; 𝒟_retain) = 𝔼_x∼𝒟_retain[ℒ_LM(θ, x) + ‖ h_θ(x) - h_θ_G_0
    (x) ‖_2^2]
 where h_θ(·) returns the residual stream hidden states for model parameters θ and ℒ_LM is the standard language modelling cross-entropy loss. Empirically, we find that pushing retain-set residual stream representations to be close to the base model θ_G_0 via the ℓ_2-norm loss in Equation <ref> maintains a high (θ_G). In Equation <ref> we include λ_TR and λ_retain as scalar weightings for the tamper-resistance loss and retain loss, respectively. We provide further details on the design of the tamper-resistance loss function in Appendix <ref> as well as an efficiency trick for sampling fine-tuning attacks for in Appendix <ref>.



























































§ EXPERIMENTS


We evaluate in weaponization knowledge restriction and harmful request refusal settings, with results shown in Table <ref> and Table <ref> respectively. We discuss the setup, baselines, and analysis for our results. In each setting, we use a specific set of training adversaries 𝒜_train and test adversaries 𝒜_test. Further experiment details are presented in Appendix <ref>.



 §.§ Weaponization Knowledge Restriction

We now describe the setup, baselines, and results for our weaponization knowledge restriction experiments, including the knowledge domains, optimizers, and evaluation details.



  
Setup.
 We focus on implementing tamper-resistant safeguards for restricting proxy weaponization knowledge about biosecurity, chemical security, and cybersecurity from Llama-3-8B-Instruct <cit.> that has been initially safeguarded via the Random Mapping method discussed in Appendix <ref>. For each weaponization domain, we assign 𝒟_𝒯ℛ to the corresponding forget set described in Appendix <ref>. We proceed to sample train-time 64-step fine-tuning attacks from different data distributions, detailed in Appendix <ref>. We use N = 750 outer loop steps, ScheduleFree AdamW <cit.> with a learning rate of 2×10^-5 as the outer loop tamper-resistance optimizer, and loss scales of λ_TR = 4.0, λ_retain = 1.0. Lastly, We evaluate Pre-Attack and Post-Attack accuracy on corresponding WMDP subjects <cit.> averaged across all adversaries in Appendix <ref>, and measure benign capabilities via the complement of subjects related to each proxy weaponization domain in MMLU <cit.>.







  
Baselines.
We evaluate two recently proposed knowledge restriction methods: RMU <cit.> and LLMU <cit.>. We also design two baseline methods for knowledge restriction: Min Posterior, which minimizes posterior loss on forget set tokens; Max Entropy, which maximizes entropy on forget set tokens. Two additional methods, MLAC <cit.> and SOPHON <cit.>, require substantial modifications for the LLM setting, so we show results on adapted versions of these baselines in <Ref>.









  
Results.
We show weaponization knowledge restriction safeguard results on Llama-3-8B-Instruct in <Ref> and Figure <ref>. These results are averaged across all 28 adversaries from Appendix <ref>. Our large-scale experiments corroborate the findings in recent work that existing LLM safeguards are extremely brittle to fine-tuning attacks. By contrast, maintains low post-attack forget accuracy across all three domains. We also observe that maintains non-negligible retain accuracy, although there is a modest trade-off, much like how adversarial robustness trades off benign accuracy. In Figure <ref>, we observe that is robust to significantly more fine-tuning attacks than all prior methods. While existing baselines break down under most attacks, obtains a post-attack forget accuracy near random chance for nearly all attacks, indicating a successful defense.

Overall, we find that provides significantly more robustness to realistic fine-tuning attacks than all prior methods, including attacks as strong as 5,000 steps of optimization on completely held-out data. These results demonstrate for the first time that obtaining strong tamper-resistance for open-weight LLMs may be possible.





 §.§ Harmful Request Refusal







 We now describe the setup, baselines, and results for our harmful request refusal experiments, including the datasets used and evaluation details. 


  
Setup.For harmful request refusal training, we seek to make existing refusal safeguards in Llama-3-8B-Instruct robust to tampering attacks. We sample train-time adversaries that perform 64-step SFT attacks using the Anthropic-HH-RLHF dataset <cit.>, following the methodology in Appendix <ref>. Similar to the weaponization knowledge restriction setting, we use N = 100 outer loop steps, ScheduleFree AdamW <cit.> with an LR of 6 × 10^-5 as the outer loop tamper-resistance optimizer, and loss scales of λ_TR = 0.1, λ_retain = 1.0. We evaluate the Post-Attack jailbreak attack success rate (ASR) on HarmBench <cit.> after the tampering attacks in Appendix <ref>, and measure benign capabilities preservation via MT-Bench <cit.>, which evaluates multi-turn conversation ability. 






  
Baselines. We consider 4 baselines alongside our model: Llama-3-8B-Instruct (Refusal Trained); Representation Rerouting (RR) <cit.> on Llama-3-8B-Instruct, which trains to push representations for harmful input prompts to be orthogonal to the original corresponding representations in Llama-3-8B-Instruct; R2D2 <cit.> on Zephyr-7B <cit.>, which performs adversarial training against GCG attacks <cit.>; and RepNoise <cit.> on Llama-2-7B <cit.>, which regularizes harmful representations to Gaussian noise.  





  
Results.
We show refusal results in <Ref>. While the Refusal Training, RR, and R2D2 baselines resist jailbreak attacks in HarmBench before tampering, we find that percentage attack success rate jumps up to above 77 after tampering, while our method only rises to 61.7. Since we apply our refusal safeguard to Llama-3-8B, it does reduce MT-Bench by 1.7. However, this exceeds the MT-Bench score of fairly capable open-weight models, indicating that benign capabilities are largely preserved. We leave the exploration of the full impact on capabilities to future work. Additional results are in <Ref>. In general, we find that our model refuses more Post-Attack jailbreaks than previous baselines, and demonstrates the flexibility of the tamper-resistance objective to accommodate the harmful request refusal setting.









 §.§ Analysis





  
Red teaming.
To stress-test the tamper-resistance of our models, we conduct an extensive suite of supervised fine-tuning attacks with  distinct adversaries. We vary the optimizer, number of optimization steps, learning rate, learning rate schedule, fine-tuning dataset, batch size, and overall fine-tuning method (e.g., full fine-tuning versus parameter-efficient fine-tuning). By default, our attacks use 1,000 fine-tuning steps, although some use 5,000 steps. Full details for these adversaries are provided in <Ref>.

We show red teaming results in Figures <ref> and <ref>. While baseline safeguards withstand fine-tuning attacks in a small number of cases, most adversaries succeed in removing the safeguards. By contrast, our safeguard is robust to a wide range of adversaries. This shows that tamper-resistance is a tractable problem on which progress can be made. However, our method is currently not robust to parameter-efficient fine-tuning (PEFT) attacks (adversaries 27 and 28), highlighting the importance of extensive red teaming when developing tamper-resistant defenses. We hypothesize that future work could easily address this issue, as we demonstrate in <Ref> that targeted patching of vulnerabilities is possible.


r0.5
    

    
    < g r a p h i c s >

    Our safeguard remains robust to fine-tuning attacks that greatly exceed the 64 steps used by train-time adversaries. For the LLMU safeguard, the adversary's loss quickly decreases into the recovery region. By contrast, maintains flatness at a high loss for all 1,000 steps.
    

    





  
Generalization to stronger test-time attacks. In Figure <ref>, we show a fine-tuning attack at an LR of 2×10^-5 that attempts to recover biosecurity knowledge on our model and a model safeguarded with LLMU. We find that the tamper-resistance of generalizes far beyond the 64 steps used by train-time adversaries. Surprisingly, we observe that the test-time adversary's cross-entropy loss does not decrease below 7 for all 1,000 steps. Moreover, the loss enters a plateau and does not decrease at all after 200 steps. In <Ref>, we find that the length and height of this plateau can be increased by increasing the number of inner loop steps during adversarial training. As a point of reference, we show the progression of the same attack on LLMU. In this case, the adversary's loss decreases to within the recovery region in under 20 steps. We note that the adversary's test loss decreasing into the recovery region does not always correspond with recovery on downstream metrics (e.g., WMDP), but often does.































    



























§ CONCLUSION


We introduced a novel method for implementing tamper-resistant safeguards for LLMs and explored applications in weaponization knowledge restriction and harmful refusal training. We compare our results to prior work in each setting, finding that our method is the first method robust under the rigorous red-teaming evaluation that we consider. More broadly, we demonstrate that progress on open-weight tamper-resistance is tractable. We believe this line of research is crucial for enabling ongoing deployment of robust, open-weight LLMs, ensuring their alignment with regulatory frameworks and preemptively addressing the risk of malicious use.





  
Acknowledgements. We thank Steven Basart and Luis Fernandez for providing valuable feedback for the paper, and Andriy Novykov for providing support for compute resources at the Center for AI Safety. We thank the Center for AI Safety for providing significant compute resources for this project. We also thank Volodymyr Kindratenko and the National Center for Supercomputing Applications (NCSA) and Illinois Campus Cluster Program (ICCP) for supporting our computing needs. This work used NVIDIA GPUs at NCSA Delta through allocations CIS230117 from the Advanced Cyberinfrastructure Coordination Ecosystem: Services & Support (ACCESS) program, which is supported by NSF Grants #2138259,
#2138286, #2138307, #2137603, and #2138296.











abbrvnat










§ LIMITATIONS




Our method for training tamper-resistant safeguards demonstrates considerable robustness against a wide range of tampering attacks, yet several avenues for improvement remain: (1) While we focus on supervised fine-tuning attacks, the broader spectrum of open-weight tampering techniques necessitates diverse future red-teaming efforts. (2) Scaling to larger models poses computational challenges that require optimization to reduce overheads. 



Tamper-resistance alone cannot fully mitigate the risks of malicious AI use. While it raises the initial costs for adversaries, it can eventually be circumvented. Once open-weight models are released, they cannot be “unreleased,” leaving any compromised defenses permanently vulnerable. Therefore, tamper-resistance should be considered a supplement to the broader effort of of improving the offense-defense balance of AI systems. Addressing these limitations will improve the robustness of LLMs to tampering and better support open-weight model developers.







§ METHOD DETAILS





 §.§ Initial Weaponization Knowledge Restriction Safeguard


Prior to tamper-resistance training, we install a safeguard that achieves surgical knowledge restriction on the target hazardous domain. Let h_θ(𝒟) denote the distribution of post-decoder layer residual stream activations for input sequences sampled from some data distribution 𝒟 and model weights θ. We define  (x) for some input sequence x, which returns fixed Gaussian-sampled vectors that are chosen via hashing the corresponding input token for each residual stream index of x in θ. As a proxy for scrubbing target representations according to downstream task labels, we propose a weaponization knowledge restriction safeguard termed Random Mapping, which maps h_θ(𝒟_TR) to random noise as follows:



    min_θ𝔼_x∼𝒟_TR[1 - | h_θ(x) ·(x)/‖ h_θ(x)‖‖(x)‖| ] + ℒ_LM(θ; 𝒟_retain)


The objective of Equation <ref> maximizes cosine similarity between row vectors in the residual stream in every layer of the LLM from h(𝒟_TR) and the hashed random vectors from (·). By providing each token's residual stream a unique random vector to push toward, the loss encourages a “re-mapping” of token representations from 𝒟_TR to the noised vectors. We include an additional term for preserving performance on 𝒟_retain via the language-modelling cross-entropy loss ℒ_LM. We show the performance of the raw Random Mapping safeguard as an ablation in Table <ref>, listed as “Excl. Adv. Training.”




 §.§ Designing the Tamper-resistance Loss




  
Weaponization knowledge restriction. For weaponization knowledge restriction, we summarize our intuition for ideal tamper-resistance loss design, corroborated by our empirical findings in Figure <ref>: we seek to flatten the adversary's loss at a high value, rather than simply raise its y-intercept.

We choose the tamper-resistance loss ℒ_TR as an entropy loss to be maximized during the adversary's cross-entropy fine-tuning trajectory, since maximizing entropy would impede the adversary's cross entropy loss from decreasing during fine-tuning. In other words, we wish to obtain θ such that after an adversary performs a fine-tuning attack on θ via a cross-entropy loss, entropy is still high. We find that this formulation achieves the desired flattening behavior, and we depict the difference in flattening between the choosing ℒ_TR to be a negative cross-entropy loss and negative entropy loss in Figure <ref>. In the lefthand plot, where ℒ_TR is a cross-entropy loss, loss only increases in the first inner loop step. In the righthand plot, where ℒ_TR is a negative entropy loss, entropy is eventually maximized in all inner loop steps. Figure <ref> also demonstrates the generalization of the flat adversary loss behavior beyond the length of the simulated fine-tuning trajectories during .



  
Harmful request refusal. For harmful request refusal, we choose ℒ_TR to be the DPO loss <cit.>, which works as follows. Given a DPO dataset containing pairs of rejected and refusal completions, the sampled  performs SFT on rejected completions, and the tamper-resistance loss ℒ_TR is a DPO loss computed on the pair chosen and rejected completions on parameter coordinates along the  trajectory. This encourages to find an initialization θ such that after a harmful fine-tuning attack, the model still prefers refusal completions over harmful completions when given an harmful prompt. While this does not necessarily encourage a flat adversary loss, we find empirically in Table <ref> that this formulation increases the average (θ_G) defined in Section <ref> after fine-tuning attacks on harmful data, detailed in Section <ref>.



We also observe that the length of the simulated adversary SFT trajectory during training affects test-time generalization in both Figure <ref> and Appendix <ref>. In particular, larger values of K result in increased tamper-resistance for longer SFT attacks. However, to maintain reasonable runtime efficiency, we need a more efficient sampling technique than simply running K independent trajectories of varying length for every outer-loop step in Algorithm <ref>, which we describe in Section <ref>.



 §.§ Efficiently Sampling Fine-tuning Attacks
 


Optimizing Equation <ref> with gradient descent requires simulating K tampering attacks for each tamper-resistance optimizer update, which is prohibitively expensive to run when the sampled  performs SFT and θ contains billions of parameters. Inspired by prior work on snapshot ensembles <cit.>, we leverage an efficiency trick: we can reuse the coordinates along steps of a single adversary fine-tuning trajectory of length K to obtain K-1 additional (though non-independent) trajectories of increasing length. Using this trick, we collect all K parameter coordinates along the trajectory into a single batch for computing the tamper-resistance losses, effectively sampling  from 𝒜_train non-IID. To further improve runtime efficiency, we do not compute the tamper-resistance loss ℒ_TR on all K steps and instead sub-sample coordinates along the trajectory for computing ℒ_TR within an adversary batch, for example every 4 adversary optimization steps. Additionally, we reduce variance in the tamper-resistance gradient by computing the tamper-resistance loss at each inner loop step on the same held-out batch, denoted as x_TR in Algorithm <ref>. 



 §.§ Implementation Details and Resource Requirements


We perform training on Llama-3-8B-Instruct <cit.> with 8 NVIDIA 80GB A100 GPUs, leveraging distributed training via FSDP <cit.>. We use ZeRO Stage 3 from DeepSpeed <cit.>, which shards optimizer states, gradients, and parameters during training. While the efficiency trick in Appendix <ref> improves runtime, we note additional considerations for conserving GPU memory. 

First, simulating fine-tuning attacks that require additional state (e.g., momentum) in the inner loop of Algorithm <ref> requires initializing a fresh optimizer for every outer loop iteration. Since we use an outer-loop optimizer that also requires maintaining state (ScheduleFree AdamW <cit.>), we move the outer loop optimizer to the CPU before instantiating inner-loop optimizers.

Second, first-order meta-learning in smaller models can typically be implemented by running multiple forward passes for each inner loop iteration, averaging losses, then backpropagating on the averaged loss term. However, because each inner-loop tamper-resistance loss term (ℒ_TR in Algorithm <ref>) is computed on a separate forward pass, this requires maintaining K computation graphs in memory. Since this is infeasible on reasonable hardware for LLMs with billions of parameters, we circumvent this inefficiency by accumulating tamper-resistance gradients in a separate data structure (g_TR in Algorithm <ref>). We note that this can be done without using additional  and  distributed operations, since tamper-resistance gradient accumulation and application to the pre-inner loop model parameters (θ_i-1 in Algorithm <ref>) can be computed solely on sharded gradients. 












§ ADDITIONAL EXPERIMENTS





 §.§ Benign Fine-Tuning






An important property of open-weight models is that they can be fine-tuned to improve performance on custom data or in specific domains. Thus, ideal tamper-resistant safeguards should allow continued fine-tuning of a model while preserving the safeguard. We evaluate whether models can be fine-tuned on data unrelated to the safeguard using economics as an example domain. Using models with biosecurity and cybersecurity safeguards, we perform supervised fine-tuning on the WMDP auxiliary economics corpora <cit.>. We fine-tune models for 800 steps using a learning rate of 2×10^-6 and a batch size of 64. For the test set, we evaluate average accuracy on the corresponding MMLU High School Macroeconomics and Microeconomics subjects before and after fine-tuning. To confirm that the safeguard remains tamper-resistant in this setting, we also evaluate accuracy on corresponding WMDP subjects before and after fine-tuning.

We show the results of this evaluation in <Ref>. We find that accuracy on economics questions can be improved by 7.5 percentage points without recovering significant hazardous knowledge. This illustrates that strong tamper-resistance can be compatible with benign model editing. Additional experiment details are discussed in Appendix <ref>.



 §.§ Varying the Train-time Inner-loop Length K



r0.5
    

    
    < g r a p h i c s >

    Comparison of a 1,000-step SFT attack against with the inner-loop length K varied between {8, 16, 32, 64}. Test-time loss plateau magnitude and duration tend to increase as K increases.
    




Recall that via the efficiency trick discussed in Appendix <ref>, a single inner loop trajectory of length K during returns the K sampled attacks in Algorithm <ref>.  We compare the test-time loss robustness as we vary the length of the inner loop K during , running fine-tuning attacks for 1,000 steps on a held-out forget dataset for biosecurity weaponization (Adversary 8 in Table <ref>). For each value of K, we observe a plateau in the test loss that drops off at later steps as K increases. This suggests that the robustness of improves as the inner loop length increases. Prior work also corroborates that increasing the inner-loop length during meta-learning increases test-time generalization <cit.>. We note the contrast to conventional meta-learning methods mentioned in Section <ref>, in which typical meta-learning applications seek optimality after as few test-time steps as possible <cit.>. Here, our results suggest that the objective is incentivized to run with as many inner loop steps as possible. We find that K=64 provides significant robustness to the range of adversaries we consider in Section <ref> and Appendix <ref>, while balancing computational efficiency as discussed in Appendix <ref>. 






 §.§ Targeted Patching of Vulnerabilities



In Section <ref>, we discuss the set of adversaries sampled during for each weaponization knowledge domain. However, we included the “Retain-set SFT followed by Weaponization-domain SFT” (R→F in Table <ref>) adversary to improve robustness to specific adversaries that broke the defense in earlier iterations of the method that did not sample this adversary. We find in Table <ref> that including the R→F adversary significantly improves robustness to more test-time attacks, indicated by an improvement in Post-Attack Forget error by over 20.0 percentage points. Importantly, we find that including the 64-step R→F adversary at train-time improves test-time robustness at for the 1000-step R→F adversary by nearly 40.0 percentage points, indicating that target patching of vulnerabilities is possible within .










 §.§ Ablations






  
Retain loss ℓ_2 representation engineering term. We observe in Table <ref> that retain-set MMLU performance is worse when excluding representation engineering <cit.> retain loss from . When including the term, retain MMLU improves by at least 5.0 percentage points. 



  
Including the initial safeguard. In Table <ref>, we examine the impact of incorporating the Random Mapping safeguard step prior to the adversarial training phase during . The Random Mapping safeguard in isolation achieves a near-random chance Pre-Attack Forget accuracy of 27.3. However, it is susceptible to fine-tuning attacks similar to other baselines in Table <ref>, indicated by a higher Post-Attack Forget accuracy of 59.4. When including the tamper-resistance adversarial training phase, we observe significantly increased tamper-resistance as the Post-Attack Forget accuracy decreases by nearly 28 percentage points. 

We also examine the impact of excluding the initial safeguarding step (“Excl. Initial Safeguard” in Table <ref>), finding that Pre-Attack and Post-Attack Forget accuracies are both higher 
without the initial safeguard. While the tamper-resistance adversarial training phase achieves moderate tamper-resistance without the initial safeguard step, we find that including the Random Mapping phase improves downstream tamper-resistance by 15.6 percentage points.





  
Varying the tamper-resistance loss scale λ_TR. We compare the downstream robustness of when varying the tamper-resistance loss weighting λ_TR between 1.0 and 4.0 in Table <ref>. We observe that when setting λ_TR = 1.0, maintains high retain MMLU accuracy at 62.5 percentage points, with moderate tamper-resistance indicated by a Post-Attack Forget accuracy of 40.8. Further increasing λ_TR to 4.0 in our final model results in a significantly improved Post-Attack Forget Accuracy of 31.3, with a partial decrease in Retain MMLU to 54.9. When varying λ_TR, we keep λ_retain constant; thus, our results indicate a clear way to increase downstream tamper-resistance by increasing the weighting of the tamper-resistance gradient during , reflecting a balance between tamper-resistance and capabilities similar to the robustness-performance tradeoff for adverarial robustness in vision models.





§ EXPERIMENT DETAILS





 §.§ Weaponization Domain Proxy Dataset Details




  
Biosecurity. We use a synthetically labeled partition of the Pile <cit.> that filters for relevance to biology and the Camel AI Biology dataset <cit.>. We generate synthetic labels for Pile token sequences using openchat-3.5 <cit.>, categorizing them as "Cellular Biology" or not. This process yields 49,984 samples: 7,558 for the forget-set (Pile-bio Forget) and 42,426 for the retain set (Pile-bio Retain). Concurrently, we pack entries from the Camel AI Biology dataset to the truncation-enabled 256 tokenization limit, resulting in 54,258 samples of about 188 words each (Camel-bio Forget). We apply the same procedure to our held-out hazardous biology dataset (identical to the WMDP biosecurity Forget-set), producing 598,933 samples of similar length (OOD Forget).



  
Chemical Security. We use a private forget dataset containing text sequences about hazardous chemical security content (Chem Forget).



  
Cybersecurity. We scrape CTF writeups on CTFtime <cit.> that are numbered between 1 and 39181, collecting cybersecurity writeups written as recently as 2024. We filter to keep writeups that contain more than 150 characters. As a result of filtering and HTTP errors while scraping, our resulting forget dataset contains slightly over 18k samples (Cyber Forget).




 §.§ Train-time Settings and Adversaries








  
Weaponization knowledge restriction. For each weaponization knowledge restriction domain, we have a corresponding retain dataset 𝒟_retain, comprised of a mix of data from the Pile-bio Retain set and Magpie-Align instruction-tuning dataset <cit.>, which we refer to as the Retain-set.  Referencing the datasets described in Appendix <ref>, we specify the adversary tamper-resistance datasets that constitute the data used by attacks sampled from 𝒜_train during training, as well as the attack setup for each domain as follows:


    
  * Biosecurity: We simulate six adversaries from the following three setups: Pile-Bio Forget-set SFT, Camel-Bio Forget-set SFT, and Retain-set SFT followed by Pile-Bio Forget-set SFT (R→F), where the switching point between Retain-set and Pile-Bio Forget-set SFT within the 64-step length trajectory is sampled from a beta distribution β(6.0,3.0). LRs are sampled from {2 × 10^-5, 4 × 10^-5}.
    
  * Chemical Security: We simulate nine adversaries from the following three setups: Chem Forget-set SFT, Retain-set SFT, and Retain-set followed by Chem Forget-set SFT, using the same switching-point sampling scheme as in the Biosecurity setting. LRs are sampled from {2 × 10^-6, 2 × 10^-5, 4 × 10^-5}.
    
  * Cybersecurity: We simulate four adversaries from the following two setups: Cyber Forget-set SFT, and Retain-set SFT followed by Cyber Forget-set SFT, using the same switching point sampling scheme as in the Chemical Security and Biosecurity settings. LRs are sampled from {2 × 10^-5, 4 × 10^-5}.


For each weaponization knowledge domain, we create 80-20 splits for adversary and held-out data of the corresponding forget sets, respectively. For Biosecurity, which uses multiple forget datasets, this involves creating 80-20 splits for each dataset, then combining the corresponding splits. The adversary data splits are used for sampled attacks from 𝒜_train, whereas the held-out split is used for computing tamper-resistance losses. The held-out splits for each domain correspond to 𝒟_TR in Section <ref>. We use minibatches from a held-out dataset for computing tamper-resistance losses rather than cycling through a single dataset, following the recommendation of <cit.>, in which first-order meta-learning without properly held-out minibatches caused a performance degradation.

All train-time adversary setups are tabulated in Table <ref>, where F-Pile, F-Chem, F-Cyber denote the respective datasets described in Appendix <ref>, and Retain denotes the mixed Pile-bio and Magpie-Align Retain-set described in Appendix <ref>. We use R→F to label the adversaries that perform Retain-set SFT followed by Forget-set SFT. The final column is an abbreviation for Finetuning Paradigm and indicates whether the SFT setup used full parameter finetuning or parameter-efficient fine-tuning (PEFT) via LoRA adapters <cit.>.





@lccccccc@
Train-time adversary setups for weaponization knowledge restriction of Biosecurity, Chemical Security, and Cybersecurity.
 


Adversary     Dataset     Opt. Steps (K)     Optimizer     LR     LR Schedule     Batch Size     FT Paradigm 



8c
  – continued from previous page 

Adversary     Dataset     Opt. Steps     Opt.     LR     LR Schedule     Batch Size     FT Paradigm 



8rContinued on next page 
 




8cBiosecurity Weaponization Restriction 

Adv 1     F-Pile     64     AdamW     2×10^-5     No Warmup     64     Full Parameter 

Adv 2     F-Pile     64     AdamW     4×10^-5     No Warmup     64     Full Parameter 

Adv 3     F-Camel     64     AdamW     2×10^-5     No Warmup     64     Full Parameter 

Adv 4     F-Camel     64     AdamW     4×10^-5     No Warmup     64     Full Parameter 

Adv 5     R→F     64     AdamW     2×10^-5     No Warmup     64     Full Parameter 

Adv 6     R→F     64     AdamW     4×10^-5     No Warmup     64     Full Parameter 

8cChemical Security Weaponization Restriction 

Adv 1     F-Chem     64     AdamW     2×10^-6     No Warmup     64     Full Parameter 

Adv 2     F-Chem     64     AdamW     2×10^-5     No Warmup     64     Full Parameter 

Adv 3     F-Chem     64     AdamW     4×10^-5     No Warmup     64     Full Parameter 

Adv 4     Retain     64     AdamW     2×10^-6     No Warmup     64     Full Parameter 

Adv 5     Retain     64     AdamW     2×10^-5     No Warmup     64     Full Parameter 

Adv 6     Retain     64     AdamW     4×10^-5     No Warmup     64     Full Parameter 

Adv 7     R→F     64     AdamW     2×10^-6     No Warmup     64     Full Parameter 

Adv 8     R→F     64     AdamW     2×10^-5     No Warmup     64     Full Parameter 

Adv 9     R→F     64     AdamW     4×10^-5     No Warmup     64     Full Parameter 

8cCybersecurity Weaponization Restriction 

Adv 1     F-Cyber     64     AdamW     2×10^-5     No Warmup     64     Full Parameter 

Adv 2     F-Cyber     64     AdamW     4×10^-5     No Warmup     64     Full Parameter 

Adv 3     R→F     64     AdamW     2×10^-5     No Warmup     64     Full Parameter 

Adv 4     R→F     64     AdamW     4×10^-5     No Warmup     64     Full Parameter 









  
Harmful request refusal. For harmful request refusal, we choose the retain dataset 𝒟_retain to be the Magpie-Align instruction-tuning dataset <cit.>. We sample train-time adversaries that perform K=64 steps of SFT on rejected completions of the Anthropic-HH-RLHF dataset <cit.> and vary the learning rate within {2 × 10^-6, 2 × 10^-5, 4 × 10^-5}. We depict this list of adversaries in Table <ref>. 

While the rejected completions from Anthropic-HH-RLHF constitute the data used for sampled attacks from 𝒜_train for harmful request refusal, we compute the tamper-resistance loss ℒ_TR as follows. Since ℒ_TR is the DPO loss <cit.> in this setting, we use the base model weights θ as the reference model and sample harmful and benign completions from a modified test split of Anthropic-HH-RLHF <cit.>, where rejected completions are replaced with refusals <cit.>. To avoid keeping the base model weights in memory and speed up training, we precompute the reference model DPO log-probabilities for the full Anthropic-HH-RLHF dataset before training. To summarize, we have that the sampled train-time adversaries perform SFT on the rejected completions from Anthropic-HH-RLHF, and tamper-resistance DPO loss is computed on the corresponding modified refusal completions; the modified refusal completions in this setting correspond to 𝒟_TR in Section <ref>.

In practice, we perform an additional 100 steps of supervised fine-tuning on the Magpie-Align dataset to improve the benign capabilities performance of the refusal model in Table <ref>.





@lccccccc@
Train-time adversary red-teaming setups harmful request refusal. “A-HH-Rejected” in the Dataset column corresponds to the adversary dataset comprised of rejected completions from the Anthropic-HH-RLHF dataset.
 


Adversary     Dataset     Opt. Steps (K)     Optimizer     LR     LR Schedule     Batch Size     FT Paradigm 



8c
  – continued from previous page 

Adversary     Dataset     Opt. Steps (K)     Optimizer     LR     LR Schedule     Batch Size     FT Paradigm 



8rContinued on next page 
 



8cHarmful Request Refusal 

Adv 1     A-HH-Rejected     64     AdamW     2×10^-5     No Warmup      64     Full Parameter 

Adv 2     A-HH-Rejected     64     AdamW     2×10^-6     No Warmup     64     Full Parameter 

Adv 3     A-HH-Rejected     64     AdamW     4×10^-5     No Warmup     64     Full Parameter 













§ RED TEAMING DETAILS






 §.§ Weaponization Knowledge Restriction


In Table <ref>, we list all test-time adversary setups for recovering Biosecurity, Chemical Security, and Cybersecurity Weaponization knowledge. 

For Biosecurity, we examine post-attack forget accuracy after fine-tuning for 1,000 steps on three data distributions: the Pile-bio Forget set and the Retain-set used during Random Mapping and , and an OOD-Forget set mentioned in Appendix <ref>, which is unseen during tamper-resistance training. We conduct an additional attack using 5,000 optimization steps on this unseen forget distribution to compose our strongest adversary. We leverage a consistent setup for Chemical Security and Cybersecurity, substituting the Pile-bio Forget set for the respective Chem and Cyber Forget sets. We exclude the unseen forget distribution adversaries for a total of 24 attacks, as we were unable to obtain an equally high quality held-out dataset for Chemical Security and Cybersecurity.

For all subjects, we compare the effect of different optimizers (AdamW, Adadelta, and Stochastic Gradient Descent with Nesterov Momentum, Schedule Free AdamW) <cit.>, learning rates (2 × 10^-6, 2 × 10^-5, 4 × 10^-5), and learning rate scheduling techniques (SGDR and 30 steps of linear warmup) <cit.>. 

Similar to Table <ref>, we use F-Pile, F-Chem, F-Cyber, OOD-F, and Retain in the Dataset Column to denote the respective datasets described in Appendices  <ref> and <ref>. Separate from the R→F train-time adversary, R→F at test-time represents the combination of 0.4 × Opt. Steps of Forget-set SFT followed by 0.6 × Opt. Steps of Retain-set SFT. In practice, we found that this fraction composed a potent attack which we sought robustness against as described in Appendix <ref>.

In the Optimizer column, Schedule Free is an abbreviation of AdamW Schedule Free and SGD Nesterov is an abbreviation of SGD with Nesterov Momentum. In cases where the adversary used parameter-efficient fine-tuning (PEFT) via LoRA adapters <cit.>, we used a LoRA config with an attention dimension, or rank, of 16, a LoRA alpha value of 32, a LoRA dropout of 0.05, on target linear modules:      

{, , , ,, , }.













@lccccccc@
Test-time adversary red-teaming setups for weaponization knowledge restriction of Biosecurity, Chemical Security, and Cybersecurity.
 


Adversary     Dataset     Opt. Steps (K)     Optimizer     LR     LR Schedule     Batch Size     FT Paradigm 



8c
  – continued from previous page 

Adversary     Dataset     Opt. Steps (K)     Optimizer     LR     LR Schedule     Batch Size     FT Paradigm 



8rContinued on next page 
 



8cBiosecurity Weaponization Restriction 

Adv 1     OOD-F     5000     AdamW     2×10^-5     No Warmup     64     Full Parameter 

Adv 2     OOD-F     5000     AdamW     4×10^-5     No Warmup     64     Full Parameter 

Adv 3     F-Pile     1000     AdamW     2×10^-5     No Warmup     64     Full Parameter 

Adv 4     F-Pile     1000     AdamW     4×10^-5     No Warmup     64     Full Parameter 

Adv 5     Retain     1000     AdamW     2×10^-5     No Warmup     64     Full Parameter 

Adv 6     Retain     1000     AdamW     4×10^-5     No Warmup     64     Full Parameter 

Adv 7     OOD-F     1000     AdamW     2×10^-5     No Warmup     64     Full Parameter 

Adv 8     OOD-F     1000     AdamW     4×10^-5     No Warmup     64     Full Parameter 

Adv 9     R→F     1000     AdamW     2×10^-5     No Warmup     64     Full Parameter 

Adv 10     R→F     1000     AdamW     4×10^-5     No Warmup     64     Full Parameter 

Adv 11     F-Pile     1000     Adadelta     2×10^-5     No Warmup     64     Full Parameter 

Adv 12     F-Pile     1000     Adadelta     4×10^-5     No Warmup     64     Full Parameter 

Adv 13     F-Pile     1000     Schedule Free     2×10^-5     No Warmup     64     Full Parameter 

Adv 14     F-Pile     1000     Schedule Free     4×10^-5     No Warmup     64     Full Parameter 

Adv 15     F-Pile     1000     SGD Nesterov     2×10^-5     No Warmup     64     Full Parameter 

Adv 16     F-Pile     1000     SGD Nesterov     4×10^-5     No Warmup     64     Full Parameter 

Adv 17     F-Pile     1000     AdamW     2×10^-6     No Warmup     64     Full Parameter 

Adv 18     F-Pile     1000     AdamW     2×10^-6     30 Steps Warmup     64     Full Parameter 

Adv 19     F-Pile     1000     AdamW     2×10^-5     30 Steps Warmup     64     Full Parameter 

Adv 20     F-Pile     1000     AdamW     4×10^-5     30 Steps Warmup     64     Full Parameter 

Adv 21     F-Pile     1000     AdamW     2×10^-5     SGDR     64     Full Parameter 

Adv 22     F-Pile     1000     AdamW     4×10^-5     SGDR     64     Full Parameter 

Adv 23     F-Pile     1000     AdamW     2×10^-5     No Warmup     32     Full Parameter 

Adv 24     F-Pile     1000     AdamW     4×10^-5     No Warmup     32     Full Parameter 

Adv 25     F-Pile     1000     AdamW     2×10^-5     No Warmup     128     Full Parameter 

Adv 26     F-Pile     1000     AdamW     4×10^-5     No Warmup     128     Full Parameter 

Adv 27     F-Pile     1000     AdamW     2×10^-5     No Warmup     64     PEFT 

Adv 28     F-Pile     1000     AdamW     4×10^-5     No Warmup     64     PEFT 

8cChemical Security Weaponization 

Adv 1     F-Chem     1000     AdamW     2×10^-5     No Warmup     64     Full Parameter 

Adv 2     F-Chem     1000     AdamW     4×10^-5     No Warmup     64     Full Parameter 

Adv 3     Retain     1000     AdamW     2×10^-5     No Warmup     64     Full Parameter 

Adv 4     Retain     1000     AdamW     4×10^-5     No Warmup     64     Full Parameter 

Adv 5     R→F     1000     AdamW     2×10^-5     No Warmup     64     Full Parameter 

Adv 6     R→F     1000     AdamW     4×10^-5     No Warmup     64     Full Parameter 

Adv 7     F-Chem     1000     Adadelta     2×10^-5     No Warmup     64     Full Parameter 

Adv 8     F-Chem     1000     Adadelta     4×10^-5     No Warmup     64     Full Parameter 

Adv 9     F-Chem     1000     ScheduleFree     2×10^-5     No Warmup     64     Full Parameter 

Adv 10     F-Chem     1000     ScheduleFree     4×10^-5     No Warmup     64     Full Parameter 

Adv 11     F-Chem     1000     SGD Nesterov     2×10^-5     No Warmup     64     Full Parameter 

Adv 12     F-Chem     1000     SGD Nesterov     4×10^-5     No Warmup     64     Full Parameter 

Adv 13     F-Chem     1000     AdamW     2×10^-6     No Warmup     64     Full Parameter 

Adv 14     F-Chem     1000     AdamW     2×10^-6     30 Steps Warmup     64     Full Parameter 

Adv 15     F-Chem     1000     AdamW     2×10^-5     30 Steps Warmup     64     Full Parameter 

Adv 16     F-Chem     1000     AdamW     4×10^-5     30 Steps Warmup     64     Full Parameter 

Adv 17     F-Chem     1000     AdamW     2×10^-5     SGDR     64     Full Parameter 

Adv 18     F-Chem     1000     AdamW     4×10^-5     SGDR     64     Full Parameter 

Adv 19     F-Chem     1000     AdamW     2×10^-5     No Warmup     32     Full Parameter 

Adv 20     F-Chem     1000     AdamW     4×10^-5     No Warmup     32     Full Parameter 

Adv 21     F-Chem     1000     AdamW     2×10^-5     No Warmup     128     Full Parameter 

Adv 22     F-Chem     1000     AdamW     4×10^-5     No Warmup     128     Full Parameter 

Adv 23     F-Chem     1000     AdamW     2×10^-5     No Warmup     64     PEFT 

Adv 24     F-Chem     1000     AdamW     4×10^-5     No Warmup     64     PEFT 

8cCybersecurity Weaponization Restriction 

Adv 1     F-Cyber     1000     AdamW     2×10^-5     No Warmup     64     Full Parameter 

Adv 2     F-Cyber     1000     AdamW     4×10^-5     No Warmup     64     Full Parameter 

Adv 3     Retain     1000     AdamW     2×10^-5     No Warmup     64     Full Parameter 

Adv 4     Retain     1000     AdamW     4×10^-5     No Warmup     64     Full Parameter 

Adv 5     R→F     1000     AdamW     2×10^-5     No Warmup     64     Full Parameter 

Adv 6     R→F     1000     AdamW     4×10^-5     No Warmup     64     Full Parameter 

Adv 7     F-Cyber     1000     Adadelta     2×10^-5     No Warmup     64     Full Parameter 

Adv 8     F-Cyber     1000     Adadelta     4×10^-5     No Warmup     64     Full Parameter 

Adv 9     F-Cyber     1000     ScheduleFree     2×10^-5     No Warmup     64     Full Parameter 

Adv 10     F-Cyber     1000     ScheduleFree     4×10^-5     No Warmup     64     Full Parameter 

Adv 11     F-Cyber     1000     SGD Nesterov     2×10^-5     No Warmup     64     Full Parameter 

Adv 12     F-Cyber     1000     SGD Nesterov     4×10^-5     No Warmup     64     Full Parameter 

Adv 13     F-Cyber     1000     AdamW     2×10^-6     No Warmup     64     Full Parameter 

Adv 14     F-Cyber     1000     AdamW     2×10^-6     30 Steps Warmup     64     Full Parameter 

Adv 15     F-Cyber     1000     AdamW     2×10^-5     30 Steps Warmup     64     Full Parameter 

Adv 16     F-Cyber     1000     AdamW     4×10^-5     30 Steps Warmup     64     Full Parameter 

Adv 17     F-Cyber     1000     AdamW     2×10^-5     SGDR     64     Full Parameter 

Adv 18     F-Cyber     1000     AdamW     4×10^-5     SGDR     64     Full Parameter 

Adv 19     F-Cyber     1000     AdamW     2×10^-5     No Warmup     32     Full Parameter 

Adv 20     F-Cyber     1000     AdamW     4×10^-5     No Warmup     32     Full Parameter 

Adv 21     F-Cyber     1000     AdamW     2×10^-5     No Warmup     128     Full Parameter 

Adv 22     F-Cyber     1000     AdamW     4×10^-5     No Warmup     128     Full Parameter 

Adv 23     F-Cyber     1000     AdamW     2×10^-5     No Warmup     64     PEFT 

Adv 24     F-Cyber     1000     AdamW     4×10^-5     No Warmup     64     PEFT 









 §.§ Harmful Request Refusal




For the harmful request refusal setting, we conduct 5 test-time adversary attacks that perform SFT for 10 epochs on a held-out toxicity dataset called Toxic-DPO v0.2, on each of the settings in Table <ref>. The dataset contains 541 user-assistant chat interactions where the assistant complies with harmful instructions. 




@lccccccc@
Test-time adversary red-teaming setups for harmful request refusal. “ToxicDPO” in the Dataset column refers to the ToxicDPOv0.2 dataset containing harmful chat completions.
 


Adversary     Dataset     Epochs     Optimizer     LR     LR Schedule     Batch Size     FT Paradigm 



8c
  – continued from previous page 

Adversary     Dataset     Epochs     Optimizer     LR     LR Schedule     Batch Size     FT Paradigm 



8rContinued on next page 
 



8cHarmful Request Refusal 

Adv 1     ToxicDPO     10     AdamW     1×10^-5     10 Steps Warmup      32     Full Parameter 

Adv 2     ToxicDPO     10     AdamW     1×10^-5     No Warmup     32     Full Parameter 

Adv 3     ToxicDPO     10     AdamW     1×10^-5     No Warmup     16     Full Parameter 

Adv 4     ToxicDPO     10     AdamW     2×10^-5     No Warmup     32     Full Parameter 

Adv 5     ToxicDPO     10     AdamW     4×10^-5     No Warmup     32     Full Parameter 

















 §.§ Additional Weaponization Knowledge Restriction Red-Teaming Results

In Figure <ref>, we characterize the post-tampering resistance of TAR and the baseline safeguards with respect to the Chemical Security and Cybersecurity domains. Similar to the Biosecurity domain, baseline safeguards break under the majority of fine-tuning attacks, while resists attacks from 23 out of 24 Chemical Security SFT adversaries and 21 out of 24 Cybersecurity SFT adversaries. 












§ BASELINE DETAILS





 §.§ Weaponization Knowledge Restriction



  
Max Entropy.

Let K be the set of all token-wise output probability distributions returned by a model θ, where k ∈ K corresponds to every position in the sequence. We maximize the average entropy of these discrete distributions in K as follows:





    ℒ_Max Entropy = ∑_k ∈𝒦 p_k log(p_k)


This is equivalent to minimizing the average Kullback-Leibler (KL) divergence <cit.> between each k and the discrete uniform distribution u(x) over the vocabulary 𝒱. For Llama-3-8B-Instruct, |𝒱| = 128256. Thus, this objective is upper-bounded by: 
    h(x) = -log(u(x)) = log(|𝒱|) ≈ 11.76

where h(x) measures the Shannon information or self-information and log(x) has base e. We apply this objective for all elements in the Forget-set and perform standard cross-entropy on the Retain-set.



  
Min Posterior.
The goal of the Min Posterior objective is to assign lower probabilities to true forget-set labels, essentially minimizing -log(1 - P(label)). Let p_i be the probability assigned to the true label for token i and 𝒱 be the model's vocabulary distribution. We define the Min Posterior objective function as follows:

    ℒ_Min Posterior = -1/|𝒱|∑_i ∈𝒱log(1 - p_i + ϵ) ·𝕀[log(p_i) ≥τ]
 
where τ is the threshold for masking out target label logits (which we set to the negative maximum entropy of the vocabulary distribution, -log|𝒱|) and 𝕀[·] is the corresponding indicator function (1 if the condition is true, 0 otherwise). We include an optional ϵ = 1×10^-12 to help with numerical stability. Similar to the Max Entropy objective, we apply this objective for all elements in the Forget-set and perform standard cross-entropy on the Retain-set.



  
RMU. We adapt RMU's implementation from <cit.> with a learning rate of 5×10^-5 and 250 unlearning steps. We use the released WMDP's unlearning datasets for Biosecurity (Bio) and Cybersecurity (Cyber) unlearning, and our private hazardous chemistry dataset for Chemical Security (Chem) unlearning. We use unlearning coefficients of 20, 30, and 50 for Bio, Cyber, and Chem respectively. We use a retain coefficient of 700 on Wikitext <cit.>.



  
LLMU.

We use a modified version of LLMU from <cit.>. Instead of computing the KL divergence to regularize retain-set logits towards the base frozen model, we employ a standard cross-entropy loss. This modification allows for memory-efficient execution on our hardware while maintaining comparable performance.




  
Hyperparameter tuning. Besides RMU, all baseline hyperparameters were chosen after a grid search across learning rates {3×10^-6, 5×10^-6, 8×10^-6, 1×10^-5}, optimization step count {600, 1000}, and warmup steps {0, 100}. We found that 600 optimization steps using the AdamW Schedule Free optimizer, at a learning rate of 1×10^-5, with 100 steps of linear warmup, and an effective batch size of 64 produced the best performance. For the Max Entropy, Min Posterior, and LLMU baselines, we train on the three corresponding forget datasets discussed in <ref>. For these baselines, we modify our Biosecurity forget corpus to be a mixture of the Pile-bio and Camel-bio Forget corpora. We use the Pile-bio Retain-set as a global Retain-set for baseline training.



 §.§ Harmful Request Refusal




  
Representation Rerouting. We use the Llama-3-8B-Instruct RR model from <cit.>, which uses a cosine distance loss to push representations for harmful inputs to become orthogonal to those of the base Llama-3-8B-Instruct model.



  
R2D2. We use the R2D2 model run on Zephyr-7B directly from <cit.>, which performs adversarial training against GCG attacks to increase jailbreak robustness.



  
RepNoise. We the RepNoise model run on Llama-2-7B directly from <cit.>, which uses a distributional loss to push representations for harmful inputs toward Gaussian noise.








 §.§ Additional Baseline Comparisons






  
MLAC-AR. Meta-Learned Adversarial Censoring (MLAC) <cit.> was originally proposed to prevent BERT-style models from learning binary classification for gender bias data. Since the approach is not immediately applicable to LLMs, we extend MLAC in a variant we call autoregressive MLAC (MLAC-AR). Since MLAC in its original formulation calls for “task-blocking” via negating the adversary's loss during the inner loop of meta-learning, we implement this by negating the cross-entropy loss of an LLM fine-tuning adversary. However, we found that this approach diverges in performance across a variety of hyperparameters, and opted to further improve performance of the MLAC-AR baseline by clamping the maximum cross-entropy loss at the value of the maximum entropy of the output vocabulary distribution, log(). We show results in Table <ref>, finding that MLAC-AR does not maintain sufficient benign capabilities performance nor uniform tamper-resistance across weaponization domains.



  
SOPHON-AR. In concurrent work, SOPHON <cit.> was introduced to prevent small diffusion models and image classifiers from learning specific data distributions. Similarly to MLAC-AR, we extend SOPHON to LLMs via SOPHON-AR, using the alternating retain loss and fine-tuning suppression loss formulation that the authors propose. Furthermore, we adapt the inverse cross-entropy loss from <cit.> , which aims to boost convergence of the fine-tuning suppression process. We find in practice that despite heavy tuning, SOPHON-AR does not converge well enough to yield a usable Pre-Attack model in Table <ref>.




 Paper Title: ProSpec RL: Plan Ahead, then Execute.


	
	
	
Enhancement of the molecular electron chirality by electronic excitation

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

§ ABSTRACT 
		Imagining potential outcomes of actions before execution helps agents make more informed decisions, a prospective thinking ability fundamental to human cognition. However, mainstream model-free Reinforcement Learning (RL) methods lack the ability to proactively envision future scenarios, plan, and guide strategies. These methods typically rely on trial and error to adjust policy functions, aiming to maximize cumulative rewards or long-term value, even if such high-reward decisions place the environment in extremely dangerous states. To address this, we propose the Prospective (ProSpec) RL method, which makes higher-value, lower-risk optimal decisions by imagining future n-stream trajectories. Specifically, ProSpec employs a dynamic model to predict future states (termed “imagined states”) based on the current state and a series of sampled actions. Furthermore, we integrate the concept of Model Predictive Control and introduce a cycle consistency constraint that allows the agent to evaluate and select the optimal actions from these trajectories. Moreover, ProSpec employs cycle consistency to mitigate two fundamental issues in RL: augmenting state reversibility to avoid irreversible events (low risk) and augmenting actions to generate numerous virtual trajectories, thereby improving data efficiency. We validated the effectiveness of our method on the DMControl benchmarks, where our approach achieved significant performance improvements. Code will be open-sourced upon acceptance.
	
	
	
	

§ INTRODUCTION

	Reinforcement Learning (RL) has proven indispensable for solving continuous decision problems <cit.>. Current RL research falls into two main categories: model-based and model-free methods. Model-based methods involve constructing a world model to simulate an interactive environment, which helps the agent predict future trajectories for optimal decision-making. However, these methods can be expensive and prone to failure due to discrepancies between the world model and reality <cit.>. On the other hand, model-free methods refine the policy function through trial-and-error learning, often outperforming model-based techniques in continuous control tasks <cit.>. However, these methods require frequent interaction with the environment, which is risky or costly in real-world applications <cit.>, and suffer from low data efficiency when learning from complex, high-dimensional observations (such as image pixels) based on limited experience.
	
	Training model-free methods using low data efficiency and sparse signals to fit high-performance feature encoders is challenging and prone to suboptimal convergence <cit.>. In contrast, humans can learn to play simple arcade games in minutes <cit.>, while even state-of-the-art RL algorithms require billions of frames to reach human-level performance <cit.>. Moreover, collecting such large amounts of interaction data is impractical, especially in domains such as autonomous driving and robot control. Currently, researchers are investigating data efficiency challenges from various angles. Including leveraging model-based world models to aid in policy learning <cit.>, employing data augmentation techniques to augment diversity in image appearances <cit.>, and introducing auxiliary tasks to assist in state representation learning <cit.>, among other approaches. While many of these methods involve creating additional training processes to accelerate agent training, the challenge of limited experience continues to impede the model's exceptional performance.
	
	Learning planning is often deemed more critical than learning policies <cit.>. In cognitive neuroscience and psychology, planning is considered a facet of prospective thinking <cit.>, which aids in predicting and adapting to future changes, envisioning potential problems or dangers, devising strategies to enhance problem-solving abilities, and setting future goals to drive action <cit.>. Ideally, prospective thinking is highly relevant to continuous decision-making tasks, where current actions can have irreversible consequences on future scenarios. However, current model-free methods lack planning capabilities, thus relying solely on trial and error may lead to unacceptable risks.
	
	In this work, we introduce ProSpec, a model-free RL method that harnesses prospective thinking to envision potential future scenarios and make optimal decisions with higher value and lower risk through planning. More importantly, ProSpec could address data efficiency concerns by generating numerous virtual trajectories through imagination. Specifically, we model prospective thinking using a Flow-based Reversible Dynamics Model (FDM) and a Model Predictive Control (MPC)-based block. The FDM predicts prospective states in latent space in the forward process. In the backward process, we enforce cycle consistency constraints on the reversibility of prospective states to prevent the failure of prospective trajectories in irreparable states. Furthermore, we augment actions in a self-supervised learning fashion to generate numerous virtual trajectories and improve the efficiency of RL data. The MPC block employs the concept of MPC to facilitate exploratory interaction between the agent and the FDM, with the objective of exploring future scenarios from diverse perspectives. We summarize the contributions of this work as follows:
	
		
		
		
		
  * We introduce a model-free method that integrates prospective thinking by incorporating a Flow-based Reversible Dynamics Model (FDM) for predicting prospective states, combined with cycle consistency to maintain state traceability. This enables the agent to explore future scenarios from various perspectives and identify optimal decisions with higher value and lower risk.
		
  * Within our reversible framework, we introduce action augmentation in a self-supervised manner, generating multiple state-action trajectories to improve data efficiency.
		
  * We validate the effectiveness of ProSpec through experiments on DeepMind benchmarks, showing significant performance improvements.
	
	
	

§ RELATED WORK

	

 §.§ Prospective thinking

	The capacity to anticipate the future represents as a crucial facet of human cognition, recently garnering substantial attention from researchers <cit.>. In psychology, prospective thinking is defined as the ability to mentally simulate the future, involving the extraction of information from episodic and semantic memory (such as detailed information about past states) as well as more abstract, schematic, and conceptual knowledge (such as envisioning general goals or events) <cit.>. Broadly, prospective thinking involves two main aspects: 1) predicting future scenarios and 2) guiding decisions for those scenarios. In deep learning, methods such as Model Predictive Control (MPC) and model-based Reinforcement Learning (RL) appear to encapsulate certain aspects of prospective thinking.
	
	Model-based RL methods aim to construct an internal model, often referred to as a "World Model," to simulate the environment. This model helps predict future states and rewards, guiding the agent's decision-making and planning. Examples include PILCO <cit.>, Dyna-Q <cit.>, and World models <cit.>. On the other hand, MPC-based methods utilize the world model to optimize trajectories within a shorter time horizon. Examples of such methods include MuZero, Eff.Zero <cit.>, and TD-MPC <cit.>. However, these approaches heavily rely on accurate environmental modeling and are limited by the complexity of state spaces.
	Consequently, model-based approaches often struggle to outperform simpler model-free methods in continuous control tasks <cit.>. Model-free methods, in contrast, learn policies through trial and error, directly interacting with the environment without the need for a World Model. This method is easy to implement and adapts well to complex environments. However, due to the frequent interaction with the environment and the lack of prediction of future states or consequences, it can lead to risky exploration of the environment <cit.>.
	
	Building upon this, we merge the planning capabilities of model-based methods to enhance model-free methods. To achieve this, we introduce a dynamics model into model-free methods, enabling the prediction of future scenarios. Additionally, we integrate the concept of MPC with value consistency to aid decision-making and guidance for future scenarios. Our objective is to leverage prospective thinking in selecting the optimal path while maintaining the trial-and-error learning approach of model-free methods. This integration aims to enhance the effectiveness of model-free methods in learning.
	

 §.§ Data-Efﬁcient Reinforcement Learning

	When confronted with a limited amount of interaction data, the performance of both model-based and model-free methods experiences a significant decline. Researchers have extensively investigated this issue from various angles <cit.>. For example, SiMPLe utilizes collected data to learn a world model and generate imagined trajectories. They achieved outstanding results in multiple games under the setting of 100k frames, albeit at the cost of several weeks of training time. Subsequent studies, such as DataEfficient Rainbow <cit.> and OTRainbow <cit.>, have demonstrated that merely increasing the number of steps in multi-step returns and conducting more frequent parameter updates can lead to improved performance under conditions of low data efficiency.
	
	Consequently, there is growing interest in using computer vision techniques to enhance representation learning in RL <cit.>. For example, DrQ <cit.> and RAD <cit.> showed that moderate image augmentation can significantly improve the efficiency of RL data, outperforming previous methods using world models in model-based approaches. Yarats et al. <cit.> introduced image reconstruction as an auxiliary loss function in RL to improve data efficiency. Inspired by contrastive learning in representation learning, CURL <cit.> incorporated contrastive loss into RL, allowing the algorithm to distinguish representations of different states through contrastive loss and promote the generation of similar embeddings for similar states. In addition, some studies aim to incorporate dynamical models to help agents predict future events, thereby generating multiple virtual trajectories to improve feature representation learning. For example, SPR <cit.>, PBL <cit.>, and SLAC <cit.>, among others. Subsequently, EfficientZero <cit.> integrated SPR's prediction loss into MuZero, improving the quality of internal state representations through self-generated predictions within the model.
	
	Recently, researchers have incorporated cycle consistency into dynamics models <cit.> to ensure that generated virtual trajectories maintain consistency and coherence with real-world behavior. The basic concept of cycle consistency involves performing a forward transformation from A to B using two transformation models, followed by a reverse transformation from A' back to B, and imposing consistency constraints between the original input A and the reconstructed input A'. For example, PlayVirtual <cit.> introduced backward dynamics models based on SPR, and established a loop through constraints between forward and backward dynamics models to improve feature representation learning via cycle consistency. VCR <cit.> introduced the concept of value consistency, where their dynamics model is constrained by Q-values computed with the real environment, enabling the model to learn representations directly related to decision-making, thereby improving model performance.
	
	These successful cases demonstrate the effectiveness of cycle consistency in improving data efficiency. However, this approach often requires additional training to construct a backward dynamics model (as seen in PlayVirtual <cit.>), or directly imposing consistency constraints on learning from real environment policies. The latter approach may be susceptible to deviations in value functions caused by high-bias environments, as observed in VCR <cit.>. Therefore, in this study, we propose a novel approach that utilizes a simple flow-based dynamics model for both forward and backward predictions. In addition, cycle consistency is employed to improve data and state, ensuring the generation of numerous imagined trajectories under low-risk constraints to mitigate data efficiency issues.
	
	

§ METHODS

	In this section, we will systematically introduce the proposed ProSpec method. First, we will review some preliminary work. Then, we will present the overall framework of ProSpec.
	

 §.§ Preliminaries: Reinforcement Learning

	Reinforcement learning (RL) typically uses Markov decision processes (MDPs) to solve continuous decision problems, denoted as <𝒮,𝒜,𝒯,ℛ,γ>. Here 𝒮 represents a finite set of states; 𝒜 is the action space; 𝒯(s_t,a_t,s_t+1)=P(s_t+1|s_t,a_t) is the dynamic function, which defines the probability of transition from state s_t to s_t+1 after taking action a_t; ℛ(s_t,a_t) is the reward function, and γ∈(0,1] is the discount factor. The goal of RL is to enable the agent to learn how to maximize the expected discounted cumulative payoff G_t=∑^∞_τ=0γ^τℛ(s_t+τ, a_t+τ) by choosing actions at each time t. The agent's decision process is guided by the policy π(a_t|s_t), which maps the current state to the action selection. The action-value function Q_π(s_t,a_t)=𝔼_π[G_t|s_t,a_t] evaluates the expected payoff of taking action a_t in state s_t and following policy π. This paper focuses on value-based RL methods, specifically Q-learning, which approximates the optimal policy π^* using the Bellman equation <cit.>. Therefore, Soft Actor-Critic (SAC) is employed as the policy algorithm, given that it is a widely utilized gradient-based algorithm for continuous control that incorporates policy entropy as an additional reward to encourage exploration. The overall training objective of SAC is as follows:
	
    𝒥_θ=	ℒ_critic+ℒ_actor+ℒ_α

	where, 	ℒ_critic, ℒ_actor and ℒ_α represent the critic loss, actor loss and temperature loss, respectively (see the supplementary material for details).
	

 §.§ Overall Framework

	
	In this paper, we introduce the ProSpec method, which aims to incorporate prospective thinking into model-free methods. ProSpec consists of two main components: first, the Flow-based Dynamics Model for predicting future scenarios, and second, the prospective unit that integrates Model Predictive Control and value consistency to guide decision-making. Figure<ref> shows the process of ProSpec RL. In the following sections, we provide detailed explanations of each component.	
	
	Flow-based Dynamics Model (FDM). The FDM predicts the future and retroactively infers previous latent states, crucial in ProSpec. Unlike PlayVirtual <cit.>, which employs separate models for these tasks, our approach employs a single model for both, avoiding training instability. We utilize Real-valued non-volume preserving (RealNVP) <cit.>, a flow-based neural network that performs bijective coupling transformations. RealNVP operates via bijective coupling layers, with each layer handling parts of the input a_t ∈ℝ^1:d and z_t ∈ℝ^d+1:D separately. It applies transformations using affine functions for scaling exp(sc_i(·)) and translation t_i(·), along with element-wise multiplication `⊙' and addition `+' operations. This enables the computation of the forward process as:
	
    [  ẑ_t+1^1:d = a_t ⊙exp(sc_2(z_t))+t_2(z_t); ẑ_t+1^d+1:D = z_t⊙exp(sc_1(a_t))+t_1(a_t) ]

	and the backward processes can be calculated as:
	
    [   z_t=(ẑ_t+1^d+1: D-t_1(ẑ_t+1^1: d)) ⊙exp(-sc_1(ẑ_t+1^1: d)); a_t=(ẑ_t+1^1: d -t_2(ẑ_t+1^d+1: D))⊙exp(-sc_2(ẑ_t+1^d+1: D)) ]

	Note that there exists a dimensional inconsistency between ẑ_t+1∈ℝ^D and z_t∈ℝ^d+1:D. Typically, a feed-forward neural network, represented as f(z_t; θ) with parameters θ, endeavors to fit training data and predict the subsequent latent state ẑ_t+1. The network's transformation can be expressed as ẑ_t+1 = W · z_t + B, where W ∈ℝ^n × n and B ∈ℝ^n × n denote the weights and biases, collectively denoted as θ = {W, B}. However, owing to the typical non-invertibility of W, conducting backward inference (i.e., directly computing z_t = W^-1(ẑ_t+1 - B)) poses challenges. To ensure reversibility, Orthogonal Weight Normalization (OWN) <cit.> is employed, supporting straightforward backward inference by enforcing orthogonality in the weights, unlike the structural design in invertible networks. Consequently, the backward process can be reformulated as:
	
    z_t=W^T(ẑ_t+1-B), W∈𝒪^n× n
	
	where the matrix family 𝒪^n × n = {W ∈ℝ^n × n| WW^T = W^-1W = I} (see <cit.> for more details). By using the combination of RealNVP and OWN, we can effectively achieve bidirectional transformation between z_t and ẑ_t+1.
	
	
	Forward States Prediction with FDM. In the MDP framework, a one-step transition (s_t, a_t, s_t+1) describes the process of transitioning from the current state s_t ∈𝒮 to the next state s_t+1 by taking action a_t ∈𝒜. Considering the crucial aspect of prospective thinking, i.e., the ability to predict future scenarios, is very important for RL tasks. Many studies focus on training dynamic models to predict future states to learn good feature representations <cit.>. In this study, the FDM h(·, ·) is used to simulate the dynamics function 𝒯 in the latent embedding space, allowing us to explore the future from multiple perspectives. As shown in Figure <ref>, the encoder f(·, ·) first converts the image-based observation/state s_0 into a latent representation z_0 = f(s_0). Then, based on the sampled action set Â(z_0) = {â_0, ⋯, â_k} from k perspectives, the FDM predicts the latent state set Ẑ_0:t for t future steps. Formally:
	
    [                                 ẑ_0=z_0=f(s_0);                             ẑ_t=h(â_t-1,ẑ_t-1); Ẑ_0:t = {h(Â_0(z_0), z_0),⋯, h(Â_t(ẑ_t),ẑ_t )} ]

	The optimization objective of the FDM h(·, ·) is to minimize the difference (error) between the predicted latent state sequence ẑ_t and the latent feature sequence z_t directly extracted from the original observation/state, where z_t = f_T(s_t). Following SPR <cit.>, we use “projection” and cosine similarity on DMControl to predict the error, aiming to achieve the following objective:
	
    ℒ_pred=
    		-∑_k=1^K(ẑ_t+k/ẑ_t+k_2)^⊤(z_t+k/z_t+k_2)

	Where z_t is referred to as the target embedding, generated by the stop-gradient version of the online encoder f(·) termed as the target encoder f_T(·).
	
	Backward Return to Previous States with FDM. The backward process of the FDM aims to derive the previous latent states z̆_t based on the given subsequent latent state z̆_t+1 and the causal action â_t. We enforce constraints on the trajectories generated by FDM, ensuring cycle consistency to prevent forward trajectories from encountering irreparable states. Given â_t and z̆_t+1, the reversibility of FDM allows us to easily trace back to the previous latent dynamics z̆_t. Formally, given the hidden state z̆_t+K and a sequence of actions â_t+K-1:t={â_t+K-1,⋯,â_t}, we iteratively compute the previous K state sequence z̆_t+k-1:t={z̆_t+k-1,⋯,z̆_t} using the inverse function h^-1(· ,· ):
	
    z̆_t+k-1={[  h^-1(â_t+k-1, ẑ_t+k),                 if k=K; h^-1(â_t+k-1, z̆_t+k),          if k=K-1,⋯,1. ].

	Using FDM, we can easily generate both a forward and a backward trajectory given the current state and a sequence of actions, forming a cycle forward-inverse trajectory. Ideally, the initial latent state z_t of this cycle trajectory should coincide with the final latent state z̆_t. Therefore, as shown in the right side of Figure <ref>, we introduce cycle consistency constraints between the initial state z_0 and the final state z̆_0 to prevent prospective trajectories from entering irreparable states. In addition, we improve actions in an unsupervised manner by generating or sampling virtual actions to create diverse virtual trajectories, which promotes data efficiency without relying solely on real trajectories for supervised learning. The inclusion of cycle consistency constraints serves a dual purpose, enhancing both the learning of feature representations and the reversibility of states to preclude irreversible events, while also elevating the training efficiency of the weight matrix W in OWN. 
	
	The optimization objective of cycle consistency is to minimize the spatial distance d(z_t, z̆_t) between two latent states z_t and z̆_t. Given the latent state representation z_t at time t, we randomly sample M actions from the action space 𝒜. Therefore, the loss function for cycle consistency is defined as:
	
    ℒ_c=1/M∑^M_m=1d(z_t,z̆_t)

	Guiding Future Scenario Decisions with Model Predictive Control. 
	In RL, policies π are formed by neural networks to approximate the optimal policy π_θ(·|s_t). In classical control tasks, however, π takes the form of trajectory optimization, often seeking a local best solution at each time step <cit.>. This involves predicting the best actions a_t:t+H over a short horizon H and executing the first action a_t, a method known as MPC:
	
    π_θ^ℳ𝒫𝒞=a_t:t+Harg max𝔼[∑^H_j=tγ^jℛ(s_j,a_j)]

	In the SAC framework, the actor selects samples from the distribution of state-action pairs (s_t, a_t) according to the current policy π: a_t ∼π_θ(·, s_t).
	To incorporate prospective thinking into model-free methods, we introduce the concept of MPC. Specifically, a_t is chosen as the first action that maximizes the long-term reward within a finite range H from j ∈ [1, n]:
	
    a_t=a_t:t+Harg max∑^H_j=tγ^jℛ(s_j,a_j)

	Due to its deep understanding of the environment model, MPC autonomously computes rewards from simulated trajectories (i.e., state-action pairs), which is distinct from model-free methods where the reward function ℛ(·,·) is provided by the environment, not computed by itself. Furthermore, it's essential to distinguish this approach from the iterative Q-value estimation in reinforcement learning, which prioritizes short-term gains over long-term cumulative rewards. 
	
	Therefore, we introduce the value consistency to guide the prospective action selection process. The evaluation of actions through prospective thinking should ideally be consistent with the Q-values determined by the RL policy π, ensuring the agent selects actions that provide the maximum long-term benefit in each state. As shown in Figure <ref>, we construct an action set Â(z_0)={â_0,..,â_t} from the current latent state z_0, which includes t potential actions. For each pair of trajectories (ẑ_i,a_i∈Â(s_t)), we use Q(ẑ_i,â_i) as a proxy for the reward function ℛ(·,·), where i ∈ [1,t]. Subsequently, after computing the cumulative discounted rewards CQ_i=∑^t_j=0γ^jQ(ẑ_j,â_j) for each trajectory set, we determine the locally optimal action a^*_0=arg max_a_0{CQ_1,⋯,CQ_k}. Ensuring the consistency of Q-values guarantees that the selected action a_t^* at each time step is the optimal choice according to the current policy, thus fulfilling the function of decision-making and guiding future scenarios.
	
		
		
	
	
	Overall Objective. As shown in Figure <ref>, the overall training objective of our approach can be summarized as follows:
	
    [ ℒ_total(θ)=𝒥_θ_RL  
    			+λ_predℒ_pred(θ)_pred
    			+λ_cℒ_c(θ)_Cycle Consistency ]

	where θ represents all model parameters used to compute the above loss functions, 𝒥_θ is the loss function for RL (see section <ref>), and λ_pred and λc are hyperparameters that weight different losses (more detailed can be seen in the supplementary material).
	
	

§ EXPERIMENTS

	In this section, we will provide a detailed introduction to the evaluation environment used in this paper, including experimental settings, experimental details, and evaluation metrics. Then, we perform a series of ablation experiments to analyze the key components of our method.
	

 §.§ Setup for Evaluation

	Environments. We conducted benchmark tests on ProSpec in environments with limited interaction, using the DeepMind Control Suite (DMControl) for continuous control testing <cit.>. Following the practices of most researchers <cit.>, we evaluated the performance after training for 100k and 500k interaction steps in the DMControl environment, selecting a total of 6 games for performance evaluation.
	
	Baseline. In this work, we consider SPR <cit.> as a strong baseline. We also selected Dreamer <cit.>, SAC+AE <cit.>, SLAC <cit.>, CURL <cit.>, DrQ <cit.>, SPR <cit.>, PlayVirtual <cit.>, and VCR <cit.> as baselines. It is worth noting that since SPR was originally designed for offline tasks, we adapted SPR for continuous tasks using a SAC-based approach (SPR^†).
	
	Implementation Datails. We used the encoder and policy network architecture of CURL <cit.> as the basic network. The contrastive loss component was removed from CURL, and we introduced BYOL <cit.>, creating a baseline solution similar to SPR. To generate imaginary trajectories in this environment, the prediction horizon was set to 6, and action generation was achieved by randomly sampling from a uniform distribution over the continuous action space. Additional implementation details are provided in the Appendix.
	
	
	Evalation Metrics. In DMControl, the maximum achievable score for each environment is capped at 1000 <cit.>. Consistent with previous studies <cit.>, we evaluated model performance across six common environments and used the median score to represent overall performance. To reduce potential variance in benchmarking results, we performed ProSpec evaluations on DMControl using 10 random seeds for each game. For the 100k step, we reran the official code for SPR and PlayVirtual with 10 random seeds and recorded the final results. We did not rerun the code for VCR as we were unable to obtain the official implementation. Due to limited computational resources, we report the 500k results directly from the respective papers. The entire implementation was using PyTorch <cit.>, with network training performed on an NVIDIA RTX 3090 GPU.
	
	

 §.§ Performance Comparison with State-of-the-Arts

	As shown in Table <ref>, our method surpassed the performance of most other methods by achieving five of the top scores in six environments, regardless of the limited 100k or 500k interactions. Furthermore, when considering the median scores, our method achieved the highest rankings in the following two aspects:
	(i) In the DMControl-100k environment with limited interaction data, our method achieves the highest median score of 807.5. This outperforms PlayVirtual by 8.32%, VCR by 3.66%, SPR by 9.64%, DrQ by 17.80%, and CURL by 44.20%.
	(ii) Our method achieves a median score of 959.5 and comes close to the maximum score of 1000 in four environments.
	These results demonstrate the exceptional performance of our method even with limited data. Furthermore, our method shows a significant improvement in performance across different environments.
	
	
	
	

 §.§ Analysis

	In this section, we performed an ablation study in the DMControl-100k to analyze the key components of ProSpec. Due to limited computing resources, we experimented using 5 random seeds.
	
	r4cm
		

		Effectiveness of FDM and Prospective Thinking.
		
		
			1l     Median Score 
 SPR^†              736.5                   

			PlayVirtual              745.5                  

			ProSpec-NP                  752.0           
 ProSpec           807.5          
 
		

	
	Effectiveness of FDM and Prospective Thinking. To explore the effect of prospective thinking on RL tasks, we evaluated ProSpec without prospective thinking (ProSpec-NP) and compared it to SPR^† and PlayVirtual. The results are detailed in Table <ref>. As can be seen: (i) ProSpec improved by 55.5 scores over ProSpec-NP; (ii) ProSpec-NP outperformed PlayVirtual by 6.5 scores; (iii) ProSpec increased PlayVirtual's median score from 752.0 to 807.5. Several factors may explain these results: First, using human-like prospective thinking, ProSpec explores the future from multiple perspectives and selects higher-value, lower-risk decisions, improving its performance by 55.5 scores over ProSpec-NP. Typcally, Prospective thinking resembles residual structures in neural networks, allowing the model to approach local optima with each decision, leading to consistent performance improvements. Second, the flow-based dynamics model provides more accurate backward predictions, while cycle consistency augments state traceability and guides low-risk decisions, improving ProSpec's performance over PlayVirtual. Finally, self-supervised methods generate multiple state-action trajectories, increasing data efficiency and allowing ProSpec to outperform SPR^†.
	
	r4cm
		

		The impact of cycle consistency on ProSpec.
		
		
			1l     Median Score 
 SPR^†              736.5                  

			ProSpec-NC                 752.8                
ProSpec           807.5          
 
		

	
	The Impact of Cycle Consistency on ProSpec. As described in Section 3.2, cycle consistency is used to address two fundamental issues in RL: augmenting state traceability to avoid irretrievable events (low risk), and generating a large number of virtual trajectories to improve data efficiency with augmented actions. We evaluated the performance of ProSpec without cycle consistency constraints (referred to as ProSpec-NC) and compared it to SPR methods. As shown in Table <ref>, regularizing the encoder with ℒ_c in ProSpec-NC resulted in a gain of 2.21%, while regularizing both the encoder and the FDM resulted in a gain of 7.27%. It is clear that cycle consistency can benefit the training of the FDM, but the greater benefit comes from the learning of the feature representations in the encoder <cit.>.
	
	r7cm
		

		Impact of perspective count k.
		
		
			                       k=1     k=3        k=6        k=9  
 Finger, spin           880.0              867.9               871.2               811.7         

			Cartpole, swingup      816.0              831.5          820.0                 800.0    

			Reacher, easy          675.8            762.6               770.3               763.6           

			Cheetah, run           349.8            388.2               384.5             419.7               

			Walker, walk           424.6            561.4               517.6               516.5               

			Ball in cup, catch     845.1            905.5               905.8               934.4             
Median Score           745.9            797.0     795.1    781.8                 
 
		
	
	Impact of Perspective Count k. We explored how the number of perspectives k affects ProSpec in lookahead tasks. Experiments were performed with different perspectives k ∈ [1, 9], and we recorded the median results under different prediction steps t ∈ [1, 15] (additional results are available in the supplementary material). Table <ref> presents these results: (i) at k = 3, ProSpec achieved peak performance, marking an improvement of 6.85% over k = 1 and 1.9% over k = 9; (ii) contrary to expectations, the performance improvement does not scale linearly with the number of perspectives. While k ≥ 3, the performance improvement reaches saturation and shows less differences. This may be because the actions sampled by the agent yield similar long-term rewards as the number of perspectives increases, resulting in insignificant differences in model performance.
	
	Impact of Prediction Steps t on ProSpec. We examined the effect of prediction steps t on ProSpec's performance, as shown in table <ref>. Similar to the previous experiment, we evaluated the median performance across different perspectives k. Our results can be summarized as follows: (i) Except for the finger-spin environments, ProSpec shows significant performance improvements in most environments. This suggests that ProSpec's prospective thinking effectively uses future information to improve the effectiveness of current decisions, and improves model performance across different scenarios. (ii) ProSpec achieves optimal performance at t = 6, showing improvements of 6.25% and 6.18% over t = 1 and t = 15, respectively. (iii) Notably, ProSpec's performance does not exhibit a linear increase with prediction steps t. Instead, beyond a certain threshold of t, performance growth begins to slow or even decline. This phenomenon may be due to the increased complexity of future states that the model attempts to capture with an excessive number of prediction steps, leading to additional prediction errors and compromising prediction accuracy.
	
	
	

§ CONCLUSION AND LIMITATION

	To imbue deep RL with cognitive abilities akin to those of humans—namely, prospective thinking—we propose a novel method named ProSpec.  ProSpec employs a reversible flow-based dynamics model to predict latent future outcomes, drawing on the concept of Model Predictive Control to guide decision-making in anticipation of future scenarios. This approach enables RL to simulate future n-stream trajectories, facilitating high-value, low-risk optimal decision-making. Additionally, ProSpec augments state traceability through cycle consistency, reducing environmental risks and leveraging this to augment actions, thus generating a wealth of virtual trajectories to improve data efficiency. We conducted a large number of experiments at DMControl and the results confirmed the effectiveness of ProSpec.  However, due to the necessity of exploring future scenarios from multiple perspectives, our method incurs additional overhead in training the agent. In addition, ProSpec is currently tailored for value-based RL methods; future research endeavors will aim to discover a universally applicable RL approach that incorporates prospective thinking.
	
	unsrtnat
	
	10
		
		mnih2015human
		Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness,
		Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg
		Ostrovski, et al.
		Human-level control through deep reinforcement learning.
		nature, 518(7540):529–533, 2015.
		
		silver2017mastering
		David Silver, Julian Schrittwieser, Karen Simonyan, Ioannis Antonoglou, Aja
		Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton,
		et al.
		Mastering the game of go without human knowledge.
		nature, 550(7676):354–359, 2017.
		
		schrittwieser2020mastering
		Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan,
		Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis,
		Thore Graepel, et al.
		Mastering atari, go, chess and shogi by planning with a learned
		model.
		Nature, 588(7839):604–609, 2020.
		
		zhihe2023dmbp
		YANG Zhihe and Yunjian Xu.
		Dmbp: Diffusion model based predictor for robust offline
		reinforcement learning against state observation perturbations.
		In The Twelfth International Conference on Learning
			Representations, 2023.
		
		hansen2022temporal
		Nicklas Hansen, Xiaolong Wang, and Hao Su.
		Temporal difference learning for model predictive control.
		arXiv preprint arXiv:2203.04955, 2022.
		
		levine2020offline
		Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu.
		Offline reinforcement learning: Tutorial, review, and perspectives on
		open problems.
		arXiv preprint arXiv:2005.01643, 2020.
		
		yarats2021improving
		Denis Yarats, Amy Zhang, Ilya Kostrikov, Brandon Amos, Joelle Pineau, and Rob
		Fergus.
		Improving sample efficiency in model-free reinforcement learning from
		images.
		In Proceedings of the AAAI Conference on Artificial
			Intelligence, volume 35, pages 10674–10681, 2021.
		
		bellemare2013arcade
		Marc G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling.
		The arcade learning environment: An evaluation platform for general
		agents.
		Journal of Artificial Intelligence Research, 47:253–279, 2013.
		
		badia2020agent57
		Adrià Puigdomènech Badia, Bilal Piot, Steven Kapturowski, Pablo
		Sprechmann, Alex Vitvitskyi, Zhaohan Daniel Guo, and Charles Blundell.
		Agent57: Outperforming the atari human benchmark.
		In International conference on machine learning, pages
		507–517. PMLR, 2020.
		
		yuan2024task
		Huining Yuan, Hongkun Dou, Xingyu Jiang, and Yue Deng.
		Task-aware world model learning with meta weighting via bi-level
		optimization.
		Advances in Neural Information Processing Systems, 36, 2024.
		
		zhu2024repo
		Chuning Zhu, Max Simchowitz, Siri Gadipudi, and Abhishek Gupta.
		Repo: Resilient model-based reinforcement learning by regularizing
		posterior predictability.
		Advances in Neural Information Processing Systems, 36, 2024.
		
		yarats2020image
		Denis Yarats, Ilya Kostrikov, and Rob Fergus.
		Image augmentation is all you need: Regularizing deep reinforcement
		learning from pixels.
		In International conference on learning representations, 2020.
		
		laskin2020reinforcement
		Misha Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind
		Srinivas.
		Reinforcement learning with augmented data.
		Advances in neural information processing systems,
		33:19884–19895, 2020.
		
		yu2021playvirtual
		Tao Yu, Cuiling Lan, Wenjun Zeng, Mingxiao Feng, Zhizheng Zhang, and Zhibo
		Chen.
		Playvirtual: Augmenting cycle-consistent virtual trajectories for
		reinforcement learning.
		Advances in Neural Information Processing Systems,
		34:5276–5289, 2021.
		
		jaderberg2016reinforcement
		Max Jaderberg, Volodymyr Mnih, Wojciech Marian Czarnecki, Tom Schaul, Joel Z
		Leibo, David Silver, and Koray Kavukcuoglu.
		Reinforcement learning with unsupervised auxiliary tasks.
		arXiv preprint arXiv:1611.05397, 2016.
		
		schwarzer2020data
		Max Schwarzer, Ankesh Anand, Rishab Goel, R Devon Hjelm, Aaron Courville, and
		Philip Bachman.
		Data-efficient reinforcement learning with self-predictive
		representations.
		arXiv preprint arXiv:2007.05929, 2020.
		
		yue2023value
		Yang Yue, Bingyi Kang, Zhongwen Xu, Gao Huang, and Shuicheng Yan.
		Value-consistent representation learning for data-efficient
		reinforcement learning.
		In Proceedings of the AAAI Conference on Artificial
			Intelligence, volume 37, pages 11069–11077, 2023.
		
		janner2019trust
		Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine.
		When to trust your model: Model-based policy optimization.
		Advances in neural information processing systems, 32, 2019.
		
		lowrey2018plan
		Kendall Lowrey, Aravind Rajeswaran, Sham Kakade, Emanuel Todorov, and Igor
		Mordatch.
		Plan online, learn offline: Efficient learning and exploration via
		model-based control.
		arXiv preprint arXiv:1811.01848, 2018.
		
		schacter2012future
		Daniel L Schacter, Donna Rose Addis, Demis Hassabis, Victoria C Martin,
		R Nathan Spreng, and Karl K Szpunar.
		The future of memory: remembering, imagining, and the brain.
		Neuron, 76(4):677–694, 2012.
		
		aspinwall2005psychology
		Lisa G Aspinwall.
		The psychology of future-oriented thinking: From achievement to
		proactive coping, adaptation, and aging.
		Motivation and emotion, 29:203–235, 2005.
		
		mckiernan2017prospective
		Peter McKiernan.
		Prospective thinking; scenario planning meets neuroscience.
		Technological Forecasting and Social Change, 124:66–76, 2017.
		
		kvavilashvili2020nature
		Lia Kvavilashvili and Jan Rummel.
		On the nature of everyday prospection: A review and theoretical
		integration of research on mind-wandering, future thinking, and prospective
		memory.
		Review of General Psychology, 24(3):210–237, 2020.
		
		schacter2007cognitive
		Daniel L Schacter and Donna Rose Addis.
		The cognitive neuroscience of constructive memory: remembering the
		past and imagining the future.
		Philosophical Transactions of the Royal Society B: Biological
			Sciences, 362(1481):773–786, 2007.
		
		hassabis2007deconstructing
		Demis Hassabis and Eleanor A Maguire.
		Deconstructing episodic memory with construction.
		Trends in cognitive sciences, 11(7):299–306, 2007.
		
		suddendorf2007evolution
		Thomas Suddendorf and Michael C Corballis.
		The evolution of foresight: What is mental time travel, and is it
		unique to humans?
		Behavioral and brain sciences, 30(3):299–313, 2007.
		
		deisenroth2011pilco
		Marc Deisenroth and Carl E Rasmussen.
		Pilco: A model-based and data-efficient approach to policy search.
		In Proceedings of the 28th International Conference on machine
			learning (ICML-11), pages 465–472, 2011.
		
		peng2018deep
		Baolin Peng, Xiujun Li, Jianfeng Gao, Jingjing Liu, Kam-Fai Wong, and Shang-Yu
		Su.
		Deep dyna-q: Integrating planning for task-completion dialogue policy
		learning.
		arXiv preprint arXiv:1801.06176, 2018.
		
		ha2018world
		David Ha and Jürgen Schmidhuber.
		World models.
		arXiv preprint arXiv:1803.10122, 2018.
		
		yu2020mopo
		Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Y Zou, Sergey
		Levine, Chelsea Finn, and Tengyu Ma.
		Mopo: Model-based offline policy optimization.
		Advances in Neural Information Processing Systems,
		33:14129–14142, 2020.
		
		van2019use
		Hado P Van Hasselt, Matteo Hessel, and John Aslanides.
		When to use parametric models in reinforcement learning?
		Advances in Neural Information Processing Systems, 32, 2019.
		
		kielak2019recent
		Kacper Piotr Kielak.
		Do recent advancements in model-based deep reinforcement learning
		really improve data efficiency?
		2019.
		
		laskin2020curl
		Michael Laskin, Aravind Srinivas, and Pieter Abbeel.
		Curl: Contrastive unsupervised representations for reinforcement
		learning.
		In International conference on machine learning, pages
		5639–5650. PMLR, 2020.
		
		guo2020bootstrap
		Zhaohan Daniel Guo, Bernardo Avila Pires, Bilal Piot, Jean-Bastien Grill,
		Florent Altché, Rémi Munos, and Mohammad Gheshlaghi Azar.
		Bootstrap latent-predictive representations for multitask
		reinforcement learning.
		In International Conference on Machine Learning, pages
		3875–3886. PMLR, 2020.
		
		lee2020stochastic
		Alex X Lee, Anusha Nagabandi, Pieter Abbeel, and Sergey Levine.
		Stochastic latent actor-critic: Deep reinforcement learning with a
		latent variable model.
		Advances in Neural Information Processing Systems, 33:741–752,
		2020.
		
		hafner2020mastering
		Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba.
		Mastering atari with discrete world models.
		arXiv preprint arXiv:2010.02193, 2020.
		
		sutton2018reinforcement
		Richard S Sutton and Andrew G Barto.
		Reinforcement learning: An introduction.
		MIT press, 2018.
		
		dinh2016density
		Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio.
		Density estimation using real nvp.
		arXiv preprint arXiv:1605.08803, 2016.
		
		huang2018orthogonal
		Lei Huang, Xianglong Liu, Bo Lang, Adams Yu, Yongliang Wang, and Bo Li.
		Orthogonal weight normalization: Solution to optimization over
		multiple dependent stiefel manifolds in deep neural networks.
		In Proceedings of the AAAI Conference on Artificial
			Intelligence, volume 32, 2018.
		
		tassa2018deepmind
		Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las
		Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al.
		Deepmind control suite.
		arXiv preprint arXiv:1801.00690, 2018.
		
		edwards2018forward
		Ashley D Edwards, Laura Downs, and James C Davidson.
		Forward-backward reinforcement learning.
		arXiv preprint arXiv:1803.10227, 2018.
		
		vinyals2019grandmaster
		Oriol Vinyals, Igor Babuschkin, Wojciech M Czarnecki, Michaël Mathieu,
		Andrew Dudzik, Junyoung Chung, David H Choi, Richard Powell, Timo Ewalds,
		Petko Georgiev, et al.
		Grandmaster level in starcraft ii using multi-agent reinforcement
		learning.
		Nature, 575(7782):350–354, 2019.
		
		wang2016dueling
		Ziyu Wang, Tom Schaul, Matteo Hessel, Hado Hasselt, Marc Lanctot, and Nando
		Freitas.
		Dueling network architectures for deep reinforcement learning.
		In International conference on machine learning, pages
		1995–2003. PMLR, 2016.
		
		hafner2019dream
		Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi.
		Dream to control: Learning behaviors by latent imagination.
		arXiv preprint arXiv:1912.01603, 2019.
		
		grill2020bootstrap
		Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre
		Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan
		Guo, Mohammad Gheshlaghi Azar, et al.
		Bootstrap your own latent-a new approach to self-supervised learning.
		Advances in neural information processing systems,
		33:21271–21284, 2020.
		
		hafner2019learning
		Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha,
		Honglak Lee, and James Davidson.
		Learning latent dynamics for planning from pixels.
		In International conference on machine learning, pages
		2555–2565. PMLR, 2019.
		
		paszke2019pytorch
		Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory
		Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al.
		Pytorch: An imperative style, high-performance deep learning library.
		Advances in neural information processing systems, 32, 2019.
		
	
	



 Paper Title: TinyChirp: Bird Song Recognition Using TinyML Models on Low-power Wireless Acoustic Sensors.


IEEEexample:BSTcontrol




TinyChirp: Bird Song Recognition Using TinyML Models on Low-power Wireless Acoustic Sensors 
Corresponding author: zhaolan.huang@fu-berlin.de. 



    
    Z. Huang1,
    A. Tousnakhoff1,
    P. Kozyr2,
    R. Rehausen3,
    F. Bießmann25
    R. Lachlan4,
    C. Adjih6,
    E. Baccelli61
    
    1Freie Universität Berlin
    2Berlin University of Applied Sciences
    
    
    
    
    
    3Woks Audio GmbH
    4Royal Holloway University of London
    
    
    
    
    
    5Einstein Center for Digital Future
    6Inria, France
    
    
    
    
    

    

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



§ ABSTRACT 
Monitoring biodiversity at scale is challenging. Detecting and identifying species in fine grained taxonomies requires highly accurate machine learning (ML) methods. Training such models requires large high quality data sets. And deploying these models to low power devices requires novel compression techniques and model architectures. While species classification methods have profited from novel data sets and advances in ML methods, in particular neural networks, deploying these state of the art models to low power devices remains difficult. Here we present a comprehensive empirical comparison of various tinyML neural network architectures and compression techniques for species classification. We focus on the example of bird song detection, more concretely a data set curated for studying the corn bunting bird species. The data set is released along with all code and experiments of this study. In our experiments we compare predictive performance, memory and time complexity of classical spectrogram based methods and recent approaches operating on raw audio signal. Our results demonstrate that our TinyChirp approach can robustly detect individual bird species with precisions over 0.98 and reducing the energy consumption such that a deployment time on a single battery charge can be extended from 2 weeks to almost an entire season of 18 weeks.  



TinyML, Microcontroller, Acoustic Sensors, Machine Learning




§ INTRODUCTION


The typical data pipeline with bio-acoustic research requires the deployment of sensors, in which each node is battery-powered and left in the field to record environmental sound, continuously, for a long period (e.g. for a whole season). Thereafter, the recorded data is manually collected by researchers on-site, from each node, and further analyzed in laboratory. In the case of avian species for instance, only the targeted bird species is relevant within the recorded data, and the rest of the data is to be discarded ultimately. This is both a cumbersome process and a waste of resources, not only at this downstream stage in the lab, but also upstream in the field, on each node, while recording large amounts of irrelevant data. In this pipeline, continuous recording of audio thus creates a bottleneck in terms of memory and energy budgets available on individual sensors – typically limited to a small battery, and an SD card, respectively, driven by rudimentary software running on a low-power microcontroller. Such hardware is very energy-efficient, but very limited in memory resources, with RAM memory budgets in the order of 500 kiloBytes <cit.>, which yields specific constraints on software embedded on such devices <cit.>.

Meanwhile, as AI achieves excellent performance in pattern recognition of audio and biosignals, more and more bioacoustic researchers leverage ML  related methods to improve accuracy and efficiency. 
However, until recently, such ML models, e.g. BirdNet <cit.> were confined to lab use only, as their resource requirements (GigaBytes of RAM) are way beyond the capacity of microcontroller-based hardware available on sensors in the field. However, recent advances in TinyML, a lively field of research targeting machine learning for microcontrollers, offer a glimpse of hope that pattern recognition of bioacoustic signals might become possible, over longer periods, as required by the aforementioned use cases. 






However, ..., which required to:
[inline]to be finish

    
  * train ML models.
    
  * deploy models on sensor node.
    
  * evaluate the resource consumption of models on boards.


This paper thus introduces ....


Paper Contributions. In this paper, we explore the possibility of using TinyML in practice, on common low-power microcontroller hardware, for a concrete use case: monitoring corn bunting birds' songs, in a rural area in the UK, over several months in a row, using a fleet of energy-efficient acoustic sensors. We aim to answer the following questions: (Q1) Can we pre-screen the audio on the low-power sensor node, so as to only store the targeted bird songs on each sensor? (Q2) How does that extend the device's lifetime, in terms of memory and energy budgets? More in detail, our contributions are as follows:


    
  * We propose a pipeline to record, recognize and store specific bird songs on low-power microcontroller-based devices, which can be further scale-up to universal species;
    
  * We developed a neural network with high accuracy on bird song recognition. Furthermore, we optimize the network with partial convolution technology to minimize the memory consumption on deployment;
    
  * We propose a two-stage classification approach to reduce computational and storage cost and enhance the overall accuracy;
    
  * We provide experimental results on both the classification performance of models and resource consumption on low-power devices;
    
  * We publish a new data set of curated audio recordings snippets of corn bunting bird songs;
    
  * We provide open source code[see <https://github.com/TinyPART/TinyBirdSounds/>] to reproduce the results on common MCU.



§ BACKGROUND & RELATED WORK




Bird Song Recognition –
Bird song detection has significantly advanced through various machine learning techniques, particularly deep learning methods. Recent research has explored different strategies for recognizing bird songs, with substantial progress in preprocessing audio data for recognition.
One prevalent preprocessing technique involves transforming the audio signal into a spectrogram using STFT <cit.>. This spectrogram can subsequently be converted into a mel spectrogram <cit.>,  by applying a mel scale, followed by a log scale transformation <cit.>. This process effectively redefines the audio recognition task as an image recognition task. CNN <cit.> and Residual Neural Networks (ResNets) <cit.> are widely utilized for such image recognition tasks in bird song detection. For instance, one study employed CNNs to recognize the songs of 17 bird species <cit.>, while another study utilized ResNets to classify the songs of 19 bird species <cit.>. BirdNET <cit.>, a well-known neural network for bird song recognition, is derived from the family of residual networks and can recognize 6,000 of the world’s most common species at the time of writing.

Another approach involves transforming the audio signal into Mel-Frequency Cepstral Coefficients (MFCCs), which are then used as input for bird song classification <cit.>. Alternatively, feature extraction using wavelet decomposition can be employed for this purpose <cit.>. Additionally, certain studies combine different architectures; for example, CNN and Transformer <cit.>, and CNN and VGG <cit.> to enhance performance. While these methods are effective for bird song recognition, their computational requirements are typically too demanding for execution on microcontrollers.

Bioacoustic Audio Datasets –
High-quality and diverse datasets are crucial for training models to accurately recognize birds and differentiate between species by their unique songs. The Xeno-canto database [see <https://www.xeno-canto.org>] is one of the largest and most comprehensive collections of bird sounds, containing over 500,000 recordings from more than 10,000 species. In <cit.>, 2,699 bird song audio recordings from the Xeno-canto database, ranging in length from 20 to 0.74 seconds, were used to classify 17 bird species.
Another extensive collection of bird audio recordings is available in the Macaulay Library [see <https://www.macaulaylibrary.org/>], renowned for its high-quality bird audio recordings and detailed metadata, which includes information about species, location, and context. The combination of the Macaulay Library and Xeno-canto audio recordings was used to train BirdNet <cit.>.
In <cit.>, the public dataset CLO-43DS, containing recordings of 43 different North American wood-warblers, was utilized. Another notable collection is the Birdsdata dataset with audio files of 20 bird species from the Beijing Academy of Artificial Intelligence (BAAI) repository, which is frequently used in research and model training <cit.>.
To complement datasets focused on bird species, the Google AudioSet <cit.> and Urbansound8K <cit.> datasets are used to include non-bird species sounds.



TinyML –
For models to operate on low-power devices, they must be compact and computationally efficient. Studies have demonstrated the use of lightweight CNNs for speech recognition and age classification <cit.>, water leakage detection <cit.>, acoustic defect localization <cit.>, fall detection for the elderly <cit.> and other tasks <cit.>. Tiny vision transformers have also been employed for classification tasks in various studies <cit.>. Moreover, quantization techniques can be applied to models to reduce their size to fit within the constraints of low-power devices <cit.>.

Creating spectrograms from audio signals can be power-intensive; hence, some studies use raw time-series data as input for neural networks <cit.>. For instance, single-channel EEG was used as input for a model based on CNN and Transformer for sleep stage classification <cit.>. Raw audio signals were used for urban sound analysis <cit.>. Time-series data offer advantages for TinyML applications, including reduced computational load and suitability for low-power devices.


Embedded Software Platforms for TinyML – The widely used model transpiler TVM (Tensor Virtual Machine <cit.>) has recently been extended with uTVM, providing automated transpilation and compilation for models output by major ML frameworks (TFML, Pytorch, etc.). As such uTVM exposes low-level routines and  optimizes these for execution on different processing units, including for targets such as a large variety of microcontrollers. Prior works such as <cit.>, <cit.> or MLPerfTiny <cit.> focused on the production, performance and analysis of standard benchmark suites of representative TinyML tasks on different microcontrollers. Conversely, prior work such as U-TOE <cit.> or RIOT-ML <cit.> provide embedded operating system integration of TinyML, facilitating TinyML benchmarking and continuous deployment over low-power wireless network links such as IEEE 802.15.4 or BLE.



§ BIOACOUSTIC MONITORING SCENARIO



As depicted in Figure <ref>, a network of battery-powered, autonomous recording units (ARUs, i.e. microcontroller-based acoustic sensors) is deployed across a monitored area. An example of ARU is given in <cit.>.

These sensors remain in the field for an entire season, typically around six months. Locations are often remote and hard to access and devices may be distributed over a relatively wide area, making visiting them time-consuming. Birds typically do not sing consistently during the whole day, while each bird has several different song types that need to be regularly sampled. Moreover, each individual bird moves around its individual territory, singing from a range of song posts, not all of which will be adequately recorded from any one location. 

On the network aspect, the sensors are distributed approximately evenly, ensuring that each sensor can be wireless connected by more than two neighboring sensors. This arrangement facilitates approximate triangulation through distance estimation based on signal strength. The placement process can be streamlined using LEDs and a basic ultra-low-power wireless protocol, such as IEEE 802.15.4 or LoRa. For space reason, we do not detail further networking aspects in the paper, but rather focus on the on-device machine learning aspects. We nevertheless keep an eye on the total memory footprint on the device, to ensure that the operating system including application TinyML code, the OS and a typical wireless low-power network stack fit in typical resource budgets on common microcontroller-based boards.

Although our work is designed around one very specific real-world research scenario (monitoring common Corn Bunting bird songs) the architectures should be readily applicable to other audio data classification tasks. Examples of related scenarios to which our method could be extended include, surveying remote areas for targeted endangered species of conservation value, or broader surveys of songbird biodiversity.  Similar devices would even be useful in the search for very elusive species such as Ivory-billed Woodpecker <cit.>. 




 §.§ Corn Bunting Monitoring Use-Case




SongBeam <cit.> microcontroller-based recorders have been used to monitor Corn Bunting (Emberiza calandra) birds in Oxfordshire, UK. 
A set of 30-40 devices have been deployed since 2022, currently comprising 3 complete breeding seasons (February - July).

SongBeam devices are based on an ARM Cortex-M microcontroller, run on 4 D-cell batteries and record 4-channel WAV files onto 128  microSD cards. To maintain such a fleet of sensors, devices are currently checked approximately every 2 weeks, especially because memory space is soon exhausted on the microSD cards. 

The above deployment has been used to produce a significant part of the dataset we present in Section <ref>.
Our experience using SongBeam and a preliminary analysis of the raw dataset we produced shows that less than 10% of recording time contains useful recordings. 
Given the extended deployment period (6 months), remaining low-power while improving storage efficiency is thus paramount. 
Moreover, current prototypes use solar power, further emphasizing that the bottleneck is downloading data from the microSD cards to free space on-board. 

This description of the limitations of using SongBeam recorders to monitor Corn Buntings is likely to apply to many biomonitoring scenarios: solar power may allow for extending deployment lifetime in remote locations, but storage of recordings is a limiting factor.






Describe Robert's scenario.
A bunch of battery-powered, microcontroller-based acoustic sensors are disseminated across an area to be monitored. The sensors are left in the field for a whole season at a time ( 6months). Power and storage efficiency are paramount.
To emphasis the network aspect, we can assume that the sensors are placed approx. evenly so that each sensor hears N>2 other sensors beacons, so that approx. triangulation is possible via (signal strength) estimated distance. Placement can be simplified via LED and a basic ultra low-power wireless protocol such as IEEE 802.15.4 or LoRa.


§ EVALUATION METRICS



On the one hand, as the basic part of bird song recording, the model should identify the target bird song as much as possible, potentially requiring a complex structure with a large size. On the other hand, the limited resource budget allows only simple models to be deployed on low-power IoT devices.

Thus we considered two orthogonal types of metrics to evaluate the prediction performance on bird song classification and resource usage on low-power IoT devices of TinyML models. Preliminary results are presented in Section <ref>.




 §.§ On-device Resource Usage


We performed extensive profiling of all models on low-power devices to capture their effective memory and time complexity in a real-world application scenario. The metrics are helpful to investigate the energy efficiency of the models, complementing the commonly used metrics for predictive performance. 

Memory (RAM) Consumption –
This metric measures the amount of dynamic memory space (primary RAM) consumed by the model during inference. It reflects the memory footprint of the model activation and is important for low-power devices that have limited memory resources. Efficient memory utilization allows for the deployment of larger and more complex models on such devices.

Storage (Flash memory) Consumption –
This metric quantifies the amount of storage space, typically in terms of Flash memory region, required to store the compute instruction and associated parameters. 

Minimizing storage consumption allows for accommodating multiple models on the device or orchestrating with other essential applications.

Computational Latency –
This metric measures the time consumption of performing inference for each input sample at the model level. It reflects the inference speed of the model on the low-power device and plays a crucial role in real-time or latency-sensitive applications. Core clock frequency, cache strategies and communication latency between memory and working core have a great impact on this indicator.

Energy Consumption – 
This metric is crucial for battery-operated and resource-constrained devices, where efficient energy usage can significantly impact device longevity and performance. Energy consumption encompasses both active power (when the device is triggered to perform tasks) and idle power (when the device is in sleep mode).



 §.§ On-device Prediction Performance



Accuracy –
It is defined as the ratio of correctly predicted instances to the total instances in the dataset. While accuracy is intuitive, it can be misleading, especially in highly imbalanced datasets. 

To address this limitation, accuracy should be combined with other metrics that offer a more detailed view of model performance.

Precision and Recall –
They are two fundamental metrics of classification models used to evaluate the ability to distinguish true positives (TP) from false positives (FP) and false negatives (FN), particularly in the context of imbalanced datasets. 


F-Score –
There is often a trade-off between precision and recall; increasing one can lead to a decrease in the other. 

As two common instances, F_1-score weighs them evenly, while F_2-score treats recall as two times more important than precision, applying in scenarios where false positives are more tolerant than false negatives.

ROC Curve –
While the above metrics are sensitive to class imbalance – which occurs in many species classification and detection tasks – there are other metrics that are more robust towards class imbalance. The receiver-operator characteristic (ROC) curve depicts the  TPR and FPR across different decision thresholds.



The AUC is a key metric derived from the ROC curve that quantifies the overall performance of the model and is invariant with respect to class imbalance. An AUC of 1.0 indicates perfect classification, while an AUC of 0.5 suggests no discriminative ability.  


move to discussion part
    - Confusion matrix
    - ROC Curve -> different cut-off value





§ METHODOLOGY


We used a combination of project-specific and publicly accessible data to train and validate our classification models. In the data pre-processing phase, the raw audio signals were further segmented, labeled, down-sampled and divided into different groups to generate classification datasets. The corresponding Mel-spectrograms of the pre-processed audio segments were also created for spectral-based methods (models). During the pre-processing stage, we conducted a pilot analysis to establish guidelines and determine specific hyper-parameters for downsampling and spectrogram generation.




 §.§ Data Acquisition


The publicly accessible data originated from Xeno-Canto, Macaulay, Google AudioSet, while the project-specific data were previously collected in a long-term research of bird song patterns. We published the resulting dataset[see <https://www.kaggle.com/datasets/zhaolanhuang/tinychirp>].



    
  * Oxfordshire Corn Buntings. This library contains recordings of corn buntings along a transect of approximately 20km in Southern England. Corn buntings sing in a mosaic-like pattern of geographical variation called dialects; our sample contains approximately 6 different dialects. The recordings were performed with directional parabolic microphones as described in <cit.>.  
    
    
    
  * Macaulay Library. This library contains the necessary audio recordings for our target species as well as the other identified species. There are a total of 796 entries of audio recordings for the target species. We filtered the audio files by the song tag and removed 16 entries due to unclear metadata and missing catalog numbers. For other species as non-target, we limit the selection to a maximum of 30 recordings per species and choose those with the highest average community rating, acquiring a total of 1468 recordings.
    
    
  * Xeno-Canto. This library also contains recordings of target species and other species. We retrieved the recordings only marked with song and rated with the best quality level, resulting in 303 recordings of target species and 9622 recordings of other species as non-target.
    
    
  * Google AudioSet. This dataset provides environmental sounds which contain non-bird sounds. 


After gathering all audio recordings from the above libraries, we used BirdNET to identify and chop corn-bunting segments from all datasets, and labeled them as target bird songs; all segments were truncated or zero-padded to the length of 3 seconds. Meanwhile, we chose other 3-s segments randomly and ensured no overlap with corn-bunting segments, with labeled as non-target bird songs. Table <ref> shows the distribution of target and non-target segments over all datasets.



 §.§ Data Pre-processing



This phase contains the following steps:


    
  * We divided the segments into training, validation and test sets with the ratio of 80:10:10.

    
  * All segments were downsampled to 16  using zero-order holder, which constituted the audio dataset.
    
    
  * We transformed the audio segments into Mel-Spectrograms, and grouped them with the same splitting ratio to form the spectral dataset.



Pilot Analysis –

We averaged the STFT spectrograms of target and non-target segments in the training dataset to investigate their frequency characteristics, as depicted in Figure <ref>. Obviously, a bright band lays on the spectrogram of target segments roughly between 4000 and 8000 , hints that a sample rate with 16  should be sufficient to preserve all frequency components of corn-bunting song according to Nyquist–Shannon theorem. A highpass filter can also be applied to eliminate low-frequency noise without damaging target songs. Thus, we downsampled the segments from 48  to 16  to reduce the resource consumption and improve the efficiency of classification phase. In the baseline model we also designed a highpass filter to enhance SNR.




Mel-Spectrogram –

Mel-spectrogram is commonly used in audio classification tasks due to its ability to closely approximate human auditory perception. To create log Mel-spectrograms, we first generated magnitude spectrograms of STFT for all downsampled audio segments. The STFT was performed with a Hann window with a width of 1024 samples and a step size of 256 samples. This process generated 184 window frames and 513 frequency bins. Thereafter, the magnitude spectrograms were mapped onto the Mel scales with 80 Mel bins. In creating the Mel-spectrograms, we employed a sampling rate of 16  and focused on the frequency range of 80-8,000 Hz. This frequency range was chosen because it encompasses the typical frequency range of most bird songs, including that of the target Corn Bunting. Lastly, we converted the Mel-spectrograms into a logarithmic scale to obtain log Mel-spectrograms, with the shape of 184 × 80. 









| Data         | # Bird | # Non-Bird |
| ————— | —— | ———- |
| Roberts         |        |            |
| Xeon-Canto      |        |            |
| Macaulay        |        |            |
| Google AudioSet |        |            |


§ BASELINE & DECISION STRATEGY


We propose an approach to optimizing the screening process by combining a high-speed signal processing step (hereafter named "baseline") with the precision of an TinyML model, ensuring reliable results while maintaining resource-consumption efficiency.



 §.§ Baseline


We developed a lightweight pre-selector signal processing step, not relying on machine learning. First, the audio segments are normalized by min-max schema. Thereafter, a 9th-order highpass Butterworth filter with the cut-off frequency of 7000  is imposed on the normalized segments, in order to suppress noise and non-relative sound in low-frequency and to preserve only target components. Finally, we calculate the signal power P of the filtered segments,


    P = 1/N∑_n=1^N x(n)^2,

where x(n) and N denote the data points and the length of a segment, respectively. 
We define two thresholds: t_low and t_high. If P is smaller than t_low, the microcontroller remains in the idle/low-power state, and the audio sample is discarded directly. Else, P is compared to t_high. Depending on this result, and on the decision strategy (see below), the audio sample is either stored or further analyzed via inference using the machine learning model.





 §.§ TinyChirp Decision Strategy





We employ a double-phase approach for enhanced accuracy and efficiency of bird song recognition, combining the baseline with a TinyML model for further verification. The process is as follows:

    
  * Step 1: Baseline Filtering – Pre-screening is conducted using the baseline pre-processing step described in Section <ref>, designed to provide a quick and efficient preliminary analysis. If P is smaller than t_low, the microcontroller remains in the idle/low-power state, the audio sample is discarded and the process aborts. Else proceed to Step 2.
    


 

    
  * Step 2: Inference-based Classification – If the power-saving flag is set and if P is smaller than t_high a higher-accuracy classification is conducted via inference  using TinyML model. If the audio sample is not classified as the target, the audio sample is discarded and the process aborts. Else proceed to Step 3.


  * Step 3: Storage – Without further processing, the audio sample is logged, i.e. stored on the SD card. 

    


Note that the decision strategy can therefore be configured with two "knobs". On one hand the power-saving flag determines skipping (or not) some of the computation and can thus decrease energy consumption. On the other hand, different TinyML models can be inserted in Step 2.



Next, we thus study two families of TinyML models for Step 2: the first category of model takes a Mel-Spectrogram as input, while the second category of model takes time-series as input. Table <ref> presents a summary of the explored TinyML model structures. 



This study mainly investigated the performance of two types of ML models – CNN and Transformer – on time-series and Mel-spectral input. Due to the compute- and memory-intensive characteristic of ML methods, two major optimization strategies were leveraged to compress the models and improve their efficiency, as discussed in Section <ref>. 

This study also presents a baseline model built with classical digital signal processing methods, as performance reference of ML-based models. Furthermore, we integrated baseline and ML-based model into ensemble one to gain a optimal compromise between accuracy and efficiency resource consumption? on low-power devices. 






§ CLASSIFICATION BASED ON SPECTROGRAMS



In contrast to the computer vision community, where the transition from frequency decomposition based feature extraction to neural feature extraction on raw image data introduced in 2012 led to a significant increase in predictive performance <cit.>, neural feature extraction on raw audio time series did not contribute to a comparable breakthrough yet. Instead most ML techniques for audio pattern recognition (somewhat surprisingly) were relying on image processing pipelines – sometimes with neural networks trained on ImageNet data. Basically, image classification is performed, using a neural network, analyzing the spectrogram rendering of the audio trace.

For this reason, we initially consider the below two models, based on state-of-the-art neural network architectures.

CNN-Mel –
This model contains two 2D convolutional layers for Mel-spectral inputs as feature extraction and two fully connected layers as classifier. Both convolutional layers utilize 4 filters with a 3 × 3 kernel size and ReLU activation function, followed by 2 × 2 max pooling to reduce spatial dimensions. The output of the classifier are normalized by Softmax function as well.

SqueezeNet-Mel – 
This model is based on SqueezeNet <cit.>, an advanced backbone focused on optimizing the efficiency of computer vision applications by strategically reducing parameters, with comparable performance to AlexNet <cit.>. It leverages a so-called Fire module to achieve higher efficiency compared to standard convolutions with only a slight decrease in accuracy. The Fire module significantly reduces the number of parameters by the squeeze layer (1×1 convolutions), which exceedingly reduces the number of channels before the more parameter-heavy 3×3 convolutions in the expand layer. This reduction in intermediate channels means fewer parameters overall.
We aligned its input shape with the Mel-spectrogram and tailored the output for binary classification. Also offering potential for further compression using quantization techniques, SqueezeNet is thus appealing given memory constraints.





Limitations of Spectrograms on Microcontrollers
Practical experience on low-power microcontroller-based devices has shown however that producing and manipulating mel-spectrograms from audio signal streams on such devices is problematic. It inserts an additional step in the processing pipeline, which increases latency and memory requirements. Prior work (such as <cit.>, section 6.2) concludes that the CPU bottleneck is not just the model inference time, but also the spectrogram calculation, which, compared to inference alone, almost doubles latency. Other previous work such as <cit.> measures on a quite powerful STM32F7 microcontroller that computing and writing in memory the mel-spectrogram of 30 columns and 40 frequency bands takes approximately 1 second.  Note that these spectrogram dimensions are much smaller than our requirements (184 columns × 80 frequency bands) hence even more latency can be expected in our case.

For these reasons, we next explore pipelines which skip the spectral pre-processing stage as described below.



SqueezeNet<cit.> is a CNN characterized by comparable performance to AlexNet while using approximately 50 times fewer parameters. The compact architecture of SqueezeNet is achieved by employing eight unique building blocks called the Fire Modules. Each Fire Module starts with a "squeeze" convolution using 1x1 filters to reduce input channel dimensions, thereby cutting down on computational load. This initial step compresses data into fewer channels. Following this, an "expand" convolution layer mixes 1x1 and 3x3 filters: the former expands channel numbers, while the latter captures spatial details. Moreover, SqueezeNet offers the potential for further compression using quantization techniques, making it an efficient choice given memory constraints.

The implemented SqueezeNet-Time-Series architecture uses a reduced number of feature maps in compare to initial SqueezeNet. This compressed model has in the first convolutional layer 0.3 times the number of filters that initial model has in this layer. In the fire modules in convolutional layers, we use the integer value of 0.3 times the number of squeeze filters given to this layer in initial model. Since the number of expand filters depends on the number of squeeze filters and is four times greater, it also undergoes compression. The structure of the fire module in SqueezeNet-Time-Series is presented in Table 2. Therefore, if the initial number of squeeze filters in the fire module is x, the compressed number of squeeze filters in this fire module will be int(0.3× x). Then both the expand layers will have 4× int(0.3× x) filters.

Additionally, since the initial model uses 2D convolutions and time-series are 1D arrays, we converted all 2D convolutional layers to 1D convolutional layers. The model was trained using time-series data for 10 epochs with the Adam optimizer and cross-entropy loss function. After training, this model was further quantized using different post-training quantization techniques.





For spectrogram approach we aimed to create an ultra-compact CNN-Mel-Spec model, comprising just two convolutional layers and two fully connected layers. The input for the model is a log Mel spectrogram with a size of 184x80. Both convolutional layers utilized 4 filters with a 3x3 kernel size and ReLU activation function, followed by max pooling with a 2x2 pooling window to reduce spatial dimensions. After the last convolutional layer and max pooling, the data were reshaped into a one-dimensional vector for input into the subsequent fully connected layer. The first fully connected layer had 8 neurons and utilized ReLU activation. The final fully connected layer contained two neurons with a linear activation function. For optimization, we employed the Adam optimizer with a learning rate of 0.001, and the model was trained over 15 epochs, using cross-entropy as the loss function. In an effort to reduce model size, we converted it to TFLite format and applied several post-training quantization techniques: float16 quantization, dynamic range quantization, and full integer quantization.













§ CLASSIFICATION BASED ON TIME-SERIES


Contrary to the models described in the previous section and inspired by the success of neural feature extraction in the computer vision domain <cit.>, the pipelines we aim at next take directly the raw audio signal time-series as input. More precisely, we designed the three models described below.

CNN-Time –
This simple model performs feature extraction with two sequential (1D) temporal convolutional layers followed by average pooling. A max pooling is inserted in-between to reduce the dimension of feature maps and enhance the non-linearity of the network. Two fully connected layers act as classifier at the end of the network, with the probability outputs normalized by Softmax function.  

Transformer-Time –
Inspired by <cit.>, we designed an efficient, attention-based model with only one temporal convolutional layer and one single-head transformer. The convolutional and pooling layers serve as feature extraction to transfer raw audio signal into embeddings for the transformer. The activations of the fully connected layer are normalized by Softmax function as well.



SqueezeNet-Time –
This model is based on SqueezeNet <cit.>. We tailored this basis to align the input shape with the audio segment, and narrowed the output channels for binary classification. To adapt the backbone to time-series input, all 2D convolutional layers are replaced by temporal convolutional layers. Furthermore, to have a more compact structure, we decreased the filter number of all convolutional layers by roughly 70%, as presented in Table <ref>.





§ ADDITIONAL MODEL OPTIMIZATIONS FOR TINYML


In this study, we used two optimization techniques to compress the models and further reduce their memory consumption: model quantization on the one hand and on the other hand partial convolution, as described below.







 §.§ Model Quantization

Quantization reduces the model size and inference time by converting the weights and activations from higher precision (e.g., 32-bit floating-point) to lower precision (e.g., 8-bit integer) <cit.>. In this study we adopted PTQ with weights and activations quantized in 8-bit integer. To avoid substantial loss of prediction performance, the training dataset was used to find the optimal scale factor and zero-point of the activations.



 §.§ Partial Convolution

We observed that the outputs of the Conv1D layer require significant memory, which impedes deployment on resource-constrained tiny devices. As presented in Table <ref>, the peak memory consumption occurs at the first Conv1D layer both in CNN-Time and Transformer-Time with 768  and 3 , respectively. This indicates that deployment on resource-constrained tiny devices is impractical.

To address this issue and inspired by <cit.>, 

we exploited the fact that average pooling can be computed iteratively point-by-point on the output channels of the final Conv1D layer over a small sliding window of inputs. Each point in the output channels depends only on a small subset (kernel window) of the outputs from the previous layer. That is, for a block of L Conv1D layers following by an average pooling layer, its output y can be iteratively computed on the input sequence x(n), n= 1 … N as following:

    y_j(k)    = y_j(k-1) + 1/N_L A^L_j(k), 
    
        k    = 1 … N_L, j = 1 … C_L . 
     
        A^l_c(n)    = ∑_i=1^C_l-1 W^l_c(i) · [ A^l-1_i(n-K^l/2) ⋯
       =4A^l-1_i(n) ⋯ A^l-1_i(n+K^l/2) ]^T 
    
        A^0(n)    = x(n) , 
    
        l    = 1 … L, c = 1 … C_l, n = 1 … N_l, N_0 = N

where y_j(N_L) is the result of the average pooling of the j-th channel; A^l_c(n) denotes the n-th point in channel c of the l-th Conv1D layer calculated with kernel weights W. Each Conv1D layer contains C_l channels with filter size of K^l and output size of N_l. Without loss of generality, the formula of partial convolution contains only the linear components and considers only one-dimension case for the sake of simplicity; it can be generalized in high-dimension and integrated with non-linear building blocks (e.g., stride, non-linear activations, pooing layers, etc.).

Figure <ref> provides a comparison between classic convolution and partial convolution. Unlike classic convolution where entire channels (C × N) are computed and stored before being processed by the next layer, partial convolution requires only a small part of the channels (C × K, K << N) for each layer. In our case with K=3 and N=48000, it can be expected roughly 16000 × smaller memory consumption, making the implementation more suitable for tiny devices. We applied this strategy to the Conv1D layers combined with average pooling in the CNN-Time and Transformer-Time architectures.










§ TINYCHIRP PERFORMANCE EVALUATION


In the following, we conducted a comprehensive evaluation of the discriminative performance of TinyChirp with different TinyML models, as well as of resource consumption and computation time on typical low-power boards based on microcontrollers.



 §.§ Classification Performance Comparison







Our evaluation process began with the generation of ROC curves and the corresponding AUC on the training dataset, providing a clear picture of the models' abilities to distinguish between classes. As shown in Figure <ref>, the spectrogram-based models – CNN-Mel and SqueezeNet-Mel – achieved the best classification performance with AUCs of 1.0 and 0.99, respectively, followed by time-series models – CNN-Time and Transformer-Time – both with 0.98 AUC. Unexpectedly, with a more complex structure, SqueezeNet-Time performs worst (lower AUC than the baseline).


We next measured accuracy, precision, recall, F_1- and F_2-score, shown in Figure <ref>. Spectrogram-based models achieve the highest metrics for the whole range of threshold values. For time-series models, Transformer-Time worked overall better than CNN-Time; Again, SqueezeNet-Time performs worst (consistent with the above ROC analysis). 


Next, we focused on finding the threshold settings for optimal  F_2-score prioritizing recall, i.e. reducing the likelihood of mistakenly discarding records of the targeted birds). The results are shown in Table <ref> for training data, thereafter verified on test (previously unseen) data as shown in Table <ref>. Note that this threshold influences energy consumption: more target bird classification leads to more energy consumption (as data needs to we stored on to SD card in this case).

These results are promising and allowing to focus only on the relevant audio segments (10% of the total recording time), as the classification performance is compelling. Thus, TinyChirp can save 90% of SD card space compared to the initial monitoring scenario (Section <ref>), potentially extending the deployment time to 18 weeks. However, complementary experiments must now be carried out on targeted hardware, so as to evaluate other key metrics: computation time, energy consumption and memory footprint on typical microcontroller-based devices. The next sections focus on that part.




presents the optimized setting with corresponding metrics on the training dataset. This approach ensured that our models were fine-tuned to prioritize recall, thereby reducing the likelihood of missing critical positive cases. Also we emphasize that the decision threshold is directly related to energy consumption in the use case considered here: more positive predictions will lead to higher energy consumption. Hence high precision is desirable for low energy consumption, but recall should be traded with caution against energy consumption.

After determining the optimal threshold setting on the training dataset, we validated this threshold using the test dataset, as described in Table <ref>. This step was crucial to ensure that the threshold not only performed well on the training data but also generalized effectively to unseen data. By confirming the optimal threshold on the test dataset, we demonstrated the robustness and reliability of our TinyML models' predictive performance.













 §.§ Performance Evaluation on Microcontrollers


Experimental Setup – For our measurements, we used a common low-power off-the-shelf board: the nRF52840 Development kit (nrf52840dk). This board is based on an ARM Cortex-M4 processor, with a typical memory budget: 1  of Flash memory and 256  of RAM. As embedded software base we used RIOT-ML <cit.>. 

Throughout the experiments, we monitored the (RAM/Flash) resource footprint using built-in diagnostic tools provided by RIOT-ML and RIOT <cit.>, as well as external measurement equipment as described below.

Energy consumption was measured using an ampermeter to gauge the board's energy efficiency. This involved recording the current draw of the MCU and the external storage system (SD card) in different stages. We used a voltage regulator to supply 3.3  DC output and a logger to capture detailed current profiles over time.

Scope of the Measurements – Besides the model inference stage, we also considered the resource footprint of pre-processing during bird song recognition. Pre-processing refers to down-sampling and calculation of Mel-spectrogram as described in Section <ref>.

Also, the energy consumption in MCU's lowest power mode (idle stage) was measured as the reference of absence of sound with sufficient intensity.



Measurement Results on Microcontrollers – The first striking observation is that  SqueezeNet-Time and SqueezeNet-Mel were not deployable due to OOM (exceded memory budget) and were thus excluded in this experiment. All other models fit in the constraints of memory and storage budget, even when discounting the memory footprint of the OS and the network stack for wireless communication. 

The resource footprint of different stages and models measured on an nrf52840dk board are shown in Table <ref>.  We observe that the model inference stage takes the bulk of the overall energy consumption, as it requires the most process time and MCU in full power mode. 
As expected, the idle stage requires the least power, as all peripherals and the MCU are in the lowest power (standby) mode (until woken up by sound with sufficient amplitude). 











Spectrogram vs Time-series – Considering only the prediction performance in Section <ref>, it would seem natural to choose CNN-Mel as the best candidate, since it also provides the best inference latency. However, taking pre-processing time into account (which for CNN-Mel includes producing a spectrogram) shows a different picture. It takes roughly 2.4 seconds in total for a 3 second audio snippet, which is dangerously near the real-time constraint and costs high power consumption. With the best prediction performance and the lowest total compute latency among time-series models, the Transformer-Time now appears like the best choice to deploy on tiny devices.



Next, we consider jointly the overall resource consumption and the classification performance to discuss variants for the decision strategy described in Section <ref>:

Baseline Only –
With the Baseline t_low (Step 1 and Step 3 only) setting, the required computational resources are minimal, while achieving a recall of 0.99; however, 
due to the low precision, many false positives are stored, leading to the lowest storage efficiency.

TinyChirp Skipping Baseline –
With this variant, the TinyML model is applied immediately, without pre-screening by the baseline (Step 2 and Step 3 Only), leading to high predictive performance and high storage efficiency, but requiring higher energy consumption. Even when using the Transformer-Time with the lowest energy footprint, this requires more than 9 times the energy required for the baseline. This variant is ideal for devices currently running out of storage space while still with substantial battery left.

TinyChirp –
This variant refers to the full series of Step 1-3 described in Section <ref>, whereby the samples not discarded by the baseline are double-checked by the TinyML model. This approach offers a good compromise of compute, storage and energy footprint, by discarding early, without impacting too much the overall prediction performance achievable by the TinyML model on its own. 

TinyChirp with Power-saving –
This variant (Step 1-3 with the power-saving flag enabled) provides a more energy-efficient twist to TinyChirp.

However, the overall precision is determined by the threshold setting t_high, which causes an extra 40% of false positives to be stored onto the SD card, leading to lower storage efficiency. This variant seems more suitable for currently running out of battery, but with sill sufficient storage resources.






§ CONCLUSION

In this paper we present an empirical study of a bioacoustic monitoring use case showing that machine learning is usable directly on autonomous recording units based on low-power microcontrollers. We focused on a concrete use-case: monitoring Corn Bunting birds in rural UK. We publish a data set of the recordings, based on which we demonstrated that TinyML can perform on-device high accuracy classification for bird song classification. Compared to the traditional approach used so far, our approach TinyChirp can extend from 2 weeks to 18 weeks – almost a full season – the time intervals until which researchers must harvest overfull data storage embedded on the recording units deployed in the field. Our experiments also observe hands-on how performing inference directly on the audio time-series is the better approach on microcontroller-based hardware, compared to inference on mel-sepctrograms typically used in audio pattern recognition. Beyond the specific use case of monitoring Corn Bunting birds' songs, based on our general-purpose TinyML approach and our open-source code implementation integrated into the RIOT operating system, TinyChirp is applicable to a wide range of similar bioacoustic use cases, using a large variety of low-power microcontroller-based hardware.




§ ACKNOWLEDGMENT

The research leading to these results partly received funding from the MESRI-BMBF German/French cybersecurity program under grant agreements No. ANR-20-CYAL-0005 and 16KIS1395K. We acknowledge the receipt of media from the Macaulay Library at the Cornell Lab of Ornithology and the Xeno-Canto database.  Felix Bießmann received support from the Einstein Center Digital Future, Berlin,
the German Research Foundation (DFG) (528483508 - FIP 12), the Federal Ministry of Economic Affairs and Climate Action (RIWWER 01MD22007H), the German Federal Ministry of the Environment (GCA 67KI2022A) and the German Federal Ministry
of Education and Research (KIP-SDM 16SV8856).


IEEEtran 









 Paper Title: StyleRF-VolVis: Style Transfer of Neural Radiance Fields for Expressive Volume Visualization.

                     















1391










IEEE VIS Paper








Representations & Interaction










  
  
  * K. Tang and C. Wang are with 
the Department of Computer Science and Engineering, University of Notre Dame, Notre Dame, IN 46556, USA.
E-mail: {ktang2, chaoli.wang}@nd.edu.



In volume visualization, visualization synthesis has attracted much attention due to its ability to generate novel visualizations without following the conventional rendering pipeline. However, existing solutions based on generative adversarial networks often require many training images and take significant training time. Still, issues such as low quality, consistency, and flexibility persist. This paper introduces StyleRF-VolVis, an innovative style transfer framework for expressive volume visualization (VolVis) via neural radiance field (NeRF). The expressiveness of StyleRF-VolVis is upheld by its ability to accurately separate the underlying scene geometry (i.e., content) and color appearance (i.e., style), conveniently modify color, opacity, and lighting of the original rendering while maintaining visual content consistency across the views, and effectively transfer arbitrary styles from reference images to the reconstructed 3D scene. To achieve these, we design a base NeRF model for scene geometry extraction, a palette color network to classify regions of the radiance field for photorealistic editing, and an unrestricted color network to lift the color palette constraint via knowledge distillation for non-photorealistic editing. We demonstrate the superior quality, consistency, and flexibility of StyleRF-VolVis by experimenting with various volume rendering scenes and reference images and comparing StyleRF-VolVis against other image-based (AdaIN), video-based (ReReVST), and NeRF-based (ARF and SNeRF) style rendering solutions.














	[ < g r a p h i c s > < g r a p h i c s >;                     ]




Given a set of multi-view DVR images produced from a volumetric dataset, StyleRF-VolVis transfers arbitrary styles from reference images to the reconstructed 3D scene, synthesizing high-quality novel views with visual content consistency.
  


















figs/figures/pictures/images/./ 


                      
                  
                    
                       




                  











         




 






 














KwInputInput                
KwOutputOutput              

   
  
 








myitemize

 
 
 
 


 










Introduction

StyleRF-VolVis: Style Transfer of Neural Radiance Fields
for Expressive Volume Visualization
    
  Kaiyuan Tang0009-0001-3512-0112 and
  Chaoli Wang0000-0002-0859-3619

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



Since the early 1990s, volume visualization (VolVis) has been a central topic in scientific visualization research. One of the most popular techniques for VolVis is direct volume rendering (DVR). It works by casting rays from the image plane to the 3D volume, gathering samples along each ray, mapping them to visual quantities (i.e., colors and opacities) via transfer function (TF) lookups, and compositing the final color for each pixel in the rendering image. Thanks to the astonishing advances in graphics hardware, DVR can be efficiently implemented to achieve interactive framerates and high-quality visualizations, providing superior capability for users to explore volumetric datasets interactively. This practice has been the norm for the past two decades. 

The recent surge of deep learning for scientific visualization research <cit.> has sprouted new opportunities for VolVis. Leveraging the capabilities of generative adversarial networks (GANs), one can train a network to synthesize rendering results under novel viewpoints, TFs, or other parameters, eliminating the need to access the original volumetric data and bypassing the conventional rendering pipeline <cit.>. These seemingly impossible advances could have shocked many researchers just several years ago but are widely understood by the research community nowadays. After all, generative AI has swept across many fields, culminating in its extraordinary power to synthesize novel images and videos from text prompts and level the playing field for the masses.  

Inspired by this fantastic trend, in this paper, we aim to design a deep learning framework to accomplish expressive visualization via style transfer for VolVis. 
Similar to prior work <cit.>, we are supposed to be given a set of 2D images captured from the sample viewpoints and rendered with a fixed, reasonably good, yet unknown TF. Also, we assume no access to the 3D volumetric dataset during training or inference. 
We aspire to develop a new model that achieves expressive visualization by meeting the following goals (G1 to G3) for end users. 

First, they can freely explore the DVR scene in excellent visual quality from previously unseen viewpoints (G1).

Second, they can further flexibly recolor the rendering results photo-realistically by editing the underlying colors, opacities, and lighting effects (G2). 

Third, and most importantly, they can even intuitively modify the rendering results non-photorealistically by transferring styles (e.g., salient brushstrokes) from any image or painting of interest to the DVR scene (G3). 

Realizing G1 is relatively straightforward. 
Unlike GAN-based techniques, state-of-the-art solutions based on neural radiance field (NeRF) can deliver high-resolution, high-quality novel view synthesis results using a set of sample images. 
Nevertheless, significant challenges exist for achieving G2 and G3 due to the considerable gaps between the given 2D images and the 3D scene we aim to reconstruct and edit. 
For G2, the main challenge lies in accurately extracting the contributing color, opacity, and lighting information for faithful downstream edits to ensure visual content consistency. 
For G3, the grand challenge is to move beyond a fixed number of colors decoded from the original rendering images and adapt to abundantly rich color patterns or textures in the reference images. 

We introduce StyleRF-VolVis, achieving style transfer of NeRFs for expressive VolVis. 

For large volumetric data, NeRF-based scene representations are space efficient and thus have implications for altering renderings in a compressive manner. 
At its core, we highlight three key components of StyleRF-VolVis to realize the stated goals (G1 to G3). 

First, we employ a base NeRF model to learn the density representation (i.e., content) of the 3D volumetric data from the collection of 2D training images (G1). 

This component accurately extracts the underlying scene geometry, paving the way for successful subsequent style editing.

Second, we design a new palette color network that classifies regions of the radiance field (RF) and extracts a color palette (i.e., appearance) from multi-view images (G2). 

This component enables photorealistic style editing by modifying the original rendering's color, opacity, and lighting while maintaining visual content consistency across the views.

Third, to remove the color palette constraint, we propose a novel knowledge distillation solution that transfers the color information from the palette color network (i.e., teacher model) to an unrestricted color network (i.e., student model) (G3). 

This component allows users to assimilate a wide spectrum of color patterns and textures from any reference image, supporting expressive visualization of the DVR scene via non-photorealistic style editing. 

Figure <ref> highlights the capability of StyleRF-VolVis on two datasets, each using two styles extracted from different reference images. 
The contributions of our work are:



  * We revisit style transfer for VolVis by presenting an innovative NeRF-based solution that supports expressive photorealistic and non-photorealistic style editing.

  * StyleRF-VolVis represents a significant leap forward for visualization synthesis of volumetric data, advancing the state-of-the-art solutions in quality, consistency, and flexibility. 


  * We show the consistency and flexibility of StyleRF-VolVis on various combinations of DVR scenes and reference images.

  * We compare StyleRF-VolVis with other image-based (AdaIN), video-based (ReReVST), and NeRF-based (ARF and SNeRF) style rendering solutions to demonstrate its superior quality. 








§ RELATED WORK


This section discusses related work on deep learning for VolVis, style transfer, NeRF for stylization, and knowledge distillation.

Deep learning for VolVis.
Over the past few years, deep learning has emerged as a promising solution for improving the DVR process. 

One application is to use deep learning models to replace the traditional DVR pipeline. 
For instance, Berger et al. <cit.> constructed a GAN to synthesize volume-rendered images by investigating the image space of volume rendering under various TFs and viewing parameters. 
Hong et al. <cit.> proposed a GAN to synthesize high-resolution images from volume data under the desired rendering effect without knowing the TF. 
He et al. <cit.> developed InsituNet, a surrogate model that correlates simulation and visualization parameters with visualization outcomes, allowing users to preview visualization results under different simulation settings with a trained model. 
Shi et al. <cit.> built a view-dependent neural-network-latent-based surrogate model that supports producing high-resolution visualization results. 
Han and Wang <cit.> presented CoordNet, a coordinate-based neural network to synthesize rendering images under novel viewpoints given a set of DVR images for training. 

Another application employs a scene representation network (SRN) to represent volumetric data, minimizing disk storage requirements and enabling interactive neural rendering without direct access to volume data. 
For example, Weiss et al. <cit.> designed a dense-grid encoding method fV-SRN that directly renders from compressed representation with no additional memory for storing volume data during rendering. 
Wu et al. <cit.> leveraged multiresolution hash grid encoding and achieved fast volume encoding and real-time rendering. 
Wurster et al. <cit.> developed APMGSRN, an adaptively placed multi-grid SRN with a domain decomposition training approach for efficient VolVis. 

Other deep learning works in the context of VolVis focus on data or visualization generation <cit.>, compression <cit.>, translation <cit.>, and completion <cit.>. 

Our work represents the first step in applying style transfer techniques to volume rendering results using a NeRF model. Compared with existing works <cit.>, StyleRF-VolVis takes a smaller set of DVR images for training and infers higher-quality results, maintaining content consistency across different viewpoints and supporting flexible style transfer. 


Style transfer.
Ever since Gatys et al. <cit.> demonstrated the capability of convolutional neural networks (CNN) in applying diverse artistic styles to natural images, neural style transfer <cit.> has emerged as a trending topic in both academia and industry. Numerous techniques have been developed to improve or extend the original algorithm.

Johnson et al. <cit.> introduced a feed-forward network to solve the slow optimization problem and achieved real-time style transfer. 
Huang et al. <cit.> further maintained speed comparable to <cit.> without sacrificing the flexibility of transferring inputs to arbitrary new styles. This was realized using an adaptive instance normalization layer that aligns the statistics information of the content and style features. 
Ruder et al. <cit.> extended style transfer to video sequences by computing optical flow to achieve consistent style transfer across frames. 
Wang et al. <cit.> presented ReReVST that relaxes the objective function of style loss to make the transfer more robust to motions in content video. 


In VolVis, Bruckner and Gröller <cit.> introduced a style TF using the lit sphere, which assigns the optical properties of voxel values with rendering styles instead of simple colors and opacities. 
However, extracting or designing the lit sphere requires artist involvement, and the rendering relies on a traditional pipeline. Our StyleRF-VolVis adopts an advanced NeRF representation, supporting end-to-end style transfer and rendering without accessing the original volume. 

NeRF for stylization.
Since the groundbreaking work of Mildenhall et al. <cit.> in 2020, NeRF has been widely used for novel view synthesis. 

We refer readers to recent survey papers <cit.> to follow the roadmap of NeRF-based applications.
While extensive studies <cit.> emphasize the quality or speed enhancement of NeRF, there is also a growing body of work focusing on editing a base NeRF to perform 3D style transfer of a scene. 
Generally, editing NeRF for style transfer can be classified into photorealistic and non-photorealistic.

One main task for photorealistic editing is recoloring the scene. 
For example, Kuang et al. <cit.> and Gong et al. <cit.> utilized optimizable base colors in a palette to fit the scene and modify the base colors to recolor the scene during inference.
Lee et al. <cit.> modified the essential weights in the color multilayer perceptron (MLP) of a trained NeRF to perform local recoloring. 

Unlike photorealistic editing, non-photorealistic editing transfers more abstract information, such as textures or brushstrokes of an artistic work. 
For instance, Chiang et al. <cit.> utilized a hypernetwork to transfer style features into the NeRF representation. 
Huang et al. <cit.> proposed a mutual learning strategy to integrate 2D stylization with NeRF geometry consistency. 
Zhang et al. <cit.> designed a nearest neighbor-based loss to stylize a trained NeRF, which performs better stylization quality compared with standard Gram matrix-based loss. 
Liu et al. <cit.> transformed the grid features within RFs to match a reference style. Even though the quality is less impressive than <cit.>, it supports zero-shot transfer to an unseen reference image.

StyleRF-VolVis differs from all these works in that it supports photorealistic and non-photorealistic editing in a single framework. We will show that conducting photorealistic editing before non-photorealistic editing yields better quality and flexibility for stylization results.





Knowledge distillation.
Hinton et al. <cit.> introduced knowledge distillation (KD), which optimizes a small network to match the predictions of a large model. KD has been used in model compression and optimization <cit.>. 
In the NeRF space, Reiser et al. <cit.> leveraged a large global pre-trained NeRF to speed up the training of multiple small local NeRFs. 
Barron et al. <cit.> improved the rendering quality with an online distillation strategy. 
Wang et al. <cit.> designed R2L to learn the light fields of a pre-trained NeRF model effectively. 
Fang et al. <cit.> proposed progressive volume distillation, which provides a systematic KD method that allows the distillation between explicit and implicit NeRF architectures. 

Instead of applying KD to optimize a small network for model compression or optimization purposes, we utilize KD to transfer the knowledge of a photo-realistically edited RF to improve the quality of later non-photorealistic editing.









§ STYLERF-VOLVIS





The objective of StyleRF-VolVis is to edit the style (i.e., appearance) of the DVR scene represented by a NeRF model without losing the original content (i.e., geometry) information. The style can be photorealistic (such as adjusting the original rendering's color, opacity, and lighting attributes) or non-photorealistic (such as altering the texture to resemble the brushstroke of an artistic image). 
Distinct from conventional 3D neural stylization methods, our StyleRF-VolVis aims to unify two types of style editing into a single framework. 
Additionally, from a visualization perspective, users design TFs to classify the regions of interest in the DVR process. 
These TFs map distinct colors and opacities to each region users specify, highlighting relevant data ranges through visual mapping. 
Such visual content is critical in obtaining helpful visualization. Our method must keep the original visual content consistent across the views to ensure the stylized visualization remains meaningful, facilitating accurate interpretation and analysis.





Figure <ref> illustrates the three-stage workflow of StyleRF-VolVis: 
(a) a base NeRF model for optimizing the density representation of the scene, 

(b) a palette color network (PCN) for supporting photorealistic style editing (PSE) and maintaining visual content consistency during style editing, and 
(c) an unrestricted color network (UCN) for enabling non-photorealistic style editing (NPSE). 
Accordingly, our method is structured into three optimization stages, each dedicated to training the parameters of one network while others are frozen for stability.

The first stage focuses on training a base NeRF model to reconstruct the scene geometry for subsequent style editing. 
The parameters of the trained base NeRF are then frozen to ensure that the following style editing does not affect the underlying scene geometry.

The second stage optimizes the PCN. 
Similar to TFs that differentiate volumetric regions with user-defined colors and opacities, the goal here is to classify regions of the RF with 

a color palette extracted from multi-view images. 
To this end, the PCN is optimized to fit input images with a weighted linear combination of extracted palette colors, given the density representation obtained from the previous stage.


By learning to represent the colors of the RF, our model achieves PSE, including modifying the original rendering's color, opacity, and lighting.


The number of palette colors inherently restricts the network's ability to utilize more diverse color patterns for NPSE, where each style often consists of multiple colors. 

To address this issue, 
in the third stage, we start with stylizing the palette colors to match the target style colors. 
Leveraging KD, we then transfer the color information from the PCN (i.e., teacher model) to a UCN (i.e., student model), representing the colors of the RF without the color palette constraint. 
In addition, we can apply different styles to different visual regions extracted from the PCN.  
As a result, we stylize the visualization scene with NPSE while maintaining the original visual content.

The rest of Section <ref> is structured as follows. In Section <ref>, we provide the preliminary knowledge of NeRF representation <cit.> and how we optimize the base NeRF model. In Sections <ref> and <ref>, we discuss how we adapt NeRF's palette-based recoloring strategy <cit.> to fit PCN and achieve PSE. The utilization of KD <cit.> to extract UCN was described in Section <ref>, followed by the NPSE details that combine stylization loss <cit.> and unsupervised segmentation techniques <cit.> in Section <ref>. Finally, we brief our interactive interface design in Section <ref>.





 §.§ Base NeRF Model
 


StyleRF-VolVis is built upon the scene representation from a base NeRF model. 
In general, a NeRF model <cit.> represents the scene with two neural functions: a density function σ(x) that maps any 3D position x to a density value σ and a color function c(x, d) that outputs RGB color c given an arbitrary 3D position x and viewing direction d. 
NeRF samples rays from the camera origin o to the rendering pixels for each training camera view d. For one sample ray r=(o,d), if we sample M points along ray r with 3D positions x_1… M at depths t_1… M. The predicted pixel color ĉ (r) of ray r is computed as

    ĉ(r)=∑^M_i=1α_i(1-ω_i)c(x_i, d),

where ω_i = exp(-(t_i-t_i-1)σ(x_i)) represents the transmittance along the ray between 
x_i and x_i-1. 

α_i=∏^i-1_j=1ω_i denotes the attenuation of the ray from its origin o to 
x_i.

We select the advanced architecture from Instant-NGP <cit.> to optimize the RF, ensuring fast convergence and efficient rendering. Specifically, Instant-NGP comprises a multiresolution hash-grid encoding and a tiny MLP for efficient optimization.
During the optimization of the base NeRF, we optimize with a loss ℒ_ defined as

    ℒ_ = ||c(r) - ĉ(r)||^2_2,

where c(r) and ĉ(r) are the ground-truth (GT) and predicted pixel colors corresponding to ray r.
After optimizing the base NeRF, we freeze its model parameters in the subsequent stages to avoid style editing influencing the scene geometry. Note that the colors of the base NeRF will not be used in rendering the stylized RF; we employ the PCN for PSE and the UCN for NPSE. 









 §.§ Palette Color Network (PCN)



After training the base NeRF model for accurate density representation, we optimize a PCN that represents the scene appearance with 
colors from a palette. 

The trained PCN can classify any 3D position within the RF with its palette color. 
The classification of the PCN preserves region classification from the user-defined TF, ensuring visual content consistency of the downstream style editing.
Figure <ref> shows the PCN architecture. 
The network is composed of two branches: 
one branch uses the lighting MLP to predict the view-dependent specular color c_s and 
the other branch uses the palette MLP to output the view-independent diffuse color c_d. 

For the specular color branch, we use a structure similar to the Instant-NGP's color function. The lighting MLP also ingests geometry features output from the density MLP and spherical harmonics (SH) encoded view direction as input. The only difference is that the output of the lighting MLP is a grayscale color instead of the original RGB color. Before optimization, we leverage the hidden parameters of the base NeRF's color function to initialize the shallow layers of the palette MLP to improve the prediction of c_s at the early stage.





For the diffuse color branch, we design a palette MLP that captures the weights ω of palette colors P for any points within the RF.
This way, we can use ω to classify points in the RF according to their primary colors, preserving visual content consistency across different views during style editing.

Similar to the recent NeRF's palette-based recoloring strategy <cit.>, we apply the RGB convex hull method <cit.> to extract colors P_ from training images. However, unlike previous palette recoloring methods that directly employ P_ to initialize the palette, we further refine P_ as it usually contains many similar colors that may significantly hinder the classification. We remove similar colors when the distance between two colors is below a certain threshold. This distance is measured by the hue component of colors in the hue-saturation-brightness (HSB) space. The saturation and brightness components of palette colors are not considered in measuring the distance as they can be easily optimized during training. Figure <ref> gives an example that compares classification and rendering results without and with palette refinement. Refer to the appendix for more details about the refinement algorithm.

Once the palette colors are initialized and the number of palette colors N_P is determined, we construct and optimize the PCN. During the optimization, the density MLP and hash-grid encoder E_σ learned in the first stage are fixed for invariant density representation. The palette MLP takes features output from the palette hash-grid encoder E_p as input and outputs three values: 
an intensity value I shared by all palette colors to counteract the range shift caused by the RGB convex hull color normalization, 
a set of weights ω that indicate the contribution of each palette color to c_d, and 
a color offset vector δ to enhance the expressiveness of palette colors. 
Given the refined palette colors P and specular color c_s, the output color c associated with position x and view direction d is computed as

    c(x,d) = c_s(x,d)+I(x)∑_i=1^N_Pω_i(x)(P_i+δ_i(x)),

where I(x) is the intensity value at position x, ω_i(x) and δ_i(x) denote the weight and color offset for palette color P_i at x, respectively.
We then predict pixel color ĉ(r) using Equation <ref> for c(x,d), and optimize network parameters and palette colors with loss ℒ_ defined as

    ℒ_=||c(r) - ĉ(r)||^2_2 + λ_δ∑_i=1^M∑_j=1^N_P||δ_j(x_i)||^2_2.

The first term of ℒ_ is the MSE loss between the GT pixel color c(r) and predicted pixel color ĉ(r). 
The second term is the offset regularization loss for sample points along ray r, aiming to suppress the magnitudes of color offsets and avoid significant palette color shiftings. We set λ_δ = 0.1 to control the regularization strength. 





























 §.§ Photorealistic Style Editing (PSE)



When inferring the PCN, the output palette weights ω(x) can classify any 3D position within the RF according to its palette color. We then tune the network values to support PSE. Interactive PSE is achieved using Instant-NGP's fast inference ability.

We use the HSB space for recoloring following <cit.>. Given the target palette colors P', we compute the difference ΔP between the original palette colors P and P' in the HSB space. Δ P_i are then added 
to P_i+δ_i(x)

in Equation <ref>. 
To achieve opacity or lighting editing, we multiply σ(x_P_i) or c_s(x_P_i, d) by a scalar value, where x_P_i denotes a point position receiving its primary color contribution from P_i. 
















 §.§ Unrestricted Color Network (UCN)



Although the PCN can represent the colors of the RF with PSE, it prohibits using diverse colors in NPSE. 
To address this issue, we utilize KD to optimize a UCN (student model) that produces similar diffuse colors within the PCN (teacher model) with no color palette constraint. 

Before distillation, we initially extract the average palette colors P̅ from each reference style to replace the original palette colors P.
This color transfer step ensures the distilled student network can match with styles, reducing optimization effort in NPSE. 
Our UCN comprises one hash-grid encoder E_u following an unrestricted MLP that outputs color c_u. 
During distillation, we first randomly sample camera views from the scene. Then we optimize E_u with volume-aligned loss ℒ_ adopted from <cit.> for each view. ℒ_ is defined as

    ℒ_ = ||E_p(x) - E_u(x)||^2_2,

where E_p(x) and E_u(x) are the output features from the palette and unrestricted encoders for each sample point x. We optimize ℒ_ for 150 iterations to initialize E_u(x) to accelerate the subsequent distillation. For the following 500 iterations, we optimize both E_u(x) and unrestricted MLP with color loss ℒ_, which is defined as

    ℒ_ = ||c̅_d(x) - c_u(x)||^2_2,

where c̅_d(x) and c_u(x) are the diffuse colors output from the PCN with the average palette colors P̅ and the UCN at x.  





 §.§ Non-Photorealistic Style Editing (NPSE)






Given one or multiple reference images, we first extract N_P desired reference styles S={S_1, S_2, ⋯, S_N_P} for N_P regions of the RF. 
When the number of reference images is less than N_P, or users are only interested in local patterns of one image, 
we can leverage unsupervised segmentation techniques such as the segment anything model (SAM) <cit.> to automatically or manually (with point prompt selection) divide a reference image into localized style regions for flexible stylization.

Figure <ref> shows our training process for non-photorealistic style optimization. We use the optimizable unrestricted color c_u(x) (Section <ref>) and frozen density σ(x) of the base NeRF (Section <ref>) to construct a new RF. 

During training, we first compute N_P rendering images R={R_1, R_2, ⋯, R_N_P} from the new RF for each camera view. 
For one sample point x, it only contributes to R_i if its palette weight ω_i(x) is greater than other palette weights.
Once we have N_P renderings R and styles S, a pre-trained VGG-16 model <cit.> computes a stylization loss based on a user-specified or randomly assigned style mapping. Specifically, we use the nearest neighbor feature matching (NNFM) loss <cit.> for each rendering and style, which is formulated as

    ℒ_(R,S)=1/(N_FN_P)∑_i=1^N_P∑_f_j∈ϕ(R_i)min_f_k∈ϕ(S_i)d(f_j,f_k),

where ϕ(R_i) and ϕ(S_i) are the rendering and style feature vectors extracted from the VGG-16 model ϕ. N_F is the number of feature vectors for ϕ(R_i), and d is the cosine distance between feature vectors f_j and f_k.





Unlike previous NeRF-based non-photorealistic stylization methods <cit.>, our style transfer process involves three unique designs for the VolVis scenario. 
First, different from traditional NeRF-based stylization that directly discards the view-dependent color to maintain visual content consistency across the views, we optimize c_u(x) (i.e., diffuse color) during training and output c_u(x)+c_s(x,d) (i.e., diffuse and specular colors) during inference. The rationale behind this design is to prevent the view-dependent color from influencing visual content consistency and preserve the specular lighting effect in the final output. 
Second, when rendering R in training, we observe that varying background colors could impact the stylization results (refer to an example shown in Figure <ref>). This phenomenon is caused by the inherent transparency in the DVR scene, which is not often present in natural scenes. Consequently, different background colors could lead to distinct renderings using the identical network, potentially incurring unsatisfactory stylization. 

To address this issue, we extract the luminance value from each style S_i as the corresponding background color for each rendering R_i.
Third, we omit the content loss between the original and stylized images due to suboptimal stylization outcomes.













 §.§ Interactive interface


Figure <ref> shows our visual interface for users to conduct PSE and NPSE for a given DVR scene. 
For PSE, users select a palette ID and interactively modify the corresponding scene region's color, opacity, and lighting. 
For NPSE, users first specify the target style for each scene region from single or multiple reference images. 
They then start optimizating UCN. 
After NPSE, users can still change the opacity and lighting of the stylized DVR scene but cannot modify the color as the stylized scene is represented with UCN. 
Refer to the accompanying video for the recorded interaction with the interface.














§ RESULTS AND DISCUSSION






 §.§ Datasets, Training, Baselines, and Metrics


Datasets and network training.
To show the quality, consistency, and flexibility of StyleRF-VolVis, we evaluate it using the datasets listed in Table <ref>. 
The visible ranges are the opacity bumps in the TF specification corresponding to distinct visual contents for generating the original rendering. 
All reference images used are copyright-free, coming from WikiArt <cit.> and Pexels <cit.>. 
We implemented StyleRF-VolVis using PyTorch and ran experiments on an NVIDIA RTX A4000 graphics card with 16 GB memory. 
We trained the base NeRF model with 30,000 iterations and PCN with 2,500 iterations using 92 DVR images with cameras evenly placed along a sphere enclosing the volume. 
The KD of UCN and NPSE does not need GT images. 
For KD, we optimized UCN with 624 iterations, where each iteration trains on one of the randomly sampled 312 camera views. 
For NPSE, we optimized UCN using the NNFM loss with 210 iterations, where each iteration trains on one of the uniformly sampled 42 camera views. 
For all three stages (base NeRF, PCN, and UCN), we applied the Adam optimizer with a learning rate of 0.01 and β_1=0.9, β_2=0.999 for training and set the batch size to 4,096 rays.
For the base NeRF model, PCN, and UCN during KD, we decayed the learning rate exponentially for every iteration until it reached 0. 
During NPSE, we fixed the learning rate and followed the setting of ARF <cit.>. 
We used the conv3 block of VGG-16 to extract features for computing the NNFM loss. 
Deferred backpropagation was employed to update UCN's parameters for the entire image under limited GPU memory. 
During inference, we generated 181 stylized images with novel views, with cameras evenly placed along a trajectory to capture the full 360-degree view. 


Baselines.
We compare four baseline methods for NPSE:



	
  * ARF <cit.> is a NeRF-based method that trains a NeRF model to represent the scene and then updates the pre-trained NeRF parameters with NNFM loss to generate stylized novel view images. 
	
  * SNeRF <cit.> is a NeRF-based method that first constructs the scene with a NeRF model and then alternatively finetunes the parameters of the pre-trained NeRF and an additional image stylization module to achieve stylization. 
	
	
  * ReReVST <cit.> is a video-based stylization method. It first performs novel view synthesis with a NeRF model, then arranges the novel view images as a video sequence, and finally performs video stylization to yield the result. 
	
  * AdaIN <cit.> is an image-based stylization method that optimizes a NeRF model for generating novel view images of the scene and then applies stylization to the NeRF synthesized images.




Note that we did not compare StyleRF-VolVis with the more recent StyleRF <cit.> work as it mainly focuses on the ability of zero-shot transfer to unseen reference images. 

StyleRF trains on all reference images from the WikiArt dataset to ensure better generalization for unseen reference images. However, this could lead to inferior stylization results compared to optimization methods like ARF and SNeRF.  
Existing visualization generation methods in scientific visualization <cit.> were also excluded because they were not designed for style transfer. 

Evalutation metrics.
Since 3D style transfer is relatively novel and scarcely explored, limited metrics are available for quantitatively assessing stylization quality. 
Consequently, our evaluation focuses on cross-view consistency of stylization results. 
Specifically, we generated testing videos in which each frame represents a rendered image from one of the 181 novel viewpoints of our stylized scene. Then, we warped one view to the other using softmax splatting <cit.> according to the optical flow estimated with RAFT <cit.>. Finally, we computed the MSE and LPIPS <cit.> scores to measure the cross-view consistency. 
Similar to <cit.>, we calculated the long-range consistency for faraway views (with an interval of 10 among the 181 inference views) and short-range consistency for adjacent views, respectively.





 §.§ Qualitative Comparison


We compare StyleRF-VolVis against baseline methods using DVR scenes of vortex, supernova, combustion, and five jets. Each scene showcases two stylization results using distinct reference images. Figure <ref> shows the qualitative comparison. Unlike all baseline methods, StyleRF-VolVis supports customized style selection and mapping to distinct visual contents, which we highlight in the DVR scene.

Overall, StyleRF-VolVis produces clearer and more consistent stylization results than other methods. 
In contrast to video- and image-based stylization methods (ReReVST and AdaIN), NeRF-based stylization methods (StyleRF-VolVis, ARF, and SNeRF) preserve better geometry consistency and are capable of separating foreground objects from the background during stylization. 
Among NeRF-based methods, SNeRF captures the overall style of the reference image but cannot assign distinct styles to different visual regions of the original scene. Thus, SNeRF's stylization is blurry and mixed, making it hard to identify the difference between individual visual contents in the original DVR scene.

For the vortex dataset, ARF shows better stylization than SNeRF in preserving visual content distinction. 

However, the uncontrollable ARF style selection process could let the model choose a negligible local style that does not match the overall style (e.g., the WikiArt-1 case) or select one style for multiple visual regions (e.g., the WikiArt-8 case). Both drawbacks could lead to poor stylization results.
In addition, in the stylization process, ARF and SNeRF discard view directions in their input to ensure cross-view stylization consistency. Consequently, they cannot preserve the view-dependent lighting of the DVR scene.

Unlike ARF and SNeRF, StyleRF-VolVis utilizes the lighting MLP to preserve lighting information of visual content while ensuring cross-view consistency during NPSE by removing view direction input in UCN optimization. 
Moreover, StyleRF-VolVis allows explicit assigning of different styles to different DVR regions. Via color transfer, we ensure that the region color always aligns with the overall style color, avoiding the disadvantage of NNFM loss, which focuses on negligible local style. 































 §.§ Quantitative Comparison


Training and inference time. 
We report the training and inference time for all NeRF-based methods in Table <ref>.
The training of StyleRF-VolVis is slower than ARF but faster than SNeRF.
Unlike ARF, StyleRF-VolVis requires additional steps to optimize PCN and apply KD to UCN before NPSE. 
Moreover, optimizing on a single camera view during NPSE requires rendering different regions independently. 
This process maps a style to a region, slowing the stylization process. 
Although the alternative training process of SNeRF is relatively straightforward, it demands significant time to train the image stylization module, leading to the longest training time. 
For inference, StyleRF-VolVis needs additional feedforward steps of the lighting MLP to provide lighting information, which takes longer than ARF or SNeRF. 







Short- and long-range consistency.
In Table <ref>, we compare short- and long-range cross-view consistency for different methods. 
We use MSE and LPIPS as the metrics. 
For each metric, we report the average value over eight stylization cases (WikiArt-1 to WikiArt-8) shown in Figure <ref>. 
StyleRF-VolVis achieves the best cross-view consistency, while AdaIN performs the worst as this image-based stylization method treats each image independently. 
ReReVST leads to comparable or even better MSEs than the NeRF-based methods because it produces stylization with relatively uniform colors, resulting in smaller pixel-wise errors. 
However, under the image-level LPIPS metric, the consistency of ReReVST significantly lags behind the NeRF-based methods.





Voting results from a user study.
As the consistency metrics do not necessarily reflect the perceived style transfer quality, we conducted a user study to measure user preference for different stylizations following the University's IRB protocol.

Per Figure <ref>, we included four datasets, each with two reference images, and considered three leading methods (ARF, SNeRF, and StyleRF-VolVis) for pairwise comparison. This leads to 24 (4×2×3) image pairs organized into eight (4×2) groups. 

We recruited 14 students from a visualization class of undergraduate, Master's, and Ph.D. students in computer science and engineering, aerospace and mechanical engineering, and psychology majors. 

The participants were briefed on the evaluation criteria before proceeding to the study. 
A full-screen display shows each pair. At the top of each display, the original DVR and reference images (with extracted style regions highlighted in different color boundaries) are shown. At the bottom, the stylized images of two methods, randomly placed on the left and right sides and labeled `A' and `B,' are presented. 
The participants were asked to decide which one (`A' or `B') achieved the better stylization outcome, and no tie was allowed.

We asked the participants to take their time, as we did not record how much time they spent on each pair, each group, or the entire study. During the evaluation, they could go back and forth to update their votes. 
We advised them that many factors should be considered, including the overall impression, content preservation, style application, and visuals (color, opacity, and lighting). 
They could decide how to weigh them, and their criteria should be consistent throughout the study.

The entire study was completed in the classroom within 10 minutes. 

The voting results are shown in Table <ref>. 
We can see that StyleRF-VolVis wins for all stylization cases except WikiArt-1, WikiArt-3, and WikiArt-8. This may be attributed to the fact that some participants prefer 
the stylization result of ARF that provides a detailed local pattern for the WikiArt-8 case or 
the stylization results of SNeRF, which contain a more global pattern in WikiArt-1 (i.e., the dark tone and contrasting colors of The Starry Night) and WikiArt-3 cases. However, StyleRF-VolVis wins for all other cases and gains the most overall preference. The cross-view consistency of StyleRF-VolVis is also better than ARF and SNeRF, as shown in Table <ref>.









 §.§ Flexibility of NPSE



Our NPSE strategy allows users to define which style should be applied to which part of a DVR scene (analogous to using the TF), which provides more flexibility than other NPSE methods. Specifically, users can apply styles obtained from one reference image for multiple stylizations, leverage styles extracted from multiple reference images for one stylization, or only stylize one part of the DVR scene. All these additions are the unique features of StyleRF-VolVis.

One reference image for multiple stylizations.
Unlike conventional NeRF-based stylization methods, which only produce one stylized scene from a single reference image, StyleRF-VolVis can generate various stylized outcomes with only one reference image.
Thanks to the advanced segmentation of PCN and SAM for the RF and reference image, we can extract multiple distinct styles and apply them to different scene regions.
By pairing styles and regions in various combinations, StyleRF-VolVis achieves controllable and diverse stylization outcomes.
Figure <ref> shows multiple stylization results of the solar plume DVR scene using the same reference image.
By comparing Pair-1 and Pair-2, we can see that when the style of one region remains unchanged, altering the style of another region keeps the unchanged style region intact. 
Further comparison of Pair-2 and Pair-3 reveals that applying the same style to different regions consistently transfers the expected texture.
These favorable features allow StyleRF-VolVis to offer users flexible stylization.
 




Multiple reference images for one stylization.
Previous approaches like StyleRF <cit.> suggest leveraging extra supervision from large models pre-trained on extensive image datasets for segmenting scene elements to support style transfer from multiple reference images to a single scene. 

However, such a strategy cannot be directly applied to VolVis scenes. 
One reason lies in the inherent difference between DVR and natural scenes. 
DVR images exhibit more complex geometric relationships, with different visual contents often nested in layers with varied transparencies, rendering 2D segmentation methods futile. 
Another reason is that large models are typically pre-trained on natural images. 
Adapting such models to process DVR images would require extensive DVR images and substantial hardware resources for fine-tuning, a demand that is impractical for user-oriented VolVis scene stylization.

Our model can segment various regions via PCN, even though separate GT rendering results for different regions are missing during optimization and additional supervision from a 2D segmentation method is lacking. 
As illustrated in Figure <ref>, StyleRF-VolVis achieves stylized results for different regions using styles from multiple reference images. 
With a collection of reference images, StyleRF-VolVis could create vast combinations of different stylization results for one DVR scene.





Partial stylization.
In some scenarios, users may wish to stylize only certain regions of the scene while keeping others unchanged.
Unlike previous NeRF-based stylization methods, StyleRF-VolVis achieves such partial stylization by combining the color representations of PCN and UCN. 
After completing the NPSE of each region, PCN can still utilize the original palette colors to represent the DVR scene. 

Users can thus specify which regions should maintain the NPSE texture and which should keep the original appearance. 
During rendering, regions that preserve the original scene's appearance utilize colors represented by PCN, while those stylized regions are processed through UCN.
Figure <ref> shows an example of partial stylization, providing extra freedom for users to achieve desirable results.









 §.§ PCN Classification and PSE


As discussed in Section <ref>, the PCN classification outcomes are crucial for style editing.
Visual contents in a DVR scene are often nested with each other. 
However, as demonstrated in Figure <ref>, even for a complex scene like the earthquake, PCN can still produce region classification closely resembling GT, even though only DVR images are used for training.
Although the colors and lighting of PCN classification results are not perfect, the boundary of each region is clear, and even fine details maintain consistency with GT. 
This ensures precise PSE and NPSE of each region in subsequent stages.

In Figure <ref>, we showcase the PSE outcomes for various DVR scenes. With the assistance of PCN classification, we can perform a distinct PSE in each region. Such flexibility further assists users in achieving desired stylization results.









 §.§ Ablation Study


Benefit of PSE to NPSE.
At the beginning of UCN optimization, we perform the color transfer 

in PCN (teacher network) to ensure that the color representation of UCN (student network) matches the selected styles before NPSE. 
This PSE step is essential for speeding up the convergence and improving the stylization quality of UCN during NPSE optimization. 
Figure <ref> compares the style transfer process without and with color transfer.
Comparing the stylization results of 42 and 168 iterations, we can see that NPSE without color transfer suffers from inaccurate style color matching at an early stage and shows unclear stylization (i.e., black borders, see arrows) at 168 iterations. In contrast, NPSE with color transfer matches the overall style color well at the early stage and reveals well-stylized texture at 168 iterations. 





View-dependent lighting in NPSE. 
Recent NeRF-based stylization methods <cit.> discard view directions in the input to ensure cross-view stylization consistency. 
However, they ignore that view-dependent lighting of the original scene is essential for maintaining consistency between the original content and the stylized scene.
In contrast, StyleRF-VolVis utilizes the additional lighting MLP optimized in the PCN training stage to preserve the DVR lighting during NPSE. 
In Figure <ref>, we show examples of vortex and supernova datasets to compare the effect of DVR lighting on NPSE outcomes. 
The results show that the stylized scene with lighting is more consistent with the original DVR scene than that without. 









 §.§ Limitations


Although StyleRF-VolVis achieves flexible, high-quality stylization of DVR scenes, it has the following limitations. 

First, given an initial TF, NeRF cannot represent the value ranges where opacity equals zero. Consequently, StyleRF-VolVis cannot perform PSE or NPSE for regions with zero opacity in the initial TF.  

Second, StyleRF-VolVis supports interactive PSE but not interactive NPSE.

Third, even though users can adjust the lighting intensity for the DVR scene, they cannot modify the lighting direction because no normal information is available.

Finally, 

if the colors among different regions within the input DVR scene are similar, color refinement may not separate these regions correctly. In such cases, users may need to manually adjust the number of palette colors and their RGB values to obtain the desired results.






§ CONCLUSIONS AND FUTURE WORK


We have presented StyleRF-VolVis, the first work in VolVis that targets style transfer in the NeRF space. 
The crux of our approach lies in the accurate extraction of content and appearance information separately from the given DVR scene and the bridging between photorealistic and non-photorealistic style editing via knowledge distillation. 
With these innovations, StyleRF-VolVis achieves high-quality, consistent, and flexible 3D style transfer outcomes with novel view synthesis. 
The efficacy of StyleRF-VolVis is demonstrated with various combinations of DVR scenes and reference images. 
Moreover, we compare StyleRF-VolVis against other image-based (AdaIN), video-based (ReReVST), and NeRF-based (ARF and SNeRF) solutions via objective and subjective evaluation to showcase its superior quality performance. 

In the future, we will extend StyleRF-VolVis to handle dynamic DVR scenes produced from time-varying datasets. 
The challenge is maintaining a consistent appearance over timesteps to achieve temporally coherent stylization. 
We will also explore StyleRF-VolVis for multivariate or ensemble VolVis, where different variables or ensembles could be mapped to visually distinct styles beyond colors for better differentiation. 
It remains to be seen what the appropriate number of styles and their mixing should be to leverage the human's visual capacity best while maintaining observation clarity.

To ease the difficulty for non-professionals using StyleRF-VolVis, we will explore integrating natural language interaction into the current graphical user interface and broaden the selection of reference images from the WikiArt collection to images created by generative AI. 

The success of StyleRF-VolVis will unfold exciting opportunities for VolVis beyond the originated scientific domain. We envision that such a solution will enable citizen science by fusing diverse disciplines, such as science and art, for the general public's exploration, understanding, and appreciation, which we would like to pursue. 





This research was supported in part by the U.S. National Science Foundation through grants IIS-1955395, IIS-2101696, OAC-2104158, and IIS-2401144, and the U.S. Department of Energy through grant DE-SC0023145. The authors would like to thank the anonymous reviewers for their insightful comments.



abbrv-doi-hyperref











§ APPENDIX








§ IMPLEMENTATION DETAILS


Palette refinement algorithm.
When initializing palette colors P before PCN training, we refine P_ extracted from the RGB convex hull method

to eliminate similar colors within the palette. 

Specifically, for any two colors in P_, if they are similar (i.e., their L1 distance of hue value in the HSB space is below a threshold T_h), we remove one of them. 
Furthermore, there is a small chance that P_ may contain gray colors corresponding to the light color in the DVR scene. 
Such gray colors are unnecessary in P as the lighting MLP has represented lighting components. 
Therefore, we remove any color in P_ if its brightness value in the HSB space is less than a threshold T_b. 
For normalized color component values in the HSB space, we empirically set T_h=0.1 and T_b=0.2 to obtain refined P.









































Luminance background.
Before NPSE, we calculate each style's luminance value L to compute the NNFM loss using VGG-16.
Following the calculation steps given in <cit.>, for each average RGB color of the selected style, we normalize each component, such as R, and convert it into the linear-scale counterpart R_

    R_ =  R^2.2.

L of the selected style can be computed as

    L = 0.2126× R_ + 0.7152× G_ + 0.0722× B_,

where the RGB component coefficients reflect the average spectral sensitivity of lighting perceived by humans. 
We use L as the corresponding background color for each style rendering. 
This way, the stylization of StyleRF-VolVis can better match the overall brightness of the selected style.










§ ADDITIONAL RESULTS


PSE after NPSE.
After NPSE, users can still apply further PSE to the stylized scene. 
When doing so, one limitation is that if a region of the NPSE scene utilizes the UCN to represent the color term, users cannot apply a color PSE to the region as the PCN does not represent the color.
However, users can modify the opacity and lighting of each region without restriction.
Figure <ref> shows examples of applying various opacity or lighting PSEs to the scene after NPSE.













Choice of λ_δ.
When optimizing the PCN to avoid palette color shiftings, we include an offset regularization loss and use λ_δ to control the regularization strength. 
We conduct an ablation study to investigate the effect of λ_δ on PCN performance.
After optimizing the PCN with different λ_δ, we compare PCN rendering results with GT using PSNR, LPIPS, and SSIM, as shown in Table <ref>. 
Figure <ref> presents the rendering results. 

We can see that the offset regularization loss is essential for PCN training. 
When offset regularization is missing (λ_δ=0), the PCN does not predict correctly, as it tries to focus more on leveraging offsets instead of palette colors to represent the DVR scene.
However, the PCN may leverage more on palette colors instead of offsets to represent colors when λ_δ gets larger, resulting in less accurate scene reconstruction. 
Therefore, we choose λ_δ = 0.1 for a good control of the regularization strength.

Iterative style refinement for NPSE.
For NPSE, StyleRF-VolVis does not support direct control over the stylization process. 
However, users can indirectly achieve their desired stylization by iterative refining the selection in the reference image.
Figure <ref> shows such an example.
Based on the style selection in (a), we apply several negative point prompts for SAM in (b) to exclude undesired styles and retrain the UCN to achieve the desired stylization.





























Comparison with NPR.
Conventional non-photorealistic rendering (NPR) of a DVR scene relies on predefined generation rules to mimic a reference style. 
Due to the complexity of hyperparameter settings, the effectiveness of these methods depends on the choice of hyperparameters to some extent.
Moreover, such an approach limits the stylization results to a single style.
In contrast, StyleRF-VolVis can transfer arbitrary styles according to different reference images within a unified framework, offering higher flexibility and robustness.
In Figure <ref>, we compare the stylization results of StyleRF-VolVis and a conventional NPR method for stipple drawing style <cit.>.
We use the stipple drawing in <cit.> as a reference image and adjust the TF and camera pose to align the viewpoint for comparison.
StyleRF-VolVis matches the overall texture of the reference style more closely.





Comparison with DVR.
In DVR, users can adjust the TF to explore the scene. 
StyleRF-VolVis can also achieve similar objectives through the PSE of PCN. 
However, PSE cannot adjust invisible parts of the scene constructed via training the multiview images. 
Despite this limitation, PCN supports random access to any position within the scene during the rendering's sampling process. 
Compared to DVR, PCN may achieve faster render speed and require smaller memory footprints for large volumes under the same TF.
In Table <ref>, we compare DVR and PCN regarding the average rendering time per image, CPU memory, and GPU memory requirement for the rotstrat dataset with different volume and image resolutions (denoted by experiment IDs). 
We run DVR using the open-source software ParaView with NVIDIA IndeX plugins.
All models converge around four minutes and occupy a storage of 168 MB.
PSNR and LPIPS values reported in Table <ref> and the difference images (with respect to the GT rendering image) shown in Figure <ref> suggest that PCN achieves acceptable accuracy but requires a smaller rendering memory footprint and faster rendering speed compared to DVR. 
Note that the NeRF-based representation is independent of the volume resolution.  
Storing the training images and the network model for large data is more space-efficient than the original volume. 
Therefore, StyleRF-VolVis provides an efficient means for altering renderings of large volumes.




 Paper Title: Improving Faithfulness of Large Language Models in Summarization via Sliding Generation and Self-Consistency.







§ INTRODUCTION


Abstractive summarization aims to generate summaries that are fluent, informative, and faithful to the source articles. Benefiting from the popularity and development of large language models (LLMs), abstractive summarization has achieved remarkable progress in fluency and coherence <cit.>. However, LLMs have propensity to generate content that contradicts or is not present in the source article <cit.>, which is commonly referred to as hallucination. Alleviating the hallucination of LLMs is a critical challenge for their reliability in real-world applications.











Many works <cit.> explore the mechanisms why LLMs exhibit hallucinations. 
<cit.> observe that the performance of LLMs significantly decreases as the length of input contexts increases, resulting in the hallucination phenomenon of LLMs being particularly serious in long document summarization. Furthermore, LLMs are sensitive to the order of context and are more likely to select information presented first or last, even for short contexts <cit.>. That is, the summaries generated by LLMs contain more content that occurs at the beginning and end of the source article, which has a detrimental effect on the summary quality of the entire article.



Recent works attempt to leverage post-processing models and the CoT (Chain of Thought) technique to improve factual consistency of summarization. Some works <cit.> employ LLMs as critic and editor models, the critic model generates editing instructions for the initial summary, and the editor model corrects factual errors by following the instructions. <cit.> propose a CoT-based method called SumCoT to elicit LLMs to generate summaries step by step. However, they do not mitigate position bias and performance degradation in long context scenarios. 

In this paper, we propose SliSum, a novel summary generation strategy that improves the faithfulness of LLMs in both short and long text summarization by sliding generation and self-consistency. As shown in Figure <ref>, our approach consists of the following three steps: (1) Sliding Generation: SliSum uses LLM to generate local summaries for overlapping windows, and local summaries contain multiple statements about the same event that may have contradictions (see Table <ref>).  
(2) Filtration: According to the principle of self-consistency, the statements generated more times by LLMs are more faithful and important to the source article <cit.>. 
Therefore, we obtain statements about the same event by lexical clustering over sentences of all local summaries and then filter out small sentence clusters and outliers so that the global summary contains only relatively more important information of the source article. (3) Aggregation: SliSum employs LLM to detect whether there are contradictions in statements about the same event and classify them into different categories based on their semantics. Finally, SliSum uses a majority voting algorithm to select the statements with the most proportion, and concatenates them to form a more faithful summary of the entire article. SliSum brings three major benefits: (1) Sliding window provides LLMs with more diverse and adequate information by splitting articles in an overlapping manner; (2) The filtration and aggregation based on self-consistency ingeniously mitigate the self-contradiction problem and further leverage the potential of LLMs to improve the faithfulness of them; (3) The combination of sliding windows and self-consistency impels LLMs to process the entire article more fairly and faithfully. Therefore, SliSum is capable of improving the faithfulness of LLMs in summarization without external resources and additional fine-tuning. 





We evaluate the effectiveness of SliSum applied to three advanced LLMs, LLaMA-2-13B <cit.>, Claude-2 <cit.> and GPT-3.5 <cit.>, on four popular summarization datasets. Extensive experiments have shown that SliSum significantly improves the faithfulness of three LLMs on the short news datasets CNN/DM and XSum, and the long scientific papers datasets arXiv and PubMed, respectively, while without sacrificing their fluency and informativeness. Besides, we conduct ablation studies to further investigate why SliSum works. We also perform quantitative investigations on hyperparameters in SliSum. Our contributions are summarized as follows.


    
  * We propose a novel summary generation architecture, SliSum, that improves faithfulness of LLMs by sliding windows and self-consistency without additional resources and fine-tuning. To the best of our knowledge, we are the first to apply overlapping context windows into LLMs for abstractive summarization.
    
  * We demonstrate that SliSum uniformly improves factual consistency of summaries generated by diverse LLMs while maintaining their fluency and informativeness, more importantly, SliSum is applicable to text of various lengths and styles.
    
  * We conduct extensive qualitative and quantitative experiments to validate the effectiveness of sliding generation and aggregation based on self-consistency and impacts of hyperparameters in SliSum on performance.






§ RELATED WORKS




  
Factual Consistency of Summarization The factual consistency in abstractive summarization has received increasing attention recently. Existing work has proposed various methods to improve the factual consistency, such as contrastive learning <cit.>, adversarial learning <cit.>, textual entailment <cit.> and post-editing <cit.>. However, these methods can not be directly applied in long document summarization due to the difficulty of modeling long texts accurately. Although recent studies have addressed this problem by leveraging Graph Neural Networks <cit.>, reinforcement learning <cit.> and structure information <cit.>, they do not improve the faithfulness of long document summarization systems. There are currently few works that focus on improving factual consistency of long document summarization. In contrast, we propose a unified architecture that consistently improves factual consistency of both short and long text summarization.





  
Mitigation of LLM Hallucination Recent studies have made several attempts to mitigate the hallucination of LLMs, including retrieval-augmented generation <cit.>, post-processing models <cit.>, prompt engineering <cit.> and self-supervised learning <cit.>. However, most of these works require training additional models or external knowledge, or lack sufficient evaluation on long contexts tasks. For example, LLM-Augmenter <cit.> acquires evidence by retrieving external knowledge for the LLM to generate candidate responses grounded in evidence. 
<cit.> propose a post-processing method that improves factuality of language models using multiagent debate. SELF-FAMILIARITY <cit.> has revealed excellent performance in short context tasks, but its advancement has not been demonstrated on long-context datasets. Different from previous works, SliSum improves faithfulness of LLMs in long context scenarios without additional fine-tuning and external resources.



  
Long Context for LLMs With the application fields of LLMs continue to expand, boosting the performance of LLMs in long context scenarios has aroused a surge of interest. <cit.> propose a decoding method, LM-Infinite, to maintain fluency and generation quality of LLMs on long sequences, which can only be applied to LLMs using relative positional encodings, but our approach is applicable to diverse LLMs, even black-box models, because it does not modify the original structure and implementation of the models. <cit.> and <cit.> compress long context by removing redundancy to enhance LLMs' perception of the key information, but there is no guarantee that compressed contexts contain faithful information to original contexts. 
<cit.> propose a parallel context windows method that LLMs process contexts located in these windows separately to mitigate the negative impact of long contexts on performance. However, parallel windows separate the semantic relationships within the article, which inevitably affects LLMs’ understanding and extraction of global information. 
We are the first to introduce overlapping sliding windows into LLMs for abstractive summarization tasks. Unlike traditional parallel windows, sliding windows overlap each other so that each segment of the source article can be located at the beginning of a certain window, allowing them to overcome their preference for position and process the entire context more fairly and faithfully. 
In previous works, the sliding window method is only applied to extractive  <cit.> and query-focused <cit.> summarization, because they do not solve the self-contradiction problem <cit.> between local summaries. SliSum addresses this problem by combining sliding windows and self-consistency and achieves better performance than existing methods.



§ APPROACH




 §.§ Sliding Generation




  
Sliding Window SliSum first divide the source article A into a list of sentences [C_1, C_2, ⋯, C_n]. The sliding window with predefined window size initially consists of several consecutive sentences [C_1, ⋯, C_t] at the beginning of the article. The number of words in the window only needs to be approximately equal to the window size. Figure <ref> shows that LLM generates a local summary for the content in the sliding window. SliSum restricts the attention of LLM within a window, alleviating the distraction issue caused by long contexts <cit.>. 
Then the sliding window moves downward by a distance of step size and LLM generates a local summary for current window again. Similarly, the distance moved does not need to be strictly equal to step size. That is, the starting index of the window advances a few sentences. Importantly, SliSum processes the source article at the sentence level. Finally, LLM iteratively generates summaries for the content in the current window until the sliding window traverses the entire article. Given an article A = [C_1, C_2, ⋯, C_n], the process of sliding generation can be described as follows:


    S_1    = SUM (W_1) = SUM ([C_p_1, ⋯, C_q_1])
    S_2    = SUM (W_2) = SUM ([C_p_2, ⋯, C_q_2])
       ⋮
    S_m    = SUM (W_m) = SUM ([C_p_m, ⋯, C_q_m])


where SUM(W_i) represents LLM summarizes the content in the window W_i, W_i consists of from p_i-th to q_i-th sentences. p_1  = 1, q_m = n, and


    p_i ≤ p_i + 1≤ q_i,    i = 1, ⋯, m - 1


Each text segment except the beginning and end can locate in different positions in the context window. SliSum balances the position distribution in the context window of text segments, eliminating the position preference of LLMs for source article.



  
Process Repeatedly Eventually SliSum will obtain a concatenation of summaries [S_1, S_2, ⋯, S_m] by sliding generation. In the sliding window method, step size must be less than or equal to window size, otherwise there will be gaps between windows. Figure <ref> shows that LLM processes the same part of the source article multiple times due to overlaps between windows. Generally, given window size L_w, step size L_s (L_s ≤ L_w) and an article whose length L_a is far larger than the window size, i.e, L_a ≫ L_w, the minimum number of times the middle part of the article is summarized by LLM:


    K = ⌊L_w/L_s⌋


It can be easily seen that if L_w is an integer multiple of L_s in sliding window method, LLM will read the middle part of the source article the same number of times, so we set L_w = K L_s in SliSum. In particular, for several windows at the beginning and end, we repeatedly generate local summaries for them so that all parts of the source article are summarized K times. 



 §.§ Events Filtering




  
Lexical Clustering The event is usually composed of entity, behavior, reason, result and other elements, so statements about different events have their own characteristic words that can be used to distinguish them. We observe that the statements generated by LLMs for the same event are usually lexically similar. Consequently, we can obtain a set of sentences about the same event by lexical similarity clustering. As shown in Figure <ref>, SliSum divides all summaries into sentences, then uses DBSCAN <cit.> algorithm to cluster them based on lexical similarity. Specifically, we use ROUGE-1 <cit.> F1 score to define the distance between two sentences C_1 and C_2:


    dist(C_1, C_2) = 1 - ℛ(C_1, C_2)


where ℛ(·) is ROUGE-1 F1 score, and is computed as follow:


    ℛ(C_1, C_2) = 2 * number of overlap words/length of  C_1 + length of  C_2


Obviously, ℛ(C_1, C_2) ∈ [0, 1]. There are two important parameters in the DBSCAN algorithm: distance threshold ε and minimum number of points (MinPts) in a cluster. Formally, let Ĉ represents a sentence, if exists a set of sentences (including Ĉ itself) ℂ = {C_1, C_2, ⋯, C_N} and the number of sentences |ℂ| ≥MinPts, every sentence C ∈ℂ satisfies


    dist(Ĉ, C) ≤ε


ℂ will be considered a cluster. A lower MinPts induces the algorithm to construct more clusters, while a higher MinPts will ensure more robust and consistent clusters. 




  
Filtering Noise Intuitively, if an event appears in the local summary of a window, it means that the event is important to the content in the window. We notice that some sentences are only important in minority windows and should not appear in the overall summary. Furthermore, LLMs occasionally generate few hallucination statements that deviate from the source article. These statements tamper with information in the source article or add information that cannot be inferred from the source article, thus they generally have low lexical similarity to factually consistent statements. These undesirable statements are treated as outliers and noise in DBSCAN algorithm. Due to MinPts determines the minimum number of points in a cluster, so we can filter out unimportant and hallucination statements by using an appropriate MinPts. 

As mentioned in Section <ref>, SliSum generate K summaries for the content of the source article by sliding generation, that is, sentences about the same event occur at most K times. Under normal circumstances, a cluster contains at most K sentences. In order to generate more concise and relative summaries, we select sentences that exist in as least MinPts local summaries, i.e, we only retains clusters whose size exceeds MinPts. Therefore, we can adjust MinPts in the range of [1, K] to filter noise and hallucination as well as improve self-consistency.




 §.§ Contradictions Detection and Sentences Aggregation




  
Sentences Selection As shown in Table <ref> and Figure <ref>, there may be contradictory statements in sentences about the same event. Due to LLMs have ability to detect self-contradictory statements without relying on additional knowledge <cit.>, SliSum first manipulates the LLM used in the sliding generation to divide sentences into different categories according to their semantics. Generally, the more times a certain response is output by LLMs, the higher the probability that the response is desirable <cit.>, since self-consistency reflects factual consistency. Consequently, SliSum can utilize majority voting algorithm to select sentences that are faithful to the source article. 

After distinguishing sentences that are high lexically similar but semantically different, sentences of the same category state the same fact. SliSum selects the sentence with the largest proportion to output to the final summary. When there are clusters of sentences with the same proportion or when it is necessary to choose sentence from a certain  cluster, the relatively last generated sentence is selected by SliSum. Because the last generated sentence is located at the beginning of the corresponding window, and LLMs have the highest factual consistency for the information occurs at the beginning of the input context <cit.>. For example, an event has five related statements [D_1, ⋯, D_5] generated sequentially, and LLM divides them into [D_2], [D_1, D_4], [D_3, D_5], SliSum selects D_5 as the most faithful statement to form the summary of the entire article.



  
Sentences Integration SliSum finally integrates all selected sentences to output the final summary of the entire article. SliSum uses LLM to generate connectives and concatenate these sentences in order, so that the order of the final summary is consistent with the source article. We determine the order of the sentences according to the order of the events they describe. In the integration stage, SliSum do not need to change any information of the selected sentences.


 


§ EXPERIMENTAL SETUP






 §.§ Datasets


We evaluate the performance of SliSum for short and long text summarization on four popular benchmark datasets: CNN/DM <cit.> and XSum <cit.> are two widely-used short news summarization datasets. PubMed and arXiv <cit.> are two scientific paper datasets for long document summarization, which are much longer than the common news articles. PubMed contains academic papers from the biotechnology domain, while arXiv contains papers from different scientific domains. 
Limited by the maximum number of requests and time cost, we randomly select 100 articles from test set of each dataset to construct our test set, respectively.









 §.§ Baselines


To validate the effectiveness of our proposed approach, we apply SliSum to three state-of-the-art instruction-tuned LLMs, including LLaMA-2 13B Chat[https://ai.meta.com/resources/models-and-libraries/llama/https://ai.meta.com/resources/models-and-libraries/llama/], Claude-2[https://www.anthropic.com/index/claude-2https://www.anthropic.com/index/claude-2] and GPT-3.5-Turbo-0613[https://openai.com/blog/chatgpthttps://openai.com/blog/chatgpt]. 







We also compare SliSum with recent faithfulness enhancement summarization models on CNN/DM and XSum, such as CLIFF <cit.> and FES <cit.>. We select GPT-3.5 with SumCoT <cit.> as LLM news summarization baseline. For arXiv and PubMed, we compare SliSum with two long document summarization baselines: FactorSum <cit.>, a factorized energy-based abstractive model that improves the performance and applicability by separate budget decisions from selecting important content in the document, and Lodoss <cit.>, an extractive architecture that learns robust sentence representations by performing summarization and segmentation simultaneously. We compare SliSum with Refine <cit.>, a LLM-based long document summarization method that joins the summary of the previous text segment with the next segment, and then iteratively generate summaries for them.







 §.§ Metrics


We evaluate the factual consistency, fluency and informativeness of summaries using four different metrics: (1) FactCC <cit.>, a weakly-supervised, model-based approach for verifying factual consistency and identifying conflicts between source documents and generated summaries, (2) SummaC <cit.> that enables natural language inference models to detect inconsistency, (3) ROUGE<cit.>, an automatic evaluation metric for the informativeness and fluency of a summary based on lexical overlap, and (4) BERTScore <cit.>, that computes a similarity score between candidate and reference summaries using contextual embeddings.



 §.§ Implementation


We run LLaMA-2-13B with Text Generation Inference[https://github.com/huggingface/text-generation-inferencehttps://github.com/huggingface/text-generation-inference] on 8 × 24GB NVIDIA GeForce RTX 3090 GPUs. For CNN/DM and XSum, we cluster and filter sentences using the DBSCAN algorithm with ε = 0.25 and MinPts = 2. The window size L_w is 150 words, and the step size L_s is 50 words. For arXiv and PubMed, we set distance threshold ε to 0.25, and MinPts to 3. The window size L_w is 750 words, and the step size L_s is 150 words. For SumCoT, we use the same prompts in paper. We set chunk size to 750 words in Refine. The prompts used for summary generation, sentences classification and sentences integration are listed in Table <ref>. Regarding other baselines used in the experiments, we use standard checkpoints provided by the authors and adopt the same configuration as in the corresponding papers, respectively.











































§ RESULTS




 §.§ Main Results
















  
Overall Performance.The experimental results of SliSum and baselines on four summarization datasets are reported in Table <ref>. 

In terms of short text summarization, Claude-2 with SliSum achieves a relative gain of 13.9% on FactCC and 15.8% on SummaC for CNN/DM respectively. Notably, Claude-2 with SliSum also performs the best on both ROUGE and BERTScore compared with other baselines, indicating that SliSum simultaneously improves informativeness and fluency of summaries. 

For long document summarization, Claude-2 with SliSum achieves a notable relative gain of 20.5% on FactCC and 22.2% on SummaC for arXiv respectively. Although FactorSum and Lodoss achieve close performance to SliSum on reference-based similarity metrics, they are significantly lower than SliSum on FactCC and SummaC. Regarding PubMed dataset, SliSum enables Claude-2 and GPT-3.5 to achieve almost the same excellent performance, and outperform other baselines on all metrics. 

Overall, LLM combined with SliSum uniformly outperforms base model itself with a wide margin on all metrics, which demonstrates the generality and effectiveness of SliSum. Specially, among three LLMs, SliSum reveals the most outstanding performance on Claude-2 and the greatest gain on LLaMA-2. Besides, SliSum applied to GPT-3.5 outperforms SumCoT and Refine on all four datasets. This shows that the superiority of our approach to other LLM-based baselines. We observe that SliSum achieves higher improvement on arXiv and PubMed than CNN/DM and XSum, and leads to greater performance gain on faithfulness metrics than on informativeness metrics. Therefore, our approach can substantially improves faithfulness of various LLMs in short and long text summarization, while maintaining their fluency and informativeness.





  
SliSum enhances the ability to process long contexts. We observe that SliSum has different performance gains on different datasets, which may be due to the style or length of the text. To eliminate the distractor of the text style and further demonstrate the effectiveness of our approach on news texts. We evaluate the factual consistency of GPT-3.5 on test samples of length ranges from 500 to 3000. We select 25 samples from CNN/DM and XSum with a length that is less than 5% different from the specified length to construct the test sets, respectively. 

As shown in Figure <ref>, compared to the original model, SliSum significantly improves the faithfulness of GPT-3.5 in news texts of various length, and the performance gain is positively correlated with the length. This is consistent with the results in Table <ref> and indicates that SliSum is capable of leading to larger benefits from increasingly long context. Besides, SliSum prevents the FactCC and R-1 score of the summaries from continually decreasing as source articles grow longer. For example, when the length of articles increases from 1000 to 3000, the FactCC score of summaries generated for CNN/DM by SliSum remains stable, while the scores of the baseline still decreases. We demonstrate that our approach enhances the ability to handle long texts in a variety of styles.



 §.§ Ablation Study


SliSum contains two essential modifications: sliding generation, filteration and aggregation based on self-consistency. In order to investigate the effect of modification designed in the SliSum, we conduct the following ablation studies by comparing SliSum with corresponding variation on GPT-3.5. 






















  
Sliding Windows vs. Parallel Windows. Parallel window(PW), a special case of sliding window(SW), partitions contexts in a non-overlapping manner. To investigate the contribution of sliding generation to SliSum, we compare comprehensive performance of SW and other baselines (also generate the summary K times to perform aggregation). Table <ref> shows that single window method is worse than others on PubMed, due to the distraction issue caused by long texts. In contrast, SW outperforms PW method on all three metrics, which demonstrates that sliding generation brings performance improvement to LLMs. This is because SW can take advantage of longer context to support generating summaries, whereas PW lack direct interaction between windows, resulting in the inability to adequately extract global important information. Specifically, benefiting from overlapping windows, SW can utilize information of  L_w + L_s(⌈ L_w / L_s ⌉ - 1) words during generating summary for each text segment, while PW can only utilize L_w words. Furthermore, Table <ref> shows that sliding windows facilitates LLMs to more fairly and faithfully process contexts of various length without position bias. The summaries generated by SliSum contain the content of each position range of the source article relatively evenly.




  
Filteration and Aggregation improve faithfulness. In order to understand the importance of filteration and aggregation based on self-consistency, we conduct ablation study by comparing SliSum to variation w/o aggregation. We remove two steps by directly generating global summaries for all local summaries (Map-Reduce <cit.>). As shown in Figure <ref>, the FactCC scores of both SW and PW w/o aggregation are lower than their counterparts. Notably, sliding windows w/o aggregation substantially diminish the factual consistency of summaries, even worse than parallel windows w/o aggregation. This indicates that aggregation is essential for SliSum, because sliding generation while providing more diverse information also leads to self-contradiction problem, which asks us to address this problem by filteration and aggregation based on self-consistency. 



 §.§ Impacts of Hyperparameters


We conduct quantitative experiments to investigate impacts of hyperparameters in SliSum on performance. Limited by computational resources and budget, we randomly select 25 articles from corresponding dataset to construct reduced test sets.





  
Ratio of L_w to L_s. Figure <ref> (left) shows that the FactCC score when varying in ratio K = ⌊ L_w / L_s⌋ from 1 to 10. As shown in Equation <ref>, K represents the minimum number of process content of the article. A large K can more adequately exploit the self-consistency of LLMs to improve performance. When the ratio grows from 1 to 5, the FactCC scores increased by 9.5% and 11.5% on arXiv and PubMed respectively, yet gradually converges when the ratio grows from 5 to 10, which indicates that processing articles too many times is extremely expensive and unnecessary. To save computational resources, we set the L_w / L_s to 5 in our experiments.











  
Window Size. Intuitively, shorter window means that LLMs can more accurately exploit information in context. But this has harmful effect on global understanding of entire article, resulting in the summary only contains locally important content. Nevertheless, a longer window will cause distraction issue and more uncertainty. As shown in Figure <ref> (right), the FactCC score of summaries initially raises as the window size grows. However, when window size increases from 750 to 1200, the performance remains stable or slightly drops. The quantitative experiments results indicate that LLMs require sufficient information to generate high-quality summaries from context, but a larger window will not further boost performance. 













  
Minimum Number of Clusters (MinPts). MinPts controls clustering and filtering of sentences about the same event. A higher MinPts forces SliSum to select sentences that are important in more windows. Table <ref> shows the FactCC and R-1 scores significantly increase with a bigger MinPts, but it will also lower R-1 score when MinPts greater than 3,  as a result of many relatively important sentences are filtered out. Therefor, we set MinPts to 3 for arXiv and PubMed. More generally, considering the balance between faithfulness and informativeness of summaries, the MinPts is set to 1/2 K for different parameter configurations.





  
Distance Threshold ε The choice of distance threshold ε depends on the data distribution. We count the average and the maximum distance between statements about the same event on the four datasets. Besides, we also calculate the average Hausdorff distance between sets of statements about different events. Given the set of statements X and Y for events A and B, the Hausdorff distance between X and Y is


    d_H(X, Y) = max{sup_x ∈ Xinf_y ∈ Yd(x, y), sup_y ∈ Yinf_x ∈ Xd(x, y)}


Table <ref> shows that the average distance between statements about the same event is around 0.25, which is much smaller than the distance between statements about different events. Hence, we set ε to 0.25 for all datasets. The results in table <ref> are consistent with our observation that statements generated by LLMs for the same event are lexically similar.



 §.§ Complexity Analysis




  
Theoretical Analysis SliSum reduces the computational complexity in the process of summary generation compared with the base model. For the input of length L, the computational complexity of most LLMs (such as LLaMA) is O(L^2). However, after applying SliSum to base models, the computational complexity of summary generation decreases to O(L). 
    
Generally, given window size L_w, step size L_s (L_s ≤ L_w) and an article whose length is L is far larger than the window size, K = L_w / L_s (see Equation <ref>). Our approach needs to process a total of (L / L_s + (K - 1)^2) text segments of length L_w, hence, the complexity of SliSum is 


    O(K × L_w × L + (K - 1)^2 × L_w^2) = O(L)


By solving the equation, we can see that when the input length is greater than 1.36 × K × L_w, the computational cost of summary generation of SliSum is less than base model.

Of course, SliSum also includes filteration and aggregation, which brings additional computational costs. However, the DBSCAN algorithm runs very fast. As for the final aggregation process, LLM usually only needs to process very few sentences. Importantly, SliSum splits the original task into smaller subtasks, so parallelization can be used to speed up inference. Overall, SliSum slightly increases the computational cost compared to the base model.





  
Quantitative Experiments In order to verify the above theoretical studies, we recorded the inference time of the LLaMA-2-13B on four datasets (including 100 articles) to evaluate the computational cost. Table <ref> indicates the LLaMA-2-13B with SliSum only took twice as long as the LLaMA-2-13B without SliSum. By contrast, as shown in Table <ref>, SliSum substantially improves faithfulness of various LLMs. Therefore, SliSum slightly increases the computational cost, but brings relatively high performance gains.



§ CONCLUSION


In this paper, we propose SliSum, a novel LLM summary generation architecture, which improves faithfulness of LLMs in both short and long text summarization without additional resources and fine-tuning. SliSum leverages sliding generation as instruments for exploiting the self-consistency of LLMs, enables LLMs to more fairly and faithfully process contexts of various length, and solve contradictions between local summaries by clustering and filtering. Extensive experiments demonstrate the effectiveness of SliSum in diverse LLMs. SliSum empowers LLMs to generate faithful summaries while maintaining their fluency and informativeness. Furthermore, we conduct ablation studies to justify the modification made in the SliSum and investigate the impact of hyperparameters in SliSum. We also demonstrate through theoretical analysis and quantitative experiments that SliSum only slightly increases the computational cost.



§ ACKNOWLEDGEMENTS

This work was supported by the Zhejiang Provincial Natural Science Foundation of China under Grant No. LZ23F020009, the NSFC project (No. 62072399), MoE Engineering Research Center of Digital Library, China Research Centre on Data and Knowledge for Engineering Sciences and Technology, and the Fundamental Research Funds for the Central Universities. We also express our sincere gratitude to anonymous reviewers for their invaluable feedback and constructive comments.




§ BIBLIOGRAPHICAL REFERENCES


lrec-coling2024-natbib















































































































    


 Paper Title: CARMIL: Context-Aware Regularization on Multiple Instance Learning models for Whole Slide Images.








CARMILT. Nait Saada and V.Di Proietto et al.

CARMIL: Context-Aware Regularization on Multiple Instance Learning models for Whole Slide Images
    Thiziri Nait SaadaThese authors contributed equally to this work Owkin, Inc., New York, NY, USA Mathematical Institute, University of Oxford, Oxford, UK naitsaadat@maths.ox.ac.uk 
 Valentina Di-Proietto[1] [2] valentina.di-proietto@owkin.com 
 Benoit Schmauch[2] benoit.schmauch@owkin.com 

Katharina Von Loga[2] katharina.vonloga@owkin.com 

Lucas Fidon[2]  lucas.fidon@owkin.com
    August, 2024
================================================================================================================================================================================================================================================================================================================================================================================================



§ ABSTRACT 
Multiple Instance Learning (MIL) models have proven effective for cancer prognosis from Whole Slide Images. 
However, the original MIL formulation incorrectly assumes the patches of the same image to be independent, leading to a loss of spatial context as information flows through the network. Incorporating contextual knowledge into predictions is particularly important given the inclination for cancerous cells to form clusters and the presence of spatial indicators for tumors. State-of-the-art methods often use attention mechanisms eventually combined with graphs to capture spatial knowledge. In this paper, we take a novel and transversal approach, addressing this issue through the lens of regularization. We propose
Context-Aware Regularization for Multiple Instance Learning (CARMIL), a versatile regularization scheme designed to seamlessly integrate spatial knowledge into any MIL model. Additionally, we present a new and generic metric to quantify the Context-Awareness of any MIL model when applied to Whole Slide Images, resolving a previously unexplored gap in the field.

The efficacy of our framework is evaluated for two survival analysis tasks on glioblastoma (TCGA GBM) and colon cancer data (TCGA COAD). 












§ INTRODUCTION




The digitization of histopathology Whole Slide Images (WSIs) and the development of deep learning methods has lead to promising computational methods for cancer prognosis.


One computational challenge is the large size of WSIs, of the order of 100,000 × 100,000 pixels. Processing images of such size with a deep neural network directly is not possible with the GPUs commonly available. 

Overcoming this problem, previous work proposes to tessellate each WSI into thousands of smaller images called tiles and global survival prediction per slide is obtained in two steps.
The tiles are first embedded into a space of lower dimension using a pre-trained feature extractor model, and a MIL model is trained to predict survival from the set of tiles embeddings of a WSI <cit.>.






One limitation of MIL is the assumption that tiles from the same WSI are independent <cit.>. In particular, MIL models fail to leverage spatial interactions between tiles and their ordering in a WSI. 

In contrast, pathologists take into account the spatial organization of WSIs in their analysis.

To tackle this issue, a correlated variant of MIL has been proposed in <cit.>.

In practice, interactions between neighboring tiles in a WSI can be modeled using a graph of tiles.


Graph Neural Network (GNN) <cit.> and Transformers with local attention <cit.> were proposed to capture correlations between neighboring tiles.

However, as these correlations are parameterized using more complex architectures involving a greater number of parameters, the effectiveness of such mechanisms might be limited by the amount of available training samples. Besides, histopathology datasets usually contain only up to a few hundred WSIs.
By contrast, incorporating spatial correlations between tiles into the model through regularization has been under-explored in computational pathology.




In this paper, we introduce the CARMIL framework to enhance the Context-Awareness of any MIL model, initially agnostic to the spatial context, using Regularization. 

In CARMIL, a spatial encoder and a spatial decoder are added between the feature extractor and the MIL model.

In addition, a Context-Aware Regularization (CAR) loss function is proposed to train the spatial decoder to reconstruct the input graph of tiles.


The proposed spatial encoder aims at distilling the spatial relation between neighboring tiles directly into the tile embeddings so that any MIL model can exploit this information.

To this end, the spatial decoder and the CAR loss are designed to encourage bringing closer the embeddings learnt by the spatial encoder for tiles that are spatially close in the original WSI.





To the best of our knowledge, such regularization strategy has so far not been studied in computational pathology. 

The nearest related work is the topologically-aware regularization for cancer blood cell classification introduced by <cit.>. 


Yet, this regularization is based on the 0th order Betti number, which only captures the connected components of the graph of tiles.

Instead, CARMIL aims to reconstruct the entire graph structure.

Our main contributions are three-fold. 

First, we introduce CARMIL, a Context-Aware Regularization module based on Graph AutoEncoders that we incorporate into the classical MIL framework through regularization.

Second, we propose DeltaCon, a novel metric to quantitatively assess the Context-Awareness of any tiles embedding. This metric allows for a better inspection of how well the embeddings conform with the original spatial arrangement of the tiles within the WSIs.

Lastly, we evaluate CAR on a set of benchmark MIL models for the task of overall survival prediction for glioblastoma and colon cancer. These models are precisely chosen because they are originally context-independent, and we show through extensive quantitative and qualitative ablation studies how CARMIL turns them into Context-Aware models and leads to improved C-index.



 











































§ METHODS





 §.§ Background: classical MIL framework for risk prediction


In this section, we summarize the main steps of the classical MIL pipeline, that we enhance with our Context-Aware Regularization in the next section.





  
Preprocessing. 
Each WSI is segmented to keep only the tissue and remove the background (e.g. using the otsu method).

The tissue parts are then tessellated into tiles of size 224 × 224 pixels taken at 20X magnitude (0.5μ m/pixel) and each WSI is reduced to a set of n tiles of tissue selected at random.

A pre-trained feature extractor is then used to map each tile into a low-dimensional tile feature space of dimension d. This results in a matrix representation X^(i)∈ℝ^n × d for the WSI indexed by i.













  
MIL model.

The MIL model f_MIL is a parametric model trained to aggregate the tile features of a WSI and estimate a risk associated with the overall survival of the patient 

    f_MIL: ℝ^n × dℝ.


The parameters θ_MIL of the MIL model f_MIL are optimized to minimize a survival loss ℒ_MIL, such as the Cox loss. Formally, this corresponds to the optimization problem


    min_θ_MIL1/N∑_i=1^N ℒ_MIL(
        f_MIL(X^(i);θ_MIL),  y^(i)),


where N is the number of training WSIs and y^(i) is the overall survival risk of patient i.

Examples of such MIL models are ABMIL <cit.>, Chowder <cit.>, and AdditiveMIL <cit.>.

It is worth noting that the spatial relationships between tiles are not used by f_MIL, since their ordering is random in X^(i).





 §.§ Context-Aware Regularization (CAR)


We propose a general approach to leverage local spatial context in computational histopathology. 

Our pipeline, which can be applied to any task or MIL model, is depicted in fig. <ref>.



  
Graph construction.
For each WSI, we build a graph G^(i)=(X^(i), A^(i)) where the vertices are the tiles features X^(i)∈ℝ^n × d, as described in the previous section, and A^(i)∈ℝ^n × n is the adjacency matrix computed based on the euclidean distance between the spatial coordinates of the tiles in the original WSI for a given number k of nearest neighbors.



  
Context-Aware Regularization.
To embed the spatial structure of the WSI into the tiles features, we introduce a spatial encoder f_E of parameters θ_E and a spatial decoder f_D of parameters θ_D before the classical MIL model f_MIL of (<ref>) as illustrated in Fig. <ref>

    [t]
        f_Eℝ^n × d×ℝ^n × n   ℝ^n × d_E
    
        (X,A)   ↦ Z,
          [t]
        f_D ℝ^n × d_E   ℝ^n × n
    
        Z    ↦Â.


 
Here, the goal of f_E is to distill the spatial information contained in an input WSI graph G=(X, A) 

into tiles features X, resulting in a new low-dimensional Context-Aware embedding Z ∈ℝ^n × d_E of the WSI. 
To achieve this goal, f_D aims at reconstructing the input adjacency matrix A from Z during training while, concurrently, f_MIL aims at predicting the patient risk from Z.

Formally, this corresponds to adding a Context-Aware Regularization term ℒ_CAR to the training optimization problem of classical MIL (<ref>):


    min_θ_MIL, θ_E, θ_D1/N∑_i=1^N( 
        (1-β) 
        ℒ_MIL(
            f_MIL(Z^(i);θ_MIL),  y^(i))
        +
        βℒ_CAR(
            f_D(Z^(i); θ_D), A^(i)) 
        )

where the same notations as in (<ref>) are used and,

    {[                                                 Z^(i) =                                        f_E(G^(i); θ_E),;                                           ℒ_CAR(Â, A) = 1/n^2∑_p,q(A_pqlog ( Â_pq ) + (1-A_pq) log (1 - Â_pq)). ].


By reformulating the training problem as a joint optimization task, the models' parameters are obtained by simultaneously solving a pair of objectives, which may initially appear orthogonal. 

When β=0, the method boils down to the classical MIL pipeline, whereas β=1 reduces the task to encoding the spatial context only as in a Graph AutoEncoder (GAE) <cit.>.

Note that the spatial decoder f_D is not required at inference.

The spatial encoder f_E and spatial decoder f_D are obtained by stacking up, respectively ℓ_E and ℓ_D, graph convolutional network (GCN) layers:



    GCN((X,A); θ)  A ReLU(X)W^θ,

    where W^θ is a learnable weight matrix and the matrix A may be preprocessed. 

If we denote the composition of a function f by itself t times as f^∘ t, then,


    Z = f_E(G;θ_E) GCN^∘ℓ_E(G;θ_E),     
        U_D(Z) = GCN^∘ℓ_D( (Z,A);θ_D).




Similarly to a GAE, the last layer of the spatial decoder is an inner-product decoder with a sigmoid function σ, that aims at reconstructing the input adjacency matrix A,

    Â= f_D(Z;θ_D) σ( U_D(Z) U_D(Z)^T).


    


















 §.§ DeltaCon for quantitative measure of Context-Awareness


Reporting task-related scores alone is not enough to quantify the amount of Context-Awareness in a tiles embedding space.

To overcome this issue, we propose a novel Context-Awareness metric based on DeltaCon <cit.>. 

Let Ẑ∈ℝ^n× p be tiles descriptors, we denote A(Ẑ)∈ℝ^n× n the adjacency matrix based on the k nearest neighbors for the euclidean distance between the components of Ẑ.

This is different from A, defined in sec. <ref>, that is based on the distance between the spatial coordinates of the tiles.

Ẑ can be any representation of the tiles, here it will either be the features X or the embeddings Z as defined in sec. <ref>.

We propose to use the DeltaCon similarity between the adjacency matrix based on tiles descriptors A(Ẑ) and the adjacency matrix based on the spatial coordinates A for the same WSI as our metric of Context-Awareness.

Since we have directed graphs and the adjacency matrices A and A(Ẑ) are non-symmetric, we propose this extension of the original DeltaCon similarity

    DeltaCon(A(Ẑ), A) = 
    1/1 + ||S(A(Ẑ)) - S(A)||_F,    
        S(A)  (𝕀 + ϵ^2 D - ϵ A)^-1,

where D D_in+ D_out, where D_in, resp. D_out, is the diagonal matrix counting the incoming, resp. outcoming, edges and ||.||_F denotes the Frobenius norm.

Intuitively, S(A)∈ℝ^n× n is a matrix capturing the degrees of similarity between all tiles when one travels in the original graph corresponding to A.

Direct neighbors are the most similar tiles, followed by neighbors of neighbors, and so on.

As a result, DeltaCon(A(Ẑ), A) is close to 1 when the neighbors in the tiles descriptors space and the spatial coordinates are almost the same and it smoothly decreases towards 0 when those neighbors become dissimilar between the tiles representations and the spatial coordinates.



















§ IMPLEMENTATION DETAILS



 §.§ Datasets for survival prediction using whole slide images




  
Glioblastoma data.
We used H&E slides of patients with glioblastoma from the datasets TCGA GBM <cit.> and TCGA LGG <cit.>.
We filtered cases according to the latest WHO classification for gliomas <cit.>. See App. <ref> for more details.








  
Colon cancer data.
We used H&E slides of patients with colon adenocarcinoma from TCGA. The TCGA COAD dataset contains a total of 431 cases from 24 centers.




 §.§ Evaluation


All models were trained and evaluated on TCGA COAD and TCGA GBM using 5-fold nested cross validation <cit.> to allow for hyperparameters tuning and assess the generalisation independently on those datasets.

Three repeats were used in the inner loop, corresponding to three different random initializations of the MIL or CARMIL models.

As a result, metric evaluation on each of the 5 test splits was performed using an ensemble of 15 models (5 inner validation splits and 3 repeats).
Ensembling was performed by averaging the risk output of the models of an ensemble. In all the tables, we show the mean (std) C-index for OS obtained using 5-fold nested cross validation. The best results are in bold.




 §.§ Preprocessing

For tissue segmentation, a 2D U-net <cit.> trained on a pancancer dataset of manually annotated WSIs is used.

Tile features of dimension d=768 are computed using a state-of-the-art self-supervised model, Phikon, trained on pancancer H&E slides from TCGA <cit.>.

The parameters of Phikon are kept frozen in all our experiments.

The preprocessing is the same for all the models we compare to.



 §.§ Deep learning training

The Cox loss was employed in the supervised training of all models, utilizing overall survival labels. This corresponds to ℒ_MIL in eq. (<ref>).

The CAR loss, denoted as ℒ_CAR in eq. (<ref>), is applied across all CAR models and we use β=0.5 in eq. (<ref>) for the total loss accross all CAR models. 
The choice of β=0.5 followed the empirical observation that ℒ_MIL and ℒ_CAR have similar range of values during the first epoch of training.

Adam optimizer <cit.> with momentum β_1=0.9 and β_2=0.999 is used for training with a learning rate on the grid {0.001, 0.003, 0.01} for all models.
The maximum learning rate value 0.01 was chosen heuristically to be the smallest value on the grid for which most MIL models diverged during training.
The number of training epochs is optimized on the grid {20, 30}.


One NVIDIA Tesla T4 GPU with 16GB of VRAM and 8 Intel(R) Xeon(R) 2.00GHz CPUs are used for training and inference of each model.





§ EXPERIMENTS







 §.§ Survival prediction performance

In tables  <ref> and <ref>, we report the performance of various MIL models for the challenging task of survival prediction on TCGA COAD and TCGA GBM.

In table <ref>, we compare classical MIL models, selected for being agnostic to the spatial context, to their performance when enhanced with our CAR scheme.

In total, CAR improved C-index values by up to 4.8 percentage points (pp) in 9/12 settings and on average the C-index increased by 1.5 pp on TCGA COAD and 0.5 pp on TCGA GBM.

In addition, we compare classical MIL models enhanced with CAR to state-of-the-art model architectures that were designed specifically to leverage graphs of tiles. 

CARMIL models perform similarly or better than these more sophisticated models in terms of C-index; see table <ref>. 











This suggests our CAR method improves performance by successfully integrating spatial context through regularization.




 §.§ Context-Awareness performance






In this section, we compile evidence supporting the successful injection of spatial information in the proposed CAR models and its associated performance benefits. First, we assess whether CAR models genuinely exploit the input graph for their predictions. To this end, we perturb the graph at inference time by randomly shuffling all the off-diagonal terms of the adjacency matrix A. If the CAR models were to fail to exploit the graph structure, this disruption would not impact their performance. However, the results reported in table <ref> indicate that graph shuffling leads to a degradation in the model's performance, showcasing that the graph structure is indeed used in the model's decision-making process.









Secondly, we assess the Context-Awareness of trained CARMIL models compared to the feature extractor, using the DeltaCon metric that we defined in sec. <ref>. 

For each WSI, the adjacency matrix computed in the embedding space learnt by the spatial encoder of fig. <ref> is compared to the original adjacency matrix that accounts for the spatial arrangement of the tiles within the slide. 

The DeltaCon provides us with a score between 0 and 1.

A higher DeltaCon score indicates a higher degree of spatial consistency in the learnt embedding space as it implies that the arrangement in the embedding space closely aligns with the original spatial organization within the slide. We average this slide-level score across the whole TCGA COAD and GBM datasets, see table <ref>. In app. <ref>, we showcase an example of a WSI illustrating how the embeddings provided by CARMIL encoders are more spatially consistent than the original features. Moreover, in fig. <ref>, we can clearly see how the spatial encoders implement more Context-Awareness than the feature extractor (dotted lines). Indeed, DeltaCon is almost always greater for CARMIL models than for the feature extractor. Therefore, these findings confirm the successful injection of spatial knowledge resulting from the CAR. Additionally, we explored the relationship between spatial information incorporated by the spatial encoder and C-index performance. In TCGA COAD, the performance seems to correlate with the level of Context-Awareness, potentially providing insights into colon cancer. However, in TCGA GBM,  it is uncertain whether adding spatial elements improves performance, and understanding how Context-Awareness influences overall performance remains an open question.









































§ CONCLUSION


The MIL framework fails to leverage spatial interaction and organization of tiles in WSIs, potentially limiting prognostic model performance. In this work, we proposed to tackle this issue by injecting spatial knowledge into the traditional MIL framework exclusively through the prism of regularization.

We introduced CARMIL, a method to embed spatial relations between tiles directly into tile features.

This addition mimics the pathologist's consideration of spatial arrangement in slide-level prognosis.

We evaluated our method on survival prediction for colon cancer and glioblastoma, showing improved performance and revealing performance declines when spatial context is disregarded. 

We also introduced a metric for quantifying Context-Awareness, hoping to help researchers assess spatial consistency more systematically. Lastly, we discussed in App. <ref> the influence of certain parameters on our method, but we did not investigate the impact of the feature extractor on overall performance. The relationship between model robustness and Context-Awareness remains an open question, suggesting avenues for future research.






We thank  Jean Philippe Vert and Eric Durand for their thorough proofreading of our paper, as well as Jean-Baptiste Schiratti and Alexandre Filiot for providing insightful feedback on various aspects of this work. We are thankful to Céline Thiriez too for her encouragement and support.






















 












§ GLIOBLASTOMA WHO 2021 CLASSIFICATION

The WHO 2021 classification defines glioblastoma as IDH-wild type and H3-wild type brain tumor with at least one of the following features: necrosis and/or microvascular proliferation, TERT promoter mutation, EGFR amplification, or concomitant gain of chromosome 7 and loss of chromosome 10.
We will refer as TCGA GBM to those glioblastoma cases from the datasets TCGA GBM and TCGA LGG after reclassification.
TCGA GBM contains 352 cases from 18 centers.

This change of classification of glioblastoma has been shown to have a negative impact on the prognostic value of previously published biomarkers <cit.>. Therefore, it is clinically important to evaluate previous and new prognostic models on glioblastoma using the new WHO classification.



§ REPRODUCIBILITY




 §.§ Construction of the spatial adjacency matrix A


Each WSI G=(X,A) consists of a set of n tiles. For each tile p, we keep track of its spatial coordinates in the 2D plane formed by the tissue region, 𝐜_p=(x_p,y_p), in addition to the d-dimensional features vector 𝐮_p produced by the feature extractor for that tile.
We thus compute the gaussian kernel, ∀ p≠ q,

    K_pq = exp(-||𝐜_p-𝐜_q||_2/2),

and we set the diagonal to 0.
Provided k the number of neighbors, we select the k nearest neighbors for each node based on the similarity matrix K. The adjacency matrix A is defined such that A_pq equals K_pq if tile q is one of the nearest neighbors of tile p, and 0 otherwise. 



 §.§ Construction of the adjacency matrix Ã in the embedding space


Adjacency matrices in the embedding space are used for the evaluation of Context-Awareness using DeltaCon.

Consider training of our model complete and the final parameters to have converged to θ_MIL, θ_E, θ_D. Given a WSI, represented as G=(X, A), our spatial encoder returns a lower dimensional vector Z=f_E(G; θ_E)∈ℝ^n× d_E. Observe that Z_p is the d_E-dimensional embedding for tile p.
Based on this vector, we construct an adjacency matrix Ã, following the same principle as before, but this time based on the affinity between embeddings rather than using the spatial coordinates of the tiles, namely, ∀ p≠ q,

    K̃_pq = exp(-||Z_p-Z_q||_2/2),

and we set the diagonal to 0. We similarly pick the k nearest neighbors for each tile and we construct the adjacency matrix Ã, that we refer to as the adjacency matrix in the embedding space. Therefore, Ã_pq is the affinity between the embeddings of the tiles p and q.



 §.§ Computation of DeltaCon



Observe that both adjacency matrices, A and Ã, are non-symmetric. Indeed, they are derived by taking the k-nearest neighbors of each node in the kernel matrices K and K̃, which is an inherently non-symmetric operation. In simple words, tile i can be connected to tile j, but
i is not necessarly a neighbor of j, meaning the induced graphs are directed. In <cit.>, the authors introduce a proxy function defined for undirected graphs with symmetric adjacency matrix A and degree matrix D as

    S(A)  (𝕀 + ϵ^2 D - ϵ A)^-1.


We extend this definition to directed graphs by considering a non-symmetric adjacency matrix A and a degree matrix that accounts for the number of edges entering or leaving each node. Namely, considering D D_in+ D_out, where D_in, resp. D_out, is the diagonal matrix counting the incoming, resp. outcoming, edges, then we can refer back to eq. (<ref>) to generalize the definition of DeltaCon similarity between directed graphs. Given A and Ã, as described in sec. <ref> and <ref>, and their associated degree matrices D and D̃, we can now quantitatively assess the amount of spatial information, or Context-Awareness, of any CARMIL model by computing,


    DeltaCon(A, Ã) 1/1+||S(A) - S(Ã) ||_F.


In table <ref>, we assess the Context-Awareness of the best-performing CARMIL model from the 15 models evaluated through nested cross-validation, see sec. <ref>. For each WSI in each dataset, we first compute the spatial adjacency matrix A with k=8 nearest neighbors, as in sec. <ref>. Then, for all rows in table <ref> except the first, the matrix Ã is computed based on the embedding vector Z from the model's encoder after training, also with k=8, see sec. <ref>. For the first row, the adjacency matrix Ã is based on the feature vector X generated by the feature extractor. The similarity score given by eq. (<ref>) is averaged across all slides within each dataset to produce the table's values, along with the C-index mean performance of the corresponding model.
In fig. <ref>, we report these same results in a scatter plot to better illustrate two important findings. First, by comparing the scores from the feature extractor to the other models on both datasets, we confirm the successful injection of Context-Awareness into our CARMIL models. As a side note, it is worth mentionning here that, even though DeltaCon lies between 0 (indicating very dissimilar graphs) to 1 (for identical graphs), the interpolation between these extremes does not seem to evenly spread out within the interval [0,1]. This would explain the small variations in our table <ref>. In fact, as evidenced in the tables of results from the original paper <cit.>, most values are closer to 0 than 1. Nevertheless, for our analysis, only the relative ordering of these quantities is key in proving the enhanced spatial awareness provided by our method.
Second, we try to correlate gains in performance with the amount of added Context-Awareness in the CARMIL model. Conclusions  such as the more spatial knowledge is incorporated, the better the model perform, are rather difficult to draw.










 §.§ Qualitative assessment of Context-Awareness in CARMIL models


Consider a WSI, from which n tiles are sampled. Passing this slide to our feature extractor, we get a vector X∈ℝ^n× d. For each tile, we compute the mean of the eucliden distance between the d-dimensional features of this very tile with the features of all its k=8 spatial nearest neighbors within the slide. This results in a vector Y^features∈ℝ^n  that accounts for how well aligned the features learnt by the feature extractor conform to the spatial arrangement of the slide. If the features exactly reflect the spatial context of each tile, the coefficients in the vector Y^features should be fairly constant and of low magnitude. Their variations are shown in fig. <ref>, where we superposed the underlying slide with the values of Y^features.
Next, we pass these features through our CARMIL encoder at inference time, following the same procedure but using the embedding vector Z∈ℝ^n × d_E – with the same notation as before. This results in a vector Y^CARMIL∈ℝ^n. We report the coefficients in fig. <ref>.
To obtain similarly scaled heatmaps, we jointly normalize the coefficients of both matrices by their maximum value, ensuring all their coefficients range between 0 and 1. Note that both tile features X and tile embeddings Z are first normalized to have magnitude 1, ensuring a fair comparison that mitigates high-dimensional concentration effects on vector norms. 

From fig. <ref>, we can clearly see that the embeddings generated by the CARMIL encoder evolve smoothly in the 2d plane, resulting in a uniformly low heatmap value across the slide. In comparison, the tile features provided by the feature extractor Phikon exhibit less alignment with the spatial organization of the tiles, providing another qualitative indication of the enhanced spatial consideration in CARMIL models.









§ DISCUSSION ON CERTAIN PARAMETERS




  
Number of tiles n. The number of tiles n that we randomly select from the tissue region of each WSI has a significant impact, and we provide an intuitive explanation for this. The amount of spatial context required to make accurate predictions is inherently tied to the number of tiles, as it determines the scale at which patterns can be grouped. Moreover, we observe that different WSIs can vary greatly in size, resulting in a wide range of number of tiles across datasets. Whilst we fixed n regardless of the WSI at stake, it would be a natural improvement to choose n for each slide in proportion to the original size of the tissue region.



  
Number of nearest neighbors k. Similar to the number of tiles, it is probably advisable to choose k, the number of nearest neighbors, in proportion to the number of tiles n, and subsequently, the size of the WSI. In our experiments, we optimized k across a fixed grid of values, independent of n.



  
Number of GCN layers ℓ_E, ℓ_D.
The number of layers ℓ_E and ℓ_D in the spatial encoder and decoder, as well as the dimensions for the projections they induce were part of our grid search. We observed that (ℓ_E, ℓ_D)=(1,1) worked best for TCGA COAD and (ℓ_E,ℓ_D)=(2,2) for TCGA GBM. The dimensions of each layer were also finetuned from a grid of fixed values, resulting in d_E=d as an optimal choice for the spatial encoder and all intermediate layers. Besides, we observed that significantly reducing the embedding dimension d_E almost systematically led to a substantial decline in performance. 



 Paper Title: Advancing Medical Image Segmentation: Morphology-Driven Learning with Diffusion Transformer.



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



§ ABSTRACT 
Understanding the morphological structure of medical images and precisely segmenting the region of interest or abnormality is an important task that can assist in diagnosis. However, the unique properties of medical imaging make clear segmentation difficult, and the high cost and time-consuming task of labeling leads to a coarse-grained representation of ground truth. Facing with these problems, we propose a novel Diffusion Transformer Segmentation (DTS) model for robust segmentation in the presence of noise.
We propose an alternative to the dominant Denoising U-Net encoder through experiments applying a transformer architecture, which captures global dependency through self-attention. Additionally, we propose k-neighbor label smoothing, reverse boundary attention, and self-supervised learning with morphology-driven learning to improve the ability to identify complex structures. Our model, which analyzes the morphological representation of images, shows better results than the previous models in various medical imaging modalities, including CT, MRI, and lesion images. Our code and dataset are publicly available at githublink






§ INTRODUCTION

Medical image segmentation is crucial in improving our understanding of complex anatomy, providing critical insights for accurate medical diagnosis and precise treatment planning. This is especially important in computed tomography(CT) scans, where the intrinsic complexity of medical images presents unique challenges that require sophisticated solutions for organ segmentation. Unlike general images, CT images are quantitative imaging, and pixel intensities are normalized to Hounsfield units (HU) values<cit.>. (e.g., air as -1000 HU, bone as +400 to +1000 HU). Therefore, clinicians must understand the quantitative meanings and select the appropriate range to enhance the visual contrast of specific tissues or organs. In particular, research is conducted to find appropriate range values for each tissue or organ in CT scans<cit.> and studies show that segmenting CT images with an inappropriate HU range normalized leads to poor performance<cit.>. ./tablesfig_drawbacks.tex This is because inappropriate normalization can occlude organs, as illustrated in Fig.<ref> (a). Additionally, clinicians can have different opinions in labeling<cit.>. Due to this, ground truths are not determinstic, and it may be difficult to obtain detailed representations of organ or lesion labels. Inaccurate manual labeling can further increase the complexity of organ segmentation, as shown in Fig.<ref> (b). We address intrinsic challenges in medical images with an architecture that combines the advantages of the adaptive and resilient Swin Transformer encoder<cit.> with the efficient decoder in UNet<cit.>. We break away from the conventional Denoising U-Net<cit.> structure because we need a model that captures a global contextual representation and can handle the various medical imaging data. In addition, we introduce three approaches to improve the segmentation process further. First, distance-aware label smoothing<cit.> is a guidance mechanism that recognizes anatomical locations in the medical image and smoothes labels by calculating location-based distances. Second, reverse boundary attention captures areas of subtle and ambiguous boundaries. This component contributes to more precise and accurate segmentation by explicitly directing the model attention to edges<cit.>, especially in the regions that have not been manually labeled. Third, self-supervised learning<cit.> allows complex features of organs to capture meaningful representations from input images in a scenario of insufficient data. We reduce reliance on labeled data and improve model adaptability to diverse and complex features of medical images. Our proposed method demonstrates generalizability beyond medical images when we evaluate it with a different domain task that can utilize morphological information. Therefore, our contribution is summarized as follows.

    

    
  * We presents a new diffusion transformer segmentation(DTS) model which performs better than previous framework(i.e. CNN based Denoising Diffusion Probabilistic Model).
    

    
  * We introduce a novel approach to address the medical image segmentation by integrating morphology-driven learning into the image processing, such as k-neighbor label smoothing, reverse boundary attention, and self-supervised learning.
    

    
  * Our model demonstrates the generality in segmentation tasks in medical modalities such as CT,  MRI, and lesion images and further suggests that this approach may be adaptable to other domains.
    





§ RELATED WORK
 

The diffusion segmentation model which applies the generative diffusion process, allows users to manipulate the ambiguity of each time step through a hierarchical structure, solving the image quality and diversity problems of existing methods, allowing the learning process to proceed stably. There is research that has notable potential applied to medical imaging<cit.>. SegDiff<cit.>, which showed consistent performance under various imaging conditions, is the first approach to solving the image segmentation problem by applying diffusion. The feature of this model is a mechanism that integrates the information of the input image and the current estimate of the segmentation map through each encoder and uses the decoder to improve the segmentation map iteratively. MedSegDiff<cit.> also applied the diffusion segmentation model to medical image segmentation. The input of the conditional image and noise segmentation map are integrated using a mechanism such as SegDiff, but high-frequency noise is constrained through the Fast Fourier transform module during the connection process. In addition, Diff-UNet<cit.> implemented the standard U-shaped architecture, which learns from the input volume in medical image segmentation effectively to extract semantic information. Here, we focus on the architecture and compare it with the existing diffusion segmentation model to demonstrate through experiments that the inductive bias, which is a major feature of CNN, can be replaced by ViT in diffusion segmentation.

Label Smoothing for Image Segmentation. Ground truth labeling for image segmentation is a time-consuming and intensive task involving experts. These processes are inherently subjective and susceptible to factors such as image quality, observer diversity, and difficulty depicting specific structures. Moreover, earlier label smoothing methods<cit.>, the inter-class relationships are usually overlooked since the labels are smoothed into one-hot encoding vectors. To address these challenges, we experimentally highlight the implementation of strategic label smoothing based on the spatial location of organs.

Reverse boundary attention refers to the integration of reverse attention<cit.>, which learns opposite concepts that are not associated with the target class in a way that substitutes existing attention mechanisms for objects, and boundary attention<cit.>, which emphasizes pixels or features of parts related to the boundary. This mechanism plays a crucial role in enhancing the performance of object segmentation in medical images. This is particularly important because medical imaging, such as CT and MRI scans, often exhibit ambiguous organ boundaries and significant amounts of noise, posing challenges for accurate segmentation. Therefore, we explore the benefits of combining unique advantages, such as a reverse boundary attention mechanism, into our framework.




§ DTS: DIFFUSION TRANSFORMER SEGMENTATION

 The diffusion model is a generative model that consists of two stages: a diffusion process and a denoising process. In the diffusion process, Gaussian noise is added incrementally to the segmentation label x_0 over a series of steps t. 
 
./tableseq_diffusion.tex
./tablesfig_architecture.tex

./tablestable_feature_extract.tex The denoising phase shown in Fig.<ref> follows the encoder-decoder network structure of standard denoising autoencoder<cit.>. As shown in Table.<ref>, we empirically suggest the possibility of replacing the latent diffusion encoder with a Swin transformer<cit.>, which has advantages such as scalability and computational efficiency when processing various images due to its hierarchical structure. Also, similar to the conditional mechanism<cit.>, our model incorporates another type of conditional encoder, τ_θ where the original image is used as input. These are demonstrated in the diffusion encoder and conditional encoder in Fig.<ref>. Our method combines information from the current estimate x_t, the image I, and the time step index t to adjust the step estimate function ϵ_θ at the input. It also takes the conditional image τ_θ(I)) and reconstructs it through a UNet decoder to produce the global feature map. Subsequently, the RBA modules facilitate the derivation of the final segmentation map, which exhibits precise edge representation, as detailed in Fig.<ref>. In conclusion, DTS(·) represents our novel diffusion transformer segmentation model, which performs segmentation by integrating the described components.



 
    ϵ_θ(x_t,I,t) = DTS((x_t,I),t, τ_θ(I))





§ MORPHOLOGY DRIVEN LEARNING


k-Neighbor Label smoothing by organ distance.
We explore medical data from body parts such as the abdomen and brain, which have organs or diseases located structurally within a compact space. As the relative positions of organs do not differ from person to person, ./tablestable_ls_scalefactor.tex we propose a k-neighbor label smoothing method that leverages the relative positions of organs for distance-aware smoothing of the labels of k neighbors for a given class or organ. In a multi-class (k>2) situation, such as in this case, there is an advantage if there is a positional relationship between them. The positional relationship refers to the relative positional relationship of organs anatomically. As shown in Fig. <ref> (b), the liver(⋆) is close to the gall bladder(▴) but relatively far from the left kidney(▪). We provided semantic information to the model based on which body structure would match this prior knowledge.  
./tableseq_k-nls.tex
./tablesfig_k-nls.tex

./tablesfig_RBA.tex
RBA: Reverse-boundary Attention.
Complex anatomy and the inherent ambiguity in defining boundaries of adjacent organs are factors that hinder accurate segmentation of organ boundaries in medical images. Considering that these factors are likely to result in false positives or missing details in the initial segmentation, our approach includes selectively dropping or reducing the prediction weights of overlooked regions. The Reverse Boundary Attention method aims to improve the prediction of segmentation models by gradually capturing and specifying areas that may have been initially ambiguous. Thus, our architecture removes previously estimated predictive areas from high-level output features where existing estimates are upsampled in deeper layers, sequentially explores these details, including areas and boundaries, and finally, improves the segmentation model predictions progressively.
./tableseq_RBA.tex

./tablesfig_ssl.tex
Self-supervised learning (SSL)
 can encode anatomical information of the human body in the image effectively. We propose three proxy tasks for learning comprehensive semantic representations within masked images without using labels. Our framework combines (1) Contrastive learning(e.g. SimCLR<cit.>), which encodes masked images to improve the ability to distinguish between different samples with hidden feature representations; (2) Masked Location Prediction, which predicts the location of the samples; and (3) Partial Reconstruct Prediction(e.g. SimMIM<cit.>), which learns the feature representation by reconstructing the masked patch area of each sub-volume.  These widely recognized self-supervised learning strategies are both straightforward and effective.

 When the input(demonstrated 2D image in Fig.<ref>) is divided into patches and then passed as input to the encoder twice, two sets of latent embeddings are obtained, and a contrastive learning is performed through constrastive loss<cit.> (Eq.<ref>). Then, masked location prediction is conducted to predict the number of randomly masked parts by dividing the ĥ_̂î image into patches from 0 to 8 (Eq.<ref>).
In addition, Partial Reconstruct Prediction is performed by masking the image of ĥ_̂ĵ, reconstructing it through a decoder, and learning the difference from the original through L_2 loss (Eq.<ref>).

./tableseq_ssl.tex




§ EXPERIMENTS

Datasets. The pre-training dataset consists of medical images sourced from partial accessible CT, MRI datasets encompassing 3,358, 6,970 subjects respectively. Notably, the pre-training step does not involve the utilization of annotations or labels on this dataset. The primary objective of this pre-training process is to enable the model to learn meaningful representations from the available image data, thus eliminating the need for manual annotation. 
The BTCV<cit.> dataset comprises 3D abdominal multi-organ CT images from 30 cases, each associated with a specific form and featuring 13 multi class segmentation objectives.
The BraTS2021<cit.> dataset includes 1,251 subjects of brain MRI images. Each image is annotated with three segmentation targets and encompasses four modalities T1, T1Gd, T2, and T2-FLAIR. 
The ISIC2018<cit.> dataset contains 2,594 dermoscopic images of skin lesions, each annotated by experts for segmentation purposes. 
The Cityscapes<cit.> dataset contains several urban street scenes for segmentation purposes and is used to test the generalization performance of our approach. It consists of 3475 semantically annotated train, val sets and 1525 test set.
Details about datasets are shown in the appendix.

Implementation Details.
Our architecture implemented in PyTorch and MONAI[https://monai.io/https://monai.io/]. For pre-training tasks, the reconstruction strategy is applied with a mask ratio of 0.4. Moving on to the fine-tuning phase, the AdamW optimizer<cit.> is used with a weight decay 1e^-3. The warm-up is set to 0.1 of the total epochs, and the learning rate undergoes linear updates following the Cosine Annealing schedule<cit.>. The loss function incorporates DICE loss<cit.>, BCE loss, and MSE loss. Random flips, rotations, intensity scaling, and shifts were applied to augment the data. We set the number of diffusion steps as 1000, and the sliding window overlap rate is 0.8 until the final prediction. Preprocessing details for each dataset are provided in the appendix.

Evaluation Metrics are important to quantify the performance of the segmentation model. Two commonly used metrics are the dice similarity coefficient<cit.>(Dice) and the Hausdorff distance<cit.>(HD). The evaluation metric are as define. Y and Ŷ represent the actual and predicted values in input units, and g' and p' represent the actual and predicted values of points on the surface.
./tableseq_metric.tex

./tablestable_ls_ablation.tex
Exploring the performance of Label Smoothing
We concentrate on the performance of k-neighbor label smoothing and explore its applicability to general datasets with structural properties. We explore its applicability to a cityscapes<cit.> dataset with structural properties by utilizing only our baseline model and label smoothing. Compared with basic label smoothing(uniform), Non-Uniform Label Smoothing(NULS), and especially Spatially Varying Label Smoothing(SVLS), which applies label smoothing to neighboring pixels using weight matrix in the form of Gaussian kernel, we can see that our performance is superior. Previous methods compensate for the label’s uncertainty in image segmentation, but our methods further estimate the positional relationship between classes to improve prediction performance with a label smoothing method, emphasizing that this can be easily applied to other tasks.

./tablestable_ssl_ablation.tex
Efficiency of Self Supervised Objectives. We conduct comprehensive ablation experiments on the BTCV dataset to evaluate the efficiency of self-supervised learning. In these experiments, we employed specific settings for calculating the loss, and the obtained results are presented in the Table.<ref>. Notably, the L_Rec is learned based on the pixel representation of input images,  L_Loc (Masked Location Prediction) is learned by recognizing the location of the masked region, and L_CL (Contrastive Learning) is focused on contrastive learning at two augmented sample level. The L_Rec lies in its important role in understanding meaningful representation learning from medical images, as shown in experimental results. By employing these three loss functions, our self-supervised learning approach aims to capture intricate details at both pixel and region levels, enhancing the model's ability to extract meaningful features from the the inputs.

./tablestable_architecture_ablation.tex
Selecting the optimal architecture
Remind the our approaches, in the case of self-supervised learning (SSL), feature representations pre-trained from the three proxy tasks are transferred to the conditional encoder to assist in understanding the original image. We experiment with the effect of freezing or leaving all weights trainable during benchmark fine-tuning. Additionally, our framework explores the ablation study on reverse boundary attention, which is integrated with the general diffusion segmentation process. We comprehensively verify the effectiveness of morphology-driven learning within the architecture to prove its hypothesis. As shown in the Table.<ref>, the single module experiments(presented in the second section) show higher performance than the scratch model, but learning by leaving the conditional encoder trainable shows a large margin in the BTCV dataset. This indicates that feature representation was achieved by aligning the learned features well with the downstream task. Our model, which comprehensively combines morphology-driven learning techniques, shows remarkable improvement in results, and our final architecture is shown in Fig. <ref>.

./tablestable_quan_result_btcv.tex
./tablesfig_qual_result_btcv.tex



§ COMPARATIVE RESULTS

As shown in Table.<ref>, we compare our model with the BTCV benchmark dataset. Compared with other models, the proposed DTS achieves the best performance and presents a higher dice result of 0.906. It can be seen that previous diffusion segmentation models show comparable performance to conventional segmentation models in relatively large organs(e.g. liver, stomach), but poor performance in small organs(e.g. esophagus, aorta). DTS surpasses the closest competing methods by an average of 2% across all classes, with an even more significant improvement of 7% specifically for gall bladder. We believe that our approach and the application of the high performance transformer architecture will lead to improved accuracy. Comprehensive qualitative results of our model, which demonstrate good segmentation performance for small organs, can be found in Fig.<ref>, highlighting our model's ability to capture details and achieve accurate boundary representations.

./tablestable_quan_result_brats_isic.tex
The results presented in Table.<ref> demonstrate that the two datasets showed optimal outcome with an average accuracy in terms of both Dice and HD score. In particular, within the ISIC dataset, solely K-neighbor label smoothing was omitted from the application. This decision was made due to the dataset has only a single label without structural position relationships between adjacent labels. Consequently, employing the K-neighbor label smoothing method in this specific scenario is unnecessary. Overall, SwinUNETR <cit.> has a competitive performance in the benchmark results. Although it uses a similar architecture as DTS, we believe that it does not lead to the robustness and of diffusion models, which excel against noise and artifacts in the input data, particularly in medical images.




§ FUTURE WORK AND CONCLUSION

Our study focuses on the advantages of morphology-driven learning for segmentation tasks, where our approach demonstrates substantial improvements. Building on these promising results, we aim to broaden the scope of our framework by applying it to other critical imaging tasks, such as classification and detection, to evaluate its effectiveness across various domains and imaging scenarios.
In conclusion, we present a novel approach to medical image segmentation. DTS suggests the potential to replace existing CNN-based down-sampling by using a Swin Transformer encoder. We believe that this model architecture enables accurate segmentation with small, detailed representations and improves performance by complementing the chronic problems of medical images with Morphological-based learning, such as k-neighbor label smoothing, reverse boundary attention and self-supervised learning. We hope that this inspires future tasks in situations with morphologically complex problems.







§ ACKNOWLEDGEMENTS

This research was supported by the MSIT(Ministry of Science and ICT), Korea, under the ITRC(Information Technology Research Center) support program(IITP-2024-2020-0-01789), and the Artificial Intelligence Convergence Innovation Human Resources Development (IITP-2024-RS-2023-00254592) supervised by the IITP(Institute for Information & Communications Technology Planning & Evaluation). Additionally, this work has been carried out with (extensive) use of the NEURON computing resource that is supported by the Korea Institute of Science and Technology Information (KISTI).




 Paper Title: VITAL: Visual Teleoperation to Enhance Robot Learning through Human-in-the-Loop Corrections.


Bimodal Scaling Law and Size Effect In Superelastic Nanopillars
    Xian Chen
    August 5, 2024
===============================================================




§ ABSTRACT 
   Imitation Learning (IL) has emerged as a powerful approach in robotics, allowing robots to acquire new skills by mimicking human actions. Despite its potential, the data collection process for IL remains a significant challenge due to the logistical difficulties and high costs associated with obtaining high-quality demonstrations. To address these issues, we propose a low-cost visual teleoperation system for bimanual manipulation tasks, called VITAL. Our approach leverages affordable hardware and visual processing techniques to collect demonstrations, which are then augmented to create extensive training datasets for imitation learning. We enhance the generalizability and robustness of the learned policies by utilizing both real and simulated environments and human-in-the-loop corrections. We evaluated our method through several rounds of experiments in simulated and real-robot settings, focusing on tasks of varying complexity, including bottle collecting, stacking objects, and hammering. Our experimental results validate the effectiveness of our approach in learning robust robot policies from simulated data, significantly improved by human-in-the-loop corrections and real-world data integration. Additionally, we demonstrate the framework's capability to generalize to new tasks, such as setting a drink tray, showcasing its adaptability and potential for handling a wide range of real-world bimanual manipulation tasks. A video of the experiments can be found at:  
    




 





§ INTRODUCTION


Imitation Learning <cit.> has emerged as a fundamental approach in robotics, enabling robots to acquire new skills by observing and mimicking human actions <cit.>. This method holds significant promise for developing complex robotic behaviors without the need for explicitly programming the robot. However, the primary challenge with imitation learning lies in the data collection process <cit.>. On the one hand, acquiring high-quality demonstrations is costly and logistically challenging. On the other hand, the need for extensive and diverse demonstrations to train models magnifies these difficulties, as each new task or environment often requires a fresh set of demonstrations <cit.>. To address these challenges, teleoperation has been proposed as a viable solution. In recent years, there has been a growing interest in the development of bimanual teleoperation interfaces that facilitate the learning and execution of various household and industrial tasks <cit.>. 


One example of such a system is the ALOHA platform <cit.>, which has garnered considerable attention from the robotics research community. Such teleoperation platforms enable the execution of various tasks such as pushing chairs, using cabinets, wiping spills, and even cooking shrimp. Despite the advantages offered by such interfaces, there remain significant challenges that hinder their widespread adoption. One of the primary issues is that such teleoperation interfaces are hardware-specific and costly. The cost of such interfaces is approximately $30K, which makes them non-scalable and expensive for many research laboratories and practical applications. To overcome these challenges, we propose a low-cost visual teleoperation interface to collect high-quality demonstrations for long-horizon bimanual manipulation tasks.


In robot manipulation tasks, IL is often trained via in-domain demonstrations, which correspond to the data obtained directly from the deployment environment <cit.>. Although it is widely accepted that gathering direct examples from the actual robot is the most efficient way to learn long-horizon manipulation tasks, we show that collecting human demonstrations in real and digital twin environments can produce better outcomes for real-world tasks. Our approach allows the creation of a large dataset of augmented demonstrations based on a few initial demonstrations. By leveraging the flexibility and functionalities of our digital twin environment, we can enhance the number of demonstrations and the generalizability of the learned policy. These augmented demonstrations are then used to train a policy, which is fine-tuned using real-world demonstrations. Experimental results show that our approach significantly improves the performance of the robot in real-world scenarios. In summary, the main contributions of this work include:


    
  * We introduce a low-cost visual teleoperation approach that enables high-quality data collection for bimanual manipulation tasks.
    
  * We leverage a digital twin of the robot to create a vast dataset of augmented demonstrations from a limited number of initial demonstrations. This enhances the scalability and generalizability of the learned policies without requiring extensive real-world demonstration.  
    
  * We incorporate human-in-the-loop corrections to refine and improve the robot's performance in complex tasks. This method allows for interactive learning and adaptation, ensuring higher success rates in real-world applications.





























§ RELATED WORK

Robot Learning with Teleoperation: IL is a powerful technique for robot learning, enabling robots to acquire new skills by mimicking human demonstrations <cit.>. However, collecting in-domain demonstrations is often difficult and not scalable due to the need for physical interaction with the robot hardware. This process is also expensive as it requires constant access to the robot. Some works use scripted “expert” in simulation to generate demonstrations, but scaling these approaches to more complex tasks is challenging <cit.>. Other IL methods leveraged visual deep neural networks to train reactive models from images captured during demonstrations <cit.>. While these models provide greater flexibility, they necessitate collecting a large number of demonstrations to learn simple pick-and-place tasks, making them user and computationally expensive <cit.>. Other approaches utilize human activity videos to guide robot learning but often fall short in training 6-DoF manipulation policies <cit.>. These methods typically require in-domain teleoperation data to bridge the gap between video demonstrations and robotic execution <cit.>.
A common problem with these IL methods is the lack of corrective feedback, leading to the accumulation of errors over the length of the demonstration, known as the compounding errors problem <cit.>. Similar to our approach, some methods have focused on dividing demonstrations into several waypoints to address this issue <cit.>. Unlike these approaches, we aim to automatically extract the waypoints of a demonstration without adding any burden on the human user. Other IL methods used bimanual teleoperation interfaces, such as ALOHA <cit.>, for mapping joints' value from master to slave robots. However, these approaches are quite expensive and not scalable as they require four robots for bimanual manipulation: two masters and two slaves. In contrast, we propose a low-cost teleoperation interface that allows controlling robot end-effectors to perform complex bimanual tasks and collect demonstrations in both real and simulation environments. Teleoperation within a digital twin of the robot enables us to scale up the training data easily without constant access to the real robot. Additionally, we introduce a human-in-the-loop correction mechanism to repair faulty demonstrations, enhancing the robot's performance in complex tasks with minimal effort. 

Sim2Real Transfer: Some approaches, such as MimicGen <cit.> and CyberDemo <cit.>, attempt to collect demonstrations in simulation to reduce the cost of collecting demonstrations on robots. However, due to the sim-to-real gap, deploying the learned policy on a real-world robot remains challenging. Other methods use VR for teleoperation and data collection <cit.>, such as DIME <cit.>. Qin et al. <cit.> and Wang et al. <cit.> proposed tracking the user's figures through an external camera for teleoperation. Approaches like DexMV <cit.>, DexVIP <cit.>, and VideoDex <cit.> use everyday videos to learn motion priors. Due to the domain gap, these methods are not directly applicable to robots and often require another round of training.
In contrast, we collect demonstrations in the digital twin of the robot and develop data augmentation techniques at the trajectory level rather than the image data level. Most of the teleoperation interfaces have been designed either for simulation <cit.> or for real robot  <cit.> setups. Unlike these approaches, we proposed a low-cost bimanual teleoperation interface that can be used to collect demonstrations using real-robot and its simulation. We showed that policies trained on a mix of simulated and real demonstrations achieve better results than those trained solely on simulation or real-robot data. 





§ METHOD





 §.§ Collecting Demonstrations through Teleoperation

r0.4


    
    < g r a p h i c s >

    Overview of the proposed low-cost teleportation interface.
    
    


    
The architecture of our low-cost visual teleoperation system is illustrated in Figure <ref>. 

We use an external RGB-D camera to capture human motions and two Bluetooth-enabled selfie sticks for commanding the left and right grippers. The captured human motion data and gripper commands are translated into real-time motor commands for the robot arm and the gripper. In particular, we utilized the Mediapipe library for skeleton tracking, specifically focusing on 24 upper body keypoints. These keypoints are crucial for capturing the operator's movements accurately. To establish a stable huamn-world reference frame, we considered the left shoulder, right shoulder, left hip, and right hip keypoints. We define the origin of this reference frame between left and right hip keypoints. This reference frame serves as the basis for calculating the pose of each hand relative to the world frame. Given the differences in body morphology and arm length between humans and the robot, two scaling factors, α for the right arm and β for the left arm,  are applied to bridge the gap, thereby making teleoperation more intuitive and precise. For controlling the grippers of the robot's arms, we utilized two Bluetooth selfie sticks. These sticks were adapted to send commands for opening and closing the grippers. 
r0.15
    

    
    < g r a p h i c s >

    Low-cost teleportation sticks.
    

    

To track the sticks' orientation, we attached AruCo markers <cit.> to the sticks as shown in Fig. <ref>.



To efficiently manage the complexity of long-horizon tasks and reduce the compounding errors, we employ task decomposition. Each task, T, is decomposed into a sequence of subtasks, S_n, T = (S_1, S_2, …, S_n). The start and end poses of the trajectory for each subtask are defined based on the current pose of the manipulator and the pose of the object that should be grasped or placed. 
 
The detection of these subtasks is facilitated by the commands given by the user for grasping or placing objects. When a user issues a command, the system identifies the relevant subtask, captures the necessary pose information, and creates a trajectory segment for that subtask. By structuring tasks this way, we ensure that the demonstration data is organized and can be effectively used for training purposes.

To replicate the physical environment in the simulation, we build a digital twin of our robot in Gazebo. This includes accurately modeling the robot and objects to ensure consistency between the simulated and real-world environments. During the demonstration phase, we performed a task in Gazebo while recording all relevant ROS topics, including the initial state of the objects, joint values, end-effector poses, outputs from RGB-D cameras (head-mounted and on the left and right arms), commands from the instructor, activity videos with corresponding 3D skeleton data, object poses, and success or failure indicators, at their respective publication rates into a ROS bag file. The comprehensive nature of this recorded data ensures that all necessary information is captured for the next stage of the methodology, where data augmentation will be applied to enhance the dataset.






 §.§ Augmenting Demonstration in Simulation


To enhance the collected demonstration and improve the robustness of the learned policies, we perform data augmentation on the trajectory level (i.e. subtasks). Our augmentation method uses several techniques to generate a diverse set of trajectories from the initial demonstrations. First, we extract waypoints from the recorded bag file by reading the positions and orientations of the robot's end effectors. These waypoints are then used to fit a polynomial trajectory. We exclude the last few points from the polynomial fitting to ensure the end of the trajectory which represents the affordance pattern is preserved. We then apply uniform sampling to the trajectory by calculating the cumulative arc length of the waypoints to ensure uniform spacing between the sampled points. 
This method helps in creating a smooth and evenly distributed trajectory, which is then segmented into smaller parts for detailed analysis and augmentation.

We first augment the trajectory by adding Gaussian noise to the waypoints. This involves introducing random variations to the trajectory points to simulate real-world imperfections and variations. By adding noise, we can create multiple versions of the trajectory, enhancing the model's ability to generalize across different scenarios. The noise can be applied to all waypoints or selectively to preserve certain critical points, such as the trajectory's endpoint.
We also perform translations and flips on the trajectory. Translating the waypoints involves shifting the entire trajectory in the x, y, and z directions, which helps in creating varied instances of the task. Flipping the trajectory along the x-axis generates a mirrored version, providing additional diversity in the dataset and making the initial trajectory of one arm useful for the other arm. Furthermore, we sample new initial points within specified bounds and connect these points to the nearest waypoints in the existing trajectory. This process involves finding the nearest index in the trajectory and creating new paths from these initial points to the rest of the trajectory. By varying the starting points, we can simulate different initial conditions for the task. Each trajectory is annotated with an arm identifier, distinguishing between left and right arm movements. This comprehensive augmentation process significantly expands the dataset, making it thousands of times larger than the original set of demonstrations. However, augmented trajectories can sometimes be imperfect and not executable due to the robot's arm kinematics, potentially leading to task failures. To ensure the reliability of the augmented data, we validate each augmented trajectory in the digital twin environment. Only successfully executed trajectories are used for training a policy. This way we ensure that the learned policies are robust and capable of handling a wide range of real-world conditions.






 §.§ Hierarchical Policy Learning




To effectively learn and execute long-horizon tasks, we formulate the learning policy as a hierarchical policy learning problem through Behavioral Cloning (BC). This approach involves training a high-level policy to select a sequence of subtasks and low-level policies to execute each subtask. 

For a given task, we assume that we can detect the sequence of subtasks based on the demonstration by detecting commands for grasp or release (indicated by pressing the bottom of the selfie stick). These commands allow us to identify the start and end poses for each subtask, facilitating the decomposition of long-horizon tasks into a series of subtasks.

The high-level policy, π_H, is responsible for selecting the appropriate subtask, S_i, based on the current state of the environment and the task progression. This high-level policy can be implemented as a state machine or a learned model. In this work, we use a state machine that receives the current state of the robot, and the pose of the objects to infer the current subtask, π_H(o, r)→ S_i, where o shows the current state of objects and r is the state of the robot. Alternatively, a learned model could receive the task ID and the state of the objects and robot, and predict the subtask.



For the low-level policy, we propose a unified multi-subtask policy learning, π_MS, that can handle all subtasks. This policy receives the predicted subtask ID, S_i, along with the start, p_s, and end poses, p_e, and generates the appropriate trajectory conditioned on the input subtask, π_MS(S_i,p_s, p_e) → trajectory. The training process begins by loading the augmented trajectory data of a task. The trajectories are padded to ensure uniform length, which is critical for batch processing in neural networks. Our model is designed to predict the intermediate points of a trajectory given the start and end points. The high-level policy coordinates these subtasks, ensuring smooth transitions and overall task completion.










 §.§ Residual Learning with Human-in-the-Loop


r0.5
    


    < g r a p h i c s >

    Overview of the residual learning based on human-in-the-loop feedback.
    
    


Despite the robustness of the trained policy, there can be scenarios where the robot fails to complete the task due to a sim-to-real gap, pose estimation errors, or small errors in executing a trajectory. To address these issues, we incorporate a residual learning approach with human-in-the-loop corrections during the policy rollouts. When the robot encounters a failure during task execution, a human operator steps in to correct the errors. The user plays the recorded rosbag file of the experiment and provides the correction feedback whenever it is needed. The correction involves measuring the 3D delta position changes of the wrist of the left and right hands of the demonstrator relative to the initial poses of the hands. These delta changes are then scaled by empirically determined factors α (for the right hand) and β (for the left hand), both < 0.1 to avoid fast movements, and added to the robot's wrist movements. This minimal adjustment helps the robot to accomplish the task with improved accuracy. If the user observes a large error in the current prediction of the policy, they can switch the system to full teleoperation mode. In this mode, the robot will directly follow the human wrist movements to accomplish the task. 

The data from these human-in-the-loop corrections is recorded into a separate dataset, 
D'. This dataset includes both residual corrections and instances where full teleoperation was required. To fine-tune the policy, training data is sampled equally from both the original dataset D and the new dataset D'. This balanced sampling ensures that the model learns from the augmented experiences and corrections provided by the human operator, thereby improving its performance and adaptability. This iterative process of learning ensures that the policy continuously improves and adapts to real-world conditions.




§ EXPERIMENTAL SETUP



We performed multiple rounds of experiments on both simulation and real-robot settings to validate our method. Specifically, we study the following questions: (Q1) Is it feasible to train robot policies using exclusively simulation data without relying on on-robot data? 
(Q2) Which model architectures and data composition ratios are most effective for policy training? Specifically, what is the optimal balance between simulated and real-world data?
(Q3) To what extent does human-in-the-loop correction enhance policy performance when teleoperation data is limited?
(Q4) Can our framework successfully execute another complex bimanual task, such as setting a drink tray, which is somewhat similar to one of the learned tasks?



 §.§ Setup
 
Our experimental setups, both in simulation and on real robots, are illustrated in Fig. <ref> and Fig. <ref>, respectively. We developed a simulation environment in Gazebo, leveraging the ODE physics engine to closely mimic the behavior of our dual-arm robot. The hardware setup in both simulation and real robot includes an Asus Xtion camera, two Universal Robots (UR5e) equipped with Robotiq 2F-140 grippers, and an interface for managing the initiation and conclusion of experiments. For the Teleoperation interface, we use a RealSense D435 RGB-D camera to capture the motion of the user. All evaluations were performed with a PC running Ubuntu 18.04 with a 3.20 GHz Intel Xeon(R) i7, and a Quadro P5000 NVIDIA. 

r0.38


    
    < g r a p h i c s >

    Three long-horizon evaluation tasks are used. (left) Bottle Collecting; (center) Stacking Pringles; (right) Hammering.
    

    

Tasks: We designed three tasks of varying difficulty levels to test the performance of the teleoperation and the robustness of the learned policies: (i) Bottle Collecting: In this task, the robot should collect multiple bottles with different shapes from different locations. (ii) Stacking Pringles: In this task, the robot must stack Pringles cans on top of each other in various locations. This requires precise control and coordination to ensure the cans are aligned and stable. 
(iii) Hammering: The most challenging task involves the robot using a hammer to strike a cylinder. This task tests the robot's ability to handle a tool effectively. To evaluate the performance of the robot on each task, we performed 100 simulation and 10 real-robot experiments. In the case of real robot experiments, we used a 3D perception system to get the pose and label of the objects <cit.>.
r0.38


    
    < g r a p h i c s >

    (left) Our real dual-arm robot setup; (right) Output of our perception including the bounding box, pose, and label of the objects.
    

    





Baselines and Metrics: To determine the best model architecture for bimanual manipulation tasks, we evaluated multiple Behavioral Cloning Recurrent Neural Network (BC-RNN) baselines. The baselines include: LSTM (Long Short-Term Memory), GRU (Gated Recurrent Unit), and Transformer-based network. PyTorch is used as a deep learning framework. We use Mean Squared Error (MSE) as the loss function for trajectory prediction. An Adam optimizer is employed to adjust the model weights during training. Early stopping is implemented to halt training if the validation loss does not improve for a specified number of epochs. Each network has been trained for 200 epochs. We split the data into training (70%), validation (15%), and test sets (15%) to evaluate the model's performance. Details of the networks are discussed as a part of supplementary materials. The Task Success Rate is used as the primary metric for evaluation. This metric is defined as the percentage of successfully completed tasks out of the total number of attempts. 





 §.§ Results



Q1 - Feasibility of Training Policies with Simulation Data:
We conducted extensive sets of experiments across all tasks by training the policy solely on augmented simulation trajectories. For each task, we collected five successful in-domain demonstrations and augmented them to 80K demonstrations. The BC-LSTM policy trained on this augmented data was then tested on both simulated and real robots. Results are summarized in Fig. <ref>. In simulation, the robot successfully performed most tasks. However, in real-world experiments, 
r0.38


 
    < g r a p h i c s >

    Results of Q1 evaluation.
    

    

certain discrepancies such as minor inaccuracies in object grasping and placement were observed. Specifically, in the bottle collecting task, the robot achieved an average task success rate of 92% in simulation and 80% in real-robot scenarios. In the stacking objects task, the robot successfully completed 84% of the experiments in simulation but only 60% in real-world tests. Similarly, for the hammering task, the success rate was 78% in simulation and 50% in real-world experiments. The main failure in bottle collecting was due to inaccurate trajectory prediction, which led to inefficiencies in approaching, grasping, or placing the bottles. The primary reasons for failures in the stacking and hammering tasks were inaccuracies in object grasping. Specifically, when the robot approached to grasp the object, the tip of the gripper touched the object, causing it to move slightly. This offset often led to failures as objects falling during the last step of stacking or missing the target when hammering. To address these challenges, we will check the impact of the model architecture and the effect of incorporating real demonstrations in addition to digital-twin demonstrations for training the policy to address sim-to-real gap in the following subsection (Q2). Furthermore, we believe incorporating feedback mechanisms through vision sensing could provide information on the exact part of the object being grasped. This feedback would enable the system to update the trajectory accordingly by learning suitable residual corrections, thereby improving the accuracy and reliability of the task execution. 













Q2 - Impact of Model Architecture and Data Composition: To understand the influence of different model architectures and the ratio of simulated to real data, we conducted experiments with various configurations. We first trained LSTM, GRU, and Transformer-based models solely on augmented simulation demonstrations and compared their performance across the proposed tasks in the simulation environment. After identifying the best-performing architecture, we trained policies using different ratios of simulated to real-robot demonstrations, including 70%-30%, 50%-50%, 30%-70%, and 0-100% mixes. For each task, we recorded five real-world demonstrations and five simulation demonstrations, which were then augmented to create a larger dataset. The final training data was sampled from these augmented demonstrations according to the specified ratios. This approach allowed us to determine the optimal balance between simulation and real-world data for training robust and effective robot policies.
The LSTM model provided a good balance of performance and training efficiency, whereas GRU models exhibited slightly lower performance but benefited from reduced computational overhead. Transformer-based models demonstrated considerable potential for handling more complex tasks; however, they demanded significantly more computational resources (see Fig. <ref>). 
r0.38


    
    < g r a p h i c s >

    

    Results of Q2 evaluation.
    

    

Upon comparing the results, we identified that the optimal ratio of simulated to real data was 70% simulated and 30% real. This ratio offered the best generalization and performance in real-world tasks, enhancing the task success rate across all evaluated tasks. Specifically, for bottle collecting, the task success rate increased from 80% to 90%; for stacking Pringles, the success rate rose from 60% to 80% (see Fig. <ref>); and for hammering, the success rate improved from 50% to 70%. This balanced approach effectively mitigated the sim-to-real gap, thereby enhancing the robustness and reliability of the learned policies.

Q3: Effectiveness of Human-in-the-Loop Corrections:
In this round of experiments, we investigated the effectiveness of human-in-the-loop correction in scenarios when the policy failed during task execution. We recorded 5 failure experiments for each task in the simulation. We then replayed the bag file and allowed the human user to correct the robot’s movements using residual learning or full teleoperation mode. The updated demonstrations were tested to ensure they could successfully accomplish the task. These corrected demonstrations were then augmented 
r0.38


    
    < g r a p h i c s >

     

    Results of Q3 evaluation.
    

    

and used for policy fine-tuning. By comparing the results, we observed that human-in-the-loop corrections significantly improved task success rates, particularly for more complex tasks such as stacking and hammering. 

The inclusion of human corrections from dataset D', resulted in an average improvement of approximately 5% in task success rates, demonstrating the value of interactive learning. Specifically, the success rate for the hammering task increased from 78% to 86%, for the stacking task from 84% to 89%, and for the bottle collecting task from 92% to 95%. These results indicate that the policy was effectively learning from the corrections, as the number of correction feedback by the human user decreased over time. 

 

r0.38
    

    
    < g r a p h i c s >

    An example of Setting a Drink Tray task: (left) initial and (right) desired objects pose.
    

    

Q4: Capability to Execute A New Complex Bimanual Task: Setting a Drink Tray
To evaluate the generalization capabilities of our framework, we tested its ability to execute a new bimanual task that shares similarities with the learned tasks: setting a drink tray (see Fig. <ref>). This task involves placing a cup and a bottle on specific markers on a tray. The tray has two circular markers: one red for placing the cup and one blue for placing the bottle. The robot must pick up a cup and a bottle from a starting position and place them accurately on the designated red and blue markers on the tray. The policy was initially trained on the previously learned object stacking task. The trained policy was then evaluated only in simulation, as we could obtain precise positions of the markers and objects. The task was considered a success if the robot accurately placed the cup and bottle on the respective markers and a failure if the distance between the object and the marker exceeded 3 cm. We repeated this experiment 
100 times, with the objects placed in random positions each time. The robot achieved a task success rate of 87%. The primary source of error in the failure cases was slight misalignment during grasping, which led to minor placement inaccuracies, despite accurate trajectory predictions. Our approach successfully generalized to the new task, demonstrating that the learned policies could be adapted to similar but distinct tasks with minimal effort. This showcases the robustness and adaptability of our teleoperation and learning framework in handling a variety of real-world scenarios. We believe that incorporating feedback mechanisms through vision, such as detecting the exact positions of the markers and adjusting the trajectory accordingly, could further enhance the accuracy and reliability of task execution.





§ CONCLUSION AND LIMITATIONS



In this work, we presented a low-cost visual teleoperation system for bimanual manipulation tasks. Our approach leverages an RGB-D camera and visual processing techniques to collect demonstration data, which is then augmented and used to train robot policies. We employed a combination of simulated and real-world data, augmented with human-in-the-loop corrections, to improve the accuracy and reliability of the learned policies. Our system was evaluated on three tasks with differing complexity levels, demonstrating its effectiveness in both simulated and real-robot environments. The experimental results validate the effectiveness of our approach in learning robust policies from simulated data, enhanced by human-in-the-loop corrections, and real-world data integration. The framework successfully generalized to a new, complex task —setting a drink tray—demonstrating the adaptability of the learned policies to similar tasks with minimal effort. 

The main limitation of this work is that for tasks requiring a high degree of precision at certain steps (e.g., staking, hammering, etc), waypoint trajectory learning may not be sufficient, and (closed-loop) vision feedback should be considered. To address this limitation, we plan to explore the integration of visual feedback mechanisms to further enhance the accuracy and reliability of task execution. By incorporating real-time visual feedback, the system can adjust and refine its actions based on the visual context, leading to more precise and robust manipulation capabilities. 








  





§ APPENDIX






§ IMPACT OF MODEL ARCHITECTURE AND DATA COMPOSITION (Q2)




















   


































 §.§ Training data composition

We trained several policies using different ratios of simulated to real-robot demonstrations, including 70%-30%, 50%-50%, 30%-70%, and 0-100% mixes. For each task, we recorded five real-world demonstrations and five simulation demonstrations, which were then augmented to create a larger dataset. The final training data was sampled from these augmented demonstrations according to the specified ratios. This approach allowed us to determine the optimal balance between simulation and real-world data for training robust and effective robot policies.







§ VISUALIZATION OF GROUND-TRUTH AND PREDICTED TRAJECTORIES 
 


We learned a policy using a single-shot demonstration and tested its performance on the validation set of the first subtask of the bottle-collecting task. Specifically, we plot 2D/3D visualization of predicted and actual trajectories for two randomly selected samples from the validation set. The network's performance is depicted through both 2D and 3D plots in Figure <ref>. Furthermore, we tested the performance of the policy on the digital twin of the robot by randomly placing a bottle object in front of the robot, retrieving the current pose of the end effector and the pose of the object, and forwarding them as input to the learned policy. We then executed the predicted trajectory and visualized the robot's performance. 

Figure <ref> shows the performance of the robot in three experiments, highlighting its ability to follow the predicted trajectory and accurately grasp the object based on the learned policy. The visualizations include both images of the robotic arm and plots of the predicted versus actual trajectories for the X, Y, and Z axes.



§ EXPERIMENTS



 §.§ Success and failure criteria for the proposed tasks


r0.35


 
    < g r a p h i c s >

    Sucesses and failures for stacking and hammering tasks.
    
    



For the collecting objects task, success is determined by placing all objects into the basket and failure is considered if one of the objects is not placed into the basket and remains outside. Figure <ref> shows examples of successful (left) and failed (right) attempts for the stacking and hammering tasks. In the stacking task, success is achieved when objects are stacked stably. In the hammering task, success criteria are based on hitting the target object accurately.  These criteria help in evaluating the performance and reliability of the robotic system in executing these tasks. In the case of setting up a drink tray task, success is considered when the robot places all items (such as a mug and a bottle) accurately on their designated spots on the tray, and failure is defined when items are not placed correctly on the tray, or the robot fails to place all items.






 §.§ Performance of the robot in simulation: known scenarios - known objects

We show the performance of the robot in the simulated environment where the robot interacts with known objects across known scenarios in Figure <ref>. The tasks include collecting bottles (top-row), stacking objects (second-row), and hammering (lower-row), all using the same objects from the collecting demonstration. The robot's ability to follow a predicted trajectory and complete these tasks accurately showcases the robustness of the learned policy.






 §.§ Performance of the robot: known scenario - new objects


In this section, we test the robot's ability to adapt to new objects in a familiar scenario. The experiment involves using bottles that differ in weight, size, and shape from the ones used in the initial demonstration. This assessment helps determine the robot's generalization capability when faced with varying object properties while performing the same task. Despite the differences in the object's physical characteristics, the robot successfully follows the predicted trajectory and completes the collection task. The sequence of snapshots in Figure <ref> demonstrates the robot's generalization ability to complete the collecting objects task with new objects.






 §.§ Performance of the robot in simulation: new scenario - new objects


We also evaluate the robot's performance in a completely new scenario using new objects. The experiment involves setting up a drink tray with various items, testing the robot's ability to adapt to a new task and handle unfamiliar objects. For this task, we utilized the stacking object policy and adjusted the subtask trajectory end poses to accomplish the task successfully.
Figure <ref> shows the sequence of snapshots detailing the robot's performance in setting up a drink tray. This image highlights the robot's capacity to generalize from its training and successfully perform tasks in new scenarios with new objects.






































 §.§ Perception of our real robot experiments

To discuss our real robot experiments, we provide an example of our experiment setup for the collecting bottles task (Fig. <ref>), for the stacking objects task (Fig. <ref>), and for the hammering task (Fig. <ref>). The robot's perception system is capable of identifying and localizing objects in its environment, which enables it to execute the task accurately.

The experiment demonstrates the integration of perception and action, where the robot uses its sensory inputs to detect various objects to guide its movements and complete the task of bottle collection.










 §.§ Failure cases during real robot experiments


During the real robot experiments, several failure cases were observed that can be attributed to the differences between simulation and real-world conditions (See Fig. <ref>). These failures provide valuable insights into the challenges of transferring learned policies from simulation to real robots. The main factors contributing to these failures include:

Gap between simulation and real robot: The objects used during the experiments in the real world differ in shape, size, weight, and material properties compared to those used in the simulation. These discrepancies can affect the robot's ability to accurately predict and execute the required trajectories, leading to unsuccessful task completion.


Inverse kinematics solver and command frequency: The inverse kinematics (IK) solver and the frequency of sending commands to the robot differ between the simulation and the real robot. In the simulation, the IK solver and command frequency are typically more idealized and may operate at higher rates, leading to smoother motions. In contrast, the real robot's controller might introduce slight delays due to various factors (e.g., network loads), causing deviations in the robot's movement and impacting task execution.

Trajectory prediction errors: Inaccuracies in the predicted trajectory can lead to the robot's end-effector being misaligned with the object, preventing successful grasping and manipulation. Poor grasping or inadequate grip strength can cause objects to slip during manipulation. 

Execution of actions: Using an open-loop controller during the execution of actions means that the robot follows the predicted trajectories blindly without feedback corrections. In real-world scenarios, this lack of feedback can lead to accumulated errors and deviations from the intended path, causing failures in task completion. In such an open-loop scenario, accurate calibration between the robot's camera and its manipulator is crucial and calibration errors can result in incorrect positioning and unsuccessful task execution.


These factors highlight the importance of addressing the sim-to-real gap and improving the robustness of robotic systems when transitioning from simulation to real-world applications. Future work should focus on integrating closed-loop feedback controllers to mitigate some of these issues and achieve more reliable performance in real-world scenarios.






