Tucker Compression for Scalable Operator Learning in Large-Scale Parametric PDE Models

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.