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
Figures & tables
Fig. 1 : (S1) Training from scratch: domain scientists must select physics-informative scientific input features and co-design the model architecture and parameters, often relying on costly grid search and ad hoc heuristics. (S2) Fine-tuning a foundation model: the pretrained backbone is fixed, and adaptation is achieved by re-mapping input features and adjusting subsets of parameters. Both workflows face common bottlenecks in feature efficiency and model/parameter efficiency , motivating the development of unified frameworks such as Lumos .
Fig. 2 : Overall workflow of Lumos .
Fig. 3 : Illustration of Input Features: the blue words represent input features regarding two examples, climate application and power grid simulation, respectively.
Symbol
Description
F
SciML model
D
Dataset
X0
Input features
X(l)
Input to the l -th layer
Ygt
Ground truth (output label)
M
Set of evaluation metrics
TABLE I : Summary of Symbols
Fig. 4 : Integration with different model layers.
Fig. 5 : Structured transition between the CONV and FC layer. When a CONV layer is followed by an FC layer, the output of the CONV layer serves as the input to the FC layer; therefore, the redundant outputs corresponding to the deactivated neurons of the CONV layer must also be removed to maintain tensor alignments.
Fig. 6 : Structured transition between the embedding and FC layer. When input features are deactivated, the corresponding weights in the following FC layer must also be removed.
Fig. 7 : Dependency between concatenation. When two outputs of two CONV layers are concatenated before being connected to another CONV layer, the indices of the CONV kernels must be remapped to maintain structural integrity.
Fig. 8 : Layer-wise pruning within a residual architecture.
Fig. 9 : Comparison between baseline SciML models and Lumos across latency, accuracy, parameter count, and FLOPs.
Metrics
OTO
L1
NEURAL
DepGraph
Lumos
One-Shot
Yes
No
No
Yes
Yes
Compression
Yes
Yes
No
Yes
Yes
Adaptiveness
No
No
No
No
Yes
Puremd(% ↑ )
72.73%
76.3%
80.10%
76.35%
76.10%
Puremd(R² ↑ )
0.9430
0.9969
0.9318
0.9937
0.9988
CFD(% ↑ )
99.10%
99.12%
99.10%
99.12%
99.10%
TABLE III : Comparison with SOTA pruning methods
Fig. 10 : Showcase I: Feature selection efficiency in three regression workloads (CFD, PureMD, and Fluid).
Fig. 11 : Peak Memory Usage and Energy Saving.
Fig. 12 : Scale Lumos up to 8 GPUs using DDP training.
Fig. 13 : Showcase II: Training loss over time for the baseline and Lumos -augmented models on the PureMD workload. Lumos achieves faster convergence, lower final loss, and greater training stability compared with the baseline training process.
Workload
Original
LUMOS
Overhead
cifar10
55
61
+10.91%
DMSNet
80
82
+2.50%
PPA
128
137
+7.03%
Molhiv
43
41
−4.65%
Slstr
129
131
+1.55%
CFD
34
37
+8.82%
TABLE IV : Training-time(minutes) overhead of integrating LUMOS.