Tucker Compression for Scalable Operator Learning in Large-Scale Parametric PDE Models
Title | Tucker Compression for Scalable Operator Learning in Large-Scale Parametric PDE Models |
Publication Type | Presentation |
Year of Publication | 2024 |
Authors | Richard Rex, Avasarala, S, Grady, T, Felix J. Herrmann |
Keywords | deep learning, FNO, hierarchical tucker tensor, Inverse problems, kronecker product, ML4SEISMIC, SLIM, two-phase flow |
Abstract | Simulating two-phase flow via PDEs is computationally expensive due to the inversion of large, ill-conditioned matrices. To accelerate these computations, we reformulate Hierarchical Tucker Tensor (HTT) decompositions into Kronecker products, enabling scalable Fourier Neural Operators (FNOs) for CO2 saturation predictions in subsurface environments. This reformulation allows efficient scaling across multiple GPUs while maintaining a large number of modes. Building on our existing matrix-free abstraction library, we extend its capabilities to support distributed tensor operators. The extended library is auto-differentiable, with customized AD rules for training complex networks. We demonstrate the performance and scalability of our approach by evaluating FNO simulations against traditional PDE solvers for predicting time-varying CO2 saturations from permeability models in large-scale subsurface environments. |
URL | https://slim.gatech.edu/Publications/Public/Conferences/ML4SEISMIC/2024/rex2024ML4SEISMIChtc |
Citation Key | rex2024ML4SEISMIChtc |