# CoFEE: Reasoning Control for LLM-Based Feature Discovery

arXiv: https://arxiv.org/abs/2604.21584 · HTML: https://mlsanity.com/paper/2604.21584/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Reasoning benchmarks and evaluation, one of 17 directions in LLM reasoning.
- Builds on 7 of its 7 references that are in the map.
- First posted on arXiv 23 Apr 2026.

## Placed in

- LLM reasoning / [Reasoning benchmarks and evaluation](https://mlsanity.com/lr/reasoning-benchmarks-and-evaluation/)

## Builds on (most cited first, 7 of 7)

- [ReAct: Synergizing Reasoning and Acting in Language Models](https://mlsanity.com/paper/2210.03629/) (2022)
- [Cognitive Behaviors that Enable Self-Improving Reasoners, or, Four Habits of Highly Effective STaRs](https://mlsanity.com/paper/2503.01307/) (2025)
- [Multimodal Chain-of-Thought Reasoning: A Comprehensive Survey](https://mlsanity.com/paper/2503.12605/) (2025)
- [Position: AI Evaluation Should Learn from How We Test Humans](https://mlsanity.com/paper/2306.10512/) (2023)
- [An Empirical Analysis of Feature Engineering for Predictive Modeling](https://mlsanity.com/paper/1701.07852/) (2017)
- [AStar: Boosting Multimodal Reasoning with Automated Structured Thinking](https://mlsanity.com/paper/2502.02339/) (2025)
- [GPTree: Towards Explainable Decision-Making via LLM-powered Decision Trees](https://mlsanity.com/paper/2411.08257/) (2024)

## Abstract (arXiv)

Feature discovery from complex unstructured data is fundamentally a reasoning problem: it requires identifying abstractions that are predictive of a target outcome while avoiding leakage, proxies, and post-outcome signals. With the introduction of ever-improving Large Language Models (LLMs), our method provides a structured method for addressing this challenge. LLMs are well suited for this task by being able to process large amounts of information, but unconstrained feature generation can lead to weak features. In this work, we study reasoning control in LLMs by inducing cognitive behaviors for improving feature discovery. We introduce CoFEE (Cognitive Feature Engineering Engine), a reasoning control framework that enforces cognitive behaviors in how the LLM reasons during feature discovery. From a machine learning perspective, these cognitive behaviors act as structured inductive biases over the space of candidate features generated by the model. These behaviors have been exploited with success in ML models, and include backward chaining from outcomes, subgoal decomposition, verification against observability and leakage criteria, and explicit backtracking of rejected reasoning paths. In a controlled comparison, we show that enforcing cognitive behaviors yields features with higher empirical predictability than those under unconstrained vanilla LLM prompts. CoFEE achieves an average Success Rate Score that is 15.2% higher than the vanilla approach, while generating 29% fewer features and reducing costs by 53.3%. Using held-out feature evaluation, we assess whether cognitively induced features generalize beyond the data used for discovery. Our results indicate that, in our evaluated setting, reasoning control is associated with improvements in quality and efficiency of LLM-based feature discovery.