Partial-Moment PINNs for Caldeira--Leggett Parameter Learning in Quantum Brownian Motion

Partial-Moment PINNs for Caldeira–Leggett Parameter Learning in Quantum Brownian Motion Krishna Bhatia

arXiv:2608.23093v1 [quant-ph] 24 Aug 2026

QuantumAI Lab, Fractal Analytics, Mumbai, India krishna.bhatia@fractal.ai

Abstract—We study parameter recovery in the Caldeira– Leggett (quantum Brownian) oscillator from partial moment traces. Our model is a moment-level PINN that predicts the five first/second moments and enforces the linear CL/HPZ ODEs by automatic differentiation. Physical structure is imposed through a PSD (Cholesky) covariance head, high-temperature CL assumptions with Dxp ≈ 0, and fluctuation–dissipation ties between Dpp and γ. On synthetic CL data with channels {µx , σxx , σxp }, the constrained variant recovers (ω, γ) accurately, stabilizes Dpp , and achieves low rollout error compared to finite differences and Kalman–EM (expectation–maximization) with exact Van Loan discretization. Fisher-style checks confirm that diffusion needs at least one variance observable, and sparse σpp “anchors” restore conditioning. We also show that the same PINN can learn timevarying HPZ coefficients. Index Terms—Physics-informed neural networks, quantum Brownian motion, partial observability

II. P ROBLEM F ORMULATION Consider the 2D phase-space state (x, p) with mass m, frequency ω > 0, and damping γ ≥ 0. Define
⊤ s(t) = µx (t), µp (t), σxx (t), σpp (t), σxp (t) . With  A=

0 −mω 2

 1/m , −γ

D=

 Dxx Dxp

 Dxp ⪰ 0, Dpp

the first and second moments obey µ = (µx , µp )⊤ ,  σ Σ̇(t) = AΣ(t) + Σ(t)A⊤ + 2D, Σ = xx σxp µ̇(t) = Aµ(t),

(1) 

σxp . σpp

(2)

Equivalently, I. I NTRODUCTION Quantum Brownian motion (QBM) in the Caldeira–Leggett (CL) model describes a system harmonic oscillator linearly coupled to an Ohmic bath [2]. In the high-temperature, Markovian limit, the Wigner picture yields an Ornstein–Uhlenbeck (OU) Fokker–Planck equation whose first and second moments {µx (t), µp (t), σxx (t), σpp (t), σxp (t)} follow a closed linear ODE in which the drift encodes the damping rate γ and frequency ω, and the diffusion encodes (Dxx , Dxp , Dpp ) [1], [3]. The textbook high-temperature fluctuation–dissipation theorem (FDT) further ties Dpp = 2mγkB T and usually takes Dxp → 0 [4]. A small position diffusion Dxx ≥ 0 is often added to restore complete positivity/Lindblad compatibility [5], [7]. In realistic experiments only a subset of these moments is observable (e.g. position mean and variance), which makes the diffusion sector weakly identifiable. Classical estimators (finite differences; Kalman filtering/EM) exist, but they are accurate only when the continuous-time dynamics are exactly discretized, e.g. via the Van Loan block exponential [8], [9]. Our goal is to show that a PINN that operates directly on the moments can incorporate the same physics, remain positive, and perform competitively—especially when channels are missing. In trapped-ion / cavity / levitated settings only position or position + cross-covariance is routinely accessible, while momentum energy readouts are sparse or calibrated. This is the regime our PINN is targeted at, not the fully observed toy CL case.

µ̇x = µp /m,

(3a)

2

(3b)

µ̇p = −mω µx − γµp , σ̇xx = 2σxp /m + 2Dxx , 2

σ̇pp = −2mω σxp − 2γσpp + 2Dpp , 2

σ̇xp = σpp /m − mω σxx − γσxp + Dxp .

(3c) (3d) (3e)

In the high-T CL limit one enforces Dpp = 2mγkB T and takes Dxp ≈ 0; a small Dxx > 0 maintains complete positivity [1], [7]. Observations are noisy and partial: yi = M s(ti ) + εi ,

2 εi ∼ N (0, σread I),

(4)

where M selects, e.g., {µx , σxx , σxp }. For baseline comparison, the exact discrete-time model with sampling step ∆t is µk+1 = Φµk ,

Σk+1 = ΦΣk Φ⊤ + Qd ,

where (Φ, Qd ) are obtained in one pass from the Van Loan block exponential [8]. III. M ETHOD : PARTIAL -M OMENT PINN S We model the five moments as smooth functions of (normalized) time with a small MLP fϕ with tanh activations and Fourier features [15] to help localize ω: (µ̂x , µ̂p , σ̂xx , σ̂pp , σ̂xp ) = fϕ (t). To guarantee PSD we parameterize a Cholesky factor   a(t) 0 L(t) = , σ̂xx = a2 +b2 , σ̂pp = c2 , b(t) c(t)

σ̂xp = bc,

so that Σ̂(t) = L(t)L(t)⊤ ⪰ 0 for all t. This was necessary because the planted ω = 1.0 produces long, low-frequency transients on [0, 20]. A. Physics residual At collocation times C we form the scaled residuals of (3b)–(3e) using autodiff to obtain f˙ϕ (t). Coefficients (γ, ω, Dxx , Dpp , Dxp ) are either global (LTI) or produced by a small 2-layer coefficient head in the Hu–Paz–Zhang (HPZ) time-varying setting. B. Loss Let obs D = {(ti , yi )}N i=1

denote the observed data, and let anchor Nanchor A = {(τj , σpp,j )}j=1

TABLE I T RAINING AND MODEL HYPERPARAMETERS USED IN THE MAIN CONSTRAINED LTI SETTING . Item

Setting

Time normalization Fourier features State MLP width / depth Observed channels Synthetic horizon / grid Noise level Collocation size Optimizer Learning rate Gradient clip Physics weight Anchor count / weight Dxx prior Residual equation weights FDT / Dxp Seeds

t′ = (t − t0 )/(t1 − t0 ) kmax = 3 (sin/cos) 256 × 4 layers (tanh) {µx , σxx , σxp } t ∈ [0, 20], N = 600 σ = 0.01 |C| ≈ 3Ntrain Adam (6000 steps) + L-BFGS (500 iters) 10−3 (Adam) ∥∇∥2 ≤ 5 λphys = 50 Nanchor = 30, λanchor = 0.8 target 10−3 , weight λxx = 1000 (1, 1, 5, 1, 1) hard-FDT, Dxp = 0 (42, 123, 999)

denote the σpp anchor set. We optimize L = λdata Ldata + λphys Lphys + λanchor Lanchor + λxx Lxx + λFDT LFDT + λss Lss . The observed-channel data term is Nobs 1 X 2 Ldata = ∥M fϕ (ti ) − yi ∥2 . |D| i=1

(5)

(6)

The physics term is formed from the scaled residuals of the moment equations in (3), using autodiff to obtain f˙ϕ (t): 5

Lphys =

1 XX wk r̃k (t)2 , |C|

(7)

t∈C k=1

where wk are fixed per-equation weights; in the constrained setting we upweight the σxx equation. To reduce diffusion non-identifiability under partial observability, we additionally use sparse σpp anchors: Nanchor
1 X anchor 2 Lanchor = σ̂pp (τj ) − σpp,j . |A| j=1

(8)

We use a small prior on Dxx : Lxx = (Dxx − 10−3 )2 .

(9)

For the main constrained LTI CL experiments we enforce the high-T relation Dpp = 2mγkB T (10) as a hard constraint with fixed T = 1 and set Dxp = 0. Accordingly, LFDT = 0 in the main constrained setting. If this tie is relaxed, we instead use the soft penalty
2 LFDT = Dpp − 2mγkB T . (11) When enabled, we also use a weak end-of-window steady-state penalty
2
2 (12) Lss = σ̂pp (t1 ) − T + σ̂xx (t1 ) − T /ω 2 . In the reported constrained LTI results, the main setting uses hard-FDT with Dxp = 0, strong σpp anchors, and a small Dxx prior.

Fig. 1. Workflow of the proposed Partial-Moment PINN for CL/HPZ parameter recovery from partial moment observations.

C. Implementation details We normalize time to t′ = (t − t0 )/(t1 − t0 ) and feed t′ through Fourier features with kmax = 3, followed by a 4-layer MLP of width 256 with tanh activations. The network outputs (µ̂x , µ̂p ) and the Cholesky-head parameters used to construct (σ̂xx , σ̂pp , σ̂xp ). For LTI runs, (ω, γ, Dxx , Dpp , Dxp ) are global trainable scalars; for HPZ-style runs, a small two-layer coefficient head maps time to the time-varying coefficients. Residuals are evaluated on a collocation set |C| ≈ 3Ntrain sampled uniformly on the training window, with per-equation scaling to prevent covariance residuals from dominating the mean residuals. The main constrained LTI setting uses the hyperparameters listed in Table I. IV. E XPERIMENTS We evaluate on two synthetic settings. LTI CL. We use planted parameters (γ, ω, Dpp , Dxp , Dxx ) = (0.25, 1.0, 0.5, 0, 10−3 ),

(13)

and integrate the moment ODEs on t ∈ [0, 20] with N = 600 RK4 steps.

HPZ-like. We fix ω and use γ(t) = γ[1 + 0.25 sin(0.35t)] ,

(14)

Dxx (t) = Dxx [1 + 0.05 cos(0.2t)] ,

(15)

with hard FDT

TABLE II PARAMETER ESTIMATES ON LTI CL DATA ( PLANTED (γ, ω, Dpp , Dxx ) = (0.25, 1.0, 0.50, 10−3 )). PINN VALUES ARE AVERAGED OVER THREE SEEDS ; STANDARD DEVIATIONS ARE REPORTED IN THE TEXT WHERE APPLICABLE . Method

Dpp (t) = 2γ(t)T,

T = 1,

m = kB = 1.

(16)

In both settings, only {µx , σxx , σxp } are observed, with i.i.d. Gaussian readout noise σ = 0.01. We use a 70%/30% train/validation split. The 30 anchors were chosen so that at least one σpp value appears every ≈ 0.7 time units on [0, 20], which was the smallest density that removed diffusion degeneracy in our Fisher check.

FD (partial) Kalman–EM (exact) PINN–A (unconstr.) PINN–B (constr.+strong)

(1) A finite-difference (FD) method estimates (γ, ω) from µ̈x + γ µ̇x + ω 2 µx = 0 and infers the diffusion terms from the restricted covariance equations. (2) A Kalman–EM estimator on the latent (x, p) state uses exact (A, Qc ) 7→ (Φ, Qd ) via the Van Loan block exponential and maps back to continuoustime parameters via log(Φ)/∆t and Qd ≈ Qc ∆t [8]–[11]. This gives us a statistically efficient baseline against which to judge the benefit of PINN constraints.

ω

Dpp

Dxx

0.1318 0.4526 0.2207 0.2058

1.7434 0.9879 0.9935 1.0235

— 0.4815 0.1651 0.4115

1.15×10−3 1.21×10−3 2.5×10−3 2.4×10−3

TABLE III P REDICTION AND ROLLOUT RMSE ( MEAN±STD ).

Method PINN–A (unconstr.)

A. Baselines

γ

PINN–B (constr.)

Network RMSE (train / test)

Rollout RMSE (train / test)

0.284 ± 0.006 / 0.283 ± 0.010 0.327 ± 0.070 / 0.396 ± 0.101

0.600 ± 0.053 / 0.768 ± 0.072 0.099 ± 0.068 / 0.110 ± 0.073

B. Metrics For each seed we report parameter estimates and then aggregate (mean ± std). Predictive quality is measured by (i) network RMSE (direct fϕ (t) vs. ground truth) and (ii) rollout RMSE, where learned coefficients are integrated in the moment ODE on the same grid. We also monitor λmin (Σ̂(t)) to ensure PSD. C. Additional diagnostics Beyond the main method comparison, we report two lightweight diagnostics that are most relevant to partial observability: (i) Fisher-conditioning across different observed subsets, and (ii) bootstrap uncertainty for the constrained PINN. A full training-time ablation over all constraints and architectural components is left for future work. V. R ESULTS We report four aspects: parameter recovery, predictive fidelity, identifiability diagnostics, and UQ. All classical baselines use exact Van Loan transitions. A. Parameter estimates (LTI) Table II compares planted parameters with FD, Kalman– EM, and two PINNs: (A) unconstrained and (B) constrained. PINN summaries are based on three random seeds. B. Predictive fidelity (RMSE) Table III shows that the constrained variant pays a small price in direct (network) RMSE but produces far more accurate ODE rollouts.

Fig. 2. PINN predictions and ODE rollouts vs. ground truth for µx , µp , and σxx on the LTI dataset. Solid black: truth. Dashed: direct PINN prediction. Dotted: rollout using learned coefficients. The constrained PINN (green/dotted red) better respects dynamics, especially for σxx .

C. Identifiability diagnostics Finite-difference Fisher analyses show that mean-only observations are extremely ill-conditioned, while adding even one variance channel drastically improves conditioning (Table IV). D. Uncertainty quantification A nonparametric bootstrap with B=10 resamples and a shortened training schedule yields the illustrative intervals in Table V; because this reduced protocol differs from the main setting used in Table II, the resulting intervals should not be

TABLE IV F ISHER CONDITION NUMBER FOR DIFFERENT OBSERVED SUBSETS ( LOWER IS BETTER ). Observed subset

cond(I)

{µx , µp } {µx , µp , σxx } {µx , σxx , σxp } (paper subset) Full {µx , µp , σxx , σpp , σxp }

8.71×1018 4.54×103 2.28×103 9.49×102

TABLE V B OOTSTRAP 95% CI S AND COVERAGE ( CONSTRAINED PINN). Param

CI (95%)

Covers truth?

γ ω Dpp Dxx

[0.2286, 0.2676] [0.8237, 0.9322] [0.4572, 0.5351] [0.0024, 0.0025]

Yes No Yes No

interpreted as confidence intervals around the Table II point estimates. Coverage is good for γ and Dpp but not for ω and Dxx , consistent with partial channels and the positivity nudge on Dxx . E. HPZ (time-varying) summary For the HPZ–hardFDT run the coefficient head learns smooth curves t 7→ γ(t), Dxx (t) and, by construction, Dpp (t) = 2γ(t)T ; the physics residual stays at ≈ 3.6 × 10−3 . Rollout RMSE is within 5% of the LTI constrained case, i.e. the method remains stable when the CL coefficients are slowly time-varying, which is the regime targeted by Hu–Paz–Zhang [3]. VI. C ONCLUSION We presented Partial-Moment PINNs for identifying CL parameters from partially observed moment trajectories. By enforcing PSD via a Cholesky covariance head, hard highT fluctuation–dissipation structure, and sparse σpp anchors, the method remains stable even when only three of the five moments are observed. Relative to the unconstrained PINN, the constrained model substantially improves dynamical fidelity (rollout RMSE ≈ 0.11 vs. ≈ 0.77) while preserving positivity, and remains competitive with exact-discrete Kalman–EM on parameter recovery. Parameter recovery is competitive for (ω, Dpp ), whereas Dxx remains statistically fragile, consistent with the Fisher identifiability analysis. The same design extends to HPZ-style time-varying coefficients. Limitations include that all results are on synthetic CL/HPZ moment traces and rely on known structural assumptions (e.g., the high-T regime and near-zero Dxp ); evaluating robustness to model mismatch is an important next step. C ODE AVAILABILITY Codes and scripts to reproduce the experiments are available at https://github.com/FraQTech/pmpinn4qbm.

ACKNOWLEDGMENT This project was funded by the QuantumAI Lab at Fractal Analytics, India. We acknowledge the use of AI model by OpenAI in assisting the manuscript preparation process of this paper R EFERENCES [1] H.-P. Breuer and F. Petruccione, The Theory of Open Quantum Systems. Oxford, U.K.: Oxford Univ. Press, 2002. [2] A. O. Caldeira and A. J. Leggett, “Path integral approach to quantum Brownian motion,” Ann. Phys., vol. 149, no. 2, pp. 374–456, 1983. [3] B. L. Hu, J. P. Paz, and Y. Zhang, “Quantum Brownian motion in a general environment: Exact master equation with nonlocal dissipation and colored noise,” Phys. Rev. D, vol. 45, pp. 2843–2861, 1992. [4] R. Kubo, “The fluctuation–dissipation theorem,” Rep. Prog. Phys., vol. 29, no. 1, pp. 255–284, 1966. [5] G. Lindblad, “On the generators of quantum dynamical semigroups,” Commun. Math. Phys., vol. 48, pp. 119–130, 1976. [6] V. Gorini, A. Kossakowski, and E. C. G. Sudarshan, “Completely positive dynamical semigroups of N -level systems,” J. Math. Phys., vol. 17, no. 5, pp. 821–825, 1976. [7] K. Jacobs, “Beyond the Brownian-motion master equation: Correcting a widely used model to ensure complete positivity,” arXiv:0807.4211, 2008. [8] C. F. Van Loan, “Computing integrals involving the matrix exponential,” IEEE Trans. Autom. Control, vol. 23, no. 3, pp. 395–404, 1978. [9] R. H. Shumway and D. S. Stoffer, “An approach to time series smoothing and forecasting using the EM algorithm,” J. Amer. Statist. Assoc., vol. 77, no. 377, pp. 159–169, 1982. [10] Z. Ghahramani and G. E. Hinton, “Parameter estimation for linear dynamical systems,” Tech. Rep. CRG-TR-96-2, Univ. of Toronto, 1996. [11] S. Särkkä and A. Solin, Applied Stochastic Differential Equations. Cambridge, U.K.: Cambridge Univ. Press, 2019. [12] M. Raissi, P. Perdikaris, and G. E. Karniadakis, “Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear PDEs,” J. Comput. Phys., vol. 378, pp. 686–707, 2019. [13] L. Lu, X. Meng, Z. Mao, and G. E. Karniadakis, “DeepXDE: A deep learning library for solving differential equations,” SIAM Rev., vol. 63, no. 1, pp. 208–228, 2021. [14] A. D. Jagtap, K. Kawaguchi, and G. E. Karniadakis, “Conservative physics-informed neural networks on discrete domains for conservation laws,” Comput. Methods Appl. Mech. Eng., vol. 365, p. 113028, 2020. [15] M. Tancik, B. Mildenhall, and J. T. Barron, “Fourier features let networks learn high frequency functions in low dimensional domains,” in Advances in Neural Information Processing Systems (NeurIPS), vol. 33, 2020. [16] B. Efron, “Bootstrap methods: Another look at the jackknife,” Ann. Statist., vol. 7, no. 1, pp. 1–26, 1979. [17] T. Kailath, A. H. Sayed, and B. Hassibi, Linear Estimation. Upper Saddle River, NJ, USA: Prentice Hall, 2000.

Related documents

• Equivariant learning of a transferable three-dimensional classical density functional#451282
• Resource-efficient equivariant quantum convolutional neural networks#149188
• Universality of the Brownian net#1110050
• Single-Particle Spectral Estimation#1157531
• Particle GFlowNets: Rethinking Generative Marginalization Models#673530