Title: Nirvana: A Specialized Generalist Model With Task-Aware Memory Mechanism
Primary area: foundation or frontier models, including LLMs
Keywords: Specialized Generalist Models, Large Language Models

Abstract: Large Language Models (LLMs) have achieved remarkable success across a wide range of general language tasks but remain constrained in specialized domains. Inspired by the brain's cognition, specialized memory mechanism can be used to enhance the model's ability on specialized tasks. Specialized Generalist Models (SGMs) aim to preserve broad capabilities while achieving expert-level performance in target domains via test-time task identification and reconfiguration. However, traditional LLM structures including Transformer, Linear Attention, and hybrid models do not employ specialized memory mechanism guided by task information. In this paper, we present Nirvana, an SGM with specialized memory mechanism, linear time complexity, and test-time task information extraction. 
Besides, we propose the Task-Aware Memory Trigger ($\textit{Trigger}$) that flexibly adjusts memory mechanism based on the current task's requirements. In Trigger, each incoming sample is treated as a self-supervised fine-tuning task, enabling Nirvana to adapt its task-related parameters on the fly to domain shifts. We also design the Specialized Memory Updater ($\textit{Updater}$) that dynamically memorizes the context guided by Trigger. We conduct experiments on both general language tasks and specialized medical tasks. 
On a variety of natural language modeling benchmarks, Nirvana achieves competitive or superior results compared to the existing LLM structures. To prove the effectiveness of Trigger on specialized tasks, we test Nirvana's performance on a challenging medical task, i.e., Magnetic Resonance Imaging (MRI). We post-train frozen Nirvana backbone with lightweight codecs on paired electromagnetic signals and MRI images. Despite the frozen Nirvana backbone, Trigger guides the model to adapt to the MRI domain with the change of task-related parameters. Nirvana achieves higher-quality MRI reconstruction compared to conventional MRI models as well as the models with traditional LLMs' backbone, and can also generate accurate preliminary clinical reports accordingly.

TL;DR: We present Nirvana, a Specialized Generalist Model with task-aware memory mechanism, linear time complexity, and test-time task information extraction.