Cell-o1: Training LLMs to Solve Single-Cell Reasoning Puzzles with Reinforcement Learning

细胞类型注释是分析单细胞RNA测序数据异质性的关键任务。尽管最近的基础模型已经实现了这一过程的自动化，但它们通常独立地注释细胞，而不考虑批量级别的细胞上下文，也不提供解释性推理。相比之下，人类专家通常基于他们的领域知识，为不同的细胞簇注释出独特的细胞类型。为了模拟这一工作流程，我们引入了CellPuzzles任务，其目标是为一批细胞分配唯一的细胞类型。这一基准测试涵盖了多样化的组织、疾病和供体条件，并要求在批量级别的细胞上下文中进行推理，以确保标签的唯一性。我们发现，现成的大规模语言模型（LLMs）在CellPuzzles任务上表现不佳，其中表现最好的基线模型（OpenAI的o1）仅能达到19.0%的批量级别准确率。为弥补这一差距，我们提出了Cell-o1，这是一个7B参数的LLM，通过监督微调蒸馏推理轨迹训练，并结合批量级别的奖励进行强化学习。Cell-o1实现了最先进的性能，比o1高出超过73%，并且在不同情境下表现出良好的泛化能力。进一步对训练动态和推理行为的分析提供了关于批量级别注释性能以及类似专家的推理能力涌现的见解。代码和数据可在https://github.com/ncbi-nlp/cell-o1获取。

该论文试图解决单细胞RNA测序数据中细胞类型注释的问题，特别是如何在批处理级别考虑细胞上下文并确保标签的独特性。这是一个相对较新的问题，因为传统的模型通常独立注释每个细胞，而未充分考虑整体的批处理级信息。

论文提出了一种名为Cell-o1的新方法，通过监督微调和基于批处理级奖励的强化学习来训练一个7B参数的语言模型。与现有方法不同，Cell-o1不仅关注单个细胞的特征，还结合了批处理级别的上下文信息，从而实现更准确的注释。这种方法模仿了人类专家的工作流程，即基于领域知识对不同细胞簇进行区分注释。

实验设计包括跨组织、疾病和供体条件的多样化基准测试，显示了Cell-o1在多种背景下的良好泛化能力。论文使用了CellPuzzles任务作为主要评估标准，并提供了详细的训练动态和推理行为分析。代码和数据已开源至GitHub (https://github.com/ncbi-nlp/cell-o1)。此外，研究发现大型语言模型（LLM）在此任务上的表现较差，表明该任务具有挑战性，值得进一步探索。

相关研究包括：1) 基于基础模型的自动化细胞类型注释方法；2) 针对单细胞RNA测序数据的无监督聚类算法；3) 结合领域知识的细胞类型注释工具。例如，《Single-cell RNA-seq data analysis with deep learning methods》探讨了深度学习在单细胞数据分析中的应用，而《Domain knowledge integration in single-cell annotation》则研究了如何将领域知识融入注释过程。