May 18 – 22, 2026
Virginia Tech
America/New_York timezone

Exploiting Kronecker Structure and Krylov Subspaces for Scalable Second-Order Optimization in Federated Physics-Informed Neural Networks

May 18, 2026, 5:00 PM
25m
Goodwin Hall 155 (Virginia Tech)

Goodwin Hall 155

Virginia Tech

Minisymposium Talk Application-Driven Family of Matrix Computations: Factorization, Inverse, Linear Solve Application-Driven Family of Matrix Computations: Factorization, Inverse, Linear Solve

Speaker

Amit Upadhyay (Indian Institute Of Technology (Indian School Of Mines) Dhanbad)

Description

Solving partial differential equations (PDEs) using distributed Physics-Informed Neural Networks (PINNs) introduces major computational challenges associated with high-dimensional curvature estimation, ill-conditioned optimization landscapes, and communication overhead in federated environments. In this work, we exploit Kronecker structure and Krylov subspace methods to develop a scalable second-order optimization framework for federated PINNs applied to the 2D Poisson equation. The proposed method combines Kronecker-Factored Approximate Curvature (K-FAC) with Conjugate Gradient (CG) iterations within a matrix-free Gauss–Newton formulation, avoiding explicit Hessian or Jacobian construction through automatic differentiation-based Jacobian-vector products. The symmetric positive definite structure induced by discretized Poisson operators supports efficient CG convergence while enabling scalable curvature-aware optimization. Within the federated learning framework, each client independently solves a local Gauss–Newton system using K-FAC approximations and Krylov iterations, transmitting only local parameter corrections to the central server. Experimental results across grid resolutions from $32\times32$ to $128\times128$ demonstrate up to $57.8\%$ relative improvement in test accuracy and $44.4\%$ reduction in required communication rounds compared with first-order Adam-based optimization. These findings illustrate how structured linear algebra techniques, particularly Kronecker factorization and Krylov subspace solvers, provide an efficient foundation for communication-efficient scientific machine learning under distributed computing constraints.

Authors

Amit Upadhyay (Indian Institute Of Technology (Indian School Of Mines) Dhanbad) Mr Dharavath Ramesh (Indian Institute Of Technology (Indian School Of Mines) Dhanbad)

Presentation materials

There are no materials yet.