A Synergistic Strategy for Data-Constrained Deep Learning in Materials Science

Abstract

Materials science research increasingly benefits from the application of machine learning methods, yet encounters fundamental challenges from data scarcity, such as limited dataset sizes and severe distribution imbalance. In this paper, we propose a hybrid framework integrating attention pooling, multi-task learning, auxiliary learning, and classification-corrected regression. Using a 2D materials dataset as a case study, our approach demonstrates significantly enhanced prediction accuracy over the baseline crystal graph convolutional neural networks (CGCNN) method. Specifically, it reduces the mean absolute error for work function prediction from 0.312 to 0.240 eV, and for band gap from 0.301 to 0.230 eV. The framework also proves effective with other graph neural network methods such as atomistic line graph neural network (ALIGNN). These gains stem from the framework's ability to exploit underlying physical correlations between material properties and atomic structures. Through extensive experiments, we demonstrate that attention pooling serves as a generally effective component for diverse property prediction tasks, particularly with small datasets, which also offers the possibility of interpretability analysis through element attention weights. Our architecture enables seamless integration with various graph-based or other end-to-end deep learning models, presenting a computationally efficient and easily implementable solution for constrained datasets in materials science.

Keywords

2D materials / attention pooling / data scarcity / graph neural networks / machine learning / multi-task learning

References

[1] K. T. Butler, D. W. Davies, H. Cartwright, O. Isayev, and A. Walsh, “Machine Learning for Molecular and Materials Science,” Nature 559 (2018): 547–555, https://www.nature.com/articles/s41586-018-0337-2.

[2] K. Choudhary, B. DeCost, C. Chen, et al., “Recent Advances and Applications of Deep Learning Methods in Materials Science,” npj Computational Materials 8 (2022): 59, https://doi.org/10.1038/s41524-022-00734-6.

[3] L. M. Ghiringhelli, J. Vybiral, S. V. Levchenko, C. Draxl, and M. Scheffler, “Big Data of Materials Science: Critical Role of the Descriptor,” Physical Review Letters 114, no. 10 (2015): 105503, https://doi.org/10.1103/physrevlett.114.105503.

[4] L. Ward, A. Agrawal, A. Choudhary, and C. Wolverton, “A General-Purpose Machine Learning Framework for Predicting Properties of Inorganic Materials,” npj Computational Materials 2 (2016): 16028, https://doi.org/10.1038/npjcompumats.2016.28.

[5] A. Jain and T. Bligaard, “Atomic-Position Independent Descriptor for Machine Learning of Material Properties,” Physical Review B: Condensed Matter 98, no. 21 (2018): 214112, https://doi.org/10.1103/physrevb.98.214112.

[6] X. Y. Ma, J. P. Lewis, Q. B. Yan, and G. Su, “Accelerated Discovery of Two-Dimensional Optoelectrosnic Octahedral Oxyhalides via High-Throughput Ab Initio Calculations and Machine Learning,” Journal of Physical Chemistry Letters 10, no. 21 (2019): 6734–6740, https://doi.org/10.1021/acs.jpclett.9b02420.

[7] S. Lu, Q. Zhou, L. Ma, Y. Guo, and J. Wang, “Rapid Discovery of Ferroelectric Photovoltaic Perovskites and Material Descriptors via Machine Learning,” Small Methods 3, no. 11 (2019): 1900360, https://doi.org/10.1002/smtd.201900360.

[8] Z. J. Zhao, S. Liu, S. Zha, et al., “Theory-Guided Design of Catalytic Materials Using Scaling Relationships and Reactivity Descriptors,” Nature Reviews Materials 4, no. 12 (2019): 792–804, https://doi.org/10.1038/s41578-019-0152-x.

[9] L. H. Mou, T. Han, P. E. S. Smith, E. Sharman, and J. Jiang, “Machine Learning Descriptors for Data-Driven Catalysis Study,” Advanced Science 10, no. 22 (2023): e2301020, https://doi.org/10.1002/advs.202301020.

[10] L. Giordano, K. Akkiraju, R. Jacobs, D. Vivona, D. Morgan, and Y. Shao-Horn, “Electronic Structure-Based Descriptors for Oxide Properties and Functions,” Accounts of Chemical Research 55, no. 3 (2022): 298–308, https://doi.org/10.1021/acs.accounts.1c00509.

[11] B. Wang and F. Zhang, “Main Descriptors to Correlate Structures With the Performances of Electrocatalysts,” Angewandte Chemie International Edition 61, no. 4 (2022): e202111026, https://doi.org/10.1002/ange.202111026.

[12] L. Himanen, M. O. J. Jäger, E. V. Morooka, et al., “DScribe: Library of Descriptors for Machine Learning in Materials Science,” Computer Physics Communications 247 (2020): 106949, https://doi.org/10.1016/j.cpc.2019.106949.

[13] T. Xie and J. C. Grossman, “Crystal Graph Convolutional Neural Networks for an Accurate and Interpretable Prediction of Material Properties,” Physical Review Letters 120, no. 14 (2018): 145301, https://doi.org/10.1103/physrevlett.120.145301.

[14] C. Chen, W. Ye, Y. Zuo, C. Zheng, and S. P. Ong, “Graph Networks as a Universal Machine Learning Framework for Molecules and Crystals,” Chemistry of Materials 31, no. 9 (2019): 3564–3572, https://doi.org/10.1021/acs.chemmater.9b01294.

[15] K. Choudhary and B. DeCost, “Atomistic Line Graph Neural Network for Improved Materials Property Predictions,” npj Computational Materials 7, no. 1 (2021): 185, https://doi.org/10.1038/s41524-021-00650-1.

[16] S. Y. Louis, Y. Zhao, A. Nasiri, et al., “Graph Convolutional Neural Networks With Global Attention for Improved Materials Property Prediction,” Physical Chemistry Chemical Physics 22 (2020): 18141–18148, https://doi.org/10.1039/D0CP01474E.

[17] K. Yan, Y. Liu, Y. Lin, and S. Ji, “Periodic Graph Transformers for Crystal Material Property Prediction,” in Proceedings of 36th International Conference on Neural Information Processing Systems (2022), 15066–15080, https://dl.acm.org/doi/10.5555/3600270.3601366.

[18] K. Yan, C. Fu, X. Qian, X. Qian, and S. Ji, “Complete and Efficient Graph Transformers for Crystal Material Property Prediction,” in 12th International Conference on Learning Representations (ICLR) (2024), 8934.

[19] H. Yu, Z. Xu, X. Qian, X. Qian, and S. Ji, “Efficient and Equivariant Graph Networks for Predicting Quantum Hamiltonian,” in Proceedings of the 40th International Conference on Machine Learning (ICML) (2023), 40412–40424, https://dl.acm.org/doi/10.5555/3618408.3620100.

[20] H. Du, J. Wang, J. Hui, L. Zhang, and H. Wang, “Densegnn: Universal and Scalable Deeper Graph Neural Networks for High-Performance Property Prediction in Crystals and Molecules,” npj Computational Materials 10, no. 1 (2024): 292, https://doi.org/10.1038/s41524-024-01444-x.

[21] S. Yang, K. Cho, A. Merchant, et al., “Scalable Diffusion for Materials Generation,” in 12th International Conference on Learning Representations (ICLR) (2024), 26525.

[22] A. Jain, S. P. Ong, G. Hautier, et al., “Commentary: The Materials Project: A Materials Genome Approach to Accelerating Materials Innovation,” APL Materials 1 (2013): 011002, https://doi.org/10.1063/1.4812323.

[23] J. E. Saal, S. Kirklin, M. Aykol, B. Meredig, and C. Wolverton, “Materials Design and Discovery With High-Throughput Density Functional Theory: The Open Quantum Materials Database (OQMD),” JOM 65, no. 11 (2013): 1501–1509, https://doi.org/10.1007/s11837-013-0755-4.

[24] S. Curtarolo, W. Setyawan, G. L. Wang, et al., “AFLOW: An Automatic Framework for High-Throughput Materials Discovery,” Computational Materials Science 58 (2012): 218–226, https://doi.org/10.1016/j.commatsci.2012.02.005.

[25] M. Liu and S. Meng, “Atomly.net Materials Database and Its Application in Inorganic Chemistry,” SCIENTIA SINICA Chimica 53 (2023): 19–25, https://doi.org/10.1360/SSC-2022-0167.

[26] K. Choudhary, K. F. Garrity, A. C. E. Reid, et al., “The Joint Automated Repository for Various Integrated Simulations (JARVIS) for Data-Driven Materials Design,” npj Computational Materials 6 (2020): 173, https://doi.org/10.1038/s41524-020-00440-1.

[27] A. Belsky, M. Hellenbrandt, V. L. Karen, and P. Luksch, “New Developments in the Inorganic Crystal Structure Database (ICSD): Accessibility in Support of Materials Research and Design,” Acta Crystallographica Section B: Structural Science 58, no. 3 (2002): 364–369, https://doi.org/10.1107/s0108768102006948.

[28] L. Sbailò, Á. Fekete, L. M. Ghiringhelli, and M. Scheffler, “The NOMAD Artificial-Intelligence Toolkit: Turning Materials-Science Data Into Knowledge and Understanding,” npj Computational Materials 8 (2022): 250, https://doi.org/10.1038/s41524-022-00935-z.

[29] S. Sarikurt, T. Kocabaş, and C. Sevik, “High-Throughput Computational Screening of 2D Materials for Thermoelectrics,” Journal of Materials Chemistry A 8 (2020): 19674–19683, https://doi.org/10.1039/D0TA04945J.

[30] T. Deng, P. Qiu, T. Yin, et al., “High-Throughput Strategies in the Discovery of Thermoelectric Materials,” Advanced Materials 36 (2024): 2311278, https://doi.org/10.1002/adma.202311278.

[31] D. Akinwande, C. Huyghebaert, C. Wang, et al., “Graphene and Two-Dimensional Materials for Silicon Technology,” Nature 573 (2019): 507–518, https://doi.org/10.1038/s41586-019-1573-9.

[32] G. Fiori, F. Bonaccorso, G. Iannaccone, et al., “Electronics Based on Two-Dimensional Materials,” Nature Nanotechnology 9, no. 10 (2014): 768–779, https://doi.org/10.1038/nnano.2014.207.

[33] K. S. Kim, J. Kwon, H. Ryu, et al., “The Future of Two-Dimensional Semiconductors Beyond Moore's Law,” Nature Nanotechnology 19 (2024): 895–906, https://doi.org/10.1038/s41565-024-01695-1.

[34] S. Haastrup, M. Strange, M. Pandey, et al., “The Computational 2D Materials Database: High-Throughput Modeling and Discovery of Atomically Thin Crystals,” 2D Materials 5, no. 4 (2018): 042002, https://doi.org/10.1088/2053-1583/aacfc1.

[35] C. Shorten and T. M. Khoshgoftaar, “A Survey on Image Data Augmentation for Deep Learning,” Journal of Big Data 6 (2019): 60, https://doi.org/10.1186/s40537-019-0197-0.

[36] X. Y. Ma, H. Y. Lyu, K. R. Hao, Z. G. Zhu, Q. B. Yan, and G. Su, “High-Efficient Ab Initio Bayesian Active Learning Method and Applications in Prediction of Two-Dimensional Functional Materials,” Nanoscale 13, no. 35 (2021): 14694–14704, https://doi.org/10.1039/d1nr03886a.

[37] V. Gupta, K. Choudhary, F. Tavazza, et al., “Cross-Property Deep Transfer Learning Framework for Enhanced Predictive Analytics on Small Materials Data,” Nature Communications 12, no. 1 (2021): 6595, https://doi.org/10.1038/s41467-021-26921-5.

[38] N. Shazeer, A. Mirhoseini, K. Maziarz, et al., “Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer,” in International Conference on Learning Representations (ICLR) (2017): 878.

[39] W. Fedus, B. Zoph, and N. Shazeer, “Switch Transformers: Scaling to Trillion Parameter Models With Simple and Efficient Sparsity,” Journal of Machine Learning Research 23 (2022): 5232–5270, https://dl.acm.org/doi/abs/10.5555/3586589.3586709.

[40] A. A. Cheenady, A. Mukherjee, R. Dongol, and K. Rajan, “Exploring Structure–Property Relationships in Sparse Data Environments Using Mixture-of-Experts Models,” MRS Bulletin 50 (2025): 32–43, https://doi.org/10.1557/s43577-024-00805-7.

[41] R. Chang, Y. X. Wang, and E. Ertekin, “Towards Overcoming Data Scarcity in Materials Science: Unifying Models and Datasets With a Mixture of Experts Framework,” npj Computational Materials 8 (2022): 242, https://doi.org/10.1038/s41524-022-00929-x.

[42] S. Li and A. Nakata, “CSIML: A Cost-Sensitive and Iterative Machine-Learning Method for Small and Imbalanced Materials Data Sets,” Chemistry Letters 53 (2024): upae090, https://doi.org/10.1093/chemle/upae090.

[43] D. Morgan and R. Jacobs, “Opportunities and Challenges for Machine Learning in Materials Science,” Annual Review of Materials Research 50 (2020): 71–103, https://doi.org/10.1146/annurev-matsci-070218-010015.

[44] P. Xu, X. Ji, M. Li, and W. Lu, “Small Data Machine Learning in Materials Science,” npj Computational Materials 9, no. 1 (2023): 42, https://doi.org/10.1038/s41524-023-01000-z.

[45] Q. Tao, J. Yu, X. Mu, et al., “Machine Learning Strategies for Small Sample Size in Materials Science,” Science China Materials 68 (2025): 387–405, https://doi.org/10.1007/s40843-024-3204-5.

[46] J. Jiang, C. Zhang, L. Ke, et al., “A Review of Machine Learning Methods for Imbalanced Data Challenges in Chemistry,” Chemical Science 16 (2025): 7637–7658, https://doi.org/10.1039/D5SC00270B.

[47] D. Bahdanau, K. Cho, and Y. Bengio, “Neural Machine Translation by Jointly Learning to Align and Translate,” in International Conference on Learning Representations (ICLR) (2015).

[48] R. Caruana, “Multitask Learning,” Machine Learning 28, no. 1 (1997): 41–75, https://doi.org/10.1023/a:1007379606734.

[49] V. Dey and X. Ning, “Enhancing Molecular Property Prediction With Auxiliary Learning and Task-Specific Adaptation,” Journal of Cheminformatics 16 (2024): 85, https://doi.org/10.1186/s13321-024-00880-7.

[50] S. M. Sze and K. K. Ng, Physics of Semiconductor Devices, 3rd ed. (Wiley, 2006): ISBN 9780471143239.

[51] H. Liu, M. Weng, Y. Huang, Y. Luo, and Y. Zheng, “Self Supervised Ensemble Learning Models for 2d Semiconductors Bandgap Prediction,” Computational Materials Science 255 (2025): 113853, https://doi.org/10.1016/j.commatsci.2025.113853.

[52] H. Sibi, J. Biju, and C. Chowdhury, “Advancing 2d Material Predictions: Superior Work Function Estimation With Atomistic Line Graph Neural Networks,” RSC Advances 14, no. 51 (2024): 38070–38078, https://doi.org/10.1039/d4ra07703b.

[53] A. Vaswani, N. Shazeer, N. Parmar, et al., “Attention Is All You Need,” in Proceedings of the 31st International Conference on Neural Information Processing Systems (NeurIPS) (2017), 6000–6010, https://dl.acm.org/doi/10.5555/3295222.3295349.

[54] J. Gasteiger, J. Groß, and S. Günnemann, “Directional Message Passing for Molecular Graphs,” in 8th International Conference on Learning Representations (ICLR) (2020), 10758.

[55] Y. Zhang and Q. Yang, “A Survey on Multi-Task Learning,” IEEE Transactions on Knowledge and Data Engineering 34 (2022): 5586–5609, https://doi.org/10.1109/TKDE.2021.3070203.

[56] S. P. Ong, W. D. Richards, A. Jain, et al., “Python Materials Genomics (Pymatgen): A Robust, Open-Source Python Library for Materials Analysis,” Computational Materials Science 68 (2013): 314–319, https://doi.org/10.1016/j.commatsci.2012.10.028.

[57] L. Pauling, “The Nature of the Chemical Bond. IV. the Energy of Single Bonds and the Relative Electronegativity of Atoms,” Journal of the American Chemical Society 54, no. 9 (1932): 3570–3582, https://doi.org/10.1021/ja01348a011.

[58] A. L. Allred and E. G. Rochow, “A Scale of Electronegativity Based on Electrostatic Force,” Journal of Inorganic and Nuclear Chemistry 5, no. 4 (1958): 264–268, https://doi.org/10.1016/0022-1902(58)80003-2.

[59] A. Kahn, “Fermi Level, Work Function and Vacuum Level,” Materials Horizons 3, no. 1 (2016): 7–10, https://doi.org/10.1039/c5mh00160a.

[60] A. Dunn, Q. Wang, A. Ganose, D. Dopp, and A. Jain, “Benchmarking Materials Property Prediction Methods: The Matbench Test Set and Automatminer Reference Model,” npj Computational Materials 6, no. 1 (2020): 138, https://doi.org/10.1038/s41524-020-00406-3.

[61] H. B. Michaelson, “The Work Function of the Elements and Its Periodicity,” Journal of Applied Physics 48, no. 11 (1977): 4729–4733, https://doi.org/10.1063/1.323539.

[62] J. Zaanen, G. A. Sawatzky, and J. W. Allen, “Band Gaps and Electronic Structure of Transition-Metal Compounds,” Physical Review Letters 55, no. 4 (1985): 418–421, https://doi.org/10.1103/physrevlett.55.418.

[63] S. Batzner, A. Musaelian, L. Sun, et al., “E(3)-Equivariant Graph Neural Networks for Data-Efficient and Accurate Interatomic Potentials,” Nature Communications 13, no. 1 (2022): 2453, https://doi.org/10.1038/s41467-022-29939-5.

Rights & permissions

2026 The Author(s). Materials Genome Engineering Advances published by Wiley-VCH GmbH on behalf of University of Science and Technology Beijing.