CoFEE: Reasoning Control for LLM-Based Feature Disc...

CoFEE: Reasoning Control for LLM-Based Feature Disc... | AI Research

Key Takeaways

CoFEE: Reasoning Control for LLM-Based Feature Discovery This paper introduces CoFEE (Cognitive Feature Engineering Engine), a framework designed to improve...
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.

Paper AbstractExpand

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.

CoFEE: Reasoning Control for LLM-Based Feature Discovery
This paper introduces CoFEE (Cognitive Feature Engineering Engine), a framework designed to improve how Large Language Models (LLMs) identify predictive features in complex, unstructured data. While LLMs are powerful, they often generate weak or irrelevant features when left to their own devices. CoFEE addresses this by enforcing specific "cognitive behaviors" through structured prompting, effectively guiding the model to reason more like a human expert during the feature discovery process.

Inducing Cognitive Reasoning

The core of CoFEE is "reasoning control," which acts as a structured inductive bias for the LLM. Rather than simply asking the model to generate features, the framework forces the LLM to follow four specific cognitive steps:

Backward Chaining: Starting from the desired outcome and working backward to identify the mechanisms that lead to it.
Subgoal Decomposition: Breaking down the complex task of feature discovery into smaller, manageable categories like market structure or team dynamics.
Verification: Systematically checking each candidate feature against criteria like observability and the risk of using "leaked" information (data that wouldn't be available before the outcome).
Backtracking: Explicitly recording and rejecting reasoning paths that fail, which helps the model avoid similar mistakes in future steps.

The CoFEE Pipeline

CoFEE operates as an agent-based system that automates the discovery process. The pipeline consists of three specialized agents: the first performs the cognitively constrained feature discovery, the second consolidates semantically similar features to reduce redundancy, and the third scores the features based on their ability to distinguish between successful and unsuccessful outcomes in a dataset of founder profiles. By using this structured approach, the system ensures that the final list of features is both high-quality and interpretable.

Performance and Efficiency

In a controlled comparison against a "vanilla" LLM approach, CoFEE demonstrated significant improvements in both quality and cost-efficiency. Features generated using CoFEE showed higher predictive power, with an average Success Rate Score 15.2% higher than those produced by the unconstrained baseline. Furthermore, because CoFEE’s structured reasoning led to more focused generation, it produced 29% fewer features overall, resulting in a 53.3% reduction in computational costs.

Limitations and Future Outlook

While the results are promising, the authors note several areas for further study. The current evaluation is limited to a single domain—venture capital—and does not yet confirm if these cognitive behaviors provide the same benefits across different industries or model architectures. Additionally, while the features show higher empirical predictability, future work is needed to assess how they perform in downstream machine learning tasks and to test the robustness of these cognitive prompts across a wider variety of settings.