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

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Fang, Yin, Jin, Qiao, Xiong, Guangzhi, Jin, Bowen, Zhong, Xianrui, Ouyang, Siru, Zhang, Aidong, Han, Jiawei, Lu, Zhiyong
Format: Preprint
Published: 2025
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866912411520335872
author Fang, Yin
Jin, Qiao
Xiong, Guangzhi
Jin, Bowen
Zhong, Xianrui
Ouyang, Siru
Zhang, Aidong
Han, Jiawei
Lu, Zhiyong
author_facet Fang, Yin
Jin, Qiao
Xiong, Guangzhi
Jin, Bowen
Zhong, Xianrui
Ouyang, Siru
Zhang, Aidong
Han, Jiawei
Lu, Zhiyong
contents Cell type annotation is a key task in analyzing the heterogeneity of single-cell RNA sequencing data. Although recent foundation models automate this process, they typically annotate cells independently, without considering batch-level cellular context or providing explanatory reasoning. In contrast, human experts often annotate distinct cell types for different cell clusters based on their domain knowledge. To mimic this workflow, we introduce the CellPuzzles task, where the objective is to assign unique cell types to a batch of cells. This benchmark spans diverse tissues, diseases, and donor conditions, and requires reasoning across the batch-level cellular context to ensure label uniqueness. We find that off-the-shelf large language models (LLMs) struggle on CellPuzzles, with the best baseline (OpenAI's o1) achieving only 19.0% batch-level accuracy. To fill this gap, we propose Cell-o1, a 7B LLM trained via supervised fine-tuning on distilled reasoning traces, followed by reinforcement learning with batch-level rewards. Cell-o1 achieves state-of-the-art performance, outperforming o1 by over 73% and generalizing well across contexts. Further analysis of training dynamics and reasoning behaviors provides insights into batch-level annotation performance and emergent expert-like reasoning. Code and data are available at https://github.com/ncbi-nlp/cell-o1.
format Preprint
id arxiv_https___arxiv_org_abs_2506_02911
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Cell-o1: Training LLMs to Solve Single-Cell Reasoning Puzzles with Reinforcement Learning
Fang, Yin
Jin, Qiao
Xiong, Guangzhi
Jin, Bowen
Zhong, Xianrui
Ouyang, Siru
Zhang, Aidong
Han, Jiawei
Lu, Zhiyong
Computation and Language
Artificial Intelligence
Computational Engineering, Finance, and Science
Human-Computer Interaction
Machine Learning
Cell type annotation is a key task in analyzing the heterogeneity of single-cell RNA sequencing data. Although recent foundation models automate this process, they typically annotate cells independently, without considering batch-level cellular context or providing explanatory reasoning. In contrast, human experts often annotate distinct cell types for different cell clusters based on their domain knowledge. To mimic this workflow, we introduce the CellPuzzles task, where the objective is to assign unique cell types to a batch of cells. This benchmark spans diverse tissues, diseases, and donor conditions, and requires reasoning across the batch-level cellular context to ensure label uniqueness. We find that off-the-shelf large language models (LLMs) struggle on CellPuzzles, with the best baseline (OpenAI's o1) achieving only 19.0% batch-level accuracy. To fill this gap, we propose Cell-o1, a 7B LLM trained via supervised fine-tuning on distilled reasoning traces, followed by reinforcement learning with batch-level rewards. Cell-o1 achieves state-of-the-art performance, outperforming o1 by over 73% and generalizing well across contexts. Further analysis of training dynamics and reasoning behaviors provides insights into batch-level annotation performance and emergent expert-like reasoning. Code and data are available at https://github.com/ncbi-nlp/cell-o1.
title Cell-o1: Training LLMs to Solve Single-Cell Reasoning Puzzles with Reinforcement Learning
topic Computation and Language
Artificial Intelligence
Computational Engineering, Finance, and Science
Human-Computer Interaction
Machine Learning
url https://arxiv.org/abs/2506.02911