# Continuous Training and Fine-tuning for Domain-Specific Language Models in
Medical Question Answering

Zhen Guo
MIT EECS
Cambridge, MA 02139
<EMAIL_ADDRESS>
Yining Hua
Harvard School of Public Health
Boston, MA 02139
<EMAIL_ADDRESS>

###### Abstract

Large language models exhibit promising general capabilities but often lack
specialized knowledge for domain-specific tasks. Developing domain experts
from a base model enables a range of applications without prohibitive training
costs. This work demonstrates a method using continuous training and
instruction fine-tuning to rapidly adapt Llama 2 base models to the Chinese
medical domain. We first conduct continuous training on 1B tokens from Chinese
medical references to teach relevant vocabulary and knowledge. The models are
then fine-tuned on 54K examples sourced from the Chinese National Medical
Licensing Examination. Experiments on Chinese medical data confirm the
effectiveness of this approach, producing a model comparable to GPT-3.5-turbo
while using way less computational resource. The resulting domain-specific
model could be useful for various Chinese medical applications. More broadly,
this provides a template for domain-specific training of large language models
in areas where pre-trained models lack the required expertise, such as law,
science, and engineering.

##  1 Introduction

Large Language Models (LLMs) have demonstrated impressive performance on a
wide range of general language tasks in recent years. Models like GPT-3, PaLM,
and Bard have achieved state-of-the-art results on applications including
translation, question answering, and program synthesis [1, 2, 3, 4], enabled
by massive training data, extensive pre-training, and rapid scaling [5, 6].
However, despite exhibiting excellent performance in the general domain, LLMs
often fail to capture the specialized expertise required for fields such as
medicine, law, or science [7, 8]. As a result, LLMs are subject to poor
performance on domain-specific applications, such as medical question
answering or legal search [9, 10, 11, 12].

Training domain-specific LLMs from scratch poses massive computational
challenges, which are prohibitively expensive for most organizations [13]. In
response, this work presents a pragmatic methodology for the rapid adaptation
of general language models to specialized domains using continual training
[14, 15, 16] and instruction fine-tuning strategies. We focus on the
healthcare and biomedicine domain, where the demand for proficient LLMs is
hampered by stringent data privacy regulations and the limited computational
resources typical of healthcare and biomedical labs [17, 18]. Concurrently,
recognizing the linguistic inequity in NLP [19, 20] \- primarily skewed
towards English - we focus on Chinese, thus broadening the global utility and
inclusivity of AI applications in healthcare. This dual challenge elevates the
task difficulty and therefore presents a valuable evaluation for our proposed
methodology.

In detail, we develop a Chinese medical LLM based on the open-source Llama
model [13]. The adaptation process is carried out in two stages. Initially,
continual pretraining is employed on a corpus of 1 billion Chinese medical
tokens derived from question-answer pairs. This aims to teach the model
general medical knowledge. Subsequently, the model is fine-tuned on authentic
medical licensing exam data with Alpaca format [21] for medical question-
answering.

The model developed with our method demonstrates capabilities comparable to
those of GPT-3.5 but using significantly reduced computational resources for
training. Compared to the original models, Llama-2-13B fine-tuned on CMExam
with the question reasoning increased from 14.1% to 38.8% in F1 score, while
Chinese-Llama-2-13B increased from 23.8% to 44.1%. Furthermore, additional
gains were achieved with continuous training on 1 billion Chinese medical
tokens before fine-tuning, with Llama-2-13B and Chinese-Llama-2-13B reaching
F1 scores of 43.3% and 45.7% respectively. This result highlights the
effectiveness of leveraging continual learning and fine-tuning to impart
domain expertise and extended vocabulary into general models, thus presenting
a viable framework for constructing specialized models for the biomedical
domain. Moreover, the methodology presented in this work could be extended to
other specialized domains. Examples include law, science, and engineering,
where base models typically lack domain-specific knowledge.

##  2 Method

###  2.1 Base models

This work utilizes two base language models. The first is Llama 2 13B, which
is pre-trained on 2 trillion tokens with a 4096 context length [13]. The
second base model is a Chinese version of Llama 2 13B, initialized from the
original Llama 2 and continuously trained on a mixture of Chinese and English
corpora [22, 23]. The Chinese data with sources such as encyclopedias and
scientific papers provides knowledge about the Chinese language. The English
sources, including datasets like SlimPajama [24] and RefinedWeb [25], maintain
existing knowledge in the base Llama model.

###  2.2 Continual training dataset

The continual training dataset includes more than 364,000 question-answer
pairs extracted from Chinese medical encyclopedias and expert articles online.
These data are part of the larger Huatuo-26M dataset [26], which contains
around 1 billion tokens. The encyclopedia excerpt focuses on formal medical
terminology and topics like andrology, gynecology, and infectious diseases
rather than consultation records. The continuous training dataset could easily
be expanded by adding more medical domain data, from open-sourced datasets in
HuggingFace or proprietary datasets [27, 28, 29], given sufficient
computational resources.

###  2.3 Fine-tuning dataset

CMExam is a large-scale Chinese medical exam dataset comprising more than
54,000 multiple-choice questions that span a wide range of medical knowledge.
The questions cover 27 disease categories and 36 clinical departments, serving
as a valuable resource for advancing medical language understanding [30].
CMExam also includes explanations for most of the questions, providing
additional knowledge for fine-tuning LLMs by enabling them to incorporate
reasoning capabilities during the training process.

While this work focuses on evaluating medical knowledge, CMExam could
potentially be combined with other resources, like dialogues from Li et al.’s
work[31], to fine-tune conversational agents for patient-doctor interactions.
However, best practices for evaluating AI chatbots remain an open research
area, as discussed in [32].

###  2.4 Evaluation datasets

To evaluate the model’s knowledge of the Chinese medical domain, we perform
greedy decoding with a maximum generation length of 256 tokens on the 6,811
question-answer pairs in the CMExam test set [30]. For base models not fine-
tuned on the CMExam training set, including Llama-13B and Chinese-Llama-13B,
we use a custom script to extract answers from the model outputs. The script
first attempts to directly match the answer options in the model output. In
case of failure, it performs a case-insensitive search for any answer strings
provided in the answer options. As a last resort, the Levenshtein distance is
used to find the closest matching answer to the first line of model
prediction. For all other models, we use the same hard match evaluation as the
original CMExam paper, comparing the ground truth answer to the predicted
answer.

To assess catastrophic forgetting of general knowledge [33], we use MMLU [8]
for English and CMMLU [34] for Chinese, with 5-shot prompting for both. These
general knowledge benchmarks allow us to measure whether medical fine-tuning
impacts the retention of broader world knowledge in LLMs. C-Eval [35] could
also provide valuable information, we did not include it in this study due to
installation issues with our cluster.

###  2.5 Hyperparameters for training

We continuously train and fine-tune the models using the HuggingFace
Transformers library on a cluster consisting of 20 nodes each with 6 NVIDIA
V100 32GB GPUs, with Fully Sharded Data Parallelism (FSDP) to distribute the
model across nodes [36]. The key hyperparameters are summarized in Tab. 1.

Table 1: Hyperparameters for Training Hyperparameter Name | Value
---|---
Learning rate | $2\times 10^{-5}$
Batch size | 120
max seq length | 4096
Number of epochs | 1
Warmup ratio | 0.03

Techniques such as mixed precision fp16 [37], gradient checkpointing [38], and
AdaFactor optimizer [39] are used to improve training efficiency.

##  3 Results

###  3.1 Training with reasoning improves model performance

We first conduct full-parameter fine-tuning on the two based models:
Llama-2-13B and Chinese-Llama-2-13B. The fine-tuning data come from the CMExam
training set, using answer-only examples or question-answer pairs with
reasoning explanations. For models in the general domain, the performances of
medical question-answering are suboptimal except for GPT-3.5-turbo and GPT-4.
As shown in Table 2, training with explanations in the answer improves the
test accuracy for both models, even when evaluating accuracy alone, suggesting
the values of reasoning data in fine-tuning.

When fine-tuned solely on the CMExam training data, Chinese-Llama-2-13B
achieves comparable performance to GPT-3.5-turbo. In the CMExam paper, a fine-
tuned Llama-7B model obtained a lower accuracy of just 18.3%. This reduced
performance could be attributed both to the smaller model size and to the use
of prompt tuning [40] and low-rank adaptation [41]. While these approaches are
parameter-efficient, full-parameter fine-tuning generally enables better
performance.

Table 2: Performance Metrics for Fine-tuning Model Type | Model | More Finetuning Dataset | CMExam
---|---|---|---
Acc (%) | F1 (%)
General Domain | GPT-3.5-turbo | - | 46.4 | 46.1
GPT-4 | - | 61.6 | 61.7
ChatGLM-7B | - | 26.3 | 25.7
Llama-2-13B | - | 21.0 | 14.1
Chinese-Llama-2-13B | - | 26.7 | 23.8
Medical Domain | Llama-2-13B | CMExam_answer_only | 37.3 | 37.0
CMExam_with_reasoning | 39.3 | 38.8
Chinese-Llama-2-13B | CMExam_answer_only | 42.9 | 42.7
CMExam_with_reasoning | 44.5 | 44.1

###  3.2 Continual training improves domain knowledge at a cost

Next, we show the training loss in Fig. 1. The perplexity loss decreases
gracefully despite the distribution shift from pre-training to continual
training data. Notably, our continual training is around 1 billion tokens,
much smaller than the 2 trillion tokens used for Llama-2 pre-training.

Figure 1: Continual training loss on encyclopedia QA dataset

In Table 3, we show the CMExam results on different continual training
checkpoints. At each checkpoint, we also fine-tune the model with question
explanations (reasoning) associated with the questions. A trend emerges where
the model accuracy on CMExam tends to increase as continual training
progresses, signifying th learning of medical domain knowledge. However,
decreases are observed in the MMLU and CMMLU metrics over time, pointing to a
phenomenon of catastrophic forgetting of general knowledge. This is
unsurprising since we did not mix in samples from the original pre-training
data due to computational constraints. As a result, the model may have
progressively replaced general knowledge with new knowledge in the medical
domain. Overall, after 1 epoch of continuous training, the model exhibits only
minor decreases in general knowledge, along with a small improvement in
healthcare and biomedical understanding. This trend of forgetting general
knowledge and gains in domain-specific knowledge could continue if training is
extended over more epochs.

Table 3: Performance Metrics for Fine-tuning Continuously Trained Checkpoints Model | Trained Checkpoint | CMExam | MMLU | CMMLU
---|---|---|---|---
Acc (%) | F1 (%) | Acc (%) | Acc (%)
Llama-2-13B | 0 | 39.3 | 38.8 | 58.6 | 42.7
750 | 43.4 | 43.0 | 57.2 | 42.1
1500 | 43.1 | 42.8 | 57.2 | 42.1
2250 | 43.6 | 42.9 | 56.4 | 42.0
3000 | 43.8 | 43.3 | 57.0 | 41.8
Chinese-Llama-2-13B | 0 | 44.5 | 44.1 | 54.4 | 46.0
750 | 44.7 | 44.7 | 54.4 | 45.7
1500 | 45.5 | 45.1 | 54.0 | 45.1
2250 | 45.9 | 45.5 | 53.1 | 45.3
3000 | 46.0 | 45.7 | 53.1 | 45.2

###  3.3 Challenges in cross-lingual fine-tuning

Table 4: Performance Metrics for Fine-tuning with Mixed Datasets Model | Checkpoint | Finetuning Dataset | CMExam
---|---|---|---
Acc (%) | F1 (%)
Llama-2-13B | 3000 | CMExam | 43.8 | 43.3
3000 | CMExam+MedQA+MedMCQA | 42.8 | 43.4
Chinese-Llama-2-13B | 3000 | CMExam | 46.0 | 45.7
3000 | CMExam+MedQA+MedMCQA | 41.3 | 43.2

Inspired by Med-PaLM 2 [29], we attempt to further improve Chinese medical QA
performance by mixing English and Chinese medical question-answering data for
fine-tuning. Specifically, we combine the MedQA [42], MedMCQA [43], and CMExam
training sets and use this mixture to fine-tune the continuously trained
Llama-2 and Chinese-Llama-2 checkpoints.

However, contrary to expectations, this multilingual fine-tuning results in
degraded performance on the CMExam benchmark, especially for the Chinese-Llama
variants, as shown in Table 4. This observation implies that although
techniques such as ensemble refinement and instruction tuning provide gains
for English medical QA with Med-PaLM 2, they do not transfer directly when
applied to Chinese medical data.

Thus, it highlights the challenges of cross-lingual transfer learning.
Although methods from English tasks can inform model development in other
languages, adaptation is still needed to account for linguistic differences.
The optimal training strategies likely vary across languages and use cases.
Simply mixing diverse multilingual data does not guarantee improved
performance. More research is required to determine robust techniques for
cross-lingual transfer that can truly augment in-language training.

##  4 Conclusion

Our experiments show the promise of continual pre-training and fine-tuning for
developing domain-specific language models, rather than training from scratch.
However, catastrophic forgetting of general knowledge remains a challenge.
Mitigating this trade-off through techniques like experience replay could
enable efficient adaptation of models to new domains while retaining diverse
capabilities. Further research into continual learning is needed to fully
realize its advantages.

##  5 Ethical statement

Continual training introduces the risk of amplifying biases in newly
incorporated datasets, necessitating periodic bias audits and fairness-aware
algorithms. Fine-tuning specialized medical tasks requires rigorous validation
by medical experts to mitigate the risk of disseminating inaccurate or
misleading medical information. The model is designed to complement, not
replace, human expertise in healthcare settings, and given its focus on the
Chinese medical language, it is fine-tuned to be sensitive to linguistic and
cultural nuances. Transparency and accountability are emphasized to ensure
ethical deployment, and ongoing ethical oversight is maintained to align the
project with principles of beneficence and nonmaleficence.

## References

  * [1] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al.  Language models are few-shot learners.  Advances in neural information processing systems, 33:1877–1901, 2020.
  * [2] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh, Kensen Shi, Sasha Tsvyashchenko, Joshua Maynez, Abhishek Rao, Parker Barnes, Yi Tay, Noam Shazeer, Vinodkumar Prabhakaran, Emily Reif, Nan Du, Ben Hutchinson, Reiner Pope, James Bradbury, Jacob Austin, Michael Isard, Guy Gur-Ari, Pengcheng Yin, Toju Duke, Anselm Levskaya, Sanjay Ghemawat, Sunipa Dev, Henryk Michalewski, Xavier Garcia, Vedant Misra, Kevin Robinson, Liam Fedus, Denny Zhou, Daphne Ippolito, David Luan, Hyeontaek Lim, Barret Zoph, Alexander Spiridonov, Ryan Sepassi, David Dohan, Shivani Agrawal, Mark Omernick, Andrew M. Dai, Thanumalayan Sankaranarayana Pillai, Marie Pellat, Aitor Lewkowycz, Erica Moreira, Rewon Child, Oleksandr Polozov, Katherine Lee, Zongwei Zhou, Xuezhi Wang, Brennan Saeta, Mark Diaz, Orhan Firat, Michele Catasta, Jason Wei, Kathy Meier-Hellstern, Douglas Eck, Jeff Dean, Slav Petrov, and Noah Fiedel.  Palm: Scaling language modeling with pathways, 2022.
  * [3] Long Ouyang, Jeff Wu, Xu Jiang, Diogo Almeida, Carroll L. Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, John Schulman, Jacob Hilton, Fraser Kelton, Luke Miller, Maddie Simens, Amanda Askell, Peter Welinder, Paul Christiano, Jan Leike, and Ryan Lowe.  Training language models to follow instructions with human feedback, 2022\.
  * [4] Wayne Xin Zhao, Kun Zhou, Junyi Li, Tianyi Tang, Xiaolei Wang, Yupeng Hou, Yingqian Min, Beichen Zhang, Junjie Zhang, Zican Dong, et al.  A survey of large language models.  arXiv preprint arXiv:2303.18223, 2023.
  * [5] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei.  Scaling laws for neural language models, 2020.
  * [6] Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al.  Training compute-optimal large language models.  arXiv preprint arXiv:2203.15556, 2022.
  * [7] Yu Gu, Robert Tinn, Hao Cheng, Michael Lucas, Naoto Usuyama, Xiaodong Liu, Tristan Naumann, Jianfeng Gao, and Hoifung Poon.  Domain-specific language model pretraining for biomedical natural language processing.  ACM Transactions on Computing for Healthcare (HEALTH), 3(1):1–23, 2021.
  * [8] Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt.  Measuring massive multitask language understanding, 2021.
  * [9] Arun James Thirunavukarasu, Darren Shu Jeng Ting, Kabilan Elangovan, Laura Gutierrez, Ting Fang Tan, and Daniel Shu Wei Ting.  Large language models in medicine.  Nature medicine, pages 1–11, 2023.
  * [10] Sangzin Ahn.  The impending impacts of large language models on medical education.  Korean Journal of Medical Education, 35(1):103, 2023.
  * [11] Shengbin Yue, Wei Chen, Siyuan Wang, Bingxuan Li, Chenchen Shen, Shujun Liu, Yuxuan Zhou, Yao Xiao, Song Yun, Xuanjing Huang, and Zhongyu Wei.  Disc-lawllm: Fine-tuning large language models for intelligent legal services, 2023.
  * [12] Ha-Thanh Nguyen.  A brief report on lawgpt 1.0: A virtual legal assistant based on gpt-3, 2023.
  * [13] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al.  Llama: Open and efficient foundation language models.  arXiv preprint arXiv:2302.13971, 2023.
  * [14] Suchin Gururangan, Ana Marasović, Swabha Swayamdipta, Kyle Lo, Iz Beltagy, Doug Downey, and Noah A. Smith.  Don’t stop pretraining: Adapt language models to domains and tasks.  In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 8342–8360, Online, July 2020. Association for Computational Linguistics.
  * [15] Zixuan Ke, Haowei Lin, Yijia Shao, Hu Xu, Lei Shu, and Bing Liu.  Continual training of language models for few-shot learning, 2022.
  * [16] Kshitij Gupta, Benjamin Thérien, Adam Ibrahim, Mats L. Richter, Quentin Anthony, Eugene Belilovsky, Irina Rish, and Timothée Lesort.  Continual pre-training of large language models: How to (re)warm your model?, 2023.
  * [17] Yining Hua, Liqin Wang, Vi Nguyen, Meghan Rieu-Werden, Alex McDowell, David W. Bates, Dinah Foer, and Li Zhou.  A deep learning approach for transgender and gender diverse patient identification in electronic health records.  Journal of Biomedical Informatics, page 104507, 2023.
  * [18] Yining Hua, Shixu Lin, Minghui Li, Yujie Zhang, Peilin Zhou, Ying-Chih Lo, Li Zhou, and Jie Yang.  Streamlining social media information retrieval for public health research with deep learning, 2023.
  * [19] Steven Bird.  Decolonising speech and language technology.  In Proceedings of the 28th International Conference on Computational Linguistics, pages 3504–3519, 2020.
  * [20] Qingcheng Zeng, Lucas Garay, Peilin Zhou, Dading Chong, Yining Hua, Jiageng Wu, Yikang Pan, Han Zhou, Rob Voigt, and Jie Yang.  Greenplm: Cross-lingual transfer of monolingual pre-trained language models at almost no cost.  In Edith Elkind, editor, Proceedings of the Thirty-Second International Joint Conference on Artificial Intelligence, IJCAI-23, pages 6290–6298. International Joint Conferences on Artificial Intelligence Organization, 8 2023.  AI for Good.
  * [21] Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto.  Stanford alpaca: An instruction-following llama model.  https://github.com/tatsu-lab/stanford_alpaca, 2023.
  * [22] Yudong Li, Yuqing Zhang, Zhe Zhao, Linlin Shen, Weijie Liu, Weiquan Mao, and Hui Zhang.  CSL: A large-scale Chinese scientific literature dataset.  In Proceedings of the 29th International Conference on Computational Linguistics, pages 3917–3923, Gyeongju, Republic of Korea, October 2022. International Committee on Computational Linguistics.
  * [23] Zhe Zhao, Yudong Li, Cheng Hou, Jing Zhao, Rong Tian, Weijie Liu, Yiren Chen, Ningyuan Sun, Haoyan Liu, Weiquan Mao, Han Guo, Weigang Guo, Taiqiang Wu, Tao Zhu, Wenhang Shi, Chen Chen, Shan Huang, Sihong Chen, Liqun Liu, Feifei Li, Xiaoshuai Chen, Xingwu Sun, Zhanhui Kang, Xiaoyong Du, Linlin Shen, and Kimmo Yan.  Tencentpretrain: A scalable and flexible toolkit for pre-training models of different modalities, 2023.
  * [24] Zhiqiang Shen, Tianhua Tao, Liqun Ma, Willie Neiswanger, Joel Hestness, Natalia Vassilieva, Daria Soboleva, and Eric Xing.  Slimpajama-dc: Understanding data combinations for llm training, 2023\.
  * [25] Guilherme Penedo, Quentin Malartic, Daniel Hesslow, Ruxandra Cojocaru, Alessandro Cappelli, Hamza Alobeidli, Baptiste Pannier, Ebtesam Almazrouei, and Julien Launay.  The refinedweb dataset for falcon llm: Outperforming curated corpora with web data, and web data only, 2023.
  * [26] Jianquan Li, Xidong Wang, Xiangbo Wu, Zhiyi Zhang, Xiaolong Xu, Jie Fu, Prayag Tiwari, Xiang Wan, and Benyou Wang.  Huatuo-26m, a large-scale chinese medical qa dataset, 2023.
  * [27] Renqian Luo, Liai Sun, Yingce Xia, Tao Qin, Sheng Zhang, Hoifung Poon, and Tie-Yan Liu.  Biogpt: generative pre-trained transformer for biomedical text generation and mining.  Briefings in Bioinformatics, 23(6), 2022.
  * [28] Rohan Anil, Andrew M. Dai, Orhan Firat, Melvin Johnson, Dmitry Lepikhin, Alexandre Passos, Siamak Shakeri, Emanuel Taropa, Paige Bailey, Zhifeng Chen, Eric Chu, Jonathan H. Clark, Laurent El Shafey, Yanping Huang, Kathy Meier-Hellstern, Gaurav Mishra, Erica Moreira, Mark Omernick, Kevin Robinson, Sebastian Ruder, Yi Tay, Kefan Xiao, Yuanzhong Xu, Yujing Zhang, Gustavo Hernandez Abrego, Junwhan Ahn, Jacob Austin, Paul Barham, Jan Botha, James Bradbury, Siddhartha Brahma, Kevin Brooks, Michele Catasta, Yong Cheng, Colin Cherry, Christopher A. Choquette-Choo, Aakanksha Chowdhery, Clément Crepy, Shachi Dave, Mostafa Dehghani, Sunipa Dev, Jacob Devlin, Mark Díaz, Nan Du, Ethan Dyer, Vlad Feinberg, Fangxiaoyu Feng, Vlad Fienber, Markus Freitag, Xavier Garcia, Sebastian Gehrmann, Lucas Gonzalez, Guy Gur-Ari, Steven Hand, Hadi Hashemi, Le Hou, Joshua Howland, Andrea Hu, Jeffrey Hui, Jeremy Hurwitz, Michael Isard, Abe Ittycheriah, Matthew Jagielski, Wenhao Jia, Kathleen Kenealy, Maxim Krikun, Sneha Kudugunta, Chang Lan, Katherine Lee, Benjamin Lee, Eric Li, Music Li, Wei Li, YaGuang Li, Jian Li, Hyeontaek Lim, Hanzhao Lin, Zhongtao Liu, Frederick Liu, Marcello Maggioni, Aroma Mahendru, Joshua Maynez, Vedant Misra, Maysam Moussalem, Zachary Nado, John Nham, Eric Ni, Andrew Nystrom, Alicia Parrish, Marie Pellat, Martin Polacek, Alex Polozov, Reiner Pope, Siyuan Qiao, Emily Reif, Bryan Richter, Parker Riley, Alex Castro Ros, Aurko Roy, Brennan Saeta, Rajkumar Samuel, Renee Shelby, Ambrose Slone, Daniel Smilkov, David R. So, Daniel Sohn, Simon Tokumine, Dasha Valter, Vijay Vasudevan, Kiran Vodrahalli, Xuezhi Wang, Pidong Wang, Zirui Wang, Tao Wang, John Wieting, Yuhuai Wu, Kelvin Xu, Yunhan Xu, Linting Xue, Pengcheng Yin, Jiahui Yu, Qiao Zhang, Steven Zheng, Ce Zheng, Weikang Zhou, Denny Zhou, Slav Petrov, and Yonghui Wu.  Palm 2 technical report, 2023.
  * [29] Karan Singhal, Tao Tu, Juraj Gottweis, Rory Sayres, Ellery Wulczyn, Le Hou, Kevin Clark, Stephen Pfohl, Heather Cole-Lewis, Darlene Neal, Mike Schaekermann, Amy Wang, Mohamed Amin, Sami Lachgar, Philip Mansfield, Sushant Prakash, Bradley Green, Ewa Dominowska, Blaise Aguera y Arcas, Nenad Tomasev, Yun Liu, Renee Wong, Christopher Semturs, S. Sara Mahdavi, Joelle Barral, Dale Webster, Greg S. Corrado, Yossi Matias, Shekoofeh Azizi, Alan Karthikesalingam, and Vivek Natarajan.  Towards expert-level medical question answering with large language models, 2023.
  * [30] Junling Liu, Peilin Zhou, Yining Hua, Dading Chong, Zhongyu Tian, Andrew Liu, Helin Wang, Chenyu You, Zhenhua Guo, Lei Zhu, and Michael Lingzhi Li.  Benchmarking large language models on cmexam – a comprehensive chinese medical exam dataset, 2023.
  * [31] Yunxiang Li, Zihan Li, Kai Zhang, Ruilong Dan, Steve Jiang, and You Zhang.  Chatdoctor: A medical chat model fine-tuned on a large language model meta-ai (llama) using medical domain knowledge, 2023.
  * [32] Wei-Lin Chiang, Zhuohan Li, Zi Lin, Ying Sheng, Zhanghao Wu, Hao Zhang, Lianmin Zheng, Siyuan Zhuang, Yonghao Zhuang, Joseph E. Gonzalez, Ion Stoica, and Eric P. Xing.  Vicuna: An open-source chatbot impressing gpt-4 with 90%* chatgpt quality, March 2023.
  * [33] James Kirkpatrick, Razvan Pascanu, Neil Rabinowitz, Joel Veness, Guillaume Desjardins, Andrei A. Rusu, Kieran Milan, John Quan, Tiago Ramalho, Agnieszka Grabska-Barwinska, Demis Hassabis, Claudia Clopath, Dharshan Kumaran, and Raia Hadsell.  Overcoming catastrophic forgetting in neural networks.  Proceedings of the National Academy of Sciences, 114(13):3521–3526, mar 2017.
  * [34] Haonan Li, Yixuan Zhang, Fajri Koto, Yifei Yang, Hai Zhao, Yeyun Gong, Nan Duan, and Timothy Baldwin.  Cmmlu: Measuring massive multitask language understanding in chinese, 2023\.
  * [35] Yuzhen Huang, Yuzhuo Bai, Zhihao Zhu, Junlei Zhang, Jinghan Zhang, Tangjun Su, Junteng Liu, Chuancheng Lv, Yikai Zhang, Jiayi Lei, Yao Fu, Maosong Sun, and Junxian He.  C-eval: A multi-level multi-discipline chinese evaluation suite for foundation models, 2023.
  * [36] Yanli Zhao, Andrew Gu, Rohan Varma, Liang Luo, Chien-Chin Huang, Min Xu, Less Wright, Hamid Shojanazeri, Myle Ott, Sam Shleifer, Alban Desmaison, Can Balioglu, Pritam Damania, Bernard Nguyen, Geeta Chauhan, Yuchen Hao, Ajit Mathews, and Shen Li.  Pytorch fsdp: Experiences on scaling fully sharded data parallel, 2023\.
  * [37] Paulius Micikevicius, Sharan Narang, Jonah Alben, Gregory Diamos, Erich Elsen, David Garcia, Boris Ginsburg, Michael Houston, Oleksii Kuchaiev, Ganesh Venkatesh, and Hao Wu.  Mixed precision training, 2018.
  * [38] Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin.  Training deep nets with sublinear memory cost, 2016.
  * [39] Noam Shazeer and Mitchell Stern.  Adafactor: Adaptive learning rates with sublinear memory cost, 2018.
  * [40] Xiang Lisa Li and Percy Liang.  Prefix-tuning: Optimizing continuous prompts for generation.  arXiv preprint arXiv:2101.00190, 2021.
  * [41] Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen.  Lora: Low-rank adaptation of large language models.  arXiv preprint arXiv:2106.09685, 2021.
  * [42] Di Jin, Eileen Pan, Nassim Oufattole, Wei-Hung Weng, Hanyi Fang, and Peter Szolovits.  What disease does this patient have? a large-scale open domain question answering dataset from medical exams.  arXiv preprint arXiv:2009.13081, 2020.
  * [43] Ankit Pal, Logesh Kumar Umapathi, and Malaikannan Sankarasubbu.  Medmcqa: A large-scale multi-subject multi-choice dataset for medical domain question answering.  In Gerardo Flores, George H Chen, Tom Pollard, Joyce C Ho, and Tristan Naumann, editors, Proceedings of the Conference on Health, Inference, and Learning, volume 174 of Proceedings of Machine Learning Research, pages 248–260. PMLR, 07–08 Apr 2022.


xx.xxxx/ACCESS.xxx.DOI

Research in this paper was supported by the Hellenic Foundation for Research
and Innovation (H.F.R.I.) under the “1st Call for H.F.R.I. Research Projects
to support Faculty Members & Researchers and the Procurement of High-Cost
Research Equipment Grant” (Project Number: 3449).

Corresponding author: Petros Giannakopoulos (e-mail<EMAIL_ADDRESS>
# Improving Post-processing of Audio Event Detectors using Reinforcement
Learning

PETROS GIANNAKOPOULOS1  AGGELOS PIKRAKIS2  AND YANNIS COTRONIS3  National and
Kapodistrian University of Athens (e-mail<EMAIL_ADDRESS> University of
Piraeus (e-mail<EMAIL_ADDRESS> National and Kapodistrian University of
Athens (e-mail<EMAIL_ADDRESS>
###### Abstract

We apply post-processing to the class probability distribution outputs of
audio event classification models and employ reinforcement learning to jointly
discover the optimal parameters for various stages of a post-processing stack,
such as the classification thresholds and the kernel sizes of median filtering
algorithms used to smooth out model predictions. To achieve this we define a
reinforcement learning environment where: 1) a state is the class probability
distribution provided by the model for a given audio sample, 2) an action is
the choice of a candidate optimal value for each parameter of the post-
processing stack, 3) the reward is based on the classification accuracy metric
we aim to optimize, which is the audio event-based macro F1-score in our case.
We apply our post-processing to the class probability distribution outputs of
two audio event classification models submitted to the DCASE Task4 2020
challenge. We find that by using reinforcement learning to discover the
optimal per-class parameters for the post-processing stack that is applied to
the outputs of audio event classification models, we can improve the audio
event-based macro F1-score (the main metric used in the DCASE challenge to
compare audio event classification accuracy) by 4-5% compared to using the
same post-processing stack with manually tuned parameters.

###### Index Terms:

audio classification, post-processing, reinforcement learning, DCASE challenge

=-15pt

Figure 1: Illustration of the proposed MDP approach to optimize the parameters
of a post-processing stack for an audio event classification model. The RL
agent receives the class probability outputs of the model and chooses the
parameters per class for each part of the post-processing stack (in this case
thresholding and median filtering). The post-processing stage with the given
parameters is subsequently applied to the model outputs and the F1-score is
updated. The agent then receives a reward based on the achieved score and
proceeds to the next audio segment.

##  I Introduction

As it is common with most audio analysis methods that operate by taking
classification decisions at the audio frame level, the goal of a final post-
processing stage is to increase the classification performance of an
inevitably imperfect model by smoothing out the sequence of predictions,
removing classification noise like false-positive or false-negative decisions
and improving the boundaries of detected audio events. A typical post-
processing stage consists of a soft output thresholding step followed by a
smoothing operation, with median filtering being a popular choice. The
parameters of such operations are usually hand-crafted following some grid
search technique. It can therefore be understood that post-processing schemes
have received significantly less attention over recent years, compared to the
core audio classification algorithms. To remove the need for such manual
parameter tuning, we propose a reinforcement learning [1] approach that
automatically computes parameter values of the post-processing pipeline
operations by treating them as learnable parameters that can be applied to the
class probability distribution outputs of audio event detection and
classification models. A widely accepted problem formulation for common audio
event detection and classification problems can be found in the DCASE Task4
challenge task descriptions, where, for example, according to the sound event
detection task specification (Task4 in recent years) [2], a model must be able
to identify the endpoints of each audio event in a given audio recording along
with the class label of each audio event. Examples of such audio events are
human speech, dog barks, vacuum cleaner sounds and so on. In general, audio
events may overlap and multiple classes are expected to be accommodated by an
algorithm ($10$ classes is a typical scenario). Following the DCASE Task4
challenge practice, audio event detection algorithms are primarily evaluated
based on the so-called event-based macro F1-score [3], according to which a
certain tolerance is permitted around the detected event endpoints, like a
$200$ ms tolerance collar on event onsets and a $200$ ms (or 20% of the events
length tolerance) collar on event offsets. In order to validate our approach,
we apply the proposed RL-based post-processing method to the class probability
outputs of two audio event detection models that were submitted to Task 4 of
the DCASE-2020 challenge: a) A baseline detector which was developed by the
organizers and achieved a $40.1\%$ event-based macro-F1 score [2] on the
evaluation set of the challenge and b) The best performing detector of the
2020 challenge, ConformerSED [4] which achieved a $47.7\%$ F1-score on the
same evaluation set. Both systems apply class probability thresholding and a
subsequent median filtering step at a post-processing stage and are,
therefore, good candidates for our method. In both challenge submissions, the
respective parameters were manually tuned by the method proposers using a
coarse grid search in the parameter space and were class-independent, i.e.,
the same values were applied to all classes. If, instead, our RL-based method
is used to automatically fine-tune a classification threshold and median
window length per audio event class, the F1-score of the baseline method
increases to $44.5\%$, and the F1-score of the best performing method
increases to $52.4\%$. We therefore observe a $4\%$ to $5\%$ performance
improvement on average, over very different methods, if the parameters of a
standard post-processing pipeline are tuned using RL, assuming that the RL
agent has access to the soft outputs of a method for a set of training
samples.

Our contribution can be summarised as follows:

  * •

We introduce a process based on reinforcement learning for optimizing the
parameters of the typical post-processing stack (thresholding and median
filtering) applied to the prediction outputs of an audio event detection
system.

  * •

We demonstrate that this reinforcement learning optimization approach of the
post-processing parameters provides superior audio event detection performance
compared to the typical grid search optimization approaches.

The code for this work is available on GitHub
111https://github.com/petrosgk/SED_RL_post_processing.

##  II Related Work

The work in [5] conducted an evaluation of post-processing algorithms that
were applied on the output of the baseline DCASE-2018 Task4 audio event
detector model and another submitted model. They concluded that post-
processing has a large impact on the final classification performance as
measured by the event-based F1-score and that post-processing algorithms
performed better when they were class-dependent. The authors in [6] proposed a
frame-wise dynamic threshold selection method for polyphonic sound event
detection and showed that it was superior to using a fixed threshold for all
audio frames. To estimate the threshold per frame, they used contour-based and
regressor-based methods. Our method adds to the post-processing approaches
outlined in these previous works by introducing RL as a way of optimizing the
parameters of post-processing, as far as both thresholding and median
filtering are concerned. Overall, however, the optimization of the post-
processing stage of audio event detection systems has not received enough
attention in literature so far, in our opinion, while it has been shown that
it can have a significant impact on the final performance of audio event
detection systems [5]. Regarding RL methods for audio analysis tasks, recent
systems primarily address the application of autonomous agents that are
capable of leveraging auditory information to perform diverse tasks given an
environment, e.g., [7, 8, 9, 10, 11], but, to the best of our knowledge, our
work is the first to formulate parameter tuning at the post-processing stage
as a RL task and employ a RL algorithm to automatically compute parameter
values for improving the final classification performance of audio event
detector systems.

##  III Method description

###  III-A Post-processing fundamentals for Audio Event Detection Methods

Given an audio recording, audio event detectors usually perform strong
labeling or weak labeling tasks. The former refers to problems where both
endpoints and class labels of audio events need to be reported, whereas the
latter refers to tasks where only the presence or absence of an audio event
type needs to be returned, thus skipping the need for endpoint detection. Note
that it is common to assume that in a multi-class, strong labelling context,
two or more events may overlap. As it is usually the case, a soft decision
output is computed per class over the duration of the recording and a
threshold is subsequently applied for each event class. Soft values above the
threshold indicate class presence in the respective frames. For the strong
labeling task, when successive frames yield values above a threshold, they are
grouped to form an event. In the case of weak labeling problems, thresholding
is used to identify which event classes are present in the audio segment by
collecting a percentage of frames that pass the threshold test per class.

As an example, if a detector model outputs posterior class probabilities at
the frame level or at the audio segment level, a threshold of $0.5$ means that
an audio event is considered to be present if the respective class probability
is higher than $0.5$. A straightforward threshold selection approach has been
to set a fixed threshold value for all classes. This type of naive threshold
selection may work on certain occasions but it will fail to provide good
results for many distributions of model outputs over the audio event classes
(Figure 2) [6]. We can also set different thresholds per event class, which
can potentially take into account variations in the model’s behaviour for
different event types. Apart from setting thresholds on the dataset level, we
can increase threshold selection granularity and set thresholds at the audio
clip level, or, in the case of strong labels, even at frame level. Our work
computes a threshold per class at the audio clip level and it is suitable for
both strong labeling and weak labeling tasks.

Figure 2: Strong (frame-level) labeling of audio events at the output of an
AED system, taken from [6]. Global fixed thresholds can lead to missed events
and wrong classifications.

Figure 3: Real (red) vs. predicted (blue) boundaries of an audio event at the
frame level, before (up) and after (down) applying median filtering. Median
filtering can be used to remove impulsive false positive or false negative
detections of frames (hills and valleys, respectively), as being part or not
of an audio event, by smoothing out the detected audio event boundaries.

Median filtering is a second post-processing step that is used especially in
strong labeling tasks to improve event detection accuracy. It can remove
certain artifacts of the event detector, like false positive or false negative
decisions that only last a few successive frames and can therefore smooth out
the detected audio event boundaries (Figure 3) and reduce over-segmentation.
The median window size can play an important role and different sizes may be
appropriate for various types of events [5], in order to balance undesirable
(excessive) smoothing of event boundaries and decision noise suppression.

###  III-B Post-processing Optimization using Reinforcement Learning

Based on the above description, a standard post-processing pipeline of an
Audio Event Detector (AED) system performing strong labeling consists of
thresholding and, optionally, median filtering. The problem is to estimate
threshold values that strike a balance between precision and recall for each
class, as well as median window sizes that remove detection artifacts without
causing excessive smoothing while taking into account the characteristics of
the system response for each audio event class.

Our method assumes that there exists a set of reasonable values for each
parameter of the post-processing pipeline, e.g. a $(0.0,1.0)$ continuous range
for the threshold selection step and discrete $[3,5,7,\ldots,21]$ values for
the median filtering window size. Contrary to the range of threshold values,
the range of median filtering window sizes is unbounded. For this reason, we
had to select a range of median filtering window sizes that are sensible and
most commonly encountered in AED systems literature. Our goal then becomes to
find “optimal”, class-dependent parameters for each audio segment in the
dataset. We can therefore consider the problem from a MDP perspective, i.e.,
as solving a MDP (Figure 1), where a state $S$ is the AED output for a given
audio segment, i.e., a matrix of class probabilities, one row per frame, like
in (Figure 2). Given $S$, we select a threshold value and a median window
length from the respective sets of possible values. If we are dealing with a
class-dependent optimization scenario, two different values will be selected
per class. Therefore, the action vector, $A$, of the RL agent, will consist of
$2\times N$ elements, where $N$ is the number of classes. If we are interested
in a simpler, class independent setting, there will only be two values per
audio clip and the action vector of the RL agent will therefore be two-
dimensional. Then the processing stack with the selected parameters is applied
to the AED outputs and the event-based macro F1-score for that audio segment
is evaluated. The higher the F1-score, the higher the reward, $R$, of taking
action $A$ in state $S$. The goal is to maximize the average reward over the
entire dataset, which directly correlates to maximizing the average F1-score.
Given that the F1-score is in the range $[0\%,100\%]$, and the reward should
be normalized in the range $[0.0,1.0]$, for training stability purposes, it
follows that the relationship between the reward $R$ and the F1-score would
simply be $R=F1/100$.

Figure 4: Neural network architecture used as policy function approximator. It
receives as input the AED frame-level event class probabilities (the state
$S$) and the outputs are vectors of length $nClasses$ for the class-dependent
thresholds and window sizes, or scalars if class-independent.

Parameter | Value
---|---
sample rate | 16000
max audio segment length (seconds) | 10
number of classes | 10
threshold range | (0.0, 1.0)
set of median filtering window sizes | [3, 5, 7, 9, 11, 13, 15, 17, 19, 21]

TABLE I: Values of parameters specific to the DCASE Task4 dataset. Used for
defining the reinforcement learning environment the PG algorithm will be
trained on.

Parameter | Value
---|---
batch size | 4
memory size | 10000 iterations
optimizer | Adam
learning rate | 0.001
update frequency | 4
discount factor ($\gamma$) | 0.99
entropy regularization | 0.001
GRU state size | 32

TABLE II: Values of hyperparameters used during training of the Policy Gradient algorithm. | Baseline | ConformerSED
---|---|---
Post-processing methods | Thresholding | Thresholding + Median Filtering | Thresholding | Thresholding + Median Filtering
| Event-based F1 (%) | Event-based F1 (%) | Event-based F1 (%) | Event-based F1 (%)
Original | 38.5 | 40.1 | 44.7 | 47.7
Class-dependent PG optimization | 42.4 | 44.5 | 49.2 | 52.4
Class-independent PG optimization | 41.3 | 43.4 | 48.0 | 51.1
TABLE III: Event-based macro F1-score for both Baseline and ConformerSED
systems, when the original parameters for thresholding and median filtering
window are used, and when optimized via Policy Gradient reinforcement learning
algorithm. We try both class-dependent and class-independent optimization. We
also compare the performance when only thresholding is used and when both
thresholding and median filtering post-processing are used. The results are
obtained on the validation set provided by the DCASE Task4 2020 challenge.

###  III-C Policy Gradient Optimization

To solve the MDP we employ Policy Gradient (PG) reinforcement learning method
[12, 13], where the policy function approximator is a simple recurrent neural
network made up from a stack of 2 GRU layers [14] (Figure 4). We chose a
recurrent architecture in order to be able to better capture time-dependent
information found in the sequences of predictions of an AED system, given that
the AED system itself usually performs classification over audio signal
feature sequences that it receives as input. The objective of the PG algorithm
is to find a policy function $\pi_{\theta}$, where $\theta$ are the weights of
the neural network, which adequately approximates the optimal policy
$\pi^{*}$. When $\pi^{*}$ is followed it yields the maximum achievable mean
reward $R$ over the dataset, which translates to the maximum achievable event-
based macro F1-score.

The PG algorithm iteratively makes $\pi_{\theta}$ converge to $\pi^{*}$ by
assigning a value to each state $S$, which represents the average reward that
it expects to get in the future by following the current policy $\pi_{\theta}$
in $S$, and it is directly related to the average expected F1-score. More
formally:

$V^{\pi}(S)=\mathbb{E}[\sum_{k=0}^{\infty}\gamma^{k}r_{t+k}]$ (1)

where $t\in[0,T]$, with $T=len(dataset)$, is the iteration index, $r_{t}$ is
the reward at iteration $t$, and $\gamma\in(0,1]$ is the future reward
discount factor, which governs how much ”weight” the algorithm should place on
long-term vs short-term rewards.

After the value of state $S$ is calculated, the next step is to estimate the
value (or advantage) of taking action $A$ in $S$ at iteration $t$ as the sum
of the expected values of future states, resulting from taking action $A$ in
state $S$, denoted here as $s_{t}$ and $a_{t}$:

$A_{t}(a_{t},s_{t})=\delta_{t}+\gamma\delta_{t+1}+\dots+\gamma^{T-t+1}\delta_{T-1},$
(2)

where

$\delta_{t}=r_{t}+\gamma V^{\pi}(s_{t+1})-V^{\pi}(s_{t}),$ (3)

At each iteration $t$ the policy approximator network outputs log-
probabilities (after the application of log-softmax) for each action,
multipled by the estimated advantage of that action:

$\log{\pi_{\theta}(a_{t}|s_{t})}$ (4)

The network outputs two actions $a_{t}$ given state $s_{t}$. In the case of
class-dependent optimization each $a_{t}$ is a vector of log-probabilities, of
length $nClasses$. Each vector entry is another vector of length $nValues$,
where $nValues$ is the number of possible values for the threshold and for the
median filtering window size. In the class-independent case $a_{t}$ is simply
a vector of length $nValues$ containing scalars. Since threshold values are
continuous, they are sampled from a uniform distribution in the range
$(0.0,1.0)$.

After multiplying the action log-probabilities (4) with the advantage estimate
for each action (2), the final objective optimized by PG for a batch of log-
probabilities is:

$L^{PG}(\theta)=\hat{\mathbb{E}}[\log{\pi_{\theta}(a_{t}|s_{t})}\hat{A_{t}}]$
(5)

where $\mathbb{E}$ is the mean over the batch. To promote exploration of the
state space and avoid convergence to a sub-optimal policy, an entropy
regularization term $H(\theta)=-\pi_{\theta}\log{\pi_{\theta}}$ is added to
the PG objective [15] so that the final objective becomes:

$L=L^{PG}(\theta)+\beta H(\theta)$ (6)

where $\beta$ is the weight of the entropy regularization term in the final
objective. In place, or in addition to entropy regularization, other methods
for facilitating exploration of the state space could be used, such as the
addition of gaussian or uniform noise to the continuous threshold predictions
of the network or randomly replacing the discrete median filtering window size
predictions of the network with other values sampled from the set of window
sizes.

##  IV Experiments

Similar to [5], we apply our method to the optimization of the post-processing
stack of two AED systems: a baseline approach and a state-of-the-art method,
to cover the two extremes of the performance spectrum. Table I shows the
values of the parameters that define the state and action spaces of the
reinforcement learning environment on which the Policy Gradient algorithm is
trained. Specifically, the state space is defined by the sample rate, maximum
length of audio segments and number of classes present in the dataset. The
sets of possible threshold values and median filter window sizes define the
action space. Table II shows the values of the Policy Gradient algorithm
hyperparameters that are used for training. Specifically, the choice of batch
size was constrained by the available system memory, or how many copies of the
dataset could fit in memory. Memory size refers to the history of training
iterations, or interactions of the RL agent with the MDP environment, that are
buffered and is also constrained by available system memory. The rest of the
parameters, such as the choice of optimizer, learning rate, update frequency
were taken from [15]. The training parameters remained identical for both AED
systems that the PG algorithm was trained on. The most important ones are the
discount factor ($\gamma$) and the entropy regularization weight. A discount
factor close to $1$ is needed to assign more importance to rewards gained in
the future, because the F1-score is only calculated once at the end of the
dataset after the predictions for all audio segments have been accumulated.
The entropy regularization weight value achieves a balance between exploration
and exploitation: a low value leads to faster convergence but less exploration
of the state space, which in turn can lead to convergence to a sub-optimal
F1-score.

The results of our experiments are grouped in Table III, for the two AED
systems that we evaluated. The first line of the table shows the event-based
macro F1-score when using the same parameters as the authors of the respective
work. The second line shows the F1-score that was achieved by optimizing the
threshold value and median filter size per class (class-dependent) using
policy gradient reinforcement learning. Finally, the third line shows the
class-independent PG optimization, where only global threshold and median
filter window values are optimized and used for all classes. We also test the
effect of median filtering on the AED performance by removing it from the
post-processing stack to only use thresholding.

For the Baseline AED system, the authors use a fixed event class probability
threshold of $0.5$ and a fixed window size of $7$ for the median filter.
Applying median filtering after thresholding improves F1-score by $1.6\%$,
from $38.5\%$ to $40.1\%$. Class-dependent optimization of thresholding and
median filtering parameters via RL leads to a F1-score improvement of $4.4\%$
over the original fixed values, from $40.1\%$ to $44.5\%$. Class-independent
optimization gives a slightly lower improvement of $3.3\%$ over the original
values.

For the ConformerSED system, the authors use a fixed threshold of $0.5$ but
perform a coarse grid search separately for each event class to tune the
median filter length. Class-dependent RL optimization computes on a class
basis better values for the median filter length, as the F1-score improves
from $47.7\%$ to $52.4\%$, an increase of $4.7\%$. Again, class-independent
optimization performs slightly worse than the class-dependent one, with an
improvement of $3.4\%$ over the original values.

Overall, a reinforcement agent that earns to compute “optimal” values for
class thresholds and median filter sizes using policy gradient optimization,
yields performance improvements for both AED methods under study. It can also
be observed that class-dependent optimization yields better results compared
to the class-independent (global) optimization case.

As PG optimizes for a reward-based objective, we can use it to indirectly
optimize a non-differentiable objective, such as the F1-score in our case, by
translating increases or decreases of that objective to positive or negative
credit assignment. Since the F1-score is in the $[0\%,100\%]$ range of values,
we simply do a straightforward linear mapping to a reward in the $[0.0,1.0]$
range.

##  V Limitations

The computational time that is required to learn the policy function with a
neural network approximator can be non-trivial, depending on the number of
audio segments. Given that each audio segment is treated as a state of the
MDP, the MDP will grow larger with the amount of audio segments. This in turn
will increase the time required for the PG algorithm to reach convergence. In
a standard experimental setup with low-cost hardware (Intel Core i9-10900K
processor, 32GB RAM) convergence was achieved after approximately $2$ hours of
processing time on the DCASE-Task4 training dataset, which consists of
approximately $2500$ audio segments. However, the computational time required
is, in our opinion, offset by the superior final classification accuracy
achieved by our method. As a final remark, note that our method needs to be
trained separately on each AED system output, i.e., we do not provide a global
RL agent for all possible AED methods.

##  VI Conclusion

In this work we introduced reinforcement learning, and in particular policy
gradient optimization, as a way of discovering optimal parameter values for
the two most commonly used stages of the post-processing stack of an audio
event detector: thresholding and median filtering. To that end, we reformulate
this parameter optimization problem as a Markov Decision Process, where: 1)
the state $S$ consists of the class probability outputs of the AED system for
a given audio segment of the dataset, 2) the actions $A$ are the values
selected as optimal by the RL algorithm for the thresholding and median
filtering parameters, either on a per-class basis or globally, 3) the reward
$R$ is the event-based macro F1-score achieved by the RL algorithm over the
entire dataset using the selected parameter values for the post-processing
stack.

We tested our method on two AED systems submitted to the DCASE Task4 (Sound
Event Detection in Domestic Environments) 2020 challenge: 1) A baseline AED
system provided by the challenge organizers, 2) The system (ConformerSED) that
achieved top ranking on the challenge, based on achieved event-based macro-F1
score. Both systems use a thresholding + median filtering post-processing
stack. We achieved an improvement of about $4\%$ to $5\%$ on average on the
event-based macro-F1 score compared to the performance of these systems with
their original parameterization of the post-processing stack.

A drawback of the proposed method is the requirement to train the PG algorithm
separately on the outputs of each AED system. The computational requirements
highly depend on the size of the dataset of audio segments. Another important
prerequisite for the success of the method is the sufficient exploration of
the state space to avoid converging on sub-optimal post-processing parameter
values, generally known as exploration vs. exploitation dilemma in RL
literature. However, when properly tuned, it can provide superior final
classification performance for an AED system compared to performing exhaustive
grid search for finding the optimal parameter values for the post-processing
stack, especially on larger datasets. An interesting avenue of future research
would be to simultaneously train the PG algorithm on the predictions of
several AED systems for a particular dataset and assess its ability to
generalize to other AED systems that it was not trained on. Another
possibility for future work would be to try other, more sophisticated than
Policy Gradient, reinforcement learning algorithms, with our method, such as
Proximal Policy Optimization (PPO) [16] and Soft Actor-Critic (SAC) [17].

## References

  * [1] R. S. Sutton and A. G. Barto, _Reinforcement learning: An introduction_. MIT press, 2018.
  * [2] N. Turpault, R. Serizel, A. Parag Shah, and J. Salamon, “Sound event detection in domestic environments with weakly labeled data and soundscape synthesis,” in _Workshop on Detection and Classification of Acoustic Scenes and Events_ , New York City, United States, October 2019.
  * [3] A. Mesaros, T. Heittola, and T. Virtanen, “Metrics for polyphonic sound event detection,” _Applied Sciences_ , vol. 6, no. 6, 2016.
  * [4] K. Miyazaki, T. Komatsu, T. Hayashi, S. Watanabe, T. Toda, and K. Takeda, “Conformer-based sound event detection with semi-supervised learning and data augmentation,” in _Detection and Classification of Acoustic Scenes and Events_ , 2020, pp. 100–104.
  * [5] L. Cances, P. Guyot, and T. Pellegrini, “Evaluation of post-processing algorithms for polyphonic sound event detection,” in _IEEE Workshop on Applications of Signal Processing to Audio and Acoustics (WASPAA)_. IEEE, 2019, pp. 318–322.
  * [6] X. Xia, R. Togneri, F. Sohel, and D. Huang, “Frame-wise dynamic threshold based polyphonic acoustic event detection,” in _INTERSPEECH_ , 2017, pp. 474–478.
  * [7] A. Woubie, A. Kanervisto, J. Karttunen, and V. Hautamaki, “Do autonomous agents benefit from hearing?” _arXiv preprint arXiv:1905.04192_ , 2019.
  * [8] S. Hegde, A. Kanervisto, and A. Petrenko, “Agents that listen: High-throughput reinforcement learning with multiple sensory systems,” in _IEEE Conference on Games (CoG)_. IEEE, 2021, pp. 1–5.
  * [9] S. Lathuilière, B. Massé, P. Mesejo, and R. Horaud, “Deep reinforcement learning for audio-visual gaze control,” in _IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS)_. IEEE, 2018, pp. 1555–1562.
  * [10] P. Giannakopoulos, A. Pikrakis, and Y. Cotronis, “A deep reinforcement learning approach to audio-based navigation in a multi-speaker environment,” in _IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_. IEEE, 2021, pp. 3475–3479.
  * [11] T. Van Nguyen, X. Dai, I. Khan, R. Thawonmas, and H. V. Pham, “A deep reinforcement learning blind ai in darefightingice,” _arXiv preprint arXiv:2205.07444_ , 2022.
  * [12] R. S. Sutton, D. McAllester, S. Singh, and Y. Mansour, “Policy gradient methods for reinforcement learning with function approximation,” _Advances in neural information processing systems_ , vol. 12, 1999.
  * [13] M. Schaarschmidt, A. Kuhnle, B. Ellis, K. Fricke, F. Gessert, and E. Yoneki, “Lift: Reinforcement learning in computer systems by learning from demonstrations,” _arXiv preprint arXiv:1808.07903_ , 2018.
  * [14] K. Cho, B. Van Merriënboer, D. Bahdanau, and Y. Bengio, “On the properties of neural machine translation: Encoder-decoder approaches,” in _Eighth Workshop on Syntax, Semantics and Structure in Statistical Translation (SSST-8)_ , 2014.
  * [15] V. Mnih, A. P. Badia, M. Mirza, A. Graves, T. Lillicrap, T. Harley, D. Silver, and K. Kavukcuoglu, “Asynchronous methods for deep reinforcement learning,” in _International Conference on Machine Learning_. PMLR, 2016, pp. 1928–1937.
  * [16] J. Schulman, F. Wolski, P. Dhariwal, A. Radford, and O. Klimov, “Proximal policy optimization algorithms,” _arXiv preprint arXiv:1707.06347_ , 2017\.
  * [17] T. Haarnoja, A. Zhou, P. Abbeel, and S. Levine, “Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor,” in _International conference on machine learning_. PMLR, 2018, pp. 1861–1870.

|  Petros Giannakopoulos is currently a PhD candidate at the Department of
Informatics and Telecommunications, National and Kapodistrian University of
Athens, Greece. His research is focused on deep reinforcement learning for
audio applications. Previously, he earned an MSc in Computer Science and a BSc
in Physics. He also has industry experience as a Machine Learning Engineer and
as a Data Scientist, working for startups and research institutions.
---|---
|  Aggelos Pikrakis received the Diploma Engineering degree from the
University of Patras, in 1993, and the Ph.D. degree in Computer Science from
the University of Athens, in 2002. He is currently an Assistant Professor at
the Department of Informatics, University of Piraeus, Greece, teaching courses
related to machine learning and audio processing. He is the co-inventor of
AI/ML patents and he has coauthored two international textbooks in the English
language and more than 50 refereed papers in international peer-reviewed
scientific journals/conferences. His research interests include audio
analysis, with an emphasis on machine learning algorithms, including deep
neural networks, hidden Markov models, and Bayesian architectures and sequence
alignment methods. His research work has received award recognition and he was
a recipient of the 2019 EURASIP Meritorious Service Award.
---|---
|  Yannis Cotronis is a Professor (retired) at the Department of Informatics
and Telecommunications, National and Kapodistrian University of Athens. He
holds a B.Sc. in Mathematics, M.Sc. and Ph.D. in Computing Science. His
research interests include software engineering for parallel programming and
e-science applications. He has taught a number of subjects in the core of
Computer Science, his expert field being Parallel Systems and Programming. He
has chaired EuroPVM-MPI 2001, Euromicro PDP 2000, 2011, 2016 and EuroMPI 2011.
---|---