Neural Representations for Computational Physics:

from Reduced Order Models to Scalable Transformers

Vedant Puri
AUG 19, 2026

Committee:

Levent Burak Kara, Yongjie Jessica Zhang, Amir Barati Farimani, Krishna Garikipati

Computer simulations are critical for industrial applications

Modern engineering is reliant on computer simulations

Predictive maintenance

Design space exploration

[2]

[1]

[3]

[1] CFD Direct / OpenFOAM – “OpenFOAM HPC on AWS with EFA”, cfd.direct
[2] EurekAlert — “New concrete system may reduce wind-turbine costs”
[3] Flow-3D, “FLOW-3D AM” product page, flow3d.com

Process optimization

Automotive engineering

Civil engineering

Advanced manufacturing

[1]

Physics-based simulations bottleneck many engineering workflows

[1] GridPro Blog — “The Art and Science of Meshing Airfoil”

[2] ResearchGate — “Transition to turbulence of Taylor-Green Vortex at different time (DNS)” (figure)
[3] ORNL / U.S. Department of Energy — “DOE and Cray deliver record-setting Frontier supercomputer at ORNL”

Numerical simulations are expensive because

[1]

Complex geometry \(\implies\) many DoFs

Multiscale physics \(\implies\) small \(\Delta t\)

[2]

Repeated large system solves

[3]

Neural signal representations learn to emulate physics from data

{\mathbf{u}}(\mathbf{x}) = \Sigma_{i=1}^{N} \mathbf{u}_i \phi_i(\mathbf{x})

(Explicit) Weighted sum of polynomial interpolants

Finite Elements

[1]

{\mathbf{u}}(\mathbf{x}) = (Z_L \circ \dotsc \circ Z_0)(\mathbf{x})

(Implicit) High-dim nonlinear feature learners

Neural Networks

[2]

Cannot learn from data

Can learn from data

Large cost per simulation

Cheap evaluation after training

High-accuracy

Problem-specific

Robust

Up to \(0.1\%\) accuracy

[1] Math StackExchange — “Interpolation in Finite Element Method”
[2] ResearchGate — “Structure of a Deep Neural Network” (figure)

\(\text{Mesh ansatz}\)

\({u}(x)=\)

\(u(x) = \)

\(\text{Neural ansatz}\)

\(\text{Physics-based}\)

\(\text{Data-driven}\)

\(\text{Numerical simulation}\)

\(\text{Reduced Order Modeling}\)

\(\text{Neural ROMs}\)

\(\text{Surrogate learning}\)

\(\text{Transformers}\)

\(\text{PINNs}\)

\(\text{Finite Elements}\)

\(\text{PCA/POD}\)

\(\text{Graph Networks}\)

Landscape of data-driven methods in computational physics

Accelerate numerical solvers with neural ROMs

Efficient attention models for large-scale surrogate learning

[1]

[3]

[2]

[1] CFD Direct / OpenFOAM, “Introduction to Computational Fluid Dynamics”

[3] Kutz, J. N. “Data-Driven Modeling & Dynamical Systems” (UW)

[2] ResearchGate, “Schematic of a Vanilla Physics-Informed Neural Network”

Outline

Scalable low-rank attention for PDE surrogates

Neural Reduced Order Modeling

Low-rank attention with dynamic routing

Accelerate PDE solves with structure learned from data.

Replace simulation with solution operator learned from data.

Improved accuracy with smarter compression.

Neural Reduced Order Modeling

Accelerate simulations with structure learned from data.

\mathbb{R}^{N_\text{FOM}}

\bar{u}(0)

\tilde{u}(0)

\tilde{u}(T)

\mathcal{M}

\bar{u}(T)

h_\text{ROM}

g_\text{ROM}

\begin{pmatrix} \hspace{0.4em} \\ \\ \\ \end{pmatrix}

Background on model order reduction

Learn low-dimensional solution representation with data + evolve with physics.

High-dimensional simulation data

Collect and compress data

Evolve ODE on low-dim manifold

\bar{u}(0)

\tilde{u}(0)

\tilde{u}(T)

\mathcal{M}

\bar{u}(T)

\begin{pmatrix} \hspace{0.4em} \\ \\ \\ \end{pmatrix}

\tilde{u}_1

\bar{u}_1

\bar{u}_T

\cdots

\begin{bmatrix} \hspace{6.5em} \\ \\ \\ \\ \\ \end{bmatrix}

\cdots

\tilde{u}_T

Cost savings from solving smaller ODE system.

\textcolor{red}{\bar{u}} \approx g(\textcolor{blue}{\tilde{u}}) ,\hspace{1em} \text{dim}(\textcolor{red}{\bar{u}}) \gg \text{dim}(\textcolor{blue}{\tilde{u}})

\partial_t \textcolor{red}{\bar{u}} + \textcolor{red}{\bar{u}} \partial_x\textcolor{red}{\bar{u}} = \nu\partial_{xx} \textcolor{red}{\bar{u}}

\frac{\partial g}{\partial \textcolor{blue}{\tilde{u}}} \partial_t \textcolor{blue}{\tilde{u}} + g(\textcolor{blue}{\tilde{u}}) \partial_x g(\textcolor{blue}{\tilde{u}}) = \nu\partial_{xx} g(\textcolor{blue}{\tilde{u}})

\partial_t u + u \partial_x u = \nu \partial_{xx}u

We accelerate traveling shock problems with SNF-ROM

2D Viscous Burgers problem \( (\mathit{Re} = 1\text{k})\)

Smooth neural field ROM (SNF-ROM)

Fast inference time - \(199\times\) speed-up
Consistently high accuracy
Robust & non-intrusive online evaluation

\(\text{Relative error: }0.37\%\)

\(\text{DoFs: }524~k \to 2\)

\(\text{Wall-time: }13.4~\text{s} \to 0.068~\text{s}\)

Large deviations!

Learning smooth latent space trajectories

\(\text{Autoencoder ROM}\)

\text{Projection}

\text{Online solve}

Evolution of ROM states

\text{FOM}

\text{POD-ROM}

\text{SNFL-ROM}

\text{CAE-ROM}

\text{SNFW-ROM}

SNF-ROM ensures accurate online dynamics evaluation.

\text{Projection}

\text{Online Solve}

We accelerate traveling shock problems with SNF-ROM

2D Viscous Burgers problem \( (\mathit{Re} = 1\text{k})\)

Smooth neural field ROM (SNF-ROM)

Fast inference time - \(199\times\) speed-up
Consistently high accuracy
Robust & non-intrusive online evaluation

\(\text{Relative error: }0.37\%\)

\(\text{DoFs: }524~k \to 2\)

\(\text{Wall-time: }13.4~\text{s} \to 0.068~\text{s}\)

High freq. noise

Non-differentiable!

Accurately capture of dynamics with smooth neural fields

\textcolor{red}\times

\mathrm{NN}(x) \approx u^{}(x) \implies \dfrac{\mathrm{d}^k}{\mathrm{d}x^k} \mathrm{NN}(x) \approx u^{(k)}(x)

\mathrm{NN}(x)

\frac{\mathrm{d}}{\mathrm{d}x} \mathrm{NN}(x)

\frac{\mathrm{d}^2}{\mathrm{d}x^2} \mathrm{NN}(x)

Large deviations!

Learning smooth latent space trajectories

\(\text{Autoencoder ROM}\)

\(\text{SNF-ROM}\)

\text{Projection}

\text{Online solve}

\text{Projection}

\text{Online solve}

Evolution of ROM states

No deviation

\text{FOM}

\text{POD-ROM}

\text{SNFL-ROM}

\text{CAE-ROM}

\text{SNFW-ROM}

SNF-ROM ensures accurate online dynamics evaluation.

Accurate capture of dynamics

Our contributions

Accurate latent space traversal

Accurate dynamics evaluation with neural ansatz

We demonstrate numerical simulation with neural ansatz

Won poster award at World Conf. Comp. Mech. 2024

Published in Journal of Comp. Phys.

Scalable low-rank attention for PDE surrogates

Surrogate models learn PDE solution operator from data

\mathcal{L}(\boldsymbol{x}, \boldsymbol{u}; \boldsymbol{\mu}) = 0

\mathcal{G}: \boldsymbol{\mu} \mapsto \boldsymbol{u}

\mathcal{G}_\theta \approx \mathcal{G}

Training

\mathcal{G}_\theta \approx \mathcal{G}

Inference

\mathcal{G}_\theta

Large training cost is amortized over several evaluations

Model learns to predict \(\boldsymbol{u}\) over a distribution of \(\boldsymbol{\mu}\)

What signal representations are optimal for surrogate learning?

\begin{bmatrix} &&&\\ &&&\\ &&& \end{bmatrix} \cdot \begin{bmatrix} \\ \underline{u}\\ \\ \end{bmatrix} = \begin{bmatrix} \\ \underline{f}\\ \\ \end{bmatrix}

\begin{bmatrix} \\ \underline{u}\\ \\ \end{bmatrix} = \begin{bmatrix} &&&\\ &&&\\ &&&\\ \end{bmatrix} \cdot \begin{bmatrix} \\ \underline{f}\\ \\ \end{bmatrix}

Solution operator: dense, global

Forward operator: sparse, structured, local

Every point depends on every other

u = \int_\Omega \mathcal{K}(x,y)f(y)dy

Transformers enable global communication

\Delta u = f

Message-passing on a dynamic all-to-all graph

Quadratic (\(\mathcal{O}(N^2)\)) cost limits scalability

[1] Vaswani et al. — “Attention Is All You Need”, NeurIPS 2017

Quadratic \((\mathcal{O}(N^2))\) cost of attention limits scalability

Infeasible for large problems

[1] Vaswani et al. — “Attention Is All You Need”, NeurIPS 2017

Over \(20~\text{s}\) per gradient step on 1M points!

Attention determines graph weights through dot-product scores \(S_{m,n}=q_m^T \cdot k_n\).

Every \(\textcolor{blue}{\text{output token}}\) produces a probability distribution over all \(\textcolor{green}{\text{inputs}}\) .

Are \(N \times N\) messages really necessary?

Smoothness implies redundancy in communication

\(\text{Fully-resolved flow-field}\)

Are \(N \times N\) messages really necessary?

Smoothness implies redundancy in communication

Method: club matching points to one cluster and communicate together

\(\text{Fully-resolved flow-field}\)

Our approach: Identify matching tokens with dot-product similarity against learned latents

Encoding: introduce \(M\) latent clusters to pool messages from matching tokens

\(M\) learned queries

Decoding: map pooled messages to matching output tokens

Our model: FLARE (Fast Low-rank Attention Routing Engine)

Encoding/decoding through dot-product scores against learned latents.

Score matrix never truly materialized.

FLARE is an encode/decode

latent-space attention mechanism.

Latent-space-attention can effect factorized low-rank self-attention.

Background on latent-space attention methods

Simplify and connect latent-attention to low-rank self-attention

Isolate gather/scatter routing as a load-bearing primitive

Multihead gather-scatter as parallel low-rank pathways

Each attention head learns independent latent queries

FLARE scales to extremely large problems on a single GPU

\(>200\times\) speedup

Extreme scalability with fused kernel implementations

Excellent memory efficency

Low-rank attention scales favorably on PDE surrogate problems

Darcy flow problem

Elasticity problem

Accuracy improves with depth

Accuracy improves with \(M\) and depth!

Low-rank attention is strong across PDE benchmarks

Relative \(L_2\) error \( (\times 10^{-3})\) (lower is better)

Pipe

Darcy

Elasticity

LPBF

DrivAerML

Airfoil

Competitive with full self-attention

Significant gains over Transolver

\(\implies\) PDE problems ok with low-rank comm.

Low-rank attention generalizes beyond PDE tasks

\texttt{INPUT:\, [MAX 4 3 [MIN 2 3 ] 1 0 [MEDIAN 1 5 8 9, 2]] \,OUTPUT: 5}

\texttt{INPUT:\, [MAX 7 [MEDIAN 1 2 3 ] [MAX 9 2 2] [MIN 2 8]] \,OUTPUT: 9}

\(\text{Accuracy (\%)}\)

FLARE is competitive with full self-attention

Our contributions

[2] Akhare et al. “High-Fidelity Industrial- Crash Dynamics”, arXiv:2605:27758

[1] Adams et al. “GeoTransolver”, arXiv:2512:20399

[3] Curtosi, et al. "Crash Assessment via Physics-Aware Attention", arxiv:2605:11784

Expose communication-rank as an axis of scaling

Robust and highly scalable low-rank attention model

Manuscript under review at TMLR

[1]

GeoFLARE: FLARE backbone + multiscale geometry feature extractor

[2]

NVIDIA

Norte Dame, GM

SEAT, IDIADA (Spain)

[3]

GeoFLARE on crash-dynamics problem

FLARE has been independently extended to complex geometries.

[2]

Low-rank attention with dynamic routing

Low-rank attention in FLARE is constrained to static routing

fixed, static queries.

Encode/decode weights are obtained by outer-product of input-projection with

Hypothesis: input-dependent queries can improve accuracy

\(S_{m,n}(X) = q_m^T \cdot k_n(X)\)

Low-rank attention in FLARE is constrained to static routing

How to produce \(M\) latents from \(N\) inputs?

fixed, static queries.

Hypothesis: input-dependent queries can improve accuracy

Encode/decode weights are obtained by outer-product of input-projection with

\(S_{m,n}(X) = q_m^T \cdot k_n(X)\)

Low-rank attention with dynamic token routing

input-dependent queries.

Dynamic query synthesis

via FLARE encoder

The residual stream is intact.

All gains from improved routing.

Hypothesis: input-dependent queries can improve accuracy

Encode/decode weights are obtained by outer-product of input-projection with

\(S_{m,n}(X) = q_m^T(X) \cdot k_n(X)\)

Dynamic routing improves accuracy on PDE benchmarks

Relative \(L_2\) error \( (\%)\) (lower is better)

Pipe

Darcy

Elasticity

DrivAerML

Airfoil

Substantial gains from dynamic routing

Dominant over Transolver family

Dynamic routing improves accuracy on Elasticity

Improved routing substitutes for depth

Dynamic routing improves accuracy on Darcy flow

Improved routing substitutes for depth

Continued improvement with \(M\)

FLARE++ scales extremely well on large problems

\(\approx1.2-1.3\times\) slower than FLARE

Inherits extreme scalability from FLARE

\(\approx1.3\times\) more params than FLARE

NASA-CRM: Large-scale external aerodynamics problem

FLARE family is dominant on this problem.

DrivAerML industrial scale benchmark (9M points)

Faster convergence with improved routing

\(\approx30\%\) improvement at \(B=2\)

Continued improvement even at \(B=20\)

Dynamic routing improves low-rank attention on LRA

\texttt{INPUT:\, [MAX 4 3 [MIN 2 3 ] 1 0 [MEDIAN 1 5 8 9, 2]] \,OUTPUT: 5}

\texttt{INPUT:\, [MAX 7 [MEDIAN 1 2 3 ] [MAX 9 2 2] [MIN 2 8]] \,OUTPUT: 9}

\(\text{Accuracy (\%)}\)

FLARE++ has highest average accuracy!

Our contributions

Future work

Substantial accuracy gains with shorter models

Low-rank attention with dynamic routing

Preprint out on arXiv. Will submit to ICLR

Extend to complex geometries: GeoFLARE++

Consider time-series PDE problems

Summary of this dissertation

Neural Reduced Order Modeling

Demonstrated numerical simulation with neural ansatz

Make attention-based PDE surrogate modeling feasible on large scale meshes

Improve accuracy with smarter compression

Scalable low-rank attention for PDE surrogates

Low-rank attention with dynamic routing

Publications

Puri, Vedant, Yongjie Jessica Zhang, and Levent Burak Kara. "FLARE++: Low-rank attention with dynamic token routing." arXiv preprint arXiv:2608.11519 (2026).
Puri, Vedant, Aditya Joglekar, Kevin Ferguson, Yu-hsuan Chen, Yongjie Jessica Zhang, and Levent Burak Kara. "FLARE: Fast Low-rank Attention Routing Engine." arXiv preprint arXiv:2508.12594 (2025).
Puri, Vedant, Aviral Prakash, Levent Burak Kara, and Yongjie Jessica Zhang. "SNF-ROM: Projection-based nonlinear reduced order modeling with smooth neural fields." Journal of Computational Physics 532 (2025): 113957.
Shankar, Varun, Vedant Puri, Ramesh Balakrishnan, Romit Maulik, and Venkatasubramanian Viswanathan. "Differentiable physics-enabled closure modeling for Burgers’ turbulence." Machine Learning: Science and Technology 4, no. 1 (2023): 015017.

Acknowledgements

Thank you to CBML, VDEL for support and guidance over the years.

Committee Members

Dr. Levent Burak Kara

Dr. Yonjie Jessica Zhang

Dr. Amir Barati Farimani

Dr. Krishna Garikipati

Dr. Varun Shankar

Dr. Ramesh Balakrishnan

Dr. Romit Maulik

Dr. Venkatsubramanium Viswanathan

Dr. Aviral Prakash

Dr. Aditya Joglekar

Dr. Kevin Ferguson

Dr. Y-hsuan Chen

Sri Datta Ganesh Bandreddi

Collaborators

Thank you

Questions?

ROMs

Manifold learning

\begin{pmatrix} \hspace{0.8em} \\ \\ \\ \\ \end{pmatrix}

\mathbb{R}^{N_\text{FOM}}

\bar{u}(0)

\bar{u}(T)

\frac{\partial \boldsymbol{u}}{\partial t} = \mathcal{L}(\boldsymbol{x}, t, \boldsymbol{u}; \boldsymbol{\mu})

Full order model (FOM)

\boldsymbol{u}(\boldsymbol{x}, t; \boldsymbol{\mu}) \approx g_\text{FOM}(\boldsymbol{x}, \textcolor{red}{\bar{u}(t; \boldsymbol{\mu}})) = \mathbf{\Phi}(\boldsymbol{x}) \cdot \textcolor{red}{\bar{u}(t; \boldsymbol{\mu}})

Manifold learning

\textcolor{orange}{N_\text{Lin-ROM}} << \textcolor{red}{N_\text{FOM}}

\begin{pmatrix} \hspace{0.8em} \\ \\ \\ \\ \end{pmatrix}

\mathbb{R}^{N_\text{FOM}}

\bar{u}(0)

\bar{u}(T)

\frac{\partial \boldsymbol{u}}{\partial t} = \mathcal{L}(\boldsymbol{x}, t, \boldsymbol{u}; \boldsymbol{\mu})

Full order model (FOM)

Linear POD-ROM

\textcolor{red}{\bar{u}(t; \boldsymbol{\mu})} \approx g'_\text{ROM}(\textcolor{orange}{\bar{u}(t; \boldsymbol{\mu})}) = \bar{u}_0 + \mathbf{P} \cdot \textcolor{orange}{\tilde{u}(t; \mathbf{\mu})}

Large Kolmogorov \(N\)-width

in advection problems

Manifold learning

\frac{\partial \boldsymbol{u}}{\partial t} = \mathcal{L}(\boldsymbol{x}, t, \boldsymbol{u}; \boldsymbol{\mu})

Full order model (FOM)

Linear POD-ROM

Nonlinear ROM

\textcolor{red}{\bar{u}(t; \boldsymbol{\mu})} \approx g'_\text{ROM}(\textcolor{orange}{\bar{u}(t; \boldsymbol{\mu})}) = \bar{u}_0 + \mathbf{P} \cdot \textcolor{orange}{\tilde{u}(t; \mathbf{\mu})}

\boldsymbol{u}(\boldsymbol{x}, t; \boldsymbol{\mu}) \approx g_\text{ROM}(\boldsymbol{x}, \textcolor{blue}{\tilde{u}(t; \boldsymbol{\mu}})) = \mathrm{NN}_\theta\left(\boldsymbol{x}, \textcolor{blue}{\tilde{u}(t; \boldsymbol{\mu}}) \right)

\textcolor{orange}{N_\text{Lin-ROM}} << \textcolor{red}{N_\text{FOM}}

\textcolor{blue}{N_\text{Nl-ROM}} \leq\,\,

\begin{pmatrix} \hspace{0.8em} \\ \\ \\ \\ \end{pmatrix}

\mathbb{R}^{N_\text{FOM}}

\bar{u}(0)

\bar{u}(T)

\({u}(x)=\)

Large Kolmogorov \(N\)-width

in advection problems

Background on model order reduction

Learn low-dimensional solution representation with data + evolve with physics.

High-dimensional simulation data

Collect and compress data

Evolve ODE on low-dim manifold

\bar{u}(0)

\tilde{u}(0)

\tilde{u}(T)

\mathcal{M}

\bar{u}(T)

\begin{pmatrix} \hspace{0.4em} \\ \\ \\ \end{pmatrix}

\tilde{u}_1

\bar{u}_1

\bar{u}_T

\cdots

\begin{bmatrix} \hspace{6.5em} \\ \\ \\ \\ \\ \end{bmatrix}

\cdots

\tilde{u}_T

Cost savings from solving smaller ODE system.

\textcolor{red}{\bar{u}} \approx g(\textcolor{blue}{\tilde{u}}) ,\hspace{1em} \text{dim}(\textcolor{red}{\bar{u}}) \gg \text{dim}(\textcolor{blue}{\tilde{u}})

\partial_t \textcolor{red}{\bar{u}} + \textcolor{red}{\bar{u}} \partial_x\textcolor{red}{\bar{u}} = \nu\partial_{xx} \textcolor{red}{\bar{u}}

\partial_t u + u \partial_x u = \nu \partial_{xx}u

Large Kolmogorov \(N\)-width in advection problems

Manifold learning via autoenoders cause deviations in physics solve

[1] Lee & Carlberg — Nonlinear manifold ROM via CNN autoencoders (JCP 2020)

\frac{\partial \textcolor{blue}{\tilde{u}} }{\partial t} = \tilde{\mathcal{L}}(t, \textcolor{blue}{\tilde{u}})

\tilde{u}(t)

\begin{pmatrix} \hspace{0.9em} \\ \\ \\ \\ \end{pmatrix}

\begin{pmatrix} \textcolor{blue}*\\ \textcolor{blue}*\\ \end{pmatrix}

\begin{pmatrix} \hspace{0.9em} \\ \\ \\ \\ \end{pmatrix}

\bar{u}(t)

\(\text{Encoder}\)

\(\text{Decoder}\)

\text{Projection}

\text{Inference}

\text{Projection}

\text{Online Solve}

Deviation in latent physics solve

\text{CAE-ROM} \,[1]

\text{FOM (ground truth)}

\text{Encoder projection}

\text{Physics solve}

\text{distribution of }\tilde{u}

2D Burgers \(\mathit{Re}=1\mathit{k}\)

\frac{\partial \textcolor{red}{\bar{u}} }{\partial t} = \bar{\mathcal{L}}(t, \textcolor{red}{\bar{u}})

Grid-dependent inference-cost

Not amenable to hyper-reduction

Smooth Neural Field ROMs (SNF-ROM)

Change the solution representation to coordinate neural fields.

\(\text{PDE Problem}\)

\((\boldsymbol{x}, t, \boldsymbol{\mu})\)

\(\text{Coordinates}\)

\(\text{Neural field MLP }(g_\theta)\)

\(\boldsymbol{x}\)

\(\boldsymbol{u}\left( \boldsymbol{x}, t; \boldsymbol{\mu} \right)\)

\(\text{ Parameters}\)

\( \text{and time}\)

\(\tilde{u}\)

Grid-independent

Supports \(\partial_x\) through autodiff

Smooth Neural Field ROMs (SNF-ROM)

\(\text{Loss } (L)\)

\(\text{Backpropagation}\)

\(\nabla_\theta L\)

\(\nabla_\varrho L\)

\(\nabla_\theta L\)

Change the solution representation to coordinate neural fields.

\(\tilde{u}\)

\(\text{PDE Problem}\)

\((\boldsymbol{x}, t, \boldsymbol{\mu})\)

\(\text{ Parameters}\)

\( \text{and time}\)

\(\text{ Intrinsic ROM manifold}\)

\tilde{\mathcal{U}} = \left\{ \tilde{u}(t; \mathbf{\mu}) |~ t,\, \mathbf{\mu} \right\}

\tilde{u}(t; \mathbf{\mu})

\(\text{Coordinates}\)

\(\text{Neural field MLP }(g_\theta)\)

\(\boldsymbol{x}\)

\(\boldsymbol{u}\left( \boldsymbol{x}, t; \boldsymbol{\mu} \right)\)

Grid-independent

Supports \(\partial_x\) through autodiff

Direct control over the complexity of latent trajecories

Experiment: 1D Viscous Burgers problem \( (\mathit{Re} = 10~{k})\)

\frac{\partial {u}}{\partial t} + {u} \frac{\partial {u}}{\partial x} = \nu \frac{\partial^2 u}{\partial x^2}, \qquad u(t=0) = u_0(\mu)

\(\text{CAE-ROM}\) [1]

\(\text{SNFL-ROM (ours)}\)

\(\text{SNFW-ROM (ours)}\)

Online dynamics solve matches learned trajectories

Online evaluation deviates!

Distribution of reduced states \((\tilde{u})\)

[1] Lee & Carlberg — Nonlinear manifold ROM via CNN autoencoders (JCP 2020)

Experiment: 2D Viscous Burgers problem \( (\mathit{Re} = 1~{k})\)

\frac{\partial \boldsymbol{u}}{\partial t} + \boldsymbol{u} \cdot \boldsymbol{\nabla}\boldsymbol{u} = \nu \Delta \boldsymbol{u}

\(\text{CAE-ROM}\) [1]

\(\text{SNFL-ROM (ours)}\)

\(\text{SNFW-ROM (ours)}\)

Relative error

[1] Lee & Carlberg — Nonlinear manifold ROM via CNN autoencoders (JCP 2020)

\([1]\)

\(\text{DoFs: }500\text{k} \to 2\)

\(\text{Time }(t)\)

\(\text{Relative Error}(t)\)

\(199\times\) speed-up

\text{FOM}

\text{POD-ROM}

\text{SNFL-ROM (ours)}

\text{CAE-ROM}

\text{SNFW-ROM (ours)}

Smooth Neural Field ROM directly controls latent trajectories

\varrho, \, \theta = \argmin_{\varrho, \, \theta}\left\{ \sum_{\boldsymbol{x}, \, t, \, \boldsymbol{\mu}} || \boldsymbol{u}(\boldsymbol{x}, t; \boldsymbol{\mu}) - g_\theta(\boldsymbol{x}, \Xi_\varrho(t; \boldsymbol{\mu})) ||_2^2 \right\}

Supervised learning problem jointly learns latent trajectories and data manifold.

\(\text{Loss } (L)\)

\(\text{Backpropagation}\)

\(\nabla_\theta L\)

\(\nabla_\varrho L\)

\(\nabla_\theta L\)

\(\text{PDE Problem}\)

\((\boldsymbol{x}, t, \boldsymbol{\mu})\)

\(\text{ Parameters}\)

\( \text{and time}\)

\(\text{ Intrinsic ROM manifold}\)

\tilde{\mathcal{U}} = \left\{ \tilde{u}(t; \mathbf{\mu}) |~ t,\, \mathbf{\mu} \right\}

\tilde{u}(t; \mathbf{\mu})

\(\text{Coordinates}\)

\(\text{Smooth neural field MLP }(g_\theta)\)

\(\tilde{u}\)

\(\boldsymbol{x}\)

\(\boldsymbol{u}\left( \boldsymbol{x}, t; \boldsymbol{\mu} \right)\)

Force \( t \mapsto \tilde{u}(t) \) to be simple, e.g., shallow MLP.

Coordinate MLPs with sinusoidal activations offers grid-independence.

Replace autoencoder with a direct prediction workflow.

Accurate derivatives require smooth neural field representations

\mathrm{NN}(x) \approx u^{}(x)

\textcolor{red}\times

\dfrac{\mathrm{d}^k}{\mathrm{d}x^k} \mathrm{NN}(x) \approx \dfrac{\mathrm{d}^k}{\mathrm{d}x^k} u(x)

SNF-ROM with Lipschitz regularization (SNFL-ROM)

\(\text{Penalize the \textcolor{blue}{Lipschitz constant} of the MLP [arXiv:2202.08345]}\)

\varrho, \, \theta = \argmin_{\varrho, \, \theta}\left\{ L_\text{data}(\varrho, \theta) + \textcolor{blue}{\alpha \bar{c}(\theta)} \right\}

\text{For MLP: } c_\theta \leq \textcolor{blue}{\bar{c}(\theta)} = \prod_{l=1}^L ||W_l||_p

||f(x_2) - f(x_1)||_p \leq \textcolor{blue}{c}||x_2 - x_1||_p

\text{change in output}

\text{change in input}

\text{For a single layer: } \textcolor{blue}{c_l} = ||W_l||_p

\(\text{[enwiki:1230354413]}\)

SNF-ROM with Weight regularization (SNFW-ROM)

\(\text{Directly penalize \textcolor{red}{high-frequency components} in }\dfrac{\text{d}}{\text{d} x}\text{NN}_\theta(x)\)

\frac{\text{d}}{\text{d} x} \mathrm{NN}_\theta(x) = \left( \prod_{l=2}^L W_l \cdot \text{diag}(\textcolor{red}{\sigma'(z_{l-1})}) \right) \cdot W_1

\textcolor{red}{ \cos\left( W_l z_{l-1} + b_l \right) }

\varrho, \, \theta = \argmin_{\varrho, \, \theta}\left\{ L_\text{data}(\varrho, \theta) + \textcolor{red}{ \frac{\gamma}{2} \sum_{l=1}^L \sum_{i,j} ||W_l^{ij}||_2^2 } \right\}

We present two approaches to learn inherently smooth and accurately differentiable neural field MLPs.

\({x}\)

\({u(x)}\)

\mathrm{NN}(x)

u(x)

\text{NN}

\frac{\mathrm{d}}{\mathrm{d}x} \mathrm{NN}(x)

\frac{\mathrm{d}^2}{\mathrm{d}x^2} \mathrm{NN}(x)

\textbf{SNFL (ours)}

\textbf{SNFW (ours)}

High freq. noise

Experiment: 1D Kuramoto-Sivashinsky problem

\frac{\partial {u}}{\partial t} + u\frac{\partial {u}}{\partial x} + \frac{\partial^2 {u}}{\partial x^2} + \nu\frac{\partial^4 {u}}{\partial x^4} = 0

SNF-ROM maintains high accuracy even with larger time-steps.

\(\text{Relative error vs time } (\Delta t = \Delta t_0)\)

\(\text{Relative error vs time } (\Delta t = 10\Delta t_0)\)

[1] Lee & Carlberg — Nonlinear manifold ROM via CNN autoencoders (JCP 2020)

Experiment: 2D Viscous Burgers problem \( (\mathit{Re} = 1~{k})\)

\frac{\partial \boldsymbol{u}}{\partial t} + \boldsymbol{u} \cdot \boldsymbol{\nabla}\boldsymbol{u} = \nu \Delta \boldsymbol{u}

\(\text{CAE-ROM}\) [1]

\(\text{SNFL-ROM (ours)}\)

\(\text{SNFW-ROM (ours)}\)

\(\text{Relative error }\)

[1] Lee & Carlberg — Nonlinear manifold ROM via CNN autoencoders (JCP 2020)

\([1]\)

FLARE

Quadratic \((\mathcal{O}(N^2))\) cost of attention limits scalability

Over \(20~\text{s}\) per gradient step on a mesh of 1m poins!

[1] Vaswani et al. — “Attention Is All You Need”, NeurIPS 2017

Situating FLARE in the space of latent attention methods

Multiple parallel low-rank pathways

Multi-head cross-attention encoder/decoder

Eliminate latent-space self-attention

Simplify and connect latent attention to low-rank

+Transolver has same nn.Linear(D,M) proj for all heads.

Isolate attention routing as the load-bearing pirmitive

Extreme scalability thanks to SDPA compatibility

Laser powder bed fusion benchmark problem

FLARE Spectral analysis

FLARE++

FLARE++ performs equally well in mixed-precision

Darcy

Elasticity

Airfoil

FLARE++ improves on FLARE!

Maintains accuracy in FP16

Competitive with full self-attention

FLARE++ External Aerodynamics

Aim 2: background on next-token prediction transformers [1]

y_t = \frac{\sum_{\tau = 1}^t\exp\left(q_t \cdot k_\tau \right) v_\tau}{\sum_{\tau = 1}^t \exp \left(q_t \cdot k_\tau \right)}

All previous key/value \(\{k_\tau, v_\tau \}_{\tau \leq t}\) must be cached on the GPU.

Major memory and latency bottleneck!

[1] Vaswani et al. — “Attention Is All You Need”, NeurIPS 2017

Training algorithm (causal masking)

Inference algorithm (recurrence relation)

Dot-products need to be recomputed for every \(q_t\).

\(\mathcal{O}(N^2)\) complexity.

Aim 2: Develop decoder version of FLARE

Linear time auto-regressive attention.

Fixed memory footprint (only store \(\mathcal{O}(M)\) cache).

Flexible latent capacity.

Advantages

Required components

Fused GPU kernels for training and inference.

Bespoke training algorithm for causal FLARE.

Extensive benchmarking and evaluation.

Z_t = \text{online\_softmax}(Z_{t-1}, k_t, v_t)\\ y_t = \text{softmax}(Q^T \cdot k_t)^T\cdot Z_t

Inference algorithm (recurrence rule)

Next-token prediction with FLARE

FLARE decoder recurrence

https://github.com/vpuri3/flash-linear-attention-pvt/blob/run/legacy/training/STORYTELLING.md

TopoFLARE

Point cloud attention fails on complex geometries

Topology-conditioned routing

GeoFLARE

FLARE has been independently extended to problems with complex geometries

[2] Akhare et al. “High-Fidelity Industrial- Crash Dynamics”, arXiv:2605:27758

[1] Adams et al. “GeoTransolver”, arXiv:2512:20399

[3] Curtosi, et al. "Crash Assessment via Physics-Aware Attention", arxiv:2605:11784

[1]

[2]

FLARE backbone + multiscale geometry feature extractor

[2]

NVIDIA

Norte Dame, GM

SEAT, IDIADA (Spain)

[3]

GeoFLARE architecture

GeoFLARE is used for crash-dynamics calculations

[2] Akhare et al. “High-Fidelity Industrial- Crash Dynamics”, arXiv:2605:27758

[1] Adams et al. “GeoTransolver”, arXiv:2512:20399

[3] GitHub:NVIDIA/physicsnemo/examples/structural_mechanics/crash

Crash-dynamics on bumper-beam problem

[2] Akhare et al. “High-Fidelity Industrial- Crash Dynamics”, arXiv:2605:27758

[1] Adams et al. “GeoTransolver”, arXiv:2512:20399

[3] GitHub:NVIDIA/physicsnemo/examples/structural_mechanics/crash

Vedant Puri defense

By Vedant Puri

Vedant Puri defense

Vedant Puri's thesis defense at Carnegie Mellon University

Vedant Puri

PhD student at Carnegie Mellon University

vpuri3.github.io