Paper deep dive
Mamba-SSM with LLM Reasoning for Biomarker Discovery: Causal Feature Refinement via Chain-of-Thought Gene Evaluation
Pushpa Kumar Balan, Aijing Feng
Intelligence
Status: succeeded | Model: google/gemini-3.1-flash-lite-preview | Prompt: intel-v1 | Confidence: 96%
Last extracted: 4/18/2026, 1:33:35 AM
Summary
The paper introduces a neuro-symbolic framework for biomarker discovery that combines Mamba-SSM gradient saliency with DeepSeek-R1 LLM chain-of-thought reasoning to filter tissue-composition confounders from RNA-seq data. The study demonstrates that LLM-filtered gene sets achieve superior downstream classification performance (AUC 0.927) compared to raw saliency (AUC 0.832) or variance-based baselines (AUC 0.903), despite a moderate recall of known BRCA biomarkers, suggesting that targeted confounder removal is more critical for performance than exhaustive recall.
Entities (5)
Relation Signals (3)
DeepSeek-R1 → rejected → FOXA1
confidence 100% · FOXA1... was present in the top-50 input (rank 49) but not selected.
DeepSeek-R1 → filtersfeaturesfrom → Mamba-SSM
confidence 95% · DeepSeek-R1 evaluates every candidate with structured CoT to produce a final 17-gene set.
Mamba-SSM → generatessaliencyfor → TCGA-BRCA
confidence 95% · We train a Mamba SSM on TCGA-BRCA RNA-seq and extract the top-50 genes by gradient saliency
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:Gradient saliency from deep sequence models surfaces candidate biomarkers efficiently, but the resulting gene lists are contaminated by tissue-composition confounders that degrade downstream classifiers. We study whether LLM chain-of-thought (CoT) reasoning can faithfully filter these confounders, and whether reasoning quality drives downstream performance. We train a Mamba SSM on TCGA-BRCA RNA-seq and extract the top-50 genes by gradient saliency; DeepSeek-R1 evaluates every candidate with structured CoT to produce a final 17-gene set. The raw 50-gene saliency set (no LLM) performs worse than a 5,000-gene variance baseline (AUC 0.832 vs. 0.903), while the LLM-filtered set surpasses it (AUC 0.927), using 294x fewer features. A faithfulness audit (COSMIC CGC, OncoKB, PAM50) reveals only 6 of 17 selected genes (35.3%) are validated BRCA biomarkers, yet 10 of 16 known BRCA genes in the input were missed - including FOXA1. This gap between downstream performance and reasoning faithfulness suggests selective faithfulness: targeted confounder removal is sufficient for performance gains even without comprehensive recall.
Tags
Links
- Source: https://arxiv.org/abs/2604.14334v1
- Canonical: https://arxiv.org/abs/2604.14334v1
Trouble viewing inline? Open PDF directly →
Full Text
21,796 characters extracted from source content.
Expand or collapse full text
Mamba-SSM with LLM Reasoning for Biomarker Discovery: Causal Feature Refinement via Chain-of-Thought Gene Evaluation Pushpa Kumar Balan, Aijing Feng Department of Computer Science and Cybersecurity University of Central Missouri pushpakumarbalan@gmail.com, feng@ucmo.edu Abstract Gradient saliency from deep sequence models surfaces candidate biomarkers efficiently, but the resulting gene lists are contaminated by tissue-composition confounders that degrade downstream classifiers. We study whether LLM chain-of-thought (CoT) reasoning can faithfully filter these confounders, and whether reasoning quality drives downstream performance. We train a Mamba SSM on TCGA-BRCA RNA-seq and extract the top-50 genes by gradient saliency; DeepSeek-R1 evaluates every candidate with structured CoT to produce a final 17-gene set. The raw 50-gene saliency set (no LLM) performs worse than a 5,000-gene variance baseline (AUC 0.832 vs. 0.903), while the LLM-filtered set surpasses it (AUC 0.927), using 294× fewer features. A faithfulness audit (COSMIC CGC, OncoKB, PAM50) reveals only 6 of 17 selected genes (35.3%) are validated BRCA biomarkers, yet 10 of 16 known BRCA genes in the input were missed—including FOXA1. This gap between downstream performance and reasoning faithfulness suggests selective faithfulness: targeted confounder removal is sufficient for performance gains even without comprehensive recall. The scope of this claim and its limitations are examined in Sec. B. Code: https://github.com/pushpakumarbalan/feature-selection 1 Introduction High-dimensional RNA-seq data (>>20,000 genes per sample) presents a severe feature-selection problem: most genes are irrelevant to the phenotype of interest, and many that appear predictive are confounders (immune infiltration, tumour purity, batch effects) rather than disease drivers (Pudjihartono et al., 2022). Standard gradient-based saliency from neural models ranks genes by their gradient magnitude, but this signal reflects what the model learned to use, not biological causality. The top-50 saliency genes from a well-trained Mamba SSM on TCGA-BRCA include muscle-specific genes (MB, UTRN), general immune markers (HLA-DRB1, ITGAL), unannotated lncRNAs, and antisense RNAs with no documented breast cancer role, all of which carry saliency simply because they co-vary with tumour samples at the RNA-seq level. This creates a natural role for LLM reasoning: the model’s encoded biomedical knowledge can in principle distinguish disease drivers (e.g., the ER-signalling gene XBP1, the EMT regulator ZEB1) from these confounders without additional data. But does the LLM’s stated reasoning actually reflect accurate biological knowledge? And is downstream performance a reliable proxy for reasoning faithfulness? We make three contributions: 1. We demonstrate that raw saliency-based feature selection hurts performance relative to a variance baseline (AUC −-0.071), while LLM-filtered selection helps (AUC ++0.024), establishing that LLM reasoning is causally necessary in this pipeline. 2. We conduct a faithfulness audit comparing the LLM’s selected gene set against curated BRCA ground-truth databases, revealing a recall of 0.375 on known BRCA genes while still achieving superior downstream AUC—a finding we term selective faithfulness: targeted removal of known non-BRCA genes is sufficient for performance gains, even without comprehensive recall of all true positives. 3. We identify a concrete failure mode: FOXA1, the most important luminal breast cancer pioneer transcription factor and a canonical PAM50 gene, was present in the input but rejected by the LLM— illustrating that LLM biomedical reasoning can be confidently wrong on well-established facts. 2 Related Work LLMs for feature selection: LLM-Select (Jeong et al., 2025) showed that zero-shot LLM feature selection using only feature names can match LASSO on tabular data in some settings. LLM-Lasso (Zhang et al., 2025) integrates domain knowledge through LLM-guided regularisation, penalising literature-supported features less. FreeForm (Lee et al., 2025) demonstrated that LLM ensembling improves variant selection in low-data genomic regimes. Our work differs by studying faithfulness—we audit whether the LLM’s stated biological rationale is correct, not just whether the output improves a metric. Faithfulness of LLM reasoning: A growing body of work questions whether CoT explanations reflect the model’s actual reasoning process (Turpin et al., 2023; Lanham et al., 2023). Most such studies use synthetic tasks with verifiable ground truth. We provide a real-world biological instance: a domain where ground truth (validated cancer driver genes) exists but is large, overlapping, and context-dependent—making faithfulness harder to assess and more practically important. SSMs for genomics: Mamba (Gu & Dao, 2024) scales linearly in sequence length, making it tractable for 20,000-dimensional gene expression vectors without attention’s quadratic cost. We use the official mamba-ssm implementation with gradient saliency (Simonyan et al., 2014) to extract a biologically plausible candidate pool before the LLM reasoning step. 3 Methodology 3.1 System Pipeline Figure 1: Complete Neuro-Symbolic Pipeline: High-dimensional genomic data is processed by the Mamba SSM to generate saliency scores, which are subsequently interpreted and filtered by the DeepSeek-R1 agentic reasoner. 3.2 Data We use TCGA-BRCA RNA-seq data: 1,095 tumour and 113 matched normal samples, ∼ 20,000 protein-coding and non-coding genes, quantified as TPM. Labels are binary: Tumour (1) / Normal (0). 3.3 Phase 1 — Mamba SSM Training Input features are log2(TPM++1)-normalised and filtered to the top-5,000 highest-variance genes. We train an OfficialMambaClassifier: a linear embedding layer projects each gene’s scalar expression value into a dmodel=128d_model=128 dimensional space; a single Mamba block (dstate=16d_state=16, dconv=4d_conv=4, expand=2=2) models long-range dependencies across the gene sequence; adaptive average pooling collapses the sequence dimension; a final linear layer with sigmoid produces the tumour probability. Training uses AdamW (lr=10−4lr=10^-4, 15 epochs, batch size 8) with class-weighted BCELoss (wnormal=Ntumour/Nnormalw_normal=N_tumour/N_normal) to handle the 8.8:1 tumour/normal imbalance. 3.4 Phase 2 — Gradient Saliency Extraction For each tumour sample, we enable gradients on the input, perform a forward pass, compute the loss, and backpropagate. The gene importance score is sj=1|T|∑i∈T|∂ℒ∂xij|s_j= 1|T| _i∈ T | ∂ x_ij |, averaged over all tumour samples T. The top-50 genes by sjs_j form the candidate pool 50G_50. 3.5 Phase 3 — Structured CoT Reasoning 50G_50 is passed to DeepSeek-R1 (7B, local via Ollama, temperature =0.3=0.3) with a structured prompt that (a) provides saliency scores, (b) explicitly states high saliency does not imply BRCA specificity, (c) requires evaluation of every gene via five rejection criteria (R1–R5) and three keep criteria (K1–K3), and (d) forbids rank-order selection. A post-hoc audit detects rank-copy solutions and hallucinated gene names. The distinction between structured rule application and genuine reasoning behaviour is examined in Sec. B. 3.6 Phase 4 — Faithfulness Audit The 17 selected genes are cross-referenced against a 101-gene ground-truth set (COSMIC CGC Tier-1, OncoKB BRCA annotations, PAM50 (Parker et al., 2009), and established pathway genes). A known non-BRCA set (muscle, neuronal, unannotated genes) provides true-negative ground truth. For each of the 50 input genes we record the ground-truth class (validated / known non-BRCA / unknown) and the LLM’s decision (selected / not selected), yielding selection-level precision, recall, and missed-gene analysis. 4 Results 4.1 Downstream Classification Benchmark Each gene set is evaluated by training an identical Mamba classifier (same architecture, same random seed, same 80/20 stratified split) on the respective feature subset. Results are in Table 1. Table 1: Classification performance on the held-out test set (20% of 1,231 samples). All conditions use the same Mamba architecture and training procedure; only the input gene set changes. Method Genes Accuracy F1 AUC B1 Variance baseline 5,000 0.8785 0.8941 0.903 B2 Mamba saliency only (no LLM) 50 0.7247 0.7813 0.832 B3 Mamba + LLM structured CoT 17 0.8907 0.9033 0.927 The comparison between B1, B2, and B3 reveals a non-monotonic relationship between gene count and performance. B2 (50 genes, no reasoning) performs worse than B1 (5,000 genes) by AUC −0.071-0.071: saliency by itself surfaces confounders that increase in-sample fit but hurt generalisation. B3 (17 genes, LLM-filtered) surpasses both by AUC +0.024+0.024 over B1, using 294× fewer features. This establishes that the LLM reasoning step is causally necessary: the performance gain is not from the Mamba architecture; it is from the removal of confounders that reasoning identifies as biologically implausible. 4.2 Faithfulness Audit Table 2 reports the selection-level faithfulness of the 17-gene output against the 101-gene ground-truth set. Table 2: Faithfulness of LLM-selected gene set against curated BRCA ground truth (COSMIC CGC + OncoKB + PAM50, N=101N=101 validated genes). The top-50 saliency input contained 16 validated BRCA genes. Metric Value Selected genes with validated BRCA evidence 6/17 (35.3%) Known non-BRCA genes incorrectly kept 3/17 (17.6%) Genes with no ground-truth label (unverifiable) 8/17 (47.1%) Known BRCA genes available in top-50 input 16 Correctly kept by LLM 6 Missed by LLM (false negatives) 10 Recall on validated input genes 0.375 Correctly kept validated genes. MLPH (luminal A marker), ZEB1 (EMT/TNBC master regulator), XBP1 (ER-stress/luminal), INPP4B (PI3K/AKT tumour suppressor), RHOB (PAM50 RhoGTPase), THY1 (breast cancer stem cell marker). Incorrectly kept non-BRCA genes. ITGAL (CD11a, immune adhesion molecule with no documented breast-specific role), LMX1B (kidney/neural transcription factor), PRKAG2-AS1 (antisense RNA with no established BRCA function)—all of which meet rejection criterion R3, R4, or R5 in the prompt, yet were kept. Critical false negative. FOXA1—the pioneer transcription factor that defines luminal lineage in breast cancer and appears in both PAM50 and multiple COSMIC entries—was present in the top-50 input (rank 49) but not selected. The LLM’s stated rationale was consistent with its “Notch pathway in cancers” explanation for RHEX (rank 1), suggesting the model was anchored on saliency-correlated reasoning rather than pathway specificity for lower-ranked genes. Selective faithfulness. Despite a recall of only 0.375 on known BRCA genes, B3 achieves AUC 0.927, exceeding the 5,000-gene baseline. We hypothesise that the 3 incorrect non-BRCA keeps (ITGAL, LMX1B, PRKAG2-AS1) contribute noise but are outnumbered by the 6 high-specificity true positives (MLPH, ZEB1, XBP1, INPP4B, RHOB, THY1). In noise-filtering tasks, precision of reasoning on true negatives (correctly rejecting confounders) matters more than recall of true positives—the LLM’s value lies primarily in confounder removal. 5 Conclusion We have presented a neuro-symbolic framework that successfully integrates Mamba SSM gradient saliency with structured LLM chain-of-thought reasoning for genomic feature selection. Our findings confirm that LLM reasoning is a causally necessary component of this pipeline, achieving significant performance gains through precision-oriented confounder removal surpassing the 5,000-gene baseline by an AUC of +0.024 while using 294x fewer features. However, the identified gap between downstream performance and biological recall (missing 62.5% of known true positives) highlights the limitations of current domain-specific reasoning and underscores that task-level metrics are an unreliable proxy for reasoning faithfulness. This motivates the next phase of this research: expanding the pipeline into a robust, real-time diagnostic platform. Future iterations will focus on training across more diverse genomic datasets to improve reasoning recall and deploying a generalized architecture that allows researchers to identify disease-specific biomarkers for any target condition in real-time. By moving beyond a controlled BRCA-specific setting, this neuro-symbolic approach aims to provide a scalable, interpretable solution for high-dimensional bioinformatics modeling. Acknowledgments This work was supported by compute credits from a Cohere Labs Research Grant and the University of Central Missouri (UCM) Graduate Student Scholarly Research Fund. The authors also thank the UCM Office of Graduate Studies for their support of this research. References Gu & Dao (2024) Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces, 2024. URL https://arxiv.org/abs/2312.00752. Jeong et al. (2025) Daniel P. Jeong, Zachary C. Lipton, and Pradeep Ravikumar. Llm-select: Feature selection with large language models, 2025. URL https://arxiv.org/abs/2407.02694. Lanham et al. (2023) Tamera Lanham, Anna Chen, Ansh Radhakrishnan, Benoit Steiner, Carson Denison, Danny Hernandez, Dustin Li, Esin Durmus, Evan Hubinger, Jackson Kernion, Kamilė Lukošiūtė, Karina Nguyen, Newton Cheng, Nicholas Joseph, Nicholas Schiefer, Oliver Rausch, Robin Larson, Sam McCandlish, Sandipan Kundu, Saurav Kadavath, Shannon Yang, Thomas Henighan, Timothy Maxwell, Timothy Telleen-Lawton, Tristan Hume, Zac Hatfield-Dodds, Jared Kaplan, Jan Brauner, Samuel R. Bowman, and Ethan Perez. Measuring faithfulness in chain-of-thought reasoning, 2023. URL https://arxiv.org/abs/2307.13702. Lee et al. (2025) Joseph Lee, Shu Yang, Jae Young Baik, Xiaoxi Liu, Zhen Tan, Dawei Li, Zixuan Wen, Bojian Hou, Duy Duong-Tran, Tianlong Chen, and Li Shen. Knowledge-driven feature selection and engineering for genotype data with large language models, 2025. URL https://arxiv.org/abs/2410.01795. Parker et al. (2009) Joel S. Parker, Michael Mullins, Maggie C.U. Cheang, Leung S., Voduc D., Vickery T., Davies S., Fauron C., He X., Hu Z., Quackenbush J.F., Stijleman I.J., Palazzo J., Marron J.S., Nobel A.B., Mardis E., Nielsen T.O., Ellis M.J., Perou C.M., and Bernard P.S. Supervised risk predictor of breast cancer based on intrinsic subtypes. Journal of Clinical Oncology, 27(8):1160–1167, 2009. doi: 10.1200/JCO.2008.18.1370. URL https://pubmed.ncbi.gov/19204204/. Pudjihartono et al. (2022) Nicholas Pudjihartono, Tayaza Fadason, Andreas W. Kempa-Liehr, and Justin M. O’Sullivan. A review of feature selection methods for machine learning-based disease risk prediction. Frontiers in Bioinformatics, 2:927312, 2022. doi: 10.3389/fbinf.2022.927312. URL https://w.frontiersin.org/journals/bioinformatics/articles/10.3389/fbinf.2022.927312/full. Simonyan et al. (2014) Karen Simonyan, Andrea Vedaldi, and Andrew Zisserman. Deep inside convolutional networks: Visualising image classification models and saliency maps, 2014. URL https://arxiv.org/abs/1312.6034. Turpin et al. (2023) Miles Turpin, Julian Michael, Ethan Perez, and Samuel R. Bowman. Language models don’t always say what they think: Unfaithful explanations in chain-of-thought prompting, 2023. URL https://arxiv.org/abs/2305.04388. Zhang et al. (2025) Erica Zhang, Ryunosuke Goto, Naomi Sagan, Jurik Mutter, Nick Phillips, Ash Alizadeh, Kangwook Lee, Jose Blanchet, Mert Pilanci, and Robert Tibshirani. Llm-lasso: A robust framework for domain-informed feature selection and regularization, 2025. URL https://arxiv.org/abs/2502.10648. Appendix A Technical Architecture and Pipeline A.1 Model Specification The Mamba-SSM architecture used for the pretext classification task is defined as follows: OfficialMambaClassifier( embedding: Linear(1, 128) mamba: Mamba(d_model=128, d_state=16, d_conv=4, expand=2) pool: AdaptiveAvgPool1d(1) fc: Linear(128, 1) sigmoid: Sigmoid() ) Input Flow: (batch, NgenesN_genes) → (batch, NgenesN_genes, 1) → (batch, NgenesN_genes, 128) → Mamba Selective Scan → Adaptive Average Pooling over NgenesN_genes → (batch, 128) → Sigmoid output. Appendix B Discussion: Reasoning and Faithfulness Does downstream performance measure reasoning faithfulness? Our results reveal a systematic divergence between downstream predictive performance and reasoning faithfulness. While the LLM-selected gene set achieves a higher AUC than the 5,000-gene variance baseline, it attains a recall of only 0.375 on validated BRCA-associated genes. Under conventional evaluation, such recall would indicate poor biological reasoning. However, the improved classification performance suggests that the primary contribution of the LLM in this pipeline is not comprehensive biomarker discovery, but precision-oriented confounder rejection. In high-dimensional genomic settings, correctly eliminating spurious, tissue-specific, or immune-correlated genes appears more consequential for generalization than exhaustive recall of all known disease drivers. This finding cautions against using task-level performance as a proxy for reasoning faithfulness. Structured CoT versus saliency-driven selection Although the LLM was constrained by explicit rejection and keep criteria, the final gene set retained a 17.6% false-positive rate, including genes without established BRCA relevance. In several cases, the model justified these selections using generic or hallucinated pathway language (e.g., invoking “NF-κ signaling” without supporting breast-specific evidence). This indicates that structured chain-of-thought prompting can reduce but not eliminate biologically implausible reasoning. Moreover, the omission of FOXA1—a canonical luminal breast cancer regulator present in the input candidate pool—demonstrates that the LLM may confidently reject well-established true positives, particularly when they appear lower in the saliency ranking. These findings underscore that structured CoT improves filtering behavior but does not guarantee faithful biological reasoning. Scope and limitations This study intentionally focuses on a single, controlled setting (TCGA-BRCA) to isolate the interaction between neural saliency and LLM-mediated reasoning. As a result, the empirical conclusions should be interpreted as case-specific rather than universally generalizable. All downstream evaluations rely on a single stratified train–test split with a fixed random seed; statistical variability across seeds and datasets remains an important direction for future work. Additionally, comparisons to classical feature-selection baselines such as LASSO or ElasticNet were not conducted, and the reported gains should therefore be interpreted relative to the evaluated neural baselines only. Finally, the term “causal necessity” is used operationally to denote necessity within this pipeline configuration, rather than a formal causal identification of reasoning processes. Decision-level faithfulness analysis Due to a parsing mismatch in llm_gene_reasoning.json, per-gene decision strings were not fully captured in this version, limiting analysis to selection-level outcomes. We plan to re-run the reasoning parser for the feature study to enable decision-level precision, recall, and consistency metrics. This will allow a finer-grained audit of which rejection criteria fail most frequently and whether specific hallucination patterns correlate with gene rank or pathway ambiguity. Appendix C Ground-Truth Gene Sources Validation of LLM-selected genes was performed against the following established sources: • COSMIC Cancer Gene Census (Tier 1): including TP53, PIK3CA, CDH1, BRCA1, ERBB2, among others. • PAM50 intrinsic subtype genes (Parker et al., 2009): ESR1, PGR, FOXA1, MLPH. • Canonical signaling pathways: estrogen receptor signaling, PI3K/AKT, EMT/TNBC-associated markers. Appendix D Performance Benchmarks and Saliency Profiles D.1 Quantitative Comparison As shown in Figure 2, our Neuro-Symbolic approach (B3) demonstrates that logical filtering of features provides superior generalizability compared to purely data-driven methods (B1 and B2). Figure 2: Performance Comparison: Accuracy, F1, and AUC metrics across experimental conditions. B3 utilizes 250x fewer genes than B1 but achieves higher predictive stability. D.2 Neural Saliency Noise Profile Figure 3 visualizes the raw gradient saliency from which the LLM must extract biological signal. The high degree of variance across non-oncogenic clusters illustrates the necessity of the symbolic filtering layer. Figure 3: Raw Gradient Saliency Heatmap for the top-50 genes in TCGA-BRCA samples, highlighting the noisy feature space prior to LLM intervention. Appendix E LLM Reasoning and Prompt Design E.1 Chain-of-Thought Qualitative Audit Figure 4 illustrates the internal reasoning blocks (<<think>>) generated by the agent. This transparency allows for a post-hoc biological audit of the selection process. Figure 4: Visualization of the Agentic Chain-of-Thought (CoT) process, mapping Mamba saliency to biological rationale. E.2 Structured Prompt Design The agent is governed by five logical rules: (i) narrow saliency calibration; (i) rank-score association; (i) mandatory per-gene justification; (iv) anti-lazy cutoff enforcement; and (v) trivial solution auditing. Appendix F Biological Ground-Truth References To evaluate the ”faithfulness” of the LLM reasoning, selections were compared against: • COSMIC CGC Tier 1: TP53, PIK3CA, CDH1, BRCA1, ERBB2. • PAM50 Intrinsic Subtype Genes: ESR1, PGR, FOXA1, MLPH. • Functional Drivers: PI3K/AKT and EMT signaling pathways.