Paper deep dive
Mamba-SSM with LLM Reasoning for Feature Selection: Faithfulness-Aware Biomarker Discovery
Pushpa Kumar Balan, Aijing Feng
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 98%
Last extracted: 6/21/2026, 11:59:00 AM
Summary
This paper presents a neuro-symbolic pipeline for biomarker discovery that combines Mamba SSM gradient saliency with DeepSeek-R1 Chain-of-Thought (CoT) reasoning. The authors demonstrate that while raw gradient saliency from Mamba models can be contaminated by tissue-composition confounders, LLM-based reasoning can effectively filter these to improve downstream classification performance (AUC 0.927 vs 0.832 for raw saliency). However, a faithfulness audit reveals 'selective faithfulness': the LLM successfully removed confounders to boost performance but failed to achieve high recall of known BRCA biomarkers (e.g., missing the canonical gene FOXA1), suggesting that task-level performance is not a perfect proxy for biological reasoning accuracy.
Entities (10)
Relation Signals (4)
DeepSeek R1 â failedtoselect â FOXA1
confidence 100% ¡ FOXA1... was present in the input but rejected by the LLM
DeepSeek R1 â filtersfeaturesfrom â Mamba-SSM
confidence 100% ¡ DeepSeek-R1 evaluates every candidate with structured CoT to produce a final 17-gene set
Mamba-SSM â generatessaliencyfor â TCGA-BRCA
confidence 100% ¡ We train a Mamba SSM on TCGA-BRCA RNA-seq and extract the top-50 genes by gradient saliency
FOXA1 â ispartof â PAM50
confidence 100% ¡ FOXA1, the most important luminal breast cancer pioneer transcription factor and a canonical PAM50 gene
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:Gradient saliency from deep sequence models surfaces candidate biomarkers efficiently, but the resulting gene lists can be contaminated by tissue-composition confounders that degrade downstream classifiers. We study whether LLM chain-of-thought (CoT) reasoning can filter these confounders, and whether reasoning quality is associated with 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. On the held-out test split, 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) shows that 6 of 17 selected genes (35.3%) are validated BRCA biomarkers, while 10 of 16 known BRCA genes present in the input were missed - including FOXA1. This divergence between downstream performance and reasoning faithfulness suggests selective faithfulness in this setting: targeted confounder removal can improve predictive performance without comprehensive recall.
Tags
Links
- Source: https://arxiv.org/abs/2604.14334v2
- Canonical: https://arxiv.org/abs/2604.14334v2
Trouble viewing inline? Open PDF directly â
Full Text
21,537 characters extracted from source content.
Expand or collapse full text
ICLR 2026 Workshop on Logical Reasoning of Large Language Models MAMBA-SSM WITH LLM REASONING FOR FEATURE SELECTION: FAITHFULNESS-AWARE BIOMARKER DISCOVERY 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 ef- ficiently, but the resulting gene lists can be contaminated by tissue-composition confounders that degrade downstream classifiers. We study whether LLM chain- of-thought (CoT) reasoning can filter these confounders, and whether reasoning quality is associated with 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. On the held-out test split, 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) shows that 6 of 17 selected genes (35.3%) are validated BRCA biomarkers, while 10 of 16 known BRCA genes present in the input were missedâincluding FOXA1. This divergence between downstream performance and reasoning faithfulness suggests selective faithfulness in this setting: targeted confounder removal can improve predictive performance without comprehensive recall. The scope of this claim and its limitations are examined in Sec. B. Code: https://github.com/pushpakumarbalan/feature-selection 1INTRODUCTION 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 (Pud- jihartono 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 empirically important 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 1 arXiv:2604.14334v2 [q-bio.QM] 17 Apr 2026 ICLR 2026 Workshop on Logical Reasoning of Large Language Models 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. 2RELATED 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, pe- nalising 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 officialmamba-ssmimplementation with gradient saliency (Simonyan et al., 2014) to extract a biologically plausible candidate pool before the LLM reasoning step. 3METHODOLOGY 3.1SYSTEM PIPELINE Figure 1: End-to-end neuro-symbolic pipeline. High-dimensional TCGA-BRCA RNA-seq is pro- cessed by Mamba SSM to produce gradient saliency scores (top-50 candidate genes), then filtered by structured DeepSeek-R1 CoT reasoning to form a 17-gene BRCA-specific subset used in downstream comparison. 2 ICLR 2026 Workshop on Logical Reasoning of Large Language Models 3.2DATA 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.3PHASE 1 â MAMBA SSM TRAINING Input features are log 2 (TPM+1)-normalised and filtered to the top-5,000 highest-variance genes. We train anOfficialMambaClassifier: a linear embedding layer projects each geneâs scalar expression value into ad model = 128dimensional space; a single Mamba block (d state = 16, d conv = 4, expand= 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 â4 , 15 epochs, batch size 8) with class-weighted BCELoss (w normal = N tumour /N normal ) to handle the 8.8:1 tumour/normal imbalance. 3.4PHASE 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 iss j = 1 |T| P iâT âL âx ij , averaged over all tumour samples T . The top-50 genes by s j form the candidate poolG 50 . 3.5PHASE 3 â STRUCTURED COT REASONING G 50 is passed to DeepSeek-R1 (7B, local via Ollama, temperature= 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 selected subset contains 17 genes and is used as the B3 input set in downstream comparison. 3.6PHASE 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 selection outcome (selected / not selected), yielding selection-level precision, recall, and missed-gene analysis. 4RESULTS 4.1DOWNSTREAM 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. MethodGenes AccuracyF1 AUC B1 Variance baseline5,0000.87850.89410.903 B2 Mamba saliency only (no LLM)500.72470.78130.832 B3 Mamba + LLM structured CoT170.8907 0.9033 0.927 3 ICLR 2026 Workshop on Logical Reasoning of Large Language Models 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: saliency by itself surfaces confounders that increase in-sample fit but hurt generalisation. B3 (17 genes, LLM-filtered) surpasses both by AUC+0.024over B1, using 294Ăfewer features. These results indicate that the LLM filtering step is empirically important in this pipeline, and that the gain is consistent with removal of biologically implausible confounders from the saliency-derived candidate set. 4.2FAITHFULNESS 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 = 101validated genes). The top-50 saliency input contained 16 validated BRCA genes. MetricValue Selected genes with validated BRCA evidence6/17 (35.3%) Known non-BRCA genes incorrectly kept3/17 (17.6%) Genes with no ground-truth label (unverifiable)8/17 (47.1%) Known BRCA genes available in top-50 input16 Correctly kept by LLM6 Missed by LLM (false negatives)10 Recall on validated input genes0.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 docu- mented breast-specific role), LMX1B (kidney/neural transcription factor), PRKAG2-AS1 (antisense RNA with no established BRCA function) were retained despite lacking established BRCA-specific evidence. 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. This highlights a failure mode in which established BRCA-relevant genes can be rejected when they appear lower in the saliency ranking. Selective faithfulness. Despite a recall of only 0.375 on known BRCA genes, B3 achieves AUC 0.927, exceeding the 5,000-gene baseline. The three incorrect non-BRCA keeps (ITGAL, LMX1B, PRKAG2-AS1) add noise, but appear to be outweighed by six high-specificity true positives (MLPH, ZEB1, XBP1, INPP4B, RHOB, THY1). In this setting, precision-oriented confounder removal appears more influential for downstream performance than exhaustive recall of all known disease genes. 5CONCLUSION We presented a neuro-symbolic framework that integrates Mamba SSM gradient saliency with structured LLM chain-of-thought reasoning for genomic feature selection in TCGA-BRCA. In the current benchmark, the LLM-filtered 17-gene set outperformed both the 5,000-gene variance baseline and the 50-gene saliency-only condition, improving AUC by +0.024 over the variance baseline while using 294Ăfewer features. These results support the practical value of precision-oriented confounder removal in this pipeline. 4 ICLR 2026 Workshop on Logical Reasoning of Large Language Models At the same time, the gap between downstream performance and biological recall (missing 62.5% of known true positives in the top-50 input) highlights limitations of current domain-specific reasoning and reinforces that task-level metrics alone are an incomplete proxy for reasoning faithfulness. Future work will evaluate this framework across additional datasets and disease contexts, improve recall- oriented reasoning constraints, and establish more robust reproducibility and validation protocols beyond the current controlled BRCA setting. 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 Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces, 2024. URL https://arxiv.org/abs/2312.00752. 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. Tamera Lanham, Anna Chen, Ansh Radhakrishnan, Benoit Steiner, Carson Denison, Danny Her- nandez, Dustin Li, Esin Durmus, Evan Hubinger, Jackson Kernion, Kamil Ě e Luko Ë si Ě ut Ě e, 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. 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. URLhttps://arxiv.org/ abs/2410.01795. 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/. 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 predic- tion. Frontiers in Bioinformatics, 2:927312, 2022.doi:10.3389/fbinf.2022.927312. URLhttps://w.frontiersin.org/journals/bioinformatics/articles/ 10.3389/fbinf.2022.927312/full. Karen Simonyan, Andrea Vedaldi, and Andrew Zisserman. Deep inside convolutional networks: Visualising image classification models and saliency maps, 2014. URLhttps://arxiv.org/ abs/1312.6034. 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. 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. URLhttps://arxiv.org/ abs/2502.10648. 5 ICLR 2026 Workshop on Logical Reasoning of Large Language Models ATECHNICAL ARCHITECTURE AND PIPELINE A.1MODEL 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,N genes )â(batch,N genes , 1)â(batch,N genes , 128)âMamba Selective Scanâ Adaptive Average Pooling over N genes â (batch, 128)â Sigmoid output. BDISCUSSION: 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 selectionAlthough 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-ÎşB 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 Duetoaparsingmismatchin llmgenereasoning.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. 6 ICLR 2026 Workshop on Logical Reasoning of Large Language Models CGROUND-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. DPERFORMANCE BENCHMARKS AND SALIENCY PROFILES D.1QUANTITATIVE 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.2NEURAL 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. ELLM REASONING AND PROMPT DESIGN E.1CHAIN-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. E.2STRUCTURED 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. 7 ICLR 2026 Workshop on Logical Reasoning of Large Language Models Figure 3: Raw Gradient Saliency Heatmap for the top-50 genes in TCGA-BRCA samples, highlight- ing the noisy feature space prior to LLM intervention. FBIOLOGICAL 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. 8 ICLR 2026 Workshop on Logical Reasoning of Large Language Models Figure 4: Visualization of the Agentic Chain-of-Thought (CoT) process, mapping Mamba saliency to biological rationale. ⢠Functional Drivers: PI3K/AKT and EMT signaling pathways. 9