[Paper Review] LUMOS: Democratizing SciML Workflows with L0-Regularized Learning for Unified Feature and Parameter Adaptation

Shouwei Gao, Xu Zheng+3 more|arXiv (Cornell University)|Feb 26, 2026

Machine Learning in Materials Science0 citations

TL;DR

Lumos unifies feature selection and model pruning using L0-regularized learning with semi-stochastic gates to automate SciML model design during training or fine-tuning. It achieves significant parameter and feature reductions with minimal overhead.

ABSTRACT

The rapid growth of scientific machine learning (SciML) has accelerated discovery across diverse domains, yet designing effective SciML models remains a challenging task. In practice, building such models often requires substantial prior knowledge and manual expertise, particularly in determining which input features to use and how large the model should be. We introduce LUMOS, an end-to-end framework based on L0-regularized learning that unifies feature selection and model pruning to democratize SciML model design. By employing semi-stochastic gating and reparameterization techniques, LUMOS dynamically selects informative features and prunes redundant parameters during training, reducing the reliance on manual tuning while maintaining predictive accuracy. We evaluate LUMOS across 13 diverse SciML workloads, including cosmology and molecular sciences, and demonstrate its effectiveness and generalizability. Experiments on 13 SciML models show that LUMOS achieves 71.45% parameter reduction and a 6.4x inference speedup on average. Furthermore, Distributed Data Parallel (DDP) training on up to eight GPUs confirms the scalability of

Motivation & Objective

• Motivate the need to automate feature selection and model efficiency in SciML workflows.
• Propose a unified framework that couples feature selection with parameter pruning during training or fine-tuning.
• Enable on-the-fly, self-guided optimization that reduces manual tuning and domain knowledge requirements.

Proposed method

• Introduce semi-stochastic gating parameters that are learned during training to gate input features and neuron outputs.
• Apply a differentiable reparameterization (hard concrete) to enable gradient-based optimization of gates.
• Integrate gate mechanisms with common SciML layers (fully connected, convolutional, GIN, GCN, attention) and align architectures through structural consistency mapper functions.
• Use threshold-based gating to prune features and parameters while maintaining model outputs via a consistency toolkit.
• Provide an end-to-end workflow that yields an optimized model with a specified feature set and pruned structure.
• Demonstrate low overhead by showing gate parameters constitute ~3% of total parameters on average.

Experimental results

Research questions

• RQ1 Can Lumos simultaneously select informative features and prune model parameters without sacrificing accuracy?
• RQ2 How does semi-stochastic gating affect training dynamics and final model size across diverse SciML workloads?
• RQ3 What is the overhead of integrating Lumos into existing SciML training or fine-tuning pipelines?
• RQ4 How well does Lumos scale when applied to multi-GPU distributed training?

Key findings

• Lumos achieves 71.45% parameter reduction across 13 SciML workloads.
• Lumos delivers an average 6.4× speedup in inference.
• Gate parameters add about 3% overhead on average.
• Lumos is effective across five layer types: fully connected, convolutional, GIN, GCN, and attention.
• DDP training on up to eight GPUs confirms scalability.