LLM-Driven Reasoning for Constraint-Aware Feature Selection in Industrial Systems

1. Problem Statement

Feature selection is a critical bottleneck in large-scale industrial machine learning systems (e.g., recommendation and notification systems). Traditional methods face significant limitations in production environments:

• Data Scarcity: Many production scenarios lack sufficient labeled training data required for statistical methods (e.g., Lasso, mutual information).
• Operational Constraints: Traditional statistical approaches often ignore business constraints, such as feature group dependencies, maintenance overhead, and capacity limits.
• Semantic Gap: Standard methods rely purely on numerical correlations, failing to leverage rich semantic metadata (feature definitions, types, and domain knowledge) available in production systems.
• Scalability: Handling thousands of heterogeneous features exceeds the context windows of standard models and makes brute-force search infeasible.

The paper proposes a need for a framework that can perform constraint-aware, context-aware feature selection using both semantic understanding and quantitative data without relying heavily on labeled training sets.

2. Methodology: Model Feature Agent (MoFA)

The authors propose MoFA, a Large Language Model (LLM)-driven framework that treats feature selection as a sequential reasoning process.

Core Workflow

MoFA operates by integrating structured prompts containing:

1. Task Objective: The specific prediction goal (e.g., engagement lift, accuracy).
2. Feature Metadata: A tuple for each feature $f_i = (name_i, description_i, meta_i)$, where $meta_i$ includes quantitative data (importance scores, correlations) and categorical data (feature groups, types).
3. Constraints: Operational requirements (e.g., "minimize feature groups," "maximize coverage").

The LLM acts as a reasoning agent ($\Phi_\theta$) that iteratively selects the next best feature $f^*_t$ based on the current set $S_t$ and the remaining candidates $A_t$, outputting both the selection and a natural language justification.

Scalability Strategy: Divide-and-Conquer

To handle industrial-scale feature spaces (thousands of features) within LLM context limits, MoFA employs a two-phase strategy:

1. Phase 1: Bucketed Parallel Selection: The total feature set $F$ is partitioned into $B$ disjoint buckets. The LLM performs sequential selection within each bucket independently to identify local candidate subsets ($S_b$).
2. Phase 2: Global Synthesis and Refinement: The local subsets are merged into a refined pool ($F_{refine}$). The LLM performs a final round of sequential reasoning on this smaller pool to select the final target size $K$. This mitigates collinearity and ensures global consistency.

3. Key Contributions

• Constraint-Aware Reasoning: Unlike previous LLM-based approaches that focus solely on predictive utility, MoFA explicitly incorporates operational constraints (e.g., feature group consolidation) into the reasoning loop.
• Hybrid Information Integration: The framework uniquely combines semantic information (natural language descriptions) with quantitative metadata (correlations, scores) to guide selection.
• Scalable Architecture: The Divide-and-Conquer approach enables the application of LLM reasoning to feature spaces with thousands of dimensions, a scale previously inaccessible to direct LLM prompting.
• Interpretability: Every selection is accompanied by a natural language rationale, making the feature selection process transparent and auditable for engineers.

4. Experimental Results

The authors evaluated MoFA across three distinct real-world industrial use cases at Meta:

A. True Interest and Time-Worthiness Prediction

• Task: Predicting if content reflects user interest or is worth their time.
• Setup: Compared MoFA against Lasso regression on 1,030 features.
• Results:
  • Accuracy: MoFA matched or slightly outperformed Lasso in small-budget scenarios (e.g., +0.63% AUC at $K=20$).
  • Operational Efficiency: MoFA significantly reduced the number of feature groups required. At $K=500$, MoFA used 85 groups compared to Lasso's 187 groups (a 54% reduction), drastically lowering cross-team maintenance overhead.

B. Value Model Enhancement (Signal Pair Selection)

• Task: Identifying high-order interaction terms (e.g., $p_{like} \times p_{click}$) for a linear value model.
• Setup: Selected top interactions from ~400 candidate pairs.
• Results: The top 3 interactions selected by MoFA were integrated into the production value model, yielding statistically significant gains in daily app sessions (+0.055%, +0.048%, +0.018%), demonstrating the discovery of non-obvious, high-value feature interactions.

C. Notification Behavior Prediction

• Task: Multi-task prediction of user click and dismissal probabilities.
• Setup: Selected 4,000 features from a pool of 8,169.
• Results:
  • MoFA outperformed the existing production model (which used 7,500+ features) and a random selection baseline.
  • Metrics: Achieved a 0.065% NE win for click prediction and 0.198% NE win for dismissal prediction compared to the production model.
  • Compared to random selection, MoFA showed even larger gains (up to 0.332% NE win), proving its ability to filter noise in high-dimensional spaces.

5. Significance and Limitations

Significance:
MoFA bridges the gap between theoretical LLM capabilities and practical industrial engineering. It demonstrates that LLMs can serve as effective "reasoning agents" for feature engineering, offering a flexible alternative to rigid statistical methods. It is particularly valuable in scenarios with limited labeled data but rich feature metadata, enabling models that are not only accurate but also operationally sustainable.

Limitations:

• Latency: The sequential, one-by-one selection process scales linearly with the number of features, which can be slow for massive feature sets (though acceptable for offline preprocessing).
• Context Limits: Even with the Divide-and-Conquer strategy, providing comprehensive metadata for thousands of features remains a challenge for LLM token limits.
• Greedy Search: The sequential nature is a greedy search, potentially missing optimal feature combinations that only exhibit synergy when selected simultaneously (non-additive interactions).

Conclusion

The paper establishes MoFA as a practical, effective framework for modern industrial ML systems. By leveraging LLM reasoning to balance predictive performance with operational constraints, it offers a scalable solution for feature selection that improves model accuracy, inference efficiency, and system maintainability.